fix: structured_output_adapter

luochen1990 · Aug 7, 2024 · a70c69b · a70c69b
1 parent fc4ab1f
commit a70c69b
Showing 1 changed file with 12 additions and 3 deletions.
diff --git a/src/ai_powered/llm/adapters/structured_output_adapter.py b/src/ai_powered/llm/adapters/structured_output_adapter.py
@@ -14,12 +14,21 @@ def _param_response_format_maker(self) -> ResponseFormat | openai.NotGiven:
             "json_schema": {
                 "name": "math_response",
                 #"strict": True if ModelFeature.strict_mode in self.model_features else None, #NOTE: https://platform.openai.com/docs/guides/structured-outputs/supported-schemas
-                "schema": self.return_schema,
+                "schema": {
+                    "type": "object",
+                    "properties": {
+                        "result": self.return_schema,
+                    },
+                    "required": ["result"],
+                }
             }
         }
 
     def _response_message_parser(self, response_message: ChatCompletionMessage) -> str:
         tool_calls = response_message.tool_calls
+        assert tool_calls is None
 
-        assert tool_calls is not None
-        return tool_calls[0].function.arguments
+        raw_resp_str = response_message.content
+        assert raw_resp_str is not None
+
+        return raw_resp_str