diff --git a/litellm/integrations/otel/model/payloads.py b/litellm/integrations/otel/model/payloads.py index 7ecbeea255c..240107251c5 100644 --- a/litellm/integrations/otel/model/payloads.py +++ b/litellm/integrations/otel/model/payloads.py @@ -760,6 +760,8 @@ def _output_choices(response: Mapping[str, object]) -> tuple[Mapping[str, object or _ocr_choices(response) or _transcription_choices(response) or _moderation_choices(response) + or _rerank_choices(response) + or _search_choices(response) or _image_choices(response) or _binary_choices(response) ) @@ -815,6 +817,32 @@ def _moderation_verdict(flagged: bool, categories: object) -> str: return f"flagged: {', '.join(hits)}" if hits else "flagged" +def _rerank_choices(response: Mapping[str, object]) -> tuple[_Choice, ...]: + return _joined_choice( + tuple( + _rerank_line(index, score, result.get("document")) + for result in _dicts(response.get("results")) + if (index := as_int(result.get("index"))) is not None + if (score := as_float(result.get("relevance_score"))) is not None + ) + ) + + +def _rerank_line(index: int, score: float, document: object) -> str: + text: Final = as_str((as_str_mapping(document) or {}).get("text")) + return f"[{index}] {score}\n{text}" if text else f"[{index}] {score}" + + +def _search_choices(response: Mapping[str, object]) -> tuple[_Choice, ...]: + return _joined_choice( + tuple( + line + for result in _dicts(response.get("results")) + if (line := "\n".join(part for key in ("title", "url", "snippet") if (part := as_str(result.get(key))))) + ) + ) + + def _image_choices(response: Mapping[str, object]) -> tuple[_Choice, ...]: return _joined_choice( tuple(summary for item in _dicts(response.get("data")) if (summary := _image_summary(item)) is not None) diff --git a/litellm/utils.py b/litellm/utils.py index 8ade6022e3e..8ce83f475a4 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -35,6 +35,7 @@ from importlib import resources from inspect import iscoroutine from io import StringIO from os.path import abspath, dirname, join +from pathlib import PurePath from types import MappingProxyType import dotenv @@ -288,7 +289,7 @@ except (ImportError, AttributeError, TypeError): claude_json_str = json.dumps(json_data) import importlib.metadata from collections.abc import AsyncIterator, Callable, Iterable, Iterator, Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Union, cast, get_args +from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, cast, runtime_checkable from typing_extensions import assert_never @@ -886,6 +887,32 @@ async def _run_success_deployment_hook_on_converted_chat_stream( ) +@runtime_checkable +class _NamedFile(Protocol): + @property + def name(self) -> object: ... + + +def _ocr_document_summary(document: object) -> str: + if not isinstance(document, Mapping): + return "default-message-value" + doc: Final = cast(Mapping[str, object], document) # cast-ok: ocr()/aocr() type the document as Mapping[str, object] + location: Final = doc.get("document_url", doc.get("image_url")) + if isinstance(location, str): + header, separator, payload = location.partition(",") + return f"{header} ({len(payload)} chars)" if separator and header.startswith("data:") else location + file_input: Final = doc.get("file") + mime_type: Final = doc.get("mime_type") + kind: Final = f"file ({mime_type})" if isinstance(mime_type, str) else "file" + if isinstance(file_input, PurePath): + return f"{kind} {file_input.name}" + if isinstance(file_input, bytes): + return f"{kind} {len(file_input)} bytes" + if isinstance(file_input, _NamedFile) and isinstance(file_input.name, str): + return f"{kind} {PurePath(file_input.name).name}" + return kind + + # Runs once per call to check if the user wants to send their data anywhere - PostHog/Sentry/Slack/etc. def function_setup( original_function: str, @@ -1126,6 +1153,17 @@ def function_setup( messages = args[0] if len(args) > 0 else kwargs["prompt"] elif call_type == CallTypes.rerank.value or call_type == CallTypes.arerank.value: messages = kwargs.get("query") + elif call_type in (CallTypes.search.value, CallTypes.asearch.value): + search_query: Final = args[0] if len(args) > 0 else kwargs.get("query") + messages = ( + "\n".join(part for part in search_query if isinstance(part, str)) + if isinstance(search_query, list) + else search_query + ) + elif call_type in (CallTypes.image_edit.value, CallTypes.aimage_edit.value): + messages = args[1] if len(args) > 1 else kwargs.get("prompt") + elif call_type in (CallTypes.ocr.value, CallTypes.aocr.value): + messages = _ocr_document_summary(args[1] if len(args) > 1 else kwargs.get("document")) elif call_type == CallTypes.atranscription.value or call_type == CallTypes.transcription.value: _file_obj: Final[FileTypes] = args[1] if len(args) > 1 else kwargs["file"] # Lazy import audio_utils.utils only when needed for transcription calls diff --git a/tests/e2e/logging/test_otel_v2_langfuse_generation_output_e2e.py b/tests/e2e/logging/test_otel_v2_langfuse_generation_output_e2e.py index 71008d94557..9e08cf21da0 100644 --- a/tests/e2e/logging/test_otel_v2_langfuse_generation_output_e2e.py +++ b/tests/e2e/logging/test_otel_v2_langfuse_generation_output_e2e.py @@ -1,13 +1,17 @@ -"""Live e2e: the OTel v2 Langfuse generation carries output for every non-chat endpoint (LIT-8309). +"""Live e2e: the OTel v2 Langfuse generation carries input and output for every non-chat endpoint. With LITELLM_OTEL_V2=true the proxy exports one generation per request to the -team's Langfuse destination. Chat, Responses, embeddings and OCR already fill -its output; this file pins the remaining five families. Each test registers a -real OpenAI deployment, drives the endpoint through the shared transport, then -reads the generation back from Langfuse and asserts its output reflects what -the caller received: the completion text, the transcript, the moderation -verdict, and for images and speech a bounded summary that never carries the -raw base64 or audio bytes. +team's Langfuse destination. Chat, Responses and embeddings already fill both +panels; this file pins the other families. LIT-8309 covers the output of +completions, images, speech, transcription and moderation; LIT-8326 covers the +rerank and search output and the OCR, image-edit and search input, which used +to read the `default-message-value` placeholder. Each test registers a real +deployment, drives the endpoint through the shared transport, then reads the +generation back from Langfuse and asserts it reflects what the caller sent and +received: the completion text, the transcript, the moderation verdict, the +ranked rerank indices and scores, the search results, the OCR document URL, +the edit prompt, and for images, speech and uploaded documents a bounded +summary that never carries the raw base64 or audio bytes. """ from __future__ import annotations @@ -25,11 +29,22 @@ from logging_client import LangfuseCreds, LangfuseObservation, LoggingClient, lo from models import ( CompletionBody, CompletionResponse, + ImageEditForm, ImageGenerationBody, ImageGenerationResponse, LiteLLMParamsBody, ModerationBody, ModerationResponse, + OcrBody, + OcrDocument, + OcrForm, + OcrResponse, + RerankBody, + SearchBody, + SearchResponse, + SearchToolBody, + SearchToolCreateBody, + SearchToolLiteLLMParamsBody, SpeechBody, TranscriptionForm, TranscriptionResponse, @@ -41,7 +56,19 @@ pytestmark = [pytest.mark.e2e, pytest.mark.otel_v2] WEATHER_WAV: Final = ( Path(__file__).resolve().parent.parent / "llm_translation" / "realtime" / "fixtures" / "weather_question_24k.wav" ) +DUMMY_PDF: Final = Path(__file__).resolve().parent.parent.parent / "llm_translation" / "fixtures" / "dummy.pdf" +DUMMY_PDF_URL: Final = ( + "https://cdn.jsdelivr.net/gh/BerriAI/litellm" + "@d769e81c90d453240c61fc572cdb27fae06a89d0" + "/tests/llm_translation/fixtures/dummy.pdf" +) +RED_SQUARE_PNG: Final = base64.b64decode( + "iVBORw0KGgoAAAANSUhEUgAAAEAAAABACAIAAAAlC+aJAAAAS0lEQVR42u3PMQ0AAAwDoPo3" + "3UrYvQQckD4XAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEB" + "AYHLAMpT0sIcNbcEAAAAAElFTkSuQmCC" +) BOUNDED_OUTPUT_CHARS: Final = 1024 +PLACEHOLDER_INPUT: Final = "default-message-value" class _OutputMessage(BaseModel): @@ -50,7 +77,14 @@ class _OutputMessage(BaseModel): content: str = "" +class _InputMessage(BaseModel): + """One user message of the Langfuse generation input; only the text is read.""" + + content: str = "" + + _OUTPUT_MESSAGES: Final = TypeAdapter(list[_OutputMessage]) +_INPUT_MESSAGES: Final = TypeAdapter(list[_InputMessage]) @pytest.fixture(scope="session") @@ -91,10 +125,50 @@ def _output_text(observation: LangfuseObservation) -> str: return "\n".join(message.content for message in messages) +def _input_text(observation: LangfuseObservation) -> str: + assert observation.input not in (None, "", [], {}), f"generation input is empty: {observation!r}" + try: + messages: Final = _INPUT_MESSAGES.validate_python(observation.input) + except ValidationError: + pytest.fail(f"generation input is not a list of user messages: {observation!r}") + assert messages, f"generation input is empty: {observation!r}" + text: Final = "\n".join(message.content for message in messages) + assert text != PLACEHOLDER_INPUT, f"generation input is the placeholder, not the request: {observation!r}" + return text + + def _openai(model: str) -> LiteLLMParamsBody: return LiteLLMParamsBody(model=model, api_key="os.environ/OPENAI_API_KEY") +def _mistral_ocr() -> LiteLLMParamsBody: + return LiteLLMParamsBody(model="mistral/mistral-ocr-latest", api_key="os.environ/MISTRAL_API_KEY") + + +def _langfuse_search_tool( + client: LoggingClient, creds: LangfuseCreds, resources: ResourceManager +) -> tuple[str, str, str]: + """A keyless DuckDuckGo search tool registered for this run, plus a key on a team whose Langfuse callback is + `creds`.""" + tool: Final = f"e2e-otel-search-{unique_marker()}" + tool_id: Final = client.proxy.create_search_tool( + SearchToolCreateBody( + search_tool=SearchToolBody( + search_tool_name=tool, + litellm_params=SearchToolLiteLLMParamsBody(search_provider="duckduckgo"), + ) + ) + ) + resources.defer(lambda: client.proxy.delete_search_tool(tool_id)) + team_id: Final = client.create_team(f"otel-search-team-{unique_marker()}", models=[tool]) + resources.defer(lambda: client.delete_team(team_id)) + client.add_team_langfuse_callback(team_id, creds) + alias: Final = f"otel-search-key-{unique_marker()}" + key: Final = client.key_with_alias(alias, models=[tool], team_id=team_id) + resources.defer(lambda: client.delete_key(key)) + return tool, key, alias + + class TestOtelV2LangfuseGenerationOutput: @pytest.mark.covers("logging.langfuse.success.logs_spend", exercised_on=["completions"]) def test_completions_output_is_the_completion_text( @@ -208,3 +282,134 @@ class TestOtelV2LangfuseGenerationOutput: assert output.startswith(verdict), ( f"generation output does not carry the moderation verdict {verdict!r}: {output!r}" ) + + @pytest.mark.covers("logging.langfuse.success.logs_spend", exercised_on=["rerank"]) + def test_rerank_output_is_the_ranked_indices_and_scores( + self, client: LoggingClient, langfuse_creds: LangfuseCreds, resources: ResourceManager + ) -> None: + model, key, alias = _langfuse_key( + client, + langfuse_creds, + resources, + LiteLLMParamsBody(model="cohere/rerank-v4.0-fast", api_key="os.environ/COHERE_API_KEY"), + ) + query: Final = f"What is the capital of France? {unique_marker()}" + started: Final = datetime.now(timezone.utc) + response: Final = unwrap( + client.proxy.rerank( + key, + RerankBody( + model=model, + query=query, + documents=["Paris is the capital of France.", "Berlin is in Germany.", "Bananas are yellow."], + top_n=2, + ), + ) + ) + ranked: Final = tuple(f"[{item.index}] {item.relevance_score}" for item in response.results) + assert len(ranked) == 2 and all(item.index is not None for item in response.results), ( + f"/v1/rerank returned no ranked results: {response!r}" + ) + + generation: Final = _generation(client, langfuse_creds, alias=alias, started=started) + assert _input_text(generation) == query + output: Final = _output_text(generation) + assert output == "\n\n".join(ranked), f"generation output is not the ranked indices and scores: {output!r}" + + @pytest.mark.covers("logging.langfuse.success.logs_spend", exercised_on=["ocr"]) + def test_ocr_input_is_the_document_url( + self, client: LoggingClient, langfuse_creds: LangfuseCreds, resources: ResourceManager + ) -> None: + model, key, alias = _langfuse_key(client, langfuse_creds, resources, _mistral_ocr()) + started: Final = datetime.now(timezone.utc) + response: Final = unwrap( + client.proxy.ocr( + key, OcrBody(model=model, document=OcrDocument(type="document_url", document_url=DUMMY_PDF_URL)) + ) + ) + assert response.pages and response.pages[0].markdown, f"/v1/ocr returned no page markdown: {response!r}" + + generation: Final = _generation(client, langfuse_creds, alias=alias, started=started) + assert _input_text(generation) == DUMMY_PDF_URL + assert response.pages[0].markdown in _output_text(generation) + + @pytest.mark.covers("logging.langfuse.success.logs_spend", exercised_on=["ocr"]) + def test_ocr_upload_input_is_a_bounded_document_summary_without_base64( + self, client: LoggingClient, langfuse_creds: LangfuseCreds, resources: ResourceManager + ) -> None: + model, key, alias = _langfuse_key(client, langfuse_creds, resources, _mistral_ocr()) + pdf: Final = DUMMY_PDF.read_bytes() + started: Final = datetime.now(timezone.utc) + response: Final = unwrap( + client.proxy.transport.upload( + "/v1/ocr", + headers=client.proxy.transport.bearer(key), + form=OcrForm(model=model), + filename=DUMMY_PDF.name, + content=pdf, + file_content_type="application/pdf", + response_type=OcrResponse, + ) + ) + assert response.pages and response.pages[0].markdown, f"/v1/ocr returned no page markdown: {response!r}" + + text: Final = _input_text(_generation(client, langfuse_creds, alias=alias, started=started)) + encoded: Final = base64.b64encode(pdf).decode() + assert text == f"data:application/pdf;base64 ({len(encoded)} chars)", ( + f"OCR upload input is not the bounded document summary: {text!r}" + ) + + @pytest.mark.covers("logging.langfuse.success.logs_spend", exercised_on=["images_edits"]) + def test_image_edit_input_is_the_edit_prompt( + self, client: LoggingClient, langfuse_creds: LangfuseCreds, resources: ResourceManager + ) -> None: + model, key, alias = _langfuse_key(client, langfuse_creds, resources, _openai("openai/gpt-image-1-mini")) + prompt: Final = f"make the square blue {unique_marker()}" + started: Final = datetime.now(timezone.utc) + response: Final = unwrap( + client.proxy.transport.upload( + "/v1/images/edits", + headers=client.proxy.transport.bearer(key), + form=ImageEditForm(model=model, prompt=prompt), + filename="red_square.png", + content=RED_SQUARE_PNG, + file_content_type="image/png", + file_field="image", + response_type=ImageGenerationResponse, + timeout=180.0, + ) + ) + assert response.data and (response.data[0].b64_json or response.data[0].url), ( + f"/v1/images/edits returned no image: {response!r}" + ) + + generation: Final = _generation(client, langfuse_creds, alias=alias, started=started) + assert _input_text(generation) == prompt + assert _output_text(generation).startswith("b64_json image (") + + @pytest.mark.covers("logging.langfuse.success.logs_spend", exercised_on=["search"]) + def test_search_input_is_the_query_and_output_the_results( + self, client: LoggingClient, langfuse_creds: LangfuseCreds, resources: ResourceManager + ) -> None: + tool, key, alias = _langfuse_search_tool(client, langfuse_creds, resources) + query: Final = "Eiffel Tower" + started: Final = datetime.now(timezone.utc) + response: Final = unwrap( + client.proxy.transport.post( + f"/v1/search/{tool}", + headers=client.proxy.transport.bearer(key), + json=SearchBody(query=query, max_results=2), + response_type=SearchResponse, + ) + ) + assert response.results and all(result.url for result in response.results), ( + f"/v1/search returned no results with a url: {response!r}" + ) + + generation: Final = _generation(client, langfuse_creds, alias=alias, started=started) + assert _input_text(generation) == query + output: Final = _output_text(generation) + assert output == "\n\n".join( + "\n".join(part for part in (result.title, result.url, result.snippet) if part) + for result in response.results + ), f"generation output is not the search results the caller received: {output!r}" diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 53fbe03c025..5278dbd287a 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -749,6 +749,12 @@ class OcrBody(BaseModel): document: OcrDocument +class OcrForm(BaseModel): + """Multipart /v1/ocr form fields; the document travels as the `file` part.""" + + model: str + + class OcrPage(BaseModel): index: int markdown: str @@ -798,6 +804,51 @@ class ImageGenerationResponse(BaseModel): data: list[ImageDatum] = [] +class ImageEditForm(BaseModel): + """POST /v1/images/edits form fields; the image travels as the `image` multipart part.""" + + model: str + prompt: str + size: str = "1024x1024" + quality: str = "low" + + +class SearchBody(BaseModel): + """POST /v1/search/{search_tool_name} body (Perplexity-compatible).""" + + query: str + max_results: int = 2 + + +class SearchResultItem(BaseModel): + title: str = "" + url: str = "" + snippet: str = "" + + +class SearchResponse(BaseModel): + results: list[SearchResultItem] = [] + + +class SearchToolLiteLLMParamsBody(BaseModel): + search_provider: str + + +class SearchToolBody(BaseModel): + search_tool_name: str + litellm_params: SearchToolLiteLLMParamsBody + + +class SearchToolCreateBody(BaseModel): + """POST /search_tools body: the tool as it would sit under `search_tools:` in the config.""" + + search_tool: SearchToolBody + + +class SearchToolCreateResponse(BaseModel): + search_tool_id: str + + # ---------- audio ---------- diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index 76c57452c69..f1981e0aa5e 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -85,6 +85,8 @@ from models import ( RerankResponse, RouterCurrentValues, RouterSettingsResponse, + SearchToolCreateBody, + SearchToolCreateResponse, SpendLogRow, SpendLogs, SpendLogsPage, @@ -859,6 +861,30 @@ class ProxyClient: response_type=NoBody, ) + def create_search_tool(self, body: SearchToolCreateBody) -> str: + """POST /search_tools: register a search tool on the running proxy and return its id + once every worker has had a config-reload window to pick it up from the DB.""" + search_tool_id: Final = unwrap( + self.transport.post( + "/search_tools", + headers=self.management_headers(), + json=body, + response_type=SearchToolCreateResponse, + ) + ).search_tool_id + settle_propagation(time.monotonic()) + return search_tool_id + + def delete_search_tool(self, search_tool_id: str) -> None: + result = self.transport.delete( + f"/search_tools/{search_tool_id}", + headers=self.management_headers(), + json=NoBody(), + response_type=NoBody, + ) + if not is_ok(result): + warnings.warn(f"delete_search_tool({search_tool_id!r}) failed: {result}", stacklevel=2) + def create_credential(self, body: CredentialCreateBody) -> None: unwrap( self.transport.post( diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py b/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py index 70ab4fe07de..b9ec625fdfc 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py @@ -1024,6 +1024,94 @@ def test_moderation_results_without_a_verdict_produce_no_output() -> None: assert data.choices_out == () +def test_rerank_results_become_ranked_indices_and_scores_with_the_document_text() -> None: + data: Final = LLMCallSpanData.from_standard_logging_payload( + _route_payload( + "arerank", + "rerank-v4.0-fast", + { + "id": "rr-1", + "results": [ + {"index": 2, "relevance_score": 0.91, "document": {"text": "Paris is the capital of France."}}, + {"index": 0, "relevance_score": 0.07}, + {"index": 1, "relevance_score": 0.02, "document": "not-a-document"}, + ], + "meta": {"billed_units": {"search_units": 1}}, + }, + ), + capture_content=True, + ) + + assert data.choices_out == (_assistant_choice("[2] 0.91\nParis is the capital of France.\n\n[0] 0.07\n\n[1] 0.02"),) + assert data.finish_reasons == () + assert data.response_id == "rr-1" + + +def test_rerank_output_follows_the_content_capture_gate() -> None: + data: Final = LLMCallSpanData.from_standard_logging_payload( + _route_payload("arerank", "rerank-v4.0-fast", {"results": [{"index": 0, "relevance_score": 0.5}]}) + ) + + assert data.choices_out == () + + +def test_rerank_results_without_an_index_and_score_produce_no_output() -> None: + data: Final = LLMCallSpanData.from_standard_logging_payload( + _route_payload( + "arerank", + "rerank-v4.0-fast", + {"results": [{"index": 0, "document": {"text": "x"}}, {"relevance_score": 0.5}, "not-a-result"]}, + ), + capture_content=True, + ) + + assert data.choices_out == () + + +def test_search_results_become_title_url_and_snippet_blocks_in_result_order() -> None: + data: Final = LLMCallSpanData.from_standard_logging_payload( + _route_payload( + "asearch", + "exa-search", + { + "object": "search", + "results": [ + {"title": "Eiffel Tower", "url": "https://example.com/eiffel", "snippet": "A lattice tower."}, + {"url": "https://example.com/bare", "date": "2024-01-01"}, + {"title": "no url", "snippet": "kept"}, + ], + }, + ), + capture_content=True, + ) + + assert data.choices_out == ( + _assistant_choice( + "Eiffel Tower\nhttps://example.com/eiffel\nA lattice tower.\n\nhttps://example.com/bare\n\nno url\nkept" + ), + ) + assert data.finish_reasons == () + + +def test_search_output_follows_the_content_capture_gate() -> None: + data: Final = LLMCallSpanData.from_standard_logging_payload( + _route_payload("asearch", "exa-search", {"results": [{"url": "https://example.com"}]}) + ) + + assert data.choices_out == () + + +def test_search_results_without_any_text_field_produce_no_output() -> None: + data: Final = LLMCallSpanData.from_standard_logging_payload( + _route_payload( + "asearch", "exa-search", {"results": [{"date": "2024-01-01"}, {"title": "", "url": None}, "not-a-result"]} + ), + capture_content=True, + ) + + assert data.choices_out == () + + def test_image_data_becomes_a_size_summary_and_never_carries_the_base64_payload() -> None: encoded: Final = "QUJDRA==" data: Final = LLMCallSpanData.from_standard_logging_payload( diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 96c420299e7..62aae7aed34 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -1,6 +1,8 @@ import asyncio +import base64 import contextlib import contextvars +import io import json import logging import os @@ -9,7 +11,8 @@ import threading from collections.abc import Callable, Iterator from concurrent.futures import Future, ThreadPoolExecutor from datetime import datetime, timedelta, timezone -from typing import Final +from pathlib import PurePath +from typing import Final, cast from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -19,9 +22,6 @@ from jsonschema import validate import litellm from litellm._internal_context import is_internal_call -from litellm.caching.caching import Cache -from litellm.caching.caching_handler import _PENDING_CACHE_WRITES -from litellm.constants import DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT from litellm._logging import ( CorrelationContextFilter, JsonFormatter, @@ -29,14 +29,18 @@ from litellm._logging import ( trace_id_var, verbose_logger, ) +from litellm.caching.caching import Cache +from litellm.caching.caching_handler import _PENDING_CACHE_WRITES +from litellm.constants import DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.get_litellm_params import get_litellm_params from litellm.litellm_core_utils.thread_pool_executor import executor as logging_executor from litellm.llms.base_llm.base_model_iterator import MockResponseIterator from litellm.proxy.utils import is_valid_api_key -from litellm.types.router import CredentialLiteLLMParams, GenericLiteLLMParams from litellm.types.integrations.custom_logger import HEADROOM_CONVERTED_STREAM_KEY +from litellm.types.router import CredentialLiteLLMParams, GenericLiteLLMParams from litellm.types.utils import ( + ADDRESSED_RESPONSE_ID_FIELD, CallTypes, Choices, Delta, @@ -46,7 +50,6 @@ from litellm.types.utils import ( PromptTokensDetailsWrapper, StreamingChoices, Usage, - ADDRESSED_RESPONSE_ID_FIELD, all_litellm_params, bedrock_batch_litellm_params, ) @@ -1638,7 +1641,6 @@ class TestProxyFunctionCalling: # For now, we expect False (current behavior), but document the limitation assert proxy_result is False, f"Current limitation: {proxy_model_with_hints} returns False without inference" - def test_litellm_utils_supports_function_calling_import(self): """Test that supports_function_calling can be imported from litellm.utils.""" try: @@ -1658,7 +1660,6 @@ class TestProxyFunctionCalling: except Exception as e: pytest.fail(f"Failed to access litellm.supports_function_calling: {e}") - def test_edge_cases_and_malformed_proxy_models(self): """Test edge cases and malformed proxy model names.""" test_cases = [ @@ -5004,7 +5005,9 @@ def _budget_reservation(callback_bound: bool = False) -> dict: _BUDGET_RESERVATION_CALL_KWARGS: Final = {"model": "gpt-4o", "messages": [{"role": "user", "content": "hi"}]} -_BUDGET_RESERVATION_REFUSAL: Final = litellm.AuthenticationError(message="bad key", llm_provider="openai", model="gpt-4o") +_BUDGET_RESERVATION_REFUSAL: Final = litellm.AuthenticationError( + message="bad key", llm_provider="openai", model="gpt-4o" +) @pytest.mark.asyncio @@ -6004,3 +6007,111 @@ def test_calculate_max_parallel_requests_precedence( ) == expected ) + + +class _NamedStream(io.BytesIO): + def __init__(self, name: str | int) -> None: + super().__init__(b"%PDF-1.4 secret document body") + self.name = name + + +def _logged_request_messages(original_function: str, *args: object, **kwargs: object) -> object: + logging_obj, _ = litellm.utils.function_setup( + original_function, + litellm.utils.Rules(), + datetime.now(), + *args, + litellm_call_id="request-text-call", + **kwargs, + ) + return logging_obj.messages + + +@pytest.mark.parametrize( + ("original_function", "args", "kwargs", "expected"), + [ + ("search", (), {"query": "Eiffel Tower"}, "Eiffel Tower"), + ("asearch", ("Eiffel Tower",), {}, "Eiffel Tower"), + ("asearch", (), {"query": ["Eiffel Tower", "Louvre"]}, "Eiffel Tower\nLouvre"), + ("asearch", (), {"query": ["Eiffel Tower", 7, None]}, "Eiffel Tower"), + ("image_edit", (), {"prompt": "make it blue", "image": b"png"}, "make it blue"), + ("aimage_edit", (b"png", "make it blue"), {}, "make it blue"), + ( + "aocr", + (), + {"document": {"type": "document_url", "document_url": "https://x.test/a.pdf"}}, + "https://x.test/a.pdf", + ), + ( + "ocr", + ("mistral-ocr-latest", {"type": "image_url", "image_url": "https://x.test/a.png"}), + {}, + "https://x.test/a.png", + ), + ( + "aocr", + (), + {"document": {"type": "document_url", "document_url": "data:application/pdf;base64,JVBERi0xLjQ="}}, + "data:application/pdf;base64 (12 chars)", + ), + ( + "aocr", + (), + {"document": {"type": "image_url", "image_url": "https://x.test/a,b.png"}}, + "https://x.test/a,b.png", + ), + ( + "aocr", + (), + {"document": {"type": "file", "file": PurePath("/tmp/hello.pdf"), "mime_type": "application/pdf"}}, + "file (application/pdf) hello.pdf", + ), + ("aocr", (), {"document": {"type": "document_url", "document_url": ""}}, ""), + ("aocr", (), {"document": {"type": "file", "file": b"%PDF"}}, "file 4 bytes"), + ("aocr", (), {"document": {"type": "file", "file": io.BytesIO(b"%PDF")}}, "file"), + ("aocr", (), {"document": {"type": "file", "file": _NamedStream("/tmp/scan.pdf")}}, "file scan.pdf"), + ( + "aocr", + (), + {"document": {"type": "file", "file": _NamedStream(3), "mime_type": "application/pdf"}}, + "file (application/pdf)", + ), + ("aocr", (), {"document": "not-a-document"}, "default-message-value"), + ], +) +def test_function_setup_logs_the_search_query_edit_prompt_and_ocr_document_summary_as_the_request( + original_function: str, args: tuple[object, ...], kwargs: dict[str, object], expected: str +) -> None: + assert _logged_request_messages(original_function, *args, **kwargs) == [{"role": "user", "content": expected}] + + +def test_search_with_a_mixed_type_query_list_still_reaches_its_own_validation_error() -> None: + mixed_query: Final = cast(list[str], ["Eiffel Tower", 7]) # cast-ok: the invalid list is the point of the test + + with pytest.raises(litellm.APIConnectionError, match="All items in query list must be strings"): + litellm.search(query=mixed_query, search_provider="duckduckgo") + + +def test_function_setup_never_logs_the_ocr_file_bytes() -> None: + content: Final = b"%PDF-1.4 secret document body" + logged: Final = _logged_request_messages("aocr", document={"type": "file", "file": content}) + + assert logged == [{"role": "user", "content": "file 29 bytes"}] + + +def test_function_setup_leaves_the_ocr_file_stream_unread_and_never_logs_its_bytes() -> None: + stream: Final = _NamedStream("/tmp/scan.pdf") + logged: Final = _logged_request_messages("aocr", document={"type": "file", "file": stream}) + + assert logged == [{"role": "user", "content": "file scan.pdf"}] + assert stream.tell() == 0 + + +def test_function_setup_never_logs_the_ocr_data_uri_payload() -> None: + payload: Final = base64.b64encode(b"%PDF-1.4 secret document body").decode() + logged: Final = _logged_request_messages( + "aocr", document={"type": "document_url", "document_url": f"data:application/pdf;base64,{payload}"} + ) + + assert logged == [{"role": "user", "content": f"data:application/pdf;base64 ({len(payload)} chars)"}] + assert payload not in str(logged)