Ignore templating for raw return requests

piercefreeman · Aug 24, 2023 · 0694734 · 0694734
1 parent 54792ea
commit 0694734
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 4 deletions.
diff --git a/gpt_json/gpt.py b/gpt_json/gpt.py
@@ -252,14 +252,17 @@ async def run(
             except (ValueError, ValidationError):
                 raise InvalidFunctionParameters(function_name, function_args_string)
 
+        raw_response = parse_obj_model(GPTMessage, response_message)
+        raw_response.allow_templating = False
+
         extracted_json, fixed_payload = self.extract_json(
             response_message, self.extract_type
         )
 
         # Cast to schema model
         if extracted_json is None:
             return RunResponse(
-                raw_response=parse_obj_model(GPTMessage, response_message),
+                raw_response=raw_response,
                 response=None,
                 fix_transforms=fixed_payload,
                 function_call=function_call,
@@ -273,7 +276,7 @@ async def run(
 
         # Allow pydantic to handle the validation
         return RunResponse(
-            raw_response=parse_obj_model(GPTMessage, response_message),
+            raw_response=raw_response,
             response=self.schema_model(**extracted_json),
             fix_transforms=fixed_payload,
             function_call=function_call,
@@ -538,6 +541,9 @@ def fill_message_template(
             SCHEMA_PROMPT_TEMPLATE_KEY: self.schema_prompt,
         }
 
+        if message.content is None or not message.allow_templating:
+            return message
+
         # Regular quotes should passthrough to the next stage, except for our special keys
         content = message.content.replace("{", "{{").replace("}", "}}")
         for key in auto_format.keys():
@@ -567,7 +573,7 @@ def trim_messages(self, messages: list[GPTMessage], n: int):
         Returns:
             list: A list of messages with a total token count less than n tokens.
         """
-        message_text = [message.content for message in messages]
+        message_text = [message.content for message in messages if message.content]
 
         enc = encoding_for_model("gpt-4")
         filtered_messages = []

diff --git a/gpt_json/models.py b/gpt_json/models.py
@@ -70,14 +70,18 @@ class GPTMessage(BaseModel):
     """
 
     role: GPTMessageRole
-    content: str
+    content: str | None
 
     # Name is only supported if we're formatting a function message
     name: str | None = None
 
     # Message from the server
     function_call: FunctionCall | None = None
 
+    # If enabled, gpt-json will attempt to format the message with the runtime variables
+    # Disable this in cases where you want the message to be formatted 1:1 with the input
+    allow_templating: bool = True
+
     @model_validator(mode="after")
     def check_name_if_function(self):
         if self.role == GPTMessageRole.FUNCTION and self.name is None: