phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2026-01-02 19:20:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix tracing and add tests

Browse source
This commit is contained in:
Dinesh Yeduguru 2025-03-12 11:29:53 -07:00
parent 714c09cd53
commit 518a5f898c
3 changed files with 170 additions and 16 deletions
Download patch file Download diff file Expand all files Collapse all files

19
llama_stack/distribution/library_client.py
View file

@ -376,18 +376,17 @@ class AsyncLlamaStackAsLibraryClient(AsyncLlamaStackClient):
body = self._convert_body(path, options.method, body)
await start_trace(options.url, {"__location__": "library_client"})
async def gen():
async for chunk in await func(**body):
data = json.dumps(convert_pydantic_to_json_value(chunk))
sse_event = f"data: {data}\n\n"
yield sse_event.encode("utf-8")
await start_trace(options.url, {"__location__": "library_client"})
try:
async for chunk in await func(**body):
data = json.dumps(convert_pydantic_to_json_value(chunk))
sse_event = f"data: {data}\n\n"
yield sse_event.encode("utf-8")
finally:
await end_trace()
try:
wrapped_gen = preserve_contexts_async_generator(gen(), [CURRENT_TRACE_CONTEXT, PROVIDER_DATA_VAR])
finally:
await end_trace()
wrapped_gen = preserve_contexts_async_generator(gen(), [CURRENT_TRACE_CONTEXT, PROVIDER_DATA_VAR])
mock_response = httpx.Response(
status_code=httpx.codes.OK,