diff --git a/examples/tutorials/10_async/00_base/030_tracing/README.md b/examples/tutorials/10_async/00_base/030_tracing/README.md index 1d91f565f..3c66248e8 100644 --- a/examples/tutorials/10_async/00_base/030_tracing/README.md +++ b/examples/tutorials/10_async/00_base/030_tracing/README.md @@ -41,6 +41,39 @@ await adk.tracing.end_span( Spans create a hierarchical view of agent execution, making it easy to see which operations take time and where errors occur. +## Token Usage & Cost Tracking + +Token usage on spans is what the backend bills from, and it reads two shapes: + +- **Per-turn aggregate** — `span.data["usage"]` + `span.data["cost_usd"]`. Emit at most + once per turn, holding that turn's own usage (not a session-cumulative total). When a + trace has an aggregate, the backend keeps it and de-dups all per-call spans against it. +- **Per-call detail** — `span.output["usage"]`. Optional; the SDK's LLM adapters + (litellm, OpenAI Agents SDK, LangGraph) emit this automatically. Summed only when no + aggregate exists in the trace. + +Record the turn rollup with `adk.tracing.turn_span()` instead of hand-writing usage keys. +It accepts the harness `TurnUsage` that every turn adapter reports (`LangGraphTurn.usage()`, +`run_turn(...).usage`, `ClaudeCodeTurn.usage()`, ...), cost included: + +```python +async with adk.tracing.turn_span( + trace_id=task.id, + name="turn", + input={"prompt": prompt}, + task_id=task.id, +) as turn: + result = await run_turn(...) + turn.output = {"response": result.final_output} + turn.record_usage(result.usage) # TurnUsage; cost_usd stamped automatically +``` + +**Never put usage on both a rollup span's `output` and its per-call children's +`output`** — that double-counts. `turn_span` writes the aggregate to `data`, so child +spans stay safe to emit. Recognized token keys: `input_tokens`/`prompt_tokens`, +`output_tokens`/`completion_tokens`, `cached_input_tokens`/`cached_tokens`, +`reasoning_tokens`; cost is `cost_usd`. + ## When to Use - Debugging complex agent behaviors - Performance optimization and bottleneck identification diff --git a/src/agentex/lib/adk/__init__.py b/src/agentex/lib/adk/__init__.py index 4d79be9dd..25b858485 100644 --- a/src/agentex/lib/adk/__init__.py +++ b/src/agentex/lib/adk/__init__.py @@ -27,7 +27,7 @@ from agentex.lib.adk._modules.state import StateModule from agentex.lib.adk._modules.streaming import StreamingModule from agentex.lib.adk._modules.tasks import TasksModule -from agentex.lib.adk._modules.tracing import TracingModule +from agentex.lib.adk._modules.tracing import TracingModule, TurnSpan # Unified harness surface (AGX1-375) from agentex.lib.core.harness import ( @@ -66,6 +66,7 @@ "tracing", "events", "agent_task_tracker", + "TurnSpan", # Checkpointing / LangGraph "create_checkpointer", "stream_langgraph_events", diff --git a/src/agentex/lib/adk/_modules/tracing.py b/src/agentex/lib/adk/_modules/tracing.py index 94bf741e4..7d49bb91c 100644 --- a/src/agentex/lib/adk/_modules/tracing.py +++ b/src/agentex/lib/adk/_modules/tracing.py @@ -20,6 +20,7 @@ TracingActivityName, ) from agentex.lib.core.tracing.tracer import AsyncTracer +from agentex.lib.core.harness.types import TurnUsage from agentex.types.span import Span from agentex.lib.utils.logging import make_logger from agentex.lib.utils.model_utils import BaseModel @@ -30,6 +31,21 @@ DEFAULT_RETRY_POLICY = RetryPolicy(maximum_attempts=1) TEMPORAL_SPAN_ACTIVITY_DROPPED_METRIC = "agentex.tracing.temporal_span_activity.dropped" +# Token key spellings the backend accepts when billing usage from spans. +RECOGNIZED_USAGE_KEYS = frozenset( + { + "input_tokens", + "prompt_tokens", + "output_tokens", + "completion_tokens", + "cached_input_tokens", + "cached_tokens", + "reasoning_tokens", + "total_tokens", + "cost_usd", + } +) + def _record_temporal_span_activity_dropped(event_type: str) -> None: try: @@ -42,6 +58,80 @@ def _record_temporal_span_activity_dropped(event_type: str) -> None: pass +class TurnSpan: + """Handle for a turn-level (rollup) span, yielded by ``TracingModule.turn_span``. + + Encapsulates the billing contract so agents cannot double-count usage: + the turn's aggregate usage goes to ``span.data["usage"]`` (+ + ``span.data["cost_usd"]``) via :meth:`record_usage`. The backend keeps the + aggregate and de-dups any per-call ``output["usage"]`` children against it. + Never hand-write usage into ``output`` on a rollup span — that is the + double-count bug this helper exists to prevent. + + All methods no-op when tracing is disabled (``span`` is None), so agent + code needs no ``if span:`` guards. + """ + + def __init__(self, span: Span | None): + self.span = span + + def record_usage( + self, + usage: TurnUsage | dict[str, Any] | None = None, + cost_usd: float | None = None, + ) -> None: + """Record the turn's aggregate usage on the span's ``data``. + + Pass the harness ``TurnUsage`` (e.g. ``LangGraphTurn.usage()`` or + ``run_turn(...).usage``) — its ``cost_usd`` is stamped automatically — + or a plain dict with backend-recognized token spellings + (``prompt_tokens``/``completion_tokens`` also work). An explicit + ``cost_usd`` argument overrides any cost carried by ``usage``. The + usage must be this turn's own tokens, not a session-cumulative total. + """ + if self.span is None: + return + + blob: dict[str, Any] + if isinstance(usage, TurnUsage): + blob = usage.model_dump(exclude_none=True) + # cost lives beside the blob as data["cost_usd"], not inside it + blob_cost = blob.pop("cost_usd", None) + if cost_usd is None: + cost_usd = blob_cost + elif usage is not None: + blob = dict(usage) + if not any(key in RECOGNIZED_USAGE_KEYS for key in blob): + logger.warning( + "TurnSpan.record_usage: usage has no recognized token keys and will " + f"not be billed. Got keys {sorted(blob)}; expected any of " + f"{sorted(RECOGNIZED_USAGE_KEYS)}." + ) + else: + blob = {} + + if self.span.data is not None and not isinstance(self.span.data, dict): + logger.warning( + f"TurnSpan.record_usage: span.data is {type(self.span.data).__name__} " + "(expected dict or None); existing data will be replaced." + ) + data = self.span.data if isinstance(self.span.data, dict) else {} + if blob: + data["usage"] = blob + if cost_usd is not None: + data["cost_usd"] = cost_usd + self.span.data = data + + @property + def output(self) -> Any: + return self.span.output if self.span is not None else None + + @output.setter + def output(self, value: Any) -> None: + if self.span is not None: + self.span.output = value + + class TracingModule: """ Module for managing tracing and span operations in Agentex. @@ -156,6 +246,47 @@ async def span( retry_policy=retry_policy, ) + @asynccontextmanager + async def turn_span( + self, + trace_id: str, + name: str, + input: list[Any] | dict[str, Any] | BaseModel | None = None, + data: list[Any] | dict[str, Any] | BaseModel | None = None, + parent_id: str | None = None, + task_id: str | None = None, + start_to_close_timeout: timedelta = timedelta(seconds=5), + heartbeat_timeout: timedelta = timedelta(seconds=5), + retry_policy: RetryPolicy = DEFAULT_RETRY_POLICY, + ) -> AsyncGenerator[TurnSpan, None]: + """Span for one agent turn, with usage recorded as the billable aggregate. + + Same lifecycle as :meth:`span`, but yields a :class:`TurnSpan` whose + ``record_usage(usage=..., cost_usd=...)`` writes the turn's rollup + usage to ``span.data`` — the shape the backend bills once per turn. + Per-call child spans (LLM adapters) may still carry + ``output["usage"]``; the backend de-dups them against this aggregate. + + Example (with a harness turn, e.g. ``LangGraphTurn`` / ``run_turn``):: + + async with adk.tracing.turn_span(trace_id=task.id, name="turn", input={...}, task_id=task.id) as turn: + result = await run_turn(...) + turn.output = {"response": result.final_output} + turn.record_usage(result.usage) # TurnUsage, cost_usd included + """ + async with self.span( + trace_id=trace_id, + name=name, + input=input, + data=data, + parent_id=parent_id, + task_id=task_id, + start_to_close_timeout=start_to_close_timeout, + heartbeat_timeout=heartbeat_timeout, + retry_policy=retry_policy, + ) as span: + yield TurnSpan(span) + async def start_span( self, trace_id: str, diff --git a/src/agentex/lib/core/services/adk/providers/litellm.py b/src/agentex/lib/core/services/adk/providers/litellm.py index a511aa7b8..416077e8f 100644 --- a/src/agentex/lib/core/services/adk/providers/litellm.py +++ b/src/agentex/lib/core/services/adk/providers/litellm.py @@ -25,6 +25,20 @@ logger = logging.make_logger(__name__) +def _stream_kwargs_with_usage(llm_config: LLMConfig) -> dict: + """Completion kwargs with usage reporting enabled on the final stream chunk. + + litellm only reports usage for streaming calls when + ``stream_options={"include_usage": True}`` is set; default it on so usage + reaches the span. Callers can still opt out by explicitly passing + ``stream_options={"include_usage": False}``. + """ + kwargs = llm_config.model_dump() + stream_options = kwargs.get("stream_options") or {} + kwargs["stream_options"] = {"include_usage": True, **stream_options} + return kwargs + + class LiteLLMService: def __init__( self, @@ -121,7 +135,11 @@ async def chat_completion_auto_send( if span: if streaming_context.task_message: - span.output = streaming_context.task_message.model_dump() + output = streaming_context.task_message.model_dump() + # Per-call usage for billing; deduped against any turn aggregate + if completion.usage is not None: + output["usage"] = completion.usage.model_dump() + span.output = output return streaming_context.task_message if streaming_context.task_message else None async def chat_completion_stream( @@ -150,18 +168,21 @@ async def chat_completion_stream( if self.llm_gateway is None: raise ValueError("LLM Gateway is not set") + completion_kwargs = _stream_kwargs_with_usage(llm_config) trace = self.tracer.trace(trace_id) async with trace.span( parent_id=parent_span_id, name="chat_completion_stream", - input=llm_config.model_dump(), + input=completion_kwargs, ) as span: # Direct streaming outside temporal - yield each chunk as it comes chunks: list[Completion] = [] - async for chunk in self.llm_gateway.acompletion_stream(**llm_config.model_dump()): + async for chunk in self.llm_gateway.acompletion_stream(**completion_kwargs): chunks.append(chunk) yield chunk if span: + # The usage-bearing final chunk survives concat, so the dumped + # completion carries usage for billing span.output = concat_completion_chunks(chunks).model_dump() async def chat_completion_stream_auto_send( @@ -190,11 +211,12 @@ async def chat_completion_stream_auto_send( if not llm_config.stream: llm_config.stream = True + completion_kwargs = _stream_kwargs_with_usage(llm_config) trace = self.tracer.trace(trace_id) async with trace.span( parent_id=parent_span_id, name="chat_completion_stream_auto_send", - input=llm_config.model_dump(), + input=completion_kwargs, ) as span: # Use streaming context manager async with self.streaming_service.streaming_task_message_context( @@ -208,8 +230,11 @@ async def chat_completion_stream_auto_send( ) as streaming_context: # Get the streaming response chunks = [] - async for response in self.llm_gateway.acompletion_stream(**llm_config.model_dump()): + async for response in self.llm_gateway.acompletion_stream(**completion_kwargs): heartbeat_if_in_workflow("chat completion streaming") + # Store every chunk for final message assembly, including + # the usage-only final chunk, which has no choices + chunks.append(response) if response.choices and len(response.choices) > 0 and response.choices[0].delta: delta = response.choices[0].delta.content if delta: @@ -223,9 +248,6 @@ async def chat_completion_stream_auto_send( ) heartbeat_if_in_workflow("content chunk streamed") - # Store the chunk for final message assembly - chunks.append(response) - # Update the final message content complete_message = concat_completion_chunks(chunks) if complete_message and complete_message.choices and complete_message.choices[0].message: @@ -246,6 +268,10 @@ async def chat_completion_stream_auto_send( if span: if streaming_context.task_message: - span.output = streaming_context.task_message.model_dump() + output = streaming_context.task_message.model_dump() + # Per-call usage for billing; deduped against any turn aggregate + if complete_message.usage is not None: + output["usage"] = complete_message.usage.model_dump() + span.output = output return streaming_context.task_message if streaming_context.task_message else None diff --git a/src/agentex/lib/core/temporal/plugins/__init__.py b/src/agentex/lib/core/temporal/plugins/__init__.py index 52ab6eac7..4da2e9ca0 100644 --- a/src/agentex/lib/core/temporal/plugins/__init__.py +++ b/src/agentex/lib/core/temporal/plugins/__init__.py @@ -11,7 +11,6 @@ Example: >>> from agentex.lib.core.temporal.plugins.openai_agents import ( ... TemporalStreamingModelProvider, - ... TemporalTracingModelProvider, ... ContextInterceptor, ... ) >>> from temporalio.contrib.openai_agents import OpenAIAgentsPlugin, ModelActivityParameters @@ -37,7 +36,6 @@ ContextInterceptor, TemporalStreamingHooks, TemporalStreamingModel, - TemporalTracingModelProvider, TemporalStreamingModelProvider, streaming_task_id, streaming_trace_id, @@ -48,11 +46,10 @@ __all__ = [ "TemporalStreamingModel", "TemporalStreamingModelProvider", - "TemporalTracingModelProvider", "ContextInterceptor", "streaming_task_id", "streaming_trace_id", "streaming_parent_span_id", "TemporalStreamingHooks", "stream_lifecycle_content", -] \ No newline at end of file +] diff --git a/src/agentex/lib/core/temporal/plugins/openai_agents/__init__.py b/src/agentex/lib/core/temporal/plugins/openai_agents/__init__.py index 7d81b37d0..453e27074 100644 --- a/src/agentex/lib/core/temporal/plugins/openai_agents/__init__.py +++ b/src/agentex/lib/core/temporal/plugins/openai_agents/__init__.py @@ -61,9 +61,6 @@ from agentex.lib.core.temporal.plugins.openai_agents.hooks.activities import ( stream_lifecycle_content, ) -from agentex.lib.core.temporal.plugins.openai_agents.models.temporal_tracing_model import ( - TemporalTracingModelProvider, -) from agentex.lib.core.temporal.plugins.openai_agents.models.temporal_streaming_model import ( TemporalStreamingModel, TemporalStreamingModelProvider, @@ -78,7 +75,6 @@ __all__ = [ "TemporalStreamingModel", "TemporalStreamingModelProvider", - "TemporalTracingModelProvider", "ContextInterceptor", "streaming_task_id", "streaming_trace_id", diff --git a/src/agentex/lib/core/temporal/plugins/openai_agents/models/__init__.py b/src/agentex/lib/core/temporal/plugins/openai_agents/models/__init__.py index bb5dc97ed..19dd967d5 100644 --- a/src/agentex/lib/core/temporal/plugins/openai_agents/models/__init__.py +++ b/src/agentex/lib/core/temporal/plugins/openai_agents/models/__init__.py @@ -4,11 +4,6 @@ capabilities to standard OpenAI models when running in Temporal workflows/activities. """ -from agentex.lib.core.temporal.plugins.openai_agents.models.temporal_tracing_model import ( - TemporalTracingModelProvider, - TemporalTracingResponsesModel, - TemporalTracingChatCompletionsModel, -) from agentex.lib.core.temporal.plugins.openai_agents.models.temporal_streaming_model import ( TemporalStreamingModel, TemporalStreamingModelProvider, @@ -17,7 +12,4 @@ __all__ = [ "TemporalStreamingModel", "TemporalStreamingModelProvider", - "TemporalTracingModelProvider", - "TemporalTracingResponsesModel", - "TemporalTracingChatCompletionsModel", -] \ No newline at end of file +] diff --git a/src/agentex/lib/core/temporal/plugins/openai_agents/models/temporal_tracing_model.py b/src/agentex/lib/core/temporal/plugins/openai_agents/models/temporal_tracing_model.py deleted file mode 100644 index e3d1b3804..000000000 --- a/src/agentex/lib/core/temporal/plugins/openai_agents/models/temporal_tracing_model.py +++ /dev/null @@ -1,402 +0,0 @@ -"""Temporal-aware tracing model provider. - -This module provides model implementations that add AgentEx tracing to standard OpenAI models -when running in Temporal workflows/activities. It uses context variables set by the Temporal -context interceptor to access task_id, trace_id, and parent_span_id. - -The key innovation is that these are thin wrappers around the standard OpenAI models, -avoiding code duplication while adding tracing capabilities. -""" -from __future__ import annotations - -import logging -from typing import Any, List, Union, Optional, override - -from agents import ( - Tool, - Model, - Handoff, - ModelTracing, - ModelResponse, - ModelSettings, - OpenAIProvider, - TResponseInputItem, - AgentOutputSchemaBase, -) -from openai import AsyncOpenAI -from openai.types.responses import ResponsePromptParam -from agents.models.openai_responses import OpenAIResponsesModel -from agents.models.openai_chatcompletions import OpenAIChatCompletionsModel - -from agentex.lib.core.tracing.tracer import AsyncTracer - -# Import AgentEx components -from agentex.lib.adk.utils._modules.client import create_async_agentex_client - -# Import context variables from the interceptor -from agentex.lib.core.temporal.plugins.openai_agents.interceptors.context_interceptor import ( - streaming_task_id, - streaming_trace_id, - streaming_parent_span_id, -) - -logger = logging.getLogger("agentex.temporal.tracing") - - -def _serialize_item(item: Any) -> dict[str, Any]: - """ - Universal serializer for any item type from OpenAI Agents SDK. - - Uses model_dump() for Pydantic models, otherwise extracts attributes manually. - Filters out internal Pydantic fields that can't be serialized. - """ - if hasattr(item, 'model_dump'): - # Pydantic model - use model_dump for proper serialization - try: - return item.model_dump(mode='json', exclude_unset=True) - except Exception: - # Fallback to dict conversion - return dict(item) if hasattr(item, '__iter__') else {} - else: - # Not a Pydantic model - extract attributes manually - item_dict = {} - for attr_name in dir(item): - if not attr_name.startswith('_') and attr_name not in ('model_fields', 'model_config', 'model_computed_fields'): - try: - attr_value = getattr(item, attr_name, None) - # Skip methods and None values - if attr_value is not None and not callable(attr_value): - # Convert to JSON-serializable format - if hasattr(attr_value, 'model_dump'): - item_dict[attr_name] = attr_value.model_dump() - elif isinstance(attr_value, (str, int, float, bool, list, dict)): - item_dict[attr_name] = attr_value - else: - item_dict[attr_name] = str(attr_value) - except Exception: - # Skip attributes that can't be accessed - pass - return item_dict - - -class TemporalTracingModelProvider(OpenAIProvider): - """Model provider that returns OpenAI models wrapped with AgentEx tracing. - - This provider extends the standard OpenAIProvider to return models that add - tracing spans around model calls when running in Temporal activities with - the context interceptor enabled. - """ - - def __init__(self, openai_client: Optional[AsyncOpenAI] = None, **kwargs): - """Initialize the tracing model provider. - - Args: - openai_client: Optional custom AsyncOpenAI client. If provided, this client - will be used for all model calls. If not provided, OpenAIProvider - will create a default client. - **kwargs: All other arguments are passed to OpenAIProvider. - """ - # Pass openai_client to parent if provided - if openai_client is not None: - super().__init__(openai_client=openai_client, **kwargs) - else: - super().__init__(**kwargs) - - # Initialize tracer for all models - agentex_client = create_async_agentex_client() - self._tracer = AsyncTracer(agentex_client) - logger.info(f"[TemporalTracingModelProvider] Initialized with AgentEx tracer, custom_client={openai_client is not None}") - - @override - def get_model(self, model_name: Optional[str]) -> Model: - """Get a model wrapped with tracing capabilities. - - Args: - model_name: The name of the model to use - - Returns: - A model instance wrapped with tracing - """ - # Get the base model from the parent provider - base_model = super().get_model(model_name) - - # Wrap with appropriate tracing wrapper based on model type - if isinstance(base_model, OpenAIResponsesModel): - logger.info(f"[TemporalTracingModelProvider] Wrapping OpenAIResponsesModel '{model_name}' with tracing") - return TemporalTracingResponsesModel(base_model, self._tracer) # type: ignore[abstract] - elif isinstance(base_model, OpenAIChatCompletionsModel): - logger.info(f"[TemporalTracingModelProvider] Wrapping OpenAIChatCompletionsModel '{model_name}' with tracing") - return TemporalTracingChatCompletionsModel(base_model, self._tracer) # type: ignore[abstract] - else: - logger.warning(f"[TemporalTracingModelProvider] Unknown model type, returning without tracing: {type(base_model)}") - return base_model - - -class TemporalTracingResponsesModel(Model): - """Wrapper for OpenAIResponsesModel that adds AgentEx tracing. - - This is a thin wrapper that adds tracing spans around the base model's - get_response() method. It reads tracing context from ContextVars set by - the Temporal context interceptor. - """ - - def __init__(self, base_model: OpenAIResponsesModel, tracer: AsyncTracer): - """Initialize the tracing wrapper. - - Args: - base_model: The OpenAI Responses model to wrap - tracer: The AgentEx tracer to use - """ - self._base_model = base_model - self._tracer = tracer - # Expose the model name for compatibility - self.model = base_model.model - - @override - async def get_response( - self, - system_instructions: Optional[str], - input: Union[str, List[TResponseInputItem]], - model_settings: ModelSettings, - tools: List[Tool], - output_schema: Optional[AgentOutputSchemaBase], - handoffs: List[Handoff], - tracing: ModelTracing, - previous_response_id: Optional[str] = None, - conversation_id: Optional[str] = None, - prompt: Optional[ResponsePromptParam] = None, - **kwargs, - ) -> ModelResponse: - """Get a response from the model with optional tracing. - - If tracing context is available from the interceptor, this wraps the - model call in a tracing span. Otherwise, it passes through to the - base model without tracing. - """ - # Try to get tracing context from ContextVars - task_id = streaming_task_id.get() - trace_id = streaming_trace_id.get() - parent_span_id = streaming_parent_span_id.get() - - # If we have tracing context, wrap with span - if trace_id and parent_span_id: - logger.debug(f"[TemporalTracingResponsesModel] Adding tracing span for task_id={task_id}, trace_id={trace_id}") - - trace = self._tracer.trace(trace_id) - - async with trace.span( - parent_id=parent_span_id, - name="model_get_response", - input={ - "model": str(self.model), - "has_system_instructions": system_instructions is not None, - "input_type": type(input).__name__, - "tools_count": len(tools) if tools else 0, - "handoffs_count": len(handoffs) if handoffs else 0, - "has_output_schema": output_schema is not None, - "model_settings": { - "temperature": model_settings.temperature, - "max_tokens": model_settings.max_tokens, - "reasoning": model_settings.reasoning, - } if model_settings else None, - }, - ) as span: - try: - # Call the base model - response = await self._base_model.get_response( - system_instructions=system_instructions, - input=input, - model_settings=model_settings, - tools=tools, - output_schema=output_schema, - handoffs=handoffs, - tracing=tracing, - previous_response_id=previous_response_id, - conversation_id=conversation_id, # type: ignore[call-arg] - prompt=prompt, - **kwargs, - ) - - # Serialize response output items for span tracing - new_items = [] - final_output = None - - if hasattr(response, 'output') and response.output: - response_output = response.output if isinstance(response.output, list) else [response.output] - - for item in response_output: - try: - item_dict = _serialize_item(item) - if item_dict: - new_items.append(item_dict) - - # Extract final_output from message type if available - if item_dict.get('type') == 'message' and not final_output: - content = item_dict.get('content', []) - if content and isinstance(content, list): - for content_part in content: - if isinstance(content_part, dict) and 'text' in content_part: - final_output = content_part['text'] - break - except Exception as e: - logger.warning(f"Failed to serialize item in temporal tracing model: {e}") - continue - - # Set span output with structured data - span.output = { # type: ignore[attr-defined] - "new_items": new_items, - "final_output": final_output, - } - - return response - - except Exception as e: - # Record error in span - span.error = str(e) # type: ignore[attr-defined] - raise - else: - # No tracing context, just pass through - logger.debug("[TemporalTracingResponsesModel] No tracing context available, calling base model directly") - return await self._base_model.get_response( - system_instructions=system_instructions, - input=input, - model_settings=model_settings, - tools=tools, - output_schema=output_schema, - handoffs=handoffs, - tracing=tracing, - previous_response_id=previous_response_id, - conversation_id=conversation_id, # type: ignore[call-arg] - prompt=prompt, - **kwargs, - ) - - -class TemporalTracingChatCompletionsModel(Model): - """Wrapper for OpenAIChatCompletionsModel that adds AgentEx tracing. - - This is a thin wrapper that adds tracing spans around the base model's - get_response() method. It reads tracing context from ContextVars set by - the Temporal context interceptor. - """ - - def __init__(self, base_model: OpenAIChatCompletionsModel, tracer: AsyncTracer): - """Initialize the tracing wrapper. - - Args: - base_model: The OpenAI ChatCompletions model to wrap - tracer: The AgentEx tracer to use - """ - self._base_model = base_model - self._tracer = tracer - # Expose the model name for compatibility - self.model = base_model.model - - @override - async def get_response( - self, - system_instructions: Optional[str], - input: Union[str, List[TResponseInputItem]], - model_settings: ModelSettings, - tools: List[Tool], - output_schema: Optional[AgentOutputSchemaBase], - handoffs: List[Handoff], - tracing: ModelTracing, - **kwargs, - ) -> ModelResponse: - """Get a response from the model with optional tracing. - - If tracing context is available from the interceptor, this wraps the - model call in a tracing span. Otherwise, it passes through to the - base model without tracing. - """ - # Try to get tracing context from ContextVars - task_id = streaming_task_id.get() - trace_id = streaming_trace_id.get() - parent_span_id = streaming_parent_span_id.get() - - # If we have tracing context, wrap with span - if trace_id and parent_span_id: - logger.debug(f"[TemporalTracingChatCompletionsModel] Adding tracing span for task_id={task_id}, trace_id={trace_id}") - - trace = self._tracer.trace(trace_id) - - async with trace.span( - parent_id=parent_span_id, - name="model_get_response", - input={ - "model": str(self.model), - "has_system_instructions": system_instructions is not None, - "input_type": type(input).__name__, - "tools_count": len(tools) if tools else 0, - "handoffs_count": len(handoffs) if handoffs else 0, - "has_output_schema": output_schema is not None, - "model_settings": { - "temperature": model_settings.temperature, - "max_tokens": model_settings.max_tokens, - } if model_settings else None, - }, - ) as span: - try: - # Call the base model - response = await self._base_model.get_response( - system_instructions=system_instructions, - input=input, - model_settings=model_settings, - tools=tools, - output_schema=output_schema, - handoffs=handoffs, - tracing=tracing, - **kwargs, - ) - - # Serialize response output items for span tracing - new_items = [] - final_output = None - - if hasattr(response, 'output') and response.output: - response_output = response.output if isinstance(response.output, list) else [response.output] - - for item in response_output: - try: - item_dict = _serialize_item(item) - if item_dict: - new_items.append(item_dict) - - # Extract final_output from message type if available - if item_dict.get('type') == 'message' and not final_output: - content = item_dict.get('content', []) - if content and isinstance(content, list): - for content_part in content: - if isinstance(content_part, dict) and 'text' in content_part: - final_output = content_part['text'] - break - except Exception as e: - logger.warning(f"Failed to serialize item in temporal tracing model: {e}") - continue - - # Set span output with structured data - span.output = { # type: ignore[attr-defined] - "new_items": new_items, - "final_output": final_output, - } - - return response - - except Exception as e: - # Record error in span - span.error = str(e) # type: ignore[attr-defined] - raise - else: - # No tracing context, just pass through - logger.debug("[TemporalTracingChatCompletionsModel] No tracing context available, calling base model directly") - return await self._base_model.get_response( - system_instructions=system_instructions, - input=input, - model_settings=model_settings, - tools=tools, - output_schema=output_schema, - handoffs=handoffs, - tracing=tracing, - **kwargs, - ) \ No newline at end of file diff --git a/src/agentex/lib/utils/completions.py b/src/agentex/lib/utils/completions.py index fe62c7d12..3cb1d7b4d 100644 --- a/src/agentex/lib/utils/completions.py +++ b/src/agentex/lib/utils/completions.py @@ -3,6 +3,7 @@ from copy import deepcopy from typing import Any from functools import reduce, singledispatch +from itertools import zip_longest from agentex.lib.types.llm_messages import ( Delta, @@ -21,7 +22,12 @@ def _concat_chunks(_a: None, b: Any): @_concat_chunks.register def _(a: Completion, b: Completion) -> Completion: - a.choices = [_concat_chunks(*c) for c in zip(a.choices, b.choices, strict=False)] + # Chunks can have unequal choices: with stream_options.include_usage the + # final chunk carries usage but no choices. Keep the unpaired side instead + # of truncating to the shorter list. + a.choices = [ + x if y is None else y if x is None else _concat_chunks(x, y) for x, y in zip_longest(a.choices, b.choices) + ] a.usage = _concat_chunks(a.usage, b.usage) return a @@ -38,6 +44,7 @@ def _(a: Choice, b: Choice) -> Choice: a.finish_reason = a.finish_reason or b.finish_reason return a + @_concat_chunks.register def _(a: Usage | None, b: Usage | None) -> Usage | None: if a is not None and b is not None: @@ -61,9 +68,7 @@ def _(a: Delta, b: Delta) -> Delta: if tool_call.index not in grouped_tool_calls: grouped_tool_calls[tool_call.index] = tool_call else: - grouped_tool_calls[tool_call.index] = _concat_chunks( - grouped_tool_calls[tool_call.index], tool_call - ) + grouped_tool_calls[tool_call.index] = _concat_chunks(grouped_tool_calls[tool_call.index], tool_call) a.tool_calls = list(grouped_tool_calls.values()) elif hasattr(b, "tool_calls") and b.tool_calls: @@ -81,11 +86,7 @@ def _(a: ToolCallRequest, b: ToolCallRequest) -> ToolCallRequest: index_val = a.index if hasattr(a, "index") and a.index is not None else b.index # Concatenate the function part - function_val = ( - _concat_chunks(a.function, b.function) - if a.function and b.function - else a.function or b.function - ) + function_val = _concat_chunks(a.function, b.function) if a.function and b.function else a.function or b.function # Set all properties a.id = id_val diff --git a/tests/lib/adk/providers/test_litellm_usage.py b/tests/lib/adk/providers/test_litellm_usage.py new file mode 100644 index 000000000..5f5d480d9 --- /dev/null +++ b/tests/lib/adk/providers/test_litellm_usage.py @@ -0,0 +1,219 @@ +"""Tests that LiteLLMService puts LLM token usage on spans for billing. + +Covers the paths that previously dropped usage: both auto_send variants (span +output was only the TaskMessage dump) and streaming (litellm omits usage unless +``stream_options.include_usage`` is set). +""" + +from __future__ import annotations + +from typing import Any +from datetime import UTC, datetime +from contextlib import asynccontextmanager +from unittest.mock import AsyncMock, MagicMock + +from agentex.types.span import Span +from agentex.types.task_message import TaskMessage +from agentex.lib.types.llm_messages import ( + Delta, + Usage, + Choice, + LLMConfig, + Completion, + AssistantMessage, +) +from agentex.types.task_message_content import TextContent +from agentex.lib.core.services.adk.providers.litellm import ( + LiteLLMService, + _stream_kwargs_with_usage, +) + + +class FakeTrace: + def __init__(self) -> None: + self.spans: list[Span] = [] + + @asynccontextmanager + async def span(self, name, parent_id=None, input=None, data=None, task_id=None): + span = Span( + id=f"span-{len(self.spans)}", + name=name, + start_time=datetime.now(UTC), + trace_id="trace-1", + parent_id=parent_id, + input=input, + ) + self.spans.append(span) + yield span + + +class FakeTracer: + def __init__(self) -> None: + self.trace_obj = FakeTrace() + + def trace(self, trace_id): + return self.trace_obj + + +def _output_dict(span: Span) -> dict[str, Any]: + assert isinstance(span.output, dict) + return span.output + + +def _make_streaming_service(): + streaming_context = MagicMock() + streaming_context.task_message = TaskMessage( + id="msg-1", + task_id="task-1", + content=TextContent(author="agent", content="", format="markdown"), + ) + streaming_context.stream_update = AsyncMock() + + @asynccontextmanager + async def fake_context(**kwargs): + yield streaming_context + + streaming_service = MagicMock() + streaming_service.streaming_task_message_context = fake_context + return streaming_service, streaming_context + + +def _make_service(llm_gateway) -> tuple[LiteLLMService, FakeTracer]: + streaming_service, _ = _make_streaming_service() + tracer = FakeTracer() + service = LiteLLMService( + agentex_client=MagicMock(), + streaming_service=streaming_service, + tracer=tracer, + llm_gateway=llm_gateway, + ) + return service, tracer + + +def _stream_gateway(chunks, captured_kwargs): + gateway = MagicMock() + + def acompletion_stream(**kwargs): + captured_kwargs.update(kwargs) + + async def stream(): + for chunk in chunks: + yield chunk + + return stream() + + gateway.acompletion_stream = acompletion_stream + return gateway + + +def _delta_chunk(content: str, role: str | None = None) -> Completion: + return Completion(choices=[Choice(index=0, delta=Delta(content=content, role=role))]) + + +def _usage_only_chunk() -> Completion: + return Completion( + choices=[], + usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15), + ) + + +class TestStreamKwargsWithUsage: + def test_defaults_include_usage_on(self): + config = LLMConfig(model="gpt-4o", messages=[], stream=True) + assert _stream_kwargs_with_usage(config)["stream_options"] == {"include_usage": True} + + def test_caller_opt_out_preserved(self): + config = LLMConfig(model="gpt-4o", messages=[], stream=True, stream_options={"include_usage": False}) + assert _stream_kwargs_with_usage(config)["stream_options"] == {"include_usage": False} + + def test_merges_with_other_stream_options(self): + config = LLMConfig(model="gpt-4o", messages=[], stream=True, stream_options={"other": 1}) + assert _stream_kwargs_with_usage(config)["stream_options"] == {"include_usage": True, "other": 1} + + +class TestChatCompletionAutoSend: + async def test_span_output_carries_usage(self): + completion = Completion( + choices=[Choice(index=0, message=AssistantMessage(content="Hello!"), finish_reason="stop")], + usage=Usage(prompt_tokens=7, completion_tokens=3, total_tokens=10), + ) + gateway = MagicMock() + gateway.acompletion = AsyncMock(return_value=completion) + service, tracer = _make_service(gateway) + + await service.chat_completion_auto_send( + task_id="task-1", + llm_config=LLMConfig(model="gpt-4o", messages=[], stream=False), + trace_id="trace-1", + ) + + span = tracer.trace_obj.spans[0] + assert _output_dict(span)["usage"] == {"prompt_tokens": 7, "completion_tokens": 3, "total_tokens": 10} + + async def test_span_output_omits_usage_when_absent(self): + completion = Completion( + choices=[Choice(index=0, message=AssistantMessage(content="Hello!"), finish_reason="stop")], + ) + gateway = MagicMock() + gateway.acompletion = AsyncMock(return_value=completion) + service, tracer = _make_service(gateway) + + await service.chat_completion_auto_send( + task_id="task-1", + llm_config=LLMConfig(model="gpt-4o", messages=[], stream=False), + trace_id="trace-1", + ) + + assert "usage" not in _output_dict(tracer.trace_obj.spans[0]) + + +class TestChatCompletionStream: + async def test_stream_requests_usage_and_span_output_carries_it(self): + captured_kwargs: dict = {} + chunks = [_delta_chunk("Hel", role="assistant"), _delta_chunk("lo!"), _usage_only_chunk()] + service, tracer = _make_service(_stream_gateway(chunks, captured_kwargs)) + + results = [] + async for chunk in service.chat_completion_stream( + llm_config=LLMConfig(model="gpt-4o", messages=[], stream=True), + trace_id="trace-1", + ): + results.append(chunk) + + assert len(results) == 3 + assert captured_kwargs["stream_options"] == {"include_usage": True} + span = tracer.trace_obj.spans[0] + assert _output_dict(span)["usage"] == {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15} + assert _output_dict(span)["choices"][0]["message"]["content"] == "Hello!" + + +class TestChatCompletionStreamAutoSend: + async def test_usage_only_final_chunk_reaches_span_output(self): + captured_kwargs: dict = {} + chunks = [_delta_chunk("Hel", role="assistant"), _delta_chunk("lo!"), _usage_only_chunk()] + service, tracer = _make_service(_stream_gateway(chunks, captured_kwargs)) + + await service.chat_completion_stream_auto_send( + task_id="task-1", + llm_config=LLMConfig(model="gpt-4o", messages=[], stream=True), + trace_id="trace-1", + ) + + assert captured_kwargs["stream_options"] == {"include_usage": True} + span = tracer.trace_obj.spans[0] + assert _output_dict(span)["usage"] == {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15} + # TaskMessage dump is still the base of the span output + assert _output_dict(span)["id"] == "msg-1" + + async def test_stream_without_usage_chunk_omits_usage(self): + captured_kwargs: dict = {} + chunks = [_delta_chunk("Hi", role="assistant")] + service, tracer = _make_service(_stream_gateway(chunks, captured_kwargs)) + + await service.chat_completion_stream_auto_send( + task_id="task-1", + llm_config=LLMConfig(model="gpt-4o", messages=[], stream=True), + trace_id="trace-1", + ) + + assert "usage" not in _output_dict(tracer.trace_obj.spans[0]) diff --git a/tests/lib/adk/test_tracing_module.py b/tests/lib/adk/test_tracing_module.py index 58d5d4a85..00e4aae65 100644 --- a/tests/lib/adk/test_tracing_module.py +++ b/tests/lib/adk/test_tracing_module.py @@ -8,7 +8,8 @@ import agentex.lib.adk._modules.tracing as _tracing_mod from agentex.types.span import Span -from agentex.lib.adk._modules.tracing import TracingModule +from agentex.lib.core.harness.types import TurnUsage +from agentex.lib.adk._modules.tracing import TurnSpan, TracingModule from agentex.lib.core.services.adk.tracing import TracingService @@ -113,10 +114,11 @@ async def test_start_span_in_workflow_returns_none_when_activity_fails(self): mock_service, module = _make_module() mock_meter = _make_metric_meter() - with patch.object(_tracing_mod, "in_temporal_workflow", return_value=True), \ - patch.object(_tracing_mod, "ActivityHelpers") as mock_helpers, \ - patch.object(_tracing_mod.workflow, "logger") as mock_logger, \ - patch.object(_tracing_mod.workflow, "metric_meter", return_value=mock_meter): + with patch.object(_tracing_mod, "in_temporal_workflow", return_value=True), patch.object( + _tracing_mod, "ActivityHelpers" + ) as mock_helpers, patch.object(_tracing_mod.workflow, "logger") as mock_logger, patch.object( + _tracing_mod.workflow, "metric_meter", return_value=mock_meter + ): mock_helpers.execute_activity = AsyncMock(side_effect=_make_activity_error()) result = await module.start_span(trace_id="trace-123", name="test-span") @@ -127,9 +129,7 @@ async def test_start_span_in_workflow_returns_none_when_activity_fails(self): description="Temporal tracing span activities dropped after fail-open", unit="1", ) - mock_meter.create_counter.return_value.add.assert_called_once_with( - 1, {"event_type": "start"} - ) + mock_meter.create_counter.return_value.add.assert_called_once_with(1, {"event_type": "start"}) mock_helpers.execute_activity.assert_called_once() mock_service.start_span.assert_not_called() @@ -138,10 +138,11 @@ async def test_end_span_in_workflow_returns_span_when_activity_fails(self): span = _make_span() mock_meter = _make_metric_meter() - with patch.object(_tracing_mod, "in_temporal_workflow", return_value=True), \ - patch.object(_tracing_mod, "ActivityHelpers") as mock_helpers, \ - patch.object(_tracing_mod.workflow, "logger") as mock_logger, \ - patch.object(_tracing_mod.workflow, "metric_meter", return_value=mock_meter): + with patch.object(_tracing_mod, "in_temporal_workflow", return_value=True), patch.object( + _tracing_mod, "ActivityHelpers" + ) as mock_helpers, patch.object(_tracing_mod.workflow, "logger") as mock_logger, patch.object( + _tracing_mod.workflow, "metric_meter", return_value=mock_meter + ): mock_helpers.execute_activity = AsyncMock(side_effect=_make_activity_error()) result = await module.end_span(trace_id="trace-123", span=span) @@ -152,18 +153,16 @@ async def test_end_span_in_workflow_returns_span_when_activity_fails(self): description="Temporal tracing span activities dropped after fail-open", unit="1", ) - mock_meter.create_counter.return_value.add.assert_called_once_with( - 1, {"event_type": "end"} - ) + mock_meter.create_counter.return_value.add.assert_called_once_with(1, {"event_type": "end"}) mock_helpers.execute_activity.assert_called_once() mock_service.end_span.assert_not_called() async def test_context_manager_skips_end_when_temporal_start_fails(self): mock_service, module = _make_module() - with patch.object(_tracing_mod, "in_temporal_workflow", return_value=True), \ - patch.object(_tracing_mod, "ActivityHelpers") as mock_helpers, \ - patch.object(_tracing_mod.workflow, "logger"): + with patch.object(_tracing_mod, "in_temporal_workflow", return_value=True), patch.object( + _tracing_mod, "ActivityHelpers" + ) as mock_helpers, patch.object(_tracing_mod.workflow, "logger"): mock_helpers.execute_activity = AsyncMock(side_effect=_make_activity_error()) async with module.span(trace_id="trace-123", name="test-span") as span: assert span is None @@ -175,8 +174,9 @@ async def test_context_manager_skips_end_when_temporal_start_fails(self): async def test_start_span_in_workflow_propagates_unexpected_errors(self): mock_service, module = _make_module() - with patch.object(_tracing_mod, "in_temporal_workflow", return_value=True), \ - patch.object(_tracing_mod, "ActivityHelpers") as mock_helpers: + with patch.object(_tracing_mod, "in_temporal_workflow", return_value=True), patch.object( + _tracing_mod, "ActivityHelpers" + ) as mock_helpers: mock_helpers.execute_activity = AsyncMock(side_effect=RuntimeError("bad response shape")) try: await module.start_span(trace_id="trace-123", name="test-span") @@ -193,11 +193,11 @@ async def test_start_span_in_workflow_propagates_cancellation(self): activity_error = _make_activity_error() mock_meter = _make_metric_meter() - with patch.object(_tracing_mod, "in_temporal_workflow", return_value=True), \ - patch.object(_tracing_mod, "ActivityHelpers") as mock_helpers, \ - patch.object(_tracing_mod, "is_cancelled_exception", return_value=True), \ - patch.object(_tracing_mod.workflow, "logger") as mock_logger, \ - patch.object(_tracing_mod.workflow, "metric_meter", return_value=mock_meter): + with patch.object(_tracing_mod, "in_temporal_workflow", return_value=True), patch.object( + _tracing_mod, "ActivityHelpers" + ) as mock_helpers, patch.object(_tracing_mod, "is_cancelled_exception", return_value=True), patch.object( + _tracing_mod.workflow, "logger" + ) as mock_logger, patch.object(_tracing_mod.workflow, "metric_meter", return_value=mock_meter): mock_helpers.execute_activity = AsyncMock(side_effect=activity_error) with pytest.raises(ActivityError): @@ -214,11 +214,11 @@ async def test_end_span_in_workflow_propagates_cancellation(self): activity_error = _make_activity_error() mock_meter = _make_metric_meter() - with patch.object(_tracing_mod, "in_temporal_workflow", return_value=True), \ - patch.object(_tracing_mod, "ActivityHelpers") as mock_helpers, \ - patch.object(_tracing_mod, "is_cancelled_exception", return_value=True), \ - patch.object(_tracing_mod.workflow, "logger") as mock_logger, \ - patch.object(_tracing_mod.workflow, "metric_meter", return_value=mock_meter): + with patch.object(_tracing_mod, "in_temporal_workflow", return_value=True), patch.object( + _tracing_mod, "ActivityHelpers" + ) as mock_helpers, patch.object(_tracing_mod, "is_cancelled_exception", return_value=True), patch.object( + _tracing_mod.workflow, "logger" + ) as mock_logger, patch.object(_tracing_mod.workflow, "metric_meter", return_value=mock_meter): mock_helpers.execute_activity = AsyncMock(side_effect=activity_error) with pytest.raises(ActivityError): @@ -258,3 +258,161 @@ async def test_span_context_manager_noop_when_no_trace_id(self): mock_service.start_span.assert_not_called() mock_service.end_span.assert_not_called() + + +class TestTurnSpan: + async def test_turn_span_records_aggregate_usage_in_data(self): + mock_service, module = _make_module() + started = _make_span(task_id="task-abc") + mock_service.start_span.return_value = started + mock_service.end_span.return_value = started + + with patch.object(_tracing_mod, "in_temporal_workflow", return_value=False): + async with module.turn_span( + trace_id="trace-123", + name="turn", + task_id="task-abc", + ) as turn: + assert isinstance(turn, TurnSpan) + turn.output = {"response": "hello"} + turn.record_usage( + usage={"input_tokens": 100, "output_tokens": 40, "total_tokens": 140}, + cost_usd=0.0125, + ) + + ended_span = mock_service.end_span.call_args.kwargs["span"] + assert ended_span.data["usage"] == { + "input_tokens": 100, + "output_tokens": 40, + "total_tokens": 140, + } + assert ended_span.data["cost_usd"] == 0.0125 + # The aggregate lives in data, never in output — output stays payload-only + assert ended_span.output == {"response": "hello"} + + async def test_turn_span_record_usage_with_turn_usage(self): + mock_service, module = _make_module() + started = _make_span() + mock_service.start_span.return_value = started + mock_service.end_span.return_value = started + + turn_usage = TurnUsage( + model="gpt-4o", + input_tokens=10, + output_tokens=5, + cached_input_tokens=2, + total_tokens=15, + cost_usd=0.5, + ) + with patch.object(_tracing_mod, "in_temporal_workflow", return_value=False): + async with module.turn_span(trace_id="trace-123", name="turn") as turn: + turn.record_usage(turn_usage) + + ended_span = mock_service.end_span.call_args.kwargs["span"] + # cost_usd is lifted out of the blob to data["cost_usd"] + assert ended_span.data["cost_usd"] == 0.5 + assert ended_span.data["usage"] == { + "model": "gpt-4o", + "input_tokens": 10, + "output_tokens": 5, + "cached_input_tokens": 2, + "total_tokens": 15, + "num_tool_calls": 0, + "num_reasoning_blocks": 0, + } + + async def test_turn_span_explicit_cost_overrides_turn_usage_cost(self): + mock_service, module = _make_module() + started = _make_span() + mock_service.start_span.return_value = started + mock_service.end_span.return_value = started + + with patch.object(_tracing_mod, "in_temporal_workflow", return_value=False): + async with module.turn_span(trace_id="trace-123", name="turn") as turn: + turn.record_usage(TurnUsage(input_tokens=1, cost_usd=0.5), cost_usd=0.75) + + ended_span = mock_service.end_span.call_args.kwargs["span"] + assert ended_span.data["cost_usd"] == 0.75 + + async def test_turn_span_warns_on_unrecognized_usage_keys(self, caplog): + mock_service, module = _make_module() + started = _make_span() + mock_service.start_span.return_value = started + mock_service.end_span.return_value = started + + with patch.object(_tracing_mod, "in_temporal_workflow", return_value=False): + async with module.turn_span(trace_id="trace-123", name="turn") as turn: + with caplog.at_level("WARNING"): + turn.record_usage(usage={"inputTokens": 10}) + + assert any("no recognized token keys" in message for message in caplog.messages) + + async def test_turn_span_preserves_existing_data(self): + mock_service, module = _make_module() + started = _make_span(data={"custom": "value"}) + mock_service.start_span.return_value = started + mock_service.end_span.return_value = started + + with patch.object(_tracing_mod, "in_temporal_workflow", return_value=False): + async with module.turn_span(trace_id="trace-123", name="turn", data={"custom": "value"}) as turn: + turn.record_usage(usage={"prompt_tokens": 3, "completion_tokens": 4}) + + ended_span = mock_service.end_span.call_args.kwargs["span"] + assert ended_span.data["custom"] == "value" + assert ended_span.data["usage"] == {"prompt_tokens": 3, "completion_tokens": 4} + + async def test_turn_span_warns_and_replaces_non_dict_data(self, caplog): + mock_service, module = _make_module() + started = _make_span(data=[{"item": 1}]) + mock_service.start_span.return_value = started + mock_service.end_span.return_value = started + + with patch.object(_tracing_mod, "in_temporal_workflow", return_value=False): + async with module.turn_span(trace_id="trace-123", name="turn") as turn: + with caplog.at_level("WARNING"): + turn.record_usage(usage={"input_tokens": 1, "output_tokens": 2}) + + assert any("existing data will be replaced" in message for message in caplog.messages) + ended_span = mock_service.end_span.call_args.kwargs["span"] + assert ended_span.data == {"usage": {"input_tokens": 1, "output_tokens": 2}} + + async def test_turn_span_dict_data_does_not_warn(self, caplog): + mock_service, module = _make_module() + started = _make_span(data={"custom": "value"}) + mock_service.start_span.return_value = started + mock_service.end_span.return_value = started + + with patch.object(_tracing_mod, "in_temporal_workflow", return_value=False): + async with module.turn_span(trace_id="trace-123", name="turn") as turn: + with caplog.at_level("WARNING"): + turn.record_usage(usage={"input_tokens": 1}) + + assert not any("existing data will be replaced" in message for message in caplog.messages) + + async def test_turn_span_cost_only(self): + mock_service, module = _make_module() + started = _make_span() + mock_service.start_span.return_value = started + mock_service.end_span.return_value = started + + with patch.object(_tracing_mod, "in_temporal_workflow", return_value=False): + async with module.turn_span(trace_id="trace-123", name="turn") as turn: + turn.record_usage(cost_usd=0.5) + + ended_span = mock_service.end_span.call_args.kwargs["span"] + assert ended_span.data == {"cost_usd": 0.5} + assert "usage" not in ended_span.data + + async def test_turn_span_noop_when_no_trace_id(self): + mock_service, module = _make_module() + + with patch.object(_tracing_mod, "in_temporal_workflow", return_value=False): + async with module.turn_span(trace_id="", name="turn") as turn: + assert turn.span is None + # Must not raise when tracing is disabled + turn.record_usage(usage={"input_tokens": 1}, cost_usd=0.1) + turn.output = {"response": "x"} + assert turn.output is None + + mock_service.start_span.assert_not_called() + mock_service.end_span.assert_not_called() diff --git a/tests/lib/core/temporal/plugins/__init__.py b/tests/lib/core/temporal/plugins/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/lib/core/temporal/plugins/openai_agents/__init__.py b/tests/lib/core/temporal/plugins/openai_agents/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/lib/core/temporal/plugins/openai_agents/test_model_usage.py b/tests/lib/core/temporal/plugins/openai_agents/test_model_usage.py new file mode 100644 index 000000000..bf3dc8006 --- /dev/null +++ b/tests/lib/core/temporal/plugins/openai_agents/test_model_usage.py @@ -0,0 +1,182 @@ +"""Tests that the openai_agents streaming model copies real token usage onto spans. + +The backend bills per-call usage from ``span.output["usage"]``; these tests +assert the streaming model writes the API-reported usage there (and into the +returned ``ModelResponse.usage``) instead of dropping it. +""" + +from __future__ import annotations + +from typing import Any +from datetime import UTC, datetime +from contextlib import asynccontextmanager +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from agents import ModelSettings +from openai.types.responses import Response, ResponseCompletedEvent +from openai.types.responses.response_usage import ( + ResponseUsage, + InputTokensDetails, + OutputTokensDetails, +) + +import agentex.lib.core.temporal.plugins.openai_agents.models.temporal_streaming_model as tsm +from agentex.types.span import Span +from agentex.lib.core.temporal.plugins.openai_agents.interceptors.context_interceptor import ( + streaming_task_id, + streaming_trace_id, + streaming_parent_span_id, +) + +pytestmark = pytest.mark.asyncio + + +class FakeTrace: + """Captures spans handed out by trace.span() so tests can inspect them.""" + + def __init__(self) -> None: + self.spans: list[Span] = [] + + @asynccontextmanager + async def span(self, name, parent_id=None, input=None, data=None, task_id=None): + span = Span( + id=f"span-{len(self.spans)}", + name=name, + start_time=datetime.now(UTC), + trace_id="trace-1", + parent_id=parent_id, + input=input, + data=data, + task_id=task_id, + ) + self.spans.append(span) + yield span + + +class FakeTracer: + def __init__(self) -> None: + self.trace_obj = FakeTrace() + + def trace(self, trace_id): + return self.trace_obj + + +@pytest.fixture +def tracing_contextvars(): + tokens = [ + streaming_task_id.set("task-1"), + streaming_trace_id.set("trace-1"), + streaming_parent_span_id.set("parent-span-1"), + ] + yield + streaming_task_id.reset(tokens[0]) + streaming_trace_id.reset(tokens[1]) + streaming_parent_span_id.reset(tokens[2]) + + +EXPECTED_USAGE_BLOB = { + "input_tokens": 120, + "output_tokens": 80, + "total_tokens": 200, + "cached_input_tokens": 30, + "reasoning_tokens": 40, +} + + +def _output_dict(span: Span) -> dict[str, Any]: + assert isinstance(span.output, dict) + return span.output + + +class FakeStream: + def __init__(self, events) -> None: + self._events = events + + def __aiter__(self): + async def gen(): + for event in self._events: + yield event + + return gen() + + +class TestTemporalStreamingModel: + async def test_streaming_model_captures_final_response_usage(self, tracing_contextvars): + usage = ResponseUsage( + input_tokens=120, + output_tokens=80, + total_tokens=200, + input_tokens_details=InputTokensDetails(cached_tokens=30), + output_tokens_details=OutputTokensDetails(reasoning_tokens=40), + ) + completed = ResponseCompletedEvent.model_construct( + type="response.completed", + response=Response.model_construct(output=[], usage=usage), + ) + + fake_tracer = FakeTracer() + openai_client = MagicMock() + openai_client.responses.create = AsyncMock(return_value=FakeStream([completed])) + + with patch.object(tsm, "create_async_agentex_client", return_value=MagicMock()): + with patch.object(tsm, "AsyncTracer", return_value=fake_tracer): + model = tsm.TemporalStreamingModel(model_name="gpt-4o", openai_client=openai_client) + + response = await model.get_response( + system_instructions=None, + input="hello", + model_settings=ModelSettings(), + tools=[], + output_schema=None, + handoffs=[], + tracing=None, + ) + + # Real usage lands on the returned ModelResponse (was zeroed before) + assert response.usage.input_tokens == 120 + assert response.usage.output_tokens == 80 + assert response.usage.total_tokens == 200 + assert response.usage.input_tokens_details.cached_tokens == 30 + assert response.usage.output_tokens_details.reasoning_tokens == 40 + + # And on the span output for billing + assert len(fake_tracer.trace_obj.spans) == 1 + span = fake_tracer.trace_obj.spans[0] + assert _output_dict(span)["usage"] == EXPECTED_USAGE_BLOB + + async def test_streaming_model_writes_zero_usage_when_api_reports_none(self, tracing_contextvars): + completed = ResponseCompletedEvent.model_construct( + type="response.completed", + response=Response.model_construct(output=[], usage=None), + ) + + fake_tracer = FakeTracer() + openai_client = MagicMock() + openai_client.responses.create = AsyncMock(return_value=FakeStream([completed])) + + with patch.object(tsm, "create_async_agentex_client", return_value=MagicMock()): + with patch.object(tsm, "AsyncTracer", return_value=fake_tracer): + model = tsm.TemporalStreamingModel(model_name="gpt-4o", openai_client=openai_client) + + response = await model.get_response( + system_instructions=None, + input="hello", + model_settings=ModelSettings(), + tools=[], + output_schema=None, + handoffs=[], + tracing=None, + ) + + # No usage from the API: the model reports zeros rather than omitting, + # so billing sums 0 instead of missing the span + assert response.usage.input_tokens == 0 + span = fake_tracer.trace_obj.spans[0] + assert _output_dict(span)["usage"] == { + "input_tokens": 0, + "output_tokens": 0, + "total_tokens": 0, + "cached_input_tokens": 0, + "reasoning_tokens": 0, + } diff --git a/tests/lib/test_claude_agents_activities.py b/tests/lib/test_claude_agents_activities.py index fa633cd6b..99982d851 100644 --- a/tests/lib/test_claude_agents_activities.py +++ b/tests/lib/test_claude_agents_activities.py @@ -59,6 +59,7 @@ sys.modules.setdefault(_name, _mock) # Also ensure parent packages exist as stubs so Python resolves the dotted path +_created_pkg_stubs: list[str] = [] for _pkg in [ "agentex.lib.core.temporal.plugins", "agentex.lib.core.temporal.plugins.claude_agents", @@ -70,6 +71,7 @@ _mod.__path__ = [] # type: ignore[attr-defined] _mod.__package__ = _pkg sys.modules[_pkg] = _mod + _created_pkg_stubs.append(_pkg) # Load activities.py directly from its file path _spec = importlib.util.spec_from_file_location( @@ -81,6 +83,13 @@ sys.modules[_spec.name] = _activities_mod _spec.loader.exec_module(_activities_mod) +# Drop the placeholder packages now that the module is loaded: their empty +# __path__ would otherwise block real imports of agentex.lib.core.temporal.* +# submodules in every test collected after this file. The loaded module keeps +# its bindings. +for _pkg in _created_pkg_stubs: + sys.modules.pop(_pkg, None) + _reconstruct_agent_defs = _activities_mod._reconstruct_agent_defs # type: ignore[attr-defined] claude_options_to_dict = _activities_mod.claude_options_to_dict # type: ignore[attr-defined] diff --git a/tests/lib/test_claude_agents_hooks.py b/tests/lib/test_claude_agents_hooks.py index 373956680..ec82b16db 100644 --- a/tests/lib/test_claude_agents_hooks.py +++ b/tests/lib/test_claude_agents_hooks.py @@ -29,6 +29,7 @@ sys.modules.setdefault(_name, MagicMock()) # Ensure parent packages exist so the dotted path resolves +_created_pkg_stubs: list[str] = [] for _pkg in [ "agentex", "agentex.lib", @@ -43,6 +44,7 @@ _mod.__path__ = [] # type: ignore[attr-defined] _mod.__package__ = _pkg sys.modules[_pkg] = _mod + _created_pkg_stubs.append(_pkg) # Real pydantic types — importable without triggering the problematic chain. from agentex.types.tool_response_content import ToolResponseContent # noqa: E402 @@ -57,6 +59,12 @@ sys.modules[_spec.name] = _hooks_mod _spec.loader.exec_module(_hooks_mod) +# Drop the placeholder packages now that the module is loaded: their empty +# __path__ would otherwise block real imports of agentex.* submodules in every +# test collected after this file. The loaded module keeps its bindings. +for _pkg in _created_pkg_stubs: + sys.modules.pop(_pkg, None) + TemporalStreamingHooks = _hooks_mod.TemporalStreamingHooks create_streaming_hooks = _hooks_mod.create_streaming_hooks diff --git a/tests/lib/utils/__init__.py b/tests/lib/utils/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/lib/utils/test_completions.py b/tests/lib/utils/test_completions.py new file mode 100644 index 000000000..3aa5f9120 --- /dev/null +++ b/tests/lib/utils/test_completions.py @@ -0,0 +1,52 @@ +from __future__ import annotations + +from agentex.lib.utils.completions import concat_completion_chunks +from agentex.lib.types.llm_messages import Delta, Usage, Choice, Completion + + +def _delta_chunk(content: str, role: str | None = None) -> Completion: + return Completion(choices=[Choice(index=0, delta=Delta(content=content, role=role))]) + + +def _usage_only_chunk(prompt: int, completion: int) -> Completion: + # stream_options.include_usage: litellm/OpenAI send a final chunk that has + # usage but an empty choices list + return Completion( + choices=[], + usage=Usage(prompt_tokens=prompt, completion_tokens=completion, total_tokens=prompt + completion), + ) + + +class TestConcatCompletionChunks: + def test_concatenates_delta_content(self): + result = concat_completion_chunks([_delta_chunk("Hel", role="assistant"), _delta_chunk("lo!")]) + + assert result.choices[0].message.content == "Hello!" + + def test_trailing_usage_only_chunk_keeps_choices_and_usage(self): + result = concat_completion_chunks( + [_delta_chunk("Hel", role="assistant"), _delta_chunk("lo!"), _usage_only_chunk(10, 5)] + ) + + assert result.choices[0].message.content == "Hello!" + assert result.usage is not None + assert result.usage.prompt_tokens == 10 + assert result.usage.completion_tokens == 5 + assert result.usage.total_tokens == 15 + + def test_usage_summed_across_chunks(self): + chunk_a = _delta_chunk("a", role="assistant") + chunk_a.usage = Usage(prompt_tokens=1, completion_tokens=2, total_tokens=3) + chunk_b = _delta_chunk("b") + chunk_b.usage = Usage(prompt_tokens=4, completion_tokens=5, total_tokens=9) + + result = concat_completion_chunks([chunk_a, chunk_b]) + + assert result.usage.prompt_tokens == 5 + assert result.usage.completion_tokens == 7 + assert result.usage.total_tokens == 12 + + def test_no_usage_chunks_leave_usage_none(self): + result = concat_completion_chunks([_delta_chunk("x", role="assistant")]) + + assert result.usage is None