mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-05 08:07:05 +00:00
fix(vector-store): embed Milvus and Azure AI Search queries through the request executor
Milvus REST and Azure AI Search still embedded the query through the SDK, so a bare Router alias as litellm_embedding_model kept failing after the executor landed for Valkey. Both now share BaseQueryEmbeddingVectorStoreConfig, which embeds through the injected executor, drops the empty litellm_embedding_config requirement, and awaits aembedding on the async path. The Router executor falls back to the SDK for models the Router does not serve, so inline provider configs such as azure/text-embedding-3-large with their own credentials keep working through the proxy. Tests fake OpenAI and Milvus at the HTTP boundary with respx instead of patching litellm.embedding.
This commit is contained in:
mateo-berri
parent
0cc0c47f8b
commit
86c8b93bf7
8 changed files with 700 additions and 271 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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 {})
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
)
|
||||
|
|
|
|||
22
uv.lock
generated
22
uv.lock
generated
|
|
@ -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]]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue