fix(persistence): postgres json calculations (#2848)

src/phoenix/server/api/input_types/SpanSort.py

@@ -5,7 +5,7 @@
 import pandas as pd
 import strawberry
 from openinference.semconv.trace import SpanAttributes
-from sqlalchemy import Integer, cast, desc, nulls_last
+from sqlalchemy import desc, nulls_last
 from strawberry import UNSET
 from typing_extensions import assert_never
 
@@ -37,15 +37,9 @@ class SpanColumn(Enum):
     SpanColumn.startTime: models.Span.start_time,
     SpanColumn.endTime: models.Span.end_time,
     SpanColumn.latencyMs: models.Span.latency_ms,
-    SpanColumn.tokenCountTotal: cast(
-        models.Span.attributes[LLM_TOKEN_COUNT_TOTAL].as_string(), Integer
-    ),
-    SpanColumn.tokenCountPrompt: cast(
-        models.Span.attributes[LLM_TOKEN_COUNT_PROMPT].as_string(), Integer
-    ),
-    SpanColumn.tokenCountCompletion: cast(
-        models.Span.attributes[LLM_TOKEN_COUNT_COMPLETION].as_string(), Integer
-    ),
+    SpanColumn.tokenCountTotal: models.Span.attributes[LLM_TOKEN_COUNT_TOTAL].as_float(),
+    SpanColumn.tokenCountPrompt: models.Span.attributes[LLM_TOKEN_COUNT_PROMPT].as_float(),
+    SpanColumn.tokenCountCompletion: models.Span.attributes[LLM_TOKEN_COUNT_COMPLETION].as_float(),
     SpanColumn.cumulativeTokenCountTotal: models.Span.cumulative_llm_token_count_prompt
     + models.Span.cumulative_llm_token_count_completion,
     SpanColumn.cumulativeTokenCountPrompt: models.Span.cumulative_llm_token_count_prompt,

src/phoenix/server/api/types/Project.py

@@ -3,7 +3,7 @@
 
 import strawberry
 from openinference.semconv.trace import SpanAttributes
-from sqlalchemy import Integer, and_, cast, func, select
+from sqlalchemy import and_, func, select
 from sqlalchemy.orm import contains_eager
 from sqlalchemy.sql.functions import coalesce
 from strawberry import ID, UNSET
@@ -116,13 +116,10 @@ async def token_count_total(
         info: Info[Context, None],
         time_range: Optional[TimeRange] = UNSET,
     ) -> int:
-        prompt = models.Span.attributes[LLM_TOKEN_COUNT_PROMPT]
-        completion = models.Span.attributes[LLM_TOKEN_COUNT_COMPLETION]
+        prompt = models.Span.attributes[LLM_TOKEN_COUNT_PROMPT].as_float()
+        completion = models.Span.attributes[LLM_TOKEN_COUNT_COMPLETION].as_float()
         stmt = (
-            select(
-                coalesce(func.sum(cast(prompt, Integer)), 0)
-                + coalesce(func.sum(cast(completion, Integer)), 0)
-            )
+            select(coalesce(func.sum(prompt), 0) + coalesce(func.sum(completion), 0))
             .join(models.Trace)
             .join(models.Project)
             .where(models.Project.name == self.name)