arize-phoenix 8.1.0__py3-none-any.whl → 8.2.1__py3-none-any.whl

phoenix/server/api/types/Span.py

@@ -1,8 +1,9 @@
 import json
-from collections.abc import Mapping, Sized
+from asyncio import gather
+from collections.abc import Mapping
 from datetime import datetime
 from enum import Enum
-from typing import TYPE_CHECKING, Any, Optional, cast
+from typing import TYPE_CHECKING, Any, Iterable, Optional, cast
 
 import numpy as np
 import strawberry
@@ -10,7 +11,7 @@ from openinference.semconv.trace import SpanAttributes
 from strawberry import ID, UNSET
 from strawberry.relay import Node, NodeID
 from strawberry.types import Info
-from typing_extensions import Annotated
+from typing_extensions import Annotated, TypeAlias
 
 import phoenix.trace.schemas as trace_schema
 from phoenix.db import models
@@ -31,7 +32,7 @@ from phoenix.server.api.types.GenerativeProvider import GenerativeProvider
 from phoenix.server.api.types.MimeType import MimeType
 from phoenix.server.api.types.SortDir import SortDir
 from phoenix.server.api.types.SpanAnnotation import SpanAnnotation, to_gql_span_annotation
-from phoenix.server.api.types.SpanIOValue import SpanIOValue
+from phoenix.server.api.types.SpanIOValue import SpanIOValue, truncate_value
 from phoenix.trace.attributes import get_attribute_value
 
 if TYPE_CHECKING:
@@ -102,50 +103,346 @@ class SpanEvent:
 class SpanAsExampleRevision(ExampleRevision): ...
 
 
+SpanRowId: TypeAlias = int
+
+
 @strawberry.type
 class Span(Node):
-    id_attr: NodeID[int]
-    db_span: strawberry.Private[models.Span]
-    name: str
-    status_code: SpanStatusCode
-    status_message: str
-    start_time: datetime
-    end_time: Optional[datetime]
-    latency_ms: Optional[float]
-    parent_id: Optional[ID] = strawberry.field(
-        description="the parent span ID. If null, it is a root span"
-    )
-    span_kind: SpanKind
-    context: SpanContext
-    attributes: str = strawberry.field(
+    span_rowid: NodeID[SpanRowId]
+    db_span: strawberry.Private[models.Span] = UNSET
+
+    def __post_init__(self) -> None:
+        if self.db_span and self.span_rowid != self.db_span.id:
+            raise ValueError("Span ID mismatch")
+
+    @strawberry.field
+    async def name(
+        self,
+        info: Info[Context, None],
+    ) -> str:
+        if self.db_span:
+            return self.db_span.name
+        value = await info.context.data_loaders.span_fields.load(
+            (self.span_rowid, models.Span.name),
+        )
+        return str(value)
+
+    @strawberry.field
+    async def status_code(
+        self,
+        info: Info[Context, None],
+    ) -> SpanStatusCode:
+        if self.db_span:
+            value = self.db_span.status_code
+        else:
+            value = await info.context.data_loaders.span_fields.load(
+                (self.span_rowid, models.Span.status_code),
+            )
+        return SpanStatusCode(value)
+
+    @strawberry.field
+    async def status_message(
+        self,
+        info: Info[Context, None],
+    ) -> str:
+        if self.db_span:
+            return self.db_span.status_message
+        value = await info.context.data_loaders.span_fields.load(
+            (self.span_rowid, models.Span.status_message),
+        )
+        return str(value)
+
+    @strawberry.field
+    async def start_time(
+        self,
+        info: Info[Context, None],
+    ) -> datetime:
+        if self.db_span:
+            return self.db_span.start_time
+        value = await info.context.data_loaders.span_fields.load(
+            (self.span_rowid, models.Span.start_time),
+        )
+        return cast(datetime, value)
+
+    @strawberry.field
+    async def end_time(
+        self,
+        info: Info[Context, None],
+    ) -> Optional[datetime]:
+        if self.db_span:
+            return self.db_span.end_time
+        value = await info.context.data_loaders.span_fields.load(
+            (self.span_rowid, models.Span.end_time),
+        )
+        return cast(datetime, value)
+
+    @strawberry.field
+    async def latency_ms(
+        self,
+        info: Info[Context, None],
+    ) -> Optional[float]:
+        if self.db_span:
+            return self.db_span.latency_ms
+        value = await info.context.data_loaders.span_fields.load(
+            (self.span_rowid, models.Span.latency_ms),
+        )
+        return cast(float, value)
+
+    @strawberry.field(
+        description="the parent span ID. If null, it is a root span",
+    )  # type: ignore
+    async def parent_id(
+        self,
+        info: Info[Context, None],
+    ) -> Optional[ID]:
+        if self.db_span:
+            value = self.db_span.parent_id
+        else:
+            value = await info.context.data_loaders.span_fields.load(
+                (self.span_rowid, models.Span.parent_id),
+            )
+        return None if value is None else ID(value)
+
+    @strawberry.field
+    async def span_kind(
+        self,
+        info: Info[Context, None],
+    ) -> SpanKind:
+        if self.db_span:
+            value = self.db_span.span_kind
+        else:
+            value = await info.context.data_loaders.span_fields.load(
+                (self.span_rowid, models.Span.span_kind),
+            )
+        return SpanKind(value)
+
+    @strawberry.field
+    async def context(
+        self,
+        info: Info[Context, None],
+    ) -> SpanContext:
+        if self.db_span:
+            trace_id = self.db_span.trace.trace_id
+            span_id = self.db_span.span_id
+        else:
+            span_id, trace_id = await gather(
+                info.context.data_loaders.span_fields.load(
+                    (self.span_rowid, models.Span.span_id),
+                ),
+                info.context.data_loaders.span_fields.load(
+                    (self.span_rowid, models.Trace.trace_id),
+                ),
+            )
+        return SpanContext(trace_id=ID(trace_id), span_id=ID(span_id))
+
+    @strawberry.field(
         description="Span attributes as a JSON string",
-    )
-    metadata: Optional[str] = strawberry.field(
+    )  # type: ignore
+    async def attributes(
+        self,
+        info: Info[Context, None],
+    ) -> str:
+        if self.db_span:
+            value = self.db_span.attributes
+        else:
+            value = await info.context.data_loaders.span_fields.load(
+                (self.span_rowid, models.Span.attributes),
+            )
+        return json.dumps(_hide_embedding_vectors(value), cls=_JSONEncoder)
+
+    @strawberry.field(
         description="Metadata as a JSON string",
-    )
-    num_documents: Optional[int]
-    token_count_total: Optional[int]
-    token_count_prompt: Optional[int]
-    token_count_completion: Optional[int]
-    input: Optional[SpanIOValue]
-    output: Optional[SpanIOValue]
-    events: list[SpanEvent]
-    cumulative_token_count_total: Optional[int] = strawberry.field(
-        description="Cumulative (prompt plus completion) token count from "
-        "self and all descendant spans (children, grandchildren, etc.)",
-    )
-    cumulative_token_count_prompt: Optional[int] = strawberry.field(
-        description="Cumulative (prompt) token count from self and all "
-        "descendant spans (children, grandchildren, etc.)",
-    )
-    cumulative_token_count_completion: Optional[int] = strawberry.field(
-        description="Cumulative (completion) token count from self and all "
+    )  # type: ignore
+    async def metadata(
+        self,
+        info: Info[Context, None],
+    ) -> Optional[str]:
+        if self.db_span:
+            value = self.db_span.metadata_
+        else:
+            value = await info.context.data_loaders.span_fields.load(
+                (self.span_rowid, models.Span.metadata_),
+            )
+        return _convert_metadata_to_string(value)
+
+    @strawberry.field
+    async def num_documents(
+        self,
+        info: Info[Context, None],
+    ) -> Optional[int]:
+        if self.db_span:
+            return self.db_span.num_documents
+        value = await info.context.data_loaders.span_fields.load(
+            (self.span_rowid, models.Span.num_documents),
+        )
+        return cast(int, value)
+
+    @strawberry.field
+    async def token_count_total(
+        self,
+        info: Info[Context, None],
+    ) -> Optional[int]:
+        if self.db_span:
+            return self.db_span.llm_token_count_total
+        value = await info.context.data_loaders.span_fields.load(
+            (self.span_rowid, models.Span.llm_token_count_total),
+        )
+        return cast(Optional[int], value)
+
+    @strawberry.field
+    async def token_count_prompt(
+        self,
+        info: Info[Context, None],
+    ) -> Optional[int]:
+        if self.db_span:
+            return self.db_span.llm_token_count_prompt
+        value = await info.context.data_loaders.span_fields.load(
+            (self.span_rowid, models.Span.llm_token_count_prompt),
+        )
+        return cast(Optional[int], value)
+
+    @strawberry.field
+    async def token_count_completion(
+        self,
+        info: Info[Context, None],
+    ) -> Optional[int]:
+        if self.db_span:
+            return self.db_span.llm_token_count_completion
+        value = await info.context.data_loaders.span_fields.load(
+            (self.span_rowid, models.Span.llm_token_count_completion),
+        )
+        return cast(Optional[int], value)
+
+    @strawberry.field
+    async def input(
+        self,
+        info: Info[Context, None],
+    ) -> Optional[SpanIOValue]:
+        if self.db_span:
+            mime_type = self.db_span.input_mime_type
+            input_value = self.db_span.input_value
+            return SpanIOValue(
+                cached_value=input_value,
+                mime_type=MimeType(mime_type),
+            )
+        mime_type, input_value_first_101_chars = await gather(
+            info.context.data_loaders.span_fields.load(
+                (self.span_rowid, models.Span.input_mime_type),
+            ),
+            info.context.data_loaders.span_fields.load(
+                (self.span_rowid, models.Span.input_value_first_101_chars),
+            ),
+        )
+        if not input_value_first_101_chars:
+            return None
+        return SpanIOValue(
+            span_rowid=self.span_rowid,
+            attr=models.Span.input_value,
+            truncated_value=truncate_value(input_value_first_101_chars),
+            mime_type=MimeType(mime_type),
+        )
+
+    @strawberry.field
+    async def output(
+        self,
+        info: Info[Context, None],
+    ) -> Optional[SpanIOValue]:
+        if self.db_span:
+            mime_type = self.db_span.output_mime_type
+            output_value = self.db_span.output_value
+            return SpanIOValue(
+                cached_value=output_value,
+                mime_type=MimeType(mime_type),
+            )
+        mime_type, output_value_first_101_chars = await gather(
+            info.context.data_loaders.span_fields.load(
+                (self.span_rowid, models.Span.output_mime_type),
+            ),
+            info.context.data_loaders.span_fields.load(
+                (self.span_rowid, models.Span.output_value_first_101_chars),
+            ),
+        )
+        if not output_value_first_101_chars:
+            return None
+        return SpanIOValue(
+            span_rowid=self.span_rowid,
+            attr=models.Span.output_value,
+            truncated_value=truncate_value(output_value_first_101_chars),
+            mime_type=MimeType(mime_type),
+        )
+
+    @strawberry.field
+    async def events(
+        self,
+        info: Info[Context, None],
+    ) -> list[SpanEvent]:
+        if self.db_span:
+            return [SpanEvent.from_dict(event) for event in self.db_span.events]
+        value = await info.context.data_loaders.span_fields.load(
+            (self.span_rowid, models.Span.events),
+        )
+        return [SpanEvent.from_dict(event) for event in value]
+
+    @strawberry.field(
+        description="Cumulative (prompt plus completion) token count from self "
+        "and all descendant spans (children, grandchildren, etc.)",
+    )  # type: ignore
+    async def cumulative_token_count_total(
+        self,
+        info: Info[Context, None],
+    ) -> Optional[int]:
+        if self.db_span:
+            return self.db_span.cumulative_llm_token_count_total
+        value = await info.context.data_loaders.span_fields.load(
+            (self.span_rowid, models.Span.cumulative_llm_token_count_total),
+        )
+        return cast(Optional[int], value)
+
+    @strawberry.field(
+        description="Cumulative (prompt) token count from self and all descendant "
+        "spans (children, grandchildren, etc.)",
+    )  # type: ignore
+    async def cumulative_token_count_prompt(
+        self,
+        info: Info[Context, None],
+    ) -> Optional[int]:
+        if self.db_span:
+            return self.db_span.cumulative_llm_token_count_prompt
+        value = await info.context.data_loaders.span_fields.load(
+            (self.span_rowid, models.Span.cumulative_llm_token_count_prompt),
+        )
+        return cast(Optional[int], value)
+
+    @strawberry.field(
+        description="Cumulative (completion) token count from self and all descendant "
+        "spans (children, grandchildren, etc.)",
+    )  # type: ignore
+    async def cumulative_token_count_completion(
+        self,
+        info: Info[Context, None],
+    ) -> Optional[int]:
+        if self.db_span:
+            return self.db_span.cumulative_llm_token_count_completion
+        value = await info.context.data_loaders.span_fields.load(
+            (self.span_rowid, models.Span.cumulative_llm_token_count_completion),
+        )
+        return cast(Optional[int], value)
+
+    @strawberry.field(
+        description="Propagated status code that percolates up error status codes from "
         "descendant spans (children, grandchildren, etc.)",
-    )
-    propagated_status_code: SpanStatusCode = strawberry.field(
-        description="Propagated status code that percolates up error status "
-        "codes from descendant spans (children, grandchildren, etc.)",
-    )
+    )  # type: ignore
+    async def propagated_status_code(
+        self,
+        info: Info[Context, None],
+    ) -> SpanStatusCode:
+        if self.db_span:
+            value = self.db_span.cumulative_error_count
+        else:
+            value = await info.context.data_loaders.span_fields.load(
+                (self.span_rowid, models.Span.cumulative_error_count),
+            )
+        return SpanStatusCode.ERROR if value else SpanStatusCode.OK
 
     @strawberry.field(
         description=(
@@ -158,7 +455,7 @@ class Span(Node):
         info: Info[Context, None],
         sort: Optional[SpanAnnotationSort] = UNSET,
     ) -> list[SpanAnnotation]:
-        span_id = self.id_attr
+        span_id = self.span_rowid
         annotations = await info.context.data_loaders.span_annotations.load(span_id)
         sort_key = SpanAnnotationColumn.name.value
         sort_descending = False
@@ -178,8 +475,11 @@ class Span(Node):
         "a list, and each evaluation is identified by its document's (zero-based) "
         "index in that list."
     )  # type: ignore
-    async def document_evaluations(self, info: Info[Context, None]) -> list[DocumentEvaluation]:
-        return await info.context.data_loaders.document_evaluations.load(self.id_attr)
+    async def document_evaluations(
+        self,
+        info: Info[Context, None],
+    ) -> list[DocumentEvaluation]:
+        return await info.context.data_loaders.document_evaluations.load(self.span_rowid)
 
     @strawberry.field(
         description="Retrieval metrics: NDCG@K, Precision@K, Reciprocal Rank, etc.",
@@ -189,10 +489,17 @@ class Span(Node):
         info: Info[Context, None],
         evaluation_name: Optional[str] = UNSET,
     ) -> list[DocumentRetrievalMetrics]:
-        if not self.num_documents:
+        num_documents = (
+            self.db_span.num_documents
+            if self.db_span
+            else await info.context.data_loaders.span_fields.load(
+                (self.span_rowid, models.Span.num_documents),
+            )
+        )
+        if not num_documents:
             return []
         return await info.context.data_loaders.document_retrieval_metrics.load(
-            (self.id_attr, evaluation_name or None, self.num_documents),
+            (self.span_rowid, evaluation_name or None, num_documents),
         )
 
     @strawberry.field(
@@ -202,15 +509,21 @@ class Span(Node):
         self,
         info: Info[Context, None],
     ) -> list["Span"]:
-        span_id = str(self.context.span_id)
-        spans = await info.context.data_loaders.span_descendants.load(span_id)
-        return [to_gql_span(span) for span in spans]
+        ids: Iterable[int] = await info.context.data_loaders.span_descendants.load(self.span_rowid)
+        return [Span(span_rowid=id_) for id_ in ids]
 
     @strawberry.field(
         description="The span's attributes translated into an example revision for a dataset",
     )  # type: ignore
-    async def as_example_revision(self, info: Info[Context, None]) -> SpanAsExampleRevision:
-        span = self.db_span
+    async def as_example_revision(
+        self,
+        info: Info[Context, None],
+    ) -> SpanAsExampleRevision:
+        span = (
+            self.db_span
+            if self.db_span
+            else await info.context.data_loaders.span_by_id.load(self.span_rowid)
+        )
 
         # Fetch annotations associated with this span
         span_annotations = await self.span_annotations(info)
@@ -244,21 +557,27 @@ class Span(Node):
     ]:  # use lazy types to avoid circular import: https://strawberry.rocks/docs/types/lazy
         from phoenix.server.api.types.Project import to_gql_project
 
-        span_id = self.id_attr
+        span_id = self.span_rowid
         project = await info.context.data_loaders.span_projects.load(span_id)
         return to_gql_project(project)
 
     @strawberry.field(description="Indicates if the span is contained in any dataset")  # type: ignore
