phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-08-02 00:34:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

add trace chunk count as span attr for async

Browse source
This commit is contained in:
Dinesh Yeduguru 2024-11-30 18:30:56 -08:00
parent e5c85b9dba
commit 4dd08e5595
1 changed files with 3 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

4
llama_stack/distribution/tracing.py
View file

@ -102,10 +102,12 @@ def trace_protocol(cls: Type[T]) -> Type[T]:
with tracing.span(f"{class_name}.{method_name}", span_attributes) as span: with tracing.span(f"{class_name}.{method_name}", span_attributes) as span:
try: try:
count = 0
async for item in method(self, *args, **kwargs): async for item in method(self, *args, **kwargs):
yield item yield item
count += 1
finally: finally:
span.set_attribute("output", "streaming output") span.set_attribute("chunk_count", count)
@wraps(method) @wraps(method)
async def async_wrapper(self: Any, *args: Any, **kwargs: Any) -> Any: async def async_wrapper(self: Any, *args: Any, **kwargs: Any) -> Any: