Add support for structured outputs

Author: RussellLuo

coagent/agents/aswarm/core.py

@@ -44,6 +44,7 @@ async def get_chat_completion(
         model_override: str,
         stream: bool,
         debug: bool,
+        response_format: dict | None = None,
     ) -> ChatCompletionMessage:
         context_variables = defaultdict(str, context_variables)
         instructions = (
@@ -65,6 +66,7 @@ async def get_chat_completion(
         create_params = {
             "model": model_override or agent.model,
             "messages": messages,
+            "response_format": response_format,
             "tools": tools or None,
             "tool_choice": agent.tool_choice,
             "stream": stream,
@@ -217,6 +219,7 @@ async def run_and_stream(
         self,
         agent: Agent,
         messages: List,
+        response_format: dict | None = None,
         context_variables: dict = {},
         model_override: str = None,
         debug: bool = False,
@@ -248,6 +251,7 @@ async def run_and_stream(
             completion = self.get_chat_completion(
                 agent=active_agent,
                 history=history,
+                response_format=response_format,
                 context_variables=context_variables,
                 model_override=model_override,
                 stream=True,

coagent/agents/chat_agent.py

@@ -241,7 +241,7 @@ async def handle_history(
     async def handle_message(
         self, msg: ChatMessage, ctx: Context
     ) -> AsyncIterator[ChatMessage]:
-        history = ChatHistory(messages=[msg])
+        history = ChatHistory(messages=[msg], response_format=msg.response_format)
         response = self._handle_history(history, ctx)
         async for resp in response:
             yield resp
@@ -260,6 +260,7 @@ async def _handle_history(
         response = self._swarm_client.run_and_stream(
             agent=swarm_agent,
             messages=[m.model_dump() for m in msg.messages],
+            response_format=msg.response_format,
             context_variables=msg.extensions,
         )
         async for resp in response:

coagent/agents/messages.py

@@ -9,6 +9,7 @@
 class ChatMessage(Message):
     role: str
     content: str
+    response_format: dict | None = None
 
     type: str = Field(default="", description="The type of the message. e.g. confirm")
     sender: str = Field(default="", description="The sending agent of the message.")
@@ -29,3 +30,4 @@ def to_llm_message(self) -> dict[str, Any]:
 
 class ChatHistory(Message):
     messages: list[ChatMessage]
+    response_format: dict | None = None

coagent/agents/model_client.py

@@ -39,6 +39,7 @@ async def acompletion(
         temperature: float = 0.1,
         tools: list | None = None,
         tool_choice: str | None = None,
+        response_format: dict | None = None,
         **kwargs,
     ):  # -> ModelResponse:
         import litellm
@@ -54,6 +55,7 @@ async def acompletion(
             api_base=self.api_base,
             api_version=self.api_version,
             api_key=self.api_key,
+            response_format=response_format,
             **kwargs,
         )
         return response