From 4eb081aaade5578335765c8aac4d1185a71cd9e6 Mon Sep 17 00:00:00 2001 From: Nightingalelyy Date: Sun, 16 Aug 2026 02:28:55 +0800 Subject: [PATCH] fix(instrumentations): repair AgentOps, AgentScope, and AgentSpec tracing --- ...20260816-agentscope-agentspec-tracing.json | 7 + .../README.md | 11 + .../_instrumentation.py | 28 +- .../tests/test_instrumentation.py | 37 +++ .../_instrumentation.py | 255 +++++++++++++++++- .../tests/test_instrumentation.py | 140 +++++++++- 6 files changed, 461 insertions(+), 17 deletions(-) create mode 100644 .release-intents/20260816-agentscope-agentspec-tracing.json diff --git a/.release-intents/20260816-agentscope-agentspec-tracing.json b/.release-intents/20260816-agentscope-agentspec-tracing.json new file mode 100644 index 00000000..62f51907 --- /dev/null +++ b/.release-intents/20260816-agentscope-agentspec-tracing.json @@ -0,0 +1,7 @@ +{ + "summary": "Fix AgentScope custom-model activation and AgentSpec trace semantics", + "packages": { + "respan-instrumentation-agentscope": "patch", + "respan-instrumentation-agentspec": "patch" + } +} diff --git a/python-sdks/instrumentations/respan-instrumentation-agentscope/README.md b/python-sdks/instrumentations/respan-instrumentation-agentscope/README.md index f003a583..e1a7d43b 100644 --- a/python-sdks/instrumentations/respan-instrumentation-agentscope/README.md +++ b/python-sdks/instrumentations/respan-instrumentation-agentscope/README.md @@ -42,6 +42,17 @@ respan = Respan( respan.flush() ``` +For custom model classes defined outside AgentScope's public model modules, +configure all instances on one instrumentor so Respan activates one lifecycle +identity while patching every distinct custom model class: + +```python +AgentScopeInstrumentor(models=[planner_model, reviewer_model, fallback_model]) +``` + +Use `model=...` for a single custom model. Do not pass `model` and `models` +together. + ## What Is Captured - Agent `reply()` and `reply_stream()` calls as `agent` spans. diff --git a/python-sdks/instrumentations/respan-instrumentation-agentscope/src/respan_instrumentation_agentscope/_instrumentation.py b/python-sdks/instrumentations/respan-instrumentation-agentscope/src/respan_instrumentation_agentscope/_instrumentation.py index 9e89346d..7d1e0ae9 100644 --- a/python-sdks/instrumentations/respan-instrumentation-agentscope/src/respan_instrumentation_agentscope/_instrumentation.py +++ b/python-sdks/instrumentations/respan-instrumentation-agentscope/src/respan_instrumentation_agentscope/_instrumentation.py @@ -8,7 +8,7 @@ import json import logging import time -from collections.abc import AsyncIterator, Mapping +from collections.abc import AsyncIterator, Mapping, Sequence from contextlib import contextmanager from contextvars import ContextVar from dataclasses import dataclass @@ -1077,12 +1077,20 @@ def __init__( *, agent: Any | None = None, model: Any | None = None, + models: Sequence[Any] | None = None, toolkit: Any | None = None, instrument_models: bool = True, instrument_tools: bool = True, ) -> None: + if model is not None and models is not None: + raise ValueError("Pass either model or models, not both") + self._agent = agent - self._model = model + self._models = ( + tuple(models) + if models is not None + else (() if model is None else (model,)) + ) self._toolkit = toolkit self._instrument_models = instrument_models self._instrument_tools = instrument_tools @@ -1205,11 +1213,17 @@ def activate(self) -> None: ) if self._instrument_models: - if self._model is not None: - patched_any |= self._patch_model_target( - self._model, - is_bound_method=True, - ) + if self._models: + seen_model_classes: set[type[Any]] = set() + for model in self._models: + model_class = type(model) + if model_class in seen_model_classes: + continue + seen_model_classes.add(model_class) + patched_any |= self._patch_model_target( + model, + is_bound_method=True, + ) else: try: model_module = self._load_model_module() diff --git a/python-sdks/instrumentations/respan-instrumentation-agentscope/tests/test_instrumentation.py b/python-sdks/instrumentations/respan-instrumentation-agentscope/tests/test_instrumentation.py index aa206e69..bc7e5085 100644 --- a/python-sdks/instrumentations/respan-instrumentation-agentscope/tests/test_instrumentation.py +++ b/python-sdks/instrumentations/respan-instrumentation-agentscope/tests/test_instrumentation.py @@ -160,6 +160,11 @@ async def __call__(self, messages, tools=None, tool_choice=None, **kwargs): return FakeChatResponse(content=[FakeToolCallBlock()]) +class FakeSecondChatModel(FakeChatModelBase): + async def __call__(self, messages, tools=None, tool_choice=None, **kwargs): + return FakeChatResponse(content=[FakeTextBlock("Second model.")]) + + class FakeKeyErrorModel(FakeChatModelBase): async def __call__(self, messages, tools=None, tool_choice=None, **kwargs): return FakeKeyErrorProxy( @@ -278,6 +283,38 @@ def test_activate_specific_instances_does_not_patch_classes(monkeypatch): assert fake_modules.agent_class.reply is original_agent_reply +def test_activate_patches_multiple_custom_model_classes_once(): + first_model = FakeChatModelBase() + duplicate_class_model = FakeChatModelBase() + second_model = FakeSecondChatModel() + original_first_call = FakeChatModelBase.__call__ + original_second_call = FakeSecondChatModel.__call__ + + instrumentor = AgentScopeInstrumentor( + agent=object(), + models=[first_model, duplicate_class_model, second_model], + instrument_tools=False, + ) + instrumentor.activate() + + assert FakeChatModelBase.__call__ is not original_first_call + assert FakeSecondChatModel.__call__ is not original_second_call + assert len(instrumentor._patches) == 2 + + instrumentor.deactivate() + + assert FakeChatModelBase.__call__ is original_first_call + assert FakeSecondChatModel.__call__ is original_second_call + + +def test_activate_rejects_model_and_models_together(): + with pytest.raises(ValueError, match="either model or models"): + AgentScopeInstrumentor( + model=FakeChatModelBase(), + models=[FakeSecondChatModel()], + ) + + def test_activate_skips_when_respan_tracing_is_disabled(monkeypatch, caplog): fake_modules = _install_fake_agentscope_modules(monkeypatch) RespanTracer(is_enabled=False) diff --git a/python-sdks/instrumentations/respan-instrumentation-agentspec/src/respan_instrumentation_agentspec/_instrumentation.py b/python-sdks/instrumentations/respan-instrumentation-agentspec/src/respan_instrumentation_agentspec/_instrumentation.py index c949304e..47ae0c47 100644 --- a/python-sdks/instrumentations/respan-instrumentation-agentspec/src/respan_instrumentation_agentspec/_instrumentation.py +++ b/python-sdks/instrumentations/respan-instrumentation-agentspec/src/respan_instrumentation_agentspec/_instrumentation.py @@ -1,21 +1,62 @@ """AgentSpec instrumentation plugin for Respan.""" +import hashlib +import json import logging from typing import Any +from uuid import UUID from opentelemetry import trace from opentelemetry.semconv.resource import ResourceAttributes from opentelemetry.semconv_ai import SpanAttributes as TLSpanAttributes from opentelemetry.sdk.trace.export import SpanProcessor +from opentelemetry.trace import SpanContext +from openinference.semconv.trace import SpanAttributes as OISpanAttributes from respan_instrumentation_openinference import OpenInferenceTranslator +from respan_sdk.constants.span_attributes import ( + RESPAN_CUSTOMER_PARAMS_ID, + RESPAN_LOG_TYPE, + RESPAN_METADATA, + RESPAN_SPAN_CUSTOM_ID, + RESPAN_SPAN_TOOL_CALLS, + RESPAN_SPAN_TOOLS, + RESPAN_THREADS_ID, + RESPAN_TRACE_GROUP_ID, +) from respan_tracing.core.tracer import RespanTracer logger = logging.getLogger(__name__) AGENTSPEC_INSTRUMENTATION_NAME = "agentspec" _ORIGINAL_ON_LLM_END_ATTR = "_respan_original_on_llm_end" +_ORIGINAL_ON_LLM_END_ASYNC_ATTR = "_respan_original_on_llm_end_async" _USAGE_PATCHED_ATTR = "_respan_usage_patched" TRACELOOP_WORKFLOW_NAME = TLSpanAttributes.TRACELOOP_WORKFLOW_NAME +TRACELOOP_ENTITY_INPUT = TLSpanAttributes.TRACELOOP_ENTITY_INPUT +TRACELOOP_ENTITY_OUTPUT = TLSpanAttributes.TRACELOOP_ENTITY_OUTPUT +OI_SESSION_ID = OISpanAttributes.SESSION_ID +PROMPT_PREFIX = f"{TLSpanAttributes.LLM_PROMPTS}." +COMPLETION_PREFIX = f"{TLSpanAttributes.LLM_COMPLETIONS}." + +_PROPAGATED_ATTRS = ( + RESPAN_CUSTOMER_PARAMS_ID, + RESPAN_THREADS_ID, + RESPAN_TRACE_GROUP_ID, + RESPAN_SPAN_CUSTOM_ID, +) +_OFF_CONTRACT_ALIASES = { + RESPAN_SPAN_TOOLS, + RESPAN_SPAN_TOOL_CALLS, + "tools", + "tool_calls", + "model", + "prompt_tokens", + "completion_tokens", + "total_request_tokens", + "span_tools", + "has_tool_calls", + "parallel_tool_calls", +} def _coerce_token_count(value: Any) -> int | None: @@ -94,14 +135,21 @@ def _patch_agentspec_langgraph_usage() -> None: except ImportError: return - handler_class = getattr(agentspec_tracing, "AgentSpecCallbackHandler", None) + handler_class = getattr( + agentspec_tracing, + "AgentSpecLlmCallbackHandler", + None, + ) or getattr(agentspec_tracing, "AgentSpecCallbackHandler", None) if handler_class is None or getattr(handler_class, _USAGE_PATCHED_ATTR, False): return - original_on_llm_end = handler_class.on_llm_end + original_on_llm_end = getattr(handler_class, "on_llm_end", None) + original_on_llm_end_async = getattr(handler_class, "on_llm_end_async", None) - def on_llm_end(self, response, *, run_id, parent_run_id=None, **kwargs): - run_id_str = str(run_id) + if original_on_llm_end is None and original_on_llm_end_async is None: + return + + def build_response_event(self, response, run_id_str): span = self.agentspec_spans_registry.get(run_id_str) if not isinstance(span, agentspec_tracing.AgentSpecLlmGenerationSpan): raise RuntimeError("LLM span not started; on_chat_model_start must run first") @@ -119,13 +167,43 @@ def on_llm_end(self, response, *, run_id, parent_run_id=None, **kwargs): input_tokens=input_tokens, output_tokens=output_tokens, ) + return span, event + + def on_llm_end(self, response, *, run_id, parent_run_id=None, **kwargs): + del parent_run_id, kwargs + run_id_str = str(run_id) + span, event = build_response_event(self, response, run_id_str) self._add_event(run_id_str, span, event) self._end_span(run_id_str, span) self.agentspec_spans_registry.pop(run_id_str, None) self.messages_in_process.pop(run_id_str, None) - setattr(handler_class, _ORIGINAL_ON_LLM_END_ATTR, original_on_llm_end) - handler_class.on_llm_end = on_llm_end + async def on_llm_end_async( + self, + response, + *, + run_id, + parent_run_id=None, + **kwargs, + ): + del parent_run_id, kwargs + run_id_str = str(run_id) + span, event = build_response_event(self, response, run_id_str) + await self._add_event_async(run_id_str, span, event) + await self._end_span_async(run_id_str, span) + self.agentspec_spans_registry.pop(run_id_str, None) + self.messages_in_process.pop(run_id_str, None) + + if original_on_llm_end is not None: + setattr(handler_class, _ORIGINAL_ON_LLM_END_ATTR, original_on_llm_end) + handler_class.on_llm_end = on_llm_end + if original_on_llm_end_async is not None: + setattr( + handler_class, + _ORIGINAL_ON_LLM_END_ASYNC_ATTR, + original_on_llm_end_async, + ) + handler_class.on_llm_end_async = on_llm_end_async setattr(handler_class, _USAGE_PATCHED_ATTR, True) @@ -135,17 +213,87 @@ def _restore_agentspec_langgraph_usage_patch() -> None: except ImportError: return - handler_class = getattr(agentspec_tracing, "AgentSpecCallbackHandler", None) + handler_class = getattr( + agentspec_tracing, + "AgentSpecLlmCallbackHandler", + None, + ) or getattr(agentspec_tracing, "AgentSpecCallbackHandler", None) if handler_class is None or not getattr(handler_class, _USAGE_PATCHED_ATTR, False): return original_on_llm_end = getattr(handler_class, _ORIGINAL_ON_LLM_END_ATTR, None) if original_on_llm_end is not None: handler_class.on_llm_end = original_on_llm_end - delattr(handler_class, _ORIGINAL_ON_LLM_END_ATTR) + delattr(handler_class, _ORIGINAL_ON_LLM_END_ATTR) + original_on_llm_end_async = getattr( + handler_class, + _ORIGINAL_ON_LLM_END_ASYNC_ATTR, + None, + ) + if original_on_llm_end_async is not None: + handler_class.on_llm_end_async = original_on_llm_end_async + delattr(handler_class, _ORIGINAL_ON_LLM_END_ASYNC_ATTR) setattr(handler_class, _USAGE_PATCHED_ATTR, False) +def _trace_id_from_session_id(session_id: Any) -> int | None: + if not session_id: + return None + try: + trace_id = UUID(str(session_id)).int + except (TypeError, ValueError, AttributeError): + trace_id = int.from_bytes( + hashlib.blake2b(str(session_id).encode(), digest_size=16).digest(), + "big", + ) + return trace_id or 1 + + +def _has_content(value: Any) -> bool: + if value is None: + return False + if isinstance(value, str): + return value.strip() not in {"", "{}", "[]", "null"} + return True + + +def _parse_structured_value(value: Any) -> Any: + if not isinstance(value, str): + return value + try: + return json.loads(value) + except (TypeError, ValueError): + return value + + +def _indexed_messages(attrs: dict[str, Any], prefix: str) -> list[dict[str, Any]]: + messages: list[dict[str, Any]] = [] + index = 0 + while True: + role_key = f"{prefix}{index}.role" + content_key = f"{prefix}{index}.content" + tool_calls_key = f"{prefix}{index}.tool_calls" + tool_call_id_key = f"{prefix}{index}.tool_call_id" + if not any( + key in attrs + for key in (role_key, content_key, tool_calls_key, tool_call_id_key) + ): + break + + message: dict[str, Any] = {} + if role_key in attrs: + message["role"] = attrs[role_key] + if content_key in attrs: + message["content"] = _parse_structured_value(attrs[content_key]) + if tool_calls_key in attrs: + message["tool_calls"] = _parse_structured_value(attrs[tool_calls_key]) + if tool_call_id_key in attrs: + message["tool_call_id"] = attrs[tool_call_id_key] + messages.append(message) + index += 1 + return messages + + class _TranslatedProcessorChain(SpanProcessor): """Run Respan's OI translator before the active export processors.""" @@ -159,6 +307,92 @@ def __init__( self._translator = translator self._processors = processors self._workflow_name = workflow_name + self._trace_state: dict[int, dict[str, Any]] = {} + + @staticmethod + def _replace_trace_id(context: Any, trace_id: int) -> SpanContext: + return SpanContext( + trace_id=trace_id, + span_id=context.span_id, + is_remote=context.is_remote, + trace_flags=context.trace_flags, + trace_state=context.trace_state, + ) + + def _normalize_trace_id(self, span: Any) -> int | None: + attributes = getattr(span, "_attributes", None) + get_span_context = getattr(span, "get_span_context", None) + if attributes is None or not callable(get_span_context): + return None + + trace_id = _trace_id_from_session_id(attributes.get(OI_SESSION_ID)) + if trace_id is None: + return None + + context = get_span_context() + if context.trace_id != trace_id: + span._context = self._replace_trace_id(context, trace_id) + + parent = getattr(span, "parent", None) + if parent is not None and parent.trace_id != trace_id: + span._parent = self._replace_trace_id(parent, trace_id) + return trace_id + + def _enrich_span_contract(self, span: Any, trace_id: int | None) -> bool: + attributes = getattr(span, "_attributes", None) + if attributes is None or trace_id is None: + return False + + attrs = dict(attributes) + state = self._trace_state.setdefault(trace_id, {}) + log_type = attrs.get(RESPAN_LOG_TYPE) + is_root = getattr(span, "parent", None) is None + + input_value = attrs.get(TRACELOOP_ENTITY_INPUT) + output_value = attrs.get(TRACELOOP_ENTITY_OUTPUT) + if log_type == "chat": + if not _has_content(input_value): + prompt_messages = _indexed_messages(attrs, PROMPT_PREFIX) + if prompt_messages: + input_value = json.dumps(prompt_messages, separators=(",", ":")) + attrs[TRACELOOP_ENTITY_INPUT] = input_value + if not _has_content(output_value): + completion_messages = _indexed_messages(attrs, COMPLETION_PREFIX) + if completion_messages: + output_payload: Any = ( + completion_messages[0] + if len(completion_messages) == 1 + else completion_messages + ) + output_value = json.dumps(output_payload, separators=(",", ":")) + attrs[TRACELOOP_ENTITY_OUTPUT] = output_value + if _has_content(input_value) and "input" not in state: + state["input"] = input_value + if _has_content(output_value): + state["output"] = output_value + + propagation = state.setdefault("propagation", {}) + for key in _PROPAGATED_ATTRS: + if key in attrs: + propagation[key] = attrs[key] + for key, value in attrs.items(): + if key.startswith(f"{RESPAN_METADATA}."): + propagation[key] = value + + if log_type == "agent" or is_root: + if not _has_content(input_value) and "input" in state: + attrs[TRACELOOP_ENTITY_INPUT] = state["input"] + if not _has_content(output_value) and "output" in state: + attrs[TRACELOOP_ENTITY_OUTPUT] = state["output"] + if is_root: + for key, value in propagation.items(): + attrs.setdefault(key, value) + + for key in _OFF_CONTRACT_ALIASES: + attrs.pop(key, None) + + span._attributes = attrs + return is_root def _set_workflow_name(self, span: Any) -> None: if not self._workflow_name: @@ -175,15 +409,20 @@ def _set_workflow_name(self, span: Any) -> None: set_attribute(TRACELOOP_WORKFLOW_NAME, self._workflow_name) def on_start(self, span, parent_context=None) -> None: + self._normalize_trace_id(span) self._set_workflow_name(span) for processor in self._processors: processor.on_start(span=span, parent_context=parent_context) def on_end(self, span) -> None: + trace_id = self._normalize_trace_id(span) self._set_workflow_name(span) self._translator.on_end(span) + is_root = self._enrich_span_contract(span, trace_id) for processor in self._processors: processor.on_end(span=span) + if is_root and trace_id is not None: + self._trace_state.pop(trace_id, None) def shutdown(self) -> None: # This chain borrows processors owned by the active tracer provider. diff --git a/python-sdks/instrumentations/respan-instrumentation-agentspec/tests/test_instrumentation.py b/python-sdks/instrumentations/respan-instrumentation-agentspec/tests/test_instrumentation.py index ff4c5a5b..5e2dcfde 100644 --- a/python-sdks/instrumentations/respan-instrumentation-agentspec/tests/test_instrumentation.py +++ b/python-sdks/instrumentations/respan-instrumentation-agentspec/tests/test_instrumentation.py @@ -1,8 +1,10 @@ +import asyncio import logging import sys from types import ModuleType, SimpleNamespace import pytest +from opentelemetry.trace import SpanContext, TraceFlags from respan_instrumentation_agentspec import AgentSpecInstrumentor from respan_instrumentation_agentspec import _instrumentation @@ -12,6 +14,15 @@ _extract_langchain_usage, _patch_agentspec_langgraph_usage, ) +from respan_sdk.constants.span_attributes import ( + RESPAN_CUSTOMER_PARAMS_ID, + RESPAN_LOG_TYPE, + RESPAN_METADATA, + RESPAN_SPAN_CUSTOM_ID, + RESPAN_SPAN_TOOLS, + RESPAN_THREADS_ID, + RESPAN_TRACE_GROUP_ID, +) from respan_tracing.core.tracer import RespanTracer @@ -205,6 +216,98 @@ def test_translated_processor_chain_does_not_shutdown_borrowed_processors(): assert export_processor.did_shutdown is False +def test_translated_processor_chain_preserves_full_trace_id_and_enriches_boundaries(): + session_id = "12345678-1234-5678-9abc-def012345678" + expected_trace_id = int(session_id.replace("-", ""), 16) + + class FakeSpan: + def __init__(self, *, span_id, parent, attributes): + self._context = SpanContext( + trace_id=1, + span_id=span_id, + is_remote=False, + trace_flags=TraceFlags(TraceFlags.SAMPLED), + ) + self._parent = parent + self._attributes = {"session.id": session_id, **attributes} + + @property + def parent(self): + return self._parent + + def get_span_context(self): + return self._context + + export_processor = FakeExportProcessor() + chain = _TranslatedProcessorChain( + translator=FakeTranslator(), + processors=(export_processor,), + ) + root = FakeSpan(span_id=1, parent=None, attributes={RESPAN_LOG_TYPE: "workflow"}) + root_parent = root.get_span_context() + agent = FakeSpan( + span_id=2, + parent=root_parent, + attributes={ + RESPAN_LOG_TYPE: "agent", + "traceloop.entity.input": "{}", + "traceloop.entity.output": "{}", + }, + ) + chat = FakeSpan( + span_id=3, + parent=agent.get_span_context(), + attributes={ + RESPAN_LOG_TYPE: "chat", + "gen_ai.prompt.0.role": "user", + "gen_ai.prompt.0.content": "hello", + "gen_ai.completion.0.role": "assistant", + "gen_ai.completion.0.content": "hi", + RESPAN_CUSTOMER_PARAMS_ID: "agentspec-user", + RESPAN_THREADS_ID: "agentspec-thread", + RESPAN_TRACE_GROUP_ID: "agentspec-group", + RESPAN_SPAN_CUSTOM_ID: "agentspec-custom", + f"{RESPAN_METADATA}.scenario": "propagated_attributes", + RESPAN_SPAN_TOOLS: "[]", + "tools": [], + }, + ) + + for span in (root, agent, chat): + chain.on_start(span) + + assert root.get_span_context().trace_id == expected_trace_id + assert agent.get_span_context().trace_id == expected_trace_id + assert agent.parent.trace_id == expected_trace_id + assert chat.get_span_context().trace_id == expected_trace_id + assert chat.parent.trace_id == expected_trace_id + assert expected_trace_id >> 64 + + chain.on_end(chat) + chain.on_end(agent) + chain.on_end(root) + + assert chat._attributes["traceloop.entity.input"] == ( + '[{"role":"user","content":"hello"}]' + ) + assert chat._attributes["traceloop.entity.output"] == ( + '{"role":"assistant","content":"hi"}' + ) + for boundary in (agent, root): + assert boundary._attributes["traceloop.entity.input"].endswith("hello\"}]") + assert boundary._attributes["traceloop.entity.output"].endswith("hi\"}") + assert root._attributes[RESPAN_CUSTOMER_PARAMS_ID] == "agentspec-user" + assert root._attributes[RESPAN_THREADS_ID] == "agentspec-thread" + assert root._attributes[RESPAN_TRACE_GROUP_ID] == "agentspec-group" + assert root._attributes[RESPAN_SPAN_CUSTOM_ID] == "agentspec-custom" + assert root._attributes[f"{RESPAN_METADATA}.scenario"] == ( + "propagated_attributes" + ) + assert RESPAN_SPAN_TOOLS not in chat._attributes + assert "tools" not in chat._attributes + assert chain._trace_state == {} + + def test_extract_langchain_usage_from_message_usage_metadata(): response = SimpleNamespace( generations=[ @@ -240,7 +343,9 @@ def test_extract_langchain_usage_from_llm_output_token_usage(): assert _extract_langchain_usage(response) == (12, 5) -def test_patch_agentspec_langgraph_usage_adds_tokens_to_response_event(monkeypatch): +def test_patch_agentspec_langgraph_usage_adds_tokens_to_sync_and_async_events( + monkeypatch, +): class FakeSpan: def __init__(self): self.events = [] @@ -262,14 +367,30 @@ def _add_event(self, run_id_str, span, event): def _end_span(self, run_id_str, span): span.did_end = True + async def _add_event_async(self, run_id_str, span, event): + self._add_event(run_id_str, span, event) + + async def _end_span_async(self, run_id_str, span): + self._end_span(run_id_str, span) + def on_llm_end(self, response, *, run_id, parent_run_id=None, **kwargs): raise AssertionError("original handler should be replaced") + async def on_llm_end_async( + self, + response, + *, + run_id, + parent_run_id=None, + **kwargs, + ): + raise AssertionError("original async handler should be replaced") + def extract_message_content_and_tool_calls(response): return "message-1", "hello", [] fake_module = ModuleType("pyagentspec.adapters.langgraph.tracing") - fake_module.AgentSpecCallbackHandler = FakeCallbackHandler + fake_module.AgentSpecLlmCallbackHandler = FakeCallbackHandler fake_module.AgentSpecLlmGenerationSpan = FakeSpan fake_module.AgentSpecLlmGenerationResponse = FakeResponseEvent fake_module._extract_message_content_and_tool_calls = ( @@ -321,6 +442,21 @@ def extract_message_content_and_tool_calls(response): assert handler.agentspec_spans_registry == {} assert handler.messages_in_process == {} + async_handler = FakeCallbackHandler() + async_span = async_handler.agentspec_spans_registry["run-1"] + asyncio.run( + async_handler.on_llm_end_async( + response, + run_id="run-1", + ) + ) + + async_event = async_span.events[0][1] + assert async_event.kwargs["input_tokens"] == 7 + assert async_event.kwargs["output_tokens"] == 3 + assert async_span.did_end is True + assert async_handler.agentspec_spans_registry == {} + def test_activate_starts_agentspec_trace_with_translated_processor_chain(monkeypatch): fake = _install_fake_modules(monkeypatch)