diff --git a/litellm/llms/base_llm/vector_store/transformation.py b/litellm/llms/base_llm/vector_store/transformation.py index 723d1375e1e..02a51a8bace 100644 --- a/litellm/llms/base_llm/vector_store/transformation.py +++ b/litellm/llms/base_llm/vector_store/transformation.py @@ -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 diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 6e5825b5ca7..cc7af345122 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -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): diff --git a/litellm/llms/valkey/vector_stores/transformation.py b/litellm/llms/valkey/vector_stores/transformation.py index 67ce02d4c69..3cbfca0f1a9 100644 --- a/litellm/llms/valkey/vector_stores/transformation.py +++ b/litellm/llms/valkey/vector_stores/transformation.py @@ -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 diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index b5bd5ed3356..2dda8bf722a 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -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" diff --git a/tests/test_litellm/llms/valkey/vector_stores/test_valkey_transformation.py b/tests/test_litellm/llms/valkey/vector_stores/test_valkey_transformation.py index 188a1982b8b..a2ee2c2bdb1 100644 --- a/tests/test_litellm/llms/valkey/vector_stores/test_valkey_transformation.py +++ b/tests/test_litellm/llms/valkey/vector_stores/test_valkey_transformation.py @@ -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"