arize-phoenix 10.0.4__py3-none-any.whl → 12.28.1__py3-none-any.whl

phoenix/server/api/types/Prompt.py

@@ -6,9 +6,11 @@ import strawberry
 from sqlalchemy import func, select
 from strawberry import UNSET
 from strawberry.relay import Connection, GlobalID, Node, NodeID
+from strawberry.scalars import JSON
 from strawberry.types import Info
 
 from phoenix.db import models
+from phoenix.db.types.identifier import Identifier as IdentifierModel
 from phoenix.server.api.context import Context
 from phoenix.server.api.exceptions import NotFound
 from phoenix.server.api.types.Identifier import Identifier
@@ -19,24 +21,96 @@ from phoenix.server.api.types.pagination import (
     connection_from_list,
 )
 
+from .PromptLabel import PromptLabel
 from .PromptVersion import (
     PromptVersion,
     to_gql_prompt_version,
 )
-from .PromptVersionTag import PromptVersionTag, to_gql_prompt_version_tag
+from .PromptVersionTag import PromptVersionTag
 
 
 @strawberry.type
 class Prompt(Node):
-    id_attr: NodeID[int]
-    source_prompt_id: Optional[GlobalID]
-    name: Identifier
-    description: Optional[str]
-    created_at: datetime
+    id: NodeID[int]
+    db_record: strawberry.Private[Optional[models.Prompt]] = None
+
+    def __post_init__(self) -> None:
+        if self.db_record and self.id != self.db_record.id:
+            raise ValueError("Prompt ID mismatch")
+
+    @strawberry.field
+    async def source_prompt_id(
+        self,
+        info: Info[Context, None],
+    ) -> Optional[GlobalID]:
+        if self.db_record:
+            source_id = self.db_record.source_prompt_id
+        else:
+            source_id = await info.context.data_loaders.prompt_fields.load(
+                (self.id, models.Prompt.source_prompt_id),
+            )
+        if not source_id:
+            return None
+        return GlobalID(Prompt.__name__, str(source_id))
+
+    @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_fields.load(
+                (self.id, models.Prompt.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_fields.load(
+                (self.id, models.Prompt.description),
+            )
+        return val
+
+    @strawberry.field
+    async def metadata(
+        self,
+        info: Info[Context, None],
+    ) -> JSON:
+        if self.db_record:
+            val = self.db_record.metadata_
+        else:
+            val = await info.context.data_loaders.prompt_fields.load(
+                (self.id, models.Prompt.metadata_),
+            )
+        return val
+
+    @strawberry.field
+    async def created_at(
+        self,
+        info: Info[Context, None],
+    ) -> datetime:
+        if self.db_record:
+            val = self.db_record.created_at
+        else:
+            val = await info.context.data_loaders.prompt_fields.load(
+                (self.id, models.Prompt.created_at),
+            )
+        return val
 
     @strawberry.field
     async def version(
-        self, info: Info[Context, None], version_id: Optional[GlobalID] = None
+        self,
+        info: Info[Context, None],
+        version_id: Optional[GlobalID] = None,
+        tag_name: Optional[Identifier] = None,
     ) -> PromptVersion:
         async with info.context.db() as session:
             if version_id:
@@ -44,15 +118,28 @@ class Prompt(Node):
                 version = await session.scalar(
                     select(models.PromptVersion).where(
                         models.PromptVersion.id == v_id,
-                        models.PromptVersion.prompt_id == self.id_attr,
+                        models.PromptVersion.prompt_id == self.id,
                     )
                 )
                 if not version:
                     raise NotFound(f"Prompt version not found: {version_id}")
+            elif tag_name:
+                try:
+                    name = IdentifierModel(tag_name)
+                except ValueError:
+                    raise NotFound(f"Prompt version tag not found: {tag_name}")
+                version = await session.scalar(
+                    select(models.PromptVersion)
+                    .where(models.PromptVersion.prompt_id == self.id)
+                    .join_from(models.PromptVersion, models.PromptVersionTag)
+                    .where(models.PromptVersionTag.name == name)
+                )
+                if not version:
+                    raise NotFound(f"This prompt has no associated versions by tag {tag_name}")
             else:
                 stmt = (
                     select(models.PromptVersion)
-                    .where(models.PromptVersion.prompt_id == self.id_attr)
+                    .where(models.PromptVersion.prompt_id == self.id)
                     .order_by(models.PromptVersion.id.desc())
                     .limit(1)
                 )
@@ -65,10 +152,11 @@ class Prompt(Node):
     async def version_tags(self, info: Info[Context, None]) -> list[PromptVersionTag]:
         async with info.context.db() as session:
             stmt = select(models.PromptVersionTag).where(
-                models.PromptVersionTag.prompt_id == self.id_attr
+                models.PromptVersionTag.prompt_id == self.id
             )
             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
@@ -89,7 +177,7 @@ class Prompt(Node):
         row_number = func.row_number().over(order_by=models.PromptVersion.id).label("row_number")
         stmt = (
             select(models.PromptVersion, row_number)
-            .where(models.PromptVersion.prompt_id == self.id_attr)
+            .where(models.PromptVersion.prompt_id == self.id)
             .order_by(models.PromptVersion.id.desc())
         )
         async with info.context.db() as session:
@@ -101,34 +189,26 @@ class Prompt(Node):
 
     @strawberry.field
     async def source_prompt(self, info: Info[Context, None]) -> Optional["Prompt"]:
-        if not self.source_prompt_id:
+        if self.db_record:
+            id_ = self.db_record.source_prompt_id
+        else:
+            id_ = await info.context.data_loaders.prompt_fields.load(
+                (self.id, models.Prompt.source_prompt_id),
+            )
+        if not id_:
             return None
+        async with info.context.db() as session:
+            source_prompt = await session.get(models.Prompt, id_)
+        if not source_prompt:
+            raise NotFound(f"Source prompt not found: {id_}")
+        return Prompt(id=source_prompt.id, db_record=source_prompt)
 
-        source_prompt_id = from_global_id_with_expected_type(
-            global_id=self.source_prompt_id, expected_type_name=Prompt.__name__
-        )
-
+    @strawberry.field
+    async def labels(self, info: Info[Context, None]) -> list["PromptLabel"]:
         async with info.context.db() as session:
-            source_prompt = await session.scalar(
-                select(models.Prompt).where(models.Prompt.id == source_prompt_id)
+            labels = await session.scalars(
+                select(models.PromptLabel)
+                .join(models.PromptPromptLabel)
+                .where(models.PromptPromptLabel.prompt_id == self.id)
             )
-            if not source_prompt:
-                raise NotFound(f"Source prompt not found: {self.source_prompt_id}")
-            return to_gql_prompt_from_orm(source_prompt)
-
-
-def to_gql_prompt_from_orm(orm_model: "models.Prompt") -> Prompt:
-    if not orm_model.source_prompt_id:
-        source_prompt_gid = None
-    else:
-        source_prompt_gid = GlobalID(
-            Prompt.__name__,
-            str(orm_model.source_prompt_id),
-        )
-    return Prompt(
-        id_attr=orm_model.id,
-        source_prompt_id=source_prompt_gid,
-        name=Identifier(orm_model.name.root),
-        description=orm_model.description,
-        created_at=orm_model.created_at,
-    )
+            return [PromptLabel(id=label.id, db_record=label) for label in labels]

phoenix/server/api/types/PromptLabel.py

@@ -1,41 +1,58 @@
 from typing import Optional
 
 import strawberry
-from sqlalchemy import select
 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
-from phoenix.server.api.types.Prompt import Prompt, to_gql_prompt_from_orm
 
 
 @strawberry.type
 class PromptLabel(Node):
-    id_attr: NodeID[int]
-    name: Identifier
-    description: Optional[str] = None
+    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 prompts(self, info: Info[Context, None]) -> list[Prompt]:
-        async with info.context.db() as session:
-            statement = (
-                select(models.Prompt)
-                .join(
-                    models.PromptPromptLabel, models.Prompt.id == models.PromptPromptLabel.prompt_id
-                )
-                .where(models.PromptPromptLabel.prompt_label_id == self.id_attr)
+    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 [
-                to_gql_prompt_from_orm(prompt_orm)
-                async for prompt_orm in await session.stream_scalars(statement)
-            ]
-
-
-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,
-    )
+        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)

phoenix/server/api/types/ServerStatus.py

@@ -0,0 +1,6 @@
+import strawberry
+
+
+@strawberry.type
+class ServerStatus:
+    insufficient_storage: bool