mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-15 23:31:29 +00:00
fix(vector-stores): surface retrieval failures to the API caller
A vector store search that fails is swallowed by the pre-call hook, so the request goes to the model with an un-augmented prompt and the caller gets a 200 answering from the model's own knowledge with no way to tell the knowledge base was skipped. Failed searches now ride the same channel their successes already use: a vector_store_search_failures entry on provider_specific_fields naming the store id, provider, and error. That is additive and always on. For callers who would rather fail than answer ungrounded, litellm_settings vector_store_search_failure_mode: error raises VectorStoreSearchError (400) instead; the default stays annotate, today's permissive behavior. The hook's outer catch-all also now names the requested vector store ids in its log line, and only wraps the augmentation itself, so the fail-closed raise is not swallowed by it.
This commit is contained in:
mateo-berri
parent
b503bcabea
commit
445ccc14b2
5 changed files with 460 additions and 158 deletions
|
|
@ -1360,6 +1360,7 @@ from .exceptions import (
|
|||
InvalidRequestError,
|
||||
BadRequestError,
|
||||
ImageFetchError,
|
||||
VectorStoreSearchError,
|
||||
NotFoundError,
|
||||
PermissionDeniedError,
|
||||
RateLimitError,
|
||||
|
|
@ -1461,9 +1462,11 @@ from .vector_stores.vector_store_registry import (
|
|||
VectorStoreRegistry,
|
||||
VectorStoreIndexRegistry,
|
||||
)
|
||||
from .types.vector_stores import VectorStoreSearchFailureMode
|
||||
|
||||
vector_store_registry: Optional[VectorStoreRegistry] = None
|
||||
vector_store_index_registry: Optional[VectorStoreIndexRegistry] = None
|
||||
vector_store_search_failure_mode: VectorStoreSearchFailureMode = "annotate"
|
||||
|
||||
### RAG ###
|
||||
from . import rag
|
||||
|
|
|
|||
|
|
@ -10,12 +10,14 @@
|
|||
## LiteLLM versions of the OpenAI Exception Types
|
||||
|
||||
import enum
|
||||
from collections.abc import Sequence
|
||||
from typing import Any, Final
|
||||
|
||||
import httpx
|
||||
import openai
|
||||
|
||||
from litellm.types.utils import LiteLLMCommonStrings
|
||||
from litellm.types.vector_stores import VectorStoreSearchFailure
|
||||
|
||||
|
||||
class RateLimitErrorCategory(str, enum.Enum):
|
||||
|
|
@ -288,6 +290,29 @@ class ImageFetchError(BadRequestError):
|
|||
)
|
||||
|
||||
|
||||
VECTOR_STORE_SEARCH_FAILED_CODE: Final = "vector_store_search_failed"
|
||||
|
||||
|
||||
class VectorStoreSearchError(BadRequestError):
|
||||
def __init__(
|
||||
self,
|
||||
failures: Sequence[VectorStoreSearchFailure],
|
||||
model: str | None = None,
|
||||
llm_provider: str | None = None,
|
||||
) -> None:
|
||||
self.failures: Final[tuple[VectorStoreSearchFailure, ...]] = tuple(failures)
|
||||
detail: Final = "; ".join(f"{failure['vector_store_id']}: {failure['error']}" for failure in self.failures)
|
||||
super().__init__(
|
||||
message=(
|
||||
"The request could not be grounded in every configured vector store. "
|
||||
f"{len(self.failures)} vector store search(es) failed: {detail}"
|
||||
),
|
||||
model=model,
|
||||
llm_provider=llm_provider,
|
||||
body={"type": "invalid_request_error", "code": VECTOR_STORE_SEARCH_FAILED_CODE},
|
||||
)
|
||||
|
||||
|
||||
class UnprocessableEntityError(openai.UnprocessableEntityError):
|
||||
def __init__(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -5,20 +5,21 @@ This hook is called before making an LLM request when a vector store is configur
|
|||
It searches the vector store for relevant context and appends it to the messages.
|
||||
"""
|
||||
|
||||
from collections.abc import Awaitable, Callable
|
||||
from collections.abc import Awaitable, Callable, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Any, Final, Protocol, cast
|
||||
from typing import TYPE_CHECKING, Any, Final, Protocol, assert_never, cast
|
||||
|
||||
import litellm
|
||||
import litellm.vector_stores
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.exceptions import VectorStoreSearchError
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.types.llms.openai import AllMessageValues, ChatCompletionUserMessage
|
||||
from litellm.types.prompts.init_prompts import PromptSpec
|
||||
from litellm.types.utils import CallTypes, StandardCallbackDynamicParams
|
||||
from litellm.types.vector_stores import (
|
||||
LiteLLM_ManagedVectorStore,
|
||||
VectorStoreResultContent,
|
||||
VectorStoreSearchFailure,
|
||||
VectorStoreSearchResponse,
|
||||
VectorStoreSearchResult,
|
||||
)
|
||||
|
|
@ -30,6 +31,8 @@ if TYPE_CHECKING:
|
|||
else:
|
||||
LiteLLMLoggingObj = Any
|
||||
|
||||
SEARCH_FAILURES_FIELD: Final = "vector_store_search_failures"
|
||||
|
||||
|
||||
class ProxyRuntime(Protocol):
|
||||
def llm_router(self) -> "Router | None": ...
|
||||
|
|
@ -54,11 +57,31 @@ class ProxyServerRuntime:
|
|||
return prisma_client
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class SearchSucceeded:
|
||||
response: VectorStoreSearchResponse
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class SearchFailed:
|
||||
failure: VectorStoreSearchFailure
|
||||
|
||||
|
||||
SearchOutcome = SearchSucceeded | SearchFailed
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class VectorStoreAugmentation:
|
||||
messages: tuple[AllMessageValues, ...]
|
||||
search_results: tuple[VectorStoreSearchResponse, ...]
|
||||
failures: tuple[VectorStoreSearchFailure, ...]
|
||||
|
||||
|
||||
class VectorStorePreCallHook(CustomLogger):
|
||||
CONTENT_PREFIX_STRING = "Context:\n\n"
|
||||
"""
|
||||
Custom logger that handles vector store searches before LLM calls.
|
||||
|
||||
|
||||
When a vector store is configured, this hook:
|
||||
1. Extracts the query from the last user message
|
||||
2. Calls litellm.vector_stores.search() to get relevant context
|
||||
|
|
@ -101,100 +124,152 @@ class VectorStorePreCallHook(CustomLogger):
|
|||
Returns:
|
||||
Tuple of (model, modified_messages, non_default_params)
|
||||
"""
|
||||
requested_vector_store_ids: Final = _requested_vector_store_ids(non_default_params)
|
||||
try:
|
||||
# Check if vector store is configured
|
||||
if litellm.vector_store_registry is None:
|
||||
return model, messages, non_default_params
|
||||
|
||||
prisma_client: Final = self.proxy_runtime.prisma_client()
|
||||
llm_router: Final = self.proxy_runtime.llm_router()
|
||||
|
||||
# Use database fallback to ensure synchronization across instances
|
||||
vector_stores_to_run: list[
|
||||
LiteLLM_ManagedVectorStore
|
||||
] = await litellm.vector_store_registry.pop_vector_stores_to_run_with_db_fallback(
|
||||
augmentation: VectorStoreAugmentation | None = await self._augment_messages(
|
||||
messages=messages,
|
||||
non_default_params=non_default_params,
|
||||
tools=tools,
|
||||
prisma_client=prisma_client,
|
||||
litellm_logging_obj=litellm_logging_obj,
|
||||
)
|
||||
|
||||
if not vector_stores_to_run:
|
||||
return model, messages, non_default_params
|
||||
|
||||
# Extract the query from the last user message
|
||||
query: Final = self._extract_query_from_messages(messages)
|
||||
|
||||
if not query:
|
||||
verbose_logger.debug("No query found in messages for vector store search")
|
||||
return model, messages, non_default_params
|
||||
|
||||
modified_messages: list[AllMessageValues] = messages.copy()
|
||||
all_search_results: Final[list[VectorStoreSearchResponse]] = []
|
||||
|
||||
for vector_store_to_run in vector_stores_to_run:
|
||||
# Get vector store id from the vector store config
|
||||
vector_store_id = vector_store_to_run.get("vector_store_id", "")
|
||||
custom_llm_provider = vector_store_to_run.get("custom_llm_provider")
|
||||
litellm_params_for_vector_store = vector_store_to_run.get("litellm_params", {}) or {}
|
||||
request_litellm_params = litellm_logging_obj.model_call_details.get("litellm_params", {})
|
||||
request_metadata = (
|
||||
request_litellm_params.get("metadata", {}) if isinstance(request_litellm_params, dict) else {}
|
||||
)
|
||||
if llm_router is not None:
|
||||
search_function = cast( # cast-ok: normalize router search callable
|
||||
Callable[..., Awaitable[VectorStoreSearchResponse]],
|
||||
llm_router.avector_store_search,
|
||||
)
|
||||
else:
|
||||
search_function = cast( # cast-ok: normalize SDK search callable
|
||||
Callable[..., Awaitable[VectorStoreSearchResponse]],
|
||||
litellm.vector_stores.asearch,
|
||||
)
|
||||
try:
|
||||
search_response = await search_function(
|
||||
**{
|
||||
"vector_store_id": vector_store_id,
|
||||
"query": query,
|
||||
"custom_llm_provider": custom_llm_provider,
|
||||
"metadata": request_metadata,
|
||||
**litellm_params_for_vector_store,
|
||||
},
|
||||
)
|
||||
except Exception as search_error:
|
||||
verbose_logger.warning(
|
||||
"Vector store search failed for vector_store_id=%s, continuing without its context: %s",
|
||||
vector_store_id,
|
||||
search_error,
|
||||
)
|
||||
continue
|
||||
|
||||
verbose_logger.debug("search_response: %s", search_response)
|
||||
|
||||
# Store search results for later use in citations
|
||||
all_search_results.append(search_response)
|
||||
|
||||
# Process search results and append as context
|
||||
modified_messages = self._append_search_results_to_messages(
|
||||
messages=modified_messages, search_response=search_response
|
||||
)
|
||||
|
||||
# Get the number of results for logging
|
||||
num_results = 0
|
||||
num_results = len(search_response.get("data", []) or [])
|
||||
verbose_logger.debug("Vector store search completed. Added context from %s results", num_results)
|
||||
|
||||
# Store search results as-is (already in OpenAI-compatible format)
|
||||
if litellm_logging_obj and all_search_results:
|
||||
litellm_logging_obj.model_call_details["search_results"] = all_search_results
|
||||
|
||||
return model, modified_messages, non_default_params
|
||||
|
||||
except Exception as e:
|
||||
verbose_logger.exception("Error in VectorStorePreCallHook: %s", e)
|
||||
# Return original parameters on error
|
||||
verbose_logger.exception(
|
||||
"Error in VectorStorePreCallHook for vector_store_ids=%s: %s",
|
||||
requested_vector_store_ids,
|
||||
e,
|
||||
)
|
||||
return model, messages, non_default_params
|
||||
|
||||
def _extract_query_from_messages(self, messages: list[AllMessageValues]) -> str | None:
|
||||
if augmentation is None:
|
||||
return model, messages, non_default_params
|
||||
|
||||
for detail, value in (
|
||||
("search_results", list(augmentation.search_results)),
|
||||
(SEARCH_FAILURES_FIELD, augmentation.failures),
|
||||
):
|
||||
if value:
|
||||
litellm_logging_obj.model_call_details[detail] = value
|
||||
|
||||
if augmentation.failures:
|
||||
match litellm.vector_store_search_failure_mode:
|
||||
case "error":
|
||||
raise VectorStoreSearchError(failures=augmentation.failures, model=model)
|
||||
case "annotate":
|
||||
pass
|
||||
case unreachable:
|
||||
assert_never(unreachable)
|
||||
|
||||
return model, list(augmentation.messages), non_default_params
|
||||
|
||||
async def _augment_messages(
|
||||
self,
|
||||
messages: Sequence[AllMessageValues],
|
||||
non_default_params: dict,
|
||||
tools: list[dict] | None,
|
||||
litellm_logging_obj: LiteLLMLoggingObj,
|
||||
) -> VectorStoreAugmentation | None:
|
||||
if litellm.vector_store_registry is None:
|
||||
return None
|
||||
|
||||
prisma_client: Final = self.proxy_runtime.prisma_client()
|
||||
llm_router: Final = self.proxy_runtime.llm_router()
|
||||
|
||||
# Use database fallback to ensure synchronization across instances
|
||||
vector_stores_to_run: Final[
|
||||
Sequence[LiteLLM_ManagedVectorStore]
|
||||
] = await litellm.vector_store_registry.pop_vector_stores_to_run_with_db_fallback(
|
||||
non_default_params=non_default_params,
|
||||
tools=tools,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
|
||||
if not vector_stores_to_run:
|
||||
return None
|
||||
|
||||
query: Final = self._extract_query_from_messages(messages)
|
||||
|
||||
if not query:
|
||||
verbose_logger.debug("No query found in messages for vector store search")
|
||||
return None
|
||||
|
||||
request_litellm_params: Final = litellm_logging_obj.model_call_details.get("litellm_params", {})
|
||||
request_metadata: Final = (
|
||||
request_litellm_params.get("metadata", {}) if isinstance(request_litellm_params, dict) else {}
|
||||
)
|
||||
search_function: Final = (
|
||||
cast( # cast-ok: normalize router search callable
|
||||
Callable[..., Awaitable[VectorStoreSearchResponse]],
|
||||
llm_router.avector_store_search,
|
||||
)
|
||||
if llm_router is not None
|
||||
else cast( # cast-ok: normalize SDK search callable
|
||||
Callable[..., Awaitable[VectorStoreSearchResponse]],
|
||||
litellm.vector_stores.asearch,
|
||||
)
|
||||
)
|
||||
|
||||
outcomes: Final = tuple(
|
||||
[
|
||||
await self._search_one(
|
||||
vector_store=vector_store_to_run,
|
||||
query=query,
|
||||
request_metadata=request_metadata,
|
||||
search_function=search_function,
|
||||
)
|
||||
for vector_store_to_run in vector_stores_to_run
|
||||
]
|
||||
)
|
||||
search_results: Final = tuple(outcome.response for outcome in outcomes if isinstance(outcome, SearchSucceeded))
|
||||
failures: Final = tuple(outcome.failure for outcome in outcomes if isinstance(outcome, SearchFailed))
|
||||
|
||||
return VectorStoreAugmentation(
|
||||
messages=self._messages_with_context(messages=messages, search_results=search_results),
|
||||
search_results=search_results,
|
||||
failures=failures,
|
||||
)
|
||||
|
||||
async def _search_one(
|
||||
self,
|
||||
vector_store: LiteLLM_ManagedVectorStore,
|
||||
query: str,
|
||||
request_metadata: Mapping[str, object],
|
||||
search_function: Callable[..., Awaitable[VectorStoreSearchResponse]],
|
||||
) -> SearchOutcome:
|
||||
vector_store_id: Final = vector_store.get("vector_store_id", "")
|
||||
custom_llm_provider: Final = vector_store.get("custom_llm_provider")
|
||||
litellm_params_for_vector_store: Final = vector_store.get("litellm_params", {}) or {}
|
||||
try:
|
||||
search_response: Final = await search_function(
|
||||
**{
|
||||
"vector_store_id": vector_store_id,
|
||||
"query": query,
|
||||
"custom_llm_provider": custom_llm_provider,
|
||||
"metadata": request_metadata,
|
||||
**litellm_params_for_vector_store,
|
||||
},
|
||||
)
|
||||
except Exception as search_error:
|
||||
verbose_logger.warning(
|
||||
"Vector store search failed for vector_store_id=%s, continuing without its context: %s",
|
||||
vector_store_id,
|
||||
search_error,
|
||||
)
|
||||
return SearchFailed(
|
||||
failure=VectorStoreSearchFailure(
|
||||
vector_store_id=vector_store_id,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
error=str(search_error),
|
||||
)
|
||||
)
|
||||
|
||||
verbose_logger.debug(
|
||||
"Vector store search completed for vector_store_id=%s. Added context from %s results",
|
||||
vector_store_id,
|
||||
len(search_response.get("data", []) or []),
|
||||
)
|
||||
return SearchSucceeded(response=search_response)
|
||||
|
||||
def _extract_query_from_messages(self, messages: Sequence[AllMessageValues]) -> str | None:
|
||||
"""
|
||||
Extract the query from the last user message.
|
||||
|
||||
|
|
@ -223,48 +298,40 @@ class VectorStorePreCallHook(CustomLogger):
|
|||
|
||||
return None
|
||||
|
||||
def _append_search_results_to_messages(
|
||||
def _messages_with_context(
|
||||
self,
|
||||
messages: list[AllMessageValues],
|
||||
search_response: VectorStoreSearchResponse,
|
||||
) -> list[AllMessageValues]:
|
||||
"""
|
||||
Append search results as context to the messages.
|
||||
messages: Sequence[AllMessageValues],
|
||||
search_results: Sequence[VectorStoreSearchResponse],
|
||||
) -> tuple[AllMessageValues, ...]:
|
||||
context_messages: Final = tuple(
|
||||
context_message
|
||||
for search_response in search_results
|
||||
if (context_message := self._context_message(search_response)) is not None
|
||||
)
|
||||
if not context_messages:
|
||||
return tuple(messages)
|
||||
return (*messages[:-1], *context_messages, *messages[-1:])
|
||||
|
||||
Args:
|
||||
messages: Original list of messages
|
||||
search_response: Response from vector store search
|
||||
|
||||
Returns:
|
||||
Modified list of messages with context appended
|
||||
"""
|
||||
search_response_data: Final[list[VectorStoreSearchResult] | None] = search_response.get("data")
|
||||
def _context_message(self, search_response: VectorStoreSearchResponse) -> AllMessageValues | None:
|
||||
"""Build the context message for one vector store's results, or None when it returned nothing usable."""
|
||||
search_response_data: Final[Sequence[VectorStoreSearchResult] | None] = search_response.get("data")
|
||||
if not search_response_data:
|
||||
return messages
|
||||
return None
|
||||
|
||||
context_content = self.CONTENT_PREFIX_STRING
|
||||
context_texts: Final = tuple(
|
||||
content_text
|
||||
for result in search_response_data
|
||||
for content_item in (result.get("content") or ())
|
||||
if (content_text := content_item.get("text"))
|
||||
)
|
||||
if not context_texts:
|
||||
return None
|
||||
|
||||
for result in search_response_data:
|
||||
result_content: list[VectorStoreResultContent] | None = result.get("content")
|
||||
if result_content:
|
||||
for content_item in result_content:
|
||||
content_text: str | None = content_item.get("text")
|
||||
if content_text:
|
||||
context_content += content_text + "\n\n"
|
||||
|
||||
# Only add context if we found any content
|
||||
if context_content != "Context:\n\n":
|
||||
# Create a copy of messages to avoid modifying the original
|
||||
modified_messages: Final = messages.copy()
|
||||
# Add context as a new message before the last user message
|
||||
context_message: Final[ChatCompletionUserMessage] = {
|
||||
"role": "user",
|
||||
"content": context_content,
|
||||
}
|
||||
modified_messages.insert(-1, cast(AllMessageValues, context_message))
|
||||
return modified_messages
|
||||
|
||||
return messages
|
||||
context_message: Final[ChatCompletionUserMessage] = {
|
||||
"role": "user",
|
||||
"content": self.CONTENT_PREFIX_STRING + "".join(f"{text}\n\n" for text in context_texts),
|
||||
}
|
||||
return cast(AllMessageValues, context_message)
|
||||
|
||||
async def async_post_call_success_deployment_hook(
|
||||
self,
|
||||
|
|
@ -287,34 +354,29 @@ class VectorStorePreCallHook(CustomLogger):
|
|||
verbose_logger.debug("No litellm_logging_obj in request_data")
|
||||
return None
|
||||
|
||||
verbose_logger.debug("model_call_details keys: %s", list(litellm_logging_obj.model_call_details.keys()))
|
||||
|
||||
# Get search results from model_call_details (already in OpenAI format)
|
||||
search_results: Final[list[VectorStoreSearchResponse] | None] = litellm_logging_obj.model_call_details.get(
|
||||
"search_results"
|
||||
search_results: Final[Sequence[VectorStoreSearchResponse] | None] = (
|
||||
litellm_logging_obj.model_call_details.get("search_results")
|
||||
)
|
||||
search_failures: Final[Sequence[VectorStoreSearchFailure] | None] = (
|
||||
litellm_logging_obj.model_call_details.get(SEARCH_FAILURES_FIELD)
|
||||
)
|
||||
|
||||
verbose_logger.debug("Search results found: %s", search_results is not None)
|
||||
|
||||
if not search_results:
|
||||
verbose_logger.debug("No search results found")
|
||||
if not search_results and not search_failures:
|
||||
verbose_logger.debug("No search results or search failures found")
|
||||
return None
|
||||
|
||||
# Add search results to response object
|
||||
if hasattr(response, "choices") and response.choices:
|
||||
for choice in response.choices:
|
||||
if hasattr(choice, "message") and choice.message:
|
||||
# Get existing provider_specific_fields or create new dict
|
||||
provider_fields = getattr(choice.message, "provider_specific_fields", None) or {}
|
||||
|
||||
# Add search results (already in OpenAI-compatible format)
|
||||
provider_fields["search_results"] = search_results
|
||||
|
||||
# Set the provider_specific_fields
|
||||
if search_results:
|
||||
provider_fields["search_results"] = search_results
|
||||
if search_failures:
|
||||
provider_fields[SEARCH_FAILURES_FIELD] = search_failures
|
||||
setattr(choice.message, "provider_specific_fields", provider_fields)
|
||||
|
||||
verbose_logger.debug("Added %s search results to response", len(search_results))
|
||||
|
||||
# Return modified response
|
||||
return response
|
||||
|
||||
|
|
@ -339,29 +401,24 @@ class VectorStorePreCallHook(CustomLogger):
|
|||
verbose_logger.debug("VectorStorePreCallHook.async_post_call_streaming_deployment_hook called")
|
||||
|
||||
# Get search results from model_call_details (already in OpenAI format)
|
||||
search_results: Final[list[VectorStoreSearchResponse] | None] = request_data.get("search_results")
|
||||
search_results: Final[Sequence[VectorStoreSearchResponse] | None] = request_data.get("search_results")
|
||||
search_failures: Final[Sequence[VectorStoreSearchFailure] | None] = request_data.get(SEARCH_FAILURES_FIELD)
|
||||
|
||||
verbose_logger.debug("Search results found for streaming chunk: %s", search_results is not None)
|
||||
|
||||
if not search_results:
|
||||
verbose_logger.debug("No search results found for streaming chunk")
|
||||
if not search_results and not search_failures:
|
||||
verbose_logger.debug("No search results or search failures found for streaming chunk")
|
||||
return response_chunk
|
||||
|
||||
# Add search results to streaming chunk
|
||||
if hasattr(response_chunk, "choices") and response_chunk.choices:
|
||||
for choice in response_chunk.choices:
|
||||
if hasattr(choice, "delta") and choice.delta:
|
||||
# Get existing provider_specific_fields or create new dict
|
||||
provider_fields = getattr(choice.delta, "provider_specific_fields", None) or {}
|
||||
|
||||
# Add search results (already in OpenAI-compatible format)
|
||||
provider_fields["search_results"] = search_results
|
||||
|
||||
# Set the provider_specific_fields
|
||||
if search_results:
|
||||
provider_fields["search_results"] = search_results
|
||||
if search_failures:
|
||||
provider_fields[SEARCH_FAILURES_FIELD] = search_failures
|
||||
choice.delta.provider_specific_fields = provider_fields
|
||||
|
||||
verbose_logger.debug("Added %s search results to streaming chunk", len(search_results))
|
||||
|
||||
# Return modified chunk
|
||||
return response_chunk
|
||||
|
||||
|
|
@ -369,3 +426,10 @@ class VectorStorePreCallHook(CustomLogger):
|
|||
verbose_logger.exception("Error adding search results to streaming chunk: %s", e)
|
||||
# Don't fail the request if search results fail to be added
|
||||
return response_chunk
|
||||
|
||||
|
||||
def _requested_vector_store_ids(non_default_params: Mapping[str, object]) -> tuple[str, ...]:
|
||||
requested: Final = non_default_params.get("vector_store_ids")
|
||||
if not isinstance(requested, (list, tuple)):
|
||||
return ()
|
||||
return tuple(str(vector_store_id) for vector_store_id in requested)
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ from enum import Enum
|
|||
from typing import Any, Literal
|
||||
|
||||
from pydantic import BaseModel
|
||||
from typing_extensions import TypedDict
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
|
||||
class SupportedVectorStoreIntegrations(str, Enum):
|
||||
|
|
@ -95,6 +95,17 @@ class VectorStoreSearchResponse(TypedDict, total=False):
|
|||
data: list[VectorStoreSearchResult] | None
|
||||
|
||||
|
||||
VectorStoreSearchFailureMode = Literal["annotate", "error"]
|
||||
|
||||
|
||||
class VectorStoreSearchFailure(TypedDict):
|
||||
"""A configured vector store whose search failed, as reported back to the API caller"""
|
||||
|
||||
vector_store_id: ReadOnly[str]
|
||||
custom_llm_provider: ReadOnly[str | None]
|
||||
error: ReadOnly[str]
|
||||
|
||||
|
||||
class VectorStoreSearchOptionalRequestParams(TypedDict, total=False):
|
||||
"""TypedDict for Optional parameters supported by the vector store search API."""
|
||||
|
||||
|
|
|
|||
|
|
@ -12,6 +12,15 @@ from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook i
|
|||
VectorStorePreCallHook,
|
||||
)
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.types.utils import (
|
||||
CallTypes,
|
||||
Choices,
|
||||
Delta,
|
||||
Message,
|
||||
ModelResponse,
|
||||
ModelResponseStream,
|
||||
StreamingChoices,
|
||||
)
|
||||
from litellm.types.vector_stores import (
|
||||
VectorStoreResultContent,
|
||||
VectorStoreSearchResponse,
|
||||
|
|
@ -36,6 +45,18 @@ def _search_response(text: str) -> VectorStoreSearchResponse:
|
|||
)
|
||||
|
||||
|
||||
def _first_message(response: ModelResponse) -> Message:
|
||||
choice = response.choices[0]
|
||||
assert isinstance(choice, Choices)
|
||||
return choice.message
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ExplodingRegistry:
|
||||
async def pop_vector_stores_to_run_with_db_fallback(self, **kwargs: object) -> list[LiteLLM_ManagedVectorStore]:
|
||||
raise RuntimeError("the registry blew up")
|
||||
|
||||
|
||||
@dataclass
|
||||
class RecordingRouter:
|
||||
failing_vector_store_ids: frozenset[str] = frozenset()
|
||||
|
|
@ -285,3 +306,181 @@ def test_the_default_runtime_follows_the_proxy_globals(monkeypatch: pytest.Monke
|
|||
|
||||
assert runtime.llm_router() is router
|
||||
assert runtime.prisma_client() is prisma
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_failing_vector_store_is_reported_back_to_the_caller(
|
||||
registry_with: RegisterStores,
|
||||
) -> None:
|
||||
"""Regression (LIT-6809): a silently dropped store left the caller with an un-augmented answer and no signal."""
|
||||
registry_with("vs-broken", "vs-healthy")
|
||||
logging_obj = FakeLoggingObj({})
|
||||
|
||||
await _run_hook(
|
||||
VectorStorePreCallHook(
|
||||
proxy_runtime=FakeProxyRuntime(router=RecordingRouter(failing_vector_store_ids=frozenset({"vs-broken"})))
|
||||
),
|
||||
["vs-broken", "vs-healthy"],
|
||||
logging_obj,
|
||||
)
|
||||
|
||||
response = ModelResponse(choices=[Choices(message=Message(content="an answer"))])
|
||||
await VectorStorePreCallHook(proxy_runtime=FakeProxyRuntime(router=None)).async_post_call_success_deployment_hook(
|
||||
request_data={"litellm_logging_obj": logging_obj},
|
||||
response=response,
|
||||
call_type=CallTypes.acompletion,
|
||||
)
|
||||
|
||||
provider_specific_fields = _first_message(response).provider_specific_fields or {}
|
||||
assert provider_specific_fields["vector_store_search_failures"] == (
|
||||
{
|
||||
"vector_store_id": "vs-broken",
|
||||
"custom_llm_provider": "bedrock",
|
||||
"error": "litellm.BadRequestError: no healthy deployments for vs-broken",
|
||||
},
|
||||
)
|
||||
assert len(provider_specific_fields["search_results"]) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_healthy_vector_store_alone_reports_no_failures(registry_with: RegisterStores) -> None:
|
||||
registry_with("vs-healthy")
|
||||
logging_obj = FakeLoggingObj({})
|
||||
|
||||
await _run_hook(
|
||||
VectorStorePreCallHook(proxy_runtime=FakeProxyRuntime(router=RecordingRouter())),
|
||||
["vs-healthy"],
|
||||
logging_obj,
|
||||
)
|
||||
|
||||
response = ModelResponse(choices=[Choices(message=Message(content="an answer"))])
|
||||
await VectorStorePreCallHook(proxy_runtime=FakeProxyRuntime(router=None)).async_post_call_success_deployment_hook(
|
||||
request_data={"litellm_logging_obj": logging_obj},
|
||||
response=response,
|
||||
call_type=CallTypes.acompletion,
|
||||
)
|
||||
|
||||
assert "vector_store_search_failures" not in (_first_message(response).provider_specific_fields or {})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_failing_vector_store_is_reported_on_the_streaming_chunk(registry_with: RegisterStores) -> None:
|
||||
registry_with("vs-broken")
|
||||
logging_obj = FakeLoggingObj({})
|
||||
|
||||
await _run_hook(
|
||||
VectorStorePreCallHook(
|
||||
proxy_runtime=FakeProxyRuntime(router=RecordingRouter(failing_vector_store_ids=frozenset({"vs-broken"})))
|
||||
),
|
||||
["vs-broken"],
|
||||
logging_obj,
|
||||
)
|
||||
|
||||
chunk = ModelResponseStream(choices=[StreamingChoices(delta=Delta(content="an answer"))])
|
||||
await VectorStorePreCallHook(
|
||||
proxy_runtime=FakeProxyRuntime(router=None)
|
||||
).async_post_call_streaming_deployment_hook(
|
||||
request_data=logging_obj.model_call_details,
|
||||
response_chunk=chunk,
|
||||
call_type=CallTypes.acompletion,
|
||||
)
|
||||
|
||||
assert (chunk.choices[0].delta.provider_specific_fields or {})["vector_store_search_failures"] == (
|
||||
{
|
||||
"vector_store_id": "vs-broken",
|
||||
"custom_llm_provider": "bedrock",
|
||||
"error": "litellm.BadRequestError: no healthy deployments for vs-broken",
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_error_mode_fails_the_request_instead_of_answering_without_the_knowledge_base(
|
||||
registry_with: RegisterStores,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""Regression (LIT-6809): opting in must turn an ungrounded answer into a 400 the caller can act on."""
|
||||
registry_with("vs-broken", "vs-healthy")
|
||||
monkeypatch.setattr(litellm, "vector_store_search_failure_mode", "error")
|
||||
|
||||
with pytest.raises(litellm.VectorStoreSearchError) as raised:
|
||||
await _run_hook(
|
||||
VectorStorePreCallHook(
|
||||
proxy_runtime=FakeProxyRuntime(
|
||||
router=RecordingRouter(failing_vector_store_ids=frozenset({"vs-broken"}))
|
||||
)
|
||||
),
|
||||
["vs-broken", "vs-healthy"],
|
||||
FakeLoggingObj({}),
|
||||
)
|
||||
|
||||
assert raised.value.status_code == 400
|
||||
assert raised.value.failures == (
|
||||
{
|
||||
"vector_store_id": "vs-broken",
|
||||
"custom_llm_provider": "bedrock",
|
||||
"error": "litellm.BadRequestError: no healthy deployments for vs-broken",
|
||||
},
|
||||
)
|
||||
assert "vs-broken: litellm.BadRequestError: no healthy deployments for vs-broken" in raised.value.message
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_error_mode_leaves_a_fully_healthy_request_alone(
|
||||
registry_with: RegisterStores,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
registry_with("vs-healthy")
|
||||
monkeypatch.setattr(litellm, "vector_store_search_failure_mode", "error")
|
||||
|
||||
_, messages, _ = await _run_hook(
|
||||
VectorStorePreCallHook(proxy_runtime=FakeProxyRuntime(router=RecordingRouter())),
|
||||
["vs-healthy"],
|
||||
FakeLoggingObj({}),
|
||||
)
|
||||
|
||||
assert messages[0]["content"] == "Context:\n\ncontext from vs-healthy\n\n"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_error_mode_does_not_swallow_the_raise_in_the_hooks_own_catch_all(
|
||||
registry_with: RegisterStores,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
warnings: list[logging.LogRecord],
|
||||
) -> None:
|
||||
"""Regression (LIT-6809): the catch-all around the hook must not turn the opted-in failure back into a 200."""
|
||||
registry_with("vs-broken")
|
||||
monkeypatch.setattr(litellm, "vector_store_search_failure_mode", "error")
|
||||
|
||||
with pytest.raises(litellm.VectorStoreSearchError):
|
||||
await _run_hook(
|
||||
VectorStorePreCallHook(
|
||||
proxy_runtime=FakeProxyRuntime(
|
||||
router=RecordingRouter(failing_vector_store_ids=frozenset({"vs-broken"}))
|
||||
)
|
||||
),
|
||||
["vs-broken"],
|
||||
FakeLoggingObj({}),
|
||||
)
|
||||
|
||||
assert [record.levelname for record in warnings] == ["WARNING"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_crash_outside_the_search_names_the_requested_vector_stores(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
warnings: list[logging.LogRecord],
|
||||
) -> None:
|
||||
"""Regression (LIT-6809): the catch-all logged no store id, so an operator could not tell which store broke."""
|
||||
monkeypatch.setattr(litellm, "vector_store_registry", ExplodingRegistry())
|
||||
|
||||
_, messages, _ = await _run_hook(
|
||||
VectorStorePreCallHook(proxy_runtime=FakeProxyRuntime(router=None)),
|
||||
["vs-one", "vs-two"],
|
||||
FakeLoggingObj({}),
|
||||
)
|
||||
|
||||
assert messages == [{"role": "user", "content": "what is litellm?"}]
|
||||
assert [record.getMessage() for record in warnings] == [
|
||||
"Error in VectorStorePreCallHook for vector_store_ids=('vs-one', 'vs-two'): the registry blew up"
|
||||
]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue