arize-phoenix 8.0.1__tar.gz → 8.2.0__tar.gz

{arize_phoenix-8.0.1 → arize_phoenix-8.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: arize-phoenix
-Version: 8.0.1
+Version: 8.2.0
 Summary: AI Observability and Evaluation
 Project-URL: Documentation, https://docs.arize.com/phoenix/
 Project-URL: Issues, https://github.com/Arize-ai/phoenix/issues
@@ -30,7 +30,6 @@ Requires-Dist: grpc-interceptor
 Requires-Dist: grpcio
 Requires-Dist: httpx
 Requires-Dist: jinja2
-Requires-Dist: jsonschema<=4.23.0,>=4.0.0
 Requires-Dist: numpy!=2.0.0
 Requires-Dist: openinference-instrumentation>=0.1.12
 Requires-Dist: openinference-semantic-conventions>=0.1.12

{arize_phoenix-8.0.1 → arize_phoenix-8.2.0}/pyproject.toml

@@ -58,7 +58,6 @@ dependencies = [
   "pydantic>=2.1.0", # exclude 2.0.* since it does not support the `json_encoders` configuration setting
   "authlib",
   "websockets",
-  "jsonschema>=4.0.0,<=4.23.0", # the upper bound is to keep us off the bleeding edge in case there's a regression since this controls what gets written to the database
   "arize-phoenix-client",
 ]
 dynamic = ["version"]

{arize_phoenix-8.0.1 → arize_phoenix-8.2.0}/src/phoenix/config.py

@@ -8,7 +8,7 @@ from enum import Enum
 from importlib.metadata import version
 from pathlib import Path
 from typing import Optional, cast, overload
-from urllib.parse import urlparse
+from urllib.parse import quote_plus, urlparse
 
 from phoenix.utilities.logging import log_a_list
 
@@ -624,11 +624,37 @@ def get_env_project_name() -> str:
 
 
 def get_env_database_connection_str() -> str:
-    env_url = getenv(ENV_PHOENIX_SQL_DATABASE_URL)
-    if env_url is None:
-        working_dir = get_working_dir()
-        return f"sqlite:///{working_dir}/phoenix.db"
-    return env_url
+    if phoenix_url := os.getenv(ENV_PHOENIX_SQL_DATABASE_URL):
+        return phoenix_url
+
+    if postgres_url := get_env_postgres_connection_str():
+        return postgres_url
+
+    working_dir = get_working_dir()
+    return f"sqlite:///{working_dir}/phoenix.db"
+
+
+def get_env_postgres_connection_str() -> Optional[str]:
+    pg_user = os.getenv("PHOENIX_POSTGRES_USER")
+    pg_password = os.getenv("PHOENIX_POSTGRES_PASSWORD")
+    pg_host = os.getenv("PHOENIX_POSTGRES_HOST")
+    pg_port = os.getenv("PHOENIX_POSTGRES_PORT")
+    pg_db = os.getenv("PHOENIX_POSTGRES_DB")
+
+    if pg_host and ":" in pg_host:
+        pg_host, parsed_port = pg_host.split(":")
+        pg_port = pg_port or parsed_port  # use the explicitly set port if provided
+
+    if pg_host and pg_user and pg_password:
+        encoded_password = quote_plus(pg_password)
+        connection_str = f"postgresql://{pg_user}:{encoded_password}@{pg_host}"
+        if pg_port:
+            connection_str = f"{connection_str}:{pg_port}"
+        if pg_db:
+            connection_str = f"{connection_str}/{pg_db}"
+
+        return connection_str
+    return None
 
 
 def get_env_database_schema() -> Optional[str]:

{arize_phoenix-8.0.1 → arize_phoenix-8.2.0}/src/phoenix/db/models.py

@@ -1,7 +1,9 @@
 from datetime import datetime, timezone
 from enum import Enum
-from typing import Any, Optional, TypedDict
+from typing import Any, Optional, Sequence, TypedDict
 
+import sqlalchemy.sql as sql
+from openinference.semconv.trace import RerankerAttributes, SpanAttributes
 from sqlalchemy import (
     JSON,
     NUMERIC,
@@ -12,6 +14,7 @@ from sqlalchemy import (
     Float,
     ForeignKey,
     Index,
+    Integer,
     MetaData,
     Null,
     String,
@@ -25,6 +28,7 @@ from sqlalchemy import (
     text,
 )
 from sqlalchemy.dialects import postgresql
+from sqlalchemy.dialects.sqlite.base import SQLiteCompiler
 from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession
 from sqlalchemy.ext.compiler import compiles
 from sqlalchemy.ext.hybrid import hybrid_property
@@ -36,6 +40,8 @@ from sqlalchemy.orm import (
     relationship,
 )
 from sqlalchemy.sql import expression
+from sqlalchemy.sql.compiler import SQLCompiler
+from sqlalchemy.sql.functions import coalesce
 
 from phoenix.config import get_env_database_schema
 from phoenix.datetime_utils import normalize_datetime
@@ -54,6 +60,18 @@ from phoenix.server.api.helpers.prompts.models import (
     is_prompt_invocation_parameters,
     is_prompt_template,
 )
+from phoenix.trace.attributes import get_attribute_value
+
+INPUT_MIME_TYPE = SpanAttributes.INPUT_MIME_TYPE.split(".")
+INPUT_VALUE = SpanAttributes.INPUT_VALUE.split(".")
+LLM_TOKEN_COUNT_TOTAL = SpanAttributes.LLM_TOKEN_COUNT_TOTAL.split(".")
+LLM_TOKEN_COUNT_PROMPT = SpanAttributes.LLM_TOKEN_COUNT_PROMPT.split(".")
+LLM_TOKEN_COUNT_COMPLETION = SpanAttributes.LLM_TOKEN_COUNT_COMPLETION.split(".")
+METADATA = SpanAttributes.METADATA.split(".")
+OUTPUT_MIME_TYPE = SpanAttributes.OUTPUT_MIME_TYPE.split(".")
+OUTPUT_VALUE = SpanAttributes.OUTPUT_VALUE.split(".")
+RERANKER_OUTPUT_DOCUMENTS = RerankerAttributes.RERANKER_OUTPUT_DOCUMENTS.split(".")
+RETRIEVAL_DOCUMENTS = SpanAttributes.RETRIEVAL_DOCUMENTS.split(".")
 
 
 class AuthMethod(Enum):
@@ -341,13 +359,11 @@ class Trace(Base):
 
     @hybrid_property
     def latency_ms(self) -> float:
-        # See https://docs.sqlalchemy.org/en/20/orm/extensions/hybrid.html
         return (self.end_time - self.start_time).total_seconds() * 1000
 
     @latency_ms.inplace.expression
     @classmethod
     def _latency_ms_expression(cls) -> ColumnElement[float]:
-        # See https://docs.sqlalchemy.org/en/20/orm/extensions/hybrid.html
         return LatencyMs(cls.start_time, cls.end_time)
 
     project: Mapped["Project"] = relationship(
@@ -404,25 +420,124 @@ class Span(Base):
 
     @hybrid_property
     def latency_ms(self) -> float:
-        # See https://docs.sqlalchemy.org/en/20/orm/extensions/hybrid.html
-        return (self.end_time - self.start_time).total_seconds() * 1000
+        return round((self.end_time - self.start_time).total_seconds() * 1000, 1)
 
     @latency_ms.inplace.expression
     @classmethod
     def _latency_ms_expression(cls) -> ColumnElement[float]:
-        # See https://docs.sqlalchemy.org/en/20/orm/extensions/hybrid.html
         return LatencyMs(cls.start_time, cls.end_time)
 
+    @hybrid_property
+    def input_value(self) -> Any:
+        return get_attribute_value(self.attributes, INPUT_VALUE)
+
+    @input_value.inplace.expression
+    @classmethod
+    def _input_value_expression(cls) -> ColumnElement[Any]:
+        return cls.attributes[INPUT_VALUE]
+
+    @hybrid_property
+    def input_value_first_101_chars(self) -> Any:
+        if (v := get_attribute_value(self.attributes, INPUT_VALUE)) is None:
+            return None
+        return str(v)[:101]
+
+    @input_value_first_101_chars.inplace.expression
+    @classmethod
+    def _input_value_first_101_chars_expression(cls) -> ColumnElement[Any]:
+        return case(
+            (
+                cls.attributes[INPUT_VALUE] != sql.null(),
+                func.substr(cls.attributes[INPUT_VALUE].as_string(), 1, 101),
+            ),
+        )
+
+    @hybrid_property
+    def input_mime_type(self) -> Any:
+        return get_attribute_value(self.attributes, INPUT_MIME_TYPE)
+
+    @input_mime_type.inplace.expression
+    @classmethod
+    def _input_mime_type_expression(cls) -> ColumnElement[Any]:
+        return cls.attributes[INPUT_MIME_TYPE]
+
+    @hybrid_property
+    def output_value(self) -> Any:
+        return get_attribute_value(self.attributes, OUTPUT_VALUE)
+
+    @output_value.inplace.expression
+    @classmethod
+    def _output_value_expression(cls) -> ColumnElement[Any]:
+        return cls.attributes[OUTPUT_VALUE]
+
+    @hybrid_property
+    def output_value_first_101_chars(self) -> Any:
+        if (v := get_attribute_value(self.attributes, OUTPUT_VALUE)) is None:
+            return None
+        return str(v)[:101]
+
+    @output_value_first_101_chars.inplace.expression
+    @classmethod
+    def _output_value_first_101_chars_expression(cls) -> ColumnElement[Any]:
+        return case(
+            (
+                cls.attributes[OUTPUT_VALUE] != sql.null(),
+                func.substr(cls.attributes[OUTPUT_VALUE].as_string(), 1, 101),
+            ),
+        )
+
+    @hybrid_property
+    def output_mime_type(self) -> Any:
+        return get_attribute_value(self.attributes, OUTPUT_MIME_TYPE)
+
+    @output_mime_type.inplace.expression
+    @classmethod
+    def _output_mime_type_expression(cls) -> ColumnElement[Any]:
+        return cls.attributes[OUTPUT_MIME_TYPE]
+
+    @hybrid_property
+    def metadata_(self) -> Any:
+        return get_attribute_value(self.attributes, METADATA)
+
+    @metadata_.inplace.expression
+    @classmethod
+    def _metadata_expression(cls) -> ColumnElement[Any]:
+        return cls.attributes[METADATA]
+
+    @hybrid_property
+    def num_documents(self) -> int:
+        if self.span_kind.upper() == "RERANKER":
+            reranker_documents = get_attribute_value(self.attributes, RERANKER_OUTPUT_DOCUMENTS)
+            return len(reranker_documents) if isinstance(reranker_documents, Sequence) else 0
+        retrieval_documents = get_attribute_value(self.attributes, RETRIEVAL_DOCUMENTS)
+        return len(retrieval_documents) if isinstance(retrieval_documents, Sequence) else 0
+
+    @num_documents.inplace.expression
+    @classmethod
+    def _num_documents_expression(cls) -> ColumnElement[int]:
+        return NumDocuments(cls.attributes, cls.span_kind)
+
     @hybrid_property
     def cumulative_llm_token_count_total(self) -> int:
         return self.cumulative_llm_token_count_prompt + self.cumulative_llm_token_count_completion
 
+    @cumulative_llm_token_count_total.inplace.expression
+    @classmethod
+    def _cumulative_llm_token_count_total_expression(cls) -> ColumnElement[int]:
+        return cls.cumulative_llm_token_count_prompt + cls.cumulative_llm_token_count_completion
+
     @hybrid_property
-    def llm_token_count_total(self) -> Optional[int]:
-        if self.llm_token_count_prompt is None and self.llm_token_count_completion is None:
-            return None
+    def llm_token_count_total(self) -> int:
         return (self.llm_token_count_prompt or 0) + (self.llm_token_count_completion or 0)
 
+    @llm_token_count_total.inplace.expression
+    @classmethod
+    def _llm_token_count_total_expression(cls) -> ColumnElement[int]:
+        return coalesce(
+            coalesce(cls.llm_token_count_prompt, 0) + coalesce(cls.llm_token_count_completion, 0),
+            0,
+        )
+
     trace: Mapped["Trace"] = relationship("Trace", back_populates="spans")
     document_annotations: Mapped[list["DocumentAnnotation"]] = relationship(back_populates="span")
     dataset_examples: Mapped[list["DatasetExample"]] = relationship(back_populates="span")
@@ -477,6 +592,33 @@ def _(element: Any, compiler: Any, **kw: Any) -> Any:
     )
 
 
+class NumDocuments(expression.FunctionElement[int]):
+    # See https://docs.sqlalchemy.org/en/20/core/compiler.html
+    inherit_cache = True
+    type = Integer()
+    name = "num_documents"
+
+
+@compiles(NumDocuments)
+def _(element: Any, compiler: SQLCompiler, **kw: Any) -> Any:
+    # See https://docs.sqlalchemy.org/en/20/core/compiler.html
+    array_length = (
+        func.json_array_length if isinstance(compiler, SQLiteCompiler) else func.jsonb_array_length
+    )
+    attributes, span_kind = list(element.clauses)
+    retrieval_docs = attributes[RETRIEVAL_DOCUMENTS]
+    num_retrieval_docs = coalesce(array_length(retrieval_docs), 0)
+    reranker_docs = attributes[RERANKER_OUTPUT_DOCUMENTS]
+    num_reranker_docs = coalesce(array_length(reranker_docs), 0)
+    return compiler.process(
+        sql.case(
+            (func.upper(span_kind) == "RERANKER", num_reranker_docs),
+            else_=num_retrieval_docs,
+        ),
+        **kw,
+    )
+
+
 class TextContains(expression.FunctionElement[str]):
     # See https://docs.sqlalchemy.org/en/20/core/compiler.html
     inherit_cache = True
@@ -763,13 +905,11 @@ class ExperimentRun(Base):
 
     @hybrid_property
     def latency_ms(self) -> float:
-        # See https://docs.sqlalchemy.org/en/20/orm/extensions/hybrid.html
         return (self.end_time - self.start_time).total_seconds() * 1000
 
     @latency_ms.inplace.expression
     @classmethod
     def _latency_expression(cls) -> ColumnElement[float]:
-        # See https://docs.sqlalchemy.org/en/20/orm/extensions/hybrid.html
         return LatencyMs(cls.start_time, cls.end_time)
 
     trace: Mapped["Trace"] = relationship(

{arize_phoenix-8.0.1 → arize_phoenix-8.2.0}/src/phoenix/server/api/context.py

@@ -38,8 +38,10 @@ from phoenix.server.api.dataloaders import (
     SessionTokenUsagesDataLoader,
     SessionTraceLatencyMsQuantileDataLoader,
     SpanAnnotationsDataLoader,
+    SpanByIdDataLoader,
     SpanDatasetExamplesDataLoader,
     SpanDescendantsDataLoader,
+    SpanFieldsDataLoader,
     SpanProjectsDataLoader,
     TokenCountDataLoader,
     TraceByTraceIdsDataLoader,
@@ -83,8 +85,10 @@ class DataLoaders:
     session_token_usages: SessionTokenUsagesDataLoader
     session_trace_latency_ms_quantile: SessionTraceLatencyMsQuantileDataLoader
     span_annotations: SpanAnnotationsDataLoader
+    span_by_id: SpanByIdDataLoader
     span_dataset_examples: SpanDatasetExamplesDataLoader
     span_descendants: SpanDescendantsDataLoader
+    span_fields: SpanFieldsDataLoader
     span_projects: SpanProjectsDataLoader
     token_counts: TokenCountDataLoader
     trace_by_trace_ids: TraceByTraceIdsDataLoader

{arize_phoenix-8.0.1 → arize_phoenix-8.2.0}/src/phoenix/server/api/dataloaders/__init__.py

@@ -26,8 +26,10 @@ from .session_num_traces_with_error import SessionNumTracesWithErrorDataLoader
 from .session_token_usages import SessionTokenUsagesDataLoader
 from .session_trace_latency_ms_quantile import SessionTraceLatencyMsQuantileDataLoader
 from .span_annotations import SpanAnnotationsDataLoader
+from .span_by_id import SpanByIdDataLoader
 from .span_dataset_examples import SpanDatasetExamplesDataLoader
 from .span_descendants import SpanDescendantsDataLoader
+from .span_fields import SpanFieldsDataLoader
 from .span_projects import SpanProjectsDataLoader
 from .token_counts import TokenCountCache, TokenCountDataLoader
 from .trace_by_trace_ids import TraceByTraceIdsDataLoader
@@ -58,6 +60,8 @@ __all__ = [
     "SessionNumTracesWithErrorDataLoader",
     "SessionTokenUsagesDataLoader",
     "SessionTraceLatencyMsQuantileDataLoader",
+    "SpanFieldsDataLoader",
+    "SpanByIdDataLoader",
     "SpanDatasetExamplesDataLoader",
     "SpanDescendantsDataLoader",
     "SpanProjectsDataLoader",

arize_phoenix-8.2.0/src/phoenix/server/api/dataloaders/span_by_id.py

@@ -0,0 +1,29 @@
+from typing import Iterable, 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
+
+SpanRowId: TypeAlias = int
+
+Key: TypeAlias = SpanRowId
+Result: TypeAlias = models.Span
+
+
+class SpanByIdDataLoader(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]]:
+        span_rowids = list(set(keys))
+        spans: dict[Key, Result] = {}
+        stmt = select(models.Span).where(models.Span.id.in_(span_rowids))
+        async with self._db() as session:
+            data = await session.stream_scalars(stmt)
+            async for span in data:
+                spans[span.id] = span
+        return [spans.get(span_rowid, ValueError("Invalid span row id")) for span_rowid in keys]

{arize_phoenix-8.0.1 → arize_phoenix-8.2.0}/src/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]

arize_phoenix-8.2.0/src/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))

{arize_phoenix-8.0.1 → arize_phoenix-8.2.0}/src/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]

{arize_phoenix-8.0.1 → arize_phoenix-8.2.0}/src/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,