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

phoenix/server/api/types/Span.py

@@ -1,36 +1,41 @@
 import json
-from collections import defaultdict
+from collections.abc import Mapping, Sized
 from datetime import datetime
 from enum import Enum
-from typing import Any, DefaultDict, Dict, List, Mapping, Optional, Sized, cast
+from typing import TYPE_CHECKING, Any, Optional, cast
 
 import numpy as np
 import strawberry
-from openinference.semconv.trace import EmbeddingAttributes, SpanAttributes
+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
 
 import phoenix.trace.schemas as trace_schema
-from phoenix.config import DEFAULT_PROJECT_NAME
-from phoenix.core.project import Project, WrappedSpan
-from phoenix.metrics.retrieval_metrics import RetrievalMetrics
+from phoenix.db import models
 from phoenix.server.api.context import Context
+from phoenix.server.api.helpers.dataset_helpers import (
+    get_dataset_example_input,
+    get_dataset_example_output,
+)
+from phoenix.server.api.input_types.InvocationParameters import InvocationParameter
+from phoenix.server.api.input_types.SpanAnnotationSort import (
+    SpanAnnotationColumn,
+    SpanAnnotationSort,
+)
 from phoenix.server.api.types.DocumentRetrievalMetrics import DocumentRetrievalMetrics
-from phoenix.server.api.types.Evaluation import DocumentEvaluation, SpanEvaluation
+from phoenix.server.api.types.Evaluation import DocumentEvaluation
+from phoenix.server.api.types.ExampleRevisionInterface import ExampleRevision
+from phoenix.server.api.types.GenerativeProvider import GenerativeProvider
 from phoenix.server.api.types.MimeType import MimeType
-from phoenix.trace.schemas import ComputedAttributes, SpanID
+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.trace.attributes import get_attribute_value
 
-EMBEDDING_EMBEDDINGS = SpanAttributes.EMBEDDING_EMBEDDINGS
-EMBEDDING_VECTOR = EmbeddingAttributes.EMBEDDING_VECTOR
-INPUT_MIME_TYPE = SpanAttributes.INPUT_MIME_TYPE
-INPUT_VALUE = SpanAttributes.INPUT_VALUE
-LLM_TOKEN_COUNT_COMPLETION = SpanAttributes.LLM_TOKEN_COUNT_COMPLETION
-LLM_TOKEN_COUNT_PROMPT = SpanAttributes.LLM_TOKEN_COUNT_PROMPT
-LLM_TOKEN_COUNT_TOTAL = SpanAttributes.LLM_TOKEN_COUNT_TOTAL
-METADATA = SpanAttributes.METADATA
-OUTPUT_MIME_TYPE = SpanAttributes.OUTPUT_MIME_TYPE
-OUTPUT_VALUE = SpanAttributes.OUTPUT_VALUE
-RETRIEVAL_DOCUMENTS = SpanAttributes.RETRIEVAL_DOCUMENTS
+if TYPE_CHECKING:
+    from phoenix.server.api.types.Project import Project
 
 
 @strawberry.enum
