From a65fba992ba89fd9d41d1a9fb38a789cab058d96 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 11 Aug 2026 13:38:39 +0000 Subject: [PATCH] chore(typing): clear basedpyright Any errors in vertex files and emulated file_search Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llms/vertex_ai/files/transformation.py | 181 ++-- .../responses/file_search/emulated_handler.py | 774 +++++++++++------- 2 files changed, 577 insertions(+), 378 deletions(-) diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index 3538fc5b1a7..fbcd1de2cfd 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -5,12 +5,13 @@ import json import os import re import time -from collections.abc import Callable, Iterable, Iterator -from typing import Any, Final +from collections.abc import Callable, Iterable, Iterator, Mapping +from typing import Final, Protocol import httpx from httpx import Headers, Response from openai.types.file_deleted import FileDeleted +from pydantic import TypeAdapter import litellm from litellm._uuid import uuid @@ -46,13 +47,14 @@ from litellm.types.files import StreamingMediaUploadConfig from litellm.types.llms.openai import ( AllMessageValues, CreateFileRequest, + FileContentRequest, FileTypes, HttpxBinaryResponseContent, OpenAICreateFileRequestOptionalParams, OpenAIFileObject, + OpenAIFilesPurpose, PathLike, ) -from litellm.types.llms.vertex_ai import GcsBucketResponse from litellm.types.utils import LlmProviders, ModelResponse from ..common_utils import VertexAIError @@ -61,6 +63,53 @@ from ..vertex_llm_base import VertexBase _GCP_LABEL_VALUE_MAX_LEN: Final = 63 _CUSTOM_ID_RAW_LABEL_PREFIX: Final = "b32_" +_JSON_OBJECT_ADAPTER: Final = TypeAdapter(dict[str, object]) +_PURPOSE_ADAPTER: Final = TypeAdapter[OpenAIFilesPurpose](OpenAIFilesPurpose) +_JSON_DECODER: Final = json.JSONDecoder() + + +class _JsonDecoder(Protocol): + def decode(self, s: str, /) -> object: ... + + +class _JsonResponse(Protocol): + def json(self) -> object: ... + + +def _json_object(value: object) -> dict[str, object]: + return _JSON_OBJECT_ADAPTER.validate_python(value) + + +def _json_object_or_empty(value: object) -> dict[str, object]: + return _json_object(value) if value else {} + + +def _parse_json_object(raw: str, decoder: _JsonDecoder = _JSON_DECODER) -> dict[str, object]: + return _json_object(decoder.decode(raw)) + + +def _response_json_object(response: _JsonResponse) -> dict[str, object]: + return _json_object(response.json()) + + +def _str_field(payload: Mapping[str, object], key: str, default: str = "") -> str: + value: Final = payload.get(key, default) + return value if isinstance(value, str) else default + + +def _int_field(payload: Mapping[str, object], key: str) -> int: + value: Final = payload.get(key, 0) + return int(value) if isinstance(value, (int, float, str)) else 0 + + +def _purpose_field(payload: Mapping[str, object], default: OpenAIFilesPurpose = "batch") -> OpenAIFilesPurpose: + return _PURPOSE_ADAPTER.validate_python(payload.get("purpose", default)) + + +def _gcs_file_id(payload: Mapping[str, object]) -> str: + gcs_id: Final = _str_field(payload, "id") + return "/".join(gcs_id.split("/")[:-1]) if gcs_id else "" + def _sanitize_gcp_label_value(value: str) -> str: """ @@ -106,7 +155,7 @@ def _decode_gcp_label_value_chunks(values: list[str]) -> str | None: return None -def _set_litellm_batch_custom_id_labels(labels: dict[str, str], custom_id: Any) -> None: +def _set_litellm_batch_custom_id_labels(labels: dict[str, str], custom_id: object) -> None: """ Store OpenAI batch custom_id for Vertex batch correlation. @@ -122,7 +171,7 @@ def _set_litellm_batch_custom_id_labels(labels: dict[str, str], custom_id: Any) labels[f"litellm_custom_id_raw_{index}"] = raw_label_chunk -def _get_litellm_batch_custom_id_from_labels(labels: dict[str, Any]) -> str: +def _get_litellm_batch_custom_id_from_labels(labels: Mapping[str, object]) -> str: """Prefer encoded custom_id when present (see _set_litellm_batch_custom_id_labels).""" raw: Final = labels.get("litellm_custom_id_raw") if raw: @@ -140,10 +189,15 @@ def _get_litellm_batch_custom_id_from_labels(labels: dict[str, Any]) -> str: return str(labels.get("litellm_custom_id", "unknown")) +# any-ok: batch bodies must reach _transform_request_body unmodified, so no lossy revalidation here +def _batch_messages(openai_request_body) -> list[AllMessageValues]: + return openai_request_body.get("messages", []) + + def _openai_batch_jsonl_entry_to_vertex_wrapped_request( - openai_entry: dict[str, Any], - map_openai_to_vertex_params: Callable[[dict[str, Any]], dict[str, Any]], -) -> dict[str, Any]: + openai_entry: Mapping[str, object], + map_openai_to_vertex_params: Callable[[dict[str, object]], dict[str, object]], +) -> dict[str, object]: """ Transforms a single OpenAI JSONL batch entry into its Vertex wrapped request. @@ -151,10 +205,10 @@ def _openai_batch_jsonl_entry_to_vertex_wrapped_request( Example Vertex jsonl {"request":{"contents": [{"role": "user", "parts": [{"text": "What is the relation between the following video and image samples?"}, {"fileData": {"fileUri": "gs://cloud-samples-data/generative-ai/video/animals.mp4", "mimeType": "video/mp4"}}, {"fileData": {"fileUri": "gs://cloud-samples-data/generative-ai/image/cricket.jpeg", "mimeType": "image/jpeg"}}]}]}} """ - openai_request_body: Final = openai_entry.get("body") or {} + openai_request_body: Final = _json_object_or_empty(openai_entry.get("body")) vertex_request_body: Final = _transform_request_body( - messages=openai_request_body.get("messages", []), - model=openai_request_body.get("model", ""), + messages=_batch_messages(openai_request_body), + model=_str_field(openai_request_body, "model"), optional_params=map_openai_to_vertex_params(openai_request_body), custom_llm_provider="vertex_ai", litellm_params={}, @@ -186,9 +240,7 @@ def _iter_openai_jsonl_lines(openai_file_content: FileTypes) -> Iterator[str]: ``str.splitlines()`` + ``line.strip()`` for ``\\n`` / ``\\r\\n`` delimited JSONL. """ - content: Any = openai_file_content - if isinstance(content, tuple): - content = content[1] + content: Final[object] = openai_file_content[1] if isinstance(openai_file_content, tuple) else openai_file_content if isinstance(content, (bytes, bytearray)): # Scan for newlines in place so a large in-memory payload is not copied @@ -220,14 +272,13 @@ def _iter_openai_jsonl_lines(openai_file_content: FileTypes) -> Iterator[str]: # object name, then the body stream), so it must rewind to 0. A # non-seekable handle would silently resume mid-stream and drop the # already-consumed first row, so reject it loudly instead. - seek: Final = getattr(content, "seek", None) - if seek is None: + if not hasattr(content, "seek"): raise ValueError( "Batch upload file handle must be seekable; got a non-seekable " "stream. Pass bytes, a path, or a seekable handle." ) try: - seek(0) + content.seek(0) except (OSError, ValueError) as e: raise ValueError( "Batch upload file handle must be seekable so it can be re-read " @@ -241,9 +292,9 @@ def _iter_openai_jsonl_lines(openai_file_content: FileTypes) -> Iterator[str]: def _iter_openai_jsonl_entries( openai_file_content: FileTypes, -) -> Iterator[dict[str, Any]]: +) -> Iterator[dict[str, object]]: for line in _iter_openai_jsonl_lines(openai_file_content): - yield json.loads(line) + yield _parse_json_object(line) class _OpenAIToVertexBatchUploadStream(BaseFileUploadStream): @@ -257,7 +308,7 @@ class _OpenAIToVertexBatchUploadStream(BaseFileUploadStream): def __init__( self, openai_file_content: FileTypes, - map_openai_to_vertex_params: Callable[[dict[str, Any]], dict[str, Any]], + map_openai_to_vertex_params: Callable[[dict[str, object]], dict[str, object]], ) -> None: self._openai_file_content = openai_file_content self._map_openai_to_vertex_params = map_openai_to_vertex_params @@ -308,17 +359,18 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): def _get_gcs_object_name_from_batch_jsonl( self, - openai_jsonl_content: list[dict[str, Any]], + openai_jsonl_content: list[dict[str, object]], ) -> str: """ Gets a unique GCS object name for the VertexAI batch prediction job named as: litellm-vertex-{model}-{uuid} """ - _model = openai_jsonl_content[0].get("body", {}).get("model", "") - if "publishers/google/models" not in _model: - _model = f"publishers/google/models/{_model}" - safe_model_path: Final = sanitize_cloud_object_path(_model, fallback="model") + model_name: Final = _str_field(_json_object_or_empty(openai_jsonl_content[0].get("body")), "model") + qualified_model: Final = ( + model_name if "publishers/google/models" in model_name else f"publishers/google/models/{model_name}" + ) + safe_model_path: Final = sanitize_cloud_object_path(qualified_model, fallback="model") object_name: Final = f"{VERTEX_AI_MANAGED_GCS_PREFIX}{safe_model_path}/{uuid.uuid4()}" return object_name @@ -343,9 +395,11 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): fallback_filename="file", ) - def _get_configured_bucket_name(self, litellm_params: dict) -> str: + def _get_configured_bucket_name(self, litellm_params: Mapping[str, object]) -> str: bucket_name: Final = ( - litellm_params.get("gcs_bucket_name") or litellm_params.get("bucket_name") or os.getenv("GCS_BUCKET_NAME") + _str_field(litellm_params, "gcs_bucket_name") + or _str_field(litellm_params, "bucket_name") + or os.getenv("GCS_BUCKET_NAME") ) if not bucket_name: raise ValueError("GCS bucket_name is required") @@ -396,8 +450,8 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): def _map_openai_to_vertex_params( self, - openai_request_body: dict[str, Any], - ) -> dict[str, Any]: + openai_request_body: dict[str, object], + ) -> dict[str, object]: """ wrapper to call VertexGeminiConfig.map_openai_params """ @@ -406,9 +460,8 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): ) config: Final = VertexGeminiConfig() - _model: Final = openai_request_body.get("model", "") - vertex_params: Final = config.map_openai_params( - model=_model, + vertex_params: Final[dict[str, object]] = config.map_openai_params( + model=_str_field(openai_request_body, "model"), non_default_params=openai_request_body, optional_params={}, drop_params=False, @@ -463,10 +516,8 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): """ Transform VertexAI File upload response into OpenAI-style FileObject """ - response_json: Final = raw_response.json() - try: - response_object: Final = GcsBucketResponse(**response_json) + response_object: Final = _response_json_object(raw_response) except Exception as e: raise VertexAIError( status_code=raw_response.status_code, @@ -474,19 +525,15 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): headers=raw_response.headers, ) - gcs_id = response_object.get("id", "") - # Remove the last numeric ID from the path - gcs_id = "/".join(gcs_id.split("/")[:-1]) if gcs_id else "" - return OpenAIFileObject( - purpose=response_object.get("purpose", "batch"), - id=f"gs://{gcs_id}", - filename=response_object.get("name", ""), + purpose=_purpose_field(response_object), + id=f"gs://{_gcs_file_id(response_object)}", + filename=_str_field(response_object, "name"), created_at=_convert_vertex_datetime_to_openai_datetime( - vertex_datetime=response_object.get("timeCreated", "") + vertex_datetime=_str_field(response_object, "timeCreated") ), status="uploaded", - bytes=int(response_object.get("size", 0)), + bytes=_int_field(response_object, "size"), object="file", ) @@ -523,18 +570,16 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): logging_obj: LiteLLMLoggingObj, litellm_params: dict, ) -> OpenAIFileObject: - response_json: Final = raw_response.json() - gcs_id = response_json.get("id", "") - gcs_id = "/".join(gcs_id.split("/")[:-1]) if gcs_id else "" + response_json: Final = _response_json_object(raw_response) return OpenAIFileObject( - id=f"gs://{gcs_id}", - bytes=int(response_json.get("size", 0)), + id=f"gs://{_gcs_file_id(response_json)}", + bytes=_int_field(response_json, "size"), created_at=_convert_vertex_datetime_to_openai_datetime( - vertex_datetime=response_json.get("timeCreated", "") + vertex_datetime=_str_field(response_json, "timeCreated") ), - filename=response_json.get("name", ""), + filename=_str_field(response_json, "name"), object="file", - purpose=response_json.get("metadata", {}).get("purpose", "batch"), + purpose=_purpose_field(_json_object_or_empty(response_json.get("metadata"))), status="processed", status_details=None, ) @@ -584,7 +629,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): def transform_file_content_request( self, - file_content_request, + file_content_request: FileContentRequest, optional_params: dict, litellm_params: dict, ) -> tuple[str, dict]: @@ -682,14 +727,15 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): # discriminating fields. Anything else (e.g. a binary file whose # first line is not valid UTF-8/JSON) raises and falls through to the # passthrough below, leaving the content untouched. - first_row: Final = json.loads(first_line) + first_row: Final = _parse_json_object(first_line) + first_row_response: Final = _json_object_or_empty(first_row.get("response")) is_vertex_batch_output: Final = ( "request" in first_row and "response" in first_row and "processed_time" in first_row and ( - "candidates" in first_row.get("response", {}) - or "promptFeedback" in first_row.get("response", {}) + "candidates" in first_row_response + or "promptFeedback" in first_row_response or bool(first_row.get("status")) ) ) @@ -723,7 +769,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): for line in itertools.chain([first_line], lines): try: openai_output = self._transform_single_vertex_batch_output_to_openai( - vertex_output=json.loads(line), + vertex_output=_parse_json_object(line), vertex_gemini_config=vertex_gemini_config, logging_obj=batch_transform_logging_obj, mock_httpx_response=mock_httpx_response, @@ -742,22 +788,22 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): def _transform_single_vertex_batch_output_to_openai( self, - vertex_output: dict[str, Any], + vertex_output: Mapping[str, object], vertex_gemini_config: VertexGeminiConfig, logging_obj: Logging, mock_httpx_response: httpx.Response, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Transform a single Vertex AI batch output line to OpenAI format. Uses the existing VertexGeminiConfig transformation for the response. """ # Extract custom_id from request labels (prefer raw for OpenAI round-trip) - request_data: Final = vertex_output.get("request", {}) - labels: Final = request_data.get("labels", {}) or {} + request_data: Final = _json_object_or_empty(vertex_output.get("request")) + labels: Final = _json_object_or_empty(request_data.get("labels")) custom_id: Final = _get_litellm_batch_custom_id_from_labels(labels) # Check if there's an error - status: Final = vertex_output.get("status", "") + status: Final = _str_field(vertex_output, "status") has_error: Final = bool(status) if has_error: @@ -772,12 +818,11 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): } # Transform successful response using existing transformation - vertex_response: Final = vertex_output.get("response", {}) + vertex_response: Final = _json_object_or_empty(vertex_output.get("response")) # Extract model from response - model = vertex_response.get("modelVersion", "gemini-1.5-flash-001") - if "@" in model: - model = model.split("@")[0] + model_version: Final = _str_field(vertex_response, "modelVersion", "gemini-1.5-flash-001") + model: Final = model_version.split("@")[0] if "@" in model_version else model_version try: # Use existing VertexGeminiConfig transformation @@ -792,7 +837,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): ) # Convert ModelResponse to dict - response_dict: Final = transformed_response.model_dump() + response_dict: Final = _json_object(transformed_response.model_dump()) # Return in OpenAI batch format return { @@ -800,7 +845,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): "custom_id": custom_id, "response": { "status_code": 200, - "request_id": response_dict.get("id", ""), + "request_id": _str_field(response_dict, "id"), "body": response_dict, }, "error": None, diff --git a/litellm/responses/file_search/emulated_handler.py b/litellm/responses/file_search/emulated_handler.py index 7854b17a06f..9d0cda5adf4 100644 --- a/litellm/responses/file_search/emulated_handler.py +++ b/litellm/responses/file_search/emulated_handler.py @@ -11,36 +11,209 @@ Flow: [file_search_call output item] + [message output item with file_citation annotations] """ -import json -import time import uuid -from collections.abc import Iterable -from typing import Any, Final, cast +from collections.abc import Iterable, Mapping, Sequence +from dataclasses import dataclass +from typing import Final, Literal, Protocol, TypeAlias, TypedDict, TypeVar + +from pydantic import BaseModel, ConfigDict, ValidationError from litellm._internal_context import is_internal_call from litellm._logging import verbose_logger -from litellm.types.llms.openai import ResponseOutputItem, ResponsesAPIResponse -from litellm.types.vector_stores import VectorStoreSearchResult - -# Keep ToolParam broad so we stay compatible with both dict and Pydantic forms -ToolParam = Any +from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig +from litellm.types.llms.openai import ( + ResponseInputParam, + ResponseOutputItem, + ResponsesAPIResponse, + ToolParam, +) +from litellm.types.vector_stores import ( + VectorStoreResultContent, + VectorStoreSearchResponse, + VectorStoreSearchResult, +) FILE_SEARCH_FUNCTION_NAME: Final = "litellm_file_search" +_ModelT = TypeVar("_ModelT", bound=BaseModel) + + +class SearchResultContentLike(Protocol): + """Attribute-based counterpart of ``VectorStoreResultContent``.""" + + @property + def text(self) -> str | None: ... + + +class SearchResultLike(Protocol): + """Attribute-based counterpart of ``VectorStoreSearchResult``.""" + + @property + def score(self) -> float | None: ... + + @property + def file_id(self) -> str | None: ... + + @property + def filename(self) -> str | None: ... + + @property + def content(self) -> Sequence[VectorStoreResultContent | SearchResultContentLike] | None: ... + + @property + def attributes(self) -> Mapping[str, object] | None: ... + + +class SearchResponseLike(Protocol): + """Attribute-based counterpart of ``VectorStoreSearchResponse``.""" + + @property + def data(self) -> Sequence["SearchResult"] | None: ... + + +SearchResult: TypeAlias = VectorStoreSearchResult | SearchResultLike + + +class FunctionToolQueriesSchema(TypedDict): + type: Literal["array"] + items: dict[str, str] + description: str + + +class FunctionToolVectorStoreIdSchema(TypedDict): + type: Literal["string"] + description: str + enum: list[str] + + +class FunctionToolProperties(TypedDict): + queries: FunctionToolQueriesSchema + vector_store_id: FunctionToolVectorStoreIdSchema + + +class FunctionToolParameters(TypedDict): + type: Literal["object"] + properties: FunctionToolProperties + required: list[str] + + +class EmulatedFileSearchTool(TypedDict): + type: Literal["function"] + name: str + description: str + parameters: FunctionToolParameters + + +class FileCitationAnnotation(TypedDict): + type: Literal["file_citation"] + index: int + file_id: str + filename: str + + +class OutputTextContent(TypedDict): + type: Literal["output_text"] + text: str + annotations: list[FileCitationAnnotation] + + +class MessageOutput(TypedDict): + type: Literal["message"] + role: Literal["assistant"] + content: list[OutputTextContent] + + +class FileSearchResultEntry(TypedDict): + file_id: str + filename: str + score: float | None + text: str + attributes: Mapping[str, object] + + +class FileSearchCallOutput(TypedDict): + type: Literal["file_search_call"] + id: str + status: Literal["completed"] + queries: list[str] + search_results: list[FileSearchResultEntry] | None + + +class FunctionCallOutput(TypedDict): + type: Literal["function_call_output"] + call_id: str + output: str + + +class FileSearchToolSpec(BaseModel): + model_config = ConfigDict(extra="ignore") + + type: Literal["file_search"] + vector_store_ids: tuple[str, ...] | None = None + + +class FileSearchFunctionCall(BaseModel): + model_config = ConfigDict(extra="ignore", from_attributes=True) + + type: Literal["function_call"] + name: str | None = None + call_id: str | None = None + id: str | None = None + arguments: str | Mapping[str, object] | None = None + + +class FileSearchArguments(BaseModel): + model_config = ConfigDict(extra="ignore") + + queries: tuple[str, ...] | str | None = None + query: str | None = None + vector_store_id: str | None = None + + +class MessageContentBlock(BaseModel): + model_config = ConfigDict(extra="ignore", from_attributes=True) + + type: str | None = None + text: str | None = None + + +class MessageOutputItem(BaseModel): + model_config = ConfigDict(extra="ignore", from_attributes=True) + + type: Literal["message"] + content: tuple[MessageContentBlock, ...] | None = None + + +class IncludeOptions(BaseModel): + model_config = ConfigDict(extra="ignore") + + include: tuple[str, ...] | None = None + + +def _validate_or_none(model: type[_ModelT], value: object) -> _ModelT | None: + try: + return model.model_validate(value) + except ValidationError: + return None + # --------------------------------------------------------------------------- # Detection # --------------------------------------------------------------------------- +def _as_file_search_tool(tool: object) -> FileSearchToolSpec | None: + return _validate_or_none(FileSearchToolSpec, tool) + + def should_use_emulated_file_search( tools: Iterable[ToolParam] | None, - provider_config: Any, # BaseResponsesAPIConfig + provider_config: BaseResponsesAPIConfig | None, ) -> bool: """Return True when there is a file_search tool and the provider can't handle it natively.""" if not tools: return False - has_fs: Final = any(isinstance(t, dict) and t.get("type") == "file_search" for t in tools) + has_fs: Final = any(_as_file_search_tool(tool) is not None for tool in tools) if not has_fs: return False return provider_config is None or not provider_config.supports_native_file_search() @@ -51,7 +224,7 @@ def should_use_emulated_file_search( # --------------------------------------------------------------------------- -def _build_function_tool(vector_store_ids: list[str]) -> dict[str, Any]: +def _build_function_tool(vector_store_ids: list[str]) -> EmulatedFileSearchTool: """ Create a Responses API function-tool definition that describes file search. The function accepts one or more natural-language queries (like OpenAI's native @@ -62,63 +235,62 @@ def _build_function_tool(vector_store_ids: list[str]) -> dict[str, Any]: NOT Chat Completion format (nested under "function"), so that the LiteLLMCompletionResponsesConfig transformation picks up name and description. """ - return { - "type": "function", - "name": FILE_SEARCH_FUNCTION_NAME, - "description": ( + return EmulatedFileSearchTool( + type="function", + name=FILE_SEARCH_FUNCTION_NAME, + description=( "Search the knowledge base for information relevant to the query. " "Use this whenever you need to look up specific facts, documents, " "or content from the vector store. You can provide multiple queries " "to search for different aspects of the information." ), - "parameters": { - "type": "object", - "properties": { - "queries": { - "type": "array", - "items": {"type": "string"}, - "description": ( + parameters=FunctionToolParameters( + type="object", + properties=FunctionToolProperties( + queries=FunctionToolQueriesSchema( + type="array", + items={"type": "string"}, + description=( "One or more search queries to look up in the vector store. " "Multiple queries help find comprehensive information from " "different angles." ), - }, - "vector_store_id": { - "type": "string", - "description": "ID of the vector store to search.", - "enum": vector_store_ids, - }, - }, - "required": ["queries"], - }, - } + ), + vector_store_id=FunctionToolVectorStoreIdSchema( + type="string", + description="ID of the vector store to search.", + enum=vector_store_ids, + ), + ), + required=["queries"], + ), + ) def _replace_file_search_tools( tools: Iterable[ToolParam] | None, -) -> tuple[list[dict[str, Any]], list[str]]: +) -> tuple[list[ToolParam | EmulatedFileSearchTool], list[str]]: """ Replace all file_search tools with a single function tool. Returns: (new_tools_list, all_vector_store_ids) """ - non_file_search: Final[list[dict[str, Any]]] = [] - vector_store_ids: Final[list[str]] = [] - - for tool in tools or []: - if isinstance(tool, dict) and tool.get("type") == "file_search": - ids = tool.get("vector_store_ids") or [] - vector_store_ids.extend(ids) - else: - non_file_search.append(tool) - - # Deduplicate while preserving order - unique_ids: Final[list[str]] = list(dict.fromkeys(vector_store_ids)) - if unique_ids: - non_file_search.append(_build_function_tool(unique_ids)) - - return non_file_search, unique_ids + parsed: Final = tuple((tool, _as_file_search_tool(tool)) for tool in tools or ()) + non_file_search: Final[list[ToolParam | EmulatedFileSearchTool]] = [ + tool for tool, file_search in parsed if file_search is None + ] + unique_ids: Final[list[str]] = list( + dict.fromkeys( + vector_store_id + for _, file_search in parsed + if file_search is not None + for vector_store_id in file_search.vector_store_ids or () + ) + ) + if not unique_ids: + return non_file_search, unique_ids + return [*non_file_search, _build_function_tool(unique_ids)], unique_ids # --------------------------------------------------------------------------- @@ -126,10 +298,31 @@ def _replace_file_search_tools( # --------------------------------------------------------------------------- +def _results_of(response: VectorStoreSearchResponse | SearchResponseLike) -> tuple[SearchResult, ...]: + results_data: Final = response.get("data") if isinstance(response, dict) else response.data + return tuple(results_data or ()) + + +async def _search_one_vector_store(vector_store_id: str, query: str) -> tuple[SearchResult, ...]: + """Run a single ``asearch`` call, returning no results when the search fails.""" + import litellm.vector_stores.main as vs_main + + try: + return _results_of(await vs_main.asearch(vector_store_id=vector_store_id, query=query)) + except Exception as exc: + verbose_logger.warning( + "file_search emulated: search failed for query='%s', vector_store_id='%s': %s", + query, + vector_store_id, + exc, + ) + return () + + async def _run_vector_searches( queries: list[str], vector_store_ids: list[str], -) -> tuple[list[str], list[VectorStoreSearchResult]]: +) -> tuple[list[str], list[SearchResult]]: """ Run `asearch` against all vector stores for all queries and collect results. @@ -140,30 +333,12 @@ async def _run_vector_searches( Returns: (queries_list, combined_results) """ - import litellm.vector_stores.main as vs_main - - all_results: Final[list[VectorStoreSearchResult]] = [] - ids_to_search: Final = vector_store_ids - - # Execute each query against all vector stores - for query in queries: - for vs_id in ids_to_search: - try: - response = await vs_main.asearch( - vector_store_id=vs_id, - query=query, - ) - results_data = response.get("data") if isinstance(response, dict) else getattr(response, "data", None) - if results_data: - all_results.extend(results_data) - except Exception as exc: - verbose_logger.warning( - "file_search emulated: search failed for query='%s', vector_store_id='%s': %s", - query, - vs_id, - exc, - ) - + all_results: Final[list[SearchResult]] = [ + result + for query in queries + for vector_store_id in vector_store_ids + for result in await _search_one_vector_store(vector_store_id=vector_store_id, query=query) + ] return queries, all_results @@ -172,78 +347,89 @@ async def _run_vector_searches( # --------------------------------------------------------------------------- -def _get_field(result: Any, key: str, default: Any = None) -> Any: - """Read a field from either a dict/TypedDict or an attribute-based object.""" +@dataclass(frozen=True, slots=True) +class SearchResultView: + """Normalized view over a dict- or attribute-shaped vector store search result.""" + + score: float | None + file_id: str + filename: str + text: str + attributes: Mapping[str, object] + + +def _content_text(item: VectorStoreResultContent | SearchResultContentLike) -> str: + text: Final = item.get("text") if isinstance(item, dict) else item.text + return text or "" + + +def _view_of(result: SearchResult) -> SearchResultView: if isinstance(result, dict): - return result.get(key, default) - return getattr(result, key, default) + score, file_id, filename = result.get("score"), result.get("file_id"), result.get("filename") + content, attributes = result.get("content"), result.get("attributes") + else: + score, file_id, filename = result.score, result.file_id, result.filename + content, attributes = result.content, result.attributes + text: Final = " ".join(chunk for chunk in (_content_text(item) for item in content or ()) if chunk) + return SearchResultView( + score=score, + file_id=file_id or "", + filename=filename or "", + text=text, + attributes=attributes or {}, + ) + + +def _result_header(index: int, view: SearchResultView) -> str: + segments: Final = ( + f"Result {index}", + view.filename or None, + f"file_id={view.file_id}" if view.file_id else None, + f"score={view.score:.3f}" if view.score is not None else None, + ) + return f"[{' | '.join(segment for segment in segments if segment)}]" def _format_search_results_as_tool_output( - results: list[VectorStoreSearchResult], + results: Sequence[SearchResult], ) -> str: """Serialize search results into a string to pass back as the tool's output.""" if not results: return "No results found in the vector store." - parts: Final[list[str]] = [] - for i, result in enumerate(results, 1): - score = _get_field(result, "score") - file_id = _get_field(result, "file_id") - filename = _get_field(result, "filename") - content_items = _get_field(result, "content") or [] - text_chunks = [c.get("text", "") if isinstance(c, dict) else getattr(c, "text", "") for c in content_items] - text = " ".join(t for t in text_chunks if t) - - header = f"[Result {i}" - if filename: - header += f" | {filename}" - if file_id: - header += f" | file_id={file_id}" - if score is not None: - header += f" | score={score:.3f}" - header += "]" - - parts.append(f"{header}\n{text}") - - return "\n\n".join(parts) + views: Final = tuple(_view_of(result) for result in results) + return "\n\n".join(f"{_result_header(index, view)}\n{view.text}" for index, view in enumerate(views, 1)) def _build_search_results_for_include( - results: list[VectorStoreSearchResult], -) -> list[dict[str, Any]]: + results: Sequence[SearchResult], +) -> list[FileSearchResultEntry]: """ Convert VectorStoreSearchResult objects to the format expected in file_search_call.search_results (mirrors OpenAI's include= format). - All chunks are returned — no deduplication by file_id — matching the + All chunks are returned, with no deduplication by file_id, matching the behaviour of OpenAI's native file_search which surfaces every relevant chunk even when multiple chunks originate from the same document. """ - formatted: Final[list[dict[str, Any]]] = [] - for result in results: - file_id = _get_field(result, "file_id") or "" - content_items = _get_field(result, "content") or [] - text_chunks = [c.get("text", "") if isinstance(c, dict) else getattr(c, "text", "") for c in content_items] - text = " ".join(t for t in text_chunks if t) - formatted.append( - { - "file_id": file_id, - "filename": _get_field(result, "filename") or "", - "score": _get_field(result, "score"), - "text": text, - "attributes": _get_field(result, "attributes") or {}, - } + return [ + FileSearchResultEntry( + file_id=view.file_id, + filename=view.filename, + score=view.score, + text=view.text, + attributes=view.attributes, ) - return formatted + for view in (_view_of(result) for result in results) + ] def _build_file_search_call_output( call_id: str, queries: list[str], - results: list[VectorStoreSearchResult] | None = None, + results: Sequence[SearchResult] | None = None, include_search_results: bool = False, -) -> dict[str, Any]: +) -> FileSearchCallOutput: """Build the file_search_call output item (mirrors OpenAI's format). Args: @@ -253,85 +439,89 @@ def _build_file_search_call_output( include_search_results: Populate search_results when the caller passed ``include=["file_search_call.results"]``. """ - search_results = None - if include_search_results and results: - search_results = _build_search_results_for_include(results) - return { - "type": "file_search_call", - "id": call_id, - "status": "completed", - "queries": queries, - "search_results": search_results, - } + search_results: Final = ( + _build_search_results_for_include(results) if include_search_results and results else None + ) + return FileSearchCallOutput( + type="file_search_call", + id=call_id, + status="completed", + queries=queries, + search_results=search_results, + ) def _build_file_citation_annotations( - results: list[VectorStoreSearchResult], + results: Sequence[SearchResult], text: str, -) -> list[dict[str, Any]]: +) -> list[FileCitationAnnotation]: """ Build file_citation annotations for the text. Each result with a file_id gets a citation at the end of the text. """ - annotations: Final[list[dict[str, Any]]] = [] - index: Final = len(text) # cite at end of text block - seen_file_ids: Final[set] = set() - - for result in results: - file_id = _get_field(result, "file_id") - filename = _get_field(result, "filename") - if not file_id or file_id in seen_file_ids: - continue - seen_file_ids.add(file_id) - annotations.append( - { - "type": "file_citation", - "index": index, - "file_id": file_id, - "filename": filename or "", - } + views: Final = tuple(_view_of(result) for result in results) + ordered_file_ids: Final = tuple(dict.fromkeys(view.file_id for view in views if view.file_id)) + filename_by_file_id: Final = {view.file_id: view.filename for view in reversed(views) if view.file_id} + return [ + FileCitationAnnotation( + type="file_citation", + index=len(text), + file_id=file_id, + filename=filename_by_file_id.get(file_id, ""), ) - - return annotations + for file_id in ordered_file_ids + ] def _build_message_output( response_text: str, - results: list[VectorStoreSearchResult], -) -> dict[str, Any]: + results: Sequence[SearchResult], +) -> MessageOutput: """Build the message output item with optional file_citation annotations.""" - annotations: Final = _build_file_citation_annotations(results, response_text) - return { - "type": "message", - "role": "assistant", - "content": [ - { - "type": "output_text", - "text": response_text, - "annotations": annotations, - } + return MessageOutput( + type="message", + role="assistant", + content=[ + OutputTextContent( + type="output_text", + text=response_text, + annotations=_build_file_citation_annotations(results, response_text), + ) ], - } + ) + + +def _output_items_of(response: ResponsesAPIResponse) -> tuple[object, ...]: + return tuple(response.output) def _extract_text_from_responses_output(response: ResponsesAPIResponse) -> str: """Pull the assistant's text from the provider's response.""" - for item in response.output: - item_type = item.get("type") if isinstance(item, dict) else getattr(item, "type", None) - if item_type == "message": - content = item.get("content") if isinstance(item, dict) else getattr(item, "content", []) - for block in content or []: - block_type = block.get("type") if isinstance(block, dict) else getattr(block, "type", None) - if block_type == "output_text": - raw = block.get("text") if isinstance(block, dict) else getattr(block, "text", "") - return str(raw) if raw is not None else "" - return "" + messages: Final = tuple( + message + for message in (_validate_or_none(MessageOutputItem, item) for item in _output_items_of(response)) + if message is not None + ) + texts: Final = tuple( + block.text or "" for message in messages for block in message.content or () if block.type == "output_text" + ) + return texts[0] if texts else "" + + +def _response_cost(hidden_params: Mapping[str, object]) -> float | None: + cost: Final = hidden_params.get("response_cost") + return float(cost) if isinstance(cost, (int, float)) else None + + +def _hidden_params_of(response: ResponsesAPIResponse) -> Mapping[str, object]: + raw: Final[Mapping[str, object] | None] = getattr(response, "_hidden_params", None) + return raw or {} def _synthesize_responses_api_response( original_response: ResponsesAPIResponse, - file_search_call_output: dict[str, Any], - message_output: dict[str, Any], + file_search_call_output: FileSearchCallOutput, + message_output: MessageOutput, first_response: ResponsesAPIResponse | None = None, ) -> ResponsesAPIResponse: """ @@ -343,30 +533,25 @@ def _synthesize_responses_api_response( synthesized _hidden_params so that billing callbacks see the total cost of both provider calls that the emulated flow makes. """ - synthesized_output: Final[list[dict[str, Any]]] = [file_search_call_output, message_output] + synthesized_output: Final[list[ResponseOutputItem | dict[str, object]]] = [ + dict(file_search_call_output), + dict(message_output), + ] synthesized: Final = ResponsesAPIResponse( - id=getattr(original_response, "id", f"resp_{uuid.uuid4().hex}"), + id=original_response.id, object="response", - created_at=getattr(original_response, "created_at", int(time.time())), + created_at=original_response.created_at, status="completed", - model=getattr(original_response, "model", ""), - output=cast(list[ResponseOutputItem | dict[str, Any]], synthesized_output), - usage=getattr(original_response, "usage", None), + model=original_response.model, + output=synthesized_output, + usage=original_response.usage, error=None, ) - if hasattr(original_response, "_hidden_params"): - hidden: Final = dict(getattr(original_response, "_hidden_params") or {}) - if first_response is not None and hasattr(first_response, "_hidden_params"): - first_hidden: Final = getattr(first_response, "_hidden_params") or {} - first_cost: Final = ( - first_hidden.get("response_cost") - if isinstance(first_hidden, dict) - else getattr(first_hidden, "response_cost", None) - ) - if first_cost is not None: - current_cost: Final = hidden.get("response_cost") if isinstance(hidden, dict) else 0 - hidden["response_cost"] = (current_cost or 0) + first_cost - synthesized._hidden_params = hidden + hidden: Final = _hidden_params_of(original_response) + first_cost: Final = _response_cost(_hidden_params_of(first_response)) if first_response is not None else None + synthesized._hidden_params = ( + {**hidden} if first_cost is None else {**hidden, "response_cost": (_response_cost(hidden) or 0) + first_cost} + ) return synthesized @@ -375,127 +560,118 @@ def _synthesize_responses_api_response( # --------------------------------------------------------------------------- -async def _call_aresponses(input, model, tools, **kwargs): # pragma: no cover – thin wrapper for patching in tests +# any-ok: thin, test-patchable seam forwarding provider pass-through params to aresponses +async def _call_aresponses(input, model: str, tools, **kwargs) -> ResponsesAPIResponse: from litellm.responses.main import aresponses - return await aresponses(input=input, model=model, tools=tools, **kwargs) + response: Final = await aresponses(input=input, model=model, tools=tools, **kwargs) + if isinstance(response, ResponsesAPIResponse): + return response + raise ValueError("emulated file_search does not support streaming responses") def _prepare_emulated_file_search_call( - kwargs: dict[str, Any], -) -> tuple[bool, dict[str, Any]]: - include_items: Final[list[str]] = list(kwargs.get("include") or []) - include_search_results: Final = "file_search_call.results" in include_items + kwargs: Mapping[str, object], +) -> tuple[bool, dict[str, object]]: + include_options: Final = _validate_or_none(IncludeOptions, kwargs) + include_search_results: Final = "file_search_call.results" in ((include_options and include_options.include) or ()) - original_stream: Final = kwargs.get("stream") - updated_kwargs = kwargs - if original_stream: - verbose_logger.debug( - "Streaming is not yet supported for emulated file_search. Disabling stream for this request." - ) - updated_kwargs = {**kwargs, "stream": False} + if not kwargs.get("stream"): + return include_search_results, {**kwargs} - return include_search_results, updated_kwargs + verbose_logger.debug("Streaming is not yet supported for emulated file_search. Disabling stream for this request.") + return include_search_results, {**kwargs, "stream": False} -def _extract_tool_call_fields(tool_call: Any, fallback_call_id: str) -> tuple[str, str]: - """Extract (call_id, raw_arguments_string) from a dict or Pydantic tool_call item.""" - if isinstance(tool_call, dict): - call_id = str(tool_call.get("call_id") or tool_call.get("id") or fallback_call_id) - raw_args = tool_call.get("arguments") or "{}" - else: - raw_call_id: Final = getattr(tool_call, "call_id", None) or getattr(tool_call, "id", None) or fallback_call_id - call_id = str(raw_call_id) - raw_args = getattr(tool_call, "arguments", "{}") or "{}" - return call_id, raw_args - - -def _resolve_queries_from_args(args: dict[str, Any], input: Any) -> list[str]: +def _resolve_queries_from_args(args: FileSearchArguments, input: str | ResponseInputParam) -> list[str]: """Pull the queries list out of parsed tool-call arguments, with backward-compat fallbacks.""" - queries_from_call: Final = args.get("queries") - if not queries_from_call: - # Fallback: check for single "query" field (backward compat) - single_query: Final = args.get("query") - return [single_query] if single_query else [str(input)] - if not isinstance(queries_from_call, list): - return [str(queries_from_call)] - return queries_from_call + if not args.queries: + return [args.query] if args.query else [str(input)] + if isinstance(args.queries, str): + return [args.queries] + return list(args.queries) + + +def _parse_file_search_arguments(raw_arguments: str | Mapping[str, object] | None) -> FileSearchArguments: + if isinstance(raw_arguments, str): + try: + return FileSearchArguments.model_validate_json(raw_arguments) + except ValidationError: + return FileSearchArguments() + return _validate_or_none(FileSearchArguments, raw_arguments) or FileSearchArguments() async def _execute_file_search_tool_calls( - file_search_calls: list[Any], + file_search_calls: Sequence[FileSearchFunctionCall], all_vs_ids: list[str], - input: Any, + input: str | ResponseInputParam, file_search_call_id: str, -) -> tuple[list[dict[str, Any]], list[str], list[VectorStoreSearchResult]]: +) -> tuple[list[FunctionCallOutput], list[str], list[SearchResult]]: """Run the vector search for each file_search tool_call and collect results.""" - tool_results: Final[list[dict[str, Any]]] = [] + tool_results: Final[list[FunctionCallOutput]] = [] all_queries: Final[list[str]] = [] - all_results: Final[list[VectorStoreSearchResult]] = [] + all_results: Final[list[SearchResult]] = [] for tool_call in file_search_calls: - call_id, raw_args = _extract_tool_call_fields(tool_call, fallback_call_id=file_search_call_id) - - try: - args = json.loads(raw_args) if isinstance(raw_args, str) else raw_args - except json.JSONDecodeError: - args = {} - - queries_from_call = _resolve_queries_from_args(args, input) - - vs_id_arg = args.get("vector_store_id") - vs_ids_for_call = [vs_id_arg] if vs_id_arg else all_vs_ids - + call_id = str(tool_call.call_id or tool_call.id or file_search_call_id) + args = _parse_file_search_arguments(tool_call.arguments) queries, results = await _run_vector_searches( - queries=queries_from_call, - vector_store_ids=vs_ids_for_call, + queries=_resolve_queries_from_args(args, input), + vector_store_ids=[args.vector_store_id] if args.vector_store_id else all_vs_ids, ) all_queries.extend(queries) all_results.extend(results) tool_results.append( - { - "type": "function_call_output", - "call_id": call_id, - "output": _format_search_results_as_tool_output(results), - } + FunctionCallOutput( + type="function_call_output", + call_id=call_id, + output=_format_search_results_as_tool_output(results), + ) ) return tool_results, all_queries, all_results +def _as_plain_input_item(item: object) -> object: + return item.model_dump(exclude_none=True) if isinstance(item, BaseModel) else item + + def _build_follow_up_input( - input: Any, + input: str | ResponseInputParam, first_response: ResponsesAPIResponse, - tool_results: list[dict[str, Any]], -) -> list[Any]: + tool_results: Sequence[FunctionCallOutput], +) -> list[object]: """Assemble the follow-up call input: original messages + first-response output + tool results. Including all output items (text blocks, reasoning, non-file-search calls) ensures providers like Anthropic that emit text before the tool call have complete conversation context. Serializes Pydantic model instances to plain dicts so the transformation layer can call .get(). """ - original_input_items: Final = ( + original_input_items: Final[list[object]] = ( list(input) if isinstance(input, (list, tuple)) else [{"role": "user", "content": str(input)}] ) - first_response_output_items: Final[list[Any]] = [] - for _item in first_response.output: - if isinstance(_item, dict): - first_response_output_items.append(_item) - elif hasattr(_item, "model_dump"): - first_response_output_items.append(_item.model_dump(exclude_none=True)) - else: - first_response_output_items.append(_item) + return [ + *original_input_items, + *(_as_plain_input_item(item) for item in _output_items_of(first_response)), + *tool_results, + ] - return original_input_items + first_response_output_items + tool_results + +def _file_search_calls_in(response: ResponsesAPIResponse) -> tuple[FileSearchFunctionCall, ...]: + return tuple( + call + for call in (_validate_or_none(FileSearchFunctionCall, item) for item in _output_items_of(response)) + if call is not None and call.name == FILE_SEARCH_FUNCTION_NAME + ) async def aresponses_with_emulated_file_search( - input: Any, + input: str | ResponseInputParam, model: str, tools: Iterable[ToolParam] | None = None, # Pass-through params — forwarded as-is to the underlying aresponses call - **kwargs: Any, + **kwargs: object, ) -> ResponsesAPIResponse: """ Emulated file_search for providers that don't support it natively. @@ -504,7 +680,7 @@ async def aresponses_with_emulated_file_search( runs vector search, and synthesizes an OpenAI-format response. """ # Determine whether caller wants search_results populated in the output. - _include_search_results, kwargs = _prepare_emulated_file_search_call(kwargs=kwargs) + include_search_results, forwarded_kwargs = _prepare_emulated_file_search_call(kwargs=kwargs) # 1. Replace file_search tools with function tool transformed_tools, all_vs_ids = _replace_file_search_tools(tools) @@ -512,42 +688,25 @@ async def aresponses_with_emulated_file_search( # 2. First provider call — provider will call the file_search function. # Mark as an internal sub-call so wrapper_async skips billing callbacks; # the parent litellm_logging_obj (propagated via kwargs) fires once at the end. - _prev_internal: Final = is_internal_call.get() + prev_internal: Final = is_internal_call.get() is_internal_call.set(True) try: - first_response: Final[ResponsesAPIResponse] = cast( - ResponsesAPIResponse, - await _call_aresponses( - input=input, - model=model, - tools=transformed_tools or None, - **kwargs, - ), + first_response: Final = await _call_aresponses( + input=input, + model=model, + tools=transformed_tools or None, + **forwarded_kwargs, ) finally: - is_internal_call.set(_prev_internal) + is_internal_call.set(prev_internal) # 3. Look for a file_search function_call in the output - file_search_calls: Final = [ - item - for item in first_response.output - if ( - isinstance(item, dict) - and item.get("type") == "function_call" - and item.get("name") == FILE_SEARCH_FUNCTION_NAME - ) - or ( - hasattr(item, "type") - and getattr(item, "type") == "function_call" - and getattr(item, "name", None) == FILE_SEARCH_FUNCTION_NAME - ) - ] + file_search_calls: Final = _file_search_calls_in(first_response) if not file_search_calls: # Provider answered without calling the tool (e.g. it had enough context). # Return as-is wrapped in OpenAI format. call_id: Final = f"fs_{uuid.uuid4().hex[:24]}" - response_text = _extract_text_from_responses_output(first_response) return _synthesize_responses_api_response( original_response=first_response, file_search_call_output=_build_file_search_call_output( @@ -556,7 +715,7 @@ async def aresponses_with_emulated_file_search( results=None, include_search_results=False, ), - message_output=_build_message_output(response_text, []), + message_output=_build_message_output(_extract_text_from_responses_output(first_response), []), ) # 4. Execute each file_search tool call @@ -579,29 +738,24 @@ async def aresponses_with_emulated_file_search( # Also an internal sub-call; billing is suppressed so the outer call fires once. is_internal_call.set(True) try: - final_response: Final[ResponsesAPIResponse] = cast( - ResponsesAPIResponse, - await _call_aresponses( - input=follow_up_input, - model=model, - tools=None, # no tools needed for the answer step - **kwargs, - ), + final_response: Final = await _call_aresponses( + input=follow_up_input, + model=model, + tools=None, # no tools needed for the answer step + **forwarded_kwargs, ) finally: - is_internal_call.set(_prev_internal) + is_internal_call.set(prev_internal) # 7. Synthesize OpenAI-format output - response_text = _extract_text_from_responses_output(final_response) - return _synthesize_responses_api_response( original_response=final_response, file_search_call_output=_build_file_search_call_output( call_id=file_search_call_id, queries=all_queries or [str(input)], results=all_results, - include_search_results=_include_search_results, + include_search_results=include_search_results, ), - message_output=_build_message_output(response_text, all_results), + message_output=_build_message_output(_extract_text_from_responses_output(final_response), all_results), first_response=first_response, )