feat(prompts): graphql types for tools, output_schema (#5849)

* feat(prompts): graphql types for tools, output_schema * fix types
Arize-ai · Dec 31, 2024 · 42625b5 · 42625b5
1 parent 7138d7b
commit 42625b5
Show file tree

Hide file tree

Showing 8 changed files with 116 additions and 60 deletions.
diff --git a/app/schema.graphql b/app/schema.graphql
@@ -1091,6 +1091,10 @@ type JSONPromptMessage {
   content: JSON!
 }
 
+type JSONSchema {
+  schema: JSON!
+}
+
 type LabelFraction {
   label: String!
   fraction: Float!
@@ -1448,8 +1452,8 @@ type PromptVersion implements Node {
   templateFormat: PromptTemplateFormat!
   template: PromptTemplate!
   invocationParameters: JSON
-  tools: JSON
-  outputSchema: JSON
+  tools: [ToolDefinition!]!
+  outputSchema: JSONSchema
   modelName: String!
   modelProvider: String!
   tags: [PromptVersionTag!]!
@@ -1850,6 +1854,10 @@ type ToolCallChunk implements ChatCompletionSubscriptionPayload {
   function: FunctionCallChunk!
 }
 
+type ToolDefinition {
+  definition: JSON!
+}
+
 type Trace implements Node {
   """The Globally Unique ID of this object"""
   id: GlobalID!

diff --git a/app/src/pages/prompt/__generated__/promptVersionLoaderQuery.graphql.ts b/app/src/pages/prompt/__generated__/promptVersionLoaderQuery.graphql.ts
diff --git a/app/src/pages/prompt/promptVersionLoader.tsx b/app/src/pages/prompt/promptVersionLoader.tsx
@@ -25,7 +25,9 @@ export async function promptVersionLoader(args: LoaderFunctionArgs) {
             description
             invocationParameters
             modelName
-            tools
+            tools {
+              definition
+            }
             user
           }
         }

diff --git a/src/phoenix/server/api/queries.py b/src/phoenix/server/api/queries.py
@@ -69,6 +69,7 @@
 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
+from phoenix.server.api.types.ToolDefinition import ToolDefinition
 from phoenix.server.api.types.Trace import to_gql_trace
 from phoenix.server.api.types.User import User, to_gql_user
 from phoenix.server.api.types.UserApiKey import UserApiKey, to_gql_api_key
@@ -562,33 +563,30 @@ async def node(self, id: GlobalID, info: Info[Context, None]) -> Node:
                     template_format=PromptTemplateFormat.MUSTACHE,
                     template=template,
                     invocation_parameters={"temperature": 0.5},
-                    tools={
-                        "_version": "tools-v1",
-                        "tools": [
-                            {
-                                "definition": {
-                                    "name": "get_current_weather",
-                                    "description": "Get the current weather in a given location",
-                                    "parameters": {
-                                        "type": "object",
-                                        "properties": {
-                                            "location": {
-                                                "type": "string",
-                                                "description": "A location in the world",
-                                            },
-                                            "unit": {
-                                                "type": "string",
-                                                "enum": ["celsius", "fahrenheit"],
-                                                "default": "fahrenheit",
-                                                "description": "The unit of temperature",
-                                            },
+                    tools=[
+                        ToolDefinition(
+                            definition={
+                                "name": "get_current_weather",
+                                "description": "Get the current weather in a given location",
+                                "parameters": {
+                                    "type": "object",
+                                    "properties": {
+                                        "location": {
+                                            "type": "string",
+                                            "description": "A location in the world",
+                                        },
+                                        "unit": {
+                                            "type": "string",
+                                            "enum": ["celsius", "fahrenheit"],
+                                            "default": "fahrenheit",
+                                            "description": "The unit of temperature",
                                         },
-                                        "required": ["location"],
                                     },
-                                }
+                                    "required": ["location"],
+                                },
                             }
-                        ],
-                    },
+                        )
+                    ],
                     output_schema=None,
                     model_name="gpt-4o",
                     model_provider="openai",
@@ -602,7 +600,7 @@ async def node(self, id: GlobalID, info: Info[Context, None]) -> Node:
                     template_format=PromptTemplateFormat.MUSTACHE,
                     template=template,
                     invocation_parameters=None,
-                    tools=None,
+                    tools=[],
                     output_schema=None,
                     model_name="gpt-4o",
                     model_provider="openai",

diff --git a/src/phoenix/server/api/types/JSONSchema.py b/src/phoenix/server/api/types/JSONSchema.py
@@ -0,0 +1,9 @@
+import strawberry
+from strawberry.scalars import JSON
+
+
+@strawberry.type
+class JSONSchema:
+    """A JSON schema definition used to guide an LLM's output"""
+
+    schema: JSON
diff --git a/src/phoenix/server/api/types/Prompt.py b/src/phoenix/server/api/types/Prompt.py
@@ -20,7 +20,9 @@
     TextPromptMessage,
 )
 
+from .JSONSchema import JSONSchema
 from .PromptVersion import PromptTemplateFormat, PromptTemplateType, PromptVersion
+from .ToolDefinition import ToolDefinition
 
 
 @strawberry.type
@@ -64,33 +66,30 @@ async def prompt_versions(
                 template_format=PromptTemplateFormat.MUSTACHE,
                 template=template,
                 invocation_parameters={"temperature": 0.5},
-                tools={
-                    "_version": "tools-v1",
-                    "tools": [
-                        {
-                            "definition": {
-                                "name": "get_current_weather",
-                                "description": "Get the current weather in a given location",
-                                "parameters": {
-                                    "type": "object",
-                                    "properties": {
-                                        "location": {
-                                            "type": "string",
-                                            "description": "A location in the world",
-                                        },
-                                        "unit": {
-                                            "type": "string",
-                                            "enum": ["celsius", "fahrenheit"],
-                                            "default": "fahrenheit",
-                                            "description": "The unit of temperature",
-                                        },
+                tools=[
+                    ToolDefinition(
+                        definition={
+                            "name": "get_current_weather",
+                            "description": "Get the current weather in a given location",
+                            "parameters": {
+                                "type": "object",
+                                "properties": {
+                                    "location": {
+                                        "type": "string",
+                                        "description": "A location in the world",
+                                    },
+                                    "unit": {
+                                        "type": "string",
+                                        "enum": ["celsius", "fahrenheit"],
+                                        "default": "fahrenheit",
+                                        "description": "The unit of temperature",
                                     },
-                                    "required": ["location"],
                                 },
-                            }
+                                "required": ["location"],
+                            },
                         }
-                    ],
-                },
+                    )
+                ],
                 model_name="gpt-4o",
                 model_provider="openai",
             ),
@@ -103,6 +102,22 @@ async def prompt_versions(
                 template=template,
                 model_name="gpt-4o",
                 model_provider="openai",
+                tools=[],
+                output_schema=JSONSchema(
+                    schema={
+                        "type": "json_schema",
+                        "json_schema": {
+                            "name": "response",
+                            "schema": {
+                                "type": "object",
+                                "properties": {},
+                                "required": [],
+                                "additionalProperties": False,
+                            },
+                            "strict": True,
+                        },
+                    }
+                ),
             ),
         ]
 

diff --git a/src/phoenix/server/api/types/PromptVersion.py b/src/phoenix/server/api/types/PromptVersion.py
@@ -11,7 +11,9 @@
 from phoenix.server.api.context import Context
 from phoenix.server.api.types.PromptVersionTemplate import PromptTemplate
 
+from .JSONSchema import JSONSchema
 from .PromptVersionTag import PromptVersionTag
+from .ToolDefinition import ToolDefinition
 
 
 @strawberry.enum
@@ -36,8 +38,8 @@ class PromptVersion(Node):
     template_format: PromptTemplateFormat
     template: PromptTemplate
     invocation_parameters: Optional[JSON] = None
-    tools: Optional[JSON] = None
-    output_schema: Optional[JSON] = None
+    tools: list[ToolDefinition]
+    output_schema: Optional[JSONSchema] = None
     model_name: str
     model_provider: str
 

diff --git a/src/phoenix/server/api/types/ToolDefinition.py b/src/phoenix/server/api/types/ToolDefinition.py
@@ -0,0 +1,9 @@
+import strawberry
+from strawberry.scalars import JSON
+
+
+@strawberry.type
+class ToolDefinition:
+    """The definition of a tool that a generative tool can invoke."""
+
+    definition: JSON