Support json_schema for openai endpoint (#432)

Author: liqul

requirements.txt

@@ -3,7 +3,8 @@ pandas>=2.0.0
 matplotlib>=3.4
 seaborn>=0.11
 python-dotenv>=1.0.0
-openai>=1.2.4
+openai>=1.42.0
+pydantic>=2.8.2
 pyyaml>=6.0
 scikit-learn>=1.2.2
 click>=8.0.1

scripts/build_executor.sh

@@ -4,7 +4,7 @@
 scriptDirectory="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
 echo "The script directory is: $scriptDirectory"
 
-version="0.2"
+version="0.3"
 imageName="taskweavercontainers/taskweaver-executor"
 imageFullName="$imageName:$version"
 

taskweaver/code_interpreter/code_interpreter/code_generator_prompt.yaml

@@ -36,7 +36,6 @@ response_json_schema: |-
                 "properties": {
                     "thought": {
                         "type": "string",
-                        "maxLength": 1000,
                         "description": "The thoughts before generating the code."
                     },
                     "reply_type": {
@@ -49,20 +48,21 @@ response_json_schema: |-
                     },
                     "reply_content": {
                         "type": "string",
-                        "minLength": 10,
                         "description": "The actual content of the response. If the reply_type is 'python', the content should be a valid python code snippet. Make sure escaping the special characters (e.g., '\\', '/', and '\"') in the strings for JSON format."
                     }
                 },
                 "required": [
                     "thought",
                     "reply_type",
                     "reply_content"
-                ]
+                ],
+                "additionalProperties": false
             }
         },
         "required": [
             "response"
-        ]
+        ],
+        "additionalProperties": false
     }
 
     

taskweaver/llm/base.py

@@ -55,7 +55,7 @@ def _configure(self) -> None:
 
         self.response_format: Optional[str] = self._get_enum(
             "response_format",
-            options=["json_object", "text"],
+            options=["json_object", "text", "json_schema"],
             default="json_object",
         )
 

taskweaver/llm/openai.py

@@ -180,6 +180,14 @@ def chat_completion(
                 response_format = kwargs["response_format"]
             elif self.config.response_format == "json_object":
                 response_format = {"type": "json_object"}
+            elif self.config.response_format == "json_schema":
+                response_format = {"type": "json_schema"}
+                assert "json_schema" in kwargs, "JSON schema is required for JSON schema response format"
+                response_format["json_schema"] = {
+                    "name": "response",
+                    "strict": True,
+                    "schema": kwargs["json_schema"],
+                }
             else:
                 response_format = None
 

taskweaver/planner/planner_prompt.yaml

@@ -150,11 +150,14 @@ response_json_schema: |-
                 "plan",
                 "current_plan_step",
                 "send_to",
-                "message"
-            ]
+                "message",
+                "review"
+            ],
+            "additionalProperties": false
         }
     },
     "required": [
         "response"
-    ]
+    ],
+    "additionalProperties": false
   }