arize-phoenix 8.0.1__py3-none-any.whl → 8.2.0__py3-none-any.whl

phoenix/server/api/dataloaders/span_descendants.py

@@ -1,18 +1,18 @@
-from random import randint
+from secrets import token_hex
+from typing import Iterable
 
 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
+SpanRowId: TypeAlias = int
 
-Key: TypeAlias = SpanId
-Result: TypeAlias = list[models.Span]
+Key: TypeAlias = SpanRowId
+Result: TypeAlias = list[SpanRowId]
 
 
 class SpanDescendantsDataLoader(DataLoader[Key, Result]):
@@ -20,16 +20,25 @@ class SpanDescendantsDataLoader(DataLoader[Key, Result]):
         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}"
+    async def _load_fn(self, keys: Iterable[Key]) -> list[Result]:
+        root_ids = (
+            select(
+                models.Span.span_id,
+                models.Span.id,
+            )
+            .where(models.Span.id.in_(set(keys)))
+            .subquery()
+        )
+        root_id_label = f"root_id_{token_hex(8)}"
+        root_rowid_label = f"root_rowid_{token_hex(8)}"
         descendant_ids = (
             select(
                 models.Span.id,
                 models.Span.span_id,
                 models.Span.parent_id.label(root_id_label),
+                root_ids.c.id.label(root_rowid_label),
             )
-            .where(models.Span.parent_id.in_(root_ids))
+            .join(root_ids, models.Span.parent_id == root_ids.c.span_id)
             .cte(recursive=True)
         )
         parent_ids = descendant_ids.alias()
@@ -38,20 +47,20 @@ class SpanDescendantsDataLoader(DataLoader[Key, Result]):
                 models.Span.id,
                 models.Span.span_id,
                 parent_ids.c[root_id_label],
+                parent_ids.c[root_rowid_label],
             ).join(
                 parent_ids,
                 models.Span.parent_id == parent_ids.c.span_id,
             )
         )
         stmt = (
-            select(descendant_ids.c[root_id_label], models.Span)
+            select(descendant_ids.c[root_rowid_label], models.Span.id)
             .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])
+            .order_by(descendant_ids.c[root_rowid_label])
         )
-        results: dict[SpanId, Result] = {key: [] for key in keys}
+        results: dict[Key, 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)
+            async for key, group in groupby(data, key=lambda d: d[0]):
+                results[key].extend(span_rowid for _, span_rowid in group)
         return [results[key].copy() for key in keys]

phoenix/server/api/dataloaders/span_fields.py

@@ -0,0 +1,76 @@
+from typing import Any, Iterable, Union
+
+from sqlalchemy import Select, select
+from sqlalchemy.orm import QueryableAttribute
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+from phoenix.server.types import DbSessionFactory
+
+SpanRowId: TypeAlias = int
+
+Key: TypeAlias = tuple[SpanRowId, QueryableAttribute[Any]]
+Result: TypeAlias = Any
+
+
+_ResultColumnPosition: TypeAlias = int
+_AttrStrIdentifier: TypeAlias = str
+
+
+class SpanFieldsDataLoader(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: Iterable[Key]) -> list[Union[Result, ValueError]]:
+        result: dict[tuple[SpanRowId, _AttrStrIdentifier], Result] = {}
+        stmt, attr_strs = _get_stmt(keys)
+        async with self._db() as session:
+            data = await session.stream(stmt)
+            async for row in data:
+                span_rowid: SpanRowId = row[0]  # models.Span's primary key
+                for i, value in enumerate(row[1:]):
+                    result[span_rowid, attr_strs[i]] = value
+        return [result.get((span_rowid, str(attr))) for span_rowid, attr in keys]
+
+
+def _get_stmt(
+    keys: Iterable[Key],
+) -> tuple[
+    Select[Any],
+    dict[_ResultColumnPosition, _AttrStrIdentifier],
+]:
+    """
+    Generate a SQLAlchemy Select statement and a mapping of attribute identifiers (from their
+    column positions in the query result starting at the second column).
+
+    This function constructs a SQLAlchemy Select statement to query the `Span` model
+    based on the provided keys. It also creates a mapping of attribute identifiers
+    to their positions in the query result (starting at the second column as the zero-th
+    position).
+
+    Args:
+        keys (list[Key]): A list of tuples, where each tuple contains an integer ID, i.e. the
+            primary key of models.Span, and a QueryableAttribute.
+
+    Returns:
+        tuple: A tuple containing:
+            - Select[Any]: A SQLAlchemy Select statement with `Span` ID and attributes.
+            - dict[int, str]: A dictionary mapping the column position--where 0-th position starts
+                at the second column (because the first column is the span's primary key)--in the
+                result to the attribute's string identifier.
+    """
+    span_rowids: set[SpanRowId] = set()
+    attrs: dict[_AttrStrIdentifier, QueryableAttribute[Any]] = {}
+    joins = set()
+    for span_rowid, attr in keys:
+        span_rowids.add(span_rowid)
+        attrs[str(attr)] = attr
+        if (entity := attr.parent.entity) is not models.Span:
+            joins.add(entity)
+    stmt = select(models.Span.id).where(models.Span.id.in_(span_rowids))
+    for table in joins:
+        stmt = stmt.join(table)
+    identifiers, columns = zip(*attrs.items())
+    return stmt.add_columns(*columns), dict(enumerate(identifiers))

phoenix/server/api/dataloaders/trace_root_spans.py

@@ -1,15 +1,17 @@
-from typing import List, Optional
+from typing import Iterable, 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]
+TraceRowId: TypeAlias = int
+SpanRowId: TypeAlias = int
+
+Key: TypeAlias = TraceRowId
+Result: TypeAlias = Optional[SpanRowId]
 
 
 class TraceRootSpansDataLoader(DataLoader[Key, Result]):
@@ -17,16 +19,13 @@ class TraceRootSpansDataLoader(DataLoader[Key, Result]):
         super().__init__(load_fn=self._load_fn)
         self._db = db
 
-    async def _load_fn(self, keys: List[Key]) -> List[Result]:
+    async def _load_fn(self, keys: Iterable[Key]) -> list[Result]:
         stmt = (
-            select(models.Span)
+            select(models.Trace.id, models.Span.id)
             .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)
-            }
+            result: dict[Key, int] = {k: v async for k, v in await session.stream(stmt)}
         return [result.get(key) for key in keys]

phoenix/server/api/mutations/chat_mutations.py

@@ -60,7 +60,7 @@ from phoenix.server.api.types.ChatCompletionSubscriptionPayload import (
 from phoenix.server.api.types.Dataset import Dataset
 from phoenix.server.api.types.DatasetVersion import DatasetVersion
 from phoenix.server.api.types.node import from_global_id_with_expected_type
-from phoenix.server.api.types.Span import Span, to_gql_span
+from phoenix.server.api.types.Span import Span
 from phoenix.server.dml_event import SpanInsertEvent
 from phoenix.trace.attributes import unflatten
 from phoenix.trace.schemas import SpanException
@@ -91,6 +91,7 @@ class ChatCompletionToolCall:
 
 @strawberry.type
 class ChatCompletionMutationPayload:
+    db_span: strawberry.Private[models.Span]
     content: Optional[str]
     tool_calls: List[ChatCompletionToolCall]
     span: Span
@@ -188,7 +189,7 @@ class ChatCompletionMutationMixin:
             session.add(experiment)
             await session.flush()
 
-        results = []
+        results: list[Union[ChatCompletionMutationPayload, BaseException]] = []
         batch_size = 3
         start_time = datetime.now(timezone.utc)
         for batch in _get_batches(revisions, batch_size):
@@ -234,19 +235,19 @@ class ChatCompletionMutationMixin:
                     error=str(result),
                 )
             else:
-                db_span = result.span.db_span
+                db_span: models.Span = result.db_span
                 experiment_run = models.ExperimentRun(
                     experiment_id=experiment.id,
                     dataset_example_id=revision.dataset_example_id,
-                    trace_id=str(result.span.context.trace_id),
+                    trace_id=db_span.trace.trace_id,
                     output=models.ExperimentRunOutput(
                         task_output=get_dataset_example_output(db_span),
                     ),
                     prompt_token_count=db_span.cumulative_llm_token_count_prompt,
                     completion_token_count=db_span.cumulative_llm_token_count_completion,
                     repetition_number=1,
-                    start_time=result.span.start_time,
-                    end_time=result.span.end_time,
+                    start_time=db_span.start_time,
+                    end_time=db_span.end_time,
                     error=str(result.error_message) if result.error_message else None,
                 )
             experiment_runs.append(experiment_run)
@@ -433,12 +434,13 @@ class ChatCompletionMutationMixin:
             session.add(span)
             await session.flush()
 
-        gql_span = to_gql_span(span)
+        gql_span = Span(span_rowid=span.id, db_span=span)
 
         info.context.event_queue.put(SpanInsertEvent(ids=(project_id,)))
 
         if status_code is StatusCode.ERROR:
             return ChatCompletionMutationPayload(
+                db_span=span,
                 content=None,
                 tool_calls=[],
                 span=gql_span,
@@ -446,6 +448,7 @@ class ChatCompletionMutationMixin:
             )
         else:
             return ChatCompletionMutationPayload(
+                db_span=span,
                 content=text_content if text_content else None,
                 tool_calls=list(tool_calls.values()),
                 span=gql_span,

phoenix/server/api/queries.py

@@ -64,7 +64,7 @@ from phoenix.server.api.types.Prompt import Prompt, to_gql_prompt_from_orm
 from phoenix.server.api.types.PromptLabel import PromptLabel, to_gql_prompt_label
 from phoenix.server.api.types.PromptVersion import PromptVersion, to_gql_prompt_version
 from phoenix.server.api.types.SortDir import SortDir
-from phoenix.server.api.types.Span import Span, to_gql_span
+from phoenix.server.api.types.Span import Span
 from phoenix.server.api.types.SystemApiKey import SystemApiKey
 from phoenix.server.api.types.Trace import to_gql_trace
 from phoenix.server.api.types.User import User, to_gql_user
@@ -483,7 +483,7 @@ class Query:
                 span = await session.scalar(span_stmt)
             if span is None:
                 raise NotFound(f"Unknown span: {id}")
-            return to_gql_span(span)
+            return Span(span_rowid=span.id, db_span=span)
         elif type_name == Dataset.__name__:
             dataset_stmt = select(models.Dataset).where(models.Dataset.id == node_id)
             async with info.context.db() as session:

phoenix/server/api/subscriptions.py

@@ -59,7 +59,7 @@ from phoenix.server.api.types.DatasetVersion import DatasetVersion
 from phoenix.server.api.types.Experiment import to_gql_experiment
 from phoenix.server.api.types.ExperimentRun import to_gql_experiment_run
 from phoenix.server.api.types.node import from_global_id_with_expected_type
-from phoenix.server.api.types.Span import to_gql_span
+from phoenix.server.api.types.Span import Span
 from phoenix.server.dml_event import SpanInsertEvent
 from phoenix.server.types import DbSessionFactory
 from phoenix.utilities.template_formatters import (
@@ -165,7 +165,7 @@ class Subscription:
             session.add(db_span)
             await session.flush()
         info.context.event_queue.put(SpanInsertEvent(ids=(playground_project_id,)))
-        yield ChatCompletionSubscriptionResult(span=to_gql_span(db_span))
+        yield ChatCompletionSubscriptionResult(span=Span(span_rowid=db_span.id, db_span=db_span))
 
     @strawberry.subscription(permission_classes=[IsNotReadOnly, IsLocked])  # type: ignore
     async def chat_completion_over_dataset(
@@ -457,7 +457,7 @@ async def _chat_completion_result_payloads(
         await session.flush()
     for example_id, span, run in results:
         yield ChatCompletionSubscriptionResult(
-            span=to_gql_span(span) if span else None,
+            span=Span(span_rowid=span.id, db_span=span) if span else None,
             experiment_run=to_gql_experiment_run(run),
             dataset_example_id=example_id,
         )

phoenix/server/api/types/Annotation.py

@@ -2,6 +2,8 @@ from typing import Optional
 
 import strawberry
 
+from phoenix.server.api.interceptor import GqlValueMediator
+
 
 @strawberry.interface
 class Annotation:
@@ -9,7 +11,8 @@ class Annotation:
         description="Name of the annotation, e.g. 'helpfulness' or 'relevance'."
     )
     score: Optional[float] = strawberry.field(
-        description="Value of the annotation in the form of a numeric score."
+        description="Value of the annotation in the form of a numeric score.",
+        default=GqlValueMediator(),
     )
     label: Optional[str] = strawberry.field(
         description="Value of the annotation in the form of a string, e.g. "

phoenix/server/api/types/DatasetExample.py

@@ -19,7 +19,7 @@ from phoenix.server.api.types.pagination import (
     CursorString,
     connection_from_list,
 )
-from phoenix.server.api.types.Span import Span, to_gql_span
+from phoenix.server.api.types.Span import Span
 
 
 @strawberry.type
@@ -52,7 +52,7 @@ class DatasetExample(Node):
         info: Info[Context, None],
     ) -> Optional[Span]:
         return (
-            to_gql_span(span)
+            Span(span_rowid=span.id, db_span=span)
             if (span := await info.context.data_loaders.dataset_example_spans.load(self.id_attr))
             else None
         )

phoenix/server/api/types/Project.py

@@ -6,7 +6,6 @@ import strawberry
 from aioitertools.itertools import islice
 from openinference.semconv.trace import SpanAttributes
 from sqlalchemy import desc, distinct, func, or_, select
-from sqlalchemy.orm import contains_eager
 from sqlalchemy.sql.elements import ColumnElement
 from sqlalchemy.sql.expression import tuple_
 from strawberry import ID, UNSET
@@ -33,7 +32,7 @@ from phoenix.server.api.types.pagination import (
 )
 from phoenix.server.api.types.ProjectSession import ProjectSession, to_gql_project_session
 from phoenix.server.api.types.SortDir import SortDir
-from phoenix.server.api.types.Span import Span, to_gql_span
+from phoenix.server.api.types.Span import Span
 from phoenix.server.api.types.Trace import Trace, to_gql_trace
 from phoenix.server.api.types.ValidationResult import ValidationResult
 from phoenix.trace.dsl import SpanFilter
@@ -184,10 +183,9 @@ class Project(Node):
         filter_condition: Optional[str] = UNSET,
     ) -> Connection[Span]:
         stmt = (
-            select(models.Span)
+            select(models.Span.id)
             .join(models.Trace)
             .where(models.Trace.project_rowid == self.id_attr)
-            .options(contains_eager(models.Span.trace).load_only(models.Trace.trace_id))
         )
         if time_range:
             if time_range.start:
@@ -232,21 +230,21 @@ class Project(Node):
         stmt = stmt.order_by(cursor_rowid_column)
         cursors_and_nodes = []
         async with info.context.db() as session:
-            span_records = await session.execute(stmt)
+            span_records = await session.stream(stmt)
             async for span_record in islice(span_records, first):
-                span = span_record[0]
-                cursor = Cursor(rowid=span.id)
+                span_rowid: int = span_record[0]
+                cursor = Cursor(rowid=span_rowid)
                 if sort_config:
                     assert len(span_record) > 1
                     cursor.sort_column = CursorSortColumn(
                         type=sort_config.column_data_type,
                         value=span_record[1],
                     )
-                cursors_and_nodes.append((cursor, to_gql_span(span)))
+                cursors_and_nodes.append((cursor, Span(span_rowid=span_rowid)))
             has_next_page = True
             try:
-                next(span_records)
-            except StopIteration:
+                await span_records.__anext__()
+            except StopAsyncIteration:
                 has_next_page = False
 
         return connection_from_cursors_and_nodes(

phoenix/server/api/types/ProjectSession.py

@@ -57,7 +57,7 @@ class ProjectSession(Node):
             return None
         return SpanIOValue(
             mime_type=MimeType(record.mime_type.value),
-            value=record.value,
+            cached_value=record.value,
         )
 
     @strawberry.field
@@ -70,7 +70,7 @@ class ProjectSession(Node):
             return None
         return SpanIOValue(
             mime_type=MimeType(record.mime_type.value),
-            value=record.value,
+            cached_value=record.value,
         )
 
     @strawberry.field