From 400dfd0583be1c91240130f1541d870da56a11e7 Mon Sep 17 00:00:00 2001 From: Shreyas Nagaraj Date: Thu, 10 Sep 2026 13:25:52 +0530 Subject: [PATCH 1/3] Instrument LiteLLM Anthropic messages create/acreate so native pass-through calls emit GenAI spans. acompletion wrapping misses Bedrock/Vertex Claude HTTP pass-through, which never re-enters acompletion. Co-authored-by: Cursor --- .../instrumentation/litellm/__init__.py | 91 ++++++++- .../litellm/litellm_instrumentation_test.py | 172 ++++++++++++++++++ 2 files changed, 262 insertions(+), 1 deletion(-) diff --git a/src/harness_sdk/instrumentation/litellm/__init__.py b/src/harness_sdk/instrumentation/litellm/__init__.py index d0220b8..9c911ca 100644 --- a/src/harness_sdk/instrumentation/litellm/__init__.py +++ b/src/harness_sdk/instrumentation/litellm/__init__.py @@ -5,6 +5,8 @@ Coverage: - litellm.completion / litellm.acompletion - litellm.embedding / litellm.aembedding + - litellm.anthropic.messages.create / acreate + (and litellm.messages.create / acreate when those names alias the same functions) Wraps the public entry points so evaluation runs on an active span before the provider call. The wrapper enriches that span with response metadata before it @@ -45,6 +47,11 @@ _LITELLM_MAIN = "litellm.main" _LITELLM_REQUEST_SPAN_NAME = "litellm_request" +_ANTHROPIC_MESSAGES_MODULES = ( + "litellm.anthropic_interface.messages", + "litellm.anthropic.messages", + "litellm.messages", +) _LITELLM_SPAN_ACTIVE: contextvars.ContextVar[bool] = contextvars.ContextVar( "harness_litellm_span_active", default=False @@ -87,6 +94,11 @@ ("aembedding", True), ) +_ANTHROPIC_MESSAGES_FUNCTIONS = ( + ("create", False), + ("acreate", True), +) + _PROVIDER_NAME_MAP = { "azure": "azure.ai.openai", "azure_ai": "azure.ai.openai", @@ -290,7 +302,13 @@ def finish_reasons(self) -> list[str]: finish_reason = _get_value(choice, "finish_reason") if finish_reason: finish_reasons.append(str(finish_reason)) - return finish_reasons + if finish_reasons: + return finish_reasons + # Anthropic Messages responses use stop_reason instead of choices. + stop_reason = self.value("stop_reason") + if stop_reason: + return [str(stop_reason)] + return [] def header_maps(self) -> list[dict[Any, Any]]: hidden = self.hidden_params() @@ -823,6 +841,74 @@ async def _async_wrapper( return _async_wrapper if is_async else _sync_wrapper +def _rebind_public_function(source_mod: Any, func_name: str) -> None: + """Copy a wrapped function onto LiteLLM's public aliases. + + wrapt patches the implementation module in place. ``from .messages import + acreate`` (and any ``litellm.messages`` alias) still holds the original + function object until rebound, the same way ``litellm.acompletion`` is + rebound after wrapping ``litellm.main.acompletion``. + """ + import litellm # pylint: disable=import-outside-toplevel + + wrapped = getattr(source_mod, func_name) + for holder in ( + getattr(getattr(litellm, "anthropic", None), "messages", None), + getattr(litellm, "anthropic", None), + getattr(litellm, "messages", None), + source_mod, + ): + if holder is not None and hasattr(holder, func_name): + setattr(holder, func_name, wrapped) + + +def _iter_anthropic_messages_modules() -> list[tuple[str, Any]]: + from importlib import import_module # pylint: disable=import-outside-toplevel + + seen: set[int] = set() + modules: list[tuple[str, Any]] = [] + for mod_name in _ANTHROPIC_MESSAGES_MODULES: + try: + mod = import_module(mod_name) + except ImportError: + continue + if id(mod) in seen: + continue + seen.add(id(mod)) + modules.append((mod_name, mod)) + return modules + + +def _wrap_anthropic_messages() -> None: + for mod_name, mod in _iter_anthropic_messages_modules(): + for func_name, is_async in _ANTHROPIC_MESSAGES_FUNCTIONS: + if not hasattr(mod, func_name): + continue + wrapt.wrap_function_wrapper( + mod_name, + func_name, + _make_wrapper(func_name, is_async), + ) + _rebind_public_function(mod, func_name) + + +def _unwrap_anthropic_messages() -> None: + for _mod_name, mod in _iter_anthropic_messages_modules(): + for func_name, _ in _ANTHROPIC_MESSAGES_FUNCTIONS: + if not hasattr(mod, func_name): + continue + try: + unwrap(mod, func_name) + _rebind_public_function(mod, func_name) + except Exception as err: # pylint: disable=broad-except + # Optional interface: skip if this LiteLLM version never wrapped it. + logger.debug( + "LiteLLM anthropic messages %s unwrap skipped: %s", + func_name, + err, + ) + + class LiteLLMInstrumentorWrapper(BaseInstrumentorWrapper): """Instrument LiteLLM with its OpenTelemetry SDK and Traceable policy evaluation.""" @@ -847,6 +933,7 @@ def instrument(self, **_kwargs: Any) -> None: ) if hasattr(litellm, func_name): setattr(litellm, func_name, getattr(main_mod, func_name)) + _wrap_anthropic_messages() LiteLLMInstrumentorWrapper._applied = True logger.debug("Traceable LiteLLM instrumentation applied.") except ImportError as err: @@ -872,6 +959,8 @@ def uninstrument(self, **_kwargs: Any) -> None: logger.error("Failed to uninstrument LiteLLM %s: %s", func_name, err) errors.append(err) + _unwrap_anthropic_messages() + _unregister_otel_callback() global _otel_logger # pylint: disable=global-statement _otel_logger = None diff --git a/test/instrumentation/litellm/litellm_instrumentation_test.py b/test/instrumentation/litellm/litellm_instrumentation_test.py index e921a73..be6898f 100644 --- a/test/instrumentation/litellm/litellm_instrumentation_test.py +++ b/test/instrumentation/litellm/litellm_instrumentation_test.py @@ -13,6 +13,7 @@ from harness_sdk.plugins.control import ControlResult, get_control_registry from harness_sdk.gen_ai.exceptions import ControlEvaluationBlocked from harness_sdk.instrumentation.litellm import LiteLLMInstrumentorWrapper +import harness_sdk.instrumentation.litellm as harness_litellm @pytest.fixture @@ -468,3 +469,174 @@ def test_litellm_mock_response_with_wrapper_enrichment(agent, exporter, litellm_ assert attrs.get("gen_ai.provider.name") == "openai" assert "gen_ai.system" not in attrs assert attrs.get("gen_ai.framework") == "litellm" + + +def _fake_anthropic_message_response(*_args, **_kwargs): + return { + "id": "msg_test", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4", + "content": [{"type": "text", "text": "hello"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 10, "output_tokens": 5}, + } + + +class _FakeAnthropicStream: + def __init__(self, chunks): + self._chunks = list(chunks) + self._index = 0 + + def __aiter__(self): + return self + + async def __anext__(self): + if self._index >= len(self._chunks): + raise StopAsyncIteration + chunk = self._chunks[self._index] + self._index += 1 + return chunk + + +def _anthropic_messages_kwargs(): + return { + "model": "bedrock/anthropic.claude-sonnet-4", + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 32, + "temperature": 0.2, + "custom_llm_provider": "bedrock", + } + + +@pytest.mark.asyncio +async def test_acreate_not_captured_when_only_acompletion_is_wrapped( # pylint: disable=unused-argument + agent, exporter, litellm_instrumentor, monkeypatch +): + # Pre-fix behavior: wrapping acompletion alone misses native Anthropic + # pass-through (acreate does not call acompletion). + monkeypatch.setattr(harness_litellm, "_ANTHROPIC_MESSAGES_FUNCTIONS", ()) + + async def _native_acreate(*_args, **_kwargs): + return _fake_anthropic_message_response() + + with patch("litellm.anthropic_interface.messages.acreate", new=_native_acreate): + litellm_instrumentor.instrument() + await litellm.anthropic.messages.acreate(**_anthropic_messages_kwargs()) + + spans = _litellm_spans(exporter.get_finished_spans()) + exporter.clear() + assert spans == [] + + +@pytest.mark.asyncio +async def test_litellm_anthropic_acreate_span_has_gen_ai_attributes( # pylint: disable=unused-argument + agent, exporter, litellm_instrumentor +): + async def _native_acreate(*_args, **_kwargs): + return _fake_anthropic_message_response() + + with patch("litellm.anthropic_interface.messages.acreate", new=_native_acreate): + litellm_instrumentor.instrument() + await litellm.anthropic.messages.acreate(**_anthropic_messages_kwargs()) + + spans = _litellm_spans(exporter.get_finished_spans()) + exporter.clear() + assert len(spans) == 1 + attrs = spans[0].attributes + assert attrs.get("gen_ai.request.model") == "bedrock/anthropic.claude-sonnet-4" + assert attrs.get("gen_ai.operation.name") == "chat" + assert attrs.get("gen_ai.provider.name") == "aws.bedrock" + assert attrs.get("gen_ai.framework") == "litellm" + assert attrs.get("gen_ai.request.max_tokens") == 32 + assert attrs.get("gen_ai.request.temperature") == 0.2 + assert attrs.get("gen_ai.response.id") == "msg_test" + assert attrs.get("gen_ai.response.model") == "claude-sonnet-4" + assert attrs.get("gen_ai.response.finish_reasons") == "['end_turn']" + assert attrs.get("gen_ai.usage.input_tokens") == 10 + assert attrs.get("gen_ai.usage.output_tokens") == 5 + + +def test_litellm_anthropic_create_span_has_gen_ai_attributes( # pylint: disable=unused-argument + agent, exporter, litellm_instrumentor +): + with patch( + "litellm.anthropic_interface.messages.create", + new=_fake_anthropic_message_response, + ): + litellm_instrumentor.instrument() + litellm.anthropic.messages.create(**_anthropic_messages_kwargs()) + + spans = _litellm_spans(exporter.get_finished_spans()) + exporter.clear() + assert len(spans) == 1 + attrs = spans[0].attributes + assert attrs.get("gen_ai.request.model") == "bedrock/anthropic.claude-sonnet-4" + assert attrs.get("gen_ai.response.finish_reasons") == "['end_turn']" + assert attrs.get("gen_ai.usage.input_tokens") == 10 + + +@pytest.mark.asyncio +async def test_litellm_anthropic_acreate_delegating_to_acompletion_emits_single_span( # pylint: disable=unused-argument + agent, exporter, litellm_instrumentor +): + async def _fake_async_completion(*_args, **_kwargs): + return _fake_model_response() + + async def _delegating_acreate(*_args, **kwargs): + return await litellm.acompletion( + model=kwargs["model"], + messages=kwargs["messages"], + ) + + with patch("litellm.main.acompletion", new=_fake_async_completion), patch( + "litellm.anthropic_interface.messages.acreate", new=_delegating_acreate + ): + litellm_instrumentor.instrument() + await litellm.anthropic.messages.acreate(**_anthropic_messages_kwargs()) + + spans = _litellm_spans(exporter.get_finished_spans()) + exporter.clear() + assert len(spans) == 1 + attrs = spans[0].attributes + assert attrs.get("gen_ai.request.model") == "bedrock/anthropic.claude-sonnet-4" + assert attrs.get("gen_ai.usage.input_tokens") == 3 + assert attrs.get("gen_ai.usage.output_tokens") == 5 + + +@pytest.mark.asyncio +async def test_litellm_anthropic_acreate_streaming_defers_until_consumed( # pylint: disable=unused-argument + agent, exporter, litellm_instrumentor +): + async def _streaming_acreate(*_args, **_kwargs): + return _FakeAnthropicStream( + [ + {"type": "content_block_delta", "delta": {"text": "hello"}}, + { + "id": "msg_stream", + "model": "claude-sonnet-4", + "stop_reason": "end_turn", + "usage": {"input_tokens": 6, "output_tokens": 2}, + }, + ] + ) + + with patch("litellm.anthropic_interface.messages.acreate", new=_streaming_acreate): + litellm_instrumentor.instrument() + stream = await litellm.anthropic.messages.acreate( + **_anthropic_messages_kwargs(), stream=True + ) + + assert len(_litellm_spans(exporter.get_finished_spans())) == 0 + + chunks = [chunk async for chunk in stream] + assert len(chunks) == 2 + + spans = _litellm_spans(exporter.get_finished_spans()) + exporter.clear() + assert len(spans) == 1 + attrs = spans[0].attributes + assert attrs.get("gen_ai.request.streaming") == "True" + assert attrs.get("gen_ai.response.finish_reasons") == "['end_turn']" + assert attrs.get("gen_ai.usage.input_tokens") == 6 + assert attrs.get("gen_ai.usage.output_tokens") == 2 From f53cb56621f5d4d51b4ea34735da92a10df407c4 Mon Sep 17 00:00:00 2001 From: Shreyas Nagaraj Date: Thu, 10 Sep 2026 13:37:11 +0530 Subject: [PATCH 2/3] Aggregate Anthropic Messages stream events so native acreate spans get usage and stop_reason. OpenAI stream_chunk_builder never sees id/tokens/stop_reason on a single last chunk for SSE or message_start/message_delta events. Co-authored-by: Cursor --- .../instrumentation/litellm/__init__.py | 132 +++++++++++++++++- .../litellm/litellm_instrumentation_test.py | 23 ++- 2 files changed, 144 insertions(+), 11 deletions(-) diff --git a/src/harness_sdk/instrumentation/litellm/__init__.py b/src/harness_sdk/instrumentation/litellm/__init__.py index 9c911ca..0de4946 100644 --- a/src/harness_sdk/instrumentation/litellm/__init__.py +++ b/src/harness_sdk/instrumentation/litellm/__init__.py @@ -512,16 +512,138 @@ def _is_stream_response(response: Any) -> bool: return hasattr(response, "__anext__") or hasattr(response, "__next__") +_ANTHROPIC_STREAM_EVENT_TYPES = frozenset( + { + "message_start", + "message_delta", + "message_stop", + "content_block_start", + "content_block_delta", + "content_block_stop", + "ping", + } +) + + +def _decode_sse_chunk(chunk: Any) -> list[Any]: + """Turn raw SSE bytes/str frames into JSON events; pass other chunks through.""" + if not isinstance(chunk, (bytes, bytearray, str)): + return [chunk] + if isinstance(chunk, (bytes, bytearray)): + try: + text = chunk.decode("utf-8") + except Exception: # pylint: disable=broad-except + return [] + else: + text = chunk + stripped = text.strip() + if not stripped: + return [] + if stripped.startswith("{") or stripped.startswith("["): + try: + parsed = json.loads(stripped) + return parsed if isinstance(parsed, list) else [parsed] + except json.JSONDecodeError: + pass + + events: list[Any] = [] + for frame in text.replace("\r\n", "\n").split("\n\n"): + data_lines = [ + line[5:].lstrip() for line in frame.split("\n") if line.startswith("data:") + ] + if not data_lines: + continue + payload = "\n".join(data_lines).strip() + if not payload or payload == "[DONE]": + continue + try: + events.append(json.loads(payload)) + except json.JSONDecodeError: + logger.debug("LiteLLM: skipped undecodable SSE payload") + return events + + +def _is_anthropic_stream_event(obj: Any) -> bool: + event_type = _get_value(obj, "type") + return isinstance(event_type, str) and event_type in _ANTHROPIC_STREAM_EVENT_TYPES + + +def _looks_like_anthropic_messages_stream(chunks: list[Any]) -> bool: + for chunk in chunks: + if isinstance(chunk, (bytes, bytearray)) and b"data:" in chunk: + return True + if isinstance(chunk, str) and "data:" in chunk: + return True + for event in _decode_sse_chunk(chunk): + if _is_anthropic_stream_event(event): + return True + return False + + +def _merge_usage_fields(target: dict[str, Any], usage: Any) -> None: + if usage is None: + return + for key in ( + "input_tokens", + "output_tokens", + "cache_creation_input_tokens", + "cache_read_input_tokens", + ): + value = _get_value(usage, key) + if value is not None: + target[key] = value + + +def _aggregate_anthropic_messages_stream(chunks: list[Any]) -> dict[str, Any]: + """Merge Anthropic Messages SSE/event chunks into one response-shaped dict. + + Native streams never put ``id`` / ``usage`` / ``stop_reason`` on a single + last chunk. ``message_start`` carries id, model, and input tokens; + ``message_delta`` carries ``delta.stop_reason`` and output tokens. + """ + aggregated: dict[str, Any] = {"usage": {}} + usage = aggregated["usage"] + for chunk in chunks: + for event in _decode_sse_chunk(chunk): + event_type = _get_value(event, "type") + if event_type == "message_start": + message = _get_value(event, "message") or {} + response_id = _get_value(message, "id") + if response_id is not None: + aggregated["id"] = response_id + model = _get_value(message, "model") + if model is not None: + aggregated["model"] = model + _merge_usage_fields(usage, _get_value(message, "usage")) + continue + if event_type == "message_delta": + delta = _get_value(event, "delta") or {} + stop_reason = _get_value(delta, "stop_reason") or _get_value( + event, "stop_reason" + ) + if stop_reason is not None: + aggregated["stop_reason"] = stop_reason + _merge_usage_fields(usage, _get_value(event, "usage")) + _merge_usage_fields(usage, _get_value(delta, "usage")) + continue + _merge_usage_fields(usage, _get_value(event, "usage")) + if not usage: + aggregated.pop("usage", None) + return aggregated + + def _aggregate_stream_response(chunks: list[Any], messages: Any) -> Any: - """Rebuild a complete ``ModelResponse`` from streamed chunks. + """Rebuild a complete response from streamed chunks. - Uses ``litellm.stream_chunk_builder`` (the same helper LiteLLM uses - internally) so usage, choices, finish_reason and content are aggregated - exactly as they would be for a non-streaming call. Falls back to the last - chunk that carries usage if the builder is unavailable or fails. + Anthropic Messages streams (SSE bytes or ``message_start`` / + ``message_delta`` events) are merged across the whole stream. OpenAI-shaped + LiteLLM streams still use ``litellm.stream_chunk_builder``, then fall back + to the last chunk that carries usage. """ if not chunks: return None + if _looks_like_anthropic_messages_stream(chunks): + return _aggregate_anthropic_messages_stream(chunks) try: import litellm # pylint: disable=import-outside-toplevel diff --git a/test/instrumentation/litellm/litellm_instrumentation_test.py b/test/instrumentation/litellm/litellm_instrumentation_test.py index be6898f..4714591 100644 --- a/test/instrumentation/litellm/litellm_instrumentation_test.py +++ b/test/instrumentation/litellm/litellm_instrumentation_test.py @@ -611,13 +611,22 @@ async def test_litellm_anthropic_acreate_streaming_defers_until_consumed( # pyl async def _streaming_acreate(*_args, **_kwargs): return _FakeAnthropicStream( [ - {"type": "content_block_delta", "delta": {"text": "hello"}}, + ( + b'data: {"type":"message_start","message":{"id":"msg_stream",' + b'"type":"message","role":"assistant","model":"claude-sonnet-4",' + b'"content":[],"usage":{"input_tokens":6,"output_tokens":1}}}\n\n' + ), { - "id": "msg_stream", - "model": "claude-sonnet-4", - "stop_reason": "end_turn", - "usage": {"input_tokens": 6, "output_tokens": 2}, + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "hello"}, }, + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 2}, + }, + {"type": "message_stop"}, ] ) @@ -630,13 +639,15 @@ async def _streaming_acreate(*_args, **_kwargs): assert len(_litellm_spans(exporter.get_finished_spans())) == 0 chunks = [chunk async for chunk in stream] - assert len(chunks) == 2 + assert len(chunks) == 4 spans = _litellm_spans(exporter.get_finished_spans()) exporter.clear() assert len(spans) == 1 attrs = spans[0].attributes assert attrs.get("gen_ai.request.streaming") == "True" + assert attrs.get("gen_ai.response.id") == "msg_stream" + assert attrs.get("gen_ai.response.model") == "claude-sonnet-4" assert attrs.get("gen_ai.response.finish_reasons") == "['end_turn']" assert attrs.get("gen_ai.usage.input_tokens") == 6 assert attrs.get("gen_ai.usage.output_tokens") == 2 From 60ea6b63e39c3ed25d0d787c6314da2f9ec1e081 Mon Sep 17 00:00:00 2001 From: Shreyas Nagaraj Date: Thu, 10 Sep 2026 18:19:53 +0530 Subject: [PATCH 3/3] Add GenAI API type to LiteLLM spans Co-authored-by: Cursor --- .../instrumentation/litellm/__init__.py | 57 +++++++++++++------ .../litellm/litellm_instrumentation_test.py | 2 + 2 files changed, 42 insertions(+), 17 deletions(-) diff --git a/src/harness_sdk/instrumentation/litellm/__init__.py b/src/harness_sdk/instrumentation/litellm/__init__.py index 0de4946..51e7a5d 100644 --- a/src/harness_sdk/instrumentation/litellm/__init__.py +++ b/src/harness_sdk/instrumentation/litellm/__init__.py @@ -47,6 +47,9 @@ _LITELLM_MAIN = "litellm.main" _LITELLM_REQUEST_SPAN_NAME = "litellm_request" +_GENAI_API_TYPE_ATTRIBUTE = "GENAI_API_TYPE" +_OPENAI_API_TYPE = "openai" +_ANTHROPIC_API_TYPE = "anthropic" _ANTHROPIC_MESSAGES_MODULES = ( "litellm.anthropic_interface.messages", "litellm.anthropic.messages", @@ -88,15 +91,15 @@ _LITELLM_PROVIDER_HEADER_PREFIX = "llm_provider-" _WRAPPED_FUNCTIONS = ( - ("completion", False), - ("acompletion", True), - ("embedding", False), - ("aembedding", True), + ("completion", False, _OPENAI_API_TYPE), + ("acompletion", True, _OPENAI_API_TYPE), + ("embedding", False, _OPENAI_API_TYPE), + ("aembedding", True, _OPENAI_API_TYPE), ) _ANTHROPIC_MESSAGES_FUNCTIONS = ( - ("create", False), - ("acreate", True), + ("create", False, _ANTHROPIC_API_TYPE), + ("acreate", True, _ANTHROPIC_API_TYPE), ) _PROVIDER_NAME_MAP = { @@ -197,6 +200,7 @@ class _PreCallSpanContext: payload: Any kwargs: dict[str, Any] call_type: str + api_type: str def _set_pre_call_request_attributes( @@ -215,6 +219,9 @@ def _set_pre_call_request_attributes( otel_logger.safe_set_attribute( span, "gen_ai.provider.name", _canonical_provider_name(provider) ) + otel_logger.safe_set_attribute( + span, _GENAI_API_TYPE_ATTRIBUTE, pre_call.api_type + ) otel_logger.safe_set_attribute(span, "gen_ai.framework", "litellm") otel_logger.safe_set_attribute( span, @@ -837,12 +844,17 @@ def _fail_pre_call_span(span: Any, exc: BaseException, *, blocked: bool = False) def _start_evaluated_span( otel_logger: Any, func_name: str, + api_type: str, args: tuple[Any, ...], kwargs: dict[str, Any], ) -> Any: model, payload = _extract_model_and_input(args, kwargs) pre_call = _PreCallSpanContext( - model=model, payload=payload, kwargs=kwargs, call_type=func_name + model=model, + payload=payload, + kwargs=kwargs, + call_type=func_name, + api_type=api_type, ) span = otel_logger.tracer.start_span(_LITELLM_REQUEST_SPAN_NAME) try: @@ -861,6 +873,7 @@ def _start_evaluated_span( class _LiteLLMSpanRun: otel_logger: Any func_name: str + api_type: str args: tuple[Any, ...] kwargs: dict[str, Any] request_model: Optional[str] = None @@ -871,7 +884,11 @@ class _LiteLLMSpanRun: def __enter__(self) -> "_LiteLLMSpanRun": self.request_model, _ = _extract_model_and_input(self.args, self.kwargs) self.span = _start_evaluated_span( - self.otel_logger, self.func_name, self.args, self.kwargs + self.otel_logger, + self.func_name, + self.api_type, + self.args, + self.kwargs, ) self.token = _activate_span(self.span) self.guard = _LITELLM_SPAN_ACTIVE.set(True) @@ -923,7 +940,9 @@ def __exit__( return False -def _make_wrapper(func_name: str, is_async: bool) -> Callable[..., Any]: +def _make_wrapper( + func_name: str, is_async: bool, api_type: str +) -> Callable[..., Any]: otel_logger = _get_otel_logger() def _sync_wrapper( @@ -937,7 +956,9 @@ def _sync_wrapper( if _LITELLM_SPAN_ACTIVE.get(): return wrapped(*args, **kwargs) - with _LiteLLMSpanRun(otel_logger, func_name, args, kwargs) as span_run: + with _LiteLLMSpanRun( + otel_logger, func_name, api_type, args, kwargs + ) as span_run: response = wrapped(*args, **kwargs) if _is_stream_response(response): return span_run.wrap_stream(response) @@ -953,7 +974,9 @@ async def _async_wrapper( if _LITELLM_SPAN_ACTIVE.get(): return await wrapped(*args, **kwargs) - with _LiteLLMSpanRun(otel_logger, func_name, args, kwargs) as span_run: + with _LiteLLMSpanRun( + otel_logger, func_name, api_type, args, kwargs + ) as span_run: response = await wrapped(*args, **kwargs) if _is_stream_response(response): return span_run.wrap_stream(response) @@ -1003,20 +1026,20 @@ def _iter_anthropic_messages_modules() -> list[tuple[str, Any]]: def _wrap_anthropic_messages() -> None: for mod_name, mod in _iter_anthropic_messages_modules(): - for func_name, is_async in _ANTHROPIC_MESSAGES_FUNCTIONS: + for func_name, is_async, api_type in _ANTHROPIC_MESSAGES_FUNCTIONS: if not hasattr(mod, func_name): continue wrapt.wrap_function_wrapper( mod_name, func_name, - _make_wrapper(func_name, is_async), + _make_wrapper(func_name, is_async, api_type), ) _rebind_public_function(mod, func_name) def _unwrap_anthropic_messages() -> None: for _mod_name, mod in _iter_anthropic_messages_modules(): - for func_name, _ in _ANTHROPIC_MESSAGES_FUNCTIONS: + for func_name, _, _ in _ANTHROPIC_MESSAGES_FUNCTIONS: if not hasattr(mod, func_name): continue try: @@ -1047,11 +1070,11 @@ def instrument(self, **_kwargs: Any) -> None: import litellm # pylint: disable=import-outside-toplevel main_mod = __import__(_LITELLM_MAIN, fromlist=["*"]) - for func_name, is_async in _WRAPPED_FUNCTIONS: + for func_name, is_async, api_type in _WRAPPED_FUNCTIONS: wrapt.wrap_function_wrapper( _LITELLM_MAIN, func_name, - _make_wrapper(func_name, is_async), + _make_wrapper(func_name, is_async, api_type), ) if hasattr(litellm, func_name): setattr(litellm, func_name, getattr(main_mod, func_name)) @@ -1072,7 +1095,7 @@ def uninstrument(self, **_kwargs: Any) -> None: errors: list[Exception] = [] mod = import_module(_LITELLM_MAIN) - for func_name, _ in _WRAPPED_FUNCTIONS: + for func_name, _, _ in _WRAPPED_FUNCTIONS: try: unwrap(mod, func_name) if hasattr(litellm, func_name): diff --git a/test/instrumentation/litellm/litellm_instrumentation_test.py b/test/instrumentation/litellm/litellm_instrumentation_test.py index 4714591..6ab8790 100644 --- a/test/instrumentation/litellm/litellm_instrumentation_test.py +++ b/test/instrumentation/litellm/litellm_instrumentation_test.py @@ -84,6 +84,7 @@ def test_litellm_completion_span_has_gen_ai_attributes(agent, exporter, litellm_ assert attrs.get("gen_ai.request.model") == "gpt-4o-mini" assert attrs.get("gen_ai.operation.name") == "chat" assert attrs.get("gen_ai.provider.name") == "openai" + assert attrs.get("GENAI_API_TYPE") == "openai" assert "gen_ai.system" not in attrs assert attrs.get("gen_ai.framework") == "litellm" assert attrs.get("gen_ai.response.model") == "gpt-4o-mini" @@ -547,6 +548,7 @@ async def _native_acreate(*_args, **_kwargs): assert attrs.get("gen_ai.request.model") == "bedrock/anthropic.claude-sonnet-4" assert attrs.get("gen_ai.operation.name") == "chat" assert attrs.get("gen_ai.provider.name") == "aws.bedrock" + assert attrs.get("GENAI_API_TYPE") == "anthropic" assert attrs.get("gen_ai.framework") == "litellm" assert attrs.get("gen_ai.request.max_tokens") == 32 assert attrs.get("gen_ai.request.temperature") == 0.2