From bb2adef7e4752a2dda6e4a685d0203d259948cc8 Mon Sep 17 00:00:00 2001
From: Dinesh Yeduguru <yvdinesh@gmail.com>
Date: Thu, 16 Jan 2025 11:03:34 -0800
Subject: [PATCH] add new endpoints

---
 llama_stack/apis/telemetry/telemetry.py        |  6 ++++++
 .../telemetry/meta_reference/telemetry.py      |  8 ++++++++
 .../utils/telemetry/sqlite_trace_store.py      | 18 +++++++++++++++++-
 3 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/llama_stack/apis/telemetry/telemetry.py b/llama_stack/apis/telemetry/telemetry.py
index 94d51f9c6..7d964dbce 100644
--- a/llama_stack/apis/telemetry/telemetry.py
+++ b/llama_stack/apis/telemetry/telemetry.py
@@ -197,6 +197,12 @@ class Telemetry(Protocol):
         order_by: Optional[List[str]] = None,
     ) -> 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")
     async def query_span_tree(
         self,
diff --git a/llama_stack/providers/inline/telemetry/meta_reference/telemetry.py b/llama_stack/providers/inline/telemetry/meta_reference/telemetry.py
index a9d94d857..f4a6df214 100644
--- a/llama_stack/providers/inline/telemetry/meta_reference/telemetry.py
+++ b/llama_stack/providers/inline/telemetry/meta_reference/telemetry.py
@@ -23,11 +23,13 @@ from llama_stack.apis.telemetry import (
     QueryCondition,
     QuerySpanTreeResponse,
     QueryTracesResponse,
+    Span,
     SpanEndPayload,
     SpanStartPayload,
     SpanStatus,
     StructuredLogEvent,
     Telemetry,
+    Trace,
     UnstructuredLogEvent,
 )
 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(
         self,
         span_id: str,
diff --git a/llama_stack/providers/utils/telemetry/sqlite_trace_store.py b/llama_stack/providers/utils/telemetry/sqlite_trace_store.py
index b0c3f7868..ea7324166 100644
--- a/llama_stack/providers/utils/telemetry/sqlite_trace_store.py
+++ b/llama_stack/providers/utils/telemetry/sqlite_trace_store.py
@@ -10,7 +10,7 @@ from typing import Dict, List, Optional, Protocol
 
 import aiosqlite
 
-from llama_stack.apis.telemetry import QueryCondition, SpanWithStatus, Trace
+from llama_stack.apis.telemetry import QueryCondition, Span, SpanWithStatus, Trace
 
 
 class TraceStore(Protocol):
@@ -167,3 +167,19 @@ class SQLiteTraceStore(TraceStore):
                     spans_by_id[span.span_id] = span
 
                 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)