scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-08 08:08:37 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 198

fix(vector_stores): join list queries and bound valkey socket timeouts

Browse source 
Review feedback: multi-string queries are now space-joined like every
other embedding-based provider instead of dropping all but the first,
and the request timeout is threaded through the direct vector store
interface into bounded socket_connect_timeout / socket_timeout values
on both redis clients so an unreachable Valkey host cannot pin proxy
workers until the OS TCP timeout.
This commit is contained in:
ryan-crabbe-berri 2026-08-15 03:24:38 -07:00
parent 10a95d0dd3
commit 27450db07e
5 changed files with 58 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm/llms/base_llm/vector_store/transformation.py
View file

@ -172,6 +172,7 @@ class BaseDirectVectorStoreConfig(BaseVectorStoreConfig):
vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams,
litellm_logging_obj: LiteLLMLoggingObj,
litellm_params: Mapping[str, object],
timeout: float | httpx.Timeout | None = None,
) -> VectorStoreSearchResponse:
pass
@ -183,6 +184,7 @@ class BaseDirectVectorStoreConfig(BaseVectorStoreConfig):
vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams,
litellm_logging_obj: LiteLLMLoggingObj,
litellm_params: Mapping[str, object],
timeout: float | httpx.Timeout | None = None,
) -> VectorStoreSearchResponse:
pass

2
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -9429,6 +9429,7 @@ class BaseLLMHTTPHandler:
vector_store_search_optional_params=vector_store_search_optional_params,
litellm_logging_obj=logging_obj,
litellm_params=dict(litellm_params), # mutable-ok: snapshot GenericLiteLLMParams into the Mapping shape
timeout=timeout,
)
if client is None or not isinstance(client, AsyncHTTPHandler):
@ -9559,6 +9560,7 @@ class BaseLLMHTTPHandler:
vector_store_search_optional_params=vector_store_search_optional_params,
litellm_logging_obj=logging_obj,
litellm_params=dict(litellm_params), # mutable-ok: snapshot GenericLiteLLMParams into the Mapping shape
timeout=timeout,
)
if client is None or not isinstance(client, HTTPHandler):

31
litellm/llms/valkey/vector_stores/transformation.py
View file

@ -36,6 +36,8 @@ if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
DEFAULT_VALKEY_PORT: Final = 6379
DEFAULT_SOCKET_CONNECT_TIMEOUT_SECONDS: Final = 5.0
DEFAULT_SOCKET_TIMEOUT_SECONDS: Final = 30.0
DEFAULT_MAX_NUM_RESULTS: Final = 10
MIN_MAX_NUM_RESULTS: Final = 1
MAX_MAX_NUM_RESULTS: Final = 50
@ -138,7 +140,18 @@ class ValkeyVectorStoreConfig(BaseDirectVectorStoreConfig):
return query
if not query:
raise ValueError("query must not be empty")
return query[0]
return " ".join(query)
@staticmethod
def _socket_timeouts(timeout: float | httpx.Timeout | None) -> tuple[float, float]:
if isinstance(timeout, httpx.Timeout):
return (
timeout.connect or DEFAULT_SOCKET_CONNECT_TIMEOUT_SECONDS,
timeout.read or DEFAULT_SOCKET_TIMEOUT_SECONDS,
)
if timeout is not None:
return (min(float(timeout), DEFAULT_SOCKET_CONNECT_TIMEOUT_SECONDS), float(timeout))
return (DEFAULT_SOCKET_CONNECT_TIMEOUT_SECONDS, DEFAULT_SOCKET_TIMEOUT_SECONDS)
@staticmethod
def _knn_limit(vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams) -> int:
@ -200,6 +213,7 @@ class ValkeyVectorStoreConfig(BaseDirectVectorStoreConfig):
vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams,
litellm_logging_obj: "LiteLLMLoggingObj",
litellm_params: Mapping[str, object],
timeout: float | httpx.Timeout | None = None,
) -> VectorStoreSearchResponse:
params: Final = _ValkeySearchParams.model_validate(litellm_params)
query_text: Final = self._query_text(query)
@ -219,7 +233,12 @@ class ValkeyVectorStoreConfig(BaseDirectVectorStoreConfig):
raw: Final = self.sync_client.ft(vector_store_id).search(knn, query_params=vec_params)
return self._to_response(raw, query_text, params.text_field)
client: Final = _import_sync_redis().from_url(params.connection_url())
connect_timeout, op_timeout = self._socket_timeouts(timeout)
client: Final = _import_sync_redis().from_url(
params.connection_url(),
socket_connect_timeout=connect_timeout,
socket_timeout=op_timeout,
)
try:
raw_result: Final = client.ft(vector_store_id).search(knn, query_params=vec_params)
return self._to_response(raw_result, query_text, params.text_field)
@ -233,6 +252,7 @@ class ValkeyVectorStoreConfig(BaseDirectVectorStoreConfig):
vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams,
litellm_logging_obj: "LiteLLMLoggingObj",
litellm_params: Mapping[str, object],
timeout: float | httpx.Timeout | None = None,
) -> VectorStoreSearchResponse:
params: Final = _ValkeySearchParams.model_validate(litellm_params)
query_text: Final = self._query_text(query)
@ -254,7 +274,12 @@ class ValkeyVectorStoreConfig(BaseDirectVectorStoreConfig):
)
return self._to_response(raw, query_text, params.text_field)
client: Final = _import_async_redis().from_url(params.connection_url())
connect_timeout, op_timeout = self._socket_timeouts(timeout)
client: Final = _import_async_redis().from_url(
params.connection_url(),
socket_connect_timeout=connect_timeout,
socket_timeout=op_timeout,
)
try:
raw_result: Final = await client.ft(vector_store_id).search( # pyright: ignore[reportGeneralTypeIssues] # types-redis 4.6 stubs shadow redis 5.3.1 and type the async client's ft() as the sync Search, so search() returns a non-awaitable Result; it is a coroutine at runtime
knn, query_params=vec_params

4
tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py
View file

@ -2110,6 +2110,7 @@ def test_vector_store_search_handler_direct_config_sync_skips_http():
custom_llm_provider="valkey",
litellm_params=GenericLiteLLMParams(valkey_host="localhost"),
logging_obj=logging_obj,
timeout=12.5,
_is_async=False,
)
@ -2119,6 +2120,7 @@ def test_vector_store_search_handler_direct_config_sync_skips_http():
call = config.sync_calls[0]
assert call["vector_store_id"] == "vs_direct"
assert call["query"] == "q"
assert call["timeout"] == 12.5
assert call["vector_store_search_optional_params"] == {"max_num_results": 4}
assert isinstance(call["litellm_params"], dict)
assert call["litellm_params"]["valkey_host"] == "localhost"
@ -2143,6 +2145,7 @@ async def test_vector_store_search_handler_direct_config_async_skips_http():
custom_llm_provider="valkey",
litellm_params=GenericLiteLLMParams(valkey_host="localhost"),
logging_obj=logging_obj,
timeout=7.0,
_is_async=True,
)
@ -2151,6 +2154,7 @@ async def test_vector_store_search_handler_direct_config_async_skips_http():
assert len(config.async_calls) == 1
assert config.async_calls[0]["query"] == ["q1", "q2"]
assert config.async_calls[0]["litellm_params"]["valkey_host"] == "localhost"
assert config.async_calls[0]["timeout"] == 7.0
pre_call_args = logging_obj.pre_call.call_args.kwargs["additional_args"]
assert pre_call_args["query"] == ["q1", "q2"]
assert pre_call_args["vector_store_id"] == "vs_direct"

28
tests/test_litellm/llms/valkey/vector_stores/test_valkey_transformation.py
View file

@ -5,6 +5,7 @@ from typing import Final
from unittest.mock import MagicMock, patch
from urllib.parse import unquote, urlsplit
import httpx
import pytest
from litellm.llms.valkey.vector_stores.transformation import (
@ -138,14 +139,29 @@ def test_sync_search_maps_response_with_inverted_score_sorted_best_first():
assert response["data"][1]["file_id"] == "doc:2"
def test_sync_search_list_query_takes_first_element():
def test_sync_search_list_query_joins_all_elements():
embedding_fn = FakeEmbeddingFn([1.0])
config = ValkeyVectorStoreConfig(sync_client=FakeRedis(), embedding_fn=embedding_fn)
response = _search(config, query=["first query", "second query"])
assert embedding_fn.captured_kwargs["input"] == ["first query"]
assert response["search_query"] == "first query"
assert embedding_fn.captured_kwargs["input"] == ["first query second query"]
assert response["search_query"] == "first query second query"
def test_socket_timeouts_default_to_bounded_values():
assert ValkeyVectorStoreConfig._socket_timeouts(None) == (5.0, 30.0)
def test_socket_timeouts_derive_from_numeric_request_timeout():
assert ValkeyVectorStoreConfig._socket_timeouts(2.0) == (2.0, 2.0)
assert ValkeyVectorStoreConfig._socket_timeouts(120.0) == (5.0, 120.0)
def test_socket_timeouts_derive_from_httpx_timeout():
timeout = httpx.Timeout(connect=3.0, read=7.0, write=1.0, pool=1.0)
assert ValkeyVectorStoreConfig._socket_timeouts(timeout) == (3.0, 7.0)
def test_sync_search_expands_embedding_config_into_kwargs():
@ -324,7 +340,7 @@ async def test_async_search_builds_knn_query_and_maps_response():
response = await config.aexecute_search_vector_store_request(
vector_store_id="my_index",
query=["async query", "ignored"],
query=["async query", "part two"],
vector_store_search_optional_params={"max_num_results": 3},
litellm_logging_obj=MagicMock(),
litellm_params={
@ -338,10 +354,10 @@ async def test_async_search_builds_knn_query_and_maps_response():
assert client.index.searched_query_params == {"vec": struct.pack("<2f", 0.5, 0.5)}
assert aembedding_fn.captured_kwargs == {
"model": "openai/text-embedding-3-small",
"input": ["async query"],
"input": ["async query part two"],
"api_key": "sk-async",
}
assert response["search_query"] == "async query"
assert response["search_query"] == "async query part two"
assert response["data"][0]["score"] == pytest.approx(0.9)
assert response["data"][0]["content"] == [{"text": "async hit", "type": "text"}]
assert response["data"][0]["file_id"] == "doc:9"