diff --git a/.agents/skills/sdk-integrations/SKILL.md b/.agents/skills/sdk-integrations/SKILL.md index 89c16c2e..793750ce 100644 --- a/.agents/skills/sdk-integrations/SKILL.md +++ b/.agents/skills/sdk-integrations/SKILL.md @@ -165,6 +165,7 @@ The integration that directly owns the model/provider API response owns token ac - Do not over-serialize. Braintrust serializes at send/log time. Integration tracing only needs readable Python dicts/lists/scalars and materialized attachments. - Keep wrapper bodies thin: prepare traced input, open span, call provider, normalize result, log. - Do not wrap `span.log(...)` / `span.set_attributes(...)` in broad try/except. Braintrust span methods are boundary-safe. +- **Rerank exception to "no silent caps":** rerank result lists can be huge and are noisy in the span UI; cap the `output` list at a fixed max (e.g. 100 entries) via a named constant. The cap is compact-span hygiene, not coverage bounding — the model already ranked everything. ### Span origin diff --git a/py/src/braintrust/integrations/auto_test_scripts/test_auto_litellm.py b/py/src/braintrust/integrations/auto_test_scripts/test_auto_litellm.py index a5f24abc..d5cdd667 100644 --- a/py/src/braintrust/integrations/auto_test_scripts/test_auto_litellm.py +++ b/py/src/braintrust/integrations/auto_test_scripts/test_auto_litellm.py @@ -32,6 +32,6 @@ spans = memory_logger.pop() assert len(spans) == 1, f"Expected 1 span, got {len(spans)}" span = spans[0] - assert span["metadata"]["provider"] == "litellm" + assert span["metadata"]["provider"] == "openai" print("SUCCESS") diff --git a/py/src/braintrust/integrations/auto_test_scripts/test_patch_litellm_aresponses.py b/py/src/braintrust/integrations/auto_test_scripts/test_patch_litellm_aresponses.py index 80da470c..496f25ec 100644 --- a/py/src/braintrust/integrations/auto_test_scripts/test_patch_litellm_aresponses.py +++ b/py/src/braintrust/integrations/auto_test_scripts/test_patch_litellm_aresponses.py @@ -30,7 +30,7 @@ async def main(): for key, value in span["metrics"].items(): assert isinstance(value, (int, float)) and not isinstance(value, bool) assert span["metadata"]["model"] == "gpt-4o-mini" - assert span["metadata"]["provider"] == "litellm" + assert span["metadata"]["provider"] == "openai" assert "What's 12 + 12?" in str(span["input"]) diff --git a/py/src/braintrust/integrations/auto_test_scripts/test_patch_litellm_responses.py b/py/src/braintrust/integrations/auto_test_scripts/test_patch_litellm_responses.py index 5f53aab8..37ef2ce0 100644 --- a/py/src/braintrust/integrations/auto_test_scripts/test_patch_litellm_responses.py +++ b/py/src/braintrust/integrations/auto_test_scripts/test_patch_litellm_responses.py @@ -26,7 +26,7 @@ for key, value in span["metrics"].items(): assert isinstance(value, (int, float)) and not isinstance(value, bool) assert span["metadata"]["model"] == "gpt-4o-mini" - assert span["metadata"]["provider"] == "litellm" + assert span["metadata"]["provider"] == "openai" assert "What's 12 + 12?" in str(span["input"]) print("SUCCESS") diff --git a/py/src/braintrust/integrations/litellm/__init__.py b/py/src/braintrust/integrations/litellm/__init__.py index 65c9b4f3..2706c6c9 100644 --- a/py/src/braintrust/integrations/litellm/__init__.py +++ b/py/src/braintrust/integrations/litellm/__init__.py @@ -1,74 +1,18 @@ """Braintrust LiteLLM integration.""" -import importlib - from .integration import LiteLLMIntegration -from .patchers import LiteLLMAcompletionPatcher, LiteLLMCompletionPatcher, wrap_litellm +from .patchers import wrap_litellm def patch_litellm() -> bool: - """Patch LiteLLM to add Braintrust tracing. - - This wraps litellm.completion, litellm.acompletion, litellm.responses, - litellm.aresponses, litellm.image_generation, litellm.aimage_generation, - litellm.embedding, litellm.aembedding, litellm.moderation, - litellm.speech, litellm.aspeech, litellm.transcription, - litellm.atranscription, litellm.rerank, and litellm.arerank to automatically - create Braintrust spans with detailed token metrics, timing, and costs. - - Returns: - True if LiteLLM was patched (or already patched), False if LiteLLM is not installed. + """Patch LiteLLM top-level callables to emit Braintrust spans. - Example: - ```python - import braintrust - braintrust.integrations.litellm.patch_litellm() - - import litellm - from braintrust import init_logger - - logger = init_logger(project="my-project") - response = litellm.completion( - model="gpt-4o-mini", - messages=[{"role": "user", "content": "Hello"}] - ) - ``` + Returns ``True`` if LiteLLM was patched (or already patched), ``False`` + if LiteLLM is not installed. """ return LiteLLMIntegration.setup() -def _is_litellm_patched() -> bool: - """Return ``True`` when the LiteLLM completion entry points are patched. - - Internal helper used by other Braintrust integrations (e.g. CrewAI, - which delegates its LLM calls to LiteLLM) to decide whether to emit - token metrics on their own LLM spans. When LiteLLM is already - patched, the Braintrust ``Completion`` span it produces is the leaf - span and must own token accounting; non-leaf wrappers should emit - timing-only metrics to avoid double-counting during trace-tree - rollup. - - Not part of the public API — cross-integration callers import this - directly by its underscore-prefixed name. - - Returns ``False`` if ``litellm`` is not importable or neither of the - ``completion`` / ``acompletion`` entry points has been patched. - """ - try: - litellm = importlib.import_module("litellm") - except ImportError: - return False - - completion = getattr(litellm, "completion", None) - acompletion = getattr(litellm, "acompletion", None) - return bool( - LiteLLMCompletionPatcher.has_patch_marker(completion) - or LiteLLMCompletionPatcher.has_patch_marker(litellm) - or LiteLLMAcompletionPatcher.has_patch_marker(acompletion) - or LiteLLMAcompletionPatcher.has_patch_marker(litellm) - ) - - __all__ = [ "LiteLLMIntegration", "patch_litellm", diff --git a/py/src/braintrust/integrations/litellm/patchers.py b/py/src/braintrust/integrations/litellm/patchers.py index 75ec6791..2c1d487e 100644 --- a/py/src/braintrust/integrations/litellm/patchers.py +++ b/py/src/braintrust/integrations/litellm/patchers.py @@ -171,37 +171,9 @@ class LiteLLMArerankPatcher(FunctionWrapperPatcher): def wrap_litellm(litellm: Any) -> Any: - """Wrap a LiteLLM module to add Braintrust tracing. + """Wrap a LiteLLM module-like object in-place with Braintrust tracing. - Unlike :func:`patch_litellm`, which patches the globally-imported ``litellm`` - module, this function instruments a specific module object (or any object - that exposes the same top-level callables such as ``completion``, - ``acompletion``, ``responses``, ``aresponses``, ``image_generation``, - ``text_completion``, ``atext_completion``, ``aimage_generation``, - ``embedding``, ``aembedding``, ``moderation``, ``amoderation``, ``speech``, - ``aspeech``, ``transcription``, ``atranscription``, ``rerank``, and - ``arerank``). Each patcher is applied idempotently — calling - ``wrap_litellm`` twice on the same object is safe. - - Args: - litellm: The ``litellm`` module or a module-like object that exposes - the standard LiteLLM top-level functions. - - Returns: - The same *litellm* object, with tracing wrappers applied in-place. - - Example:: - - import litellm - from braintrust.integrations.litellm import wrap_litellm - - wrap_litellm(litellm) - - # All subsequent calls are automatically traced. - response = litellm.completion( - model="gpt-4o-mini", - messages=[{"role": "user", "content": "Hello"}], - ) + Idempotent; safe to call twice on the same object. """ for patcher in _ALL_LITELLM_PATCHERS: patcher.wrap_target(litellm) diff --git a/py/src/braintrust/integrations/litellm/test_litellm.py b/py/src/braintrust/integrations/litellm/test_litellm.py index 33e97ca5..634347c3 100644 --- a/py/src/braintrust/integrations/litellm/test_litellm.py +++ b/py/src/braintrust/integrations/litellm/test_litellm.py @@ -8,7 +8,6 @@ from braintrust.integrations.litellm import patch_litellm from braintrust.integrations.test_utils import ( assert_metrics_are_valid, - run_in_subprocess, verify_autoinstrument_script, ) from braintrust.test_helpers import assert_dict_matches, init_test_logger @@ -72,7 +71,7 @@ def test_litellm_completion_metrics(memory_logger) -> None: metrics = span["metrics"] assert_metrics_are_valid(metrics, start, end) assert span["metadata"]["model"] == TEST_MODEL - assert span["metadata"]["provider"] == "litellm" + assert span["metadata"]["provider"] == "openai" assert TEST_PROMPT in str(span["input"]) @@ -96,7 +95,7 @@ async def test_litellm_acompletion_metrics(memory_logger): metrics = span["metrics"] assert_metrics_are_valid(metrics, start, end) assert span["metadata"]["model"] == TEST_MODEL - assert span["metadata"]["provider"] == "litellm" + assert span["metadata"]["provider"] == "openai" assert TEST_PROMPT in str(span["input"]) @@ -119,7 +118,7 @@ def test_litellm_text_completion_metrics(memory_logger) -> None: metrics = span["metrics"] assert_metrics_are_valid(metrics, start, end) assert span["metadata"]["model"] == TEST_TEXT_MODEL - assert span["metadata"]["provider"] == "litellm" + assert span["metadata"]["provider"] == "openai" assert TEST_PROMPT in str(span["input"]) assert "text" in span["output"][0] @@ -144,7 +143,7 @@ async def test_litellm_atext_completion_metrics(memory_logger): metrics = span["metrics"] assert_metrics_are_valid(metrics, start, end) assert span["metadata"]["model"] == TEST_TEXT_MODEL - assert span["metadata"]["provider"] == "litellm" + assert span["metadata"]["provider"] == "openai" assert TEST_PROMPT in str(span["input"]) assert "text" in span["output"][0] @@ -186,7 +185,7 @@ def test_litellm_completion_streaming_sync(memory_logger): metrics = span["metrics"] assert_metrics_are_valid(metrics, start, end) assert span["metadata"]["model"] == TEST_MODEL - assert span["metadata"]["provider"] == "litellm" + assert span["metadata"]["provider"] == "openai" assert TEST_PROMPT in str(span["input"]) assert "24" in str(span["output"]) or "twenty-four" in str(span["output"]).lower() @@ -220,7 +219,7 @@ async def test_litellm_acompletion_streaming_async(memory_logger): metrics = span["metrics"] assert_metrics_are_valid(metrics, start, end) assert span["metadata"]["model"] == TEST_MODEL - assert span["metadata"]["provider"] == "litellm" + assert span["metadata"]["provider"] == "openai" assert TEST_PROMPT in str(span["input"]) @@ -250,7 +249,7 @@ def test_litellm_responses_metrics(memory_logger): metrics = span["metrics"] assert_metrics_are_valid(metrics, start, end) assert span["metadata"]["model"] == TEST_MODEL - assert span["metadata"]["provider"] == "litellm" + assert span["metadata"]["provider"] == "openai" assert TEST_PROMPT in str(span["input"]) @@ -281,7 +280,7 @@ async def test_litellm_aresponses_metrics(memory_logger): metrics = span["metrics"] assert_metrics_are_valid(metrics, start, end) assert span["metadata"]["model"] == TEST_MODEL - assert span["metadata"]["provider"] == "litellm" + assert span["metadata"]["provider"] == "openai" assert TEST_PROMPT in str(span["input"]) @@ -303,7 +302,7 @@ def test_litellm_embeddings(memory_logger): span = spans[0] assert span assert span["metadata"]["model"] == "text-embedding-ada-002" - assert span["metadata"]["provider"] == "litellm" + assert span["metadata"]["provider"] == "openai" assert "This is a test" in str(span["input"]) @@ -323,7 +322,7 @@ async def test_litellm_aembedding(memory_logger): span = spans[0] assert span assert span["metadata"]["model"] == "text-embedding-ada-002" - assert span["metadata"]["provider"] == "litellm" + assert span["metadata"]["provider"] == "openai" assert "This is a test" in str(span["input"]) @@ -344,7 +343,7 @@ def test_litellm_moderation(memory_logger): span = spans[0] assert span assert span["metadata"]["model"] == "omni-moderation-latest" - assert span["metadata"]["provider"] == "litellm" + assert span["metadata"]["provider"] == "openai" assert "This is a test message" in str(span["input"]) @@ -362,7 +361,7 @@ async def test_litellm_amoderation(memory_logger): assert len(spans) == 1 span = spans[0] assert span["metadata"]["model"] == "omni-moderation-latest" - assert span["metadata"]["provider"] == "litellm" + assert span["metadata"]["provider"] == "openai" assert "This is a test message" in str(span["input"]) @@ -386,7 +385,7 @@ def test_litellm_image_generation(memory_logger): assert len(spans) == 1 span = spans[0] assert span["metadata"]["model"] == "gpt-image-1-mini" - assert span["metadata"]["provider"] == "litellm" + assert span["metadata"]["provider"] == "openai" assert span["input"] == prompt assert span["output"]["images_count"] == 1 assert span["metrics"]["duration"] >= 0 @@ -413,7 +412,7 @@ async def test_litellm_aimage_generation(memory_logger): assert len(spans) == 1 span = spans[0] assert span["metadata"]["model"] == "gpt-image-1-mini" - assert span["metadata"]["provider"] == "litellm" + assert span["metadata"]["provider"] == "openai" assert span["input"] == prompt assert span["output"]["images_count"] == 1 assert span["metrics"]["duration"] >= 0 @@ -457,7 +456,7 @@ def test_litellm_transcription(memory_logger): assert len(spans) == 1 span = spans[0] assert span["metadata"]["model"] == "whisper-1" - assert span["metadata"]["provider"] == "litellm" + assert span["metadata"]["provider"] == "openai" assert isinstance(span["input"]["file"], Attachment) assert span["input"]["file"].reference["filename"] == "test_audio.wav" assert span["input"]["file"].reference["content_type"] in ("audio/x-wav", "audio/wav") # OS-dependent @@ -479,7 +478,7 @@ async def test_litellm_atranscription(memory_logger): assert len(spans) == 1 span = spans[0] assert span["metadata"]["model"] == "whisper-1" - assert span["metadata"]["provider"] == "litellm" + assert span["metadata"]["provider"] == "openai" assert isinstance(span["input"]["file"], Attachment) assert span["input"]["file"].reference["filename"] == "test_audio.wav" assert span["input"]["file"].reference["content_type"] in ("audio/x-wav", "audio/wav") # OS-dependent @@ -506,7 +505,7 @@ def test_litellm_speech(memory_logger): assert span["metadata"]["model"] == "tts-1" assert span["metadata"]["voice"] == "alloy" assert span["metadata"]["response_format"] == "mp3" - assert span["metadata"]["provider"] == "litellm" + assert span["metadata"]["provider"] == "openai" assert span["input"] == "Hello, this is a test." _assert_speech_output_attachment(span) @@ -532,7 +531,7 @@ async def test_litellm_aspeech(memory_logger): assert span["metadata"]["model"] == "tts-1" assert span["metadata"]["voice"] == "alloy" assert span["metadata"]["response_format"] == "mp3" - assert span["metadata"]["provider"] == "litellm" + assert span["metadata"]["provider"] == "openai" assert span["input"] == "Hello, this is a test." _assert_speech_output_attachment(span) @@ -636,7 +635,7 @@ async def test_litellm_async_parallel_requests(memory_logger): # Verify each span has proper data for i, span in enumerate(spans): assert span["metadata"]["model"] == TEST_MODEL - assert span["metadata"]["provider"] == "litellm" + assert span["metadata"]["provider"] == "openai" assert prompts[i] in str(span["input"]) assert_metrics_are_valid(span["metrics"]) @@ -687,7 +686,7 @@ def test_litellm_tool_calls(memory_logger): "span_attributes": {"type": "llm", "name": "Completion"}, "metadata": { "model": TEST_MODEL, - "provider": "litellm", + "provider": "openai", "tools": lambda tools_list: ( len(tools_list) == 1 and any(tool.get("function", {}).get("name") == "get_weather" for tool in tools_list) @@ -790,45 +789,6 @@ def test_patch_litellm_responses(): verify_autoinstrument_script("test_patch_litellm_responses.py") -def test_is_litellm_patched_internal_helper(): - """Internal ``_is_litellm_patched()`` flips True after setup and honors manual wrap. - - Run in a subprocess so that marker attributes set by other tests in this - module do not leak into the assertion. - """ - result = run_in_subprocess( - """ - from braintrust.integrations.litellm import _is_litellm_patched, patch_litellm, wrap_litellm - assert _is_litellm_patched() is False, 'expected unpatched at startup' - assert patch_litellm() is True - assert _is_litellm_patched() is True, 'expected True after patch_litellm()' - # idempotent wrap on the already-patched module stays True. - import litellm - wrap_litellm(litellm) - assert _is_litellm_patched() is True - print("SUCCESS") - """ - ) - assert result.returncode == 0, result.stderr - assert "SUCCESS" in result.stdout - - -def test_is_litellm_patched_returns_false_without_litellm_installed(): - """``_is_litellm_patched()`` must not raise when ``litellm`` is absent.""" - result = run_in_subprocess( - """ - import sys - # Simulate a missing litellm install. - sys.modules['litellm'] = None - from braintrust.integrations.litellm import _is_litellm_patched - assert _is_litellm_patched() is False - print("SUCCESS") - """ - ) - assert result.returncode == 0, result.stderr - assert "SUCCESS" in result.stdout - - def test_patch_litellm_aresponses(): """Test that patch_litellm() patches aresponses (subprocess to avoid global state pollution).""" verify_autoinstrument_script("test_patch_litellm_aresponses.py") @@ -901,10 +861,10 @@ def test_litellm_openrouter_no_booleans_in_metrics(memory_logger): spans = memory_logger.pop() assert len(spans) == 1 metrics = spans[0]["metrics"] - metadata_keys = set(spans[0]["metadata"]) - assert "api_key" not in metadata_keys + metadata = spans[0]["metadata"] + assert "api_key" not in metadata + assert metadata["provider"] == "openrouter" - # No boolean values should be in metrics for key, value in metrics.items(): assert not isinstance(value, bool) assert "is_byok" not in metrics @@ -934,7 +894,7 @@ def test_litellm_rerank(memory_logger): span = spans[0] assert span["span_attributes"]["name"] == "Rerank" assert span["span_attributes"]["type"] == "llm" - assert span["metadata"]["provider"] == "litellm" + assert span["metadata"]["provider"] == "cohere" assert span["metadata"]["model"] == RERANK_MODEL assert span["metadata"]["top_n"] == 2 assert span["metadata"]["document_count"] == 3 @@ -979,7 +939,7 @@ async def test_litellm_arerank(memory_logger): assert len(spans) == 1 span = spans[0] assert span["span_attributes"]["name"] == "Rerank" - assert span["metadata"]["provider"] == "litellm" + assert span["metadata"]["provider"] == "cohere" assert span["metadata"]["model"] == RERANK_MODEL assert span["metadata"]["top_n"] == 2 assert span["metadata"]["document_count"] == 3 diff --git a/py/src/braintrust/integrations/litellm/tracing.py b/py/src/braintrust/integrations/litellm/tracing.py index 750ad2d3..7b934d91 100644 --- a/py/src/braintrust/integrations/litellm/tracing.py +++ b/py/src/braintrust/integrations/litellm/tracing.py @@ -1,7 +1,9 @@ """LiteLLM tracing helpers — spans, metadata extraction, stream handling.""" +import logging import time from collections.abc import AsyncGenerator, Generator +from functools import lru_cache from types import TracebackType from typing import Any @@ -16,10 +18,14 @@ ) from braintrust.logger import Span from braintrust.logger import start_span as _bt_start_span +from braintrust.span_types import SpanTypeAttribute +from braintrust.util import clean_nones, is_numeric, merge_dicts _INSTRUMENTATION = "litellm-auto" +_log = logging.getLogger(__name__) + def start_span(*args, **kwargs): internal = dict(kwargs.get("internal") or {}) @@ -28,10 +34,6 @@ def start_span(*args, **kwargs): return _bt_start_span(*args, **kwargs) -from braintrust.span_types import SpanTypeAttribute -from braintrust.util import clean_nones, is_numeric, merge_dicts - - # LiteLLM's representation to Braintrust's representation TOKEN_NAME_MAP: dict[str, str] = { # chat API @@ -50,6 +52,42 @@ def start_span(*args, **kwargs): } +_RERANK_MAX_RESULTS = 100 + +# litellm reports the legacy /v1/completions endpoint under its own routing +# key, but pricing/billing is against the underlying provider — normalize so +# metadata.provider stays consistent across chat and text-completion calls. +_PROVIDER_ALIASES: dict[str, str] = { + "text-completion-openai": "openai", + "azure_text": "azure", + "text-completion-codestral": "codestral", + "text-completion-inception": "inception", +} + + +def _get_field(obj: Any, key: str, default: Any = None) -> Any: + """Read *key* off *obj* whether it's a dict or an SDK/pydantic object.""" + if obj is None: + return default + if isinstance(obj, dict): + return obj.get(key, default) + return getattr(obj, key, default) + + +@lru_cache(maxsize=256) +def _resolve_provider(model: str) -> str: + # Lazy import: tracing.py loads eagerly via auto.py even when litellm isn't installed. + try: + import litellm + + _, provider, _, _ = litellm.get_llm_provider(model) + if isinstance(provider, str) and provider: + return _PROVIDER_ALIASES.get(provider, provider) + except Exception: + _log.debug("litellm.get_llm_provider failed for model=%r", model, exc_info=True) + return "litellm" + + # --------------------------------------------------------------------------- # Async response wrapper (preserves async context manager / iterator behavior) # --------------------------------------------------------------------------- @@ -198,10 +236,9 @@ def _completion_wrapper_impl(wrapped, args, kwargs, *, input_key: str): should_end = False return _handle_completion_streaming(completion_response, span, start, is_async=False) - log_response = _try_to_dict(completion_response) - metrics = _parse_metrics_from_usage(log_response.get("usage", {})) + metrics = _parse_metrics_from_usage(_get_field(completion_response, "usage")) metrics["time_to_first_token"] = time.time() - start - span.log(metrics=metrics, output=log_response["choices"]) + span.log(metrics=metrics, output=_get_field(completion_response, "choices")) return completion_response finally: if should_end: @@ -225,10 +262,9 @@ async def _acompletion_wrapper_impl(wrapped, args, kwargs, *, input_key: str): should_end = False return _handle_completion_streaming(completion_response, span, start, is_async=True) - log_response = _try_to_dict(completion_response) - metrics = _parse_metrics_from_usage(log_response.get("usage", {})) + metrics = _parse_metrics_from_usage(_get_field(completion_response, "usage")) metrics["time_to_first_token"] = time.time() - start - span.log(metrics=metrics, output=log_response["choices"]) + span.log(metrics=metrics, output=_get_field(completion_response, "choices")) return completion_response finally: if should_end: @@ -236,27 +272,22 @@ async def _acompletion_wrapper_impl(wrapped, args, kwargs, *, input_key: str): def _completion_wrapper(wrapped, instance, args, kwargs): - """wrapt wrapper for litellm.completion.""" return _completion_wrapper_impl(wrapped, args, kwargs, input_key="messages") def _text_completion_wrapper(wrapped, instance, args, kwargs): - """wrapt wrapper for litellm.text_completion.""" return _completion_wrapper_impl(wrapped, args, kwargs, input_key="prompt") async def _acompletion_wrapper_async(wrapped, instance, args, kwargs): - """wrapt wrapper for litellm.acompletion.""" return await _acompletion_wrapper_impl(wrapped, args, kwargs, input_key="messages") async def _atext_completion_wrapper_async(wrapped, instance, args, kwargs): - """wrapt wrapper for litellm.atext_completion.""" return await _acompletion_wrapper_impl(wrapped, args, kwargs, input_key="prompt") def _responses_wrapper(wrapped, instance, args, kwargs): - """wrapt wrapper for litellm.responses.""" updated_span_payload = _update_span_payload_from_params(kwargs, input_key="input") is_streaming = kwargs.get("stream", False) @@ -272,19 +303,17 @@ def _responses_wrapper(wrapped, instance, args, kwargs): if is_streaming: should_end = False return _handle_responses_streaming(response, span, start, is_async=False) - else: - log_response = _try_to_dict(response) - metrics = _parse_metrics_from_usage(log_response.get("usage", {})) - metrics["time_to_first_token"] = time.time() - start - span.log(metrics=metrics, output=log_response["output"]) - return response + + metrics = _parse_metrics_from_usage(_get_field(response, "usage")) + metrics["time_to_first_token"] = time.time() - start + span.log(metrics=metrics, output=_get_field(response, "output")) + return response finally: if should_end: span.end() async def _aresponses_wrapper_async(wrapped, instance, args, kwargs): - """wrapt wrapper for litellm.aresponses.""" updated_span_payload = _update_span_payload_from_params(kwargs, input_key="input") is_streaming = kwargs.get("stream", False) @@ -300,12 +329,11 @@ async def _aresponses_wrapper_async(wrapped, instance, args, kwargs): if is_streaming: should_end = False return _handle_responses_streaming(response, span, start, is_async=True) - else: - log_response = _try_to_dict(response) - metrics = _parse_metrics_from_usage(log_response.get("usage", {})) - metrics["time_to_first_token"] = time.time() - start - span.log(metrics=metrics, output=log_response["output"]) - return response + + metrics = _parse_metrics_from_usage(_get_field(response, "usage")) + metrics["time_to_first_token"] = time.time() - start + span.log(metrics=metrics, output=_get_field(response, "output")) + return response finally: if should_end: span.end() @@ -329,25 +357,18 @@ def _image_attachment_from_base64( return resolved_attachment, len(resolved_attachment.attachment.data) -def _extract_image_generation_output(response: dict[str, Any]) -> dict[str, Any]: - images = [] - output_format = response.get("output_format") +def _extract_image_generation_output(response: Any) -> Any: + images: list[dict[str, Any]] = [] + output_format = _get_field(response, "output_format") - for index, image in enumerate(response.get("data") or []): - image_dict = _try_to_dict(image) - if not isinstance(image_dict, dict): - continue + for index, image in enumerate(_get_field(response, "data") or []): + image_entry = clean_nones({"revised_prompt": _get_field(image, "revised_prompt")}) - image_entry = clean_nones( - { - "revised_prompt": image_dict.get("revised_prompt"), - } - ) + url = _get_field(image, "url") + if isinstance(url, str): + image_entry["image_url"] = {"url": url} - if isinstance(image_dict.get("url"), str): - image_entry["image_url"] = {"url": image_dict["url"]} - - b64_json = image_dict.get("b64_json") + b64_json = _get_field(image, "b64_json") resolved_attachment, image_size_bytes = _image_attachment_from_base64( b64_json, output_format=output_format, @@ -364,21 +385,25 @@ def _extract_image_generation_output(response: dict[str, Any]) -> dict[str, Any] return clean_nones( { - "created": response.get("created"), - "background": response.get("background"), + "created": _get_field(response, "created"), + "background": _get_field(response, "background"), "output_format": output_format, - "quality": response.get("quality"), - "size": response.get("size"), + "quality": _get_field(response, "quality"), + "size": _get_field(response, "size"), "images_count": len(images), "images": images, } ) +def _log_image_generation(span: Span, image_response: Any, start: float) -> None: + metrics = _timing_metrics(start, time.time()) + metrics.update(_parse_metrics_from_usage(_get_field(image_response, "usage"))) + span.log(metrics=metrics, output=_extract_image_generation_output(image_response)) + + def _image_generation_wrapper(wrapped, instance, args, kwargs): - """wrapt wrapper for litellm.image_generation.""" updated_span_payload = _update_span_payload_from_params(kwargs, input_key="prompt") - with start_span( **merge_dicts( dict(name="Image Generation", span_attributes={"type": SpanTypeAttribute.LLM}), updated_span_payload @@ -386,21 +411,12 @@ def _image_generation_wrapper(wrapped, instance, args, kwargs): ) as span: start = time.time() image_response = wrapped(*args, **kwargs) - log_response = _try_to_dict(image_response) - metrics = _timing_metrics(start, time.time()) - if isinstance(log_response, dict): - metrics.update(_parse_metrics_from_usage(log_response.get("usage", {}))) - span.log( - metrics=metrics, - output=_extract_image_generation_output(log_response) if isinstance(log_response, dict) else log_response, - ) + _log_image_generation(span, image_response, start) return image_response async def _aimage_generation_wrapper_async(wrapped, instance, args, kwargs): - """wrapt wrapper for litellm.aimage_generation.""" updated_span_payload = _update_span_payload_from_params(kwargs, input_key="prompt") - with start_span( **merge_dicts( dict(name="Image Generation", span_attributes={"type": SpanTypeAttribute.LLM}), updated_span_payload @@ -408,67 +424,49 @@ async def _aimage_generation_wrapper_async(wrapped, instance, args, kwargs): ) as span: start = time.time() image_response = await wrapped(*args, **kwargs) - log_response = _try_to_dict(image_response) - metrics = _timing_metrics(start, time.time()) - if isinstance(log_response, dict): - metrics.update(_parse_metrics_from_usage(log_response.get("usage", {}))) - span.log( - metrics=metrics, - output=_extract_image_generation_output(log_response) if isinstance(log_response, dict) else log_response, - ) + _log_image_generation(span, image_response, start) return image_response +def _log_embedding_response(span: Span, embedding_response: Any) -> None: + data = _get_field(embedding_response, "data") or [] + first = data[0] if data else None + length = len(_get_field(first, "embedding") or []) + span.log( + metrics=_parse_metrics_from_usage(_get_field(embedding_response, "usage")), + output={"embedding_length": length}, + ) + + def _embedding_wrapper(wrapped, instance, args, kwargs): - """wrapt wrapper for litellm.embedding.""" updated_span_payload = _update_span_payload_from_params(kwargs, input_key="input") - with start_span( **merge_dicts(dict(name="Embedding", span_attributes={"type": SpanTypeAttribute.LLM}), updated_span_payload) ) as span: embedding_response = wrapped(*args, **kwargs) - log_response = _try_to_dict(embedding_response) - usage = log_response.get("usage") - metrics = _parse_metrics_from_usage(usage) - span.log( - metrics=metrics, - output={"embedding_length": len(log_response["data"][0]["embedding"])}, - ) + _log_embedding_response(span, embedding_response) return embedding_response async def _aembedding_wrapper_async(wrapped, instance, args, kwargs): - """wrapt wrapper for litellm.aembedding.""" updated_span_payload = _update_span_payload_from_params(kwargs, input_key="input") - with start_span( **merge_dicts(dict(name="Embedding", span_attributes={"type": SpanTypeAttribute.LLM}), updated_span_payload) ) as span: embedding_response = await wrapped(*args, **kwargs) - log_response = _try_to_dict(embedding_response) - usage = log_response.get("usage") - metrics = _parse_metrics_from_usage(usage) - span.log( - metrics=metrics, - output={"embedding_length": len(log_response["data"][0]["embedding"])}, - ) + _log_embedding_response(span, embedding_response) return embedding_response def _log_moderation_response(span: Span, moderation_response: Any) -> None: - log_response = _try_to_dict(moderation_response) - usage = log_response.get("usage") - metrics = _parse_metrics_from_usage(usage) span.log( - metrics=metrics, - output=log_response["results"], + metrics=_parse_metrics_from_usage(_get_field(moderation_response, "usage")), + output=_get_field(moderation_response, "results"), ) def _moderation_wrapper(wrapped, instance, args, kwargs): - """wrapt wrapper for litellm.moderation.""" updated_span_payload = _update_span_payload_from_params(kwargs, input_key="input") - with start_span( **merge_dicts(dict(name="Moderation", span_attributes={"type": SpanTypeAttribute.LLM}), updated_span_payload) ) as span: @@ -478,9 +476,7 @@ def _moderation_wrapper(wrapped, instance, args, kwargs): async def _amoderation_wrapper_async(wrapped, instance, args, kwargs): - """wrapt wrapper for litellm.amoderation.""" updated_span_payload = _update_span_payload_from_params(kwargs, input_key="input") - with start_span( **merge_dicts(dict(name="Moderation", span_attributes={"type": SpanTypeAttribute.LLM}), updated_span_payload) ) as span: @@ -490,9 +486,7 @@ async def _amoderation_wrapper_async(wrapped, instance, args, kwargs): def _speech_wrapper(wrapped, instance, args, kwargs): - """wrapt wrapper for litellm.speech.""" updated_span_payload = _update_span_payload_from_params(kwargs, input_key="input") - with start_span( **merge_dicts(dict(name="Speech", span_attributes={"type": SpanTypeAttribute.LLM}), updated_span_payload) ) as span: @@ -510,9 +504,7 @@ def _speech_wrapper(wrapped, instance, args, kwargs): async def _aspeech_wrapper_async(wrapped, instance, args, kwargs): - """wrapt wrapper for litellm.aspeech.""" updated_span_payload = _update_span_payload_from_params(kwargs, input_key="input") - with start_span( **merge_dicts(dict(name="Speech", span_attributes={"type": SpanTypeAttribute.LLM}), updated_span_payload) ) as span: @@ -530,74 +522,58 @@ async def _aspeech_wrapper_async(wrapped, instance, args, kwargs): def _rerank_wrapper(wrapped, instance, args, kwargs): - """wrapt wrapper for litellm.rerank.""" updated_span_payload = _update_rerank_span_payload_from_params(kwargs) - with start_span( **merge_dicts(dict(name="Rerank", span_attributes={"type": SpanTypeAttribute.LLM}), updated_span_payload) ) as span: start = time.time() rerank_response = wrapped(*args, **kwargs) - log_response = _try_to_dict(rerank_response) metrics = _timing_metrics(start, time.time()) - metrics.update(_parse_rerank_metrics(log_response)) - span.log( - metrics=metrics, - output=_extract_rerank_output(log_response), - ) + metrics.update(_parse_rerank_metrics(rerank_response)) + span.log(metrics=metrics, output=_extract_rerank_output(rerank_response)) return rerank_response async def _arerank_wrapper_async(wrapped, instance, args, kwargs): - """wrapt wrapper for litellm.arerank.""" updated_span_payload = _update_rerank_span_payload_from_params(kwargs) - with start_span( **merge_dicts(dict(name="Rerank", span_attributes={"type": SpanTypeAttribute.LLM}), updated_span_payload) ) as span: start = time.time() rerank_response = await wrapped(*args, **kwargs) - log_response = _try_to_dict(rerank_response) metrics = _timing_metrics(start, time.time()) - metrics.update(_parse_rerank_metrics(log_response)) - span.log( - metrics=metrics, - output=_extract_rerank_output(log_response), - ) + metrics.update(_parse_rerank_metrics(rerank_response)) + span.log(metrics=metrics, output=_extract_rerank_output(rerank_response)) return rerank_response def _transcription_wrapper(wrapped, instance, args, kwargs): - """wrapt wrapper for litellm.transcription.""" updated_span_payload = _update_audio_span_payload_from_params(kwargs) - with start_span( **merge_dicts( dict(name="Transcription", span_attributes={"type": SpanTypeAttribute.LLM}), updated_span_payload ) ) as span: transcription_response = wrapped(*args, **kwargs) - log_response = _try_to_dict(transcription_response) - usage = log_response.get("usage") if isinstance(log_response, dict) else None - metrics = _parse_metrics_from_usage(usage) - span.log(metrics=metrics, output=_extract_transcription_text(log_response)) + span.log( + metrics=_parse_metrics_from_usage(_get_field(transcription_response, "usage")), + output=_extract_transcription_text(transcription_response), + ) return transcription_response async def _atranscription_wrapper_async(wrapped, instance, args, kwargs): - """wrapt wrapper for litellm.atranscription.""" updated_span_payload = _update_audio_span_payload_from_params(kwargs) - with start_span( **merge_dicts( dict(name="Transcription", span_attributes={"type": SpanTypeAttribute.LLM}), updated_span_payload ) ) as span: transcription_response = await wrapped(*args, **kwargs) - log_response = _try_to_dict(transcription_response) - usage = log_response.get("usage") if isinstance(log_response, dict) else None - metrics = _parse_metrics_from_usage(usage) - span.log(metrics=metrics, output=_extract_transcription_text(log_response)) + span.log( + metrics=_parse_metrics_from_usage(_get_field(transcription_response, "usage")), + output=_extract_transcription_text(transcription_response), + ) return transcription_response @@ -817,7 +793,8 @@ def _postprocess_responses_streaming_results(all_results: list[Any]) -> dict[str def _build_span_metadata(params: dict[str, Any], model: Any) -> dict[str, Any]: """Build span metadata from known-safe LiteLLM request parameters plus provider/model.""" metadata = {key: value for key, value in params.items() if key in _SAFE_METADATA_KEYS} - metadata.update(provider="litellm", model=model) + provider = _resolve_provider(model) if isinstance(model, str) and model else "litellm" + metadata.update(provider=provider, model=model) return metadata @@ -879,12 +856,9 @@ def _update_audio_span_payload_from_params(params: dict[str, Any]) -> dict[str, def _extract_transcription_text(response: Any) -> str | None: - """Extract text output from a LiteLLM transcription response.""" - if isinstance(response, dict): - return response.get("text") if isinstance(response, str): return response.strip() - return getattr(response, "text", None) + return _get_field(response, "text") def _parse_metrics_from_usage(usage: Any) -> dict[str, Any]: @@ -914,67 +888,49 @@ def _parse_metrics_from_usage(usage: Any) -> dict[str, Any]: def _parse_rerank_metrics(response: Any) -> dict[str, Any]: """Parse token / billed-unit metrics from a LiteLLM rerank response. - LiteLLM follows Cohere's rerank response shape: usage lives under - ``meta.billed_units`` (the authoritative billing counter) and - ``meta.tokens``. ``billed_units`` wins when both are present. + Why: LiteLLM follows Cohere's shape — usage lives under ``meta.billed_units`` + (authoritative billing counter) and ``meta.tokens``. ``billed_units`` wins + when both are present. """ metrics: dict[str, Any] = {} - response_dict = _try_to_dict(response) - if not isinstance(response_dict, dict): + meta = _get_field(response, "meta") + if meta is None: return metrics - meta = _try_to_dict(response_dict.get("meta")) - if not isinstance(meta, dict): - return metrics - - tokens_block = _try_to_dict(meta.get("tokens")) - if isinstance(tokens_block, dict): - for src_key, dst_key in _RERANK_TOKENS_MAP.items(): - value = tokens_block.get(src_key) + def _read(block: Any, mapping: dict[str, str]) -> None: + if block is None: + return + for src_key, dst_key in mapping.items(): + value = _get_field(block, src_key) if is_numeric(value): metrics[dst_key] = value - # ``billed_units`` is Cohere's authoritative billing counter and - # intentionally overrides values from ``tokens`` when both are present. - billed = _try_to_dict(meta.get("billed_units")) - if isinstance(billed, dict): - for src_key, dst_key in _RERANK_BILLED_UNITS_MAP.items(): - value = billed.get(src_key) - if is_numeric(value): - metrics[dst_key] = value + _read(_get_field(meta, "tokens"), _RERANK_TOKENS_MAP) + _read(_get_field(meta, "billed_units"), _RERANK_BILLED_UNITS_MAP) if "tokens" not in metrics and "prompt_tokens" in metrics and "completion_tokens" in metrics: metrics["tokens"] = metrics["prompt_tokens"] + metrics["completion_tokens"] - return metrics def _extract_rerank_output(response: Any) -> list[dict[str, Any]] | None: """Return the ranked ``{index, relevance_score}`` list from a rerank response. - The document payload (returned when ``return_documents=True``) is dropped - on purpose to keep the span compact and avoid logging the raw corpus. - Results are capped at 100 entries. + The document payload (from ``return_documents=True``) is dropped and the + list is capped at ``_RERANK_MAX_RESULTS`` to keep the span compact. """ - response_dict = _try_to_dict(response) - if not isinstance(response_dict, dict): - return None - - results = response_dict.get("results") + results = _get_field(response, "results") if not isinstance(results, list): return None out: list[dict[str, Any]] = [] - for item in results[:100]: + for item in results[:_RERANK_MAX_RESULTS]: if item is None: continue - item_dict = _try_to_dict(item) - if not isinstance(item_dict, dict): - continue out.append( { - "index": item_dict.get("index"), - "relevance_score": item_dict.get("relevance_score"), + "index": _get_field(item, "index"), + "relevance_score": _get_field(item, "relevance_score"), } ) return out