phoenix-oss/llama-stack
1
0
Fork 0
forked from phoenix-oss/llama-stack-mirror
Code Pull requests Releases Packages 2 Activity Actions

Update Telemetry API so OpenAPI generation can work (#640)

Browse source 
We cannot use recursive types because not only does our OpenAPI
generator not like them, even if it did, it is not easy for all client
languages to automatically construct proper APIs (especially considering
garbage collection) around them. For now, we can return a `Dict[str,
SpanWithStatus]` instead of `SpanWithChildren` and rely on the client to
reconstruct the tree.

Also fixed a super subtle issue with the OpenAPI generation process
(monkey-patching of json_schema_type wasn't working because of import
reordering.)
This commit is contained in:
Ashwin Bharambe 2024-12-16 13:00:14 -08:00 • committed by GitHub
parent 78e2bfbe7a
commit 2e5bfcd42a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
10 changed files with 349 additions and 473 deletions
Download patch file Download diff file Expand all files Collapse all files

23
llama_stack/providers/utils/telemetry/sqlite_trace_store.py
View file

@ -6,11 +6,11 @@
import json
from datetime import datetime
from typing import List, Optional, Protocol
from typing import Dict, List, Optional, Protocol
import aiosqlite
from llama_stack.apis.telemetry import QueryCondition, SpanWithChildren, Trace
from llama_stack.apis.telemetry import QueryCondition, SpanWithStatus, Trace
class TraceStore(Protocol):
@ -27,7 +27,7 @@ class TraceStore(Protocol):
span_id: str,
attributes_to_return: Optional[List[str]] = None,
max_depth: Optional[int] = None,
) -> SpanWithChildren: ...
) -> Dict[str, SpanWithStatus]: ...
class SQLiteTraceStore(TraceStore):
@ -114,7 +114,7 @@ class SQLiteTraceStore(TraceStore):
span_id: str,
attributes_to_return: Optional[List[str]] = None,
max_depth: Optional[int] = None,
) -> SpanWithChildren:
) -> Dict[str, SpanWithStatus]:
# Build the attributes selection
attributes_select = "s.attributes"
if attributes_to_return:
@ -143,6 +143,7 @@ class SQLiteTraceStore(TraceStore):
ORDER BY depth, start_time
"""
spans_by_id = {}
async with aiosqlite.connect(self.conn_string) as conn:
conn.row_factory = aiosqlite.Row
async with conn.execute(query, (span_id, max_depth, max_depth)) as cursor:
@ -151,12 +152,8 @@ class SQLiteTraceStore(TraceStore):
if not rows:
raise ValueError(f"Span {span_id} not found")
# Build span tree
spans_by_id = {}
root_span = None
for row in rows:
span = SpanWithChildren(
span = SpanWithStatus(
span_id=row["span_id"],
trace_id=row["trace_id"],
parent_span_id=row["parent_span_id"],
@ -165,14 +162,8 @@ class SQLiteTraceStore(TraceStore):
end_time=datetime.fromisoformat(row["end_time"]),
attributes=json.loads(row["filtered_attributes"]),
status=row["status"].lower(),
children=[],
)
spans_by_id[span.span_id] = span
if span.span_id == span_id:
root_span = span
elif span.parent_span_id in spans_by_id:
spans_by_id[span.parent_span_id].children.append(span)
return root_span
return spans_by_id