修复tool监控丢失的问题

raychen911 · raychen911 · commit 2d38d1806463 · 2026-04-24T11:29:02.000+08:00
diff --git a/trpc_agent_sdk/agents/core/_llm_processor.py b/trpc_agent_sdk/agents/core/_llm_processor.py
@@ -79,20 +79,36 @@ async def call_llm_async(self,
                 return
 
             # Step 2: Call the model and process responses with telemetry tracing.
-            # Avoid start_as_current_span in async generators because cancellation can
-            # close the generator from a different context, which may trigger
-            # "Token was created in a different Context" during detach.
-            span = tracer.start_span('call_llm')
-            try:
+            with tracer.start_as_current_span('call_llm'):
                 event_id = Event.new_id()
                 final_llm_response = None
+                aggregated_raw_function_calls: list[dict] = []
+                aggregated_event_function_calls: list[dict] = []
+
+                def _append_function_calls(target: list[dict], calls: list) -> None:
+                    for call in calls or []:
+                        # Keep only telemetry-safe fields for trace attributes.
+                        target.append({
+                            "id": getattr(call, "id", None),
+                            "name": getattr(call, "name", None),
+                            "args": getattr(call, "args", None),
+                        })
 
                 async for llm_response in self.model.generate_async(request, stream=stream, ctx=context):
+                    # Collect raw model-level function calls from every chunk.
+                    raw_calls = []
+                    if llm_response.content and llm_response.content.parts:
+                        for part in llm_response.content.parts:
+                            if part.function_call:
+                                raw_calls.append(part.function_call)
+                    _append_function_calls(aggregated_raw_function_calls, raw_calls)
+
                     # Create Event directly from LlmResponse
                     event = self._create_event_from_response(context, event_id, llm_response)
 
                     # Process response with planner if available
                     event = self._process_planning_response(event, context)
+                    _append_function_calls(aggregated_event_function_calls, event.get_function_calls())
 
                     # Track the latest non-partial response for tracing
                     # In streaming mode, only the final (non-partial) response
@@ -111,9 +127,9 @@ async def call_llm_async(self,
                                    event_id,
                                    request,
                                    final_llm_response,
-                                   instruction_metadata=instruction_metadata)
-            finally:
-                span.end()
+                                   instruction_metadata=instruction_metadata,
+                                   stream_function_calls_raw=aggregated_raw_function_calls,
+                                   stream_function_calls_post_planner=aggregated_event_function_calls)
 
         except Exception as ex:  # pylint: disable=broad-except
             logger.error("LLM call failed for agent %s: %s", author, ex)
