[Enhancement] Update SDK to return BaseResponse version 3.0 and new trace tree format

agenta-cli/agenta/sdk/tracing/inline.py

@@ -992,50 +992,13 @@ def parse_inline_trace(
     ### services.observability.service.query() ###
     ##############################################
 
-    LEGACY = True
     inline_trace = None
 
-    if LEGACY:
-        legacy_spans = [
-            _parse_to_legacy_span(span_dto) for span_dto in span_idx.values()
-        ]
-
-        root_span = agenta_span_dtos[0]
-
-        trace_id = root_span.root.id.hex
-        latency = root_span.time.span / 1_000_000
-        cost = root_span.metrics.get("acc", {}).get("costs", {}).get("total", 0.0)
-        tokens = {
-            "prompt_tokens": root_span.metrics.get("acc", {})
-            .get("tokens", {})
-            .get("prompt", 0),
-            "completion_tokens": root_span.metrics.get("acc", {})
-            .get("tokens", {})
-            .get("completion", 0),
-            "total_tokens": root_span.metrics.get("acc", {})
-            .get("tokens", {})
-            .get("total", 0),
-        }
-
-        spans = [
-            loads(span.model_dump_json(exclude_none=True)) for span in legacy_spans
-        ]
-
-        inline_trace = {
-            "trace_id": trace_id,
-            "latency": latency,
-            "cost": cost,
-            "usage": tokens,
-            "spans": spans,
-        }
-
-    else:
-        spans = [
-            loads(span_dto.model_dump_json(exclude_none=True))
-            for span_dto in agenta_span_dtos
-        ]
-
-        inline_trace = spans  # turn into Agenta Model ?
+    spans = [
+        loads(span_dto.model_dump_json(exclude_none=True))
+        for span_dto in agenta_span_dtos
+    ]
+    inline_trace = spans  # NOTE (author: @jp-agenta): turn into Agenta Model ?
 
     return inline_trace
 
@@ -1120,98 +1083,6 @@ class LlmTokens(BaseModel):
     total_tokens: Optional[int] = 0
 
 
-class CreateSpan(BaseModel):
-    id: str
-    app_id: str
-    variant_id: Optional[str] = None
-    variant_name: Optional[str] = None
-    inputs: Optional[Dict[str, Optional[Any]]] = None
-    internals: Optional[Dict[str, Optional[Any]]] = None
-    outputs: Optional[Union[str, Dict[str, Optional[Any]], List[Any]]] = None
-    config: Optional[Dict[str, Optional[Any]]] = None
-    environment: Optional[str] = None
-    tags: Optional[List[str]] = None
-    token_consumption: Optional[int] = None
-    name: str
-    parent_span_id: Optional[str] = None
-    attributes: Optional[Dict[str, Optional[Any]]] = None
-    spankind: str
-    status: str
-    user: Optional[str] = None
-    start_time: datetime
-    end_time: datetime
-    tokens: Optional[LlmTokens] = None
-    cost: Optional[float] = None
-
-
-def _parse_to_legacy_span(span: SpanDTO) -> CreateSpan:
-    attributes = None
-    if span.otel:
-        attributes = span.otel.attributes or {}
-
-        if span.otel.events:
-            for event in span.otel.events:
-                if event.name == "exception":
-                    attributes.update(**event.attributes)
-
-    legacy_span = CreateSpan(
-        id=span.node.id.hex[:24],
-        spankind=span.node.type,
-        name=span.node.name,
-        #
-        status=span.status.code.name,
-        #
-        start_time=span.time.start,
-        end_time=span.time.end,
-        #
-        parent_span_id=span.parent.id.hex[:24] if span.parent else None,
-        #
-        inputs=span.data.get("inputs") if span.data else {},
-        internals=span.data.get("internals") if span.data else {},
-        outputs=span.data.get("outputs") if span.data else {},
-        #
-        environment=span.meta.get("environment") if span.meta else None,
-        config=span.meta.get("configuration") if span.meta else None,
-        #
-        tokens=(
-            LlmTokens(
-                prompt_tokens=span.metrics.get("acc", {})
-                .get("tokens", {})
-                .get("prompt", 0.0),
-                completion_tokens=span.metrics.get("acc", {})
-                .get("tokens", {})
-                .get("completion", 0.0),
-                total_tokens=span.metrics.get("acc", {})
-                .get("tokens", {})
-                .get("total", 0.0),
-            )
-            if span.metrics
-            else None
-        ),
-        cost=(
-            span.metrics.get("acc", {}).get("costs", {}).get("total", 0.0)
-            if span.metrics
-            else None
-        ),
-        #
-        app_id=(
-            span.refs.get("application", {}).get("id", "missing-app-id")
-            if span.refs
-            else "missing-app-id"
-        ),
-        #
-        attributes=attributes,
-        #
-        variant_id=None,
-        variant_name=None,
-        tags=None,
-        token_consumption=None,
-        user=None,
-    )
-
-    return legacy_span
-
-
 TYPES_WITH_COSTS = [
     "embedding",
     "query",

agenta-cli/agenta/sdk/types.py

@@ -23,9 +23,9 @@ class LLMTokenUsage(BaseModel):
 
 
 class BaseResponse(BaseModel):
-    version: Optional[str] = "2.0"
+    version: Optional[str] = "3.0"
     data: Optional[Union[str, Dict[str, Any]]]
-    trace: Optional[Dict[str, Any]]
+    trace: Optional[List[Dict[str, Any]]]
 
 
 class DictInput(dict):

agenta-web/src/components/Playground/Views/TestView.tsx

@@ -174,7 +174,7 @@ interface BoxComponentProps {
     additionalData: {
         cost: number | null
         latency: number | null
-        usage: {completion_tokens: number; prompt_tokens: number; total_tokens: number} | null
+        usage: {completion: number; prompt: number; total: number} | null
     }
     onInputParamChange: (paramName: string, newValue: any) => void
     onRun: () => void
@@ -336,7 +336,7 @@ const BoxComponent: React.FC<BoxComponentProps> = ({
             )}
             {additionalData?.cost || additionalData?.latency ? (
                 <Space className="flex items-center gap-3">
-                    <span>Tokens: {formatTokenUsage(additionalData?.usage?.total_tokens)}</span>
+                    <span>Tokens: {formatTokenUsage(additionalData?.usage?.total)}</span>
                     <span>Cost: {formatCurrency(additionalData?.cost)}</span>
                     <span>Latency: {formatLatency(additionalData?.latency)}</span>
                     {traceSpans?.spans?.length && isDemo() && (
@@ -397,19 +397,18 @@ const App: React.FC<TestViewProps> = ({
         Array<{
             cost: number | null
             latency: number | null
-            usage: {completion_tokens: number; prompt_tokens: number; total_tokens: number} | null
+            usage:
+                | {completion?: number; prompt?: number; total: number}
+                | {completion_tokens?: number; prompt_tokens?: number; total_tokens: number}
+                | null
         }>
     >(testList.map(() => ({cost: null, latency: null, usage: null})))
     const [traceSpans, setTraceSpans] = useState<
         | {
               trace_id: string
               cost?: number
               latency?: number
-              usage?: {
-                  completion_tokens: number
-                  prompt_tokens: number
-                  total_tokens: number
-              }
+              usage?: {completion: number; prompt: number; total: number}
               spans?: BaseResponseSpans[]
           }
         | undefined
@@ -567,10 +566,10 @@ const App: React.FC<TestViewProps> = ({
             // Check result type
             // String, FuncResponse or BaseResponse
             if (typeof result === "string") {
-                res = {version: "2.0", data: result} as BaseResponse
+                res = {version: "3.0", data: result} as BaseResponse
                 setResultForIndex(getStringOrJson(res.data), index)
             } else if (isFuncResponse(result)) {
-                res = {version: "2.0", data: result.message} as BaseResponse
+                res = {version: "3.0", data: result.message} as BaseResponse
                 setResultForIndex(getStringOrJson(res.data), index)
 
                 const {message, cost, latency, usage} = result
@@ -583,13 +582,13 @@ const App: React.FC<TestViewProps> = ({
                 res = result as BaseResponse
                 setResultForIndex(getStringOrJson(res.data), index)
 
-                const {data, trace} = result
+                const trace = result.trace
                 setAdditionalDataList((prev) => {
                     const newDataList = [...prev]
                     newDataList[index] = {
-                        cost: trace?.cost || null,
-                        latency: trace?.latency || null,
-                        usage: trace?.usage || null,
+                        cost: trace[0]?.metrics?.acc?.costs?.total || null,
+                        latency: trace[0]?.time?.span / 1_000_000 || null,
+                        usage: trace[0]?.metrics?.acc?.tokens || null,
                     }
                     return newDataList
                 })

agenta-web/src/lib/Types.ts

@@ -493,7 +493,7 @@ export type ComparisonResultRow = {
 export type RequestMetadata = {
     cost: number
     latency: number
-    usage: {completion_tokens?: number; prompt_tokens?: number; total_tokens: number}
+    usage: {completion?: number; prompt?: number; total: number}
 }
 
 export type WithPagination<T> = {
@@ -547,13 +547,7 @@ export type FuncResponse = {
 export type BaseResponse = {
     version: string
     data: string | Record<string, any>
-    trace?: {
-        trace_id: string
-        cost?: number
-        latency?: number
-        usage?: {completion_tokens: number; prompt_tokens: number; total_tokens: number}
-        spans?: BaseResponseSpans[]
-    }
+    trace?: Record<string, any>[]
 }
 
 export type BaseResponseSpans = {

agenta-web/src/lib/helpers/playgroundResp.ts

@@ -5,5 +5,5 @@ export function isFuncResponse(res: any): res is FuncResponse {
 }
 
 export function isBaseResponse(res: any): res is BaseResponse {
-    return res && res?.version === "2.0"
+    return res && res?.version
 }