From 6834e84b8b4717674a35a6d4a65ba04c7d77e98a Mon Sep 17 00:00:00 2001 From: Nightingalelyy Date: Mon, 17 Aug 2026 00:17:27 +0800 Subject: [PATCH] fix(instrumentations): repair Elasticsearch, Google ADK, and Google GenAI tracing --- ...elasticsearch-google-adk-google-genai.json | 8 + .../pyproject.toml | 4 + .../_processor.py | 147 ++++++++-- .../tests/test_instrumentation.py | 267 ++++++++++++++---- .../_instrumentation.py | 76 +++-- .../_otel_emitter.py | 41 ++- .../_translator.py | 32 ++- .../tests/test_instrumentation.py | 93 +++++- 8 files changed, 533 insertions(+), 135 deletions(-) create mode 100644 .release-intents/20260817-elasticsearch-google-adk-google-genai.json diff --git a/.release-intents/20260817-elasticsearch-google-adk-google-genai.json b/.release-intents/20260817-elasticsearch-google-adk-google-genai.json new file mode 100644 index 00000000..2b25d5eb --- /dev/null +++ b/.release-intents/20260817-elasticsearch-google-adk-google-genai.json @@ -0,0 +1,8 @@ +{ + "summary": "Restore canonical Elasticsearch, Google ADK, and Google GenAI tracing semantics and repeatable example validation", + "packages": { + "respan-instrumentation-elasticsearch": "patch", + "respan-instrumentation-google-adk": "patch", + "respan-instrumentation-google-genai": "patch" + } +} diff --git a/python-sdks/instrumentations/respan-instrumentation-elasticsearch/pyproject.toml b/python-sdks/instrumentations/respan-instrumentation-elasticsearch/pyproject.toml index 1da9aff1..c2946cac 100644 --- a/python-sdks/instrumentations/respan-instrumentation-elasticsearch/pyproject.toml +++ b/python-sdks/instrumentations/respan-instrumentation-elasticsearch/pyproject.toml @@ -26,6 +26,10 @@ markers = [ "integration: live network tests (requires an Elasticsearch server)", ] +[tool.poetry.group.dev.dependencies] +pytest = ">=8.0.0,<9.0.0" +pytest-asyncio = ">=0.24.0,<2.0.0" + [build-system] requires = ["poetry-core"] build-backend = "poetry.core.masonry.api" diff --git a/python-sdks/instrumentations/respan-instrumentation-google-adk/src/respan_instrumentation_google_adk/_processor.py b/python-sdks/instrumentations/respan-instrumentation-google-adk/src/respan_instrumentation_google_adk/_processor.py index 85e89b91..245c4c9f 100644 --- a/python-sdks/instrumentations/respan-instrumentation-google-adk/src/respan_instrumentation_google_adk/_processor.py +++ b/python-sdks/instrumentations/respan-instrumentation-google-adk/src/respan_instrumentation_google_adk/_processor.py @@ -4,6 +4,7 @@ import json import logging +import threading from collections import defaultdict from typing import Any @@ -25,9 +26,15 @@ ) from opentelemetry.semconv_ai import SpanAttributes as TLSpanAttributes from respan_instrumentation_openinference._translator import OpenInferenceTranslator -from respan_sdk.constants.llm_logging import LOG_TYPE_CHAT, LOG_TYPE_TOOL +from respan_sdk.constants.llm_logging import ( + LOG_TYPE_AGENT, + LOG_TYPE_CHAT, + LOG_TYPE_TOOL, + LOG_TYPE_WORKFLOW, +) from respan_sdk.constants.span_attributes import ( RESPAN_LOG_TYPE, + RESPAN_SESSION_ID, RESPAN_SPAN_HANDOFFS, RESPAN_SPAN_TOOL_CALLS, RESPAN_SPAN_TOOLS, @@ -60,6 +67,10 @@ _GOOGLE_ADK_LLM_RESPONSE = "gcp.vertex.agent.llm_response" _GOOGLE_ADK_TOOL_CALL_ARGS = "gcp.vertex.agent.tool_call_args" _GOOGLE_ADK_TOOL_RESPONSE = "gcp.vertex.agent.tool_response" +_SESSION_ID_ATTRS = ( + "session.id", + "gen_ai.conversation.id", +) _OFF_CONTRACT_ALIASES = { "model", @@ -85,6 +96,43 @@ def _safe_json_str(value: Any) -> str: return json.dumps(value, default=str) +def _span_trace_id(span: ReadableSpan) -> int | None: + context = getattr(span, "context", None) + if context is None: + get_span_context = getattr(span, "get_span_context", None) + if callable(get_span_context): + context = get_span_context() + trace_id = getattr(context, "trace_id", None) + return trace_id if isinstance(trace_id, int) and trace_id else None + + +def _first_user_prompt(attrs: dict[str, Any]) -> str | None: + prefix = f"{TLSpanAttributes.LLM_PROMPTS}." + indexes = sorted( + { + int(parts[0]) + for key in attrs + if key.startswith(prefix) + and len((parts := key[len(prefix) :].split(".", 1))) == 2 + and parts[0].isdigit() + } + ) + for index in indexes: + role = attrs.get(f"{prefix}{index}.role") + content = attrs.get(f"{prefix}{index}.content") + if role == "user" and content not in (None, ""): + return str(content) + return None + + +def _session_id(raw_attrs: dict[str, Any]) -> str | None: + for key in _SESSION_ID_ATTRS: + value = raw_attrs.get(key) + if value not in (None, ""): + return str(value) + return None + + def _parse_json(value: Any) -> Any: if not isinstance(value, str): return value @@ -120,7 +168,7 @@ def _collect_message_buckets( for key, value in attrs.items(): if not key.startswith(prefix): continue - rest = key[len(prefix):] + rest = key[len(prefix) :] parts = rest.split(".", 1) if not parts[0].isdigit() or len(parts) == 1: continue @@ -170,18 +218,20 @@ def _tool_call_signature(tool_call: dict[str, Any]) -> str: return json.dumps(normalized, default=str, sort_keys=True, separators=(",", ":")) -def _extract_tool_calls_from_message(raw: dict[str, Any]) -> list[dict[str, Any]] | None: +def _extract_tool_calls_from_message( + raw: dict[str, Any], +) -> list[dict[str, Any]] | None: tool_call_buckets: dict[int, dict[str, Any]] = defaultdict(dict) for field_key, field_val in raw.items(): if not field_key.startswith(_OI_MESSAGE_TOOL_CALLS_PREFIX): continue - rest = field_key[len(_OI_MESSAGE_TOOL_CALLS_PREFIX):] + rest = field_key[len(_OI_MESSAGE_TOOL_CALLS_PREFIX) :] parts = rest.split(".", 1) if not parts[0].isdigit() or len(parts) == 1: continue tc_field = parts[1] if tc_field.startswith(_OI_TOOL_CALL_PREFIX): - tc_field = tc_field[len(_OI_TOOL_CALL_PREFIX):] + tc_field = tc_field[len(_OI_TOOL_CALL_PREFIX) :] tool_call_buckets[int(parts[0])][tc_field] = field_val result: list[dict[str, Any]] = [] @@ -226,19 +276,19 @@ def _extract_message_content(raw: dict[str, Any]) -> Any: content_blocks: dict[int, dict[str, Any]] = defaultdict(dict) for field_key, field_val in raw.items(): if field_key.startswith(_OI_MESSAGE_CONTENT_PREFIX): - idx_str = field_key[len(_OI_MESSAGE_CONTENT_PREFIX):] + idx_str = field_key[len(_OI_MESSAGE_CONTENT_PREFIX) :] if idx_str.isdigit(): indexed_content.append((int(idx_str), field_val)) continue if not field_key.startswith(_OI_MESSAGE_CONTENTS_PREFIX): continue - rest = field_key[len(_OI_MESSAGE_CONTENTS_PREFIX):] + rest = field_key[len(_OI_MESSAGE_CONTENTS_PREFIX) :] parts = rest.split(".", 1) if not parts[0].isdigit() or len(parts) == 1: continue block_field = parts[1] if block_field.startswith(_OI_MESSAGE_CONTENT_BLOCK_PREFIX): - block_field = block_field[len(_OI_MESSAGE_CONTENT_BLOCK_PREFIX):] + block_field = block_field[len(_OI_MESSAGE_CONTENT_BLOCK_PREFIX) :] _set_nested_value(content_blocks[int(parts[0])], block_field, field_val) if content_blocks: @@ -335,9 +385,13 @@ def _adk_part_text_tool_calls( function_response = part.get("function_response") if isinstance(function_response, dict): - response = function_response.get("response") - if response is not None: - tool_response = _safe_json_str(response) + response_payload = { + key: function_response[key] + for key in ("id", "name", "response") + if function_response.get(key) is not None + } + if response_payload: + tool_response = _safe_json_str(response_payload) text_value = "\n".join(text_parts) if text_parts else None return text_value, tool_calls or None, tool_response @@ -360,11 +414,17 @@ def _set_adk_content_message( text = tool_response target = f"{target_prefix}.{index}" - attrs.setdefault(f"{target}.role", role) + if tool_response is not None: + attrs[f"{target}.role"] = role + else: + attrs.setdefault(f"{target}.role", role) if attrs.get(f"{target}.role") == "model": attrs[f"{target}.role"] = "assistant" if text is not None: - attrs.setdefault(f"{target}.content", text) + if tool_response is not None: + attrs[f"{target}.content"] = text + else: + attrs.setdefault(f"{target}.content", text) if tool_calls is not None: attrs.setdefault(f"{target}.tool_calls", _safe_json_str(tool_calls)) @@ -404,13 +464,17 @@ def _apply_google_adk_payload_fallbacks( if model: attrs.setdefault(TLSpanAttributes.LLM_REQUEST_MODEL, model) + inserted_system_prompt = False config = request.get("config") if isinstance(config, dict): system_instruction = config.get("system_instruction") if isinstance(system_instruction, str) and system_instruction: system_prompt = f"{TLSpanAttributes.LLM_PROMPTS}.0" - attrs.setdefault(f"{system_prompt}.role", "system") - attrs.setdefault(f"{system_prompt}.content", system_instruction) + system_content_key = f"{system_prompt}.content" + if system_content_key not in attrs: + attrs[f"{system_prompt}.role"] = "system" + attrs[system_content_key] = system_instruction + inserted_system_prompt = True tools = _extract_adk_tools(config) if tools is not None: @@ -421,9 +485,10 @@ def _apply_google_adk_payload_fallbacks( contents = request.get("contents") if isinstance(contents, list): - start_index = 1 if isinstance( - attrs.get(f"{TLSpanAttributes.LLM_PROMPTS}.0.content"), str - ) else 0 + first_prompt_role = attrs.get(f"{TLSpanAttributes.LLM_PROMPTS}.0.role") + start_index = ( + 1 if inserted_system_prompt or first_prompt_role == "system" else 0 + ) for offset, content in enumerate(contents): _set_adk_content_message( attrs, @@ -459,7 +524,9 @@ def _apply_google_adk_payload_fallbacks( total_tokens = usage.get("total_token_count") if isinstance(total_tokens, int): - attrs.setdefault(TLSpanAttributes.LLM_USAGE_TOTAL_TOKENS, total_tokens) + attrs.setdefault( + TLSpanAttributes.LLM_USAGE_TOTAL_TOKENS, total_tokens + ) content = response.get("content") if isinstance(content, dict): @@ -490,7 +557,7 @@ def _is_indexed_structured_message_attr(key: str) -> bool: prefixed = f"{prefix}." if not key.startswith(prefixed): continue - parts = key[len(prefixed):].split(".") + parts = key[len(prefixed) :].split(".") if len(parts) < 3 or not parts[0].isdigit(): return False return parts[1] in {"tool_calls", "function_call"} @@ -511,9 +578,7 @@ def _stringify_structured_message_values(attrs: dict[str, Any]) -> None: def _cleanup_google_adk_attrs(attrs: dict[str, Any], *, is_chat_span: bool) -> None: if is_chat_span and attrs.get(GEN_AI_SYSTEM) in (None, "", "gcp.vertex.agent"): provider = attrs.get(GEN_AI_PROVIDER_NAME) - attrs[GEN_AI_SYSTEM] = ( - str(provider).lower() if provider else "google" - ) + attrs[GEN_AI_SYSTEM] = str(provider).lower() if provider else "google" elif not is_chat_span: attrs.pop(GEN_AI_SYSTEM, None) @@ -521,6 +586,8 @@ def _cleanup_google_adk_attrs(attrs: dict[str, Any], *, is_chat_span: bool) -> N for key in (TLSpanAttributes.TRACELOOP_SPAN_KIND, *_OFF_CONTRACT_ALIASES): attrs.pop(key, None) + for key in _SESSION_ID_ATTRS: + attrs.pop(key, None) prefixes_to_remove = ( _GOOGLE_ADK_PREFIX, @@ -558,6 +625,9 @@ def _normalize_google_adk_attrs( target_prefix=TLSpanAttributes.LLM_COMPLETIONS, ) _apply_google_adk_payload_fallbacks(raw_attrs, attrs) + session_id = _session_id(raw_attrs) + if session_id is not None: + attrs.setdefault(RESPAN_SESSION_ID, session_id) _cleanup_google_adk_attrs(attrs, is_chat_span=is_chat_span) @@ -566,6 +636,8 @@ class GoogleADKSpanProcessor(SpanProcessor): def __init__(self) -> None: self._translator = OpenInferenceTranslator() + self._trace_context: dict[int, dict[str, str]] = {} + self._context_lock = threading.Lock() def on_start( self, @@ -582,9 +654,38 @@ def on_end(self, span: ReadableSpan) -> None: self._translator.on_end(span) translated_attrs = dict(getattr(span, "_attributes", None) or {}) _normalize_google_adk_attrs(raw_attrs, translated_attrs) + + trace_id = _span_trace_id(span) + log_type = translated_attrs.get(RESPAN_LOG_TYPE) + if trace_id is not None and log_type == LOG_TYPE_CHAT: + prompt = _first_user_prompt(translated_attrs) + session_id = translated_attrs.get(RESPAN_SESSION_ID) + if prompt is not None or session_id not in (None, ""): + with self._context_lock: + context = self._trace_context.setdefault(trace_id, {}) + if prompt is not None: + context.setdefault("prompt", prompt) + if session_id not in (None, ""): + context.setdefault("session_id", str(session_id)) + elif trace_id is not None and log_type in (LOG_TYPE_AGENT, LOG_TYPE_WORKFLOW): + with self._context_lock: + context = self._trace_context.pop(trace_id, {}) + if translated_attrs.get(TLSpanAttributes.TRACELOOP_ENTITY_INPUT) in ( + None, + "", + ) and context.get("prompt"): + translated_attrs[TLSpanAttributes.TRACELOOP_ENTITY_INPUT] = ( + _safe_json_str({"prompt": context["prompt"]}) + ) + if translated_attrs.get(RESPAN_SESSION_ID) in (None, "") and context.get( + "session_id" + ): + translated_attrs[RESPAN_SESSION_ID] = context["session_id"] span._attributes = translated_attrs def shutdown(self) -> None: + with self._context_lock: + self._trace_context.clear() return None def force_flush(self, timeout_millis: int = 30000) -> bool: diff --git a/python-sdks/instrumentations/respan-instrumentation-google-adk/tests/test_instrumentation.py b/python-sdks/instrumentations/respan-instrumentation-google-adk/tests/test_instrumentation.py index 35ecade8..42f84feb 100644 --- a/python-sdks/instrumentations/respan-instrumentation-google-adk/tests/test_instrumentation.py +++ b/python-sdks/instrumentations/respan-instrumentation-google-adk/tests/test_instrumentation.py @@ -37,11 +37,12 @@ def add_span_processor(self, processor): class FakeSpan: - def __init__(self, attrs, name="call_llm"): + def __init__(self, attrs, name="call_llm", trace_id=None): self.name = name self._attributes = dict(attrs) self.attributes = self._attributes self.instrumentation_scope = SimpleNamespace(name=GOOGLE_ADK_SCOPE_NAME) + self.context = SimpleNamespace(trace_id=trace_id) def _install_fake_modules(monkeypatch): @@ -234,45 +235,51 @@ def import_module_raises(module_name): def test_processor_promotes_google_adk_payloads_and_strips_local_noise(): - span = FakeSpan({ - "openinference.span.kind": "LLM", - "gen_ai.system": "gcp.vertex.agent", - "llm.provider": "google", - "gcp.vertex.agent.llm_request": json.dumps({ - "model": "openai/gpt-4o", - "config": { - "system_instruction": "You are concise.", - "tools": [ - { - "function_declarations": [ + span = FakeSpan( + { + "openinference.span.kind": "LLM", + "gen_ai.system": "gcp.vertex.agent", + "llm.provider": "google", + "gcp.vertex.agent.llm_request": json.dumps( + { + "model": "openai/gpt-4o", + "config": { + "system_instruction": "You are concise.", + "tools": [ { - "name": "get_weather", - "description": "Get weather.", - "parameters": {"type": "OBJECT"}, + "function_declarations": [ + { + "name": "get_weather", + "description": "Get weather.", + "parameters": {"type": "OBJECT"}, + } + ] } - ] - } - ], - }, - "contents": [ + ], + }, + "contents": [ + { + "role": "user", + "parts": [{"text": "Weather in Tokyo?"}], + } + ], + } + ), + "gcp.vertex.agent.llm_response": json.dumps( { - "role": "user", - "parts": [{"text": "Weather in Tokyo?"}], + "content": { + "role": "model", + "parts": [{"text": "It is sunny."}], + }, + "usage_metadata": { + "prompt_token_count": 12, + "candidates_token_count": 5, + "total_token_count": 17, + }, } - ], - }), - "gcp.vertex.agent.llm_response": json.dumps({ - "content": { - "role": "model", - "parts": [{"text": "It is sunny."}], - }, - "usage_metadata": { - "prompt_token_count": 12, - "candidates_token_count": 5, - "total_token_count": 17, - }, - }), - }) + ), + } + ) GoogleADKSpanProcessor().on_end(span) @@ -303,19 +310,21 @@ def test_processor_promotes_google_adk_payloads_and_strips_local_noise(): def test_processor_promotes_openinference_message_content_blocks(): - span = FakeSpan({ - "openinference.span.kind": "LLM", - "llm.input_messages.0.message.role": "user", - "llm.input_messages.0.message.contents.0.message_content.type": "text", - "llm.input_messages.0.message.contents.0.message_content.text": "Hello", - "llm.output_messages.0.message.role": "model", - "llm.output_messages.0.message.contents.0.message_content.type": "text", - "llm.output_messages.0.message.contents.0.message_content.text": "Hi there", - "llm.output_messages.0.message.tool_calls.0.tool_call.function.name": "lookup", - "llm.output_messages.0.message.tool_calls.0.tool_call.function.arguments": ( - '{"query":"weather"}' - ), - }) + span = FakeSpan( + { + "openinference.span.kind": "LLM", + "llm.input_messages.0.message.role": "user", + "llm.input_messages.0.message.contents.0.message_content.type": "text", + "llm.input_messages.0.message.contents.0.message_content.text": "Hello", + "llm.output_messages.0.message.role": "model", + "llm.output_messages.0.message.contents.0.message_content.type": "text", + "llm.output_messages.0.message.contents.0.message_content.text": "Hi there", + "llm.output_messages.0.message.tool_calls.0.tool_call.function.name": "lookup", + "llm.output_messages.0.message.tool_calls.0.tool_call.function.arguments": ( + '{"query":"weather"}' + ), + } + ) GoogleADKSpanProcessor().on_end(span) @@ -339,13 +348,16 @@ def test_processor_promotes_openinference_message_content_blocks(): def test_processor_cleans_google_adk_tool_attrs(): - span = FakeSpan({ - "openinference.span.kind": "TOOL", - "tool.name": "get_weather", - "gen_ai.tool.name": "get_weather", - "gcp.vertex.agent.tool_call_args": '{"city":"Paris"}', - "gcp.vertex.agent.tool_response": '{"result":"sunny"}', - }, name="execute_tool get_weather") + span = FakeSpan( + { + "openinference.span.kind": "TOOL", + "tool.name": "get_weather", + "gen_ai.tool.name": "get_weather", + "gcp.vertex.agent.tool_call_args": '{"city":"Paris"}', + "gcp.vertex.agent.tool_response": '{"result":"sunny"}', + }, + name="execute_tool get_weather", + ) GoogleADKSpanProcessor().on_end(span) @@ -358,3 +370,146 @@ def test_processor_cleans_google_adk_tool_attrs(): assert "gen_ai.tool.name" not in span._attributes assert "gcp.vertex.agent.tool_call_args" not in span._attributes assert "gcp.vertex.agent.tool_response" not in span._attributes + + +def test_processor_preserves_tool_result_identity_in_chat_history(): + span = FakeSpan( + { + "openinference.span.kind": "LLM", + "llm.input_messages.0.message.role": "tool", + "llm.input_messages.0.message.content": '{"result":"sunny"}', + "gcp.vertex.agent.llm_request": json.dumps( + { + "model": "test-model", + "contents": [ + { + "role": "user", + "parts": [ + { + "function_response": { + "id": "call-weather-1", + "name": "get_weather", + "response": {"result": "sunny"}, + } + } + ], + } + ], + } + ), + } + ) + + GoogleADKSpanProcessor().on_end(span) + + assert span._attributes["gen_ai.prompt.0.role"] == "tool" + assert json.loads(span._attributes["gen_ai.prompt.0.content"]) == { + "id": "call-weather-1", + "name": "get_weather", + "response": {"result": "sunny"}, + } + + +def test_processor_aligns_adk_history_after_existing_system_prompt(): + span = FakeSpan( + { + "openinference.span.kind": "LLM", + "llm.input_messages.0.message.role": "system", + "llm.input_messages.0.message.content": "Use the weather tool.", + "llm.input_messages.1.message.role": "user", + "llm.input_messages.1.message.content": "Weather in Paris?", + "llm.input_messages.2.message.role": "assistant", + "llm.input_messages.2.message.tool_calls.0.tool_call.id": ( + "call-weather-1" + ), + "llm.input_messages.2.message.tool_calls.0.tool_call.function.name": ( + "get_weather" + ), + "llm.input_messages.2.message.tool_calls.0.tool_call.function.arguments": ( + '{"city":"Paris"}' + ), + "llm.input_messages.3.message.role": "tool", + "llm.input_messages.3.message.content": '{"result":"sunny"}', + "gcp.vertex.agent.llm_request": json.dumps( + { + "model": "test-model", + "contents": [ + {"role": "user", "parts": [{"text": "Weather in Paris?"}]}, + { + "role": "model", + "parts": [ + { + "function_call": { + "id": "call-weather-1", + "name": "get_weather", + "args": {"city": "Paris"}, + } + } + ], + }, + { + "role": "user", + "parts": [ + { + "function_response": { + "id": "call-weather-1", + "name": "get_weather", + "response": {"result": "sunny"}, + } + } + ], + }, + ], + } + ), + } + ) + + GoogleADKSpanProcessor().on_end(span) + + assert span._attributes["gen_ai.prompt.0.role"] == "system" + assert span._attributes["gen_ai.prompt.1.role"] == "user" + assert span._attributes["gen_ai.prompt.2.role"] == "assistant" + assert span._attributes["gen_ai.prompt.3.role"] == "tool" + assert json.loads(span._attributes["gen_ai.prompt.3.content"]) == { + "id": "call-weather-1", + "name": "get_weather", + "response": {"result": "sunny"}, + } + assert "gen_ai.prompt.4.role" not in span._attributes + + +def test_processor_promotes_session_and_carries_prompt_to_agent(): + processor = GoogleADKSpanProcessor() + llm_span = FakeSpan( + { + "openinference.span.kind": "LLM", + "session.id": "session-123", + "llm.input_messages.0.message.role": "user", + "llm.input_messages.0.message.content": "Plan a Tokyo day trip.", + "llm.output_messages.0.message.role": "model", + "llm.output_messages.0.message.content": "Here is a plan.", + }, + trace_id=1234, + ) + agent_span = FakeSpan( + { + "openinference.span.kind": "AGENT", + "input.value": "", + "output.value": '{"answer":"Here is a plan."}', + }, + name="travel_agent", + trace_id=1234, + ) + + processor.on_end(llm_span) + processor.on_end(agent_span) + + assert llm_span._attributes["respan.sessions.session_identifier"] == "session-123" + assert agent_span._attributes["respan.sessions.session_identifier"] == "session-123" + assert json.loads(agent_span._attributes["traceloop.entity.input"]) == { + "prompt": "Plan a Tokyo day trip." + } + assert "session.id" not in llm_span._attributes + assert "gen_ai.request.model" not in agent_span._attributes + assert processor._trace_context == {} diff --git a/python-sdks/instrumentations/respan-instrumentation-google-genai/src/respan_instrumentation_google_genai/_instrumentation.py b/python-sdks/instrumentations/respan-instrumentation-google-genai/src/respan_instrumentation_google_genai/_instrumentation.py index 422aa3d2..a6738e0e 100644 --- a/python-sdks/instrumentations/respan-instrumentation-google-genai/src/respan_instrumentation_google_genai/_instrumentation.py +++ b/python-sdks/instrumentations/respan-instrumentation-google-genai/src/respan_instrumentation_google_genai/_instrumentation.py @@ -47,37 +47,60 @@ def _load_models_classes() -> tuple[type[Any], type[Any]]: def _emit_span_safely( *, - kwargs: dict[str, Any], + request_kwargs: dict[str, Any], start_ns: int, response_or_chunks: Any = None, error_message: str | None = None, status_code: int = 200, + is_streaming: bool = False, ) -> None: emit_generate_content_span( - request_kwargs=request_kwargs_from_call(kwargs=kwargs), + request_kwargs=request_kwargs, start_ns=start_ns, response_or_chunks=response_or_chunks, error_message=error_message, status_code=status_code, + is_streaming=is_streaming, ) +def _status_code_from_exception(exc: BaseException) -> int: + candidates = [ + getattr(exc, "status_code", None), + getattr(exc, "code", None), + getattr(getattr(exc, "response", None), "status_code", None), + getattr(getattr(exc, "response", None), "status", None), + ] + for arg in getattr(exc, "args", ()): + if isinstance(arg, dict): + candidates.extend((arg.get("status_code"), arg.get("code"))) + for candidate in candidates: + try: + status_code = int(candidate) + except (TypeError, ValueError): + continue + if 400 <= status_code <= 599: + return status_code + return 500 + + def _wrap_sync_generate_content(original: Any) -> Any: def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any: start_ns = time.time_ns() + request_kwargs = request_kwargs_from_call(args=args, kwargs=kwargs) try: response = original(self, *args, **kwargs) except Exception as exc: _emit_span_safely( - kwargs=kwargs, + request_kwargs=request_kwargs, start_ns=start_ns, error_message=str(exc), - status_code=500, + status_code=_status_code_from_exception(exc), ) raise _emit_span_safely( - kwargs=kwargs, + request_kwargs=request_kwargs, start_ns=start_ns, response_or_chunks=response, ) @@ -89,7 +112,7 @@ def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any: def _instrument_sync_stream( *, iterator: Iterator[Any], - kwargs: dict[str, Any], + request_kwargs: dict[str, Any], start_ns: int, ) -> Iterator[Any]: chunks: list[Any] = [] @@ -99,37 +122,41 @@ def _instrument_sync_stream( yield chunk except Exception as exc: _emit_span_safely( - kwargs=kwargs, + request_kwargs=request_kwargs, start_ns=start_ns, response_or_chunks=chunks, error_message=str(exc), - status_code=500, + status_code=_status_code_from_exception(exc), + is_streaming=True, ) raise else: _emit_span_safely( - kwargs=kwargs, + request_kwargs=request_kwargs, start_ns=start_ns, response_or_chunks=chunks, + is_streaming=True, ) def _wrap_sync_generate_content_stream(original: Any) -> Any: def wrapper(self: Any, *args: Any, **kwargs: Any) -> Iterator[Any]: start_ns = time.time_ns() + request_kwargs = request_kwargs_from_call(args=args, kwargs=kwargs) try: iterator = original(self, *args, **kwargs) except Exception as exc: _emit_span_safely( - kwargs=kwargs, + request_kwargs=request_kwargs, start_ns=start_ns, error_message=str(exc), - status_code=500, + status_code=_status_code_from_exception(exc), + is_streaming=True, ) raise return _instrument_sync_stream( iterator=iterator, - kwargs=kwargs, + request_kwargs=request_kwargs, start_ns=start_ns, ) @@ -139,19 +166,20 @@ def wrapper(self: Any, *args: Any, **kwargs: Any) -> Iterator[Any]: def _wrap_async_generate_content(original: Any) -> Any: async def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any: start_ns = time.time_ns() + request_kwargs = request_kwargs_from_call(args=args, kwargs=kwargs) try: response = await original(self, *args, **kwargs) except Exception as exc: _emit_span_safely( - kwargs=kwargs, + request_kwargs=request_kwargs, start_ns=start_ns, error_message=str(exc), - status_code=500, + status_code=_status_code_from_exception(exc), ) raise _emit_span_safely( - kwargs=kwargs, + request_kwargs=request_kwargs, start_ns=start_ns, response_or_chunks=response, ) @@ -163,7 +191,7 @@ async def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any: async def _instrument_async_stream( *, async_iterator: AsyncIterator[Any], - kwargs: dict[str, Any], + request_kwargs: dict[str, Any], start_ns: int, ) -> AsyncIterator[Any]: chunks: list[Any] = [] @@ -173,37 +201,41 @@ async def _instrument_async_stream( yield chunk except Exception as exc: _emit_span_safely( - kwargs=kwargs, + request_kwargs=request_kwargs, start_ns=start_ns, response_or_chunks=chunks, error_message=str(exc), - status_code=500, + status_code=_status_code_from_exception(exc), + is_streaming=True, ) raise else: _emit_span_safely( - kwargs=kwargs, + request_kwargs=request_kwargs, start_ns=start_ns, response_or_chunks=chunks, + is_streaming=True, ) def _wrap_async_generate_content_stream(original: Any) -> Any: async def wrapper(self: Any, *args: Any, **kwargs: Any) -> AsyncIterator[Any]: start_ns = time.time_ns() + request_kwargs = request_kwargs_from_call(args=args, kwargs=kwargs) try: async_iterator = await original(self, *args, **kwargs) except Exception as exc: _emit_span_safely( - kwargs=kwargs, + request_kwargs=request_kwargs, start_ns=start_ns, error_message=str(exc), - status_code=500, + status_code=_status_code_from_exception(exc), + is_streaming=True, ) raise return _instrument_async_stream( async_iterator=async_iterator, - kwargs=kwargs, + request_kwargs=request_kwargs, start_ns=start_ns, ) diff --git a/python-sdks/instrumentations/respan-instrumentation-google-genai/src/respan_instrumentation_google_genai/_otel_emitter.py b/python-sdks/instrumentations/respan-instrumentation-google-genai/src/respan_instrumentation_google_genai/_otel_emitter.py index d919613e..204a54c3 100644 --- a/python-sdks/instrumentations/respan-instrumentation-google-genai/src/respan_instrumentation_google_genai/_otel_emitter.py +++ b/python-sdks/instrumentations/respan-instrumentation-google-genai/src/respan_instrumentation_google_genai/_otel_emitter.py @@ -8,6 +8,10 @@ from opentelemetry import context as context_api from opentelemetry import trace +from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import ( + GEN_AI_USAGE_INPUT_TOKENS, + GEN_AI_USAGE_OUTPUT_TOKENS, +) from opentelemetry.semconv_ai import LLMRequestTypeValues, SpanAttributes from respan_instrumentation_google_genai._constants import ( @@ -26,8 +30,6 @@ MODEL_KEY, PROMPT_TOKEN_COUNT_KEY, ROLE_KEY, - TOOL_CALLS_OVERRIDE_ATTR, - TOOLS_OVERRIDE_ATTR, TOTAL_TOKEN_COUNT_KEY, ) from respan_instrumentation_google_genai._translator import ( @@ -48,10 +50,11 @@ LLM_USAGE_COMPLETION_TOKENS, LLM_USAGE_PROMPT_TOKENS, RESPAN_LOG_TYPE, - RESPAN_SPAN_TOOL_CALLS, - RESPAN_SPAN_TOOLS, ) -from respan_sdk.utils.data_processing.id_processing import format_span_id, format_trace_id +from respan_sdk.utils.data_processing.id_processing import ( + format_span_id, + format_trace_id, +) from respan_tracing.utils.span_factory import build_readable_span, inject_span logger = logging.getLogger(__name__) @@ -71,13 +74,13 @@ def _current_trace_parent_ids() -> tuple[str | None, str | None]: return format_trace_id(trace_id=trace_id), format_span_id(span_id=span_id) -def _base_attrs() -> dict[str, Any]: +def _base_attrs(*, is_streaming: bool) -> dict[str, Any]: attrs = { GEN_AI_SYSTEM: GOOGLE_GENAI_SYSTEM_NAME, LLM_REQUEST_TYPE: LLMRequestTypeValues.CHAT.value, SpanAttributes.TRACELOOP_ENTITY_NAME: GOOGLE_GENAI_CHAT_SPAN_NAME, SpanAttributes.TRACELOOP_ENTITY_PATH: GOOGLE_GENAI_CHAT_SPAN_NAME, - SpanAttributes.TRACELOOP_SPAN_KIND: LLMRequestTypeValues.CHAT.value, + SpanAttributes.LLM_IS_STREAMING: is_streaming, RESPAN_LOG_TYPE: LOG_TYPE_CHAT, } workflow_name = context_api.get_value(SpanAttributes.TRACELOOP_ENTITY_NAME) @@ -114,17 +117,17 @@ def _set_output_attrs( tool_calls = extract_tool_calls(response_or_chunks=response_or_chunks) if tool_calls: - attrs[RESPAN_SPAN_TOOL_CALLS] = safe_json(value=tool_calls) - attrs[TOOL_CALLS_OVERRIDE_ATTR] = tool_calls - attrs[GEN_AI_COMPLETION_TOOL_CALLS_ATTR] = tool_calls + attrs[GEN_AI_COMPLETION_TOOL_CALLS_ATTR] = safe_json(value=tool_calls) usage = extract_usage(response_or_chunks=response_or_chunks) if PROMPT_TOKEN_COUNT_KEY in usage: attrs[LLM_USAGE_PROMPT_TOKENS] = usage[PROMPT_TOKEN_COUNT_KEY] + attrs[GEN_AI_USAGE_INPUT_TOKENS] = usage[PROMPT_TOKEN_COUNT_KEY] if CANDIDATES_TOKEN_COUNT_KEY in usage: attrs[LLM_USAGE_COMPLETION_TOKENS] = usage[CANDIDATES_TOKEN_COUNT_KEY] + attrs[GEN_AI_USAGE_OUTPUT_TOKENS] = usage[CANDIDATES_TOKEN_COUNT_KEY] if TOTAL_TOKEN_COUNT_KEY in usage: - attrs["gen_ai.usage.total_tokens"] = usage[TOTAL_TOKEN_COUNT_KEY] + attrs[SpanAttributes.LLM_USAGE_TOTAL_TOKENS] = usage[TOTAL_TOKEN_COUNT_KEY] def _set_request_attrs(attrs: dict[str, Any], request_kwargs: dict[str, Any]) -> None: @@ -136,8 +139,6 @@ def _set_request_attrs(attrs: dict[str, Any], request_kwargs: dict[str, Any]) -> tools = extract_tools(config=config) if tools: tools_json = safe_json(value=tools) - attrs[RESPAN_SPAN_TOOLS] = tools_json - attrs[TOOLS_OVERRIDE_ATTR] = tools attrs[LLM_REQUEST_FUNCTIONS_ATTR] = tools_json _set_input_attrs( @@ -151,8 +152,9 @@ def build_generate_content_attrs( *, request_kwargs: dict[str, Any], response_or_chunks: Any = None, + is_streaming: bool = False, ) -> dict[str, Any]: - attrs = _base_attrs() + attrs = _base_attrs(is_streaming=is_streaming) _set_request_attrs(attrs=attrs, request_kwargs=request_kwargs) if response_or_chunks is not None: _set_output_attrs(attrs=attrs, response_or_chunks=response_or_chunks) @@ -166,16 +168,27 @@ def emit_generate_content_span( response_or_chunks: Any = None, error_message: str | None = None, status_code: int = 200, + is_streaming: bool = False, ) -> None: """Build a ReadableSpan for a Google Gen AI generation and inject it.""" try: attrs = build_generate_content_attrs( request_kwargs=request_kwargs, response_or_chunks=response_or_chunks, + is_streaming=is_streaming, ) if error_message: attrs["error.message"] = error_message attrs.setdefault("status_code", status_code if status_code >= 400 else 500) + attrs.setdefault( + SpanAttributes.TRACELOOP_ENTITY_OUTPUT, + safe_json( + value={ + "error": error_message, + "status_code": status_code if status_code >= 400 else 500, + } + ), + ) trace_id, parent_id = _current_trace_parent_ids() span = build_readable_span( diff --git a/python-sdks/instrumentations/respan-instrumentation-google-genai/src/respan_instrumentation_google_genai/_translator.py b/python-sdks/instrumentations/respan-instrumentation-google-genai/src/respan_instrumentation_google_genai/_translator.py index 8cace651..2ab27126 100644 --- a/python-sdks/instrumentations/respan-instrumentation-google-genai/src/respan_instrumentation_google_genai/_translator.py +++ b/python-sdks/instrumentations/respan-instrumentation-google-genai/src/respan_instrumentation_google_genai/_translator.py @@ -197,7 +197,9 @@ def _normalize_parts(parts: Any) -> Any: return normalized_parts -def _normalize_content(content: Any, *, default_role: str = USER_ROLE) -> dict[str, Any]: +def _normalize_content( + content: Any, *, default_role: str = USER_ROLE +) -> dict[str, Any]: if isinstance(content, str): return {ROLE_KEY: default_role, CONTENT_KEY: content} if _is_part_like(content) and not _is_content_like(content): @@ -216,7 +218,10 @@ def normalize_input_messages(contents: Any, config: Any = None) -> list[dict[str messages.append( _normalize_content(system_instruction, default_role=SYSTEM_ROLE) if _is_content_like(system_instruction) - else {ROLE_KEY: SYSTEM_ROLE, CONTENT_KEY: _normalize_parts(system_instruction)} + else { + ROLE_KEY: SYSTEM_ROLE, + CONTENT_KEY: _normalize_parts(system_instruction), + } ) if contents is None: @@ -313,7 +318,11 @@ def extract_usage(response_or_chunks: Any) -> dict[str, int]: def _iter_response_contents(response_or_chunks: Any) -> Iterable[Any]: - chunks = response_or_chunks if isinstance(response_or_chunks, list) else [response_or_chunks] + chunks = ( + response_or_chunks + if isinstance(response_or_chunks, list) + else [response_or_chunks] + ) for response in chunks: if response is None: continue @@ -425,14 +434,21 @@ def extract_tools(config: Any) -> list[dict[str, Any]]: for field_name in BUILTIN_TOOL_FIELDS: value = _field(tool, field_name) if value is not None: - normalized_tools.append({TYPE_KEY: field_name, field_name: _dump_value(value)}) + normalized_tools.append( + {TYPE_KEY: field_name, field_name: _dump_value(value)} + ) return normalized_tools -def request_kwargs_from_call(kwargs: dict[str, Any]) -> dict[str, Any]: +def request_kwargs_from_call( + *, + args: tuple[Any, ...] = (), + kwargs: dict[str, Any], +) -> dict[str, Any]: + positional = dict(zip((MODEL_KEY, CONTENTS_KEY, CONFIG_KEY), args)) return { - MODEL_KEY: kwargs.get(MODEL_KEY), - CONTENTS_KEY: kwargs.get(CONTENTS_KEY), - CONFIG_KEY: kwargs.get(CONFIG_KEY), + MODEL_KEY: kwargs.get(MODEL_KEY, positional.get(MODEL_KEY)), + CONTENTS_KEY: kwargs.get(CONTENTS_KEY, positional.get(CONTENTS_KEY)), + CONFIG_KEY: kwargs.get(CONFIG_KEY, positional.get(CONFIG_KEY)), } diff --git a/python-sdks/instrumentations/respan-instrumentation-google-genai/tests/test_instrumentation.py b/python-sdks/instrumentations/respan-instrumentation-google-genai/tests/test_instrumentation.py index ca6c3f41..cd8a37e4 100644 --- a/python-sdks/instrumentations/respan-instrumentation-google-genai/tests/test_instrumentation.py +++ b/python-sdks/instrumentations/respan-instrumentation-google-genai/tests/test_instrumentation.py @@ -19,15 +19,16 @@ GOOGLE_GENAI_MODELS_MODULE, MODELS_CLASS_NAME, ) -from respan_instrumentation_google_genai._otel_emitter import build_generate_content_attrs +from respan_instrumentation_google_genai._otel_emitter import ( + build_generate_content_attrs, +) +from respan_instrumentation_google_genai._translator import request_kwargs_from_call from respan_sdk.constants.llm_logging import LOG_TYPE_CHAT from respan_sdk.constants.span_attributes import ( LLM_REQUEST_MODEL, LLM_USAGE_COMPLETION_TOKENS, LLM_USAGE_PROMPT_TOKENS, RESPAN_LOG_TYPE, - RESPAN_SPAN_TOOL_CALLS, - RESPAN_SPAN_TOOLS, ) @@ -83,7 +84,9 @@ def captured_spans(monkeypatch: pytest.MonkeyPatch) -> list[Any]: @pytest.fixture() def fake_google_genai(monkeypatch: pytest.MonkeyPatch) -> tuple[type[Any], type[Any]]: class Models: - def generate_content(self, *, model: str, contents: Any, config: Any = None) -> Obj: + def generate_content( + self, *, model: str, contents: Any, config: Any = None + ) -> Obj: return make_response(text=f"{model}: {contents}", usage=make_usage()) def generate_content_stream( @@ -163,7 +166,13 @@ def test_activate_patches_sync_generate_content_and_emits_chat_span( assert attrs["gen_ai.completion.0.content"] == "gemini-2.5-flash: Say hello" assert attrs[LLM_USAGE_PROMPT_TOKENS] == 3 assert attrs[LLM_USAGE_COMPLETION_TOKENS] == 4 - assert json.loads(attrs[RESPAN_SPAN_TOOLS])[0]["function"]["name"] == "weather_tool" + assert ( + json.loads(attrs[SpanAttributes.LLM_REQUEST_FUNCTIONS])[0]["function"]["name"] + == "weather_tool" + ) + assert "tools" not in attrs + assert "respan.span.tools" not in attrs + assert SpanAttributes.TRACELOOP_SPAN_KIND not in attrs instrumentor.deactivate() @@ -183,7 +192,9 @@ def test_thinking_tokens_are_included_in_output_usage() -> None: assert attrs[LLM_USAGE_PROMPT_TOKENS] == 100 assert attrs[LLM_USAGE_COMPLETION_TOKENS] == 850 - assert attrs["gen_ai.usage.total_tokens"] == 950 + assert attrs[SpanAttributes.LLM_USAGE_TOTAL_TOKENS] == 950 + assert attrs["gen_ai.usage.input_tokens"] == 100 + assert attrs["gen_ai.usage.output_tokens"] == 850 def test_stream_emits_one_span_after_iterator_is_consumed( @@ -207,6 +218,7 @@ def test_stream_emits_one_span_after_iterator_is_consumed( assert attrs["gen_ai.completion.0.content"] == "Hello world" assert attrs[LLM_USAGE_PROMPT_TOKENS] == 5 assert attrs[LLM_USAGE_COMPLETION_TOKENS] == 6 + assert attrs[SpanAttributes.LLM_IS_STREAMING] is True instrumentor.deactivate() @@ -234,7 +246,12 @@ async def run() -> None: assert [chunk.text for chunk in chunks] == ["async ", "stream"] assert len(captured_spans) == 2 assert captured_spans[0]._attributes[LLM_USAGE_PROMPT_TOKENS] == 7 - assert captured_spans[1]._attributes["gen_ai.completion.0.content"] == "async stream" + assert ( + captured_spans[1]._attributes["gen_ai.completion.0.content"] + == "async stream" + ) + assert captured_spans[0]._attributes[SpanAttributes.LLM_IS_STREAMING] is False + assert captured_spans[1]._attributes[SpanAttributes.LLM_IS_STREAMING] is True instrumentor.deactivate() @@ -261,8 +278,7 @@ def test_active_workflow_name_is_attached_to_injected_chat_span() -> None: context_api.detach(token) assert ( - attrs[SpanAttributes.TRACELOOP_WORKFLOW_NAME] - == "google_genai_generate_content" + attrs[SpanAttributes.TRACELOOP_WORKFLOW_NAME] == "google_genai_generate_content" ) @@ -285,7 +301,7 @@ def test_automatic_function_calling_history_promotes_tool_calls() -> None: response_or_chunks=response, ) - tool_calls = json.loads(attrs[RESPAN_SPAN_TOOL_CALLS]) + tool_calls = json.loads(attrs["gen_ai.completion.0.tool_calls"]) assert tool_calls == [ { "id": "call_1", @@ -296,7 +312,9 @@ def test_automatic_function_calling_history_promotes_tool_calls() -> None: }, } ] - assert attrs["gen_ai.completion.0.tool_calls"][0]["function"]["name"] == "get_weather" + assert tool_calls[0]["function"]["name"] == "get_weather" + assert "tool_calls" not in attrs + assert "respan.span.tool_calls" not in attrs def test_error_path_emits_failed_span( @@ -320,11 +338,59 @@ def raise_error(self: Any, *, model: str, contents: Any, config: Any = None) -> span = captured_spans[0] assert span.status.status_code.name == "ERROR" assert span._attributes["error.message"] == "boom" + assert json.loads(span._attributes[SpanAttributes.TRACELOOP_ENTITY_OUTPUT]) == { + "error": "boom", + "status_code": 500, + } assert span._attributes[LLM_REQUEST_MODEL] == "gemini-2.5-flash" instrumentor.deactivate() +def test_error_path_preserves_provider_status_code( + monkeypatch: pytest.MonkeyPatch, + fake_google_genai: tuple[type[Any], type[Any]], + captured_spans: list[Any], +) -> None: + Models, _ = fake_google_genai + + class APIError(RuntimeError): + def __init__(self) -> None: + super().__init__("model unavailable") + self.response = Obj(status_code=503) + + def raise_error(self: Any, *args: Any, **kwargs: Any) -> Any: + raise APIError() + + monkeypatch.setattr(Models, GENERATE_CONTENT_METHOD_NAME, raise_error) + instrumentor = GoogleGenAIInstrumentor() + instrumentor.activate() + + with pytest.raises(APIError, match="model unavailable"): + Models().generate_content("gemini-3-flash-preview", "fail") + + attrs = captured_spans[0]._attributes + assert attrs["status_code"] == 503 + assert attrs[LLM_REQUEST_MODEL] == "gemini-3-flash-preview" + assert "fail" in attrs[SpanAttributes.TRACELOOP_ENTITY_INPUT] + assert ( + json.loads(attrs[SpanAttributes.TRACELOOP_ENTITY_OUTPUT])["status_code"] == 503 + ) + + instrumentor.deactivate() + + +def test_request_kwargs_support_positional_sdk_arguments() -> None: + assert request_kwargs_from_call( + args=("gemini-3-flash-preview", "Hello", {"temperature": 0}), + kwargs={}, + ) == { + "model": "gemini-3-flash-preview", + "contents": "Hello", + "config": {"temperature": 0}, + } + + def test_deactivate_restores_original_methods( fake_google_genai: tuple[type[Any], type[Any]], ) -> None: @@ -343,4 +409,7 @@ def test_deactivate_restores_original_methods( assert getattr(Models, GENERATE_CONTENT_METHOD_NAME) is original_sync assert getattr(Models, GENERATE_CONTENT_STREAM_METHOD_NAME) is original_sync_stream assert getattr(AsyncModels, GENERATE_CONTENT_METHOD_NAME) is original_async - assert getattr(AsyncModels, GENERATE_CONTENT_STREAM_METHOD_NAME) is original_async_stream + assert ( + getattr(AsyncModels, GENERATE_CONTENT_STREAM_METHOD_NAME) + is original_async_stream + )