arize-phoenix 3.16.1__py3-none-any.whl → 7.7.0__py3-none-any.whl

phoenix/server/api/dataloaders/session_token_usages.py

@@ -0,0 +1,41 @@
+from sqlalchemy import func, select
+from sqlalchemy.sql.functions import coalesce
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+from phoenix.server.types import DbSessionFactory
+from phoenix.trace.schemas import TokenUsage
+
+Key: TypeAlias = int
+Result: TypeAlias = TokenUsage
+
+
+class SessionTokenUsagesDataLoader(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]:
+        stmt = (
+            select(
+                models.Trace.project_session_rowid.label("id_"),
+                func.sum(coalesce(models.Span.cumulative_llm_token_count_prompt, 0)).label(
+                    "prompt"
+                ),
+                func.sum(coalesce(models.Span.cumulative_llm_token_count_completion, 0)).label(
+                    "completion"
+                ),
+            )
+            .join_from(models.Span, models.Trace)
+            .where(models.Span.parent_id.is_(None))
+            .where(models.Trace.project_session_rowid.in_(keys))
+            .group_by(models.Trace.project_session_rowid)
+        )
+        async with self._db() as session:
+            result: dict[Key, TokenUsage] = {
+                id_: TokenUsage(prompt=prompt, completion=completion)
+                async for id_, prompt, completion in await session.stream(stmt)
+                if id_ is not None
+            }
+        return [result.get(key, TokenUsage()) for key in keys]

phoenix/server/api/dataloaders/session_trace_latency_ms_quantile.py

@@ -0,0 +1,55 @@
+from collections import defaultdict
+from typing import Optional
+
+import numpy as np
+from aioitertools.itertools import groupby
+from sqlalchemy import select
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+from phoenix.server.types import DbSessionFactory
+
+SessionId: TypeAlias = int
+Probability: TypeAlias = float
+QuantileValue: TypeAlias = float
+
+Key: TypeAlias = tuple[SessionId, Probability]
+Result: TypeAlias = Optional[QuantileValue]
+ResultPosition: TypeAlias = int
+
+DEFAULT_VALUE: Result = None
+
+
+class SessionTraceLatencyMsQuantileDataLoader(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]:
+        results: list[Result] = [DEFAULT_VALUE] * len(keys)
+        argument_position_map: defaultdict[
+            SessionId, defaultdict[Probability, list[ResultPosition]]
+        ] = defaultdict(lambda: defaultdict(list))
+        session_rowids = {session_id for session_id, _ in keys}
+        for position, (session_id, probability) in enumerate(keys):
+            argument_position_map[session_id][probability].append(position)
+        stmt = (
+            select(
+                models.Trace.project_session_rowid,
+                models.Trace.latency_ms,
+            )
+            .where(models.Trace.project_session_rowid.in_(session_rowids))
+            .order_by(models.Trace.project_session_rowid)
+        )
+        async with self._db() as session:
+            data = await session.stream(stmt)
+            async for project_session_rowid, group in groupby(
+                data, lambda row: row.project_session_rowid
+            ):
+                session_latencies = [row.latency_ms for row in group]
+                for probability, positions in argument_position_map[project_session_rowid].items():
+                    quantile_value = np.quantile(session_latencies, probability)
+                    for position in positions:
+                        results[position] = quantile_value
+        return results

phoenix/server/api/dataloaders/span_annotations.py

