add new endpoints

This commit is contained in:
Dinesh Yeduguru 2025-01-16 11:03:34 -08:00
parent a9192e252c
commit bb2adef7e4
3 changed files with 31 additions and 1 deletions

View file

@ -197,6 +197,12 @@ class Telemetry(Protocol):
order_by: Optional[List[str]] = None, order_by: Optional[List[str]] = None,
) -> QueryTracesResponse: ... ) -> QueryTracesResponse: ...
@webmethod(route="/telemetry/traces/{trace_id}", method="GET")
async def get_trace(self, trace_id: str) -> Trace: ...
@webmethod(route="/telemetry/traces/{trace_id}/spans/{span_id}", method="GET")
async def get_span(self, trace_id: str, span_id: str) -> Span: ...
@webmethod(route="/telemetry/spans/{span_id}/tree", method="GET") @webmethod(route="/telemetry/spans/{span_id}/tree", method="GET")
async def query_span_tree( async def query_span_tree(
self, self,

View file

@ -23,11 +23,13 @@ from llama_stack.apis.telemetry import (
QueryCondition, QueryCondition,
QuerySpanTreeResponse, QuerySpanTreeResponse,
QueryTracesResponse, QueryTracesResponse,
Span,
SpanEndPayload, SpanEndPayload,
SpanStartPayload, SpanStartPayload,
SpanStatus, SpanStatus,
StructuredLogEvent, StructuredLogEvent,
Telemetry, Telemetry,
Trace,
UnstructuredLogEvent, UnstructuredLogEvent,
) )
from llama_stack.distribution.datatypes import Api from llama_stack.distribution.datatypes import Api
@ -251,6 +253,12 @@ class TelemetryAdapter(TelemetryDatasetMixin, Telemetry):
) )
) )
async def get_trace(self, trace_id: str) -> Trace:
return await self.trace_store.get_trace(trace_id)
async def get_span(self, trace_id: str, span_id: str) -> Span:
return await self.trace_store.get_span(trace_id, span_id)
async def query_span_tree( async def query_span_tree(
self, self,
span_id: str, span_id: str,

View file

@ -10,7 +10,7 @@ from typing import Dict, List, Optional, Protocol
import aiosqlite import aiosqlite
from llama_stack.apis.telemetry import QueryCondition, SpanWithStatus, Trace from llama_stack.apis.telemetry import QueryCondition, Span, SpanWithStatus, Trace
class TraceStore(Protocol): class TraceStore(Protocol):
@ -167,3 +167,19 @@ class SQLiteTraceStore(TraceStore):
spans_by_id[span.span_id] = span spans_by_id[span.span_id] = span
return spans_by_id return spans_by_id
async def get_trace(self, trace_id: str) -> Trace:
query = "SELECT * FROM traces WHERE trace_id = ?"
async with aiosqlite.connect(self.conn_string) as conn:
conn.row_factory = aiosqlite.Row
async with conn.execute(query, (trace_id,)) as cursor:
row = await cursor.fetchone()
return Trace(**row)
async def get_span(self, trace_id: str, span_id: str) -> Span:
query = "SELECT * FROM spans WHERE trace_id = ? AND span_id = ?"
async with aiosqlite.connect(self.conn_string) as conn:
conn.row_factory = aiosqlite.Row
async with conn.execute(query, (trace_id, span_id)) as cursor:
row = await cursor.fetchone()
return Span(**row)