scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-08 22:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 78

refactor(vector_stores): route MongoDB query embeddings through the shared executor

Browse source 
The base vector store interface grew an embedding_executor argument, and
litellm.vector_stores.search now always passes one. MongoDB still carried its
own embedding_fn/aembedding_fn constructor seam, so every search through the
public entry point failed with an unexpected keyword argument.

Drop the local seam in favour of the shared executor: one path instead of two,
and the unit tests now drive the same seam production uses.
This commit is contained in:
Yuneng Jiang 2026-09-04 14:20:28 -07:00
parent 431579dc16
commit 4774a426c5
2 changed files with 76 additions and 39 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

37
litellm/llms/mongodb/vector_stores/transformation.py
View file

@ -11,15 +11,18 @@ where the id names the index; the database and collection it covers come from
litellm_params.
"""
from collections.abc import Awaitable, Callable, Mapping, Sequence
from collections.abc import Callable, Mapping, Sequence
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, NoReturn
import httpx
from pydantic import BaseModel, ConfigDict
import litellm
from litellm.llms.base_llm.vector_store.transformation import BaseDirectVectorStoreConfig
from litellm.llms.base_llm.vector_store.transformation import (
BaseDirectVectorStoreConfig,
LiteLLMVectorStoreEmbeddingExecutor,
VectorStoreEmbeddingExecutor,
)
from litellm.llms.mongodb.common_utils import (
DEFAULT_CONNECT_TIMEOUT_MS,
DEFAULT_SERVER_SELECTION_TIMEOUT_MS,
@ -139,17 +142,13 @@ _KNOWN_MONGODB_PARAMS: Final = frozenset(
class MongoDBVectorStoreConfig(BaseDirectVectorStoreConfig):
def __init__(
self,
embedding_fn: Callable[..., EmbeddingResponse] | None = None,
aembedding_fn: Callable[..., Awaitable[EmbeddingResponse]] | None = None,
embedding_executor: VectorStoreEmbeddingExecutor | None = None,
sync_client_factory: Callable[[MongoClientKey], object] | None = None,
async_client_factory: Callable[[MongoClientKey], object] | None = None,
) -> None:
super().__init__()
self.embedding_fn: Final[Callable[..., EmbeddingResponse]] = (
embedding_fn if embedding_fn is not None else litellm.embedding
)
self.aembedding_fn: Final[Callable[..., Awaitable[EmbeddingResponse]]] = (
aembedding_fn if aembedding_fn is not None else litellm.aembedding
self.embedding_executor: Final[VectorStoreEmbeddingExecutor] = (
embedding_executor if embedding_executor is not None else LiteLLMVectorStoreEmbeddingExecutor()
)
self.sync_client_factory: Final[Callable[[MongoClientKey], object]] = (
sync_client_factory if sync_client_factory is not None else get_sync_client
@ -350,6 +349,7 @@ class MongoDBVectorStoreConfig(BaseDirectVectorStoreConfig):
vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams,
litellm_logging_obj: "LiteLLMLoggingObj",
litellm_params: Mapping[str, object],
embedding_executor: VectorStoreEmbeddingExecutor | None = None,
timeout: float | httpx.Timeout | None = None,
) -> VectorStoreSearchResponse:
self._reject_unknown_params(litellm_params)
@ -359,10 +359,10 @@ class MongoDBVectorStoreConfig(BaseDirectVectorStoreConfig):
database: Final = params.require_database()
collection: Final = params.require_collection()
embedding_response: Final = self.embedding_fn(
model=params.require_embedding_model(),
input=[query_text], # mutable-ok: litellm.embedding's input contract is a list
**(params.litellm_embedding_config or _EMPTY_EMBEDDING_CONFIG),
embedding_response: Final = (embedding_executor or self.embedding_executor).embed(
params.require_embedding_model(),
query_text,
params.litellm_embedding_config or _EMPTY_EMBEDDING_CONFIG,
)
pipeline: Final = self._pipeline(
vector_store_id, self._embedding_vector(embedding_response), params, vector_store_search_optional_params
@ -392,6 +392,7 @@ class MongoDBVectorStoreConfig(BaseDirectVectorStoreConfig):
vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams,
litellm_logging_obj: "LiteLLMLoggingObj",
litellm_params: Mapping[str, object],
embedding_executor: VectorStoreEmbeddingExecutor | None = None,
timeout: float | httpx.Timeout | None = None,
) -> VectorStoreSearchResponse:
self._reject_unknown_params(litellm_params)
@ -401,10 +402,10 @@ class MongoDBVectorStoreConfig(BaseDirectVectorStoreConfig):
database: Final = params.require_database()
collection: Final = params.require_collection()
embedding_response: Final = await self.aembedding_fn(
model=params.require_embedding_model(),
input=[query_text], # mutable-ok: litellm.embedding's input contract is a list
**(params.litellm_embedding_config or _EMPTY_EMBEDDING_CONFIG),
embedding_response: Final = await (embedding_executor or self.embedding_executor).aembed(
params.require_embedding_model(),
query_text,
params.litellm_embedding_config or _EMPTY_EMBEDDING_CONFIG,
)
pipeline: Final = self._pipeline(
vector_store_id, self._embedding_vector(embedding_response), params, vector_store_search_optional_params

78
tests/test_litellm/llms/mongodb/vector_stores/test_mongodb_transformation.py
View file

@ -111,27 +111,27 @@ class FakeClient:
return self.database
class FakeEmbeddingFn:
class FakeEmbeddingExecutor:
def __init__(self, embedding):
self.embedding = embedding
self.captured_kwargs = None
self.captured = None
def __call__(self, **kwargs):
self.captured_kwargs = kwargs
def _respond(self, model, query, configuration):
self.captured = SimpleNamespace(model=model, query=query, configuration=configuration)
return SimpleNamespace(data=[{"embedding": self.embedding}] if self.embedding is not None else [])
def embed(self, model, query, configuration):
return self._respond(model, query, configuration)
class FakeAsyncEmbeddingFn(FakeEmbeddingFn):
async def __call__(self, **kwargs):
self.captured_kwargs = kwargs
return SimpleNamespace(data=[{"embedding": self.embedding}] if self.embedding is not None else [])
async def aembed(self, model, query, configuration):
return self._respond(model, query, configuration)
def _config(documents=(), embedding=(0.1, 0.2, 0.3), error=None, search_indexes=None):
collection = FakeCollection(list(documents), error, search_indexes)
client = FakeClient(collection)
config = MongoDBVectorStoreConfig(
embedding_fn=FakeEmbeddingFn(list(embedding) if embedding is not None else None),
embedding_executor=FakeEmbeddingExecutor(list(embedding) if embedding is not None else None),
sync_client_factory=lambda key: client,
)
return config, client, collection
@ -141,7 +141,7 @@ def _async_config(documents=(), embedding=(0.1, 0.2, 0.3), error=None, search_in
collection = FakeAsyncCollection(list(documents), error, search_indexes)
client = FakeClient(collection)
config = MongoDBVectorStoreConfig(
aembedding_fn=FakeAsyncEmbeddingFn(list(embedding) if embedding is not None else None),
embedding_executor=FakeEmbeddingExecutor(list(embedding) if embedding is not None else None),
async_client_factory=lambda key: client,
)
return config, client, collection
@ -252,22 +252,20 @@ def test_num_candidates_below_the_limit_or_above_the_ceiling_is_rejected(configu
def test_list_query_is_joined_into_one_embedding_input():
config, _, _ = _config()
embedding_fn = config.embedding_fn
_search(config, query=["deep", "space", "rescue"])
assert embedding_fn.captured_kwargs["input"] == ["deep space rescue"]
assert config.embedding_executor.captured.query == "deep space rescue"
def test_embedding_config_is_expanded_into_the_embedding_call():
config, _, _ = _config()
embedding_fn = config.embedding_fn
_search(config, litellm_params={"litellm_embedding_config": {"api_base": "https://example.test", "timeout": 7}})
assert embedding_fn.captured_kwargs["api_base"] == "https://example.test"
assert embedding_fn.captured_kwargs["timeout"] == 7
assert embedding_fn.captured_kwargs["model"] == "openai/text-embedding-ada-002"
captured = config.embedding_executor.captured
assert captured.configuration == {"api_base": "https://example.test", "timeout": 7}
assert captured.model == "openai/text-embedding-ada-002"
def test_response_maps_documents_to_openai_shaped_results():
@ -530,7 +528,7 @@ def test_search_fails_when_the_embedding_model_returns_nothing():
def test_validation_runs_before_any_connection_is_opened():
opened = []
config = MongoDBVectorStoreConfig(
embedding_fn=FakeEmbeddingFn([0.1]),
embedding_executor=FakeEmbeddingExecutor([0.1]),
sync_client_factory=lambda key: opened.append(key) or FakeClient(FakeCollection([])),
)
@ -973,7 +971,7 @@ class TestEmptyResultsAreDisambiguated:
collection = ExplodingCollection([], None, [])
config = MongoDBVectorStoreConfig(
embedding_fn=FakeEmbeddingFn([0.1]),
embedding_executor=FakeEmbeddingExecutor([0.1]),
sync_client_factory=lambda key: FakeClient(collection),
)
@ -1157,7 +1155,7 @@ class TestClientConstructionFailures:
raise error
return MongoDBVectorStoreConfig(
embedding_fn=FakeEmbeddingFn([0.1, 0.2, 0.3]), sync_client_factory=factory
embedding_executor=FakeEmbeddingExecutor([0.1, 0.2, 0.3]), sync_client_factory=factory
)
def _async_config_that_fails_to_connect(self, error):
@ -1165,7 +1163,7 @@ class TestClientConstructionFailures:
raise error
return MongoDBVectorStoreConfig(
aembedding_fn=FakeAsyncEmbeddingFn([0.1, 0.2, 0.3]), async_client_factory=factory
embedding_executor=FakeEmbeddingExecutor([0.1, 0.2, 0.3]), async_client_factory=factory
)
def test_a_malformed_uri_is_a_bad_request_not_a_500(self):
@ -1215,7 +1213,7 @@ class TestSelfManagedDeploymentsAreFirstClass:
raise error
return MongoDBVectorStoreConfig(
embedding_fn=FakeEmbeddingFn([0.1, 0.2, 0.3]), sync_client_factory=factory
embedding_executor=FakeEmbeddingExecutor([0.1, 0.2, 0.3]), sync_client_factory=factory
)
def test_a_plain_mongodb_uri_without_srv_or_credentials_is_accepted(self):
@ -1405,3 +1403,41 @@ class TestUnreadableTlsFilesAreDiagnosed:
translated = translate_mongo_error(OSError("socket hung up"), index_name=INDEX, database="db", collection="c")
assert not isinstance(translated, BadRequestError)
class TestTheCallerSuppliedEmbeddingExecutorIsUsed:
"""litellm.vector_stores.search always hands a direct provider an embedding_executor, so the
provider has to accept it and route the query through it rather than its own default."""
def test_the_supplied_executor_produces_the_query_vector(self):
config, _, collection = _config(embedding=(0.9, 0.9, 0.9), search_indexes=READY_INDEX)
caller = FakeEmbeddingExecutor([0.4, 0.5, 0.6])
config.execute_search_vector_store_request(
vector_store_id=INDEX,
query="a lone astronaut",
vector_store_search_optional_params={},
litellm_logging_obj=MagicMock(),
litellm_params=BASE_PARAMS,
embedding_executor=caller,
)
assert caller.captured.query == "a lone astronaut"
assert _stage(collection, "$vectorSearch")["queryVector"] == (0.4, 0.5, 0.6)
@pytest.mark.asyncio
async def test_the_supplied_executor_produces_the_query_vector_on_the_async_path(self):
config, _, collection = _async_config(embedding=(0.9, 0.9, 0.9), search_indexes=READY_INDEX)
caller = FakeEmbeddingExecutor([0.4, 0.5, 0.6])
await config.aexecute_search_vector_store_request(
vector_store_id=INDEX,
query="a lone astronaut",
vector_store_search_optional_params={},
litellm_logging_obj=MagicMock(),
litellm_params=BASE_PARAMS,
embedding_executor=caller,
)
assert caller.captured.query == "a lone astronaut"
assert _stage(collection, "$vectorSearch")["queryVector"] == (0.4, 0.5, 0.6)