@@ -0,0 +1,26 @@
+from collections import defaultdict
+
+from sqlalchemy import select
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db.models import SpanAnnotation as ORMSpanAnnotation
+from phoenix.server.types import DbSessionFactory
+
+Key: TypeAlias = int
+Result: TypeAlias = list[ORMSpanAnnotation]
+
+
+class SpanAnnotationsDataLoader(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_annotations_by_id: defaultdict[Key, Result] = defaultdict(list)
+        async with self._db() as session:
+            async for span_annotation in await session.stream_scalars(
+                select(ORMSpanAnnotation).where(ORMSpanAnnotation.span_rowid.in_(keys))
+            ):
+                span_annotations_by_id[span_annotation.span_rowid].append(span_annotation)
+        return [span_annotations_by_id[key] for key in keys]

phoenix/server/api/dataloaders/span_dataset_examples.py

@@ -0,0 +1,31 @@
+from sqlalchemy import select
+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 = list[models.DatasetExample]
+
+
+class SpanDatasetExamplesDataLoader(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_rowids = keys
+        async with self._db() as session:
+            dataset_examples: dict[Key, list[models.DatasetExample]] = {
+                span_rowid: [] for span_rowid in span_rowids
+            }
+            async for span_rowid, dataset_example in await session.stream(
+                select(models.Span.id, models.DatasetExample)
+                .select_from(models.Span)
+                .join(models.DatasetExample, models.DatasetExample.span_rowid == models.Span.id)
+                .where(models.Span.id.in_(span_rowids))
+            ):
+                dataset_examples[span_rowid].append(dataset_example)
+        return [dataset_examples.get(span_rowid, []) for span_rowid in keys]

phoenix/server/api/dataloaders/span_descendants.py

@@ -0,0 +1,57 @@
+from random import randint
+
+from aioitertools.itertools import groupby
+from sqlalchemy import select
+from sqlalchemy.orm import joinedload
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+from phoenix.server.types import DbSessionFactory
+
+SpanId: TypeAlias = str
+
+Key: TypeAlias = SpanId
+Result: TypeAlias = list[models.Span]
+
+
+class SpanDescendantsDataLoader(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]:
+        root_ids = set(keys)
+        root_id_label = f"root_id_{randint(0, 10**6):06}"
+        descendant_ids = (
+            select(
+                models.Span.id,
+                models.Span.span_id,
+                models.Span.parent_id.label(root_id_label),
+            )
+            .where(models.Span.parent_id.in_(root_ids))
+            .cte(recursive=True)
+        )
+        parent_ids = descendant_ids.alias()
+        descendant_ids = descendant_ids.union_all(
+            select(
+                models.Span.id,
+                models.Span.span_id,
+                parent_ids.c[root_id_label],
+            ).join(
+                parent_ids,
+                models.Span.parent_id == parent_ids.c.span_id,
+            )
+        )
+        stmt = (
+            select(descendant_ids.c[root_id_label], models.Span)
+            .join(descendant_ids, models.Span.id == descendant_ids.c.id)
+            .options(joinedload(models.Span.trace, innerjoin=True).load_only(models.Trace.trace_id))
+            .order_by(descendant_ids.c[root_id_label])
+        )
+        results: dict[SpanId, Result] = {key: [] for key in keys}
+        async with self._db() as session:
+            data = await session.stream(stmt)
+            async for root_id, group in groupby(data, key=lambda d: d[0]):
+                results[root_id].extend(span for _, span in group)
+        return [results[key].copy() for key in keys]

phoenix/server/api/dataloaders/span_projects.py

@@ -0,0 +1,33 @@
+from typing import Union
+
+from sqlalchemy import select
+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 = models.Project
+
+
+class SpanProjectsDataLoader(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[Union[Result, ValueError]]:
+        span_ids = list(set(keys))
+        async with self._db() as session:
+            projects = {
+                span_id: project
+                async for span_id, project in await session.stream(
+                    select(models.Span.id, models.Project)
+                    .select_from(models.Span)
+                    .join(models.Trace, models.Span.trace_rowid == models.Trace.id)
+                    .join(models.Project, models.Trace.project_rowid == models.Project.id)
+                    .where(models.Span.id.in_(span_ids))
+                )
+            }
+        return [projects.get(span_id) or ValueError("Invalid span ID") for span_id in keys]

phoenix/server/api/dataloaders/token_counts.py

@@ -0,0 +1,124 @@
+from collections import defaultdict
+from datetime import datetime
+from typing import Any, Literal, Optional
+
+from cachetools import LFUCache, TTLCache
+from sqlalchemy import Select, func, select
+from sqlalchemy.sql.functions import coalesce
+from strawberry.dataloader import AbstractCache, DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+from phoenix.server.api.dataloaders.cache import TwoTierCache
+from phoenix.server.api.input_types.TimeRange import TimeRange
+from phoenix.server.types import DbSessionFactory
+from phoenix.trace.dsl import SpanFilter
+
+Kind: TypeAlias = Literal["prompt", "completion", "total"]
+ProjectRowId: TypeAlias = int
+TimeInterval: TypeAlias = tuple[Optional[datetime], Optional[datetime]]
+FilterCondition: TypeAlias = Optional[str]
+TokenCount: TypeAlias = int
+
+Segment: TypeAlias = tuple[TimeInterval, FilterCondition]
+Param: TypeAlias = tuple[ProjectRowId, Kind]
+
+Key: TypeAlias = tuple[Kind, ProjectRowId, Optional[TimeRange], FilterCondition]
+Result: TypeAlias = TokenCount
+ResultPosition: TypeAlias = int
+DEFAULT_VALUE: Result = 0
+
+
+def _cache_key_fn(key: Key) -> tuple[Segment, Param]:
+    kind, project_rowid, time_range, filter_condition = key
+    interval = (
+        (time_range.start, time_range.end) if isinstance(time_range, TimeRange) else (None, None)
+    )
+    return (interval, filter_condition), (project_rowid, kind)
+
+
+_Section: TypeAlias = ProjectRowId
+_SubKey: TypeAlias = tuple[TimeInterval, FilterCondition, Kind]
+
+
+class TokenCountCache(
+    TwoTierCache[Key, Result, _Section, _SubKey],
+):
+    def __init__(self) -> None:
+        super().__init__(
+            # TTL=3600 (1-hour) because time intervals are always moving forward, but
+            # interval endpoints are rounded down to the hour by the UI, so anything
+            # older than an hour most likely won't be a cache-hit anyway.
+            main_cache=TTLCache(maxsize=64, ttl=3600),
+            sub_cache_factory=lambda: LFUCache(maxsize=2 * 2 * 3),
+        )
+
+    def _cache_key(self, key: Key) -> tuple[_Section, _SubKey]:
+        (interval, filter_condition), (project_rowid, kind) = _cache_key_fn(key)
+        return project_rowid, (interval, filter_condition, kind)
+
+
+class TokenCountDataLoader(DataLoader[Key, Result]):
+    def __init__(
+        self,
+        db: DbSessionFactory,
+        cache_map: Optional[AbstractCache[Key, Result]] = None,
+    ) -> None:
+        super().__init__(
+            load_fn=self._load_fn,
+            cache_key_fn=_cache_key_fn,
+            cache_map=cache_map,
+        )
+        self._db = db
+
+    async def _load_fn(self, keys: list[Key]) -> list[Result]:
+        results: list[Result] = [DEFAULT_VALUE] * len(keys)
+        arguments: defaultdict[
+            Segment,
+            defaultdict[Param, list[ResultPosition]],
+        ] = defaultdict(lambda: defaultdict(list))
+        for position, key in enumerate(keys):
+            segment, param = _cache_key_fn(key)
+            arguments[segment][param].append(position)
+        async with self._db() as session:
+            for segment, params in arguments.items():
+                stmt = _get_stmt(segment, *params.keys())
+                data = await session.stream(stmt)
+                async for project_rowid, prompt, completion, total in data:
+                    for position in params[(project_rowid, "prompt")]:
+                        results[position] = prompt
+                    for position in params[(project_rowid, "completion")]:
+                        results[position] = completion
+                    for position in params[(project_rowid, "total")]:
+                        results[position] = total
+        return results
+
+
+def _get_stmt(
+    segment: Segment,
+    *params: Param,
+) -> Select[Any]:
+    (start_time, end_time), filter_condition = segment
+    prompt = coalesce(func.sum(models.Span.llm_token_count_prompt), 0)
+    completion = coalesce(func.sum(models.Span.llm_token_count_completion), 0)
+    total = prompt + completion
+    pid = models.Trace.project_rowid
+    stmt: Select[Any] = (
+        select(
+            pid,
+            prompt.label("prompt"),
+            completion.label("completion"),
+            total.label("total"),
+        )
+        .join_from(models.Trace, models.Span)
+        .group_by(pid)
+    )
+    if start_time:
+        stmt = stmt.where(start_time <= models.Span.start_time)
+    if end_time:
+        stmt = stmt.where(models.Span.start_time < end_time)
+    if filter_condition:
+        sf = SpanFilter(filter_condition)
+        stmt = sf(stmt)
+    stmt = stmt.where(pid.in_([rowid for rowid, _ in params]))
+    return stmt

phoenix/server/api/dataloaders/trace_by_trace_ids.py

@@ -0,0 +1,25 @@
+from typing import List, Optional
+
+from sqlalchemy import select
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+from phoenix.server.types import DbSessionFactory
+
+Key: TypeAlias = str
+Result: TypeAlias = Optional[models.Trace]
+
+
+class TraceByTraceIdsDataLoader(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]:
+        stmt = select(models.Trace).where(models.Trace.trace_id.in_(keys))
+        async with self._db() as session:
+            result: dict[Key, models.Trace] = {
+                trace.trace_id: trace async for trace in await session.stream_scalars(stmt)
+            }
+        return [result.get(trace_id) for trace_id in keys]

phoenix/server/api/dataloaders/trace_root_spans.py

@@ -0,0 +1,32 @@
+from typing import List, Optional
+
+from sqlalchemy import select
+from sqlalchemy.orm import contains_eager
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+from phoenix.server.types import DbSessionFactory
+
+Key: TypeAlias = int
+Result: TypeAlias = Optional[models.Span]
+
+
+class TraceRootSpansDataLoader(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]:
+        stmt = (
+            select(models.Span)
+            .join(models.Trace)
+            .where(models.Span.parent_id.is_(None))
+            .where(models.Trace.id.in_(keys))
+            .options(contains_eager(models.Span.trace).load_only(models.Trace.trace_id))
+        )
+        async with self._db() as session:
+            result: dict[Key, models.Span] = {
+                span.trace_rowid: span async for span in await session.stream_scalars(stmt)
+            }
+        return [result.get(key) for key in keys]

phoenix/server/api/dataloaders/user_roles.py

@@ -0,0 +1,30 @@
+from collections import defaultdict
+from typing import Optional
+
+from sqlalchemy import select
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+from phoenix.server.types import DbSessionFactory
+
+UserRoleId: TypeAlias = int
+Key: TypeAlias = UserRoleId
+Result: TypeAlias = Optional[models.UserRole]
+
+
+class UserRolesDataLoader(DataLoader[Key, Result]):
+    """DataLoader that batches together user roles by their ids."""
+
+    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]:
+        user_roles_by_id: defaultdict[Key, Result] = defaultdict(None)
+        async with self._db() as session:
+            data = await session.stream_scalars(select(models.UserRole))
+            async for user_role in data:
+                user_roles_by_id[user_role.id] = user_role
+
+        return [user_roles_by_id.get(role_id) for role_id in keys]

phoenix/server/api/dataloaders/users.py

@@ -0,0 +1,33 @@
+from collections import defaultdict
+from typing import Optional
+
+from sqlalchemy import select
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+from phoenix.server.types import DbSessionFactory
+
+UserId: TypeAlias = int
+Key: TypeAlias = UserId
+Result: TypeAlias = Optional[models.User]
+
+
+class UsersDataLoader(DataLoader[Key, Result]):
+    """DataLoader that batches together users by their ids."""
+
+    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]:
+        user_ids = list(set(keys))
+        users_by_id: defaultdict[Key, Result] = defaultdict(None)
+        async with self._db() as session:
+            data = await session.stream_scalars(
+                select(models.User).where(models.User.id.in_(user_ids))
+            )
+            async for user in data:
+                users_by_id[user.id] = user
+
+        return [users_by_id.get(user_id) for user_id in keys]

phoenix/server/api/exceptions.py

@@ -0,0 +1,48 @@
+from graphql.error import GraphQLError
+from strawberry.extensions import MaskErrors
+
+
+class CustomGraphQLError(Exception):
+    """
+    An error that represents an expected error scenario in a GraphQL resolver.
+    """
+
+
+class BadRequest(CustomGraphQLError):
+    """
+    An error raised due to a malformed or invalid request.
+    """
+
+
+class NotFound(CustomGraphQLError):
+    """
+    An error raised when the requested resource is not found.
+    """
+
+
+class Unauthorized(CustomGraphQLError):
+    """
+    An error raised when login fails or a user or other entity is not authorized
+    to access a resource.
+    """
+
+
+class Conflict(CustomGraphQLError):
+    """
+    An error raised when a mutation cannot be completed due to a conflict with
+    the current state of one or more resources.
+    """
+
+
+def get_mask_errors_extension() -> MaskErrors:
+    return MaskErrors(
+        should_mask_error=_should_mask_error,
+        error_message="an unexpected error occurred",
+    )
+
+
+def _should_mask_error(error: GraphQLError) -> bool:
+    """
+    Masks unexpected errors raised from GraphQL resolvers.
+    """
+    return not isinstance(error.original_error, CustomGraphQLError)

phoenix/server/api/helpers/__init__.py

@@ -0,0 +1,12 @@
+from collections.abc import Iterable
+from typing import Optional, TypeVar
+
+T = TypeVar("T")
+
+
+def ensure_list(obj: Optional[Iterable[T]]) -> list[T]:
+    if isinstance(obj, list):
+        return obj
+    if isinstance(obj, Iterable):
+        return list(obj)
+    return []