Feat: structured output

openhands-sdk/openhands/sdk/agent/agent.py

@@ -1210,6 +1210,10 @@ def _get_action_event(
                 return
 
             arguments = fix_malformed_tool_arguments(arguments, tool.action_type)
+            if tool.response_schema is not None:
+                arguments = fix_malformed_tool_arguments(
+                    arguments, tool.response_schema
+                )
             normalized_tool_call = tool_call.model_copy(
                 update={
                     "name": tool_name,

openhands-sdk/openhands/sdk/agent/utils.py

@@ -21,13 +21,15 @@
     overload,
 )
 
+from pydantic import BaseModel
+
 from openhands.sdk.context.condenser.base import CondenserBase
 from openhands.sdk.context.view import View
 from openhands.sdk.conversation.types import ConversationTokenCallbackType
 from openhands.sdk.event.base import LLMConvertibleEvent
 from openhands.sdk.event.condenser import Condensation
 from openhands.sdk.llm import LLM, LLMResponse, Message
-from openhands.sdk.tool import Action, ToolDefinition
+from openhands.sdk.tool import ToolDefinition
 
 
 if TYPE_CHECKING:
@@ -91,8 +93,41 @@ def _is_chunked_str_field(value: Any, expected_origins: list[Any]) -> bool:
     )
 
 
+def _json_schema_expected_types(
+    schema: dict[str, Any], defs: dict[str, Any]
+) -> list[Any]:
+    ref = schema.get("$ref")
+    if isinstance(ref, str) and ref.startswith("#/$defs/"):
+        target = defs.get(ref.rsplit("/", 1)[-1])
+        if isinstance(target, dict):
+            return _json_schema_expected_types(target, defs)
+
+    expected: list[Any] = []
+    for key in ("anyOf", "oneOf"):
+        for option in schema.get(key, []):
+            if isinstance(option, dict):
+                expected.extend(_json_schema_expected_types(option, defs))
+
+    json_types = schema.get("type")
+    if isinstance(json_types, str):
+        json_types = [json_types]
+    if isinstance(json_types, list):
+        type_map = {
+            "array": list,
+            "boolean": bool,
+            "integer": int,
+            "number": float,
+            "object": dict,
+            "string": str,
+        }
+        expected.extend(
+            type_map[json_type] for json_type in json_types if json_type in type_map
+        )
+    return expected
+
+
 def fix_malformed_tool_arguments(
-    arguments: dict[str, Any], action_type: type[Action]
+    arguments: dict[str, Any], action_type: type[BaseModel] | dict[str, Any]
 ) -> dict[str, Any]:
     """Fix malformed tool arguments emitted by some LLMs under native fn calling.
 
@@ -131,7 +166,7 @@ def fix_malformed_tool_arguments(
 
     Args:
         arguments: The parsed arguments dict from json.loads(tool_call.arguments).
-        action_type: The action type that defines the expected schema.
+        action_type: The model or JSON Schema defining expected arguments.
 
     Returns:
         The arguments dict with JSON strings decoded where appropriate.
@@ -141,34 +176,36 @@ def fix_malformed_tool_arguments(
 
     fixed_arguments = arguments.copy()
 
-    # Use model_fields to properly handle aliases and inherited fields
-    for field_name, field_info in action_type.model_fields.items():
-        # Check both the field name and its alias (if any)
-        data_key = field_info.alias if field_info.alias else field_name
+    if isinstance(action_type, dict):
+        defs = action_type.get("$defs", {})
+        field_types = [
+            (field_name, _json_schema_expected_types(field_schema, defs))
+            for field_name, field_schema in action_type.get("properties", {}).items()
+            if isinstance(field_schema, dict)
+        ]
+    else:
+        field_types = []
+        for field_name, field_info in action_type.model_fields.items():
+            data_key = field_info.alias if field_info.alias else field_name
+            expected_type = field_info.annotation
+            if get_origin(expected_type) is Annotated:
+                type_args = get_args(expected_type)
+                expected_type = type_args[0] if type_args else expected_type
+
+            origin = get_origin(expected_type)
+            if origin is Union or origin is types.UnionType:
+                type_args = get_args(expected_type)
+                expected_origins = [get_origin(arg) or arg for arg in type_args]
+            else:
+                expected_origins = [origin or expected_type]
+            field_types.append((data_key, expected_origins))
+
+    for data_key, expected_origins in field_types:
         if data_key not in fixed_arguments:
             continue
 
         value = fixed_arguments[data_key]
 
-        expected_type = field_info.annotation
-
-        # Unwrap Annotated types - only the first arg is the actual type
-        if get_origin(expected_type) is Annotated:
-            type_args = get_args(expected_type)
-            expected_type = type_args[0] if type_args else expected_type
-
-        # Get the origin of the expected type (e.g., list from list[str])
-        origin = get_origin(expected_type)
-
-        # For Union types, we need to check all union members
-        if origin is Union or origin is types.UnionType:
-            # For Union types, check each union member
-            type_args = get_args(expected_type)
-            expected_origins = [get_origin(arg) or arg for arg in type_args]
-        else:
-            # For non-Union types, just check the origin
-            expected_origins = [origin or expected_type]
-
         # Rejoin a str-only field that a model chunked into a JSON array.
         if _is_chunked_str_field(value, expected_origins):
             fixed_arguments[data_key] = "".join(value)

openhands-sdk/openhands/sdk/mcp/tool.py

@@ -291,9 +291,10 @@ def action_from_arguments(self, arguments: dict[str, Any]) -> MCPToolAction:
         Raises:
             ValidationError: If the arguments do not conform to the tool schema.
         """
-        # Drop None-valued keys before validation to avoid type errors
-        # on optional fields
-        prefiltered_args = {k: v for k, v in (arguments or {}).items() if v is not None}
+        tool_arguments, structured_output = self._split_response_arguments(arguments)
+        prefiltered_args = {
+            key: value for key, value in tool_arguments.items() if value is not None
+        }
         # Validate against the dynamically created action type (from MCP schema)
         mcp_action_type = _create_mcp_action_type(self.mcp_tool)
         validated = mcp_action_type.model_validate(prefiltered_args)
@@ -308,7 +309,9 @@ def action_from_arguments(self, arguments: dict[str, Any]) -> MCPToolAction:
             exclude_none=True,
             exclude=exclude_fields,
         )
-        return MCPToolAction(data=sanitized)
+        action = MCPToolAction(data=sanitized)
+        action._structured_output = structured_output
+        return action
 
     @classmethod
     def create(
@@ -408,6 +411,7 @@ def _get_tool_schema(
                 ),
             }
 
+        schema = self._merge_response_schema(schema)
         _prioritize_schema_fields(
             schema=schema,
             priority=("security_risk", "summary"),

openhands-sdk/openhands/sdk/tool/client_tool.py

@@ -298,6 +298,7 @@ def _get_tool_schema(
                     required = merged.setdefault("required", [])
                     if meta not in required:
                         required.append(meta)
+        merged = self._merge_response_schema(merged)
 
         from openhands.sdk.tool.tool import _prioritize_schema_fields
 

openhands-sdk/openhands/sdk/tool/registry.py

@@ -155,7 +155,16 @@ def resolve_tool(
     if resolver is None:
         raise KeyError(f"ToolDefinition '{tool_spec.name}' is not registered")
 
-    return resolver(tool_spec.params, conv_state)
+    params = dict(tool_spec.params)
+    response_schema = params.pop("response_schema", None)
+    tools = resolver(params, conv_state)
+    if response_schema is not None:
+        if len(tools) != 1:
+            raise ValueError(
+                "response_schema requires a spec that resolves to exactly one tool"
+            )
+        tools = [tools[0].set_response_schema(response_schema)]
+    return tools
 
 
 def list_registered_tools() -> list[str]:

openhands-sdk/openhands/sdk/tool/schema.py

@@ -3,7 +3,7 @@
 from collections.abc import Sequence
 from typing import TYPE_CHECKING, Any, ClassVar, TypeVar
 
-from pydantic import ConfigDict, Field, create_model
+from pydantic import ConfigDict, Field, PrivateAttr, create_model
 from rich.text import Text
 
 from openhands.sdk.llm import ImageContent, TextContent
@@ -331,6 +331,12 @@ def from_mcp_schema(
 class Action(Schema, ABC):
     """Base schema for input action."""
 
+    _structured_output: dict[str, Any] | None = PrivateAttr(default=None)
+
+    @property
+    def structured_output(self) -> dict[str, Any] | None:
+        return self._structured_output
+
     @property
     def visualize(self) -> Text:
         """Return Rich Text representation of this action.

openhands-sdk/openhands/sdk/tool/spec.py

@@ -1,6 +1,12 @@
 from typing import Any
 
-from pydantic import BaseModel, Field, field_validator
+from pydantic import (
+    BaseModel,
+    Field,
+    SerializationInfo,
+    field_serializer,
+    field_validator,
+)
 
 
 class Tool(BaseModel):
@@ -37,3 +43,18 @@ def validate_name(cls, v: str) -> str:
     def validate_params(cls, v: dict[str, Any] | None) -> dict[str, Any]:
         """Convert None params to empty dict."""
         return v if v is not None else {}
+
+    @field_serializer("params")
+    def _serialize_params(
+        self, params: dict[str, Any], info: SerializationInfo
+    ) -> dict[str, Any]:
+        """Serialize Pydantic response schemas as JSON Schema."""
+        response_schema = params.get("response_schema")
+        if info.mode != "json" or not (
+            isinstance(response_schema, type) and issubclass(response_schema, BaseModel)
+        ):
+            return params
+        return {
+            **params,
+            "response_schema": response_schema.model_json_schema(),
+        }