-    async def contained_in_dataset(self, info: Info[Context, None]) -> bool:
-        examples = await info.context.data_loaders.span_dataset_examples.load(self.id_attr)
+    async def contained_in_dataset(
+        self,
+        info: Info[Context, None],
+    ) -> bool:
+        examples = await info.context.data_loaders.span_dataset_examples.load(self.span_rowid)
         return bool(examples)
 
     @strawberry.field(description="Invocation parameters for the span")  # type: ignore
-    async def invocation_parameters(self, info: Info[Context, None]) -> list[InvocationParameter]:
+    async def invocation_parameters(
+        self,
+        info: Info[Context, None],
+    ) -> list[InvocationParameter]:
         from phoenix.server.api.helpers.playground_clients import OpenAIStreamingClient
         from phoenix.server.api.helpers.playground_registry import PLAYGROUND_CLIENT_REGISTRY
 
-        db_span = self.db_span
+        db_span: models.Span = await info.context.data_loaders.span_by_id.load(self.span_rowid)
         attributes = db_span.attributes
         llm_provider = GenerativeProvider.get_model_provider_from_attributes(attributes)
         if llm_provider is None:
@@ -288,68 +607,6 @@ class Span(Node):
         ]
 
 
-def to_gql_span(span: models.Span) -> Span:
-    events: list[SpanEvent] = list(map(SpanEvent.from_dict, span.events))
-    input_value = get_attribute_value(span.attributes, INPUT_VALUE)
-    if input_value is not None:
-        input_value = str(input_value)
-    assert input_value is None or isinstance(input_value, str)
-    output_value = get_attribute_value(span.attributes, OUTPUT_VALUE)
-    if output_value is not None:
-        output_value = str(output_value)
-    assert output_value is None or isinstance(output_value, str)
-    retrieval_documents = get_attribute_value(span.attributes, RETRIEVAL_DOCUMENTS)
-    num_documents = len(retrieval_documents) if isinstance(retrieval_documents, Sized) else None
-    return Span(
-        id_attr=span.id,
-        db_span=span,
-        name=span.name,
-        status_code=SpanStatusCode(span.status_code),
-        status_message=span.status_message,
-        parent_id=cast(Optional[ID], span.parent_id),
-        span_kind=SpanKind(span.span_kind),
-        start_time=span.start_time,
-        end_time=span.end_time,
-        latency_ms=span.latency_ms,
-        context=SpanContext(
-            trace_id=cast(ID, span.trace.trace_id),
-            span_id=cast(ID, span.span_id),
-        ),
-        attributes=json.dumps(_hide_embedding_vectors(span.attributes), cls=_JSONEncoder),
-        metadata=_convert_metadata_to_string(get_attribute_value(span.attributes, METADATA)),
-        num_documents=num_documents,
-        token_count_total=span.llm_token_count_total,
-        token_count_prompt=span.llm_token_count_prompt,
-        token_count_completion=span.llm_token_count_completion,
-        cumulative_token_count_total=span.cumulative_llm_token_count_prompt
-        + span.cumulative_llm_token_count_completion,
-        cumulative_token_count_prompt=span.cumulative_llm_token_count_prompt,
-        cumulative_token_count_completion=span.cumulative_llm_token_count_completion,
-        propagated_status_code=(
-            SpanStatusCode.ERROR
-            if span.cumulative_error_count
-            else SpanStatusCode(span.status_code)
-        ),
-        events=events,
-        input=(
-            SpanIOValue(
-                mime_type=MimeType(get_attribute_value(span.attributes, INPUT_MIME_TYPE)),
-                value=input_value,
-            )
-            if input_value is not None
-            else None
-        ),
-        output=(
-            SpanIOValue(
-                mime_type=MimeType(get_attribute_value(span.attributes, OUTPUT_MIME_TYPE)),
-                value=output_value,
-            )
-            if output_value is not None
-            else None
-        ),
-    )
-
-
 def _hide_embedding_vectors(attributes: Mapping[str, Any]) -> Mapping[str, Any]:
     if not (
         isinstance(em := attributes.get("embedding"), dict)

phoenix/server/api/types/SpanIOValue.py

@@ -1,15 +1,48 @@
+from typing import Any
+
 import strawberry
+from sqlalchemy.orm import QueryableAttribute
+from strawberry import UNSET, Info
+from typing_extensions import TypeAlias
 
+from phoenix.server.api.context import Context
 from phoenix.server.api.types.MimeType import MimeType
 
+SpanRowId: TypeAlias = int
+
 
 @strawberry.type
 class SpanIOValue:
+    span_rowid: strawberry.Private[SpanRowId] = UNSET
+    attr: strawberry.Private[QueryableAttribute[Any]] = UNSET
+    cached_value: strawberry.Private[str] = UNSET
     mime_type: MimeType
-    value: str
+    truncated_value: str = strawberry.field(
+        default=UNSET,
+        description="Truncated value up to 100 characters, appending '...' if truncated.",
+    )
+
+    def __post_init__(self) -> None:
+        if self.cached_value is not UNSET:
+            self.truncated_value = truncate_value(self.cached_value)
+        elif self.span_rowid is UNSET or self.attr is UNSET or self.truncated_value is UNSET:
+            raise ValueError(
+                "SpanIOValue must be initialized with either 'cached_value' or "
+                "'truncated_value' and 'id_' and 'attr'."
+            )
+
+    @strawberry.field
+    async def value(
+        self,
+        info: Info[Context, None],
+    ) -> str:
+        if self.cached_value is not UNSET:
+            return self.cached_value
+        if not self.truncated_value:
+            return ""
+        io_value = await info.context.data_loaders.span_fields.load((self.span_rowid, self.attr))
+        return "" if io_value is None else str(io_value)
+
 
-    @strawberry.field(
-        description="Truncate value up to `chars` characters, appending '...' if truncated.",
-    )  # type: ignore
-    def truncated_value(self, chars: int = 100) -> str:
-        return f"{self.value[: max(0, chars - 3)]}..." if len(self.value) > chars else self.value
+def truncate_value(value: str, chars: int = 100) -> str:
+    return f"{value[: max(0, chars - 3)]}..." if len(value) > chars else value