fix(sdk): Fixed token propagation in spans

Agenta-AI · Jun 6, 2024 · 12bde7c · 12bde7c
1 parent a0cf1e8
commit 12bde7c
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/agenta-cli/agenta/sdk/tracing/llm_tracing.py b/agenta-cli/agenta/sdk/tracing/llm_tracing.py
@@ -176,6 +176,8 @@ def _update_span_cost(self, span: CreateSpan, cost: Optional[float]):
                 span.cost += cost
 
     def _update_span_tokens(self, span: CreateSpan, tokens: Optional[dict]):
+        if isinstance(tokens, LlmTokens):
+            tokens = tokens.dict()
         if tokens is not None and isinstance(tokens, dict):
             if span.tokens is None:
                 span.tokens = LlmTokens(**tokens)