added tracing

Author: michael-chou359

examples/tutorials/00_sync/040_pydantic_ai/project/acp.py

@@ -6,22 +6,37 @@
 
 from __future__ import annotations
 
+import os
 from typing import AsyncGenerator
 
 from dotenv import load_dotenv
 
 load_dotenv()
 
+import agentex.lib.adk as adk
 from project.agent import create_agent
-from agentex.lib.adk import convert_pydantic_ai_to_agentex_events
+from agentex.lib.adk import (
+    create_pydantic_ai_tracing_handler,
+    convert_pydantic_ai_to_agentex_events,
+)
 from agentex.lib.types.acp import SendMessageParams
+from agentex.lib.types.tracing import SGPTracingProcessorConfig
 from agentex.lib.utils.logging import make_logger
 from agentex.lib.sdk.fastacp.fastacp import FastACP
 from agentex.types.task_message_update import TaskMessageUpdate
 from agentex.types.task_message_content import TaskMessageContent
+from agentex.lib.core.tracing.tracing_processor_manager import add_tracing_processor_config
 
 logger = make_logger(__name__)
 
+add_tracing_processor_config(
+    SGPTracingProcessorConfig(
+        sgp_api_key=os.environ.get("SGP_API_KEY", ""),
+        sgp_account_id=os.environ.get("SGP_ACCOUNT_ID", ""),
+        sgp_base_url=os.environ.get("SGP_CLIENT_BASE_URL", ""),
+    )
+)
+
 acp = FastACP.create(acp_type="sync")
 
 _agent = None
@@ -41,10 +56,23 @@ async def handle_message_send(
 ) -> TaskMessageContent | list[TaskMessageContent] | AsyncGenerator[TaskMessageUpdate, None]:
     """Handle incoming messages from Agentex, streaming tokens and tool calls."""
     agent = get_agent()
+    task_id = params.task.id
 
     user_message = params.content.content
-    logger.info(f"Processing message for task {params.task.id}")
+    logger.info(f"Processing message for task {task_id}")
 
-    async with agent.run_stream_events(user_message) as stream:
-        async for event in convert_pydantic_ai_to_agentex_events(stream):
-            yield event
+    async with adk.tracing.span(
+        trace_id=task_id,
+        task_id=task_id,
+        name="message",
+        input={"message": user_message},
+        data={"__span_type__": "AGENT_WORKFLOW"},
+    ) as turn_span:
+        tracing_handler = create_pydantic_ai_tracing_handler(
+            trace_id=task_id,
+            parent_span_id=turn_span.id if turn_span else None,
+            task_id=task_id,
+        )
+        async with agent.run_stream_events(user_message) as stream:
+            async for event in convert_pydantic_ai_to_agentex_events(stream, tracing_handler=tracing_handler):
+                yield event

examples/tutorials/10_async/00_base/110_pydantic_ai/project/acp.py

@@ -15,7 +15,10 @@
 
 import agentex.lib.adk as adk
 from project.agent import create_agent
-from agentex.lib.adk import stream_pydantic_ai_events
+from agentex.lib.adk import (
+    stream_pydantic_ai_events,
+    create_pydantic_ai_tracing_handler,
+)
 from agentex.lib.types.acp import SendEventParams, CancelTaskParams, CreateTaskParams
 from agentex.lib.types.fastacp import AsyncACPConfig
 from agentex.lib.types.tracing import SGPTracingProcessorConfig
@@ -62,12 +65,18 @@ async def handle_task_event_send(params: SendEventParams):
 
     async with adk.tracing.span(
         trace_id=task_id,
+        task_id=task_id,
         name="message",
         input={"message": user_message},
         data={"__span_type__": "AGENT_WORKFLOW"},
     ) as turn_span:
+        tracing_handler = create_pydantic_ai_tracing_handler(
+            trace_id=task_id,
+            parent_span_id=turn_span.id if turn_span else None,
+            task_id=task_id,
+        )
         async with agent.run_stream_events(user_message) as stream:
-            final_output = await stream_pydantic_ai_events(stream, task_id)
+            final_output = await stream_pydantic_ai_events(stream, task_id, tracing_handler=tracing_handler)
 
         if turn_span:
             turn_span.output = {"final_output": final_output}

examples/tutorials/10_async/10_temporal/110_pydantic_ai/project/agent.py

@@ -26,7 +26,10 @@
 from pydantic_ai.durable_exec.temporal import TemporalAgent
 
 from project.tools import get_weather
-from agentex.lib.adk import stream_pydantic_ai_events
+from agentex.lib.adk import (
+    stream_pydantic_ai_events,
+    create_pydantic_ai_tracing_handler,
+)
 
 MODEL_NAME = "openai:gpt-4o-mini"
 SYSTEM_PROMPT = """You are a helpful AI assistant with access to tools.
@@ -50,6 +53,9 @@ class TaskDeps(BaseModel):
     """
 
     task_id: str
+    # When set, the event handler nests per-tool-call spans under this
+    # span. Typically the ID of the per-turn span opened by the workflow.
+    parent_span_id: str | None = None
 
 
 def _build_base_agent() -> Agent[TaskDeps, str]:
@@ -76,9 +82,19 @@ async def event_handler(
     Pydantic AI calls this with the live event stream as soon as the model
     activity begins emitting parts. Because the handler runs inside the
     activity (not the workflow), it can freely make non-deterministic
-    Redis writes.
+    Redis writes — including the tracing HTTP calls that record per-tool-call
+    spans under the workflow's per-turn span (when ``parent_span_id`` is set).
     """
-    await stream_pydantic_ai_events(events, run_context.deps.task_id)
+    tracing_handler = create_pydantic_ai_tracing_handler(
+        trace_id=run_context.deps.task_id,
+        parent_span_id=run_context.deps.parent_span_id,
+        task_id=run_context.deps.task_id,
+    )
+    await stream_pydantic_ai_events(
+        events,
+        run_context.deps.task_id,
+        tracing_handler=tracing_handler,
+    )
 
 
 # Construct the durable agent at module load time so that the

examples/tutorials/10_async/10_temporal/110_pydantic_ai/project/workflow.py

@@ -74,6 +74,7 @@ async def on_task_event_send(self, params: SendEventParams) -> None:
 
         async with adk.tracing.span(
             trace_id=params.task.id,
+            task_id=params.task.id,
             name=f"Turn {self._turn_number}",
             input={"message": params.event.content.content},
         ) as span:
@@ -86,7 +87,10 @@ async def on_task_event_send(self, params: SendEventParams) -> None:
             # temporal_agent pushes deltas to Redis so the UI sees tokens.
             result = await temporal_agent.run(
                 params.event.content.content,
-                deps=TaskDeps(task_id=params.task.id),
+                deps=TaskDeps(
+                    task_id=params.task.id,
+                    parent_span_id=span.id if span else None,
+                ),
             )
             if span:
                 span.output = {"final_output": result.output}

src/agentex/lib/adk/__init__.py

@@ -11,6 +11,7 @@
 from agentex.lib.adk._modules._langgraph_sync import convert_langgraph_to_agentex_events
 from agentex.lib.adk._modules._pydantic_ai_async import stream_pydantic_ai_events
 from agentex.lib.adk._modules._pydantic_ai_sync import convert_pydantic_ai_to_agentex_events
+from agentex.lib.adk._modules._pydantic_ai_tracing import create_pydantic_ai_tracing_handler
 from agentex.lib.adk._modules.events import EventsModule
 from agentex.lib.adk._modules.messages import MessagesModule
 from agentex.lib.adk._modules.state import StateModule
@@ -42,17 +43,15 @@
     "tracing",
     "events",
     "agent_task_tracker",
-
     # Checkpointing / LangGraph
     "create_checkpointer",
     "create_langgraph_tracing_handler",
     "stream_langgraph_events",
     "convert_langgraph_to_agentex_events",
-
     # Pydantic AI
     "stream_pydantic_ai_events",
     "convert_pydantic_ai_to_agentex_events",
-
+    "create_pydantic_ai_tracing_handler",
     # Providers
     "providers",
     # Utils

src/agentex/lib/adk/_modules/_pydantic_ai_async.py

@@ -11,15 +11,37 @@
 helper's convention). To stream tool-call argument tokens, see the sync
 converter at ``agentex.lib.adk._modules._pydantic_ai_sync`` which yields
 ``ToolRequestDelta`` events.
+
+Tracing is opt-in via a ``tracing_handler`` parameter — see
+``create_pydantic_ai_tracing_handler`` in
+``agentex.lib.adk._modules._pydantic_ai_tracing``.
 """
 
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from agentex.lib.adk._modules._pydantic_ai_tracing import (
+        AgentexPydanticAITracingHandler,
+    )
+
 
-async def stream_pydantic_ai_events(stream, task_id: str) -> str:
+async def stream_pydantic_ai_events(
+    stream,
+    task_id: str,
+    tracing_handler: "AgentexPydanticAITracingHandler | None" = None,
+) -> str:
     """Stream Pydantic AI events to Agentex via Redis.
 
     Args:
         stream: Async iterator yielded by ``agent.run_stream_events(...)``.
         task_id: The Agentex task ID to stream messages to.
+        tracing_handler: Optional handler from
+            ``create_pydantic_ai_tracing_handler(...)``. When provided, each
+            tool call in the run is also recorded as an Agentex child span
+            beneath the handler's configured ``parent_span_id``. Streaming
+            behavior is unchanged when omitted.
 
     Returns:
         The accumulated text content of the **last** text part in the run.
@@ -197,6 +219,12 @@ async def _close_reasoning():
                             author="agent",
                         ),
                     )
+                    if tracing_handler is not None and tool_call_id:
+                        await tracing_handler.on_tool_start(
+                            tool_call_id=tool_call_id,
+                            tool_name=tool_name,
+                            arguments=args,
+                        )
 
             elif isinstance(event, FunctionToolResultEvent):
                 await _close_text()
@@ -221,6 +249,11 @@ async def _close_reasoning():
                         author="agent",
                     ),
                 )
+                if tracing_handler is not None and tool_call_id:
+                    await tracing_handler.on_tool_end(
+                        tool_call_id=tool_call_id,
+                        result=content_str,
+                    )
 
             # FunctionToolCallEvent / FinalResultEvent / AgentRunResultEvent
             # are intentionally ignored — same as the sync converter.

src/agentex/lib/adk/_modules/_pydantic_ai_sync.py

@@ -21,9 +21,14 @@ async def handle_message_send(params):
 from __future__ import annotations
 
 import json
-from typing import Any, AsyncIterator
+from typing import TYPE_CHECKING, Any, AsyncIterator
 
 from pydantic_ai.run import AgentRunResultEvent
+
+if TYPE_CHECKING:
+    from agentex.lib.adk._modules._pydantic_ai_tracing import (
+        AgentexPydanticAITracingHandler,
+    )
 from pydantic_ai.messages import (
     TextPart,
     PartEndEvent,
@@ -98,6 +103,7 @@ def _tool_return_content(result: ToolReturnPart | Any) -> Any:
 
 async def convert_pydantic_ai_to_agentex_events(
     stream_response: AsyncIterator[Any],
+    tracing_handler: "AgentexPydanticAITracingHandler | None" = None,
 ) -> AsyncIterator[StreamTaskMessageStart | StreamTaskMessageDelta | StreamTaskMessageFull | StreamTaskMessageDone]:
     """Convert a Pydantic AI agent event stream into Agentex stream events.
 
@@ -120,6 +126,11 @@ async def convert_pydantic_ai_to_agentex_events(
         stream_response: The async iterator yielded by Pydantic AI's
             ``agent.run_stream_events(...)`` context manager (or a stream of
             ``AgentStreamEvent`` items received in an ``event_stream_handler``).
+        tracing_handler: Optional handler from
+            ``create_pydantic_ai_tracing_handler(...)``. When provided, each
+            tool call in the run is also recorded as an Agentex child span
+            beneath the handler's configured ``parent_span_id``. Streaming
+            behavior is unchanged when omitted.
 
     Yields:
         Agentex ``StreamTaskMessage*`` events suitable for forwarding back over
@@ -265,13 +276,34 @@ async def convert_pydantic_ai_to_agentex_events(
             if message_index is None:
                 continue
             yield StreamTaskMessageDone(type="done", index=message_index)
+            # Tool-call parts end with the model's full args known. Open a
+            # tracing child span for the tool execution now; close it when
+            # FunctionToolResultEvent arrives below.
+            if tracing_handler is not None and isinstance(event.part, ToolCallPart) and event.part.tool_call_id:
+                args: dict[str, Any] | str | None
+                raw_args = event.part.args
+                if isinstance(raw_args, dict):
+                    args = dict(raw_args)
+                elif isinstance(raw_args, str):
+                    try:
+                        args = json.loads(raw_args) if raw_args else {}
+                    except json.JSONDecodeError:
+                        args = {"_raw": raw_args}
+                else:
+                    args = {}
+                await tracing_handler.on_tool_start(
+                    tool_call_id=event.part.tool_call_id,
+                    tool_name=event.part.tool_name,
+                    arguments=args,
+                )
 
         elif isinstance(event, FunctionToolResultEvent):
             result = event.part
             tool_call_id = result.tool_call_id
             tool_name = getattr(result, "tool_name", "") or ""
             message_index = next_message_index
             next_message_index += 1
+            content_payload = _tool_return_content(result)
             yield StreamTaskMessageFull(
                 type="full",
                 index=message_index,
@@ -280,9 +312,14 @@ async def convert_pydantic_ai_to_agentex_events(
                     author="agent",
                     tool_call_id=tool_call_id,
                     name=tool_name,
-                    content=_tool_return_content(result),
+                    content=content_payload,
                 ),
             )
+            if tracing_handler is not None and tool_call_id:
+                await tracing_handler.on_tool_end(
+                    tool_call_id=tool_call_id,
+                    result=content_payload,
+                )
 
         elif isinstance(event, (FunctionToolCallEvent, FinalResultEvent, AgentRunResultEvent)):
             # Already covered by PartStart/PartDelta/PartEnd events above, or