@@ -41,18 +46,22 @@ class SpanKind(Enum):
     NB: this is actively under construction
     """
 
-    chain = trace_schema.SpanKind.CHAIN
-    tool = trace_schema.SpanKind.TOOL
-    llm = trace_schema.SpanKind.LLM
-    retriever = trace_schema.SpanKind.RETRIEVER
-    embedding = trace_schema.SpanKind.EMBEDDING
-    agent = trace_schema.SpanKind.AGENT
-    reranker = trace_schema.SpanKind.RERANKER
-    unknown = trace_schema.SpanKind.UNKNOWN
+    chain = "CHAIN"
+    tool = "TOOL"
+    llm = "LLM"
+    retriever = "RETRIEVER"
+    embedding = "EMBEDDING"
+    agent = "AGENT"
+    reranker = "RERANKER"
+    evaluator = "EVALUATOR"
+    guardrail = "GUARDRAIL"
+    unknown = "UNKNOWN"
 
     @classmethod
     def _missing_(cls, v: Any) -> Optional["SpanKind"]:
-        return None if v else cls.unknown
+        if v and isinstance(v, str) and v.isascii() and not v.isupper():
+            return cls(v.upper())
+        return cls.unknown
 
 
 @strawberry.type
@@ -61,17 +70,11 @@ class SpanContext:
     span_id: ID
 
 
-@strawberry.type
-class SpanIOValue:
-    mime_type: MimeType
-    value: str
-
-
 @strawberry.enum
 class SpanStatusCode(Enum):
-    OK = trace_schema.SpanStatusCode.OK
-    ERROR = trace_schema.SpanStatusCode.ERROR
-    UNSET = trace_schema.SpanStatusCode.UNSET
+    OK = "OK"
+    ERROR = "ERROR"
+    UNSET = "UNSET"
 
     @classmethod
     def _missing_(cls, v: Any) -> Optional["SpanStatusCode"]:
@@ -85,19 +88,24 @@ class SpanEvent:
     timestamp: datetime
 
     @staticmethod
-    def from_event(
-        event: trace_schema.SpanEvent,
+    def from_dict(
+        event: Mapping[str, Any],
     ) -> "SpanEvent":
         return SpanEvent(
-            name=event.name,
-            message=cast(str, event.attributes.get(trace_schema.EXCEPTION_MESSAGE) or ""),
-            timestamp=event.timestamp,
+            name=event["name"],
+            message=cast(str, event["attributes"].get(trace_schema.EXCEPTION_MESSAGE) or ""),
+            timestamp=datetime.fromisoformat(event["timestamp"]),
         )
 
 
 @strawberry.type
-class Span:
-    project: strawberry.Private[Project]
+class SpanAsExampleRevision(ExampleRevision): ...
+
+
+@strawberry.type
+class Span(Node):
+    id_attr: NodeID[int]
+    db_span: strawberry.Private[models.Span]
     name: str
     status_code: SpanStatusCode
     status_message: str
@@ -121,7 +129,7 @@ class Span:
     token_count_completion: Optional[int]
     input: Optional[SpanIOValue]
     output: Optional[SpanIOValue]
-    events: List[SpanEvent]
+    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.)",
@@ -140,23 +148,27 @@ class Span:
     )
 
     @strawberry.field(
-        description="Evaluations associated with the span, e.g. if the span is "
-        "an LLM, an evaluation may assess the helpfulness of its response with "
-        "respect to its input."
+        description=(
+            "Annotations associated with the span. This encompasses both "
+            "LLM and human annotations."
+        )
     )  # type: ignore
-    def span_evaluations(
+    async def span_annotations(
         self,
         info: Info[Context, None],
-    ) -> List[SpanEvaluation]:
-        if not (traces := info.context.traces) or not (
-            project := traces.get_project(DEFAULT_PROJECT_NAME)
-        ):
-            return []
-        span_id = SpanID(str(self.context.span_id))
-        return [
-            SpanEvaluation.from_pb_evaluation(evaluation)
-            for evaluation in project.get_evaluations_by_span_id(span_id)
-        ]
+        sort: Optional[SpanAnnotationSort] = UNSET,
+    ) -> list[SpanAnnotation]:
+        span_id = self.id_attr
+        annotations = await info.context.data_loaders.span_annotations.load(span_id)
+        sort_key = SpanAnnotationColumn.name.value
+        sort_descending = False
+        if sort:
+            sort_key = sort.col.value
+            sort_descending = sort.dir is SortDir.desc
+        annotations.sort(
+            key=lambda annotation: getattr(annotation, sort_key), reverse=sort_descending
+        )
+        return [to_gql_span_annotation(annotation) for annotation in annotations]
 
     @strawberry.field(
         description="Evaluations of the documents associated with the span, e.g. "
@@ -166,76 +178,131 @@ class Span:
         "a list, and each evaluation is identified by its document's (zero-based) "
         "index in that list."
     )  # type: ignore
-    def document_evaluations(
-        self,
-        info: Info[Context, None],
-    ) -> List[DocumentEvaluation]:
-        if not (traces := info.context.traces) or not (
-            project := traces.get_project(DEFAULT_PROJECT_NAME)
-        ):
-            return []
-        span_id = SpanID(str(self.context.span_id))
-        return [
-            DocumentEvaluation.from_pb_evaluation(evaluation)
-            for evaluation in project.get_document_evaluations_by_span_id(span_id)
-        ]
+    async def document_evaluations(self, info: Info[Context, None]) -> list[DocumentEvaluation]:
+        return await info.context.data_loaders.document_evaluations.load(self.id_attr)
 
     @strawberry.field(
         description="Retrieval metrics: NDCG@K, Precision@K, Reciprocal Rank, etc.",
     )  # type: ignore
-    def document_retrieval_metrics(
+    async def document_retrieval_metrics(
         self,
         info: Info[Context, None],
         evaluation_name: Optional[str] = UNSET,
-    ) -> List[DocumentRetrievalMetrics]:
+    ) -> list[DocumentRetrievalMetrics]:
         if not self.num_documents:
             return []
-        span_id = SpanID(str(self.context.span_id))
-        all_document_evaluation_names = self.project.get_document_evaluation_names(span_id)
-        if not all_document_evaluation_names:
-            return []
-        if evaluation_name is UNSET:
-            evaluation_names = all_document_evaluation_names
-        elif evaluation_name not in all_document_evaluation_names:
-            return []
-        else:
-            evaluation_names = [evaluation_name]
-        retrieval_metrics = []
-        for name in evaluation_names:
-            evaluation_scores = self.project.get_document_evaluation_scores(
-                span_id=span_id,
-                evaluation_name=name,
-                num_documents=self.num_documents,
-            )
-            retrieval_metrics.append(
-                DocumentRetrievalMetrics(
-                    evaluation_name=name,
-                    metrics=RetrievalMetrics(evaluation_scores),
-                )
-            )
-        return retrieval_metrics
+        return await info.context.data_loaders.document_retrieval_metrics.load(
+            (self.id_attr, evaluation_name or None, self.num_documents),
+        )
 
     @strawberry.field(
         description="All descendant spans (children, grandchildren, etc.)",
     )  # type: ignore
-    def descendants(
+    async def descendants(
+        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]
+
+    @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
+
+        # Fetch annotations associated with this span
+        span_annotations = await self.span_annotations(info)
+        annotations = dict()
+        for annotation in span_annotations:
+            annotations[annotation.name] = {
+                "label": annotation.label,
+                "score": annotation.score,
+                "explanation": annotation.explanation,
+                "metadata": annotation.metadata,
+                "annotator_kind": annotation.annotator_kind.value,
+            }
+        # Merge annotations into the metadata
+        metadata = {
+            "span_kind": span.span_kind,
+            **({"annotations": annotations} if annotations else {}),
+        }
+
+        return SpanAsExampleRevision(
+            input=get_dataset_example_input(span),
+            output=get_dataset_example_output(span),
+            metadata=metadata,
+        )
+
+    @strawberry.field(description="The project that this span belongs to.")  # type: ignore
+    async def project(
         self,
         info: Info[Context, None],
-    ) -> List["Span"]:
+    ) -> Annotated[
+        "Project", strawberry.lazy("phoenix.server.api.types.Project")
+    ]:  # 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
+        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)
+        return bool(examples)
+
+    @strawberry.field(description="Invocation parameters for the span")  # type: ignore
+    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
+        attributes = db_span.attributes
+        llm_provider = GenerativeProvider.get_model_provider_from_attributes(attributes)
+        if llm_provider is None:
+            return []
+        llm_model = get_attribute_value(attributes, SpanAttributes.LLM_MODEL_NAME)
+        invocation_parameters = get_attribute_value(
+            attributes, SpanAttributes.LLM_INVOCATION_PARAMETERS
+        )
+        if invocation_parameters is None:
+            return []
+        invocation_parameters = json.loads(invocation_parameters)
+        # find the client class for the provider, if there is no client class or provider,
+        # return openai as default
+        client_class = PLAYGROUND_CLIENT_REGISTRY.get_client(llm_provider, llm_model)
+        if not client_class:
+            client_class = OpenAIStreamingClient
+        supported_invocation_parameters = client_class.supported_invocation_parameters()
+        # filter supported invocation parameters down to those whose canonical_name is in the
+        # invocation_parameters keys
         return [
-            to_gql_span(span, self.project)
-            for span in self.project.get_descendant_spans(SpanID(self.context.span_id))
+            ip
+            for ip in supported_invocation_parameters
+            if (
+                ip.canonical_name in invocation_parameters
+                or ip.invocation_name in invocation_parameters
+            )
         ]
 
 
-def to_gql_span(span: WrappedSpan, project: Project) -> "Span":
-    events: List[SpanEvent] = list(map(SpanEvent.from_event, span.events))
-    input_value = cast(Optional[str], span.attributes.get(INPUT_VALUE))
-    output_value = cast(Optional[str], span.attributes.get(OUTPUT_VALUE))
-    retrieval_documents = span.attributes.get(RETRIEVAL_DOCUMENTS)
+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(
-        project=project,
+        id_attr=span.id,
+        db_span=span,
         name=span.name,
         status_code=SpanStatusCode(span.status_code),
         status_message=span.status_message,
@@ -243,50 +310,30 @@ def to_gql_span(span: WrappedSpan, project: Project) -> "Span":
         span_kind=SpanKind(span.span_kind),
         start_time=span.start_time,
         end_time=span.end_time,
-        latency_ms=cast(Optional[float], span[ComputedAttributes.LATENCY_MS]),
+        latency_ms=span.latency_ms,
         context=SpanContext(
-            trace_id=cast(ID, span.context.trace_id),
-            span_id=cast(ID, span.context.span_id),
+            trace_id=cast(ID, span.trace.trace_id),
+            span_id=cast(ID, span.span_id),
         ),
-        attributes=json.dumps(
-            _nested_attributes(_hide_embedding_vectors(span.attributes)),
-            cls=_JSONEncoder,
-        ),
-        metadata=_convert_metadata_to_string(span.attributes.get(METADATA)),
+        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=cast(
-            Optional[int],
-            span.attributes.get(LLM_TOKEN_COUNT_TOTAL),
-        ),
-        token_count_prompt=cast(
-            Optional[int],
-            span.attributes.get(LLM_TOKEN_COUNT_PROMPT),
-        ),
-        token_count_completion=cast(
-            Optional[int],
-            span.attributes.get(LLM_TOKEN_COUNT_COMPLETION),
-        ),
-        cumulative_token_count_total=cast(
-            Optional[int],
-            span[ComputedAttributes.CUMULATIVE_LLM_TOKEN_COUNT_TOTAL],
-        ),
-        cumulative_token_count_prompt=cast(
-            Optional[int],
-            span[ComputedAttributes.CUMULATIVE_LLM_TOKEN_COUNT_PROMPT],
-        ),
-        cumulative_token_count_completion=cast(
-            Optional[int],
-            span[ComputedAttributes.CUMULATIVE_LLM_TOKEN_COUNT_COMPLETION],
-        ),
+        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[ComputedAttributes.CUMULATIVE_ERROR_COUNT]
+            if span.cumulative_error_count
             else SpanStatusCode(span.status_code)
         ),
         events=events,
         input=(
             SpanIOValue(
-                mime_type=MimeType(span.attributes.get(INPUT_MIME_TYPE)),
+                mime_type=MimeType(get_attribute_value(span.attributes, INPUT_MIME_TYPE)),
                 value=input_value,
             )
             if input_value is not None
@@ -294,7 +341,7 @@ def to_gql_span(span: WrappedSpan, project: Project) -> "Span":
         ),
         output=(
             SpanIOValue(
-                mime_type=MimeType(span.attributes.get(OUTPUT_MIME_TYPE)),
+                mime_type=MimeType(get_attribute_value(span.attributes, OUTPUT_MIME_TYPE)),
                 value=output_value,
             )
             if output_value is not None
@@ -303,6 +350,29 @@ def to_gql_span(span: WrappedSpan, project: Project) -> "Span":
     )
 
 
+def _hide_embedding_vectors(attributes: Mapping[str, Any]) -> Mapping[str, Any]:
+    if not (
+        isinstance(em := attributes.get("embedding"), dict)
+        and isinstance(embeddings := em.get("embeddings"), list)
+        and embeddings
+    ):
+        return attributes
+    embeddings = embeddings.copy()
+    for i, embedding in enumerate(embeddings):
+        if not (
+            isinstance(embedding, dict)
+            and isinstance(emb := embedding.get("embedding"), dict)
+            and isinstance(vector := emb.get("vector"), list)
+            and vector
+        ):
+            continue
+        embeddings[i] = {
+            **embedding,
+            "embedding": {**emb, "vector": f"<{len(vector)} dimensional vector>"},
+        }
+    return {**attributes, "embedding": {**em, "embeddings": embeddings}}
+
+
 class _JSONEncoder(json.JSONEncoder):
     def default(self, obj: Any) -> Any:
         if isinstance(obj, datetime):
@@ -318,39 +388,6 @@ class _JSONEncoder(json.JSONEncoder):
         return super().default(obj)
 
 
-def _trie() -> DefaultDict[str, Any]:
-    return defaultdict(_trie)
-
-
-def _nested_attributes(
-    attributes: Mapping[str, Any],
-) -> DefaultDict[str, Any]:
-    nested_attributes = _trie()
-    for attribute_name, attribute_value in attributes.items():
-        trie = nested_attributes
-        keys = attribute_name.split(".")
-        for key in keys[:-1]:
-            trie = trie[key]
-        trie[keys[-1]] = attribute_value
-    return nested_attributes
-
-
-def _hide_embedding_vectors(
-    attributes: Mapping[str, Any],
-) -> Dict[str, Any]:
-    _attributes = dict(attributes)
-    if not (embeddings := _attributes.get(EMBEDDING_EMBEDDINGS)):
-        return _attributes
-    _embeddings = []
-    for embedding in embeddings:
-        _embedding = dict(embedding)
-        if vector := _embedding.get(EMBEDDING_VECTOR):
-            _embedding[EMBEDDING_VECTOR] = f"<{len(vector)} dimensional vector>"
-        _embeddings.append(_embedding)
-    _attributes[EMBEDDING_EMBEDDINGS] = _embeddings
-    return _attributes
-
-
 def _convert_metadata_to_string(metadata: Any) -> Optional[str]:
     """
     Converts metadata to a string representation.
