arize-phoenix 5.5.2__py3-none-any.whl → 5.7.0__py3-none-any.whl

phoenix/server/api/routers/v1/spans.py

@@ -1,5 +1,6 @@
+from collections.abc import AsyncIterator
 from datetime import datetime, timezone
-from typing import Any, AsyncIterator, Dict, List, Literal, Optional
+from typing import Any, Literal, Optional
 
 from fastapi import APIRouter, HTTPException, Query
 from pydantic import Field
@@ -28,16 +29,16 @@ router = APIRouter(tags=["spans"])
 
 
 class SpanQuery(V1RoutesBaseModel):
-    select: Optional[Dict[str, Any]] = None
-    filter: Optional[Dict[str, Any]] = None
-    explode: Optional[Dict[str, Any]] = None
-    concat: Optional[Dict[str, Any]] = None
-    rename: Optional[Dict[str, Any]] = None
-    index: Optional[Dict[str, Any]] = None
+    select: Optional[dict[str, Any]] = None
+    filter: Optional[dict[str, Any]] = None
+    explode: Optional[dict[str, Any]] = None
+    concat: Optional[dict[str, Any]] = None
+    rename: Optional[dict[str, Any]] = None
+    index: Optional[dict[str, Any]] = None
 
 
 class QuerySpansRequestBody(V1RoutesBaseModel):
-    queries: List[SpanQuery]
+    queries: list[SpanQuery]
     start_time: Optional[datetime] = None
     end_time: Optional[datetime] = None
     limit: int = DEFAULT_SPAN_LIMIT
@@ -153,7 +154,7 @@ class SpanAnnotation(V1RoutesBaseModel):
     result: Optional[SpanAnnotationResult] = Field(
         default=None, description="The result of the annotation"
     )
-    metadata: Optional[Dict[str, Any]] = Field(
+    metadata: Optional[dict[str, Any]] = Field(
         default=None, description="Metadata for the annotation"
     )
 
@@ -171,15 +172,15 @@ class SpanAnnotation(V1RoutesBaseModel):
         )
 
 
-class AnnotateSpansRequestBody(RequestBody[List[SpanAnnotation]]):
-    data: List[SpanAnnotation]
+class AnnotateSpansRequestBody(RequestBody[list[SpanAnnotation]]):
+    data: list[SpanAnnotation]
 
 
 class InsertedSpanAnnotation(V1RoutesBaseModel):
     id: str = Field(description="The ID of the inserted span annotation")
 
 
-class AnnotateSpansResponseBody(ResponseBody[List[InsertedSpanAnnotation]]):
+class AnnotateSpansResponseBody(ResponseBody[list[InsertedSpanAnnotation]]):
     pass
 
 

phoenix/server/api/routers/v1/traces.py

@@ -1,6 +1,6 @@
 import gzip
 import zlib
-from typing import Any, Dict, List, Literal, Optional
+from typing import Any, Literal, Optional
 
 from fastapi import APIRouter, BackgroundTasks, Header, HTTPException, Query
 from google.protobuf.message import DecodeError
@@ -110,7 +110,7 @@ class TraceAnnotation(V1RoutesBaseModel):
     result: Optional[TraceAnnotationResult] = Field(
         default=None, description="The result of the annotation"
     )
-    metadata: Optional[Dict[str, Any]] = Field(
+    metadata: Optional[dict[str, Any]] = Field(
         default=None, description="Metadata for the annotation"
     )
 
@@ -128,15 +128,15 @@ class TraceAnnotation(V1RoutesBaseModel):
         )
 
 
-class AnnotateTracesRequestBody(RequestBody[List[TraceAnnotation]]):
-    data: List[TraceAnnotation] = Field(description="The trace annotations to be upserted")
+class AnnotateTracesRequestBody(RequestBody[list[TraceAnnotation]]):
+    data: list[TraceAnnotation] = Field(description="The trace annotations to be upserted")
 
 
 class InsertedTraceAnnotation(V1RoutesBaseModel):
     id: str = Field(description="The ID of the inserted trace annotation")
 
 
-class AnnotateTracesResponseBody(ResponseBody[List[InsertedTraceAnnotation]]):
+class AnnotateTracesResponseBody(ResponseBody[list[InsertedTraceAnnotation]]):
     pass
 
 

phoenix/server/api/routers/v1/utils.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, Generic, List, Optional, TypedDict, TypeVar, Union
+from typing import Any, Generic, Optional, TypedDict, TypeVar, Union
 
 from typing_extensions import TypeAlias, assert_never
 
@@ -6,8 +6,8 @@ from .pydantic_compat import V1RoutesBaseModel
 
 StatusCode: TypeAlias = int
 DataType = TypeVar("DataType")
-Responses: TypeAlias = Dict[
-    Union[int, str], Dict[str, Any]
+Responses: TypeAlias = dict[
+    Union[int, str], dict[str, Any]
 ]  # input type for the `responses` parameter of a fastapi route
 
 
@@ -44,12 +44,12 @@ class PaginatedResponseBody(V1RoutesBaseModel, Generic[DataType]):
     # Don't use """ for this docstring or it will be included as a description
     # in the generated OpenAPI schema.
 
-    data: List[DataType]
+    data: list[DataType]
     next_cursor: Optional[str]
 
 
 def add_errors_to_responses(
-    errors: List[Union[StatusCode, StatusCodeWithDescription]],
+    errors: list[Union[StatusCode, StatusCodeWithDescription]],
     /,
     *,
     responses: Optional[Responses] = None,

phoenix/server/api/schema.py

@@ -1,17 +1,49 @@
+from itertools import chain
+from typing import Any, Iterable, Iterator, Optional, Union
+
 import strawberry
+from strawberry.extensions import SchemaExtension
+from strawberry.types.base import StrawberryObjectDefinition, StrawberryType
 
 from phoenix.server.api.exceptions import get_mask_errors_extension
 from phoenix.server.api.mutations import Mutation
 from phoenix.server.api.queries import Query
 from phoenix.server.api.subscriptions import Subscription
-
-# This is the schema for generating `schema.graphql`.
-# See https://strawberry.rocks/docs/guides/schema-export
-# It should be kept in sync with the server's runtime-initialized
-# instance. To do so, search for the usage of `strawberry.Schema(...)`.
-schema = strawberry.Schema(
-    query=Query,
-    mutation=Mutation,
-    extensions=[get_mask_errors_extension()],
-    subscription=Subscription,
+from phoenix.server.api.types.ChatCompletionSubscriptionPayload import (
+    ChatCompletionSubscriptionPayload,
 )
+
+
+def build_graphql_schema(
+    extensions: Optional[Iterable[Union[type[SchemaExtension], SchemaExtension]]] = None,
+) -> strawberry.Schema:
+    """
+    Builds a strawberry schema.
+    """
+    return strawberry.Schema(
+        query=Query,
+        mutation=Mutation,
+        extensions=chain(extensions or [], [get_mask_errors_extension()]),
+        subscription=Subscription,
+        types=_implementing_types(ChatCompletionSubscriptionPayload),
+    )
+
+
+def _implementing_types(interface: Any) -> Iterator[StrawberryType]:
+    """
+    Iterates over strawberry types implementing the given strawberry interface.
+    """
+    assert isinstance(
+        strawberry_definition := getattr(interface, "__strawberry_definition__", None),
+        StrawberryObjectDefinition,
+    )
+    assert strawberry_definition.is_interface
+    for subcls in interface.__subclasses__():
+        if isinstance(
+            getattr(subcls, "__strawberry_definition__", None),
+            StrawberryObjectDefinition,
+        ):
+            yield subcls
+
+
+_EXPORTED_GRAPHQL_SCHEMA = build_graphql_schema()  # used to export the GraphQL schema to file