arize-phoenix 10.14.0__py3-none-any.whl → 11.0.0__py3-none-any.whl

phoenix/server/api/dataloaders/span_costs.py

@@ -0,0 +1,35 @@
+from typing import Optional
+
+from sqlalchemy import select
+from sqlalchemy.orm import joinedload, load_only
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+from phoenix.server.types import DbSessionFactory
+
+SpanID: TypeAlias = int
+Key: TypeAlias = SpanID
+Result: TypeAlias = Optional[models.SpanCost]
+
+
+class SpanCostsDataLoader(DataLoader[Key, Result]):
+    def __init__(self, db: DbSessionFactory) -> None:
+        super().__init__(load_fn=self._load_fn)
+        self._db = db
+
+    async def _load_fn(self, keys: list[Key]) -> list[Result]:
+        span_ids = list(set(keys))
+        async with self._db() as session:
+            costs = {
+                span.id: span.span_cost
+                async for span in await session.stream_scalars(
+                    select(models.Span)
+                    .where(models.Span.id.in_(span_ids))
+                    .options(
+                        load_only(models.Span.id),
+                        joinedload(models.Span.span_cost),
+                    )
+                )
+            }
+        return [costs.get(span_id) for span_id in keys]

phoenix/server/api/dataloaders/types.py

@@ -0,0 +1,29 @@
+from dataclasses import dataclass, field
+from functools import cached_property
+from typing import Optional
+
+
+@dataclass(frozen=True)
+class CostBreakdown:
+    tokens: Optional[float] = None
+    cost: Optional[float] = None
+
+    @cached_property
+    def cost_per_token(self) -> Optional[float]:
+        if self.tokens and self.cost:
+            return self.cost / self.tokens
+        return None
+
+
+@dataclass(frozen=True)
+class SpanCostSummary:
+    prompt: CostBreakdown = field(default_factory=CostBreakdown)
+    completion: CostBreakdown = field(default_factory=CostBreakdown)
+    total: CostBreakdown = field(default_factory=CostBreakdown)
+
+
+@dataclass(frozen=True)
+class SpanCostDetailSummaryEntry:
+    token_type: str
+    is_prompt: bool
+    value: CostBreakdown = field(default_factory=CostBreakdown)