@@ -362,3 +399,11 @@ def _convert_metadata_to_string(metadata: Any) -> Optional[str]:
         return json.dumps(metadata)
     except Exception:
         return str(metadata)
+
+
+INPUT_MIME_TYPE = SpanAttributes.INPUT_MIME_TYPE
+INPUT_VALUE = SpanAttributes.INPUT_VALUE
+METADATA = SpanAttributes.METADATA
+OUTPUT_MIME_TYPE = SpanAttributes.OUTPUT_MIME_TYPE
+OUTPUT_VALUE = SpanAttributes.OUTPUT_VALUE
+RETRIEVAL_DOCUMENTS = SpanAttributes.RETRIEVAL_DOCUMENTS

phoenix/server/api/types/SpanAnnotation.py

@@ -0,0 +1,43 @@
+from typing import Optional
+
+import strawberry
+from strawberry import Private
+from strawberry.relay import GlobalID, Node, NodeID
+from strawberry.scalars import JSON
+
+from phoenix.db import models
+
+from .Annotation import Annotation
+from .AnnotatorKind import AnnotatorKind
+
+
+@strawberry.type
+class SpanAnnotation(Node, Annotation):
+    id_attr: NodeID[int]
+    annotator_kind: AnnotatorKind
+    metadata: JSON
+    span_rowid: Private[Optional[int]]
+
+    @strawberry.field
+    async def span_id(self) -> GlobalID:
+        from phoenix.server.api.types.Span import Span
+
+        return GlobalID(type_name=Span.__name__, node_id=str(self.span_rowid))
+
+
+def to_gql_span_annotation(
+    annotation: models.SpanAnnotation,
+) -> SpanAnnotation:
+    """
+    Converts an ORM span annotation to a GraphQL SpanAnnotation.
+    """
+    return SpanAnnotation(
+        id_attr=annotation.id,
+        span_rowid=annotation.span_rowid,
+        name=annotation.name,
+        annotator_kind=AnnotatorKind(annotation.annotator_kind),
+        label=annotation.label,
+        score=annotation.score,
+        explanation=annotation.explanation,
+        metadata=annotation.metadata_,
+    )

phoenix/server/api/types/SpanIOValue.py

@@ -0,0 +1,15 @@
+import strawberry
+
+from phoenix.server.api.types.MimeType import MimeType
+
+
+@strawberry.type
+class SpanIOValue:
+    mime_type: MimeType
+    value: str
+
+    @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

phoenix/server/api/types/SystemApiKey.py

@@ -0,0 +1,9 @@
+import strawberry
+from strawberry.relay import Node, NodeID
+
+from .ApiKey import ApiKey
+
+
+@strawberry.type
+class SystemApiKey(ApiKey, Node):
+    id_attr: NodeID[int]

phoenix/server/api/types/TemplateLanguage.py

@@ -0,0 +1,10 @@
+from enum import Enum
+
+import strawberry
+
+
+@strawberry.enum
+class TemplateLanguage(Enum):
+    NONE = "NONE"
+    MUSTACHE = "MUSTACHE"
+    F_STRING = "F_STRING"