mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-06 08:16:43 +00:00
fix(vector_stores): keep a lost MongoDB connection retryable and bound the client cache by use
litellm only retries 408, 409, 429 and 5xx, so classifying a dropped connection as a 400 turned one replica set failover into a permanently failed search. It is a 503 now, with the message still naming the misconfigurations that also close a connection. The client cache skipped insertion once it held 32 entries, so any store added after that rebuilt its client on every search, paying an SRV lookup, a TLS handshake and topology discovery each time. It evicts the least recently used entry instead, which only drops the cache's own reference. Also trims the explanatory comments to the one-line form the repo asks for.
This commit is contained in:
Yuneng Jiang
parent
38cd1bff7b
commit
da58c0c6d5
3 changed files with 133 additions and 75 deletions
|
|
@ -1,22 +1,16 @@
|
|||
"""Shared helpers for MongoDB integrations.
|
||||
|
||||
pymongo ships in the optional ``mongodb`` extra, so every import of it is
|
||||
deferred to call time and raises an actionable error when it is absent.
|
||||
|
||||
Clients are cached per connection because building one costs an SRV lookup, a
|
||||
TLS handshake and topology discovery: measured at ~890ms against a remote deployment versus
|
||||
~80ms on a warm client, so a client per search would dominate query latency.
|
||||
"""
|
||||
"""Shared helpers for the MongoDB integrations. pymongo lives in the optional ``mongodb`` extra,
|
||||
so every import of it is deferred to call time."""
|
||||
|
||||
import asyncio
|
||||
import weakref
|
||||
from asyncio import AbstractEventLoop
|
||||
from collections import OrderedDict
|
||||
from collections.abc import Callable, Mapping
|
||||
from dataclasses import dataclass
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final, TypeAlias
|
||||
from typing import TYPE_CHECKING, Final, TypeAlias, TypeVar
|
||||
|
||||
from litellm.exceptions import BadRequestError, Timeout
|
||||
from litellm.exceptions import BadRequestError, ServiceUnavailableError, Timeout
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from pymongo import AsyncMongoClient, MongoClient
|
||||
|
|
@ -30,8 +24,7 @@ MONGODB_PROVIDER: Final = "mongodb"
|
|||
|
||||
|
||||
def config_error(message: str) -> BadRequestError:
|
||||
"""Misconfiguration is the caller's to fix, so it maps to 400 rather than the 500
|
||||
a bare ValueError would become once litellm.exception_type wraps it."""
|
||||
"""400 rather than the 500 a bare ValueError becomes once litellm.exception_type wraps it."""
|
||||
return BadRequestError(message=message, model=None, llm_provider=MONGODB_PROVIDER)
|
||||
|
||||
|
||||
|
|
@ -39,6 +32,11 @@ def timeout_error(message: str) -> Timeout:
|
|||
return Timeout(message=message, model=None, llm_provider=MONGODB_PROVIDER)
|
||||
|
||||
|
||||
def unavailable_error(message: str) -> ServiceUnavailableError:
|
||||
"""litellm only retries 408, 409, 429 and 5xx, so a 400 here would make a failover permanent."""
|
||||
return ServiceUnavailableError(message=message, model=None, llm_provider=MONGODB_PROVIDER)
|
||||
|
||||
|
||||
DEFAULT_CONNECT_TIMEOUT_MS: Final = 10_000
|
||||
DEFAULT_SOCKET_TIMEOUT_MS: Final = 30_000
|
||||
DEFAULT_SERVER_SELECTION_TIMEOUT_MS: Final = 10_000
|
||||
|
|
@ -59,12 +57,26 @@ class MongoClientKey:
|
|||
SyncClientFactory: TypeAlias = Callable[..., "MongoClient"]
|
||||
AsyncClientFactory: TypeAlias = Callable[..., "AsyncMongoClient"]
|
||||
|
||||
_K = TypeVar("_K")
|
||||
_V = TypeVar("_V")
|
||||
|
||||
_AsyncClientCacheKey: TypeAlias = tuple[MongoClientKey, int]
|
||||
# CPython recycles id() aggressively, so the id alone would hand a new loop a closed loop's client
|
||||
_AsyncClientEntry: TypeAlias = tuple["weakref.ref[AbstractEventLoop]", "AsyncMongoClient"]
|
||||
|
||||
_sync_clients: Final[dict[MongoClientKey, "MongoClient"]] = {} # mutable-ok: process-level client cache
|
||||
_async_clients: Final[dict[_AsyncClientCacheKey, _AsyncClientEntry]] = {} # mutable-ok: same cache, per loop
|
||||
_SyncClientCache: TypeAlias = "OrderedDict[MongoClientKey, MongoClient]"
|
||||
_AsyncClientCache: TypeAlias = "OrderedDict[_AsyncClientCacheKey, _AsyncClientEntry]"
|
||||
|
||||
_sync_clients: Final[_SyncClientCache] = OrderedDict() # mutable-ok: process-level client cache
|
||||
_async_clients: Final[_AsyncClientCache] = OrderedDict() # mutable-ok: same cache, per loop
|
||||
|
||||
|
||||
def _store_bounded(cache: "OrderedDict[_K, _V]", cache_key: "_K", value: "_V") -> None:
|
||||
"""Eviction only drops this cache's reference; an in-flight search keeps its client alive."""
|
||||
cache[cache_key] = value # mutable-ok: an LRU cache is mutable state by definition
|
||||
cache.move_to_end(cache_key)
|
||||
while len(cache) > _MAX_CACHED_CLIENTS:
|
||||
cache.popitem(last=False)
|
||||
|
||||
|
||||
def import_sync_mongo_client() -> "type[MongoClient]":
|
||||
|
|
@ -95,22 +107,19 @@ def _client_kwargs(key: MongoClientKey) -> Mapping[str, object]:
|
|||
|
||||
|
||||
def get_sync_client(key: MongoClientKey, client_class: SyncClientFactory | None = None) -> "MongoClient":
|
||||
"""``client_class`` is the injection seam the tests build fake clients through; left unset the
|
||||
real pymongo class is imported at call time, keeping pymongo out of import-time dependencies."""
|
||||
cached: Final = _sync_clients.get(key)
|
||||
if cached is not None:
|
||||
_sync_clients.move_to_end(key)
|
||||
return cached
|
||||
build: Final = client_class if client_class is not None else import_sync_mongo_client()
|
||||
client: Final = build(key.connection_string, **_client_kwargs(key))
|
||||
if len(_sync_clients) < _MAX_CACHED_CLIENTS:
|
||||
_sync_clients[key] = client
|
||||
_store_bounded(_sync_clients, key, client)
|
||||
return client
|
||||
|
||||
|
||||
def _purge_dead_loops() -> None:
|
||||
"""The cached client holds its loop object alive, so a closed loop's entry would otherwise pin
|
||||
that client and its sockets for the life of the process. Callers that run one loop per search
|
||||
(``asyncio.run`` in a script) reach the cap this way and never release what is behind it."""
|
||||
"""A cached client holds its loop alive, so a closed loop's entry would pin that client and its
|
||||
sockets for the life of the process."""
|
||||
for stale in tuple(
|
||||
cache_key
|
||||
for cache_key, (loop_ref, _) in _async_clients.items()
|
||||
|
|
@ -125,12 +134,12 @@ def get_async_client(key: MongoClientKey, client_class: AsyncClientFactory | Non
|
|||
loop_key: Final = (key, id(loop))
|
||||
cached: Final = _async_clients.get(loop_key)
|
||||
if cached is not None and cached[0]() is loop:
|
||||
_async_clients.move_to_end(loop_key)
|
||||
return cached[1]
|
||||
_purge_dead_loops()
|
||||
build: Final = client_class if client_class is not None else import_async_mongo_client()
|
||||
client: Final = build(key.connection_string, **_client_kwargs(key))
|
||||
if len(_async_clients) < _MAX_CACHED_CLIENTS or loop_key in _async_clients:
|
||||
_async_clients[loop_key] = (weakref.ref(loop), client)
|
||||
_store_bounded(_async_clients, loop_key, (weakref.ref(loop), client))
|
||||
return client
|
||||
|
||||
|
||||
|
|
@ -157,9 +166,8 @@ def _index_hint(index_name: str, database: str, collection: str) -> str:
|
|||
|
||||
|
||||
def missing_index_error(index_name: str, database: str, collection: str) -> BadRequestError:
|
||||
"""$vectorSearch against a missing index, database or collection returns zero documents
|
||||
instead of failing, so an empty result set is checked against the index catalogue and
|
||||
turned into this rather than being reported as 'no matches'."""
|
||||
"""$vectorSearch against a missing index, database or collection returns zero documents rather
|
||||
than failing, so an empty result set is checked against the catalogue and reported as this."""
|
||||
return config_error(
|
||||
f"{_index_hint(index_name, database, collection)} A vector search against a database, "
|
||||
"collection or index that does not exist returns no results rather than an error, so this "
|
||||
|
|
@ -175,10 +183,7 @@ def index_not_ready_error(index_name: str, database: str, collection: str, statu
|
|||
|
||||
|
||||
def translate_mongo_error(error: Exception, index_name: str, database: str, collection: str) -> Exception:
|
||||
"""Turn a driver failure into a message that names the misconfiguration, never a silent empty result.
|
||||
|
||||
Returns the exception to raise so callers keep the original as ``__cause__``.
|
||||
"""
|
||||
"""Returns the exception to raise, so callers keep the driver error as ``__cause__``."""
|
||||
try:
|
||||
from pymongo.errors import (
|
||||
ConfigurationError,
|
||||
|
|
@ -205,14 +210,16 @@ def translate_mongo_error(error: Exception, index_name: str, database: str, coll
|
|||
f"The MongoDB vector search against '{database}.{collection}' timed out before returning. "
|
||||
f"Driver detail: {error}"
|
||||
)
|
||||
# ServerSelectionTimeoutError and NetworkTimeout both sit under ConnectionFailure, so this
|
||||
# only sees what those two branches left: a dropped or refused connection
|
||||
# ServerSelectionTimeoutError and NetworkTimeout also subclass ConnectionFailure, so this only
|
||||
# sees what those branches left
|
||||
if isinstance(error, ConnectionFailure):
|
||||
return config_error(
|
||||
f"The connection to '{database}.{collection}' was refused or dropped. On Atlas this is "
|
||||
"usually a connection string with no username and password, or a TLS failure, so confirm "
|
||||
"the URI is the one Atlas shows under Connect, Drivers. On a self-managed deployment, check "
|
||||
f"that mongod is listening on the host and port in the URI. Driver detail: {error}"
|
||||
return unavailable_error(
|
||||
f"The connection to '{database}.{collection}' was dropped or refused. That is usually a "
|
||||
"replica set failover or a restarted node, so the search is worth retrying. If it keeps "
|
||||
"happening: on Atlas the usual cause is a connection string with no username and password, "
|
||||
"or a TLS failure, so confirm the URI is the one Atlas shows under Connect, Drivers; on a "
|
||||
"self-managed deployment, check that mongod is listening on the host and port in the URI. "
|
||||
f"Driver detail: {error}"
|
||||
)
|
||||
if isinstance(error, OperationFailure):
|
||||
code: Final = error.code
|
||||
|
|
@ -267,16 +274,14 @@ def translate_mongo_error(error: Exception, index_name: str, database: str, coll
|
|||
)
|
||||
if isinstance(error, InvalidOperation):
|
||||
return config_error(f"The MongoDB client was already closed or is unusable. Driver detail: {error}")
|
||||
# A tlsCAFile or tlsCertificateKeyFile the process cannot open raises OSError from the TLS setup
|
||||
# rather than a PyMongoError, and those options are how self-managed deployments present a private CA
|
||||
# An unreadable tlsCAFile or tlsCertificateKeyFile raises OSError, not a PyMongoError
|
||||
if isinstance(error, OSError) and error.filename:
|
||||
return config_error(
|
||||
f"'{error.filename}', named by a TLS option in mongodb_connection_string, could not be read. "
|
||||
"Check that tlsCAFile and tlsCertificateKeyFile point at files this process can open; inside "
|
||||
f"a container that is the path in the container, not on the host. Driver detail: {error}"
|
||||
)
|
||||
# pymongo raises a plain ValueError, not a PyMongoError, for an unusable port, which an unescaped
|
||||
# ':' in a password also produces, and which would otherwise reach the caller as a 500
|
||||
# pymongo raises a plain ValueError, not a PyMongoError, for an unusable port
|
||||
if isinstance(error, ValueError):
|
||||
return config_error(
|
||||
"The host and port in mongodb_connection_string could not be parsed. If the port is a "
|
||||
|
|
|
|||
|
|
@ -1,15 +1,5 @@
|
|||
"""MongoDB vector store provider, for Atlas and self-managed deployments alike.
|
||||
|
||||
MongoDB Vector Search has no HTTP query API (the Data API and HTTPS Endpoints are
|
||||
end-of-life), so this config extends BaseDirectVectorStoreConfig and runs the
|
||||
``$vectorSearch`` aggregation itself through pymongo instead of shaping an httpx
|
||||
request. mongod serves that stage identically whether mongot runs under Atlas or
|
||||
beside a self-managed deployment, so one code path covers both.
|
||||
|
||||
``vector_store_id`` is the search index name, matching the Valkey provider
|
||||
where the id names the index; the database and collection it covers come from
|
||||
litellm_params.
|
||||
"""
|
||||
"""MongoDB Vector Search has no HTTP query API, so this is a direct provider that runs the
|
||||
``$vectorSearch`` aggregation through pymongo. ``vector_store_id`` is the search index name."""
|
||||
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from types import MappingProxyType
|
||||
|
|
@ -159,9 +149,8 @@ class MongoDBVectorStoreConfig(BaseDirectVectorStoreConfig):
|
|||
|
||||
@staticmethod
|
||||
def _reject_unknown_params(litellm_params: Mapping[str, object]) -> None:
|
||||
"""The params model ignores unrelated keys because litellm_params carries plenty of them,
|
||||
which would otherwise turn a mistyped mongodb_collection into 'mongodb_collection is
|
||||
required' pointing at a key the reader can see they have set."""
|
||||
"""Without this a mistyped mongodb_collection reads as 'mongodb_collection is required',
|
||||
naming a key the reader can see they have set."""
|
||||
unknown: Final = sorted(
|
||||
key for key in litellm_params if key.startswith(_MONGODB_PARAM_PREFIX) and key not in _KNOWN_MONGODB_PARAMS
|
||||
)
|
||||
|
|
@ -268,8 +257,7 @@ class MongoDBVectorStoreConfig(BaseDirectVectorStoreConfig):
|
|||
|
||||
@classmethod
|
||||
def _field_value(cls, document: Mapping[str, object], dotted_path: str) -> str | None:
|
||||
"""None means the path is absent from the document, which is what separates a
|
||||
mistyped mongodb_text_field from a document whose text is genuinely empty."""
|
||||
"""None means absent, which is what separates a mistyped field from genuinely empty text."""
|
||||
head, _, rest = dotted_path.partition(".")
|
||||
if head not in document:
|
||||
return None
|
||||
|
|
@ -297,9 +285,8 @@ class MongoDBVectorStoreConfig(BaseDirectVectorStoreConfig):
|
|||
def _raise_for_missing_text_field(
|
||||
cls, documents: Sequence[Mapping[str, object]], text_field: str, database: str, collection: str
|
||||
) -> None:
|
||||
"""$vectorSearch happily matches documents that carry no text at all, so a mistyped
|
||||
mongodb_text_field returns well-scored results whose content is empty and feeds an empty
|
||||
context to the model. Every matched document lacking the field is the misconfiguration."""
|
||||
"""$vectorSearch matches documents carrying no text, so a mistyped mongodb_text_field
|
||||
returns well-scored results with empty content instead of failing."""
|
||||
if documents and all(cls._field_value(document, text_field) is None for document in documents):
|
||||
raise config_error(
|
||||
f"None of the {len(documents)} matched documents in '{database}.{collection}' has a "
|
||||
|
|
@ -323,9 +310,8 @@ class MongoDBVectorStoreConfig(BaseDirectVectorStoreConfig):
|
|||
def _raise_for_unusable_index(
|
||||
catalogue: Sequence[Mapping[str, object]], index_name: str, database: str, collection: str
|
||||
) -> None:
|
||||
"""An empty result set is ambiguous: mongod returns zero documents both for a query that
|
||||
genuinely matched nothing and for a missing database, collection or index. Only the second
|
||||
is a misconfiguration, so the index catalogue decides which one happened."""
|
||||
"""mongod returns zero documents both for a query that matched nothing and for a missing
|
||||
database, collection or index, so the catalogue decides which one happened."""
|
||||
if not catalogue:
|
||||
raise missing_index_error(index_name, database, collection)
|
||||
entry: Final = catalogue[0]
|
||||
|
|
|
|||
|
|
@ -8,10 +8,12 @@ from unittest.mock import MagicMock, patch
|
|||
import httpx
|
||||
import pytest
|
||||
|
||||
from litellm.exceptions import BadRequestError, Timeout
|
||||
import litellm
|
||||
from litellm.exceptions import BadRequestError, ServiceUnavailableError, Timeout
|
||||
from litellm.llms.mongodb.common_utils import (
|
||||
_MAX_CACHED_CLIENTS,
|
||||
_async_clients,
|
||||
_sync_clients,
|
||||
MongoClientKey,
|
||||
index_not_ready_error,
|
||||
missing_index_error,
|
||||
|
|
@ -643,6 +645,37 @@ class TestClientCache:
|
|||
assert first.connection_string == CONNECTION_STRING
|
||||
|
||||
|
||||
def _fill_cache(self):
|
||||
for slot in range(_MAX_CACHED_CLIENTS):
|
||||
get_sync_client(self._key(f"mongodb://cold-{slot}:27017"), RecordingClient)
|
||||
|
||||
def test_a_store_added_after_the_cache_filled_is_still_cached(self):
|
||||
"""Rebuilding a client costs an SRV lookup, a TLS handshake and topology discovery, so a
|
||||
store that misses the cache on every single search pays that on every search."""
|
||||
self._fill_cache()
|
||||
latecomer = self._key("mongodb://latecomer:27017")
|
||||
|
||||
first = get_sync_client(latecomer, RecordingClient)
|
||||
|
||||
assert get_sync_client(latecomer, RecordingClient) is first
|
||||
|
||||
def test_the_cache_evicts_the_least_recently_used_client(self):
|
||||
self._fill_cache()
|
||||
oldest = self._key("mongodb://cold-0:27017")
|
||||
newest = self._key(f"mongodb://cold-{_MAX_CACHED_CLIENTS - 1}:27017")
|
||||
kept = get_sync_client(newest, RecordingClient)
|
||||
|
||||
get_sync_client(self._key("mongodb://latecomer:27017"), RecordingClient)
|
||||
|
||||
assert get_sync_client(newest, RecordingClient) is kept
|
||||
assert oldest not in _sync_clients
|
||||
|
||||
def test_the_cache_never_grows_past_its_cap(self):
|
||||
for slot in range(_MAX_CACHED_CLIENTS * 3):
|
||||
get_sync_client(self._key(f"mongodb://host-{slot}:27017"), RecordingClient)
|
||||
|
||||
assert len(_sync_clients) == _MAX_CACHED_CLIENTS
|
||||
|
||||
def test_a_new_loop_never_inherits_a_closed_loop_client(self):
|
||||
"""CPython recycles id() so aggressively that a fresh event loop almost always lands on
|
||||
the id of one already collected: measured at 37 of 40 rounds. Keying the cache on the id
|
||||
|
|
@ -757,17 +790,51 @@ class TestErrorTranslation:
|
|||
|
||||
assert "rejected the credentials" in str(translated)
|
||||
|
||||
def test_a_dropped_connection_is_a_400_not_an_unhandled_driver_error(self):
|
||||
"""AutoReconnect sits under ConnectionFailure alongside the two timeout classes, and Atlas
|
||||
answers a URI with no credentials by closing the connection rather than failing auth. Left
|
||||
untranslated it is not a litellm exception type, so it reaches the caller as a 500."""
|
||||
def test_a_dropped_connection_stays_retryable(self):
|
||||
"""A replica set failover reaches the driver as AutoReconnect. litellm only retries 408,
|
||||
409, 429 and 5xx, so classifying it as a client error would turn one failover into a
|
||||
permanently failed search."""
|
||||
from pymongo.errors import AutoReconnect
|
||||
|
||||
translated = self._translate(AutoReconnect("connection closed"))
|
||||
|
||||
assert litellm._should_retry(translated.status_code)
|
||||
assert "dropped or refused" in str(translated)
|
||||
|
||||
def test_a_dropped_connection_still_names_the_misconfigurations_behind_it(self):
|
||||
"""Atlas answers a URI with no credentials by closing the connection rather than failing
|
||||
auth, so the retryable message still has to name that."""
|
||||
from pymongo.errors import AutoReconnect
|
||||
|
||||
translated = self._translate(AutoReconnect("connection closed"))
|
||||
|
||||
assert isinstance(translated, BadRequestError)
|
||||
assert "refused or dropped" in str(translated)
|
||||
assert "no username and password" in str(translated)
|
||||
assert "mongod is listening" in str(translated)
|
||||
|
||||
def test_the_retryable_classification_survives_the_public_sdk_error_wrapper(self):
|
||||
"""litellm.exception_type only passes its own exception types through; anything else becomes
|
||||
an APIConnectionError and a 500, which would drop the retryable classification."""
|
||||
from pymongo.errors import AutoReconnect
|
||||
|
||||
translated = self._translate(AutoReconnect("connection closed"))
|
||||
|
||||
wrapped = litellm.exception_type(
|
||||
model=None,
|
||||
original_exception=translated,
|
||||
custom_llm_provider="mongodb",
|
||||
completion_kwargs={},
|
||||
extra_kwargs={},
|
||||
)
|
||||
|
||||
assert isinstance(wrapped, ServiceUnavailableError)
|
||||
assert litellm._should_retry(wrapped.status_code)
|
||||
|
||||
def test_a_pool_wait_queue_timeout_stays_retryable(self):
|
||||
from pymongo.errors import WaitQueueTimeoutError
|
||||
|
||||
translated = self._translate(WaitQueueTimeoutError("timed out waiting for a connection"))
|
||||
|
||||
assert litellm._should_retry(translated.status_code)
|
||||
|
||||
def test_server_selection_timeout_still_wins_over_the_connection_branch(self):
|
||||
from pymongo.errors import ServerSelectionTimeoutError
|
||||
|
|
@ -775,7 +842,7 @@ class TestErrorTranslation:
|
|||
translated = self._translate(ServerSelectionTimeoutError("no servers"))
|
||||
|
||||
assert isinstance(translated, Timeout)
|
||||
assert "refused or dropped" not in str(translated)
|
||||
assert "dropped or refused" not in str(translated)
|
||||
|
||||
def test_network_timeout_still_wins_over_the_connection_branch(self):
|
||||
from pymongo.errors import NetworkTimeout
|
||||
|
|
@ -783,7 +850,7 @@ class TestErrorTranslation:
|
|||
translated = self._translate(NetworkTimeout("socket timed out"))
|
||||
|
||||
assert isinstance(translated, Timeout)
|
||||
assert "refused or dropped" not in str(translated)
|
||||
assert "dropped or refused" not in str(translated)
|
||||
|
||||
def test_an_unescaped_password_character_is_a_400_not_a_500(self):
|
||||
"""pymongo's URI parser raises a plain ValueError, not a PyMongoError, for an unusable port,
|
||||
|
|
@ -1239,7 +1306,7 @@ class TestSelfManagedDeploymentsAreFirstClass:
|
|||
|
||||
config = self._config_that_fails_to_connect(ConnectionFailure("connection closed"))
|
||||
|
||||
with pytest.raises(BadRequestError) as excinfo:
|
||||
with pytest.raises(ServiceUnavailableError) as excinfo:
|
||||
_search(config)
|
||||
|
||||
assert "self-managed" in str(excinfo.value)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue