feat(playground): tool call backend (#5027)

Arize-ai · Oct 16, 2024 · 5485c03 · 5485c03
1 parent 45973b7
commit 5485c03
Show file tree

Hide file tree

Showing 5 changed files with 229 additions and 28 deletions.
diff --git a/app/schema.graphql b/app/schema.graphql
@@ -65,11 +65,14 @@ enum AuthMethod {
 
 union Bin = NominalBin | IntervalBin | MissingValueBin
 
+union ChatCompletionChunk = TextChunk | ToolCallChunk
+
 input ChatCompletionInput {
   messages: [ChatCompletionMessageInput!]!
   model: GenerativeModelInput!
   invocationParameters: InvocationParameters!
-  apiKey: String
+  tools: [JSON!]
+  apiKey: String = null
 }
 
 input ChatCompletionMessageInput {
@@ -817,6 +820,11 @@ type ExportedFile {
   fileName: String!
 }
 
+type FunctionCallChunk {
+  name: String!
+  arguments: String!
+}
+
 type Functionality {
   """Model inferences are available for analysis"""
   modelInferences: Boolean!
@@ -903,6 +911,7 @@ input InvocationParameters {
   topP: Float
   stop: [String!]
   seed: Int
+  toolChoice: JSON
 }
 
 """
@@ -1425,7 +1434,7 @@ enum SpanStatusCode {
 }
 
 type Subscription {
-  chatCompletion(input: ChatCompletionInput!): String!
+  chatCompletion(input: ChatCompletionInput!): ChatCompletionChunk!
 }
 
 type SystemApiKey implements ApiKey & Node {
@@ -1445,6 +1454,10 @@ type SystemApiKey implements ApiKey & Node {
   id: GlobalID!
 }
 
+type TextChunk {
+  content: String!
+}
+
 input TimeRange {
   """The start of the time range"""
   start: DateTime!
@@ -1462,6 +1475,11 @@ type TimeSeriesDataPoint {
   value: Float
 }
 
+type ToolCallChunk {
+  id: String!
+  function: FunctionCallChunk!
+}
+
 type Trace implements Node {
   """The Globally Unique ID of this object"""
   id: GlobalID!

diff --git a/app/src/pages/playground/PlaygroundOutput.tsx b/app/src/pages/playground/PlaygroundOutput.tsx
@@ -115,7 +115,19 @@ function useChatCompletionSubscription({
               invocationParameters: $invocationParameters
               apiKey: $apiKey
             }
-          )
+          ) {
+            __typename
+            ... on TextChunk {
+              content
+            }
+            ... on ToolCallChunk {
+              id
+              function {
+                name
+                arguments
+              }
+            }
+          }
         }
       `,
       variables: params,
@@ -200,7 +212,10 @@ function PlaygroundOutputText(props: PlaygroundInstanceProps) {
     },
     runId: instance.activeRunId,
     onNext: (response) => {
-      setOutput((acc) => acc + response.chatCompletion);
+      const chatCompletion = response.chatCompletion;
+      if (chatCompletion.__typename === "TextChunk") {
+        setOutput((acc) => acc + chatCompletion.content);
+      }
     },
     onCompleted: () => {
       markPlaygroundInstanceComplete(props.playgroundInstanceId);

diff --git a/app/src/pages/playground/__generated__/PlaygroundOutputSubscription.graphql.ts b/app/src/pages/playground/__generated__/PlaygroundOutputSubscription.graphql.ts
diff --git a/src/phoenix/server/api/input_types/InvocationParameters.py b/src/phoenix/server/api/input_types/InvocationParameters.py
@@ -2,6 +2,7 @@
 
 import strawberry
 from strawberry import UNSET
+from strawberry.scalars import JSON
 
 
 @strawberry.input
@@ -16,3 +17,4 @@ class InvocationParameters:
     top_p: Optional[float] = UNSET
     stop: Optional[List[str]] = UNSET
     seed: Optional[int] = UNSET
+    tool_choice: Optional[JSON] = UNSET