From 4dd08e559525ef9af1111c255db5461bd788981a Mon Sep 17 00:00:00 2001 From: Dinesh Yeduguru Date: Sat, 30 Nov 2024 18:30:56 -0800 Subject: [PATCH] add trace chunk count as span attr for async --- llama_stack/distribution/tracing.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/llama_stack/distribution/tracing.py b/llama_stack/distribution/tracing.py index 34ebc6c90..0e27b775a 100644 --- a/llama_stack/distribution/tracing.py +++ b/llama_stack/distribution/tracing.py @@ -102,10 +102,12 @@ def trace_protocol(cls: Type[T]) -> Type[T]: with tracing.span(f"{class_name}.{method_name}", span_attributes) as span: try: + count = 0 async for item in method(self, *args, **kwargs): yield item + count += 1 finally: - span.set_attribute("output", "streaming output") + span.set_attribute("chunk_count", count) @wraps(method) async def async_wrapper(self: Any, *args: Any, **kwargs: Any) -> Any: