arize-phoenix 12.7.1__py3-none-any.whl → 12.9.0__py3-none-any.whl

phoenix/server/api/types/PromptLabel.py

@@ -2,23 +2,57 @@ from typing import Optional
 
 import strawberry
 from strawberry.relay import Node, NodeID
+from strawberry.types import Info
 
 from phoenix.db import models
+from phoenix.server.api.context import Context
 from phoenix.server.api.types.Identifier import Identifier
 
 
 @strawberry.type
 class PromptLabel(Node):
-    id_attr: NodeID[int]
-    name: Identifier
-    description: Optional[str] = None
-    color: str
-
-
-def to_gql_prompt_label(label_orm: models.PromptLabel) -> PromptLabel:
-    return PromptLabel(
-        id_attr=label_orm.id,
-        name=Identifier(label_orm.name),
-        description=label_orm.description,
-        color=label_orm.color,
-    )
+    id: NodeID[int]
+    db_record: strawberry.Private[Optional[models.PromptLabel]] = None
+
+    def __post_init__(self) -> None:
+        if self.db_record and self.id != self.db_record.id:
+            raise ValueError("PromptLabel ID mismatch")
+
+    @strawberry.field
+    async def name(
+        self,
+        info: Info[Context, None],
+    ) -> Identifier:
+        if self.db_record:
+            val = self.db_record.name
+        else:
+            val = await info.context.data_loaders.prompt_label_fields.load(
+                (self.id, models.PromptLabel.name),
+            )
+        return Identifier(val)
+
+    @strawberry.field
+    async def description(
+        self,
+        info: Info[Context, None],
+    ) -> Optional[str]:
+        if self.db_record:
+            val = self.db_record.description
+        else:
+            val = await info.context.data_loaders.prompt_label_fields.load(
+                (self.id, models.PromptLabel.description),
+            )
+        return val
+
+    @strawberry.field
+    async def color(
+        self,
+        info: Info[Context, None],
+    ) -> str:
+        if self.db_record:
+            val = self.db_record.color
+        else:
+            val = await info.context.data_loaders.prompt_label_fields.load(
+                (self.id, models.PromptLabel.color),
+            )
+        return val

phoenix/server/api/types/PromptVersion.py

@@ -1,5 +1,5 @@
 from datetime import datetime
-from typing import Optional
+from typing import TYPE_CHECKING, Annotated, Optional
 
 import strawberry
 from sqlalchemy import select
@@ -18,7 +18,7 @@ from phoenix.server.api.helpers.prompts.models import (
     denormalize_tools,
     get_raw_invocation_parameters,
 )
-from phoenix.server.api.types.PromptVersionTag import PromptVersionTag, to_gql_prompt_version_tag
+from phoenix.server.api.types.PromptVersionTag import PromptVersionTag
 from phoenix.server.api.types.PromptVersionTemplate import (
     PromptTemplate,
     to_gql_template_from_orm,
@@ -26,7 +26,9 @@ from phoenix.server.api.types.PromptVersionTemplate import (
 
 from .ResponseFormat import ResponseFormat
 from .ToolDefinition import ToolDefinition
-from .User import User, to_gql_user
+
+if TYPE_CHECKING:
+    from .User import User
 
 
 @strawberry.type
@@ -53,16 +55,17 @@ class PromptVersion(Node):
                 models.PromptVersionTag.prompt_version_id == self.id_attr
             )
             return [
-                to_gql_prompt_version_tag(tag) async for tag in await session.stream_scalars(stmt)
+                PromptVersionTag(id=tag.id, db_record=tag)
+                async for tag in await session.stream_scalars(stmt)
             ]
 
     @strawberry.field
-    async def user(self, info: Info[Context, None]) -> Optional[User]:
+    async def user(self) -> Optional[Annotated["User", strawberry.lazy(".User")]]:
         if self.user_id is None:
             return None
-        async with info.context.db() as session:
-            user = await session.get(models.User, self.user_id)
-        return to_gql_user(user) if user is not None else None
+        from .User import User
+
+        return User(id=self.user_id)
 
     @strawberry.field
     async def previous_version(self, info: Info[Context, None]) -> Optional["PromptVersion"]:

phoenix/server/api/types/PromptVersionTag.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from typing import TYPE_CHECKING, Annotated, Optional
 
 import strawberry
 from strawberry import Info
@@ -7,34 +7,74 @@ from strawberry.relay import GlobalID, Node, NodeID
 from phoenix.db import models
 from phoenix.server.api.context import Context
 from phoenix.server.api.types.Identifier import Identifier
-from phoenix.server.api.types.User import User, to_gql_user
+from phoenix.server.api.types.User import User
+
+if TYPE_CHECKING:
+    from .User import User
 
 
 @strawberry.type
 class PromptVersionTag(Node):
-    id_attr: NodeID[int]
-    user_id: strawberry.Private[Optional[int]]
-    prompt_version_id: GlobalID
-    name: Identifier
-    description: Optional[str] = None
+    id: NodeID[int]
+    db_record: strawberry.Private[Optional[models.PromptVersionTag]] = None
+
+    def __post_init__(self) -> None:
+        if self.db_record and self.id != self.db_record.id:
+            raise ValueError("PromptVersionTag ID mismatch")
+
+    @strawberry.field
+    async def prompt_version_id(
+        self,
+        info: Info[Context, None],
+    ) -> GlobalID:
+        from phoenix.server.api.types.PromptVersion import PromptVersion
+
+        if self.db_record:
+            version_id = self.db_record.prompt_version_id
+        else:
+            version_id = await info.context.data_loaders.prompt_version_tag_fields.load(
+                (self.id, models.PromptVersionTag.prompt_version_id),
+            )
+        return GlobalID(PromptVersion.__name__, str(version_id))
 
     @strawberry.field
-    async def user(self, info: Info[Context, None]) -> Optional[User]:
-        if self.user_id is None:
+    async def name(
+        self,
+        info: Info[Context, None],
+    ) -> Identifier:
+        if self.db_record:
+            val = self.db_record.name
+        else:
+            val = await info.context.data_loaders.prompt_version_tag_fields.load(
+                (self.id, models.PromptVersionTag.name),
+            )
+        return Identifier(val.root)
+
+    @strawberry.field
+    async def description(
+        self,
+        info: Info[Context, None],
+    ) -> Optional[str]:
+        if self.db_record:
+            val = self.db_record.description
+        else:
+            val = await info.context.data_loaders.prompt_version_tag_fields.load(
+                (self.id, models.PromptVersionTag.description),
+            )
+        return val
+
+    @strawberry.field
+    async def user(
+        self, info: Info[Context, None]
+    ) -> Optional[Annotated["User", strawberry.lazy(".User")]]:
+        if self.db_record:
+            user_id = self.db_record.user_id
+        else:
+            user_id = await info.context.data_loaders.prompt_version_tag_fields.load(
+                (self.id, models.PromptVersionTag.user_id),
+            )
+        if user_id is None:
             return None
-        async with info.context.db() as session:
-            user = await session.get(models.User, self.user_id)
-        return to_gql_user(user) if user is not None else None
-
-
-def to_gql_prompt_version_tag(prompt_version_tag: models.PromptVersionTag) -> PromptVersionTag:
-    from phoenix.server.api.types.PromptVersion import PromptVersion
-
-    version_gid = GlobalID(PromptVersion.__name__, str(prompt_version_tag.prompt_version_id))
-    return PromptVersionTag(
-        id_attr=prompt_version_tag.id,
-        prompt_version_id=version_gid,
-        name=Identifier(prompt_version_tag.name.root),
-        description=prompt_version_tag.description,
-        user_id=prompt_version_tag.user_id,
-    )
+        from .User import User
+
+        return User(id=user_id)