refactor input types

app/schema.graphql

@@ -145,15 +145,13 @@ type ChatCompletionMutationPayload {
 
 union ChatCompletionMutationPayloadChatCompletionMutationError = ChatCompletionMutationPayload | ChatCompletionMutationError
 
-input ChatCompletionOverDatasetInput {
+input ChatCompletionOverDatasetConfig {
   messages: [ChatCompletionMessageInput!]!
   model: GenerativeModelInput!
   invocationParameters: [InvocationParameterInput!]! = []
   tools: [JSON!]
   apiKey: String = null
   templateLanguage: TemplateLanguage!
-  datasetId: GlobalID!
-  datasetVersionId: GlobalID = null
   experimentName: String = null
   experimentDescription: String = null
   experimentMetadata: JSON = {}
@@ -173,20 +171,24 @@ type ChatCompletionOverDatasetMutationPayload {
 }
 
 type ChatCompletionSubscriptionError implements ChatCompletionSubscriptionPayload {
+  configIndex: Int
   datasetExampleId: GlobalID
   message: String!
 }
 
 type ChatCompletionSubscriptionExperiment implements ChatCompletionSubscriptionPayload {
+  configIndex: Int
   datasetExampleId: GlobalID
   experiment: Experiment!
 }
 
 interface ChatCompletionSubscriptionPayload {
+  configIndex: Int
   datasetExampleId: GlobalID
 }
 
 type ChatCompletionSubscriptionResult implements ChatCompletionSubscriptionPayload {
+  configIndex: Int
   datasetExampleId: GlobalID
   span: Span
   experimentRun: ExperimentRun
@@ -197,6 +199,12 @@ type ChatCompletionToolCall {
   function: ChatCompletionFunctionCall!
 }
 
+input ChatCompletionsOverDatasetInput {
+  configs: [ChatCompletionOverDatasetConfig!]!
+  datasetId: GlobalID!
+  datasetVersionId: GlobalID = null
+}
+
 input ClearProjectInput {
   id: GlobalID!
 
@@ -939,6 +947,7 @@ type FloatInvocationParameter implements InvocationParameterBase {
 }
 
 type FunctionCallChunk implements ChatCompletionSubscriptionPayload {
+  configIndex: Int
   datasetExampleId: GlobalID
   name: String!
   arguments: String!
@@ -1171,7 +1180,7 @@ type Mutation {
   patchUser(input: PatchUserInput!): UserMutationPayload!
   patchViewer(input: PatchViewerInput!): UserMutationPayload!
   deleteUsers(input: DeleteUsersInput!): Void
-  chatCompletionOverDataset(input: ChatCompletionOverDatasetInput!): ChatCompletionOverDatasetMutationPayload!
+  chatCompletionsOverDataset(input: ChatCompletionsOverDatasetInput!): ChatCompletionOverDatasetMutationPayload!
   chatCompletion(input: ChatCompletionInput!): ChatCompletionMutationPayload!
 }
 
@@ -1676,7 +1685,7 @@ type StringListInvocationParameter implements InvocationParameterBase {
 
 type Subscription {
   chatCompletion(input: ChatCompletionInput!): ChatCompletionSubscriptionPayload!
-  chatCompletionOverDataset(input: ChatCompletionOverDatasetInput!): ChatCompletionSubscriptionPayload!
+  chatCompletionsOverDataset(input: ChatCompletionsOverDatasetInput!): ChatCompletionSubscriptionPayload!
 }
 
 type SystemApiKey implements ApiKey & Node {
@@ -1703,6 +1712,7 @@ enum TemplateLanguage {
 }
 
 type TextChunk implements ChatCompletionSubscriptionPayload {
+  configIndex: Int
   datasetExampleId: GlobalID
   content: String!
 }
@@ -1731,6 +1741,7 @@ type TokenUsage {
 }
 
 type ToolCallChunk implements ChatCompletionSubscriptionPayload {
+  configIndex: Int
   datasetExampleId: GlobalID
   id: String!
   function: FunctionCallChunk!

app/src/pages/playground/PlaygroundDatasetExamplesTable.tsx

@@ -95,7 +95,7 @@ import {
 const PAGE_SIZE = 10;
 
 type ChatCompletionOverDatasetMutationPayload = Extract<
-  PlaygroundDatasetExamplesTableMutation$data["chatCompletionOverDataset"],
+  PlaygroundDatasetExamplesTableMutation$data["chatCompletionsOverDataset"],
   { __typename: "ChatCompletionOverDatasetMutationPayload" }
 >;
 
@@ -474,7 +474,7 @@ export function PlaygroundDatasetExamplesTable({
         if (response == null) {
           return;
         }
-        const chatCompletion = response.chatCompletionOverDataset;
+        const chatCompletion = response.chatCompletionsOverDataset;
         switch (chatCompletion.__typename) {
           case "ChatCompletionSubscriptionExperiment":
             updateInstance({
@@ -565,13 +565,13 @@ export function PlaygroundDatasetExamplesTable({
         updateInstance({
           instanceId,
           patch: {
-            experimentId: response.chatCompletionOverDataset.experimentId,
+            experimentId: response.chatCompletionsOverDataset.experimentId,
           },
         });
         setExampleDataForInstance({
           instanceId,
           data: createExampleResponsesForInstance(
-            response.chatCompletionOverDataset
+            response.chatCompletionsOverDataset
           ),
         });
       },
@@ -966,9 +966,9 @@ export function PlaygroundDatasetExamplesTable({
 
 graphql`
   subscription PlaygroundDatasetExamplesTableSubscription(
-    $input: ChatCompletionOverDatasetInput!
+    $input: ChatCompletionsOverDatasetInput!
   ) {
-    chatCompletionOverDataset(input: $input) {
+    chatCompletionsOverDataset(input: $input) {
       __typename
       ... on TextChunk {
         content
@@ -1016,9 +1016,9 @@ graphql`
 
 graphql`
   mutation PlaygroundDatasetExamplesTableMutation(
-    $input: ChatCompletionOverDatasetInput!
+    $input: ChatCompletionsOverDatasetInput!
   ) {
-    chatCompletionOverDataset(input: $input) {
+    chatCompletionsOverDataset(input: $input) {
       __typename
       experimentId
       examples {

app/src/pages/playground/PlaygroundDatasetExamplesTableContext.tsx

@@ -8,13 +8,13 @@ import { PartialOutputToolCall } from "./PlaygroundToolCall";
 type InstanceId = number;
 export type ExampleId = string;
 type ChatCompletionSubscriptionResult = Extract<
-  PlaygroundDatasetExamplesTableSubscription$data["chatCompletionOverDataset"],
+  PlaygroundDatasetExamplesTableSubscription$data["chatCompletionsOverDataset"],
   { __typename: "ChatCompletionSubscriptionResult" }
 >;
 export type Span = NonNullable<ChatCompletionSubscriptionResult["span"]>;
 
 type ToolCallChunk = Extract<
-  PlaygroundDatasetExamplesTableSubscription$data["chatCompletionOverDataset"],
+  PlaygroundDatasetExamplesTableSubscription$data["chatCompletionsOverDataset"],
   { __typename: "ToolCallChunk" }
 >;
 

app/src/pages/playground/playgroundUtils.ts

@@ -38,7 +38,7 @@ import { safelyParseJSON } from "@phoenix/utils/jsonUtils";
 
 import { TemplateLanguages } from "../../components/templateEditor/constants";
 
-import { ChatCompletionOverDatasetInput } from "./__generated__/PlaygroundDatasetExamplesTableSubscription.graphql";
+import { ChatCompletionsOverDatasetInput } from "./__generated__/PlaygroundDatasetExamplesTableSubscription.graphql";
 import {
   ChatCompletionInput,
   ChatCompletionMessageInput,
@@ -1028,16 +1028,20 @@ export const getChatCompletionOverDatasetInput = ({
   instanceId: number;
   credentials: CredentialsState;
   datasetId: string;
-}): ChatCompletionOverDatasetInput => {
+}): ChatCompletionsOverDatasetInput => {
   const baseChatCompletionVariables = getBaseChatCompletionInput({
     playgroundStore,
     instanceId,
     credentials,
   });
 
   return {
-    ...baseChatCompletionVariables,
-    templateLanguage: playgroundStore.getState().templateLanguage,
+    configs: [
+      {
+        ...baseChatCompletionVariables,
+        templateLanguage: playgroundStore.getState().templateLanguage,
+      },
+    ],
     datasetId,
   };
 };

src/phoenix/server/api/helpers/playground_spans.py

@@ -34,7 +34,7 @@
 from phoenix.server.api.helpers.dataset_helpers import get_dataset_example_output
 from phoenix.server.api.input_types.ChatCompletionInput import (
     ChatCompletionInput,
-    ChatCompletionOverDatasetInput,
+    ChatCompletionOverDatasetConfig,
 )
 from phoenix.server.api.types.ChatCompletionMessageRole import ChatCompletionMessageRole
 from phoenix.server.api.types.ChatCompletionSubscriptionPayload import (
@@ -63,7 +63,7 @@ class streaming_llm_span:
     def __init__(
         self,
         *,
-        input: Union[ChatCompletionInput, ChatCompletionOverDatasetInput],
+        input: Union[ChatCompletionInput, ChatCompletionOverDatasetConfig],
         messages: list[ChatCompletionMessage],
         invocation_parameters: Mapping[str, Any],
         attributes: Optional[dict[str, Any]] = None,
@@ -254,7 +254,7 @@ def llm_tools(tools: list[JSONScalarType]) -> Iterator[tuple[str, Any]]:
 
 
 def input_value_and_mime_type(
-    input: Union[ChatCompletionInput, ChatCompletionOverDatasetInput],
+    input: Union[ChatCompletionInput, ChatCompletionOverDatasetConfig],
 ) -> Iterator[tuple[str, Any]]:
     assert (api_key := "api_key") in (input_data := jsonify(input))
     disallowed_keys = {"api_key", "invocation_parameters"}

src/phoenix/server/api/input_types/ChatCompletionInput.py

@@ -24,15 +24,20 @@ class ChatCompletionInput:
 
 
 @strawberry.input
-class ChatCompletionOverDatasetInput:
+class ChatCompletionOverDatasetConfig:
     messages: list[ChatCompletionMessageInput]
     model: GenerativeModelInput
     invocation_parameters: list[InvocationParameterInput] = strawberry.field(default_factory=list)
     tools: Optional[list[JSON]] = UNSET
     api_key: Optional[str] = strawberry.field(default=None)
     template_language: TemplateLanguage
-    dataset_id: GlobalID
-    dataset_version_id: Optional[GlobalID] = None
     experiment_name: Optional[str] = None
     experiment_description: Optional[str] = None
     experiment_metadata: Optional[JSON] = strawberry.field(default_factory=dict)
+
+
+@strawberry.input
+class ChatCompletionsOverDatasetInput:
+    configs: list[ChatCompletionOverDatasetConfig]
+    dataset_id: GlobalID
+    dataset_version_id: Optional[GlobalID] = None