From 3a2314dcefbc24f3fa5d089d5e566a1af49e4f13 Mon Sep 17 00:00:00 2001 From: ehhuang Date: Sat, 29 Mar 2025 14:55:31 -0700 Subject: [PATCH] fix(telemetry): library client does not log span (#1833) --- llama_stack/distribution/library_client.py | 3 ++- tests/integration/telemetry/test_telemetry.py | 7 ------- 2 files changed, 2 insertions(+), 8 deletions(-) diff --git a/llama_stack/distribution/library_client.py b/llama_stack/distribution/library_client.py index 565f22ae0..c0143363d 100644 --- a/llama_stack/distribution/library_client.py +++ b/llama_stack/distribution/library_client.py @@ -328,8 +328,9 @@ class AsyncLlamaStackAsLibraryClient(AsyncLlamaStackClient): body = self._convert_body(path, options.method, body) + await start_trace(route, {"__location__": "library_client"}) + async def gen(): - await start_trace(route, {"__location__": "library_client"}) try: async for chunk in await func(**body): data = json.dumps(convert_pydantic_to_json_value(chunk)) diff --git a/tests/integration/telemetry/test_telemetry.py b/tests/integration/telemetry/test_telemetry.py index 10cbb2eeb..c46de3742 100644 --- a/tests/integration/telemetry/test_telemetry.py +++ b/tests/integration/telemetry/test_telemetry.py @@ -7,16 +7,10 @@ import time from uuid import uuid4 -import pytest from llama_stack_client import Agent -from llama_stack.distribution.library_client import LlamaStackAsLibraryClient - def test_agent_query_spans(llama_stack_client, text_model_id): - if isinstance(llama_stack_client, LlamaStackAsLibraryClient): - pytest.mark.xfail(reason="Need to fix LlamaStackAsLibraryClient to log spans") - agent = Agent(llama_stack_client, model=text_model_id, instructions="You are a helpful assistant") session_id = agent.create_session(f"test-session-{uuid4()}") agent.create_turn( @@ -41,7 +35,6 @@ def test_agent_query_spans(llama_stack_client, text_model_id): ], attributes_to_return=["input", "output"], ): - print(span.attributes) if span.attributes["output"] != "no shields": agent_logs.append(span.attributes)