feat: Add PromptTemplate type

app/schema.graphql

@@ -1086,6 +1086,11 @@ type JSONInvocationParameter implements InvocationParameterBase {
   defaultValue: JSON
 }
 
+type JSONPromptMessage {
+  role: PromptMessageRole!
+  content: JSON!
+}
+
 type LabelFraction {
   label: String!
   fraction: Float!
@@ -1377,6 +1382,11 @@ type Prompt implements Node {
   promptVersions(first: Int = 50, last: Int, after: String, before: String): PromptVersionConnection!
 }
 
+type PromptChatTemplateV1 {
+  version: String!
+  messages: [TextPromptMessageJSONPromptMessage!]!
+}
+
 """A connection to a list of items."""
 type PromptConnection {
   """Pagination data for this connection"""
@@ -1395,6 +1405,12 @@ type PromptEdge {
   node: Prompt!
 }
 
+enum PromptMessageRole {
+  USER
+  SYSTEM
+  AI
+}
+
 type PromptResponse {
   """The prompt submitted to the LLM"""
   prompt: String
@@ -1403,6 +1419,10 @@ type PromptResponse {
   response: String
 }
 
+type PromptStringTemplate {
+  template: String!
+}
+
 enum PromptTemplateFormat {
   MUSTACHE
   FSTRING
@@ -1414,14 +1434,16 @@ enum PromptTemplateType {
   CHAT
 }
 
+union PromptTemplateVersion = PromptStringTemplate | PromptChatTemplateV1
+
 type PromptVersion implements Node {
   """The Globally Unique ID of this object"""
   id: GlobalID!
   user: String
   description: String!
   templateType: PromptTemplateType!
   templateFormat: PromptTemplateFormat!
-  template: JSON!
+  template: PromptTemplateVersion!
   invocationParameters: JSON
   tools: JSON
   outputSchema: JSON
@@ -1783,6 +1805,13 @@ type TextChunk implements ChatCompletionSubscriptionPayload {
   content: String!
 }
 
+type TextPromptMessage {
+  role: PromptMessageRole!
+  content: String!
+}
+
+union TextPromptMessageJSONPromptMessage = TextPromptMessage | JSONPromptMessage
+
 input TimeRange {
   """The start of the time range"""
   start: DateTime!

src/phoenix/server/api/helpers/prompthub/models.py

@@ -0,0 +1,51 @@
+from enum import Enum
+from typing import Any, Union
+
+import strawberry
+from pydantic import BaseModel
+
+JSONSerializable = Union[None, bool, int, float, str, dict[str, Any], list[Any]]
+
+
+@strawberry.enum
+class PromptMessageRole(str, Enum):
+    USER = "user"
+    SYSTEM = "system"
+    AI = "ai"  # E.g. the assistant. Normalize to AI for consistency.
+
+
+class TextPromptMessage(BaseModel):
+    role: PromptMessageRole
+    content: str
+
+
+class JSONPromptMessage(BaseModel):
+    role: PromptMessageRole
+    content: JSONSerializable
+
+
+class PromptChatTemplateV1(BaseModel):
+    _version: str = "messages-v1"
+    template: list[Union[TextPromptMessage, JSONPromptMessage]]
+
+
+class PromptStringTemplate(BaseModel):
+    template: str
+
+
+# TODO: Figure out enums, maybe just store whole tool blobs
+# class PromptToolParameter(BaseModel):
+#     name: str
+#     type: str
+#     description: str
+#     required: bool
+#     default: str
+
+
+class PromptToolDefinition(BaseModel):
+    definition: JSONSerializable
+
+
+class PromptTools(BaseModel):
+    _version: str = "tools-v1"
+    tools: list[PromptToolDefinition]

src/phoenix/server/api/queries.py

@@ -61,6 +61,11 @@
     PromptTemplateType,
     PromptVersion,
 )
+from phoenix.server.api.types.PromptVersionTemplate import (
+    PromptChatTemplateV1,
+    PromptMessageRole,
+    TextPromptMessage,
+)
 from phoenix.server.api.types.SortDir import SortDir
 from phoenix.server.api.types.Span import Span, to_gql_span
 from phoenix.server.api.types.SystemApiKey import SystemApiKey
@@ -539,29 +544,24 @@ async def node(self, id: GlobalID, info: Info[Context, None]) -> Node:
                 created_at=datetime.now(),
             )
         elif type_name == PromptVersion.__name__:
+            template = PromptChatTemplateV1(
+                messages=[
+                    TextPromptMessage(
+                        role=PromptMessageRole.USER,
+                        content="Hello what's the weather in Antarctica like?",
+                    )
+                ]
+            )
+
             if node_id == 2:
                 return PromptVersion(
                     id_attr=2,
                     user="alice",
                     description="A dummy prompt version",
                     template_type=PromptTemplateType.CHAT,
                     template_format=PromptTemplateFormat.MUSTACHE,
-                    template={
-                        "_version": "messages-v1",
-                        "messages": [
-                            {"role": "system", "content": "You are a helpful assistant"},
-                            {
-                                "role": "user",
-                                "content": "Hello what's the weather in {{location}} like?",
-                            },
-                            {"role": "ai", "content": "Looking up the weather in {{location}}..."},
-                        ],
-                    },
-                    invocation_parameters={
-                        "temperature": 0.5,
-                        "model": "gpt-4o",
-                        "max_tokens": 100,
-                    },
+                    template=template,
+                    invocation_parameters={"temperature": 0.5},
                     tools={
                         "_version": "tools-v1",
                         "tools": [
@@ -600,15 +600,7 @@ async def node(self, id: GlobalID, info: Info[Context, None]) -> Node:
                     description="A dummy prompt version",
                     template_type=PromptTemplateType.CHAT,
                     template_format=PromptTemplateFormat.MUSTACHE,
-                    template={
-                        "_version": "messages-v1",
-                        "messages": [
-                            {
-                                "role": "user",
-                                "content": "Hello what's the weather in {{location}} like?",
-                            }
-                        ],
-                    },
+                    template=template,
                     invocation_parameters=None,
                     tools=None,
                     output_schema=None,

src/phoenix/server/api/types/Prompt.py

@@ -9,7 +9,16 @@
 from strawberry.types import Info
 
 from phoenix.server.api.context import Context
-from phoenix.server.api.types.pagination import ConnectionArgs, CursorString, connection_from_list
+from phoenix.server.api.types.pagination import (
+    ConnectionArgs,
+    CursorString,
+    connection_from_list,
+)
+from phoenix.server.api.types.PromptVersionTemplate import (
+    PromptChatTemplateV1,
+    PromptMessageRole,
+    TextPromptMessage,
+)
 
 from .PromptVersion import PromptTemplateFormat, PromptTemplateType, PromptVersion
 
@@ -37,29 +46,24 @@ async def prompt_versions(
             before=before if isinstance(before, CursorString) else None,
         )
 
+        template = PromptChatTemplateV1(
+            messages=[
+                TextPromptMessage(
+                    role=PromptMessageRole.USER,
+                    content="Hello what's the weather in Antarctica like?",
+                )
+            ]
+        )
+
         dummy_data = [
             PromptVersion(
                 id_attr=2,
                 user="alice",
                 description="A dummy prompt version",
                 template_type=PromptTemplateType.CHAT,
                 template_format=PromptTemplateFormat.MUSTACHE,
-                template={
-                    "_version": "messages-v1",
-                    "messages": [
-                        {"role": "system", "content": "You are a helpful assistant"},
-                        {
-                            "role": "user",
-                            "content": "Hello what's the weather in {{location}} like?",
-                        },
-                        {"role": "ai", "content": "Looking up the weather in {{location}}..."},
-                    ],
-                },
-                invocation_parameters={
-                    "temperature": 0.5,
-                    "model": "gpt-4o",
-                    "max_tokens": 100,
-                },
+                template=template,
+                invocation_parameters={"temperature": 0.5},
                 tools={
                     "_version": "tools-v1",
                     "tools": [
@@ -96,15 +100,7 @@ async def prompt_versions(
                 description="A dummy prompt version",
                 template_type=PromptTemplateType.CHAT,
                 template_format=PromptTemplateFormat.MUSTACHE,
-                template={
-                    "_version": "messages-v1",
-                    "messages": [
-                        {
-                            "role": "user",
-                            "content": "Hello what's the weather in {{location}} like?",
-                        }
-                    ],
-                },
+                template=template,
                 model_name="gpt-4o",
                 model_provider="openai",
             ),

src/phoenix/server/api/types/PromptVersion.py

@@ -7,6 +7,8 @@
 from strawberry.relay import Node, NodeID
 from strawberry.scalars import JSON
 
+from phoenix.server.api.types.PromptVersionTemplate import PromptTemplate
+
 
 @strawberry.enum
 class PromptTemplateType(str, Enum):
@@ -28,7 +30,7 @@ class PromptVersion(Node):
     description: str
     template_type: PromptTemplateType
     template_format: PromptTemplateFormat
-    template: JSON
+    template: PromptTemplate
     invocation_parameters: Optional[JSON] = None
     tools: Optional[JSON] = None
     output_schema: Optional[JSON] = None

src/phoenix/server/api/types/PromptVersionTemplate.py

@@ -0,0 +1,38 @@
+# Part of the Phoenix PromptHub feature set
+
+from typing import Union
+
+import strawberry
+from strawberry.scalars import JSON
+
+from phoenix.server.api.helpers.prompthub.models import (
+    PromptMessageRole,
+)
+
+
+@strawberry.type
+class TextPromptMessage:
+    role: PromptMessageRole
+    content: str
+
+
+@strawberry.type
+class JSONPromptMessage:
+    role: PromptMessageRole
+    content: JSON
+
+
+@strawberry.type
+class PromptChatTemplateV1:
+    version: str = "messages-v1"
+    messages: list[Union[TextPromptMessage, JSONPromptMessage]]
+
+
+@strawberry.type
+class PromptStringTemplate:
+    template: str
+
+
+PromptTemplate = strawberry.union(
+    "PromptTemplateVersion", (PromptStringTemplate, PromptChatTemplateV1)
+)