diff --git a/trpc_agent_sdk/agents/core/_tools_processor.py b/trpc_agent_sdk/agents/core/_tools_processor.py
@@ -221,7 +221,10 @@ async def execute_tools_async(
                     state_end.update(merged_event.actions.state_delta)
 
                 # Add merged tool call tracing
-                with tracer.start_as_current_span("execute_tool (merged)"):
+                with tracer.start_as_current_span(
+                    "execute_tool (merged)",
+                    attributes={"gen_ai.operation.name": "execute_tool"},
+                ):
                     trace_merged_tool_calls(
                         response_event_id=merged_event.id,
                         function_response_event=merged_event,
@@ -285,8 +288,17 @@ async def _execute_tool(self, tool_call: FunctionCall, tool: BaseTool, context:
             Event: The result of tool execution
         """
 
-        # Wrap tool execution in telemetry span
-        with tracer.start_as_current_span(f"execute_tool {tool.name}"):
+        # Wrap tool execution in telemetry span.
+        # Pass initial attributes so the Galileo sampler can make a sampling
+        # decision at span-creation time (before trace_tool_call sets them).
+        with tracer.start_as_current_span(
+            f"execute_tool {tool.name}",
+            attributes={
+                "gen_ai.operation.name": "execute_tool",
+                "gen_ai.tool.name": tool.name,
+                "gen_ai.tool.description": tool.description or "",
+            },
+        ):
             # Capture state before tool execution
             state_begin = dict(context.session.state)
 
diff --git a/trpc_agent_sdk/filter/_run_filter.py b/trpc_agent_sdk/filter/_run_filter.py
@@ -51,8 +51,7 @@ async def run_stream_filters(ctx: AgentContext, req: Any, filters: list[BaseFilt
     if handle is None:
         raise ValueError("handle must be provided")
     current_handle = partial(stream_handler_adapter, handle)
-    filters.reverse()
-    for filter in filters:
+    for filter in reversed(filters):
         current_handle = partial(filter.run_stream, ctx, req, current_handle)
     async for event in current_handle():
         yield event.rsp
@@ -95,9 +94,8 @@ async def run_filters(ctx: AgentContext, req: Any, filters: list[BaseFilter],
     """
     if handle is None:
         raise ValueError("handle must be provided")
-    filters.reverse()
     current_handle = partial(coroutine_handler_adapter, handle)
-    for filter in filters:
+    for filter in reversed(filters):
         current_handle = partial(filter.run, ctx, req, current_handle)
     rsp, error = await current_handle()
     if error:
diff --git a/trpc_agent_sdk/storage/_sql_common.py b/trpc_agent_sdk/storage/_sql_common.py
@@ -44,6 +44,64 @@
 from trpc_agent_sdk.types import GroundingMetadata
 from trpc_agent_sdk.types import GenerateContentResponseUsageMetadata
 
+LoadDialectHook = Callable[[TypeDecorator, Dialect], Any]
+ProcessBindHook = Callable[[TypeDecorator, Any, Dialect], Any]
+ProcessResultHook = Callable[[TypeDecorator, Any, Dialect], Any]
+
+
+class TypeDecoratorHookRegistry:
+    """Global hook registry for SQLAlchemy TypeDecorator callbacks."""
+
+    _load_dialect_hooks: dict[type[TypeDecorator], list[LoadDialectHook]] = {}
+    _process_bind_hooks: dict[type[TypeDecorator], list[ProcessBindHook]] = {}
+    _process_result_hooks: dict[type[TypeDecorator], list[ProcessResultHook]] = {}
+
+    @classmethod
+    def register_load_dialect_hook(cls, decorator_cls: type[TypeDecorator], hook: LoadDialectHook) -> None:
+        """Register hook for ``load_dialect_impl``."""
+        cls._load_dialect_hooks.setdefault(decorator_cls, []).append(hook)
+
+    @classmethod
+    def register_process_bind_hook(cls, decorator_cls: type[TypeDecorator], hook: ProcessBindHook) -> None:
+        """Register hook for ``process_bind_param``."""
+        cls._process_bind_hooks.setdefault(decorator_cls, []).append(hook)
+
+    @classmethod
+    def register_process_result_hook(cls, decorator_cls: type[TypeDecorator], hook: ProcessResultHook) -> None:
+        """Register hook for ``process_result_value``."""
+        cls._process_result_hooks.setdefault(decorator_cls, []).append(hook)
+
+    @classmethod
+    def run_load_dialect_hooks(cls, decorator: TypeDecorator, dialect: Dialect) -> Any:
+        """Run load hooks and return first override result."""
+        for hook in cls._load_dialect_hooks.get(type(decorator), []):
+            result = hook(decorator, dialect)
+            if result is not None:
+                return result
+        return None
+
+    @classmethod
+    def run_process_bind_hooks(cls, decorator: TypeDecorator, value: Any, dialect: Dialect) -> Any:
+        """Run bind hooks and return first override result."""
+        for hook in cls._process_bind_hooks.get(type(decorator), []):
+            result = hook(decorator, value, dialect)
+            if result is not None:
+                return result
+        return None
+
+    @classmethod
+    def run_process_result_hooks(cls, decorator: TypeDecorator, value: Any, dialect: Dialect) -> Any:
+        """Run result hooks and return first override result."""
+        for hook in cls._process_result_hooks.get(type(decorator), []):
+            result = hook(decorator, value, dialect)
+            if result is not None:
+                return result
+        return None
+
+
+# Global class object used as unified registration entry.
+GLOBAL_TYPE_DECORATOR_HOOK_REGISTRY = TypeDecoratorHookRegistry
+
 
 def decode_content(content: Optional[dict[str, Any]]) -> Optional[Content]:
     """Decode a content object from a JSON dictionary.
@@ -119,13 +177,19 @@ class DynamicJSON(TypeDecorator):
     impl = Text  # Default implementation is TEXT
 
     def load_dialect_impl(self, dialect: Dialect) -> TypeDecorator:
+        hook_result = GLOBAL_TYPE_DECORATOR_HOOK_REGISTRY.run_load_dialect_hooks(self, dialect)
+        if hook_result is not None:
+            return hook_result
         if dialect.name == "postgresql":
             return dialect.type_descriptor(postgresql.JSONB)  # type: ignore
         if dialect.name == "mysql":
             return dialect.type_descriptor(mysql.LONGTEXT)  # type: ignore
         return dialect.type_descriptor(Text)  # Default to Text for other dialects # type: ignore
 
     def process_bind_param(self, value: Any, dialect: Dialect) -> Any:
+        hook_result = GLOBAL_TYPE_DECORATOR_HOOK_REGISTRY.run_process_bind_hooks(self, value, dialect)
+        if hook_result is not None:
+            return hook_result
         if value is not None:
             if dialect.name == "postgresql":
                 return value  # JSONB handles dict directly
@@ -134,6 +198,9 @@ def process_bind_param(self, value: Any, dialect: Dialect) -> Any:
         return value
 
     def process_result_value(self, value: Any, dialect: Dialect) -> Any:
+        hook_result = GLOBAL_TYPE_DECORATOR_HOOK_REGISTRY.run_process_result_hooks(self, value, dialect)
+        if hook_result is not None:
+            return hook_result
         if value is not None:
             if dialect.name == "postgresql":
                 return value  # JSONB returns dict directly
@@ -157,6 +224,9 @@ def __init__(self, length: Optional[int] = None, *args: Any, **kwargs: Any) -> N
         self.length = length
 
     def load_dialect_impl(self, dialect: Dialect) -> TypeDecorator:
+        hook_result = GLOBAL_TYPE_DECORATOR_HOOK_REGISTRY.run_load_dialect_hooks(self, dialect)
+        if hook_result is not None:
+            return hook_result
         if dialect.name == "mysql":
             # Use VARCHAR with utf8mb4 charset and utf8mb4_unicode_ci collation
             return dialect.type_descriptor(mysql.VARCHAR(self.length, charset='utf8mb4',
@@ -166,6 +236,18 @@ def load_dialect_impl(self, dialect: Dialect) -> TypeDecorator:
             return dialect.type_descriptor(String(self.length))
         return dialect.type_descriptor(String())
 
+    def process_bind_param(self, value: Any, dialect: Dialect) -> Any:
+        hook_result = GLOBAL_TYPE_DECORATOR_HOOK_REGISTRY.run_process_bind_hooks(self, value, dialect)
+        if hook_result is not None:
+            return hook_result
+        return value
+
+    def process_result_value(self, value: Any, dialect: Dialect) -> Any:
+        hook_result = GLOBAL_TYPE_DECORATOR_HOOK_REGISTRY.run_process_result_hooks(self, value, dialect)
+        if hook_result is not None:
+            return hook_result
+        return value
+
 
 class PreciseTimestamp(TypeDecorator):
     """Represents a timestamp precise to the microsecond."""
@@ -174,30 +256,54 @@ class PreciseTimestamp(TypeDecorator):
     cache_ok = True
 
     def load_dialect_impl(self, dialect: Dialect) -> TypeDecorator:
+        hook_result = GLOBAL_TYPE_DECORATOR_HOOK_REGISTRY.run_load_dialect_hooks(self, dialect)
+        if hook_result is not None:
+            return hook_result
         if dialect.name == "mysql":
             return dialect.type_descriptor(mysql.DATETIME(fsp=6))
         return self.impl
 
+    def process_bind_param(self, value: Any, dialect: Dialect) -> Any:
+        hook_result = GLOBAL_TYPE_DECORATOR_HOOK_REGISTRY.run_process_bind_hooks(self, value, dialect)
+        if hook_result is not None:
+            return hook_result
+        return value
+
+    def process_result_value(self, value: Any, dialect: Dialect) -> Any:
+        hook_result = GLOBAL_TYPE_DECORATOR_HOOK_REGISTRY.run_process_result_hooks(self, value, dialect)
+        if hook_result is not None:
+            return hook_result
+        return value
+
 
 class DynamicPickleType(TypeDecorator):
     """Represents a type that can be pickled."""
 
     impl = PickleType
 
     def load_dialect_impl(self, dialect: Dialect) -> TypeDecorator:
+        hook_result = GLOBAL_TYPE_DECORATOR_HOOK_REGISTRY.run_load_dialect_hooks(self, dialect)
+        if hook_result is not None:
+            return hook_result
         if dialect.name == "spanner+spanner":
             return dialect.type_descriptor(SpannerPickleType)  # type: ignore
         if dialect.name == "mysql":
             return dialect.type_descriptor(mysql.LONGBLOB)  # type: ignore
         return self.impl
 
     def process_bind_param(self, value: Any, dialect: Dialect) -> Any:
+        hook_result = GLOBAL_TYPE_DECORATOR_HOOK_REGISTRY.run_process_bind_hooks(self, value, dialect)
+        if hook_result is not None:
+            return hook_result
         if value is not None:
             if dialect.name == "spanner+spanner":
                 return pickle.dumps(value)
         return value
 
     def process_result_value(self, value: Any, dialect: Dialect) -> Any:
+        hook_result = GLOBAL_TYPE_DECORATOR_HOOK_REGISTRY.run_process_result_hooks(self, value, dialect)
+        if hook_result is not None:
+            return hook_result
         if value is not None:
             if dialect.name == "spanner+spanner":
                 return pickle.loads(value)
diff --git a/trpc_agent_sdk/telemetry/_trace.py b/trpc_agent_sdk/telemetry/_trace.py
@@ -372,6 +372,8 @@ def trace_call_llm(
     llm_request: LlmRequest,
     llm_response: LlmResponse,
     instruction_metadata: Optional[InstructionMetadata] = None,
+    stream_function_calls_raw: Optional[list[dict[str, Any]]] = None,
+    stream_function_calls_post_planner: Optional[list[dict[str, Any]]] = None,
 ):
     """Traces a call to the LLM.
 
@@ -387,6 +389,10 @@ def trace_call_llm(
             with ``name``, ``version``, and ``labels`` attributes. When
             provided, these are written directly to the call_llm span for
             precise instruction-to-generation association.
+        stream_function_calls_raw: Optional function calls collected from all
+            raw LLM stream chunks.
+        stream_function_calls_post_planner: Optional function calls collected
+            from post-planner events emitted during stream processing.
     """
     global _trpc_agent_span_name  # pylint: disable=invalid-name
     span = trace.get_current_span()
@@ -414,6 +420,26 @@ def trace_call_llm(
         llm_response_json,
     )
 
+    if stream_function_calls_raw:
+        span.set_attribute(
+            f"{_trpc_agent_span_name}.stream_function_calls.raw",
+            _safe_json_serialize(stream_function_calls_raw),
+        )
+        span.set_attribute(
+            f"{_trpc_agent_span_name}.stream_function_calls.raw_count",
+            len(stream_function_calls_raw),
+        )
+
+    if stream_function_calls_post_planner:
+        span.set_attribute(
+            f"{_trpc_agent_span_name}.stream_function_calls.post_planner",
+            _safe_json_serialize(stream_function_calls_post_planner),
+        )
+        span.set_attribute(
+            f"{_trpc_agent_span_name}.stream_function_calls.post_planner_count",
+            len(stream_function_calls_post_planner),
+        )
+
     if llm_response.usage_metadata is not None:
         usage = llm_response.usage_metadata
         if usage.prompt_token_count and usage.total_token_count:
