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

restore trace context in event loop of lib cli

Browse source
This commit is contained in:
Dinesh Yeduguru 2025-03-11 15:02:34 -07:00
parent 0cac53ab30
commit 9d4716521d
1 changed files with 8 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

8
llama_stack/distribution/library_client.py
View file

@ -46,6 +46,7 @@ from llama_stack.distribution.stack import (
)
from llama_stack.distribution.utils.exec import in_notebook
from llama_stack.providers.utils.telemetry.tracing import (
CURRENT_TRACE_CONTEXT,
end_trace,
setup_logger,
start_trace,
@ -55,6 +56,8 @@ logger = logging.getLogger(__name__)
T = TypeVar("T")
trace_context = None
def convert_pydantic_to_json_value(value: Any) -> Any:
if isinstance(value, Enum):
@ -156,8 +159,11 @@ class LlamaStackAsLibraryClient(LlamaStackClient):
def sync_generator():
try:
global trace_context
async_stream = loop.run_until_complete(self.async_client.request(*args, **kwargs))
while True:
if trace_context:
CURRENT_TRACE_CONTEXT.set(trace_context)
chunk = loop.run_until_complete(async_stream.__anext__())
yield chunk
except StopAsyncIteration:
@ -376,6 +382,8 @@ class AsyncLlamaStackAsLibraryClient(AsyncLlamaStackClient):
async def gen():
await start_trace(options.url, {"__location__": "library_client"})
global trace_context
trace_context = CURRENT_TRACE_CONTEXT.get()
try:
async for chunk in await func(**body):
data = json.dumps(convert_pydantic_to_json_value(chunk))