Use new types in node query

Arize-ai · Dec 19, 2024 · 7a90962 · 7a90962
1 parent 6b3a116
commit 7a90962
Showing 1 changed file with 19 additions and 18 deletions.
diff --git a/src/phoenix/server/api/queries.py b/src/phoenix/server/api/queries.py
@@ -61,6 +61,12 @@
     PromptTemplateType,
     PromptVersion,
 )
+from phoenix.server.api.types.PromptVersionTemplate import (
+    PromptMessageRole,
+    PromptMessagesTemplateV1,
+    PromptMessagesTemplateV1GQL,
+    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,22 +545,25 @@ async def node(self, id: GlobalID, info: Info[Context, None]) -> Node:
                 created_at=datetime.now(),
             )
         elif type_name == PromptVersion.__name__:
+            template_model = PromptMessagesTemplateV1(
+                template=[
+                    TextPromptMessage(
+                        role=PromptMessageRole.USER,
+                        content="Hello what's the weather in Antarctica like?",
+                    )
+                ]
+            )
+
+            template_gql = PromptMessagesTemplateV1GQL.from_model(template_model)
+
             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": "user",
-                                "content": "Hello what's the weather in Antarctica like?",
-                            }
-                        ],
-                    },
+                    template=template_gql,
                     invocation_parameters={"temperature": 0.5},
                     tools={
                         "_version": "tools-v1",
@@ -594,15 +603,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 Antarctica like?",
-                            }
-                        ],
-                    },
+                    template=template_gql,
                     invocation_parameters=None,
                     tools=None,
                     output_schema=None,