remove template options

Arize-ai · Dec 24, 2024 · 81b2b11 · 81b2b11
1 parent b7fe57c
commit 81b2b11
Show file tree

Hide file tree

Showing 8 changed files with 34 additions and 50 deletions.
diff --git a/app/schema.graphql b/app/schema.graphql
@@ -106,7 +106,8 @@ input ChatCompletionInput {
   invocationParameters: [InvocationParameterInput!]! = []
   tools: [JSON!]
   apiKey: String = null
-  template: TemplateOptions
+  templateLanguage: TemplateLanguage
+  templateVariables: JSON
 }
 
 input ChatCompletionMessageInput {
@@ -1701,11 +1702,6 @@ enum TemplateLanguage {
   F_STRING
 }
 
-input TemplateOptions {
-  variables: JSON!
-  language: TemplateLanguage!
-}
-
 type TextChunk implements ChatCompletionSubscriptionPayload {
   datasetExampleId: GlobalID
   content: String!

diff --git a/app/src/pages/playground/__generated__/PlaygroundOutputMutation.graphql.ts b/app/src/pages/playground/__generated__/PlaygroundOutputMutation.graphql.ts
diff --git a/app/src/pages/playground/__generated__/PlaygroundOutputSubscription.graphql.ts b/app/src/pages/playground/__generated__/PlaygroundOutputSubscription.graphql.ts
diff --git a/app/src/pages/playground/playgroundUtils.ts b/app/src/pages/playground/playgroundUtils.ts
@@ -1010,10 +1010,8 @@ export const getChatCompletionInput = ({
 
   return {
     ...baseChatCompletionVariables,
-    template: {
-      variables: variablesMap,
-      language: templateLanguage,
-    },
+    templateVariables: variablesMap,
+    templateLanguage,
   };
 };
 

diff --git a/src/phoenix/server/api/input_types/ChatCompletionInput.py b/src/phoenix/server/api/input_types/ChatCompletionInput.py
@@ -10,7 +10,6 @@
 from .ChatCompletionMessageInput import ChatCompletionMessageInput
 from .GenerativeModelInput import GenerativeModelInput
 from .InvocationParameters import InvocationParameterInput
-from .TemplateOptions import TemplateOptions
 
 
 @strawberry.input
@@ -20,7 +19,8 @@ class ChatCompletionInput:
     invocation_parameters: list[InvocationParameterInput] = strawberry.field(default_factory=list)
     tools: Optional[list[JSON]] = UNSET
     api_key: Optional[str] = strawberry.field(default=None)
-    template: Optional[TemplateOptions] = UNSET
+    template_language: Optional[TemplateLanguage] = UNSET
+    template_variables: Optional[JSON] = UNSET
 
 
 @strawberry.input

diff --git a/src/phoenix/server/api/input_types/TemplateOptions.py b/src/phoenix/server/api/input_types/TemplateOptions.py
diff --git a/src/phoenix/server/api/mutations/chat_mutations.py b/src/phoenix/server/api/mutations/chat_mutations.py
@@ -46,7 +46,6 @@
     ChatCompletionInput,
     ChatCompletionOverDatasetInput,
 )
-from phoenix.server.api.input_types.TemplateOptions import TemplateOptions
 from phoenix.server.api.subscriptions import (
     _default_playground_experiment_description,
     _default_playground_experiment_metadata,
@@ -209,10 +208,8 @@ async def chat_completion_over_dataset(
                             messages=input.messages,
                             tools=input.tools,
                             invocation_parameters=input.invocation_parameters,
-                            template=TemplateOptions(
-                                language=input.template_language,
-                                variables=revision.input,
-                            ),
+                            template_language=input.template_language,
+                            template_variables=revision.input,
                         ),
                     )
                     for revision in batch
@@ -316,11 +313,17 @@ async def _chat_completion(
             )
             for message in input.messages
         ]
-        if template_options := input.template:
-            messages = list(_formatted_messages(messages, template_options))
-            attributes.update(
-                {PROMPT_TEMPLATE_VARIABLES: safe_json_dumps(template_options.variables)}
+        if (template_language := input.template_language) and (
+            template_variables := input.template_variables
+        ):
+            messages = list(
+                _formatted_messages(
+                    messages=messages,
+                    template_language=template_language,
+                    template_variables=template_variables,
+                )
             )
+            attributes.update({PROMPT_TEMPLATE_VARIABLES: safe_json_dumps(template_variables)})
 
         invocation_parameters = llm_client.construct_invocation_parameters(
             input.invocation_parameters
@@ -459,20 +462,21 @@ async def _chat_completion(
 
 def _formatted_messages(
     messages: Iterable[ChatCompletionMessage],
-    template_options: TemplateOptions,
+    template_language: TemplateLanguage,
+    template_variables: dict[str, Any],
 ) -> Iterator[ChatCompletionMessage]:
     """
     Formats the messages using the given template options.
     """
-    template_formatter = _template_formatter(template_language=template_options.language)
+    template_formatter = _template_formatter(template_language=template_language)
     (
         roles,
         templates,
         tool_call_id,
         tool_calls,
     ) = zip(*messages)
     formatted_templates = map(
-        lambda template: template_formatter.format(template, **template_options.variables),
+        lambda template: template_formatter.format(template, **template_variables),
         templates,
     )
     formatted_messages = zip(roles, formatted_templates, tool_call_id, tool_calls)

diff --git a/src/phoenix/server/api/subscriptions.py b/src/phoenix/server/api/subscriptions.py
@@ -120,15 +120,17 @@ async def chat_completion(
             for message in input.messages
         ]
         attributes = None
-        if template_options := input.template:
+        if (template_language := input.template_language) and (
+            template_variables := input.template_variables
+        ):
             messages = list(
                 _formatted_messages(
                     messages=messages,
-                    template_language=template_options.language,
-                    template_variables=template_options.variables,
+                    template_language=template_language,
+                    template_variables=template_variables,
                 )
             )
-            attributes = {PROMPT_TEMPLATE_VARIABLES: safe_json_dumps(template_options.variables)}
+            attributes = {PROMPT_TEMPLATE_VARIABLES: safe_json_dumps(template_variables)}
         invocation_parameters = llm_client.construct_invocation_parameters(
             input.invocation_parameters
         )