InftyAI
diff --git a/‎.env.example‎
Lines changed: 1 addition & 0 deletions b/‎.env.example‎
Lines changed: 1 addition & 0 deletions
diff --git a/‎.env.integration-test‎
Lines changed: 2 additions & 1 deletion b/‎.env.integration-test‎
Lines changed: 2 additions & 1 deletion
diff --git a/‎alphatrion/envs.py‎
Lines changed: 1 addition & 1 deletion b/‎alphatrion/envs.py‎
Lines changed: 1 addition & 1 deletion
diff --git a/‎alphatrion/server/graphql/resolvers.py‎
Lines changed: 73 additions & 5 deletions b/‎alphatrion/server/graphql/resolvers.py‎
Lines changed: 73 additions & 5 deletions
diff --git a/‎alphatrion/server/graphql/types.py‎
Lines changed: 14 additions & 0 deletions b/‎alphatrion/server/graphql/types.py‎
Lines changed: 14 additions & 0 deletions
diff --git a/‎alphatrion/storage/runtime.py‎
Lines changed: 11 additions & 1 deletion b/‎alphatrion/storage/runtime.py‎
Lines changed: 11 additions & 1 deletion
diff --git a/‎alphatrion/storage/sqlstore.py‎
Lines changed: 11 additions & 0 deletions b/‎alphatrion/storage/sqlstore.py‎
Lines changed: 11 additions & 0 deletions
@@ -15,3 +15,4 @@ ALPHATRION_ARTIFACT_INSECURE=false
 # Tracing configurations
 ALPHATRION_ENABLE_TRACING=true
 ALPHATRION_CLICKHOUSE_INIT_TABLES=true
+ALPHATRION_CLICKHOUSE_ENABLE_BATCH=true
@@ -5,4 +5,5 @@ ALPHATRION_ARTIFACT_INSECURE=true
 ALPHATRION_LOG_LEVEL=INFO
 ALPHATRION_AUTO_CLEANUP=true
 ALPHATRION_ENABLE_TRACING=true
-ALPHATRION_CLICKHOUSE_INIT_TABLES=true
+ALPHATRION_CLICKHOUSE_INIT_TABLES=true
+ALPHATRION_CLICKHOUSE_ENABLE_BATCH=true
@@ -14,7 +14,7 @@
 CLICKHOUSE_DATABASE = "ALPHATRION_CLICKHOUSE_DATABASE"
 CLICKHOUSE_USERNAME = "ALPHATRION_CLICKHOUSE_USERNAME"
 CLICKHOUSE_PASSWORD = "ALPHATRION_CLICKHOUSE_PASSWORD"
-INIT_CLICKHOUSE_TABLES = "ALPHATRION_INIT_CLICKHOUSE_TABLES"
+CLICKHOUSE_ENABLE_BATCH = "ALPHATRION_CLICKHOUSE_ENABLE_BATCH"
 
 # Dashboard only related envs
 DASHBOARD_USER_ID = "ALPHATRION_DASHBOARD_USER_ID"
 
@@ -220,13 +220,23 @@ def get_run(id: strawberry.ID) -> Run | None:
         metadb = runtime.storage_runtime().metadb
         run = metadb.get_run(run_id=uuid.UUID(id))
         if run:
+            meta = run.meta or {}
+
+            # Aggregate and cache tokens for completed runs.
+            # It could be slow for the first time.
+            if Status(run.status) == Status.COMPLETED and "total_tokens" not in meta:
+                token_data = GraphQLResolvers.aggregate_run_tokens(run_id=id)
+                if token_data["total_tokens"] > 0:
+                    meta.update(token_data)
+                    metadb.update_run(run_id=uuid.UUID(id), meta=meta)
+
             return Run(
                 id=run.uuid,
                 team_id=run.team_id,
                 user_id=run.user_id,
                 project_id=run.project_id,
                 experiment_id=run.experiment_id,
-                meta=run.meta,
+                meta=meta,
                 status=GraphQLStatusEnum[Status(run.status).name],
                 created_at=run.created_at,
             )
@@ -250,6 +260,24 @@ def list_exp_metrics(experiment_id: strawberry.ID) -> list[Metric]:
             for m in metrics
         ]
 
+    @staticmethod
+    def list_run_metrics(run_id: strawberry.ID) -> list[Metric]:
+        metadb = runtime.storage_runtime().metadb
+        metrics = metadb.list_metrics_by_run_id(run_id=run_id)
+        return [
+            Metric(
+                id=m.uuid,
+                key=m.key,
+                value=m.value,
+                team_id=m.team_id,
+                project_id=m.project_id,
+                experiment_id=m.experiment_id,
+                run_id=m.run_id,
+                created_at=m.created_at,
+            )
+            for m in metrics
+        ]
+
     @staticmethod
     def total_projects(team_id: strawberry.ID) -> int:
         metadb = runtime.storage_runtime().metadb
