arize-phoenix 5.6.0__py3-none-any.whl → 5.8.0__py3-none-any.whl

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