From b3e149334a6f4bdfbe4222b2965fec52641d9ac0 Mon Sep 17 00:00:00 2001 From: Dinesh Yeduguru Date: Tue, 26 Nov 2024 15:41:08 -0800 Subject: [PATCH] fixes --- llama_stack/distribution/tracing.py | 10 ++++++---- .../inline/agents/meta_reference/agent_instance.py | 9 ++++++++- .../inline/meta_reference/telemetry/console.py | 9 ++++++++- 3 files changed, 22 insertions(+), 6 deletions(-) diff --git a/llama_stack/distribution/tracing.py b/llama_stack/distribution/tracing.py index 7737b2abc..34ebc6c90 100644 --- a/llama_stack/distribution/tracing.py +++ b/llama_stack/distribution/tracing.py @@ -155,17 +155,19 @@ def trace_protocol(cls: Type[T]) -> Type[T]: traced_methods = {} for parent in cls_child.__mro__[1:]: # Skip the class itself for name, method in vars(parent).items(): - if inspect.isfunction(method) and method._trace_input: - traced_methods[name] = method._trace_input + if inspect.isfunction(method) and getattr( + method, "_trace_input", None + ): # noqa: B009 + traced_methods[name] = getattr(method, "_trace_input") # noqa: B009 # Trace child class methods if their name matches a traced parent method for name, method in vars(cls_child).items(): if inspect.isfunction(method) and not name.startswith("_"): if name in traced_methods: # Copy the trace configuration from the parent - method._trace_input = traced_methods[name] + setattr(method, "_trace_input", traced_methods[name]) # noqa: B010 - cls_child.__dict__[name] = trace_method(method) + setattr(cls_child, name, trace_method(method)) # noqa: B010 # Set the new __init_subclass__ cls.__init_subclass__ = classmethod(__init_subclass__) diff --git a/llama_stack/providers/inline/agents/meta_reference/agent_instance.py b/llama_stack/providers/inline/agents/meta_reference/agent_instance.py index dddb34b5d..49149d886 100644 --- a/llama_stack/providers/inline/agents/meta_reference/agent_instance.py +++ b/llama_stack/providers/inline/agents/meta_reference/agent_instance.py @@ -568,7 +568,13 @@ class ChatAgent(ShieldRunnerMixin): ) ) - with tracing.span("tool_execution"): + with tracing.span( + "tool_execution", + { + "tool_name": tool_call.tool_name, + "input": message.model_dump_json(), + }, + ) as span: result_messages = await execute_tool_call_maybe( self.tools_dict, [message], @@ -577,6 +583,7 @@ class ChatAgent(ShieldRunnerMixin): len(result_messages) == 1 ), "Currently not supporting multiple messages" result_message = result_messages[0] + span.set_attribute("output", result_message.model_dump_json()) yield AgentTurnResponseStreamChunk( event=AgentTurnResponseEvent( diff --git a/llama_stack/providers/inline/meta_reference/telemetry/console.py b/llama_stack/providers/inline/meta_reference/telemetry/console.py index d8ef49481..264b82b69 100644 --- a/llama_stack/providers/inline/meta_reference/telemetry/console.py +++ b/llama_stack/providers/inline/meta_reference/telemetry/console.py @@ -5,7 +5,9 @@ # the root directory of this source tree. import json -from typing import Optional +from typing import List, Optional + +from llama_stack.apis.telemetry.telemetry import Trace from .config import LogFormat @@ -52,6 +54,11 @@ class ConsoleTelemetryImpl(Telemetry): async def get_trace(self, trace_id: str) -> Trace: raise NotImplementedError() + async def get_traces_for_session( + self, session_id: str, lookback: str = "1h", limit: int = 100 + ) -> List[Trace]: + raise NotImplementedError() + COLORS = { "reset": "\033[0m",