From 445ccc14b29bb4a93e3b21bf907a7b2a26da507a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 00:53:55 -0700 Subject: [PATCH] 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. --- litellm/__init__.py | 3 + litellm/exceptions.py | 25 ++ .../vector_store_pre_call_hook.py | 378 ++++++++++-------- litellm/types/vector_stores.py | 13 +- .../test_vector_store_pre_call_hook.py | 199 +++++++++ 5 files changed, 460 insertions(+), 158 deletions(-) diff --git a/litellm/__init__.py b/litellm/__init__.py index 41a3789ab0d..76daf771c4b 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -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 diff --git a/litellm/exceptions.py b/litellm/exceptions.py index 286f7528896..4387370e58c 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -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, diff --git a/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py b/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py index 12ff38ce4ba..d5c893c4b20 100644 --- a/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py +++ b/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py @@ -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) diff --git a/litellm/types/vector_stores.py b/litellm/types/vector_stores.py index 474c652ff3a..9051287b30a 100644 --- a/litellm/types/vector_stores.py +++ b/litellm/types/vector_stores.py @@ -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.""" diff --git a/tests/test_litellm/integrations/vector_store_integrations/test_vector_store_pre_call_hook.py b/tests/test_litellm/integrations/vector_store_integrations/test_vector_store_pre_call_hook.py index ae5cffd8ab0..64ee9474e44 100644 --- a/tests/test_litellm/integrations/vector_store_integrations/test_vector_store_pre_call_hook.py +++ b/tests/test_litellm/integrations/vector_store_integrations/test_vector_store_pre_call_hook.py @@ -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" + ]