@@ -373,8 +401,48 @@ async def get_artifact_content(
             raise RuntimeError(f"Failed to get artifact content: {e}") from e
 
     @staticmethod
-    def list_traces(run_id: strawberry.ID) -> list[Span]:
-        """List all traces/spans for a specific run."""
+    def aggregate_run_tokens(run_id: strawberry.ID) -> dict[str, int]:
+        """Aggregate token usage from all traces for a run."""
+        from alphatrion import envs
+
+        # Check if tracing is enabled
+        if os.getenv(envs.ENABLE_TRACING, "false").lower() != "true":
+            return {"total_tokens": 0, "input_tokens": 0, "output_tokens": 0}
+
+        try:
+            trace_store = runtime.storage_runtime().tracestore
+            spans = trace_store.get_spans_by_run_id(uuid.UUID(run_id))
+            trace_store.close()
+
+            total_tokens = 0
+            input_tokens = 0
+            output_tokens = 0
+
+            for span in spans:
+                span_attrs = span.get("SpanAttributes", {})
+
+                # Aggregate tokens from LLM spans
+                if "llm.usage.total_tokens" in span_attrs:
+                    total_tokens += int(span_attrs["llm.usage.total_tokens"])
+                if "gen_ai.usage.input_tokens" in span_attrs:
+                    input_tokens += int(span_attrs["gen_ai.usage.input_tokens"])
+                if "gen_ai.usage.output_tokens" in span_attrs:
+                    output_tokens += int(span_attrs["gen_ai.usage.output_tokens"])
+
+            return {
+                "total_tokens": total_tokens,
+                "input_tokens": input_tokens,
+                "output_tokens": output_tokens,
+            }
+        except Exception as e:
+            import logging
+
+            logging.error(f"Failed to aggregate tokens for run {run_id}: {e}")
+            return {"total_tokens": 0, "input_tokens": 0, "output_tokens": 0}
+
+    @staticmethod
+    def list_spans(run_id: strawberry.ID) -> list[Span]:
+        """List all spans for a specific run."""
         from alphatrion import envs
 
         # Check if tracing is enabled
@@ -385,12 +453,12 @@ def list_traces(run_id: strawberry.ID) -> list[Span]:
             trace_store = runtime.storage_runtime().tracestore
 
             # Get traces from ClickHouse
-            traces = trace_store.get_traces_by_run_id(uuid.UUID(run_id))
+            raw_spans = trace_store.get_spans_by_run_id(uuid.UUID(run_id))
             trace_store.close()
 
             # Convert to GraphQL Span objects
             spans = []
-            for t in traces:
+            for t in raw_spans:
                 # Convert events
                 events = []
                 if t.get("Events"):
 
@@ -129,6 +129,20 @@ class Run:
     status: GraphQLStatusEnum
     created_at: datetime
 
+    @strawberry.field
+    def metrics(self) -> list["Metric"]:
+        """Get metrics for this run."""
+        from alphatrion.server.graphql.resolvers import GraphQLResolvers
+
+        return GraphQLResolvers.list_run_metrics(run_id=self.id)
+
+    @strawberry.field
+    def spans(self) -> list["Span"]:
+        """Get spans for this run."""
+        from alphatrion.server.graphql.resolvers import GraphQLResolvers
+
+        return GraphQLResolvers.list_spans(run_id=str(self.id))
+
 
 @strawberry.type
 class Metric:
 
@@ -2,6 +2,7 @@
 import os
 
 from opentelemetry import trace
+from opentelemetry.sdk.trace import TracerProvider
 from traceloop.sdk import Traceloop
 
 from alphatrion import envs
@@ -38,10 +39,13 @@ def __init__(self):
                 == "true",
             )
 
+            enable_batch = (
+                os.getenv(envs.CLICKHOUSE_ENABLE_BATCH, "true").lower() == "true"
+            )
             Traceloop.init(
                 app_name="alphatrion",
                 exporter=ClickHouseSpanExporter(self.tracestore),
-                disable_batch=False,  # Enable batching
+                disable_batch=not enable_batch,
                 telemetry_enabled=False,
             )
 
@@ -60,6 +64,12 @@ def metadb(self):
     def tracestore(self):
         return self._tracestore
 
+    def flush(self):
+        if self._tracestore:
+            tracer_provider = trace.get_tracer_provider()
+            if isinstance(tracer_provider, TracerProvider):
+                tracer_provider.force_flush(timeout_millis=5000)
+
 
 def init():
     """
 
@@ -695,3 +695,14 @@ def list_metrics_by_experiment_id(self, experiment_id: uuid.UUID) -> list[Metric
         )
         session.close()
         return metrics
+
+    def list_metrics_by_run_id(self, run_id: uuid.UUID) -> list[Metric]:
+        session = self._session()
+        metrics = (
+            session.query(Metric)
+            .filter(Metric.run_id == run_id)
+            .order_by(Metric.created_at.asc())
+            .all()
+        )
+        session.close()
+        return metrics