Add structured outputs for openai

Diegi97 · Diegi97 · commit 98be67fd2ca7 · 2025-02-19T17:48:21.000+01:00
diff --git a/chat_rag/chat_rag/llms/base_llm.py b/chat_rag/chat_rag/llms/base_llm.py
@@ -1,4 +1,5 @@
-from typing import Dict, List, Optional, Union, Tuple, Callable
+from typing import Callable, Dict, List, Optional, Tuple, Union
+
 from chat_rag.llms.types import Message
 
 
@@ -43,3 +44,18 @@ async def agenerate(
         tool_choice: str = None,
     ) -> Message:
         pass
+
+    def parse(
+        self,
+        messages: List[Dict[str, str]],
+        schema: Dict,
+    ) -> Message:
+        raise NotImplementedError("This LLM does not support enforced structured output.")
+
+    async def aparse(
+        self,
+        messages: List[Dict[str, str]],
+        schema: Dict,
+    ) -> Message:
+        raise NotImplementedError("This LLM does not support enforced structured output.")
+    
diff --git a/chat_rag/chat_rag/llms/openai_client.py b/chat_rag/chat_rag/llms/openai_client.py
@@ -1,7 +1,8 @@
-from typing import Callable, Dict, List, Union
 import json
+from typing import Callable, Dict, List, Union
 
 from openai import AsyncOpenAI, OpenAI
+from openai.lib._pydantic import _ensure_strict_json_schema
 
 from chat_rag.llms.types import Content, Message, ToolUse, Usage
 
@@ -27,7 +28,9 @@ def _format_tools(
         """
         Format the tools from a openai dict or a callable function to the OpenAI format.
         """
-        tools_formatted, tool_choice = format_tools(tools, tool_choice, mode=Mode.OPENAI_TOOLS)
+        tools_formatted, tool_choice = format_tools(
+            tools, tool_choice, mode=Mode.OPENAI_TOOLS
+        )
 
         # If the tool_choice is a named tool, then apply correct formatting
         if tool_choice in [tool["function"]["name"] for tool in tools_formatted]:
@@ -278,3 +281,67 @@ async def agenerate(
         )
 
         return self._map_openai_message(response)
+
+    def parse(
+        self,
+        messages: List[Union[Dict, Message]],
+        schema: Dict,
+        **kwargs,
+    ) -> Message:
+        """
+        Parse the response from the model into a structured format.
+        Parameters
+        ----------
+        messages : List[Tuple[str, str]]
+            The messages to use for the prompt. Pair of (role, message).
+        schema : Dict
+            The schema to use for the response. It must be a pydantic model json schema, it can be generated using the `model_json_schema` method of the pydantic model.
+        Returns
+        -------
+        Message
+            The parsed message.
+        """
+        messages = self._format_messages(messages)
+        response_format = {
+            "type": "json_schema",
+            "json_schema": {
+                "schema": _ensure_strict_json_schema(schema, path=(), root=schema),
+                "name": schema["title"],
+                "strict": True,
+            },
+        }
+
+        response = self.client.beta.chat.completions.parse(
+            model=self.llm_name,
+            messages=messages,
+            response_format=response_format,
+        )
+
+        return json.loads(response.choices[0].message.content)
+    
+    async def aparse(
+        self,
+        messages: List[Union[Dict, Message]],
+        schema: Dict,
+        **kwargs,
+    ) -> Message:
+        """
+        Parse the response from the model into a structured format.
+        """
+        messages = self._format_messages(messages)
+        response_format = {
+            "type": "json_schema",
+            "json_schema": {
+                "schema": _ensure_strict_json_schema(schema, path=(), root=schema),
+                "name": schema["title"],
+                "strict": True,
+            },
+        }
+
+        response = await self.aclient.beta.chat.completions.parse(
+            model=self.llm_name,
+            messages=messages,
+            response_format=response_format,
+        )
+
+        return json.loads(response.choices[0].message.content)
diff --git a/sdk/examples/agent_example/fsm_definition.py b/sdk/examples/agent_example/fsm_definition.py
@@ -21,6 +21,7 @@ async def send_greeting(sdk: ChatFAQSDK, ctx: dict):
 async def send_answer(sdk: ChatFAQSDK, ctx: dict):
     agent = Agent(
         sdk=sdk,
+        model_name=MODEL_NAME,
         tools=[get_weather],
         system_instruction="You are a knowledgeable weather assistant. Use provided tools when necessary."
     )

Original file line number	Diff line number	Diff line change
`@@ -21,6 +21,7 @@ async def send_greeting(sdk: ChatFAQSDK, ctx: dict):`
`21`	`21`	`async def send_answer(sdk: ChatFAQSDK, ctx: dict):`
`22`	`22`	`agent = Agent(`
`23`	`23`	`sdk=sdk,`
	`24`	`+ model_name=MODEL_NAME,`
`24`	`25`	`tools=[get_weather],`
`25`	`26`	`system_instruction="You are a knowledgeable weather assistant. Use provided tools when necessary."`
`26`	`27`	`)`