await end_trace in libcli

This commit is contained in:
Dinesh Yeduguru 2024-12-09 15:54:42 -08:00
parent bc1fddf1df
commit 7615da78b8

View file

@ -253,7 +253,7 @@ class AsyncLlamaStackAsLibraryClient(AsyncLlamaStackClient):
body = self._convert_body(path, body) body = self._convert_body(path, body)
return await func(**body) return await func(**body)
finally: finally:
end_trace() await end_trace()
async def _call_streaming(self, path: str, method: str, body: dict = None): async def _call_streaming(self, path: str, method: str, body: dict = None):
await start_trace(path, {"__location__": "library_client"}) await start_trace(path, {"__location__": "library_client"})
@ -266,7 +266,7 @@ class AsyncLlamaStackAsLibraryClient(AsyncLlamaStackClient):
async for chunk in await func(**body): async for chunk in await func(**body):
yield chunk yield chunk
finally: finally:
end_trace() await end_trace()
def _convert_body(self, path: str, body: Optional[dict] = None) -> dict: def _convert_body(self, path: str, body: Optional[dict] = None) -> dict:
if not body: if not body: