From 5bbdcf70f61b3121d29f7343b2a79c1f62026f18 Mon Sep 17 00:00:00 2001
From: Ashwin Bharambe <ashwin.bharambe@gmail.com>
Date: Mon, 16 Dec 2024 11:39:33 -0800
Subject: [PATCH] Use a Dict instead of tree for spans

---
 llama_stack/apis/telemetry/telemetry.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/llama_stack/apis/telemetry/telemetry.py b/llama_stack/apis/telemetry/telemetry.py
index 12ec5f1d9..23a475bff 100644
--- a/llama_stack/apis/telemetry/telemetry.py
+++ b/llama_stack/apis/telemetry/telemetry.py
@@ -150,8 +150,7 @@ class EvalTrace(BaseModel):
 
 
 @json_schema_type
-class SpanWithChildren(Span):
-    children: List["SpanWithChildren"] = Field(default_factory=list)
+class SpanWithStatus(Span):
     status: Optional[SpanStatus] = None
 
 
@@ -192,7 +191,7 @@ class Telemetry(Protocol):
         span_id: str,
         attributes_to_return: Optional[List[str]] = None,
         max_depth: Optional[int] = None,
-    ) -> SpanWithChildren: ...
+    ) -> Dict[str, SpanWithStatus]: ...
 
     @webmethod(route="/telemetry/query-spans", method="POST")
     async def query_spans(