diff --git a/CHANGELOG.md b/CHANGELOG.md index f130afd15..bd88f691c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -44,6 +44,8 @@ to include examples, links to docs, or any other relevant information. ### Fixed +- `contrib.google_adk_agents`: agents with an `output_schema` no longer fail every workflow task + when calling the model. The schema type is now sent to the model activity as its JSON schema. - `contrib.deepagents`: prevent duplicate input messages after continue-as-new. - `DataConverter.payload_converter` and current workflow and activity payload converter accessors now return the configured converter without SDK-internal transfer type conversion. diff --git a/temporalio/contrib/google_adk_agents/_model.py b/temporalio/contrib/google_adk_agents/_model.py index 1992d0f4c..71e4a641c 100644 --- a/temporalio/contrib/google_adk_agents/_model.py +++ b/temporalio/contrib/google_adk_agents/_model.py @@ -5,6 +5,8 @@ from google.adk.models import BaseLlm, LLMRegistry from google.adk.models.llm_request import LlmRequest from google.adk.models.llm_response import LlmResponse +from google.genai import types +from pydantic import TypeAdapter import temporalio.workflow from temporalio import activity, workflow @@ -91,6 +93,25 @@ async def invoke_model_streaming( return responses +def _with_serializable_response_schema(llm_request: LlmRequest) -> LlmRequest: + """Return the request with a ``response_schema`` that can be serialized. + + ADK stores an agent's ``output_schema`` on the request as a Python type + (for example a Pydantic model class), which the payload converter cannot + serialize. google-genai and ADK's LiteLlm both turn such a type into its + JSON schema before calling the model, so sending the JSON schema instead + is equivalent. + """ + schema = llm_request.config.response_schema + if schema is None or isinstance(schema, (dict, types.Schema)): + return llm_request + request = llm_request.model_copy() + request.config = llm_request.config.model_copy( + update={"response_schema": TypeAdapter(schema).json_schema()} + ) + return request + + class TemporalModel(BaseLlm): """A Temporal-based LLM model that executes model invocations as activities.""" @@ -183,6 +204,7 @@ async def generate_content_async( if agent_name: config["summary"] = agent_name + llm_request = _with_serializable_response_schema(llm_request) if stream: if self._streaming_topic is None: raise ApplicationError( diff --git a/tests/contrib/google_adk_agents/test_google_adk_agents.py b/tests/contrib/google_adk_agents/test_google_adk_agents.py index fac3138d2..37677ee7d 100644 --- a/tests/contrib/google_adk_agents/test_google_adk_agents.py +++ b/tests/contrib/google_adk_agents/test_google_adk_agents.py @@ -43,6 +43,7 @@ from opentelemetry.sdk.trace.export import SimpleSpanProcessor from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter from opentelemetry.trace import get_tracer_provider, set_tracer_provider +from pydantic import BaseModel, TypeAdapter import temporalio.contrib.google_adk_agents.workflow from temporalio import activity, workflow @@ -53,6 +54,9 @@ TemporalMcpToolSetProvider, TemporalModel, ) +from temporalio.contrib.google_adk_agents._model import ( + _with_serializable_response_schema, +) from temporalio.contrib.opentelemetry import OpenTelemetryPlugin, create_tracer_provider from temporalio.worker import Worker from temporalio.workflow import ActivityConfig @@ -1168,3 +1172,109 @@ async def my_activity(city: str, count: int = 1) -> str: assert params == ["city", "count"] assert sig.parameters["city"].annotation is str assert sig.parameters["count"].default == 1 + + +class CityWeather(BaseModel): + city: str + temperature_c: float + + +class OutputSchemaModel(TestModel): + def responses(self) -> list[LlmResponse]: + return [ + LlmResponse( + content=Content( + role="model", + parts=[Part(text='{"city": "Paris", "temperature_c": 17.5}')], + ) + ) + ] + + @classmethod + def supported_models(cls) -> list[str]: + return ["output_schema_model"] + + +@workflow.defn +class OutputSchemaAgentWorkflow: + @workflow.run + async def run(self, prompt: str) -> dict[str, Any] | None: + agent = LlmAgent( + name="output_schema_agent", + model=TemporalModel("output_schema_model"), + output_schema=CityWeather, + output_key="weather", + ) + runner = InMemoryRunner(agent=agent, app_name="output_schema_app") + session = await runner.session_service.create_session( + app_name="output_schema_app", user_id="test" + ) + async with Aclosing( + runner.run_async( + user_id="test", + session_id=session.id, + new_message=types.Content(role="user", parts=[types.Part(text=prompt)]), + ) + ) as agen: + async for _ in agen: + pass + + final_session = await runner.session_service.get_session( + app_name="output_schema_app", user_id="test", session_id=session.id + ) + return final_session.state.get("weather") if final_session else None + + +@pytest.mark.asyncio +async def test_agent_with_output_schema(client: Client): + LLMRegistry.register(OutputSchemaModel) + + new_config = client.config() + new_config["plugins"] = [GoogleAdkPlugin()] + client = Client(**new_config) + + async with Worker( + client, + task_queue="adk-task-queue-output-schema", + workflows=[OutputSchemaAgentWorkflow], + max_cached_workflows=0, + ): + result = await client.execute_workflow( + OutputSchemaAgentWorkflow.run, + "What is the weather in Paris?", + id=f"output-schema-agent-workflow-{uuid.uuid4()}", + task_queue="adk-task-queue-output-schema", + execution_timeout=timedelta(seconds=60), + ) + + assert result == {"city": "Paris", "temperature_c": 17.5} + + +@pytest.mark.parametrize("schema", [CityWeather, list[CityWeather]]) +def test_output_schema_type_sent_as_json_schema(schema: Any) -> None: + request = LlmRequest( + model="gemini-2.0-flash", + contents=[Content(role="user", parts=[Part(text="hello")])], + config=types.GenerateContentConfig(), + ) + request.set_output_schema(schema) + + converted = _with_serializable_response_schema(request) + + assert request.config.response_schema is schema + assert converted.config.response_mime_type == "application/json" + converter = GoogleAdkPlugin()._configure_data_converter(None) + payloads = converter.payload_converter.to_payloads([converted]) + serialized = json.loads(payloads[0].data) + assert serialized["config"]["response_schema"] == TypeAdapter(schema).json_schema() + + +def test_json_output_schema_left_unchanged() -> None: + request = LlmRequest( + model="gemini-2.0-flash", + contents=[Content(role="user", parts=[Part(text="hello")])], + config=types.GenerateContentConfig(), + ) + request.set_output_schema(CityWeather.model_json_schema()) + + assert _with_serializable_response_schema(request) is request