diff --git a/litellm/llms/azure_ai/vector_stores/transformation.py b/litellm/llms/azure_ai/vector_stores/transformation.py index 5e61d0a1dd9..044b8f5243c 100644 --- a/litellm/llms/azure_ai/vector_stores/transformation.py +++ b/litellm/llms/azure_ai/vector_stores/transformation.py @@ -1,10 +1,13 @@ +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final import httpx -import litellm from litellm.llms.azure.common_utils import BaseAzureLLM -from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig +from litellm.llms.base_llm.vector_store.transformation import ( + BaseQueryEmbeddingVectorStoreConfig, + VectorStoreEmbeddingExecutor, +) from litellm.types.router import GenericLiteLLMParams from litellm.types.vector_stores import ( BaseVectorStoreAuthCredentials, @@ -25,7 +28,7 @@ else: LiteLLMLoggingObj = Any -class AzureAIVectorStoreConfig(BaseVectorStoreConfig, BaseAzureLLM): +class AzureAIVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig, BaseAzureLLM): """ Configuration for Azure AI Search Vector Store @@ -109,82 +112,71 @@ class AzureAIVectorStoreConfig(BaseVectorStoreConfig, BaseAzureLLM): def transform_search_vector_store_request( self, vector_store_id: str, - query: str | list[str], + query: str | Sequence[str], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, api_base: str, litellm_logging_obj: LiteLLMLoggingObj, - litellm_params: dict, - extra_body: dict[str, Any] | None = None, - ) -> tuple[str, dict[str, Any]]: - """ - Transform search request for Azure AI Search API + litellm_params: Mapping[str, object], + extra_body: Mapping[str, object] | None = None, + embedding_executor: VectorStoreEmbeddingExecutor | None = None, + ) -> tuple[str, dict[str, object]]: + query_text: Final = self.query_text(query) + query_vector: Final = self.embed_query(query_text, litellm_params, embedding_executor) + return self._search_request( + vector_store_id, + query_text, + query_vector, + vector_store_search_optional_params, + api_base, + litellm_logging_obj, + litellm_params, + ) - Generates embeddings using litellm.embeddings and constructs Azure AI Search request - """ - # Convert query to string if it's a list - if isinstance(query, list): - query = " ".join(query) + async def atransform_search_vector_store_request( + self, + vector_store_id: str, + query: str | Sequence[str], + vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + api_base: str, + litellm_logging_obj: LiteLLMLoggingObj, + litellm_params: Mapping[str, object], + extra_body: Mapping[str, object] | None = None, + embedding_executor: VectorStoreEmbeddingExecutor | None = None, + ) -> tuple[str, dict[str, object]]: + query_text: Final = self.query_text(query) + query_vector: Final = await self.aembed_query(query_text, litellm_params, embedding_executor) + return self._search_request( + vector_store_id, + query_text, + query_vector, + vector_store_search_optional_params, + api_base, + litellm_logging_obj, + litellm_params, + ) - # Get embedding model from litellm_params (required) - embedding_model: Final = litellm_params.get("litellm_embedding_model") - if not embedding_model: - raise ValueError( - "embedding_model is required in litellm_params for Azure AI Search. " - "Example: litellm_params['embedding_model'] = 'azure/text-embedding-3-large'" - ) - - embedding_config: Final = litellm_params.get("litellm_embedding_config", {}) - if not embedding_config: - raise ValueError( - "embedding_config is required in litellm_params for Azure AI Search. " - "Example: litellm_params['embedding_config'] = {'api_base': 'https://krris-mh44uf7y-eastus2.cognitiveservices.azure.com/', 'api_key': 'os.environ/AZURE_API_KEY', 'api_version': '2025-09-01'}" - ) - - # Get vector field name (defaults to contentVector) + @staticmethod + def _search_request( + vector_store_id: str, + query_text: str, + query_vector: Sequence[float], + vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + api_base: str, + litellm_logging_obj: LiteLLMLoggingObj, + litellm_params: Mapping[str, object], + ) -> tuple[str, dict[str, object]]: vector_field: Final = litellm_params.get("azure_search_vector_field", "contentVector") - - # Get top_k (number of results to return) top_k: Final = vector_store_search_optional_params.get("top_k", 10) - - # Generate embedding for the query using litellm.embeddings - try: - embedding_response: Final = litellm.embedding( - model=embedding_model, - input=[query], - **embedding_config, - ) - query_vector: Final = embedding_response.data[0]["embedding"] - except Exception as e: - raise Exception(f"Failed to generate embedding for query: {e}") - - # Azure AI Search endpoint for search - index_name: Final = vector_store_id # vector_store_id is the index name - url: Final = f"{api_base}/indexes/{index_name}/docs/search?api-version=2024-07-01" - - # Build the request body for Azure AI Search with vector search - request_body: Final = { - "search": "*", # Get all documents (filtered by vector similarity) - "vectorQueries": [ - { - "vector": query_vector, - "fields": vector_field, - "kind": "vector", - "k": top_k, # Number of nearest neighbors to return - } - ], - "select": "id,content", # Fields to return (customize based on schema) + litellm_logging_obj.model_call_details["input"] = query_text + litellm_logging_obj.model_call_details["embedding_model"] = litellm_params.get("litellm_embedding_model") + litellm_logging_obj.model_call_details["top_k"] = top_k + return f"{api_base}/indexes/{vector_store_id}/docs/search?api-version=2024-07-01", { + "search": "*", + "vectorQueries": [{"vector": query_vector, "fields": vector_field, "kind": "vector", "k": top_k}], + "select": "id,content", "top": top_k, } - ######################################################### - # Update logging object with details of the request - ######################################################### - litellm_logging_obj.model_call_details["input"] = query - litellm_logging_obj.model_call_details["embedding_model"] = embedding_model - litellm_logging_obj.model_call_details["top_k"] = top_k - - return url, request_body - def transform_search_vector_store_response( self, response: httpx.Response, litellm_logging_obj: LiteLLMLoggingObj ) -> VectorStoreSearchResponse: diff --git a/litellm/llms/base_llm/vector_store/transformation.py b/litellm/llms/base_llm/vector_store/transformation.py index e9c925448a8..95863266bf7 100644 --- a/litellm/llms/base_llm/vector_store/transformation.py +++ b/litellm/llms/base_llm/vector_store/transformation.py @@ -3,9 +3,11 @@ from __future__ import annotations from abc import abstractmethod from collections.abc import Mapping, Sequence from dataclasses import dataclass +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, NoReturn, Protocol, runtime_checkable import httpx +from pydantic import TypeAdapter from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import EmbeddingResponse @@ -65,7 +67,7 @@ class RouterVectorStoreEmbeddingExecutor: router: Router metadata: Mapping[str, object] - def _embedding_kwargs(self, configuration: Mapping[str, object]) -> dict[str, object]: + def _embedding_kwargs(self, configuration: Mapping[str, object]) -> Mapping[str, object]: configured_metadata: Final = configuration.get("metadata") metadata: Final = { **(configured_metadata if isinstance(configured_metadata, Mapping) else {}), @@ -76,18 +78,32 @@ class RouterVectorStoreEmbeddingExecutor: "metadata": metadata, } + def _router_serves(self, model: str) -> bool: + team_id: Final = self.metadata.get("user_api_key_team_id") + resolved: Final = self.router.resolved_litellm_models(model, team_id if isinstance(team_id, str) else None) + deployment_models: Final = ( + deployment.get("litellm_params", {}).get("model") for deployment in self.router.get_model_list() or () + ) + return bool(resolved) or model in deployment_models + def embed(self, model: str, query: str, configuration: Mapping[str, object]) -> EmbeddingResponse: + embedding_kwargs: Final = self._embedding_kwargs(configuration) + if not self._router_serves(model): + return LiteLLMVectorStoreEmbeddingExecutor().embed(model, query, embedding_kwargs) return self.router.embedding( # pyright: ignore[reportUnknownMemberType] # Router embedding input retains a legacy untyped list model=model, input=[query], # mutable-ok: Router embedding requires a mutable input list - **self._embedding_kwargs(configuration), # pyright: ignore[reportArgumentType] # provider kwargs are intentionally dynamic + **embedding_kwargs, # pyright: ignore[reportArgumentType] # provider kwargs are intentionally dynamic ) async def aembed(self, model: str, query: str, configuration: Mapping[str, object]) -> EmbeddingResponse: + embedding_kwargs: Final = self._embedding_kwargs(configuration) + if not self._router_serves(model): + return await LiteLLMVectorStoreEmbeddingExecutor().aembed(model, query, embedding_kwargs) return await self.router.aembedding( # pyright: ignore[reportUnknownMemberType] # Router embedding input retains a legacy untyped list model=model, input=[query], # mutable-ok: Router embedding requires a mutable input list - **self._embedding_kwargs(configuration), # pyright: ignore[reportArgumentType] # provider kwargs are intentionally dynamic + **embedding_kwargs, # pyright: ignore[reportArgumentType] # provider kwargs are intentionally dynamic ) @@ -221,6 +237,103 @@ class BaseVectorStoreConfig: return 0.0, 0.0 +_EMPTY_EMBEDDING_CONFIGURATION: Final[Mapping[str, object]] = MappingProxyType({}) +_QUERY_VECTOR: Final = TypeAdapter(list[float]) + + +class BaseQueryEmbeddingVectorStoreConfig(BaseVectorStoreConfig): + @abstractmethod + def transform_search_vector_store_request( + self, + vector_store_id: str, + query: str | Sequence[str], + vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + api_base: str, + litellm_logging_obj: LiteLLMLoggingObj, + litellm_params: Mapping[str, object], + extra_body: Mapping[str, object] | None = None, + embedding_executor: VectorStoreEmbeddingExecutor | None = None, + ) -> tuple[str, dict[str, object]]: + pass + + async def atransform_search_vector_store_request( + self, + vector_store_id: str, + query: str | Sequence[str], + vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + api_base: str, + litellm_logging_obj: LiteLLMLoggingObj, + litellm_params: Mapping[str, object], + extra_body: Mapping[str, object] | None = None, + embedding_executor: VectorStoreEmbeddingExecutor | None = None, + ) -> tuple[str, dict[str, object]]: + return self.transform_search_vector_store_request( + vector_store_id=vector_store_id, + query=query, + vector_store_search_optional_params=vector_store_search_optional_params, + api_base=api_base, + litellm_logging_obj=litellm_logging_obj, + litellm_params=litellm_params, + extra_body=extra_body, + embedding_executor=embedding_executor, + ) + + @staticmethod + def query_text(query: str | Sequence[str]) -> str: + return query if isinstance(query, str) else " ".join(query) + + @staticmethod + def query_embedding_model(litellm_params: Mapping[str, object]) -> str: + embedding_model: Final = litellm_params.get("litellm_embedding_model") + if isinstance(embedding_model, str) and embedding_model: + return embedding_model + raise ValueError( + "litellm_embedding_model is required in litellm_params for this vector store. " + "Example: litellm_params['litellm_embedding_model'] = 'openai/text-embedding-3-small'" + ) + + @staticmethod + def query_embedding_configuration(litellm_params: Mapping[str, object]) -> Mapping[str, object]: + configuration: Final = litellm_params.get("litellm_embedding_config") + if isinstance(configuration, Mapping): + return {str(key): value for key, value in configuration.items()} # pyright: ignore[reportUnknownVariableType, reportUnknownArgumentType] # litellm_params is an untyped dict, keys are re-validated as str here + return _EMPTY_EMBEDDING_CONFIGURATION + + def embed_query( + self, + query_text: str, + litellm_params: Mapping[str, object], + embedding_executor: VectorStoreEmbeddingExecutor | None, + ) -> Sequence[float]: + model: Final = self.query_embedding_model(litellm_params) + configuration: Final = self.query_embedding_configuration(litellm_params) + executor: Final = ( + embedding_executor if embedding_executor is not None else LiteLLMVectorStoreEmbeddingExecutor() + ) + try: + response: Final = executor.embed(model, query_text, configuration) + except Exception as e: + raise Exception(f"Failed to generate embedding for query: {e}") + return _QUERY_VECTOR.validate_python(response.data[0]["embedding"]) # pyright: ignore[reportUnknownMemberType] # EmbeddingResponse.data is an untyped list, the vector is validated here + + async def aembed_query( + self, + query_text: str, + litellm_params: Mapping[str, object], + embedding_executor: VectorStoreEmbeddingExecutor | None, + ) -> Sequence[float]: + model: Final = self.query_embedding_model(litellm_params) + configuration: Final = self.query_embedding_configuration(litellm_params) + executor: Final = ( + embedding_executor if embedding_executor is not None else LiteLLMVectorStoreEmbeddingExecutor() + ) + try: + response: Final = await executor.aembed(model, query_text, configuration) + except Exception as e: + raise Exception(f"Failed to generate embedding for query: {e}") + return _QUERY_VECTOR.validate_python(response.data[0]["embedding"]) # pyright: ignore[reportUnknownMemberType] # EmbeddingResponse.data is an untyped list, the vector is validated here + + class BaseDirectVectorStoreConfig(BaseVectorStoreConfig): """ Base config for vector store providers whose datastore has no HTTP API diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 118656b81a2..c0ef7456680 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -69,6 +69,7 @@ from litellm.llms.base_llm.skills.transformation import BaseSkillsAPIConfig from litellm.llms.base_llm.text_to_speech.transformation import BaseTextToSpeechConfig from litellm.llms.base_llm.vector_store.transformation import ( BaseDirectVectorStoreConfig, + BaseQueryEmbeddingVectorStoreConfig, BaseVectorStoreConfig, VectorStoreEmbeddingExecutor, ) @@ -9728,8 +9729,7 @@ class BaseLLMHTTPHandler: litellm_params=dict(litellm_params), ) - # Check if provider has async transform method - if hasattr(vector_store_provider_config, "atransform_search_vector_store_request"): + if isinstance(vector_store_provider_config, BaseQueryEmbeddingVectorStoreConfig): ( url, request_body, @@ -9741,12 +9741,13 @@ class BaseLLMHTTPHandler: litellm_logging_obj=logging_obj, litellm_params=dict(litellm_params), extra_body=extra_body, + embedding_executor=embedding_executor, ) else: ( url, request_body, - ) = vector_store_provider_config.transform_search_vector_store_request( + ) = await vector_store_provider_config.atransform_search_vector_store_request( vector_store_id=vector_store_id, query=query, vector_store_search_optional_params=vector_store_search_optional_params, @@ -9857,18 +9858,33 @@ class BaseLLMHTTPHandler: litellm_params=dict(litellm_params), ) - ( - url, - request_body, - ) = vector_store_provider_config.transform_search_vector_store_request( - vector_store_id=vector_store_id, - query=query, - vector_store_search_optional_params=vector_store_search_optional_params, - api_base=api_base, - litellm_logging_obj=logging_obj, - litellm_params=dict(litellm_params), - extra_body=extra_body, - ) + if isinstance(vector_store_provider_config, BaseQueryEmbeddingVectorStoreConfig): + ( + url, + request_body, + ) = vector_store_provider_config.transform_search_vector_store_request( + vector_store_id=vector_store_id, + query=query, + vector_store_search_optional_params=vector_store_search_optional_params, + api_base=api_base, + litellm_logging_obj=logging_obj, + litellm_params=dict(litellm_params), + extra_body=extra_body, + embedding_executor=embedding_executor, + ) + else: + ( + url, + request_body, + ) = vector_store_provider_config.transform_search_vector_store_request( + vector_store_id=vector_store_id, + query=query, + vector_store_search_optional_params=vector_store_search_optional_params, + api_base=api_base, + litellm_logging_obj=logging_obj, + litellm_params=dict(litellm_params), + extra_body=extra_body, + ) all_optional_params: Final[dict[str, object]] = dict(litellm_params) all_optional_params.update(vector_store_search_optional_params or {}) diff --git a/litellm/llms/milvus/vector_stores/transformation.py b/litellm/llms/milvus/vector_stores/transformation.py index 34f0cd854c4..b0291c692d5 100644 --- a/litellm/llms/milvus/vector_stores/transformation.py +++ b/litellm/llms/milvus/vector_stores/transformation.py @@ -1,9 +1,12 @@ +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final import httpx -import litellm -from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig +from litellm.llms.base_llm.vector_store.transformation import ( + BaseQueryEmbeddingVectorStoreConfig, + VectorStoreEmbeddingExecutor, +) from litellm.secret_managers.main import get_secret_str from litellm.types.router import GenericLiteLLMParams from litellm.types.vector_stores import ( @@ -36,7 +39,7 @@ MILVUS_OPTIONAL_PARAMS: Final = { } -class MilvusVectorStoreConfig(BaseVectorStoreConfig): +class MilvusVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig): """ Configuration for Milvus Vector Store @@ -117,77 +120,77 @@ class MilvusVectorStoreConfig(BaseVectorStoreConfig): def transform_search_vector_store_request( self, vector_store_id: str, - query: str | list[str], + query: str | Sequence[str], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, api_base: str, litellm_logging_obj: LiteLLMLoggingObj, - litellm_params: dict, - extra_body: dict[str, Any] | None = None, - ) -> tuple[str, dict[str, Any]]: - """ - Transform search request for Azure AI Search API + litellm_params: Mapping[str, object], + extra_body: Mapping[str, object] | None = None, + embedding_executor: VectorStoreEmbeddingExecutor | None = None, + ) -> tuple[str, dict[str, object]]: + query_text: Final = self.query_text(query) + query_vector: Final = self.embed_query(query_text, litellm_params, embedding_executor) + return self._search_request( + vector_store_id, + query_text, + query_vector, + vector_store_search_optional_params, + api_base, + litellm_logging_obj, + litellm_params, + ) - Generates embeddings using litellm.embeddings and constructs Azure AI Search request - """ - # Convert query to string if it's a list - if isinstance(query, list): - query = " ".join(query) + async def atransform_search_vector_store_request( + self, + vector_store_id: str, + query: str | Sequence[str], + vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + api_base: str, + litellm_logging_obj: LiteLLMLoggingObj, + litellm_params: Mapping[str, object], + extra_body: Mapping[str, object] | None = None, + embedding_executor: VectorStoreEmbeddingExecutor | None = None, + ) -> tuple[str, dict[str, object]]: + query_text: Final = self.query_text(query) + query_vector: Final = await self.aembed_query(query_text, litellm_params, embedding_executor) + return self._search_request( + vector_store_id, + query_text, + query_vector, + vector_store_search_optional_params, + api_base, + litellm_logging_obj, + litellm_params, + ) - # Get embedding model from litellm_params (required) - embedding_model: Final = litellm_params.get("litellm_embedding_model") - if not embedding_model: - raise ValueError( - "embedding_model is required in litellm_params for Milvus. You can call any litellm embedding model." - "Example: litellm_params['embedding_model'] = 'azure/text-embedding-3-large'" + @staticmethod + def _search_request( + vector_store_id: str, + query_text: str, + query_vector: Sequence[float], + vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + api_base: str, + litellm_logging_obj: LiteLLMLoggingObj, + litellm_params: Mapping[str, object], + ) -> tuple[str, dict[str, object]]: + scope: Final = { + key: value + for key, value in ( + ("dbName", litellm_params.get("milvus_db_name")), + ("partitionNames", litellm_params.get("milvus_partition_names")), ) - - embedding_config: Final = litellm_params.get("litellm_embedding_config", {}) - if not embedding_config: - raise ValueError( - "embedding_config is required in litellm_params for Milvus. You can call any litellm embedding model." - "Example: litellm_params['embedding_config'] = {'api_base': 'https://krris-mh44uf7y-eastus2.cognitiveservices.azure.com/', 'api_key': 'os.environ/AZURE_API_KEY', 'api_version': '2025-09-01'}" - ) - - # Get top_k (number of results to return) - # Generate embedding for the query using litellm.embeddings - try: - embedding_response: Final = litellm.embedding( - model=embedding_model, - input=[query], - **embedding_config, - ) - query_vector: Final = embedding_response.data[0]["embedding"] - except Exception as e: - raise Exception(f"Failed to generate embedding for query: {e}") - - # Azure AI Search endpoint for search - index_name: Final = vector_store_id # vector_store_id is the index name - url: Final = f"{api_base}/v2/vectordb/entities/search" - - # Build the request body for Azure AI Search with vector search - request_body: Final[dict[str, Any]] = { - "collectionName": index_name, + if value + } + litellm_logging_obj.model_call_details["input"] = query_text + litellm_logging_obj.model_call_details["embedding_model"] = litellm_params.get("litellm_embedding_model") + return f"{api_base}/v2/vectordb/entities/search", { + "collectionName": vector_store_id, "data": [query_vector], "annsField": "book_intro_vector", **vector_store_search_optional_params, + **scope, } - db_name: Final = litellm_params.get("milvus_db_name") - if db_name: - request_body["dbName"] = db_name - - partition_names: Final = litellm_params.get("milvus_partition_names") - if partition_names: - request_body["partitionNames"] = partition_names - - ######################################################### - # Update logging object with details of the request - ######################################################### - litellm_logging_obj.model_call_details["input"] = query - litellm_logging_obj.model_call_details["embedding_model"] = embedding_model - - return url, request_body - def transform_search_vector_store_response( self, response: httpx.Response, litellm_logging_obj: LiteLLMLoggingObj ) -> VectorStoreSearchResponse: diff --git a/tests/router_unit_tests/test_router_embedding_integration.py b/tests/router_unit_tests/test_router_embedding_integration.py index d5e0c750d88..96c6ce3708e 100644 --- a/tests/router_unit_tests/test_router_embedding_integration.py +++ b/tests/router_unit_tests/test_router_embedding_integration.py @@ -5,16 +5,57 @@ These tests simulate real-world scenarios where headers and configuration need to be properly propagated through the router to the LLM API. """ +import json from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest +import respx +import litellm from litellm import Router from litellm.llms.base_llm.vector_store.transformation import ( LiteLLMVectorStoreEmbeddingExecutor, RouterVectorStoreEmbeddingExecutor, ) -from litellm.types.utils import EmbeddingResponse + +QUERY_VECTOR = [0.5, -0.25, 0.125] +OPENAI_EMBEDDINGS_URL = "https://api.openai.com/v1/embeddings" +STORE_EMBEDDINGS_URL = "https://embedding.example/v1/embeddings" + + +def _mock_embedding_route(respx_mock: respx.MockRouter, url: str) -> respx.Route: + return respx_mock.post(url).mock( + return_value=httpx.Response( + 200, + json={ + "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": QUERY_VECTOR}], + "model": "text-embedding-3-small", + "usage": {"prompt_tokens": 2, "total_tokens": 2}, + }, + ) + ) + + +def _sent(route: respx.Route, index: int) -> tuple[str, str, list[str]]: + request = route.calls[index].request + body = json.loads(request.read()) + return request.headers["authorization"], body["model"], body["input"] + + +def _alias_router() -> Router: + return Router( + model_list=[ + { + "model_name": "team-alias", + "litellm_params": { + "model": "openai/text-embedding-3-small", + "api_key": "deployment-key", + }, + } + ] + ) class TestRouterEmbeddingIntegration: @@ -70,48 +111,23 @@ class TestRouterEmbeddingIntegration: ) @pytest.mark.asyncio - async def test_vector_store_embedding_executors_cover_sdk_and_router_paths(self): - response = EmbeddingResponse(data=[{"embedding": [0.1], "index": 0, "object": "embedding"}]) + async def test_vector_store_embedding_executors_cover_sdk_and_router_paths( + self, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch + ): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + openai_route = _mock_embedding_route(respx_mock, OPENAI_EMBEDDINGS_URL) + store_route = _mock_embedding_route(respx_mock, STORE_EMBEDDINGS_URL) sdk_executor = LiteLLMVectorStoreEmbeddingExecutor() - with ( - patch("litellm.embedding", return_value=response) as embedding, - patch("litellm.aembedding", new=AsyncMock(return_value=response)) as aembedding, - ): - assert sdk_executor.embed("openai/model", "sync", {"api_key": "explicit"}) is response - assert await sdk_executor.aembed("openai/model", "async", {"api_key": "explicit"}) is response + sync_response = sdk_executor.embed("openai/text-embedding-3-small", "sync", {"api_key": "explicit"}) + async_response = await sdk_executor.aembed("openai/text-embedding-3-small", "async", {"api_key": "explicit"}) - embedding.assert_called_once_with(model="openai/model", input=["sync"], api_key="explicit") - aembedding.assert_awaited_once_with(model="openai/model", input=["async"], api_key="explicit") + assert sync_response.data[0]["embedding"] == QUERY_VECTOR + assert async_response.data[0]["embedding"] == QUERY_VECTOR + assert _sent(openai_route, 0) == ("Bearer explicit", "text-embedding-3-small", ["sync"]) + assert _sent(openai_route, 1) == ("Bearer explicit", "text-embedding-3-small", ["async"]) - mock_router = MagicMock() - mock_router.embedding.return_value = response - router_executor = RouterVectorStoreEmbeddingExecutor( - router=mock_router, - metadata={"user_api_key_team_id": "team-a"}, - ) - assert router_executor.embed("team-alias", "query", {}) is response - mock_router.embedding.assert_called_once_with( - model="team-alias", - input=["query"], - metadata={"user_api_key_team_id": "team-a"}, - ) - - alias_router = Router( - model_list=[ - { - "model_name": "team-alias", - "litellm_params": { - "model": "openai/text-embedding-3-small", - "api_key": "deployment-key", - }, - } - ] - ) - alias_executor = RouterVectorStoreEmbeddingExecutor( - router=alias_router, - metadata={"user_api_key_team_id": "team-a"}, - ) explicit_config = { "api_base": "https://embedding.example/v1", "api_key": "store-key", @@ -121,29 +137,66 @@ class TestRouterEmbeddingIntegration: }, "model": "untrusted-model", } + mock_router = MagicMock() + mock_router.embedding.return_value = sync_response + router_executor = RouterVectorStoreEmbeddingExecutor( + router=mock_router, + metadata={"user_api_key_team_id": "team-a"}, + ) + assert router_executor.embed("team-alias", "query", explicit_config) is sync_response + mock_router.embedding.assert_called_once_with( + model="team-alias", + input=["query"], + api_base="https://embedding.example/v1", + api_key="store-key", + metadata={"configured": True, "user_api_key_team_id": "team-a"}, + ) - with ( - patch("litellm.embedding", return_value=response) as explicit_embedding, - patch("litellm.aembedding", new=AsyncMock(return_value=response)) as explicit_aembedding, - ): - assert alias_executor.embed("team-alias", "sync query", explicit_config) is response - assert await alias_executor.aembed("team-alias", "async query", explicit_config) is response + alias_executor = RouterVectorStoreEmbeddingExecutor( + router=_alias_router(), + metadata={"user_api_key_team_id": "team-a"}, + ) + sync_alias = alias_executor.embed("team-alias", "sync query", explicit_config) + async_alias = await alias_executor.aembed("team-alias", "async query", explicit_config) - sync_kwargs = explicit_embedding.call_args.kwargs - assert sync_kwargs["model"] == "openai/text-embedding-3-small" - assert sync_kwargs["input"] == ["sync query"] - assert sync_kwargs["api_base"] == "https://embedding.example/v1" - assert sync_kwargs["api_key"] == "store-key" - assert sync_kwargs["metadata"]["configured"] is True - assert sync_kwargs["metadata"]["user_api_key_team_id"] == "team-a" + assert sync_alias.data[0]["embedding"] == QUERY_VECTOR + assert async_alias.data[0]["embedding"] == QUERY_VECTOR + assert openai_route.call_count == 2 + assert _sent(store_route, 0) == ("Bearer store-key", "text-embedding-3-small", ["sync query"]) + assert _sent(store_route, 1) == ("Bearer store-key", "text-embedding-3-small", ["async query"]) - async_kwargs = explicit_aembedding.await_args.kwargs - assert async_kwargs["model"] == "openai/text-embedding-3-small" - assert async_kwargs["input"] == ["async query"] - assert async_kwargs["api_base"] == "https://embedding.example/v1" - assert async_kwargs["api_key"] == "store-key" - assert async_kwargs["metadata"]["configured"] is True - assert async_kwargs["metadata"]["user_api_key_team_id"] == "team-a" + @pytest.mark.asyncio + async def test_router_executor_falls_back_to_sdk_for_models_the_router_does_not_serve( + self, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch + ): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + store_route = _mock_embedding_route(respx_mock, STORE_EMBEDDINGS_URL) + executor = RouterVectorStoreEmbeddingExecutor( + router=_alias_router(), + metadata={"user_api_key_team_id": "team-a"}, + ) + inline_config = {"api_base": "https://embedding.example/v1", "api_key": "store-key"} + + sync_response = executor.embed("openai/text-embedding-3-large", "sync query", inline_config) + async_response = await executor.aembed("openai/text-embedding-3-large", "async query", inline_config) + + assert sync_response.data[0]["embedding"] == QUERY_VECTOR + assert async_response.data[0]["embedding"] == QUERY_VECTOR + assert _sent(store_route, 0) == ("Bearer store-key", "text-embedding-3-large", ["sync query"]) + assert _sent(store_route, 1) == ("Bearer store-key", "text-embedding-3-large", ["async query"]) + + def test_router_executor_routes_deployment_model_names_through_the_router( + self, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch + ): + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + openai_route = _mock_embedding_route(respx_mock, OPENAI_EMBEDDINGS_URL) + executor = RouterVectorStoreEmbeddingExecutor(router=_alias_router(), metadata={}) + + response = executor.embed("openai/text-embedding-3-small", "query", {}) + + assert response.data[0]["embedding"] == QUERY_VECTOR + assert _sent(openai_route, 0) == ("Bearer deployment-key", "text-embedding-3-small", ["query"]) def test_embedding_with_deployment_specific_headers(self): """ @@ -251,9 +304,7 @@ class TestRouterEmbeddingIntegration: router = Router( model_list=model_list, - default_litellm_params={ - "metadata": {"environment": "test", "service": "embedding-service"} - }, + default_litellm_params={"metadata": {"environment": "test", "service": "embedding-service"}}, ) with patch("litellm.embedding") as mock_embedding: @@ -369,9 +420,7 @@ class TestRouterEmbeddingIntegration: # Make multiple calls and verify headers are always present for i in range(5): with patch("litellm.embedding") as mock_embedding: - mock_embedding.return_value = MagicMock( - data=[{"embedding": [0.1, 0.2]}] - ) + mock_embedding.return_value = MagicMock(data=[{"embedding": [0.1, 0.2]}]) router.embedding(model="shared-embedding-model", input=[f"test {i}"]) @@ -456,9 +505,7 @@ class TestRouterEmbeddingIntegration: router = Router( model_list=model_list, - default_litellm_params={ - "headers": {"X-Custom-Azure-Header": "azure-value"} - }, + default_litellm_params={"headers": {"X-Custom-Azure-Header": "azure-value"}}, ) with patch("litellm.embedding") as mock_embedding: diff --git a/tests/vector_store_tests/test_azure_ai_vector_store.py b/tests/vector_store_tests/test_azure_ai_vector_store.py index 58e45f259ab..d1fc8436fc9 100644 --- a/tests/vector_store_tests/test_azure_ai_vector_store.py +++ b/tests/vector_store_tests/test_azure_ai_vector_store.py @@ -1,10 +1,19 @@ -import pytest -import litellm import json import os +from unittest.mock import MagicMock + +import httpx +import pytest +import respx + +import litellm +from litellm.llms.azure_ai.vector_stores.transformation import AzureAIVectorStoreConfig +from litellm.types.utils import EmbeddingResponse +from litellm.vector_stores import ( + asearch as vector_store_asearch, +) from litellm.vector_stores import ( search as vector_store_search, - asearch as vector_store_asearch, ) @@ -30,10 +39,108 @@ async def test_basic_search_vector_store(sync_mode): if sync_mode: response = vector_store_search(query=default_query, **base_request_args) else: - response = await vector_store_asearch( - query=default_query, **base_request_args - ) + response = await vector_store_asearch(query=default_query, **base_request_args) except litellm.InternalServerError: pytest.skip("Skipping test due to litellm.InternalServerError") print("litellm response=", json.dumps(response, indent=4, default=str)) + + +class RecordingEmbeddingExecutor: + def __init__(self, response): + self.response = response + self.calls = [] + + def embed(self, model, query, configuration): + self.calls.append((model, query, dict(configuration))) + return self.response + + async def aembed(self, model, query, configuration): + self.calls.append((model, query, dict(configuration))) + return self.response + + +ALIAS_QUERY_VECTOR = [0.5, -0.25, 0.125] +ALIAS_EMBEDDING_RESPONSE = EmbeddingResponse( + data=[{"embedding": ALIAS_QUERY_VECTOR, "index": 0, "object": "embedding"}] +) +STORE_EMBEDDINGS_URL = "https://embedding.example/v1/embeddings" + + +def _transform_kwargs(executor): + logging_obj = MagicMock() + logging_obj.model_call_details = {} + return { + "vector_store_id": "my-vector-index", + "query": "what is azure search?", + "vector_store_search_optional_params": {"top_k": 2}, + "api_base": "https://azure-kb-search.search.windows.net", + "litellm_logging_obj": logging_obj, + "litellm_params": { + "litellm_embedding_model": "multilingual-e5-large", + "azure_search_vector_field": "embedding", + }, + "embedding_executor": executor, + } + + +@pytest.mark.asyncio +async def test_transform_uses_injected_executor_without_embedding_config(respx_mock: respx.MockRouter): + executor = RecordingEmbeddingExecutor(ALIAS_EMBEDDING_RESPONSE) + config = AzureAIVectorStoreConfig() + transform_kwargs = _transform_kwargs(executor) + + url, sync_body = config.transform_search_vector_store_request(**transform_kwargs) + _, async_body = await config.atransform_search_vector_store_request(**transform_kwargs) + + assert respx_mock.calls.call_count == 0 + assert executor.calls == [("multilingual-e5-large", "what is azure search?", {})] * 2 + assert ( + url == "https://azure-kb-search.search.windows.net/indexes/my-vector-index/docs/search?api-version=2024-07-01" + ) + assert sync_body == async_body + assert sync_body["vectorQueries"] == [ + {"vector": ALIAS_QUERY_VECTOR, "fields": "embedding", "kind": "vector", "k": 2} + ] + assert sync_body["top"] == 2 + logging_details = transform_kwargs["litellm_logging_obj"].model_call_details + assert logging_details["embedding_model"] == "multilingual-e5-large" + assert logging_details["top_k"] == 2 + + +def test_transform_falls_back_to_sdk_embedding_without_executor( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + embedding_route = respx_mock.post(STORE_EMBEDDINGS_URL).mock( + return_value=httpx.Response( + 200, + json={ + "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": ALIAS_QUERY_VECTOR}], + "model": "text-embedding-3-small", + "usage": {"prompt_tokens": 2, "total_tokens": 2}, + }, + ) + ) + transform_kwargs = _transform_kwargs(None) + transform_kwargs["litellm_params"] = { + "litellm_embedding_model": "openai/text-embedding-3-small", + "litellm_embedding_config": {"api_base": "https://embedding.example/v1", "api_key": "store-key"}, + } + + _, body = AzureAIVectorStoreConfig().transform_search_vector_store_request(**transform_kwargs) + + embedding_request = embedding_route.calls.last.request + assert embedding_request.headers["authorization"] == "Bearer store-key" + assert json.loads(embedding_request.read())["input"] == ["what is azure search?"] + assert body["vectorQueries"][0]["vector"] == ALIAS_QUERY_VECTOR + assert body["vectorQueries"][0]["fields"] == "contentVector" + + +def test_transform_requires_embedding_model(): + transform_kwargs = _transform_kwargs(RecordingEmbeddingExecutor(ALIAS_EMBEDDING_RESPONSE)) + transform_kwargs["litellm_params"] = {"litellm_embedding_config": {"api_key": "store-key"}} + + with pytest.raises(ValueError, match="litellm_embedding_model is required"): + AzureAIVectorStoreConfig().transform_search_vector_store_request(**transform_kwargs) diff --git a/tests/vector_store_tests/test_milvus_vector_store.py b/tests/vector_store_tests/test_milvus_vector_store.py index 6627f6006d1..ea3c1883e46 100644 --- a/tests/vector_store_tests/test_milvus_vector_store.py +++ b/tests/vector_store_tests/test_milvus_vector_store.py @@ -3,16 +3,19 @@ Tests for Milvus Vector Store """ import json -import os from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest +import respx import litellm +from litellm import Router +from litellm.llms.milvus.vector_stores.transformation import MilvusVectorStoreConfig +from litellm.types.utils import EmbeddingResponse from litellm.vector_stores import asearch as vector_store_asearch from litellm.vector_stores import search as vector_store_search - # Mock response from actual Milvus API MOCK_MILVUS_SEARCH_RESPONSE = { "code": 0, @@ -98,7 +101,7 @@ class TestMilvusVectorStore: mock_response.json.return_value = MOCK_MILVUS_SEARCH_RESPONSE mock_response.text = json.dumps(MOCK_MILVUS_SEARCH_RESPONSE) - with patch("litellm.embedding") as mock_embedding: + with patch("litellm.aembedding", new_callable=AsyncMock) as mock_embedding: mock_embedding.return_value = MOCK_EMBEDDING_RESPONSE with patch( @@ -147,16 +150,10 @@ class TestMilvusVectorStore: else: # Fallback: check for json kwarg or in args request_data = call_args.kwargs.get("json") - if ( - request_data is None - and len(call_args.args) > 0 - and isinstance(call_args.args[0], dict) - ): + if request_data is None and len(call_args.args) > 0 and isinstance(call_args.args[0], dict): request_data = call_args.args[0] - assert ( - request_data is not None - ), f"Could not extract request data. Call args: {call_args}" + assert request_data is not None, f"Could not extract request data. Call args: {call_args}" print("Request data:", json.dumps(request_data, indent=2, default=str)) # Validate request structure @@ -213,9 +210,7 @@ class TestMilvusVectorStore: with patch("litellm.embedding") as mock_embedding: mock_embedding.return_value = MOCK_EMBEDDING_RESPONSE - with patch( - "litellm.llms.custom_httpx.http_handler.HTTPHandler.post" - ) as mock_post: + with patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post") as mock_post: mock_post.return_value = mock_response # Make the search request @@ -252,16 +247,10 @@ class TestMilvusVectorStore: else: # Fallback: check for json kwarg or in args request_data = call_args.kwargs.get("json") - if ( - request_data is None - and len(call_args.args) > 0 - and isinstance(call_args.args[0], dict) - ): + if request_data is None and len(call_args.args) > 0 and isinstance(call_args.args[0], dict): request_data = call_args.args[0] - assert ( - request_data is not None - ), f"Could not extract request data. Call args: {call_args}" + assert request_data is not None, f"Could not extract request data. Call args: {call_args}" # Validate request structure assert "collectionName" in request_data @@ -316,11 +305,7 @@ class TestMilvusVectorStore: if request_data_str: return json.loads(request_data_str) request_data = call_args.kwargs.get("json") - if ( - request_data is None - and len(call_args.args) > 0 - and isinstance(call_args.args[0], dict) - ): + if request_data is None and len(call_args.args) > 0 and isinstance(call_args.args[0], dict): request_data = call_args.args[0] return request_data @@ -334,9 +319,7 @@ class TestMilvusVectorStore: with patch("litellm.embedding") as mock_embedding: mock_embedding.return_value = MOCK_EMBEDDING_RESPONSE - with patch( - "litellm.llms.custom_httpx.http_handler.HTTPHandler.post" - ) as mock_post: + with patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post") as mock_post: mock_post.return_value = mock_response vector_store_search( @@ -375,9 +358,7 @@ class TestMilvusVectorStore: with patch("litellm.embedding") as mock_embedding: mock_embedding.return_value = MOCK_EMBEDDING_RESPONSE - with patch( - "litellm.llms.custom_httpx.http_handler.HTTPHandler.post" - ) as mock_post: + with patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post") as mock_post: mock_post.return_value = mock_response vector_store_search( @@ -413,9 +394,7 @@ class TestMilvusVectorStore: with patch("litellm.embedding") as mock_embedding: mock_embedding.return_value = MOCK_EMBEDDING_RESPONSE - with patch( - "litellm.llms.custom_httpx.http_handler.HTTPHandler.post" - ) as mock_post: + with patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post") as mock_post: mock_post.return_value = mock_response vector_store_search( @@ -492,3 +471,175 @@ if __name__ == "__main__": test.test_basic_search_with_mock_sync() print("\n✅ All mock tests passed!") + + +class RecordingEmbeddingExecutor: + def __init__(self, response): + self.response = response + self.calls = [] + + def embed(self, model, query, configuration): + self.calls.append((model, query, dict(configuration))) + return self.response + + async def aembed(self, model, query, configuration): + self.calls.append((model, query, dict(configuration))) + return self.response + + +ALIAS_QUERY_VECTOR = [0.5, -0.25, 0.125] +ALIAS_EMBEDDING_RESPONSE = EmbeddingResponse( + data=[{"embedding": ALIAS_QUERY_VECTOR, "index": 0, "object": "embedding"}] +) +OPENAI_EMBEDDINGS_URL = "https://api.openai.com/v1/embeddings" +MILVUS_SEARCH_URL = "https://milvus.example/v2/vectordb/entities/search" +ALIAS_SEARCH_KWARGS = { + "query": "what is machine learning?", + "vector_store_id": "book_2", + "custom_llm_provider": "milvus", + "api_base": "https://milvus.example", + "api_key": "mock_milvus_api_key", + "litellm_embedding_model": "multilingual-e5-large", + "milvus_text_field": "book_intro_text", +} + + +def _alias_router(): + return Router( + model_list=[ + { + "model_name": "multilingual-e5-large", + "litellm_params": { + "model": "openai/text-embedding-3-small", + "api_key": "deployment-key", + }, + } + ] + ) + + +def _mock_embedding_route(respx_mock: respx.MockRouter) -> respx.Route: + return respx_mock.post(OPENAI_EMBEDDINGS_URL).mock( + return_value=httpx.Response( + 200, + json={ + "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": ALIAS_QUERY_VECTOR}], + "model": "text-embedding-3-small", + "usage": {"prompt_tokens": 2, "total_tokens": 2}, + }, + ) + ) + + +def _mock_search_route(respx_mock: respx.MockRouter) -> respx.Route: + return respx_mock.post(MILVUS_SEARCH_URL).mock(return_value=httpx.Response(200, json=MOCK_MILVUS_SEARCH_RESPONSE)) + + +def _assert_alias_resolved(embedding_route: respx.Route, search_route: respx.Route, response): + embedding_request = embedding_route.calls.last.request + assert embedding_request.headers["authorization"] == "Bearer deployment-key" + embedding_body = json.loads(embedding_request.read()) + assert embedding_body["model"] == "text-embedding-3-small" + assert embedding_body["input"] == ["what is machine learning?"] + search_request = search_route.calls.last.request + assert search_request.headers["authorization"] == "Bearer mock_milvus_api_key" + assert json.loads(search_request.read())["data"] == [ALIAS_QUERY_VECTOR] + assert len(response["data"]) == len(MOCK_MILVUS_SEARCH_RESPONSE["data"]) + assert response["data"][0]["content"][0]["text"] == MOCK_MILVUS_SEARCH_RESPONSE["data"][0]["book_intro_text"] + + +def test_router_search_resolves_bare_embedding_alias_sync( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + embedding_route = _mock_embedding_route(respx_mock) + search_route = _mock_search_route(respx_mock) + + response = _alias_router().vector_store_search(**ALIAS_SEARCH_KWARGS) + + _assert_alias_resolved(embedding_route, search_route, response) + + +@pytest.mark.asyncio +async def test_router_search_resolves_bare_embedding_alias_async( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + embedding_route = _mock_embedding_route(respx_mock) + search_route = _mock_search_route(respx_mock) + + response = await _alias_router().avector_store_search(**ALIAS_SEARCH_KWARGS) + + _assert_alias_resolved(embedding_route, search_route, response) + + +@pytest.mark.asyncio +async def test_transform_uses_injected_executor_without_embedding_config(respx_mock: respx.MockRouter): + executor = RecordingEmbeddingExecutor(ALIAS_EMBEDDING_RESPONSE) + config = MilvusVectorStoreConfig() + logging_obj = MagicMock() + logging_obj.model_call_details = {} + transform_kwargs = { + "vector_store_id": "book_2", + "query": ["what is", "milvus?"], + "vector_store_search_optional_params": {"limit": 3}, + "api_base": "https://milvus.example", + "litellm_logging_obj": logging_obj, + "litellm_params": {"litellm_embedding_model": "multilingual-e5-large", "milvus_db_name": "docs"}, + "embedding_executor": executor, + } + + url, sync_body = config.transform_search_vector_store_request(**transform_kwargs) + _, async_body = await config.atransform_search_vector_store_request(**transform_kwargs) + + assert respx_mock.calls.call_count == 0 + assert executor.calls == [("multilingual-e5-large", "what is milvus?", {})] * 2 + assert url == MILVUS_SEARCH_URL + assert sync_body == async_body + assert sync_body == { + "collectionName": "book_2", + "data": [ALIAS_QUERY_VECTOR], + "annsField": "book_intro_vector", + "limit": 3, + "dbName": "docs", + } + assert logging_obj.model_call_details["input"] == "what is milvus?" + assert logging_obj.model_call_details["embedding_model"] == "multilingual-e5-large" + + +def test_transform_falls_back_to_sdk_embedding_without_executor_or_config( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setenv("OPENAI_API_KEY", "env-key") + embedding_route = _mock_embedding_route(respx_mock) + logging_obj = MagicMock() + logging_obj.model_call_details = {} + + _, body = MilvusVectorStoreConfig().transform_search_vector_store_request( + vector_store_id="book_2", + query="q", + vector_store_search_optional_params={}, + api_base="https://milvus.example", + litellm_logging_obj=logging_obj, + litellm_params={"litellm_embedding_model": "openai/text-embedding-3-small"}, + ) + + embedding_request = embedding_route.calls.last.request + assert embedding_request.headers["authorization"] == "Bearer env-key" + assert json.loads(embedding_request.read())["input"] == ["q"] + assert body["data"] == [ALIAS_QUERY_VECTOR] + + +def test_transform_requires_embedding_model(): + with pytest.raises(ValueError, match="litellm_embedding_model is required"): + MilvusVectorStoreConfig().transform_search_vector_store_request( + vector_store_id="book_2", + query="q", + vector_store_search_optional_params={}, + api_base="https://milvus.example", + litellm_logging_obj=MagicMock(), + litellm_params={"litellm_embedding_config": {"api_key": "store-key"}}, + embedding_executor=RecordingEmbeddingExecutor(ALIAS_EMBEDDING_RESPONSE), + ) diff --git a/uv.lock b/uv.lock index 8d886044083..27be919eea1 100644 --- a/uv.lock +++ b/uv.lock @@ -9441,19 +9441,19 @@ wheels = [ [[package]] name = "tornado" -version = "6.5.8" +version = "6.5.7" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/10/d3/343e5bb989d6515b1646cf3d40135d73f3d5e45339bded401b56cdac24dd/tornado-6.5.8.tar.gz", hash = "sha256:9452e1b208a8bd771e2cb1f2ff564985b9b214bdebbe622793e1799e0a6bd23f", size = 520493, upload-time = "2026-08-07T02:12:42.971Z" } +sdist = { url = "https://files.pythonhosted.org/packages/64/24/95ec527ad67b76d59299e5465b3935d05e4294b7e0290a3924b7487df30b/tornado-6.5.7.tar.gz", hash = "sha256:66c513a76cda70d53907bc27cf1447557699c2e95aa48ba27a442ff61c3ddfc2", size = 519252, upload-time = "2026-06-08T17:34:51.232Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/f2/d5/007086fd8df5489338e204f65adce33fd4f21a4999dbb2b9cff2f897b5f4/tornado-6.5.8-cp39-abi3-macosx_10_9_universal2.whl", hash = "sha256:cc6aa787d7cfab7c3d35189dc7a56fbd2399a569624c730c6b55b3d6531d0403", size = 449487, upload-time = "2026-08-07T02:12:28.682Z" }, - { url = "https://files.pythonhosted.org/packages/70/c8/5a24a99495903f594f6a199dd7beead1cbc0a13e2cb9102727bcaaf2a997/tornado-6.5.8-cp39-abi3-macosx_10_9_x86_64.whl", hash = "sha256:9715b5eb79735b2bcd454ce216a9275b7c0470e64ea1bf5742f78b2f72b26eeb", size = 447649, upload-time = "2026-08-07T02:12:30.306Z" }, - { url = "https://files.pythonhosted.org/packages/6e/de/f2e733f386b85962d1b1dc82cd63d169b5b4580062b35397eac9244a41fe/tornado-6.5.8-cp39-abi3-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:547d63f450d570c14fe0e8db2cfb14c9bbd1c2503b4a6612586267955aa47b58", size = 450707, upload-time = "2026-08-07T02:12:31.95Z" }, - { url = "https://files.pythonhosted.org/packages/0b/94/20efeee9a01c141e9ac47c397f81679dfda24b32768fc4fff24e76d36c2c/tornado-6.5.8-cp39-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:7e2360a0ffbe145eca8af0b19cb7203d79b1a98dd4cccdd6b368f6f49c2e3808", size = 451677, upload-time = "2026-08-07T02:12:33.512Z" }, - { url = "https://files.pythonhosted.org/packages/42/ec/a96ccb8ccf0de2b7bc2c5fa1608a4803735018242e90c4882365a9fd418f/tornado-6.5.8-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:5d242290bdf7ab3151bc1065fdd75c0dcc21cbc7b49f22a4c56329c2d6566d22", size = 451510, upload-time = "2026-08-07T02:12:35.346Z" }, - { url = "https://files.pythonhosted.org/packages/29/b5/93185859245ad3f00e62175f29607346788b696369347f0146e0421286bb/tornado-6.5.8-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:7b94ff0e128fe0542f3bd331fb44d06260fc4ac16881545159f34ef08aad4195", size = 450917, upload-time = "2026-08-07T02:12:36.963Z" }, - { url = "https://files.pythonhosted.org/packages/97/cf/fe33cf062834487d34d1559746a4a12521033c22645b6d74d4bca702e018/tornado-6.5.8-cp39-abi3-win32.whl", hash = "sha256:67832909c4779c64942380cb5f044a5c6163d00831472d80e25e115de9917836", size = 451952, upload-time = "2026-08-07T02:12:38.512Z" }, - { url = "https://files.pythonhosted.org/packages/cb/e1/468ad54333e92ccb62627e62cb88e5fc14a2171daa67ed47b1b8542d5b86/tornado-6.5.8-cp39-abi3-win_amd64.whl", hash = "sha256:11881db6b7c168494be2c2d12e65931451bdf7ee718535418ae1d8855dd5a0ee", size = 452391, upload-time = "2026-08-07T02:12:39.971Z" }, - { url = "https://files.pythonhosted.org/packages/ad/3e/cd5e4f06e34cde33b8ef66cf36aa2b5ad46354cc1af7d2136bbe365fee1d/tornado-6.5.8-cp39-abi3-win_arm64.whl", hash = "sha256:68a7468c7e289f8514d7d664101753903217eff1bb6822c6b5994a0b5f5bcb26", size = 451411, upload-time = "2026-08-07T02:12:41.469Z" }, + { url = "https://files.pythonhosted.org/packages/02/dc/c7043cab6fed8ae159fc1923ce829ada35c4dbd797d408a43858ffaf9639/tornado-6.5.7-cp39-abi3-macosx_10_9_universal2.whl", hash = "sha256:148b2eb15c2c765a50796172c1e499649b35f30d2e3c3d3e15913cfa56bfb163", size = 448543, upload-time = "2026-06-08T17:34:38.052Z" }, + { url = "https://files.pythonhosted.org/packages/92/4f/090b1431e5a43df696feceffc268c5383cc079ecb5f08ce58f917109aafe/tornado-6.5.7-cp39-abi3-macosx_10_9_x86_64.whl", hash = "sha256:9da38de27f1da3b78a966f0dae12b5a1ea9afe72ca805d84ff06508272ddf100", size = 446707, upload-time = "2026-06-08T17:34:39.594Z" }, + { url = "https://files.pythonhosted.org/packages/37/d8/ef374952fd5da67d4463122c2b8e5a96536ec10b4b339254c6dcde81d01c/tornado-6.5.7-cp39-abi3-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:8d759e71906ee783f8867b93bf26a265743da4c1e2f4a018464c1ba019862972", size = 449774, upload-time = "2026-06-08T17:34:41.204Z" }, + { url = "https://files.pythonhosted.org/packages/35/37/d434c73f4c6e014b745b9b37085f34f40c022f007efff3d7fe65991899f3/tornado-6.5.7-cp39-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:8a46347a18f23fb92b396beebe0fb78f61dda0cc302445202c16203d8a18848b", size = 450745, upload-time = "2026-06-08T17:34:42.531Z" }, + { url = "https://files.pythonhosted.org/packages/b6/2b/56b9aff361d7f1ab728a805ec7d7ea835f8807afa9f5cc690ea0e630efb9/tornado-6.5.7-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:7778b30bef919231265e91c69963ce0f49a1e9c07ac900bbe75b19ce2575ba92", size = 450578, upload-time = "2026-06-08T17:34:43.787Z" }, + { url = "https://files.pythonhosted.org/packages/02/30/a7444fb23aa76860a14198fab96ac79f1866b0a6e19e26c4381b0938e50f/tornado-6.5.7-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:e726f0c75da7726eec023aa62751ff8878bd2737e34fbdd33b1ae5897d2200f5", size = 449985, upload-time = "2026-06-08T17:34:45.326Z" }, + { url = "https://files.pythonhosted.org/packages/5c/42/5f0e56c01e8d9d36f4e23f367b85ae6cae0c1ecddd5e6977d8388ad27488/tornado-6.5.7-cp39-abi3-win32.whl", hash = "sha256:f8de3bf12d3efdd0cbe7c8887868198f8a91415e3f29fcf258d9b8eb7b1d9ae4", size = 451047, upload-time = "2026-06-08T17:34:46.784Z" }, + { url = "https://files.pythonhosted.org/packages/c9/a4/b393076ffb21b469eec5b328a0534cf03a3b90bfc6b1f09507cdd075d938/tornado-6.5.7-cp39-abi3-win_amd64.whl", hash = "sha256:de942f843533a039ef9fa3d9c88c7cd8a7c94553fb5ad0154270989b3d99a2c4", size = 451485, upload-time = "2026-06-08T17:34:48.248Z" }, + { url = "https://files.pythonhosted.org/packages/71/2e/7b1c769803121b809112cf9a00681c472eae1d80e32d7ec0e0bd61d0d0e1/tornado-6.5.7-cp39-abi3-win_arm64.whl", hash = "sha256:ff934fce95643af5f11efdae618eaa73d469dc588641e5c8d19295a0c65c4796", size = 450506, upload-time = "2026-06-08T17:34:49.702Z" }, ] [[package]]