diff --git a/src/google/adk/agents/llm/task/_finish_task_tool.py b/src/google/adk/agents/llm/task/_finish_task_tool.py index cd747e26f0b..db45863c692 100644 --- a/src/google/adk/agents/llm/task/_finish_task_tool.py +++ b/src/google/adk/agents/llm/task/_finish_task_tool.py @@ -16,6 +16,7 @@ from __future__ import annotations +import copy from typing import Any from typing import Optional from typing import TYPE_CHECKING @@ -92,12 +93,14 @@ def __init__( self.output_schema: SchemaType = ( output_schema if output_schema is not None else _DefaultTaskOutput ) - self._adapter: TypeAdapter[Any] = TypeAdapter(self.output_schema) - raw_schema = self._adapter.json_schema() + if isinstance(self.output_schema, (dict, types.Schema)): + self._adapter: Optional[TypeAdapter[Any]] = None + else: + self._adapter = TypeAdapter(self.output_schema) # FunctionDeclaration parameters must be a JSON object schema. # If the schema is already an object (e.g. BaseModel), use it directly. # Otherwise wrap it in an object with a single key. - self._wrapper_key: str | None = get_output_wrapper_key(raw_schema) + self._wrapper_key: str | None = get_output_wrapper_key(self.output_schema) description = ( 'Signal that this agent has completed its delegated task. Call this' @@ -114,7 +117,10 @@ def __init__( @override def _get_declaration(self) -> Optional[types.FunctionDeclaration]: """Get the function declaration for this tool.""" - raw_schema = self._adapter.json_schema() + if self._adapter is not None: + raw_schema = self._adapter.json_schema() + else: + raw_schema = copy.deepcopy(schema_to_json_schema(self.output_schema)) if self._wrapper_key: # Extract $defs to the root level so $ref pointers remain valid # after wrapping the schema inside an object property. @@ -186,8 +192,11 @@ async def run_async( """ try: raw_value = args.get(self._wrapper_key) if self._wrapper_key else args - validated = self._adapter.validate_python(raw_value) - validated_output = self._adapter.dump_python(validated, mode='json') + if self._adapter is not None: + validated = self._adapter.validate_python(raw_value) + validated_output = self._adapter.dump_python(validated, mode='json') + else: + validated_output = raw_value except ValidationError as e: return { 'error': ( diff --git a/src/google/adk/utils/_schema_utils.py b/src/google/adk/utils/_schema_utils.py index 5bebe3a3f69..e8a8e8a85f5 100644 --- a/src/google/adk/utils/_schema_utils.py +++ b/src/google/adk/utils/_schema_utils.py @@ -130,6 +130,10 @@ def schema_to_json_schema(schema: SchemaType) -> dict[str, Any]: """ if isinstance(schema, dict): return schema + if isinstance(schema, types.Schema): + raw_schema = schema.model_dump(mode="json", exclude_none=True) + lowercase_schema_types(raw_schema) + return raw_schema return TypeAdapter(schema).json_schema() diff --git a/tests/unittests/agents/llm/task/test_finish_task_tool.py b/tests/unittests/agents/llm/task/test_finish_task_tool.py index 9ce4025cbe7..ab4242a4e07 100644 --- a/tests/unittests/agents/llm/task/test_finish_task_tool.py +++ b/tests/unittests/agents/llm/task/test_finish_task_tool.py @@ -19,8 +19,10 @@ from typing import Any from unittest import mock +from google.adk.agents import LlmAgent from google.adk.agents.llm.task._finish_task_tool import FINISH_TASK_TOOL_NAME from google.adk.agents.llm.task._finish_task_tool import FinishTaskTool +from google.genai import types from pydantic import BaseModel import pytest @@ -72,6 +74,28 @@ def test_init_with_output_schema(self): assert 'Signal that this agent has completed' in tool.description assert 'output data' in tool.description + def test_init_with_raw_dict_output_schema(self): + """Test FinishTaskTool initialization with a raw dict output schema.""" + schema = {"type": "object", "properties": {"answer": {"type": "string"}}} + agent = _make_task_agent(output_schema=schema) + tool = FinishTaskTool(task_agent=agent) + assert tool.name == FINISH_TASK_TOOL_NAME + assert tool.output_schema == schema + assert tool._adapter is None + assert tool._wrapper_key is None + + def test_init_with_types_schema_output_schema(self): + """Test FinishTaskTool initialization with a types.Schema output schema.""" + schema = types.Schema( + type="OBJECT", properties={"answer": types.Schema(type="STRING")} + ) + agent = _make_task_agent(output_schema=schema) + tool = FinishTaskTool(task_agent=agent) + assert tool.name == FINISH_TASK_TOOL_NAME + assert tool.output_schema == schema + assert tool._adapter is None + assert tool._wrapper_key is None + def test_get_declaration_without_output_schema(self): """Test function declaration generation without output schema.""" agent = _make_task_agent() @@ -97,6 +121,39 @@ def test_get_declaration_with_output_schema(self): assert 'result' in schema.get('properties', {}) assert 'count' in schema.get('properties', {}) + def test_get_declaration_with_raw_dict_output_schema(self): + """Test function declaration generation with raw dict output schema.""" + schema = { + 'type': 'object', + 'properties': {'answer': {'type': 'string'}}, + 'required': ['answer'], + } + agent = _make_task_agent(output_schema=schema) + tool = FinishTaskTool(task_agent=agent) + declaration = tool._get_declaration() + + assert declaration is not None + assert declaration.name == FINISH_TASK_TOOL_NAME + assert declaration.parameters_json_schema == schema + # Verify original dict was not mutated + assert tool.output_schema == schema + + def test_get_declaration_with_types_schema_output_schema(self): + """Test function declaration generation with types.Schema output schema.""" + schema = types.Schema( + type='OBJECT', properties={'answer': types.Schema(type='STRING')} + ) + agent = _make_task_agent(output_schema=schema) + tool = FinishTaskTool(task_agent=agent) + declaration = tool._get_declaration() + + assert declaration is not None + assert declaration.name == FINISH_TASK_TOOL_NAME + assert declaration.parameters_json_schema == { + 'type': 'object', + 'properties': {'answer': {'type': 'string'}}, + } + @pytest.mark.asyncio async def test_run_async_returns_confirmation(self): """Test that run_async returns a confirmation message.""" @@ -245,6 +302,16 @@ class TestFinishTaskToolAllSchemaTypes: [ (SampleOutputSchema, None), (dict[str, Any], None), + ({'type': 'object', 'properties': {'x': {'type': 'string'}}}, None), + ({'type': 'string'}, 'result'), + ( + types.Schema( + type='OBJECT', + properties={'x': types.Schema(type='STRING')}, + ), + None, + ), + (types.Schema(type='STRING'), 'result'), (str, 'result'), (int, 'result'), (bool, 'result'), @@ -256,6 +323,10 @@ class TestFinishTaskToolAllSchemaTypes: ids=[ 'BaseModel', 'dict', + 'raw_dict_object', + 'raw_dict_string', + 'types_schema_object', + 'types_schema_string', 'str', 'int', 'bool', @@ -305,6 +376,21 @@ def test_get_declaration_wrapped_schema(self, output_schema): {'result': 'done', 'count': 5}, ), (dict[str, Any], {'key': 'value'}, {'key': 'value'}), + ( + {'type': 'object', 'properties': {'answer': {'type': 'string'}}}, + {'answer': 'hello'}, + {'answer': 'hello'}, + ), + ({'type': 'string'}, {'result': 'hello'}, 'hello'), + ( + types.Schema( + type='OBJECT', + properties={'answer': types.Schema(type='STRING')}, + ), + {'answer': 'hello'}, + {'answer': 'hello'}, + ), + (types.Schema(type='STRING'), {'result': 'hello'}, 'hello'), (str, {'result': 'hello'}, 'hello'), (int, {'result': 42}, 42), (bool, {'result': True}, True), @@ -321,6 +407,10 @@ def test_get_declaration_wrapped_schema(self, output_schema): ids=[ 'BaseModel', 'dict', + 'raw_dict_object', + 'raw_dict_string', + 'types_schema_object', + 'types_schema_string', 'str', 'int', 'bool', @@ -361,3 +451,37 @@ def test_get_declaration_list_basemodel_defs_at_root(self): assert 'SampleOutputSchema' in schema['$defs'] # The nested result schema should not contain $defs assert '$defs' not in schema['properties']['result'] + + +class TestFinishTaskToolLlmAgentIntegration: + """Integration tests constructing LlmAgent with mode='task'.""" + + def test_task_agent_with_raw_dict_output_schema(self): + schema = {'type': 'object', 'properties': {'answer': {'type': 'string'}}} + agent = LlmAgent( + name='task_worker', + model='gemini-2.5-flash', + output_schema=schema, + mode='task', + ) + assert agent.output_schema == schema + finish_tools = [t for t in agent.tools if isinstance(t, FinishTaskTool)] + assert len(finish_tools) == 1 + assert finish_tools[0]._adapter is None + assert finish_tools[0]._wrapper_key is None + + def test_task_agent_with_types_schema_output_schema(self): + schema = types.Schema( + type='OBJECT', properties={'answer': types.Schema(type='STRING')} + ) + agent = LlmAgent( + name='task_worker', + model='gemini-2.5-flash', + output_schema=schema, + mode='task', + ) + assert agent.output_schema == schema + finish_tools = [t for t in agent.tools if isinstance(t, FinishTaskTool)] + assert len(finish_tools) == 1 + assert finish_tools[0]._adapter is None + assert finish_tools[0]._wrapper_key is None diff --git a/tests/unittests/utils/test_schema_utils.py b/tests/unittests/utils/test_schema_utils.py index 719c404eb2c..9d3a1338d37 100644 --- a/tests/unittests/utils/test_schema_utils.py +++ b/tests/unittests/utils/test_schema_utils.py @@ -466,6 +466,17 @@ def test_list_of_basemodel_schema_becomes_an_array_of_objects(self): ref = result["items"]["$ref"].rsplit("/", 1)[-1] assert result["$defs"][ref]["properties"]["name"]["type"] == "string" + def test_types_schema_is_converted_to_json_schema(self): + schema = types.Schema( + type="OBJECT", + properties={"answer": types.Schema(type="STRING")}, + ) + result = schema_to_json_schema(schema) + assert result == { + "type": "object", + "properties": {"answer": {"type": "string"}}, + } + class TestLowercaseSchemaTypes: """Tests for lowercase_schema_types function."""