Improve ChatAgent to handle StructuredOutput

Author: RussellLuo

coagent/agents/chat_agent.py

@@ -12,7 +12,7 @@
 
 from .aswarm import Agent as SwarmAgent, Swarm
 from .aswarm.util import function_to_jsonschema
-from .messages import ChatMessage, ChatHistory
+from .messages import ChatMessage, ChatHistory, StructuredOutput
 from .model_client import default_model_client, ModelClient
 from .util import is_user_confirmed
 
@@ -233,21 +233,38 @@ async def agent(self, agent_type: str) -> AsyncIterator[ChatMessage]:
     async def handle_history(
         self, msg: ChatHistory, ctx: Context
     ) -> AsyncIterator[ChatMessage]:
-        response = self._handle_history(msg, ctx)
+        response = self._handle_history(msg)
         async for resp in response:
             yield resp
 
     @handler
     async def handle_message(
         self, msg: ChatMessage, ctx: Context
     ) -> AsyncIterator[ChatMessage]:
-        history = ChatHistory(messages=[msg], response_format=msg.response_format)
-        response = self._handle_history(history, ctx)
+        history = ChatHistory(messages=[msg])
+        response = self._handle_history(history)
         async for resp in response:
             yield resp
 
+    @handler
+    async def handle_structured_output(
+        self, msg: StructuredOutput, ctx: Context
+    ) -> AsyncIterator[ChatMessage]:
+        match msg.input:
+            case ChatMessage():
+                history = ChatHistory(messages=[msg.input])
+                response = self._handle_history(history, msg.output_schema)
+                async for resp in response:
+                    yield resp
+            case ChatHistory():
+                response = self._handle_history(msg.input, msg.output_schema)
+                async for resp in response:
+                    yield resp
+
     async def _handle_history(
-        self, msg: ChatHistory, ctx: Context
+        self,
+        msg: ChatHistory,
+        response_format: dict | None = None,
     ) -> AsyncIterator[ChatMessage]:
         # For now, we assume that the agent is processing messages sequentially.
         self._history: ChatHistory = msg
@@ -260,7 +277,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,
+            response_format=response_format,
             context_variables=msg.extensions,
         )
         async for resp in response:

coagent/agents/messages.py

@@ -1,15 +1,14 @@
 from __future__ import annotations
 
-from typing import Any
+from typing import Any, Type
 
-from pydantic import Field
-from coagent.core import Message
+from coagent.core import logger, Message
+from pydantic import BaseModel, Field, field_validator, field_serializer
 
 
 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.")
@@ -30,4 +29,57 @@ def to_llm_message(self) -> dict[str, Any]:
 
 class ChatHistory(Message):
     messages: list[ChatMessage]
-    response_format: dict | None = None
+
+
+class StructuredOutput(Message):
+    input: ChatMessage | ChatHistory = Field(..., description="Input message.")
+    output_type: Type[BaseModel] | None = Field(
+        None,
+        description="Output schema specified as a Pydantic model. Equivalent to OpenAI's `response_format`.",
+    )
+    output_schema: dict | None = Field(
+        None,
+        description="Output schema specified as a dict. Setting this suppresses `output_type`.",
+    )
+
+    @field_serializer("input")
+    def serialize_input(self, value: Message, _info) -> dict:
+        data = value.model_dump(exclude_defaults=True)
+        data["__message_type__"] = value.__class__.__name__
+        return data
+
+    @field_validator("input", mode="before")
+    @classmethod
+    def validate_input(cls, value: Message | dict) -> Message:
+        if isinstance(value, dict):
+            message_type = value.pop("__message_type__", None)
+            match message_type:
+                # Only support ChatMessage and ChatHistory for now.
+                case "ChatMessage":
+                    return ChatMessage.model_validate(value)
+                case "ChatHistory":
+                    return ChatHistory.model_validate(value)
+        return value
+
+    @field_serializer("output_type")
+    def serialize_output_type(self, value: Type[BaseModel] | None, _info) -> None:
+        # Always return None for `output_type` since it will be converted to `output_schema`.
+        return None
+
+    @field_serializer("output_schema")
+    def serialize_output_schema(self, value: dict | None, _info) -> dict | None:
+        if self.output_type:
+            if value:
+                logger.warning("Setting output_schema suppresses output_type")
+                return value
+            return type_to_response_format_param(self.output_type)
+
+        return value
+
+
+def type_to_response_format_param(
+    response_format: Type[BaseModel] | dict | None,
+) -> dict | None:
+    import litellm.utils
+
+    return litellm.utils.type_to_response_format_param(response_format)

coagent/agents/util.py

@@ -1,22 +1,13 @@
 import asyncio
 import functools
-from typing import AsyncIterator, Type
+from typing import AsyncIterator
 
 from coagent.core import logger
-from pydantic import BaseModel
 
 from .messages import ChatMessage
 from .model_client import default_model_client, ModelClient
 
 
-def type_to_response_format_param(
-    response_format: Type[BaseModel] | dict | None,
-) -> dict | None:
-    import litellm.utils
-
-    return litellm.utils.type_to_response_format_param(response_format)
-
-
 async def chat(
     messages: list[ChatMessage], client: ModelClient = default_model_client
 ) -> ChatMessage: