mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-25 01:02:15 +00:00
fix(otel v2): map rerank and search output and the OCR, image edit and search input onto the Langfuse generation (#42444)
* fix(otel v2): map rerank and search output and the OCR, image edit and search input onto the Langfuse generation Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(otel v2): summarize OCR data URIs by media type and size and log an empty document URL as empty Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(otel v2): keep URL-less search results, name OCR file streams and skip non-str query parts when logging Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore(otel v2): drop the unused typing imports and the decorative section divider Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yucheng <yucheng@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
deba473821
commit
88a4cbdd7b
7 changed files with 565 additions and 18 deletions
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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}"
|
||||
|
|
|
|||
|
|
@ -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 ----------
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue