feat: model distribution (InftyAI#189)

* update helm chart

Signed-off-by: kerthcet <kerthcet@gmail.com>

* add artifacts

Signed-off-by: kerthcet <kerthcet@gmail.com>

* optimize

Signed-off-by: kerthcet <kerthcet@gmail.com>

* add artifacts page

Signed-off-by: kerthcet <kerthcet@gmail.com>

* add model distribution

Signed-off-by: kerthcet <kerthcet@gmail.com>

* fix error

Signed-off-by: kerthcet <kerthcet@gmail.com>

* fix lint

Signed-off-by: kerthcet <kerthcet@gmail.com>

---------

Signed-off-by: kerthcet <kerthcet@gmail.com>

Author: kerthcet

alphatrion/server/graphql/resolvers.py

@@ -24,6 +24,7 @@
     GraphQLStatusEnum,
     Label,
     Metric,
+    ModelDistribution,
     RemoveUserFromTeamInput,
     Run,
     Span,
@@ -263,12 +264,28 @@ def aggregate_team_tokens(team_id: strawberry.ID) -> dict[str, int]:
             return {"total_tokens": 0, "input_tokens": 0, "output_tokens": 0}
 
         trace_store = runtime.storage_runtime().tracestore
-        result = trace_store.get_llm_spans_by_team_id(team_id=team_id)
-        # get_llm_spans_by_team_id returns a list with one dict
+        result = trace_store.get_llm_tokens_by_team_id(team_id=team_id)
+        # get_llm_tokens_by_team_id returns a list with one dict
         if result and len(result) > 0:
             return result[0]
         return {"total_tokens": 0, "input_tokens": 0, "output_tokens": 0}
 
+    @staticmethod
+    def aggregate_model_distributions(
+        team_id: strawberry.ID,
+    ) -> list[ModelDistribution]:
+        from alphatrion import envs
+
+        if os.getenv(envs.ENABLE_TRACING, "false").lower() != "true":
+            return []
+
+        trace_store = runtime.storage_runtime().tracestore
+        result = trace_store.get_model_distributions_by_team_id(team_id=team_id)
+        return [
+            ModelDistribution(model=item["model"], count=item["count"])
+            for item in result
+        ]
+
     @staticmethod
     def list_exps_by_timeframe(
         team_id: strawberry.ID,

alphatrion/server/graphql/types.py

@@ -13,6 +13,12 @@ class TokenStats:
     output_tokens: int
 
 
+@strawberry.type
+class ModelDistribution:
+    model: str
+    count: int
+
+
 @strawberry.type
 class Team:
     id: strawberry.ID
@@ -45,6 +51,12 @@ def aggregated_tokens(self) -> TokenStats:
             output_tokens=token_data["output_tokens"],
         )
 
+    @strawberry.field
+    def model_distributions(self) -> list["ModelDistribution"]:
+        from .resolvers import GraphQLResolvers
+
+        return GraphQLResolvers.aggregate_model_distributions(team_id=self.id)
+
     @strawberry.field
     def exps_by_timeframe(
         self, start_time: datetime, end_time: datetime

alphatrion/storage/tracestore.py

@@ -331,7 +331,7 @@ def get_llm_spans_by_exp_id(self, exp_id: uuid.UUID) -> list[dict[str, Any]]:
                 logger.error(f"Failed to get spans by exp_id: {e}")
                 return []
 
-    def get_llm_spans_by_team_id(self, team_id: uuid.UUID) -> list[dict[str, Any]]:
+    def get_llm_tokens_by_team_id(self, team_id: uuid.UUID) -> list[dict[str, Any]]:
         """Get all LLM spans for a specific team_id.
 
         Args:
@@ -365,6 +365,46 @@ def get_llm_spans_by_team_id(self, team_id: uuid.UUID) -> list[dict[str, Any]]:
                 logger.error(f"Failed to get daily token usage: {e}")
                 return []
 
+    def get_model_distributions_by_team_id(
+        self, team_id: uuid.UUID
+    ) -> list[dict[str, Any]]:
+        """Get model distribution (count of requests per model) for a specific team.
+
+        Args:
+            team_id: The team ID to filter by
+
+        Returns:
+            List of dicts with keys: model, count
+        """
+        with self._lock:
+            try:
+                query = f"""
+                SELECT
+                    coalesce(
+                        SpanAttributes['gen_ai.request.model'],
+                        SpanAttributes['gen_ai.response.model'],
+                        'unknown'
+                    ) as model,
+                    COUNT(*) as count
+                FROM {self.database}.otel_spans
+                WHERE TeamId = '{team_id}'
+                  AND SemanticKind = 'llm'
+                GROUP BY model
+                ORDER BY count DESC
+                """
+
+                result = self.client.query(query)
+                return [
+                    {
+                        "model": row["model"],
+                        "count": int(row["count"]),
+                    }
+                    for row in result.named_results()
+                ]
+            except Exception as e:
+                logger.error(f"Failed to get model distributions: {e}")
+                return []
+
     def get_daily_token_usage(
         self, team_id: uuid.UUID, days: int = 30
     ) -> list[dict[str, Any]]: