From 7c30b202335fcf6054be9ebf254406d469a354e3 Mon Sep 17 00:00:00 2001 From: pcerypeng Date: Thu, 23 Jul 2026 17:39:05 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E6=94=AF=E6=8C=81=20OpenAI=20Responses?= =?UTF-8?q?=20API=E3=80=81Agent=20Node=20HITL=20=E5=A4=9A=E8=BD=AE?= =?UTF-8?q?=E4=BA=A4=E4=BA=92=E5=8F=8A=E5=B7=A5=E5=85=B7=E9=94=99=E8=AF=AF?= =?UTF-8?q?=E6=A3=80=E6=B5=8B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 新增功能: - OpenAI Responses API 适配 (非流式/流式), 支持 reasoning、tool calls、logprobs - Agent 节点多轮 HITL 机制, 通过 interrupt bridge 桥接子 agent LongRunningEvent - AG-UI GraphAgent checkpoint 保护, 防止工具结果恢复时覆盖 LangGraph checkpoint 代码修复 (代码审查): - _constants.py: 将 STATE_KEY_PENDING_AGENT_NODE_HITL 加入 UNSAFE_STATE_KEYS, 防止含敏感工具参数的 child_state 通过 completion 事件对外暴露 - _openai_model.py: 非流式 Responses 路径分离 http_options 传参, 与流式路径 保持一致, 修复 extra_body 被误传为 responses.create 顶层参数的 Bug - _llm_agent.py: 删除重复的 logger.debug 行; 在长运行工具与并行工具批混用时 发出警告, 提示同批其它工具结果不会被 LLM 进一步处理 - _long_running_tool.py: 提取 TOOL_ERROR_CODE_* 共享常量, 消除与 _tools_processor.py 中的硬编码字符串重复 - _agui_agent.py: 优化 GraphAgent checkpoint 检测——复用 _ensure_session_exists 返回的 session 消除额外 DB 查询, 用模块级常量替代硬编码前缀, 改为同步方法 - _session_manager.py: 新增 session_service 公共属性, 消除私有成员访问 新增测试: - AgentNode HITL 中断后进程重启恢复场景 (SqlSessionService) - STATE_KEY_PENDING_AGENT_NODE_HITL unsafe 归类验证 - 多轮 HITL 客户端使用过期 function_call.id 提交 resume 时不静默完成 - 非流式 Responses 路径 http_options 含 extra_body/extra_headers/timeout 时参数分离 - GraphAgent checkpoint 保护测试 - test_constants.py 更新覆盖新增 unsafe key --- docs/mkdocs/en/graph.md | 4 + docs/mkdocs/en/model.md | 26 + docs/mkdocs/zh/graph.md | 4 + docs/mkdocs/zh/model.md | 25 + pyproject.toml | 2 +- tests/agents/test_llm_agent_ext.py | 169 ++++ tests/models/test_openai_responses_model.py | 860 ++++++++++++++++++ tests/server/ag_ui/_core/test_agui_agent.py | 41 + tests/tools/test_function_tool.py | 6 +- tests/tools/test_long_running_tool.py | 61 +- .../graph/test_agent_node_hitl.py | 495 ++++++++++ tests/trpc_agent_dsl/graph/test_constants.py | 3 + tests/trpc_agent_dsl/graph/test_events.py | 20 + trpc_agent_sdk/agents/_llm_agent.py | 120 +-- .../agents/core/_tools_processor.py | 18 +- trpc_agent_sdk/dsl/graph/_constants.py | 40 +- trpc_agent_sdk/dsl/graph/_events/_builder.py | 9 +- trpc_agent_sdk/dsl/graph/_graph_agent.py | 46 +- .../dsl/graph/_node_action/_agent.py | 163 +++- trpc_agent_sdk/models/_openai_model.py | 754 ++++++++++++--- .../server/ag_ui/_core/_agui_agent.py | 49 +- .../server/ag_ui/_core/_session_manager.py | 9 + trpc_agent_sdk/tools/__init__.py | 157 ++-- trpc_agent_sdk/tools/_function_tool.py | 11 +- trpc_agent_sdk/tools/_long_running_tool.py | 41 +- 25 files changed, 2748 insertions(+), 385 deletions(-) create mode 100644 tests/models/test_openai_responses_model.py create mode 100644 tests/trpc_agent_dsl/graph/test_agent_node_hitl.py diff --git a/docs/mkdocs/en/graph.md b/docs/mkdocs/en/graph.md index 83ed3038..e93c6120 100644 --- a/docs/mkdocs/en/graph.md +++ b/docs/mkdocs/en/graph.md @@ -429,6 +429,10 @@ Common options: - input_mapper / output_mapper: Parent-child state mapping (explicit configuration recommended) - config / callbacks: Same as add_node +When a child Agent emits a `LongRunningEvent`, `add_agent_node` promotes it to a parent GraphAgent `interrupt`: the parent graph does not execute downstream nodes, and the Runner event preserves the original tool name and arguments. After the client submits the matching `FunctionResponse`, the parent graph resumes the current Agent node and the SDK maps the response back to the child's original function call. The graph continues only after the child Agent reaches a final result. This supports multiple HITL rounds in one node and persists child state in the SessionService-backed checkpoint for process-restart recovery. + +When the Agent node is a TeamAgent, the Leader can use the same HITL flow. Regular Team Members still must not be configured with or invoke `LongRunningFunctionTool`. + GraphAgent does not require (nor support) registering Agent nodes via sub_agents; composition relationships are handled uniformly through add_agent_node. ## Advanced Usage diff --git a/docs/mkdocs/en/model.md b/docs/mkdocs/en/model.md index 5b9618e7..539dd617 100644 --- a/docs/mkdocs/en/model.md +++ b/docs/mkdocs/en/model.md @@ -121,6 +121,32 @@ model = OpenAIModel( ) ``` +#### Responses API + +`OpenAIModel` uses Chat Completions by default. Enable the Responses API explicitly for OpenAI or compatible +providers that expose `/v1/responses`: + +```python +model = OpenAIModel( + model_name="your-responses-model", + api_key="your-api-key", + base_url="https://api.openai.com/v1", + use_responses_api=True, + responses_api_params={ + "store": False, + "reasoning": {"summary": "auto"}, + }, +) +``` + +The switch is opt-in so existing OpenAI-compatible providers continue to use Chat Completions. The adapter maps +conversation history, function calls and `function_call_output` items, structured output, semantic streaming events, +reasoning summaries, and token usage into the existing tRPC-Agent types. When `store=False`, the SDK automatically +requests `reasoning.encrypted_content` so reasoning items can be replayed with tool outputs in the next turn. + +`responses_api_params` accepts Responses-only fields such as `store`, `reasoning`, `include`, and `truncation`. +`model`, `input`, and `stream` are managed by `OpenAIModel` and cannot be overridden there. + #### Advanced Usage Since version `1.1.10`, `OpenAIModel` supports passing a shared HTTP client provider to enable connection reuse. By default, `OpenAIModel` creates a temporary HTTP client for each model-service request. If you want to reuse connections, use the following configuration: diff --git a/docs/mkdocs/zh/graph.md b/docs/mkdocs/zh/graph.md index f5683e0e..15df04ed 100644 --- a/docs/mkdocs/zh/graph.md +++ b/docs/mkdocs/zh/graph.md @@ -428,6 +428,10 @@ graph.add_agent_node( - input_mapper / output_mapper：父子状态映射（推荐显式配置） - config / callbacks：同 add_node +当子 Agent 发出 `LongRunningEvent` 时，`add_agent_node` 会将其提升为父 GraphAgent 的 `interrupt`：父图不会执行后继节点，Runner 返回的事件保留原工具名和参数。客户端提交对应 `FunctionResponse` 后，父图恢复当前 Agent 节点，SDK 将响应映射回子 Agent 的原始 function call；只有子 Agent 最终完成后父图才继续。该机制支持同一节点多轮 HITL，并将 child state 写入 SessionService-backed checkpoint，服务重启后仍可恢复。 + +TeamAgent 作为 Agent 节点时同样支持 Leader 发起 HITL；普通 Team Member 仍不允许配置或调用 `LongRunningFunctionTool`。 + GraphAgent 不需要（也不支持）通过 sub_agents 注册 Agent 节点；组合关系统一用 add_agent_node 完成。 ## 进阶用法 diff --git a/docs/mkdocs/zh/model.md b/docs/mkdocs/zh/model.md index 16e2320b..6209891b 100644 --- a/docs/mkdocs/zh/model.md +++ b/docs/mkdocs/zh/model.md @@ -121,6 +121,31 @@ model = OpenAIModel( ) ``` +#### Responses API + +`OpenAIModel` 默认仍使用 Chat Completions。对于 OpenAI 或提供 `/v1/responses` 的兼容服务，需要显式开启： + +```python +model = OpenAIModel( + model_name="your-responses-model", + api_key="your-api-key", + base_url="https://api.openai.com/v1", + use_responses_api=True, + responses_api_params={ + "store": False, + "reasoning": {"summary": "auto"}, + }, +) +``` + +该开关默认关闭，现有 OpenAI-compatible 服务不会改变调用路径。适配层会把多轮消息、函数调用及 +`function_call_output`、结构化输出、语义化流式事件、reasoning summary 和 token usage 映射为现有 +tRPC-Agent 类型。设置 `store=False` 时，SDK 会自动请求 `reasoning.encrypted_content`，以便下一轮 +连同工具结果一起回放 reasoning item。 + +`responses_api_params` 可传入 `store`、`reasoning`、`include`、`truncation` 等 Responses 专用字段； +`model`、`input`、`stream` 由 `OpenAIModel` 管理，不能在此覆盖。 + #### 高级用法 从版本 `1.1.10`之后 OpenAIModel 支持传入共享的 http client 来解决连接复用的场景，当前的 OpenAIModel 默认每次都会创建临时的 http client 去访问模型服务；如果期望连接复用可以使用如下的方式 diff --git a/pyproject.toml b/pyproject.toml index db0e7828..33c466a5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,7 +25,7 @@ classifiers = [ ] dependencies = [ "pydantic>=2.11.3", - "openai>=1.3.0", + "openai>=1.66.0", "mcp>=1.10.1", "aiohttp", "httpx>=0.27.0", diff --git a/tests/agents/test_llm_agent_ext.py b/tests/agents/test_llm_agent_ext.py index 5dc3ad32..daae4961 100644 --- a/tests/agents/test_llm_agent_ext.py +++ b/tests/agents/test_llm_agent_ext.py @@ -409,5 +409,174 @@ async def run(): assert events[0].error_code == "build_error" +# --------------------------------------------------------------------------- +# _run_async_impl — LongRunningFunctionTool integration +# --------------------------------------------------------------------------- + + +class TestRunAsyncImplLongRunningTool: + """Cover the long-running tool path in _run_async_impl (lines 562-620).""" + + def test_yields_long_running_event_on_tool_completion(self, invocation_context): + """When LLM calls a LongRunningFunctionTool, a LongRunningEvent is yielded.""" + from trpc_agent_sdk.events import LongRunningEvent + from trpc_agent_sdk.tools import LongRunningFunctionTool + from trpc_agent_sdk.types import FunctionCall, FunctionResponse + + async def long_op(status: str) -> dict: + """A long-running operation.""" + return {"status": status, "result": "done"} + + long_tool = LongRunningFunctionTool(long_op) + agent = _agent(tools=[long_tool]) + invocation_context.agent = agent + + fc = FunctionCall(id="call_lr_1", name="long_op", args={"status": "check"}) + fr = FunctionResponse(id="call_lr_1", name="long_op", response={"status": "done", "result": "ok"}) + llm_event = Event( + invocation_id="inv-1", + author="test_agent", + content=Content( + parts=[Part.from_function_call(name="long_op", args={"status": "check"})], + role="model", + ), + ) + llm_event.content.parts[0].function_call.id = "call_lr_1" + + tool_event = Event( + invocation_id="inv-1", + author="test_agent", + content=Content(parts=[Part(function_response=fr)], role="user"), + ) + + async def run(): + events = [] + with patch( + "trpc_agent_sdk.agents._llm_agent.default_request_processor" + ) as mock_rp: + mock_rp.build_request = AsyncMock(return_value=None) + with patch( + "trpc_agent_sdk.agents._llm_agent.LlmProcessor" + ) as mock_lp_cls: + mock_lp = MagicMock() + + async def fake_call_llm_async(request, ctx=None, stream=True): + for ev in [llm_event]: + yield ev + + mock_lp.call_llm_async = fake_call_llm_async + mock_lp_cls.return_value = mock_lp + with patch.object( + agent, "_get_extended_tools_processor" + ) as mock_get_tp: + mock_tp = MagicMock() + mock_tp.find_tool = AsyncMock(return_value=long_tool) + + async def fake_execute_tools_async(tool_calls, context): + for ev in [tool_event]: + yield ev + + mock_tp.execute_tools_async = fake_execute_tools_async + mock_get_tp.return_value = mock_tp + async for event in agent._run_async_impl(invocation_context): + events.append(event) + return events + + events = asyncio.run(run()) + long_running_events = [e for e in events if isinstance(e, LongRunningEvent)] + assert len(long_running_events) == 1 + lre = long_running_events[0] + assert lre.function_call.id == "call_lr_1" + assert lre.function_call.name == "long_op" + assert lre.function_response.id == "call_lr_1" + + def test_long_running_tool_error_does_not_yield_long_running_event(self, invocation_context): + """is_tool_execution_error suppresses LongRunningEvent; tool error is yielded normally.""" + from trpc_agent_sdk.tools import LongRunningFunctionTool + from trpc_agent_sdk.types import FunctionCall, FunctionResponse + + async def long_op(status: str) -> dict: + """A long-running operation.""" + return {"status": status} + + long_tool = LongRunningFunctionTool(long_op) + agent = _agent(tools=[long_tool]) + invocation_context.agent = agent + + fr = FunctionResponse(id="call_err_1", name="long_op", response={"error": "fail"}) + llm_tool_call_event = Event( + invocation_id="inv-1", + author="test_agent", + content=Content( + parts=[Part.from_function_call(name="long_op", args={"status": "check"})], + role="model", + ), + ) + llm_tool_call_event.content.parts[0].function_call.id = "call_err_1" + + llm_final_event = Event( + invocation_id="inv-1", + author="test_agent", + content=Content(parts=[Part.from_text(text="done")], role="model"), + ) + + tool_event = Event( + invocation_id="inv-1", + author="test_agent", + error_code="tool_execution_error", + error_message="Tool failed", + content=Content(parts=[Part(function_response=fr)], role="user"), + ) + + async def run(): + events = [] + with patch( + "trpc_agent_sdk.agents._llm_agent.default_request_processor" + ) as mock_rp: + mock_rp.build_request = AsyncMock(return_value=None) + with patch( + "trpc_agent_sdk.agents._llm_agent.LlmProcessor" + ) as mock_lp_cls: + mock_lp = MagicMock() + + call_count = 0 + + async def fake_call_llm_async(request, ctx=None, stream=True): + nonlocal call_count + call_count += 1 + if call_count == 1: + for ev in [llm_tool_call_event]: + yield ev + else: + for ev in [llm_final_event]: + yield ev + + mock_lp.call_llm_async = fake_call_llm_async + mock_lp_cls.return_value = mock_lp + with patch.object( + agent, "_get_extended_tools_processor" + ) as mock_get_tp: + mock_tp = MagicMock() + mock_tp.find_tool = AsyncMock(return_value=long_tool) + + async def fake_execute_tools_async(tool_calls, context): + for ev in [tool_event]: + yield ev + + mock_tp.execute_tools_async = fake_execute_tools_async + mock_get_tp.return_value = mock_tp + async for event in agent._run_async_impl(invocation_context): + events.append(event) + return events + + events = asyncio.run(run()) + from trpc_agent_sdk.events import LongRunningEvent + long_running_events = [e for e in events if isinstance(e, LongRunningEvent)] + assert len(long_running_events) == 0 + # The regular tool event and final text event are still yielded + regular_events = [e for e in events if not isinstance(e, LongRunningEvent)] + assert len(regular_events) >= 1 + + if __name__ == "__main__": pytest.main([__file__, "-v"]) diff --git a/tests/models/test_openai_responses_model.py b/tests/models/test_openai_responses_model.py new file mode 100644 index 00000000..ecbb30dd --- /dev/null +++ b/tests/models/test_openai_responses_model.py @@ -0,0 +1,860 @@ +# Tencent is pleased to support the open source community by making tRPC-Agent-Python available. +# +# Copyright (C) 2026 Tencent. All rights reserved. +# +# tRPC-Agent-Python is licensed under Apache-2.0. + +from unittest.mock import AsyncMock, Mock, patch + +import pytest +from trpc_agent_sdk.models import LlmRequest, OpenAIModel +from trpc_agent_sdk.types import Content, GenerateContentConfig, HttpOptions, Part + + +def _model(**kwargs): + """Create an OpenAIModel with test defaults.""" + kwargs.setdefault("model_name", "gpt-4") + kwargs.setdefault("api_key", "test_key") + return OpenAIModel(**kwargs) + + +def _request(contents, config=None, streaming_tool_names=None): + """Create an LlmRequest for Responses API tests.""" + request = LlmRequest(contents=contents, config=config, tools_dict={}) + if streaming_tool_names is not None: + request.streaming_tool_names = streaming_tool_names + return request + + +# --------------------------------------------------------------------------- +# Responses API +# --------------------------------------------------------------------------- + + +class TestOpenAIResponsesAPI: + """Tests for the opt-in OpenAI Responses transport.""" + + def test_is_disabled_by_default_and_rejects_managed_overrides(self): + model = _model() + assert model.use_responses_api is False + + with pytest.raises(ValueError, match="cannot override managed parameters: input"): + _model(use_responses_api=True, responses_api_params={"input": "override"}) + + def test_converts_tool_history_and_function_definitions(self): + model = _model(use_responses_api=True) + messages = [ + { + "role": "assistant", + "content": "Checking", + "tool_calls": [ + { + "id": "call_weather", + "type": "function", + "function": { + "name": "weather", + "arguments": '{"city":"Shenzhen"}', + }, + } + ], + }, + {"role": "tool", "tool_call_id": "call_weather", "content": '{"temperature":30}'}, + ] + + items = model._convert_messages_to_responses_input(messages) + assert items == [ + {"role": "assistant", "content": "Checking"}, + { + "type": "function_call", + "call_id": "call_weather", + "name": "weather", + "arguments": '{"city":"Shenzhen"}', + }, + { + "type": "function_call_output", + "call_id": "call_weather", + "output": '{"temperature":30}', + }, + ] + tools = model._convert_tools_to_responses_format( + [ + { + "type": "function", + "function": { + "name": "weather", + "description": "Get weather", + "parameters": {"type": "object", "properties": {}}, + }, + } + ] + ) + assert tools == [ + { + "type": "function", + "name": "weather", + "description": "Get weather", + "parameters": {"type": "object", "properties": {}}, + } + ] + + def test_converts_multimodal_input(self): + model = _model(use_responses_api=True) + + items = model._convert_messages_to_responses_input( + [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "Inspect this image", + }, + { + "type": "image_url", + "image_url": { + "url": "data:image/png;base64,AAAA", + "detail": "high", + }, + }, + ], + } + ] + ) + + assert items == [ + { + "role": "user", + "content": [ + { + "type": "input_text", + "text": "Inspect this image", + }, + { + "type": "input_image", + "image_url": "data:image/png;base64,AAAA", + "detail": "high", + }, + ], + } + ] + + def test_preserves_empty_reasoning_item_for_tool_continuation(self): + model = _model(use_responses_api=True) + raw_reasoning = { + "id": "rs_empty", + "type": "reasoning", + "encrypted_content": "encrypted-reasoning", + "summary": [], + } + + response = model._create_responses_response( + { + "id": "resp_tool", + "status": "completed", + "output": [ + raw_reasoning, + { + "type": "function_call", + "call_id": "call_1", + "name": "lookup", + "arguments": "{}", + }, + ], + } + ) + + thought = response.content.parts[0] + assert thought.thought is True + assert thought.text == "" + replay_request = _request([Content(parts=response.content.parts, role="model")]) + replay_items = model._convert_messages_to_responses_input(model._format_messages(replay_request)) + assert replay_items[0] == raw_reasoning + assert replay_items[1]["type"] == "function_call" + + def test_converts_structured_output_and_failed_response(self): + model = _model(use_responses_api=True) + text = model._convert_response_format_to_responses( + { + "type": "json_schema", + "json_schema": { + "name": "answer", + "schema": {"type": "object", "properties": {}}, + "strict": True, + }, + } + ) + assert text == { + "format": { + "type": "json_schema", + "name": "answer", + "schema": {"type": "object", "properties": {}}, + "strict": True, + }, + } + + response = model._create_responses_response( + { + "id": "resp_failed", + "status": "failed", + "error": {"code": "server_error", "message": "upstream failed"}, + "output": [], + } + ) + assert response.error_code == "server_error" + assert response.error_message == "upstream failed" + + def test_logprobs_uses_the_installed_responses_client_shape(self): + model = _model(use_responses_api=True) + params = model._convert_api_params_to_responses( + { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Hello"}], + "stream": False, + "logprobs": True, + "top_logprobs": 3, + } + ) + + class TopLogprobsResponses: + async def create(self, *, top_logprobs, **kwargs): + del top_logprobs, kwargs + + class TopLogprobsClient: + responses = TopLogprobsResponses() + + prepared = model._prepare_responses_api_params(TopLogprobsClient(), params) + assert prepared["top_logprobs"] == 3 + assert "_trpc_responses_logprobs_request" not in prepared + + def test_logprobs_fails_clearly_when_the_installed_client_lacks_support(self): + model = _model(use_responses_api=True) + params = model._convert_api_params_to_responses( + { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Hello"}], + "stream": False, + "logprobs": True, + "top_logprobs": 3, + } + ) + + class LegacyResponses: + async def create(self, *, model, input, stream): + del model, input, stream + + class LegacyClient: + responses = LegacyResponses() + + with pytest.raises(ValueError, match="upgrade openai or disable logprobs"): + model._prepare_responses_api_params(LegacyClient(), params) + + @pytest.mark.asyncio + async def test_non_streaming_uses_responses_create_and_maps_output(self): + model = _model( + use_responses_api=True, + responses_api_params={"store": False, "truncation": "auto"}, + ) + request = _request( + [Content(parts=[Part.from_text(text="Hello")], role="user")], + GenerateContentConfig(max_output_tokens=128), + ) + response = Mock() + response.model_dump.return_value = { + "id": "resp_123", + "status": "completed", + "output": [ + { + "id": "rs_123", + "type": "reasoning", + "encrypted_content": "encrypted-reasoning", + "summary": [{"type": "summary_text", "text": "Check context"}], + }, + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "Hello back"}], + }, + { + "type": "function_call", + "call_id": "call_123", + "name": "lookup", + "arguments": '{"query":"hello"}', + }, + ], + "usage": { + "input_tokens": 10, + "output_tokens": 7, + "total_tokens": 17, + "input_tokens_details": {"cached_tokens": 4}, + "output_tokens_details": {"reasoning_tokens": 2}, + }, + } + captured = {} + + async def create(**kwargs): + captured.update(kwargs) + return response + + with patch.object(model, "_create_async_client") as client_factory: + client = AsyncMock() + client.responses.create = create + client_factory.return_value = client + responses = [item async for item in model.generate_async(request, stream=False)] + + assert captured["model"] == "gpt-4" + assert captured["input"] == [{"role": "user", "content": "Hello"}] + assert captured["max_output_tokens"] == 128 + assert captured["store"] is False + assert captured["truncation"] == "auto" + assert captured["include"] == ["reasoning.encrypted_content"] + assert "messages" not in captured + assert "max_completion_tokens" not in captured + result = responses[0] + assert result.response_id == "resp_123" + assert [part.text for part in result.content.parts if part.text] == ["Check context", "Hello back"] + assert result.content.parts[0].thought is True + assert result.content.parts[-1].function_call.id == "call_123" + assert result.usage_metadata.prompt_token_count == 10 + assert result.usage_metadata.candidates_token_count == 7 + assert result.usage_metadata.thoughts_token_count == 2 + assert result.usage_metadata.cache_read_input_tokens == 4 + + replay_request = _request([Content(parts=result.content.parts, role="model")]) + replay_items = model._convert_messages_to_responses_input(model._format_messages(replay_request)) + assert replay_items[0] == { + "id": "rs_123", + "type": "reasoning", + "encrypted_content": "encrypted-reasoning", + "summary": [{"type": "summary_text", "text": "Check context"}], + } + assert replay_items[1] == {"role": "assistant", "content": "Hello back"} + assert replay_items[2]["type"] == "function_call" + + @pytest.mark.asyncio + async def test_non_streaming_passes_http_options_separately(self): + """Non-streaming Responses path must pass http_options (extra_body, + extra_headers, timeout) as separate kwargs to responses.create, + not merged into api_params (which would cause extra_body to be + treated as an unknown top-level parameter). + """ + model = _model(use_responses_api=True) + request = _request( + [Content(parts=[Part.from_text(text="Hello")], role="user")], + GenerateContentConfig( + max_output_tokens=128, + http_options=HttpOptions( + headers={"X-Custom-Header": "test-value"}, + timeout=5000, + extra_body={"custom_param": "custom_value"}, + ), + ), + ) + response = Mock() + response.model_dump.return_value = { + "id": "resp_456", + "status": "completed", + "output": [ + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "Hi there"}], + }, + ], + "usage": {"input_tokens": 5, "output_tokens": 3, "total_tokens": 8}, + } + captured = {} + + async def create(**kwargs): + captured.update(kwargs) + return response + + with patch.object(model, "_create_async_client") as client_factory: + client = AsyncMock() + client.responses.create = create + client_factory.return_value = client + responses = [item async for item in model.generate_async(request, stream=False)] + + # Core API params must be present. + assert captured["model"] == "gpt-4" + assert captured["input"] == [{"role": "user", "content": "Hello"}] + assert captured["max_output_tokens"] == 128 + + # http_options must be passed as separate kwargs, not merged into + # the api_params dict that goes through _prepare_responses_api_params. + assert captured["extra_headers"] == {"X-Custom-Header": "test-value"} + assert captured["extra_body"] == {"custom_param": "custom_value"} + assert captured["timeout"] == 5.0 # 5000ms / 1000 + + # The response should be properly mapped. + result = responses[0] + assert result.response_id == "resp_456" + assert [part.text for part in result.content.parts if part.text] == ["Hi there"] + + @pytest.mark.asyncio + async def test_streaming_maps_text_reasoning_tool_calls_and_usage(self): + model = _model(use_responses_api=True) + request = _request( + [Content(parts=[Part.from_text(text="Use a tool")], role="user")], + streaming_tool_names={"lookup"}, + ) + events = [ + {"type": "response.created", "response": {"id": "resp_stream"}}, + {"type": "response.reasoning_summary_text.delta", "delta": "Thinking"}, + {"type": "response.output_text.delta", "delta": "Working"}, + { + "type": "response.output_item.added", + "item": { + "type": "function_call", + "id": "fc_1", + "call_id": "call_1", + "name": "lookup", + "arguments": "", + }, + }, + {"type": "response.function_call_arguments.delta", "item_id": "fc_1", "delta": '{"q":'}, + {"type": "response.function_call_arguments.delta", "item_id": "fc_1", "delta": '"x"}'}, + { + "type": "response.completed", + "response": { + "id": "resp_stream", + "status": "completed", + "output": [ + {"type": "reasoning", "summary": [{"type": "summary_text", "text": "Thinking"}]}, + { + "type": "message", + "content": [{"type": "output_text", "text": "Working"}], + }, + { + "type": "function_call", + "id": "fc_1", + "call_id": "call_1", + "name": "lookup", + "arguments": '{"q":"x"}', + }, + ], + "usage": {"input_tokens": 5, "output_tokens": 4, "total_tokens": 9}, + }, + }, + ] + + async def stream_events(): + for event in events: + yield event + + with patch.object(model, "_create_async_client") as client_factory: + client = AsyncMock() + client.responses.create = AsyncMock(return_value=stream_events()) + client_factory.return_value = client + responses = [item async for item in model.generate_async(request, stream=True)] + + assert responses[0].content.parts[0].thought is True + assert responses[0].content.parts[0].text == "Thinking" + assert responses[1].content.parts[0].text == "Working" + tool_deltas = [item.content.parts[0].function_call.args["tool_streaming_args"] for item in responses[2:-1]] + assert tool_deltas == ['{"q":', '"x"}'] + final = responses[-1] + assert final.partial is False + assert final.response_id == "resp_stream" + assert final.content.parts[-1].function_call.args == {"q": "x"} + assert final.usage_metadata.total_token_count == 9 + + @pytest.mark.asyncio + async def test_streaming_error_event_is_not_reported_as_success(self): + model = _model(use_responses_api=True) + request = _request([Content(parts=[Part.from_text(text="Hello")], role="user")]) + + async def stream_events(): + yield { + "type": "error", + "response_id": "resp_error", + "code": "server_error", + "message": "upstream unavailable", + } + + with patch.object(model, "_create_async_client") as client_factory: + client = AsyncMock() + client.responses.create = AsyncMock(return_value=stream_events()) + client_factory.return_value = client + responses = [item async for item in model.generate_async(request, stream=True)] + + assert responses[-1].response_id == "resp_error" + assert responses[-1].error_code == "server_error" + assert responses[-1].error_message == "upstream unavailable" + + @pytest.mark.asyncio + async def test_streaming_uses_arguments_done_payload_in_fallback(self): + model = _model(use_responses_api=True) + request = _request([Content(parts=[Part.from_text(text="Use a tool")], role="user")]) + + async def stream_events(): + yield { + "type": "response.output_item.added", + "item": { + "type": "function_call", + "id": "fc_1", + "call_id": "call_1", + "name": "lookup", + "arguments": "", + }, + } + yield { + "type": "response.function_call_arguments.done", + "item_id": "fc_1", + "arguments": '{"query":"done"}', + } + + with patch.object(model, "_create_async_client") as client_factory: + client = AsyncMock() + client.responses.create = AsyncMock(return_value=stream_events()) + client_factory.return_value = client + responses = [item async for item in model.generate_async(request, stream=True)] + + final = responses[-1] + assert final.partial is False + assert final.content.parts[0].function_call.id == "call_1" + assert final.content.parts[0].function_call.args == {"query": "done"} + + def test_incomplete_response_maps_reason_to_error(self): + model = _model(use_responses_api=True) + + response = model._create_responses_response( + { + "id": "resp_incomplete", + "status": "incomplete", + "incomplete_details": { + "reason": "max_output_tokens", + }, + "output": [], + } + ) + + assert response.error_code == "incomplete" + assert response.error_message == "max_output_tokens" + + # ------------------------------------------------------------------ + # _responses_error — cancelled / incomplete_details fallback + # ------------------------------------------------------------------ + + def test_cancelled_response_maps_status_to_error(self): + model = _model(use_responses_api=True) + response = model._create_responses_response( + { + "id": "resp_cancelled", + "status": "cancelled", + "output": [], + } + ) + assert response.error_code == "cancelled" + assert response.error_message == "cancelled" + + def test_incomplete_details_non_dict_fallback(self): + model = _model(use_responses_api=True) + response = model._create_responses_response( + { + "id": "resp_incomplete", + "status": "incomplete", + "incomplete_details": "timeout", + "output": [], + } + ) + assert response.error_code == "incomplete" + assert response.error_message == "timeout" + + # ------------------------------------------------------------------ + # _create_responses_response — refusal text / non-dict arguments + # ------------------------------------------------------------------ + + def test_message_with_refusal_text(self): + model = _model(use_responses_api=True) + response = model._create_responses_response( + { + "id": "resp_refusal", + "status": "completed", + "output": [ + { + "type": "message", + "role": "assistant", + "content": [{"type": "refusal", "refusal": "I cannot answer that."}], + } + ], + } + ) + assert response.content.parts[0].text == "I cannot answer that." + + def test_function_call_non_dict_arguments_skipped(self): + model = _model(use_responses_api=True) + response = model._create_responses_response( + { + "id": "resp_bad_args", + "status": "completed", + "output": [ + { + "type": "function_call", + "call_id": "call_bad", + "name": "broken_tool", + "arguments": "[not a dict]", + } + ], + } + ) + # Non-dict arguments are skipped, so no parts from function_call + assert response.content is None + + # ------------------------------------------------------------------ + # _convert_messages_to_responses_input — assistant / unknown type + # ------------------------------------------------------------------ + + def test_converts_assistant_message_to_output_text(self): + model = _model(use_responses_api=True) + items = model._convert_messages_to_responses_input( + [ + { + "role": "assistant", + "content": [{"type": "text", "text": "Hello from assistant"}], + } + ] + ) + assert items == [ + { + "role": "assistant", + "content": [{"type": "output_text", "text": "Hello from assistant"}], + } + ] + + def test_unknown_content_type_is_passthrough(self): + model = _model(use_responses_api=True) + items = model._convert_messages_to_responses_input( + [ + { + "role": "user", + "content": [{"type": "custom_block", "data": "raw"}], + } + ] + ) + assert items == [ + { + "role": "user", + "content": [{"type": "custom_block", "data": "raw"}], + } + ] + + def test_converts_non_function_tool_to_responses_format(self): + model = _model(use_responses_api=True) + tools = model._convert_tools_to_responses_format( + [ + {"type": "web_search", "name": "search"}, + { + "type": "function", + "function": { + "name": "calc", + "description": "Calculate", + "parameters": {"type": "object", "properties": {}}, + }, + }, + ] + ) + assert tools == [ + {"type": "web_search", "name": "search"}, + { + "type": "function", + "name": "calc", + "description": "Calculate", + "parameters": {"type": "object", "properties": {}}, + }, + ] + + # ------------------------------------------------------------------ + # _prepare_responses_api_params — logprobs structured support + # ------------------------------------------------------------------ + + def test_logprobs_uses_structured_logprobs_when_supported(self): + model = _model(use_responses_api=True) + params = model._convert_api_params_to_responses( + { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Hello"}], + "stream": False, + "logprobs": True, + "top_logprobs": 3, + } + ) + + class LogprobsResponses: + async def create(self, *, logprobs, **kwargs): + del logprobs, kwargs + + class LogprobsClient: + responses = LogprobsResponses() + + prepared = model._prepare_responses_api_params(LogprobsClient(), params) + assert prepared["logprobs"] == {"enabled": True, "top_logprobs": 3} + + def test_logprobs_raises_on_inspect_failure(self): + model = _model(use_responses_api=True) + params = model._convert_api_params_to_responses( + { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Hello"}], + "stream": False, + "logprobs": True, + "top_logprobs": 3, + } + ) + + class BrokenResponses: + create = "not_callable" + + class BrokenClient: + responses = BrokenResponses() + + with pytest.raises(ValueError, match="Unable to determine Responses logprobs support"): + model._prepare_responses_api_params(BrokenClient(), params) + + # ------------------------------------------------------------------ + # _generate_responses_stream — fallback when no completed event + # ------------------------------------------------------------------ + + @pytest.mark.asyncio + async def test_streaming_fallback_when_no_completed_event(self): + """Streaming without response.completed builds final from accumulated text.""" + model = _model(use_responses_api=True) + request = _request( + [Content(parts=[Part.from_text(text="Hello")], role="user")], + streaming_tool_names={"search"}, + ) + + async def stream_events(): + yield {"type": "response.created", "response": {"id": "resp_fallback"}} + yield {"type": "response.reasoning_summary_text.delta", "delta": "Hmm"} + yield {"type": "response.output_text.delta", "delta": "Answer"} + yield { + "type": "response.output_item.added", + "item": { + "type": "function_call", + "id": "fc_fb", + "call_id": "call_fb", + "name": "search", + "arguments": "", + }, + } + yield { + "type": "response.function_call_arguments.delta", + "item_id": "fc_fb", + "delta": '{"q":"x"}', + } + # No response.completed — triggers fallback + + with patch.object(model, "_create_async_client") as client_factory: + client = AsyncMock() + client.responses.create = AsyncMock(return_value=stream_events()) + client_factory.return_value = client + responses = [item async for item in model.generate_async(request, stream=True)] + + final = responses[-1] + assert final.partial is False + assert final.response_id == "resp_fallback" + assert final.content is not None + texts = [p.text for p in final.content.parts if p.text] + assert "Hmm" in texts or "Answer" in texts + + @pytest.mark.asyncio + async def test_streaming_arguments_done_with_function_call_item(self): + """response.function_call_arguments.done with item.type=function_call uses the item.""" + model = _model(use_responses_api=True) + request = _request([Content(parts=[Part.from_text(text="Use a tool")], role="user")]) + + async def stream_events(): + yield {"type": "response.created", "response": {"id": "resp_done_item"}} + yield { + "type": "response.output_item.added", + "item": { + "type": "function_call", + "id": "fc_item", + "call_id": "call_item", + "name": "lookup", + "arguments": "", + }, + } + yield { + "type": "response.function_call_arguments.done", + "item_id": "fc_item", + "item": { + "type": "function_call", + "id": "fc_item", + "call_id": "call_item", + "name": "lookup", + "arguments": '{"result":"from_item"}', + }, + } + yield { + "type": "response.completed", + "response": { + "id": "resp_done_item", + "status": "completed", + "output": [ + { + "type": "function_call", + "id": "fc_item", + "call_id": "call_item", + "name": "lookup", + "arguments": '{"result":"from_item"}', + } + ], + }, + } + + with patch.object(model, "_create_async_client") as client_factory: + client = AsyncMock() + client.responses.create = AsyncMock(return_value=stream_events()) + client_factory.return_value = client + responses = [item async for item in model.generate_async(request, stream=True)] + + final = responses[-1] + assert final.partial is False + # Function call from item should be in final output + assert final.content is not None + + @pytest.mark.asyncio + async def test_streaming_response_incomplete_event(self): + """response.incomplete event sets completed_response.""" + model = _model(use_responses_api=True) + request = _request([Content(parts=[Part.from_text(text="Hello")], role="user")]) + + async def stream_events(): + yield {"type": "response.created", "response": {"id": "resp_inc"}} + yield { + "type": "response.incomplete", + "response": { + "id": "resp_inc", + "status": "incomplete", + "incomplete_details": {"reason": "max_output_tokens"}, + "output": [], + }, + } + + with patch.object(model, "_create_async_client") as client_factory: + client = AsyncMock() + client.responses.create = AsyncMock(return_value=stream_events()) + client_factory.return_value = client + responses = [item async for item in model.generate_async(request, stream=True)] + + final = responses[-1] + assert final.error_code == "incomplete" + assert final.error_message == "max_output_tokens" + + # ------------------------------------------------------------------ + # _convert_api_params_to_responses — max_tokens fallback + # ------------------------------------------------------------------ + + def test_max_tokens_falls_back_to_max_output_tokens(self): + model = _model(use_responses_api=True) + params = model._convert_api_params_to_responses( + { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Hi"}], + "stream": False, + "max_tokens": 256, + } + ) + assert params["max_output_tokens"] == 256 + assert "max_tokens" not in params diff --git a/tests/server/ag_ui/_core/test_agui_agent.py b/tests/server/ag_ui/_core/test_agui_agent.py index 08bb62eb..b247efd6 100644 --- a/tests/server/ag_ui/_core/test_agui_agent.py +++ b/tests/server/ag_ui/_core/test_agui_agent.py @@ -1490,6 +1490,9 @@ async def mock_run_async(**kwargs): events.append(queue.get_nowait()) assert events[-1] is None # sentinel assert len(events) >= 2 # at least one event + None + agui_agent._session_manager.update_session_state.assert_awaited_once_with( + "thread-1", "test_app", "test_user", {}, + ) async def test_handles_error_and_puts_error_event(self, agui_agent, mock_agent): queue = asyncio.Queue() @@ -1577,6 +1580,44 @@ async def mock_run_async(**kwargs): while not queue.empty(): events.append(queue.get_nowait()) assert events[-1] is None + # Non-graph AG-UI users retain the established state-sync behaviour. + agui_agent._session_manager.update_session_state.assert_awaited_once() + + async def test_preserves_graph_checkpoint_on_tool_result_submission(self, agui_agent, mock_agent): + from trpc_agent_sdk.events import Event + from trpc_agent_sdk import types + + queue = asyncio.Queue() + trpc_event = Event( + invocation_id="inv-1", author="agent", + content=types.Content(role="model", parts=[types.Part(text="done")]), + partial=False, timestamp=1000.0, + ) + agui_agent._session_manager.get_or_create_session = AsyncMock( + return_value=Mock(state={"_trpc_graph_checkpoints": {"thread-1": {}}}), + ) + agui_agent._session_manager.update_session_state = AsyncMock(return_value=True) + agui_agent._session_manager.get_session_state = AsyncMock(return_value={}) + + tc = _make_tool_call(tc_id="tc-1", name="search") + inp = _make_input(messages=[ + _make_assistant_message(tool_calls=[tc]), + _make_tool_message(content='{"ok": true}', tool_call_id="tc-1"), + ]) + with patch("trpc_agent_sdk.server.ag_ui._core._agui_agent.Runner") as MockRunner: + mock_runner = Mock() + + async def mock_run_async(**kwargs): + yield trpc_event + + mock_runner.run_async = mock_run_async + MockRunner.return_value = mock_runner + await agui_agent._run_trpc_in_background( + input=inp, agent=mock_agent, user_id="test_user", + app_name="test_app", event_queue=queue, + ) + + agui_agent._session_manager.update_session_state.assert_not_awaited() async def test_handles_lro_events(self, agui_agent, mock_agent): from trpc_agent_sdk.events import LongRunningEvent diff --git a/tests/tools/test_function_tool.py b/tests/tools/test_function_tool.py index c550ba98..88ca1efd 100644 --- a/tests/tools/test_function_tool.py +++ b/tests/tools/test_function_tool.py @@ -6,15 +6,14 @@ from __future__ import annotations -import asyncio -from typing import Any -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import MagicMock import pytest from pydantic import BaseModel from trpc_agent_sdk.context import InvocationContext from trpc_agent_sdk.tools._function_tool import FunctionTool +from trpc_agent_sdk.tools._function_tool import ToolArgumentErrorResponse # --- Test functions --- @@ -158,6 +157,7 @@ async def test_missing_mandatory_args(self, mock_context): args={}, ) assert isinstance(result, dict) + assert isinstance(result, ToolArgumentErrorResponse) assert "error" in result assert "param1" in result["error"] diff --git a/tests/tools/test_long_running_tool.py b/tests/tools/test_long_running_tool.py index a32af212..c0b9dea2 100644 --- a/tests/tools/test_long_running_tool.py +++ b/tests/tools/test_long_running_tool.py @@ -1,57 +1,16 @@ -# Tencent is pleased to support the open source community by making tRPC-Agent-Python available. -# -# Copyright (C) 2026 Tencent. All rights reserved. -# -# tRPC-Agent-Python is licensed under Apache-2.0. +from trpc_agent_sdk.events import Event +from trpc_agent_sdk.tools import is_tool_execution_error -from __future__ import annotations -import pytest +def test_framework_argument_error_is_not_a_long_running_success(): + assert is_tool_execution_error(Event(error_code="tool_argument_error")) -from trpc_agent_sdk.tools._long_running_tool import LongRunningFunctionTool +def test_business_tool_errors_are_not_classified_as_invocation_errors(): + assert not is_tool_execution_error(Event()) + assert not is_tool_execution_error(Event(content=None, custom_metadata={"result": "json parse error"})) -def long_func(query: str) -> str: - """A long running function.""" - return f"result-{query}" - -def long_func_no_doc(query: str) -> str: - return f"result-{query}" - - -long_func_no_doc.__doc__ = None - - -class TestLongRunningFunctionToolInit: - - def test_init(self): - tool = LongRunningFunctionTool(long_func) - assert tool.is_long_running is True - assert tool.name == "long_func" - - def test_init_with_invalid_filters_raises(self): - with pytest.raises(ValueError, match="not found"): - LongRunningFunctionTool(long_func, filters_name=["nonexistent"]) - - -class TestLongRunningFunctionToolGetDeclaration: - - def test_declaration_appends_note(self): - tool = LongRunningFunctionTool(long_func) - decl = tool._get_declaration() - assert decl is not None - assert "long-running operation" in decl.description - assert "A long running function." in decl.description - - def test_declaration_with_no_existing_description(self): - tool = LongRunningFunctionTool(long_func_no_doc) - decl = tool._get_declaration() - assert decl is not None - # When description is empty string '', the instruction is set with lstrip - assert "long-running operation" in (decl.description or "") - - def test_declaration_contains_do_not_call_again(self): - tool = LongRunningFunctionTool(long_func) - decl = tool._get_declaration() - assert "Do not call this tool again" in decl.description +def test_framework_execution_errors_are_classified(): + assert is_tool_execution_error(Event(error_code="tool_execution_error")) + assert is_tool_execution_error(Event(error_code="tool_not_found")) diff --git a/tests/trpc_agent_dsl/graph/test_agent_node_hitl.py b/tests/trpc_agent_dsl/graph/test_agent_node_hitl.py new file mode 100644 index 00000000..8443e1ad --- /dev/null +++ b/tests/trpc_agent_dsl/graph/test_agent_node_hitl.py @@ -0,0 +1,495 @@ +# Tencent is pleased to support the open source community by making tRPC-Agent-Python available. +# +# Copyright (C) 2026 Tencent. All rights reserved. +# +# tRPC-Agent-Python is licensed under Apache-2.0. +"""End-to-end HITL propagation tests for GraphAgent agent nodes.""" + +from typing import AsyncGenerator +from typing import List + +import pytest +from trpc_agent_sdk.agents import BaseAgent +from trpc_agent_sdk.dsl.graph import END +from trpc_agent_sdk.dsl.graph import START +from trpc_agent_sdk.dsl.graph import GraphAgent +from trpc_agent_sdk.dsl.graph import State +from trpc_agent_sdk.dsl.graph import StateGraph +from trpc_agent_sdk.events import Event +from trpc_agent_sdk.events import LongRunningEvent +from trpc_agent_sdk.models import LLMModel +from trpc_agent_sdk.models import LlmRequest +from trpc_agent_sdk.models import LlmResponse +from trpc_agent_sdk.runners import Runner +from trpc_agent_sdk.sessions import InMemorySessionService +from trpc_agent_sdk.sessions import SqlSessionService +from trpc_agent_sdk.teams import TeamAgent +from trpc_agent_sdk.teams.core import TEAM_STATE_KEY +from trpc_agent_sdk.types import Content +from trpc_agent_sdk.types import FunctionCall +from trpc_agent_sdk.types import FunctionResponse +from trpc_agent_sdk.types import Part + + +class HitlState(State, total=False): + after_child: bool + + +class StubModel(LLMModel): + @classmethod + def supported_models(cls) -> List[str]: + return [r"stub-model"] + + async def _generate_async_impl( + self, + request: LlmRequest, + stream: bool = False, + ctx=None, + ) -> AsyncGenerator[LlmResponse, None]: + del request, stream, ctx + yield LlmResponse(content=None) + + def validate_request(self, request: LlmRequest) -> None: + del request + + +class TwoRoundClarifyingAgent(BaseAgent): + """Child agent that requires two FunctionResponses before completing.""" + + received_call_ids: list[str] = [] + + async def _run_async_impl(self, ctx) -> AsyncGenerator[Event, None]: + response = self._function_response(ctx.user_content) + if response is not None: + self.received_call_ids.append(str(response.id)) + if response.response.get("answer") == "done": + yield Event( + invocation_id=ctx.invocation_id, + author=self.name, + branch=ctx.branch, + content=Content(role="model", parts=[Part.from_text(text="clarification-complete")]), + ) + return + call_id = "child-question-2" + else: + call_id = "child-question-1" + + call = FunctionCall( + id=call_id, + name="ask_clarification", + args={"round": 1 if call_id.endswith("1") else 2}, + ) + yield LongRunningEvent( + invocation_id=ctx.invocation_id, + author=self.name, + branch=ctx.branch, + function_call=call, + function_response=FunctionResponse( + id=call.id, + name=call.name, + response={"status": "pending"}, + ), + ) + + @staticmethod + def _function_response(content: Content | None) -> FunctionResponse | None: + if content is None: + return None + for part in content.parts or []: + if part.function_response is not None: + return part.function_response + return None + + +class TeamHitlLeader(BaseAgent): + """Leader double that verifies TeamAgent receives the original response ID.""" + + received_call_ids: list[str] = [] + + async def _run_async_impl(self, ctx) -> AsyncGenerator[Event, None]: + response = TwoRoundClarifyingAgent._function_response(ctx.user_content) + if response is not None: + self.received_call_ids.append(str(response.id)) + yield Event( + invocation_id=ctx.invocation_id, + author=self.name, + branch=ctx.branch, + content=Content(role="model", parts=[Part.from_text(text="team-clarification-complete")]), + ) + return + + call = FunctionCall( + id="team-question-1", + name="ask_clarification", + args={"question": "approve team plan?"}, + ) + yield LongRunningEvent( + invocation_id=ctx.invocation_id, + author=self.name, + branch=ctx.branch, + function_call=call, + function_response=FunctionResponse( + id=call.id, + name=call.name, + response={"status": "pending"}, + ), + ) + + +def _user_text(text: str) -> Content: + return Content(role="user", parts=[Part.from_text(text=text)]) + + +def _tool_response(event: LongRunningEvent, answer: str) -> Content: + return Content( + role="user", + parts=[ + Part( + function_response=FunctionResponse( + id=event.function_call.id, + name=event.function_call.name, + response={"answer": answer}, + ) + ) + ], + ) + + +async def _run(runner: Runner, message: Content) -> list[Event]: + return [ + event + async for event in runner.run_async( + user_id="user-1", + session_id="session-1", + new_message=message, + ) + ] + + +@pytest.mark.asyncio +async def test_agent_node_long_running_interrupts_parent_and_resumes_multiple_rounds(): + child = TwoRoundClarifyingAgent(name="clarifier") + + async def after_child(state: HitlState) -> dict[str, bool]: + del state + return {"after_child": True} + + graph = StateGraph(HitlState) + graph.add_agent_node("clarify", child, isolated_messages=True) + graph.add_node("after_child", after_child) + graph.add_edge(START, "clarify") + graph.add_edge("clarify", "after_child") + graph.add_edge("after_child", END) + + service = InMemorySessionService() + runner = Runner( + app_name="agent-node-hitl-test", + agent=GraphAgent(name="workflow", graph=graph.compile()), + session_service=service, + close_session_service_on_close=False, + ) + + first_events = await _run(runner, _user_text("start")) + first_pending = next(event for event in first_events if isinstance(event, LongRunningEvent)) + assert first_pending.function_call.name == "ask_clarification" + assert first_pending.function_call.args == {"round": 1, "status": "pending"} + first_session = await service.get_session( + app_name="agent-node-hitl-test", + user_id="user-1", + session_id="session-1", + ) + assert first_session is not None + assert first_session.state.get("after_child") is not True + + second_events = await _run(runner, _tool_response(first_pending, "again")) + second_pending = next(event for event in second_events if isinstance(event, LongRunningEvent)) + assert second_pending.function_call.name == "ask_clarification" + second_session = await service.get_session( + app_name="agent-node-hitl-test", + user_id="user-1", + session_id="session-1", + ) + assert second_session is not None + assert second_session.state.get("after_child") is not True + + await _run(runner, _tool_response(second_pending, "done")) + completed_session = await service.get_session( + app_name="agent-node-hitl-test", + user_id="user-1", + session_id="session-1", + ) + assert completed_session is not None + assert completed_session.state.get("after_child") is True + assert child.received_call_ids == ["child-question-1", "child-question-2"] + + await runner.close() + await service.close() + + +@pytest.mark.asyncio +async def test_agent_node_hitl_survives_service_restart(tmp_path): + """AgentNode HITL state must persist across process restart via SqlSessionService. + + This verifies that STATE_KEY_PENDING_AGENT_NODE_HITL is written to + session.state through the interrupt bridge event's state_delta, so that + a fresh Runner reading from the same SqlSessionService can resume the + pending HITL round. + """ + from trpc_agent_sdk.dsl.graph._constants import STATE_KEY_PENDING_AGENT_NODE_HITL + + def build_runner(service: SqlSessionService) -> tuple[Runner, TwoRoundClarifyingAgent]: + child = TwoRoundClarifyingAgent(name="clarifier") + graph = StateGraph(HitlState) + graph.add_agent_node("clarify", child, isolated_messages=True) + graph.add_edge(START, "clarify") + graph.add_edge("clarify", END) + return ( + Runner( + app_name="agent-node-restart-test", + agent=GraphAgent(name="workflow", graph=graph.compile()), + session_service=service, + close_session_service_on_close=False, + ), + child, + ) + + db_url = f"sqlite:///{tmp_path / 'agent-node-restart.sqlite'}" + service = SqlSessionService(db_url, is_async=False) + runner, child = build_runner(service) + + first_events = await _run(runner, _user_text("start")) + first_pending = next(event for event in first_events if isinstance(event, LongRunningEvent)) + assert first_pending.function_call.name == "ask_clarification" + + first_session = await service.get_session( + app_name="agent-node-restart-test", + user_id="user-1", + session_id="session-1", + ) + assert first_session is not None + # The pending HITL state must be persisted in session.state so that a + # fresh Runner can pick it up. + assert STATE_KEY_PENDING_AGENT_NODE_HITL in first_session.state + pending = first_session.state[STATE_KEY_PENDING_AGENT_NODE_HITL] + assert pending["node_id"] == "clarify" + assert pending["current"]["function_call"]["id"] == "child-question-1" + + # Simulate process restart: close runner + service, re-open from same DB. + await runner.close() + await service.close() + + service = SqlSessionService(db_url, is_async=False) + runner, resumed_child = build_runner(service) + + # Resume with the first round's response. + second_events = await _run(runner, _tool_response(first_pending, "again")) + second_pending = next(event for event in second_events if isinstance(event, LongRunningEvent)) + + second_session = await service.get_session( + app_name="agent-node-restart-test", + user_id="user-1", + session_id="session-1", + ) + assert second_session is not None + assert STATE_KEY_PENDING_AGENT_NODE_HITL in second_session.state + pending2 = second_session.state[STATE_KEY_PENDING_AGENT_NODE_HITL] + assert len(pending2["completed"]) == 1 + assert pending2["current"]["function_call"]["id"] == "child-question-2" + + # Complete the second round. + await _run(runner, _tool_response(second_pending, "done")) + completed_session = await service.get_session( + app_name="agent-node-restart-test", + user_id="user-1", + session_id="session-1", + ) + assert completed_session is not None + # After completion, the pending HITL state should be cleared. + hitl_value = completed_session.state.get(STATE_KEY_PENDING_AGENT_NODE_HITL) + assert hitl_value is None + assert resumed_child.received_call_ids == ["child-question-1", "child-question-2"] + + await runner.close() + await service.close() + + +@pytest.mark.asyncio +async def test_agent_node_hitl_state_key_is_unsafe(): + """STATE_KEY_PENDING_AGENT_NODE_HITL must be in UNSAFE_STATE_KEYS so it + is filtered out of the final_state/state_delta exposed via completion + events (it may contain sensitive tool arguments and child state). + """ + from trpc_agent_sdk.dsl.graph._constants import STATE_KEY_PENDING_AGENT_NODE_HITL + from trpc_agent_sdk.dsl.graph._constants import is_unsafe_state_key + + assert is_unsafe_state_key(STATE_KEY_PENDING_AGENT_NODE_HITL) is True + + +@pytest.mark.asyncio +async def test_team_agent_leader_hitl_survives_service_restart(tmp_path): + async def after_team(state: HitlState) -> dict[str, bool]: + del state + return {"after_child": True} + + def build_runner(service: SqlSessionService) -> tuple[Runner, TeamHitlLeader]: + member = TwoRoundClarifyingAgent(name="unused_member") + team = TeamAgent( + name="development_team", + model=StubModel(model_name="stub-model"), + members=[member], + ) + leader = TeamHitlLeader(name="development_team") + team.__pydantic_private__["_leader_agent"] = leader + graph = StateGraph(HitlState) + graph.add_agent_node("development", team, isolated_messages=True) + graph.add_node("after_team", after_team) + graph.add_edge(START, "development") + graph.add_edge("development", "after_team") + graph.add_edge("after_team", END) + return ( + Runner( + app_name="agent-node-hitl-test", + agent=GraphAgent(name="team_workflow", graph=graph.compile()), + session_service=service, + close_session_service_on_close=False, + ), + leader, + ) + + db_url = f"sqlite:///{tmp_path / 'agent-node-hitl.sqlite'}" + service = SqlSessionService(db_url, is_async=False) + runner, _ = build_runner(service) + + first_events = await _run(runner, _user_text("start team")) + pending = next(event for event in first_events if isinstance(event, LongRunningEvent)) + first_session = await service.get_session( + app_name="agent-node-hitl-test", + user_id="user-1", + session_id="session-1", + ) + assert first_session is not None + assert first_session.state.get("after_child") is not True + assert TEAM_STATE_KEY in first_session.state + + await runner.close() + await service.close() + + service = SqlSessionService(db_url, is_async=False) + runner, resumed_leader = build_runner(service) + await _run(runner, _tool_response(pending, "done")) + completed_session = await service.get_session( + app_name="agent-node-hitl-test", + user_id="user-1", + session_id="session-1", + ) + assert completed_session is not None + assert completed_session.state.get("after_child") is True + assert resumed_leader.received_call_ids == ["team-question-1"] + + await runner.close() + await service.close() + + +@pytest.mark.asyncio +async def test_agent_node_multiround_hitl_resume_with_correct_order(): + """Multi-round HITL resume must replay completed rounds in order. + + After the first round interrupts, the client resumes with a FunctionResponse. + The second round interrupts again. The client must then resume with the + second round's response. This test verifies that submitting a stale + (already-completed round's) response does not silently complete the graph, + and that subsequently submitting the correct round's response succeeds. + """ + child = TwoRoundClarifyingAgent(name="clarifier") + + async def after_child(state: HitlState) -> dict[str, bool]: + del state + return {"after_child": True} + + graph = StateGraph(HitlState) + graph.add_agent_node("clarify", child, isolated_messages=True) + graph.add_node("after_child", after_child) + graph.add_edge(START, "clarify") + graph.add_edge("clarify", "after_child") + graph.add_edge("after_child", END) + + service = InMemorySessionService() + runner = Runner( + app_name="agent-node-hitl-order-test", + agent=GraphAgent(name="workflow", graph=graph.compile()), + session_service=service, + close_session_service_on_close=False, + ) + + # Round 1: interrupt with child-question-1 (wrapped in a synthesized + # graph-level interrupt bridge function_call.id). + first_events = await _run(runner, _user_text("start")) + first_pending = next(event for event in first_events if isinstance(event, LongRunningEvent)) + assert first_pending.function_call.name == "ask_clarification" + assert first_pending.function_call.args == {"round": 1, "status": "pending"} + + # Resume round 1 with "again" → round 2 interrupts with child-question-2. + second_events = await _run(runner, _tool_response(first_pending, "again")) + second_pending = next(event for event in second_events if isinstance(event, LongRunningEvent)) + assert second_pending.function_call.name == "ask_clarification" + assert second_pending.function_call.args == {"round": 2, "status": "pending"} + + # Now try to resume round 2 using the FIRST round's function_call.id. + # This is a "stale" resume — the client submitted a response for an + # already-completed round. The graph's _extract_resume_command builds a + # Command with the stale function_response.id; LangGraph will not find + # a matching interrupt and the resume value will be ignored or cause an + # error. We assert that the graph does NOT silently complete with + # incorrect state. + stale_response = Content( + role="user", + parts=[ + Part( + function_response=FunctionResponse( + id=first_pending.function_call.id, + name=first_pending.function_call.name, + response={"answer": "stale"}, + ) + ) + ], + ) + stale_events = await _run(runner, stale_response) + # The graph should not have completed successfully with a stale resume. + completed_session = await service.get_session( + app_name="agent-node-hitl-order-test", + user_id="user-1", + session_id="session-1", + ) + assert completed_session is not None + # after_child should NOT be True because we didn't properly complete round 2. + assert completed_session.state.get("after_child") is not True + + # After the stale resume the graph may have re-interrupted (the stale + # value was consumed by the current round's interrupt, causing the child + # to ask another question). We verify the key invariant: the graph did + # NOT silently complete. A subsequent proper resume should still be able + # to drive the graph to completion. + # + # Find the latest pending LongRunningEvent from the stale run and resume + # it with "done" to complete the flow. + stale_pending = next( + (event for event in reversed(stale_events) if isinstance(event, LongRunningEvent)), + None, + ) + if stale_pending is not None: + await _run(runner, _tool_response(stale_pending, "done")) + else: + # If the stale resume did not produce a new interrupt, try the + # original second_pending. + await _run(runner, _tool_response(second_pending, "done")) + + completed_session = await service.get_session( + app_name="agent-node-hitl-order-test", + user_id="user-1", + session_id="session-1", + ) + assert completed_session is not None + assert completed_session.state.get("after_child") is True + + await runner.close() + await service.close() diff --git a/tests/trpc_agent_dsl/graph/test_constants.py b/tests/trpc_agent_dsl/graph/test_constants.py index 5dc561a2..44b91c50 100644 --- a/tests/trpc_agent_dsl/graph/test_constants.py +++ b/tests/trpc_agent_dsl/graph/test_constants.py @@ -39,6 +39,7 @@ STATE_KEY_NODE_RESPONSES, STATE_KEY_ONE_SHOT_MESSAGES, STATE_KEY_ONE_SHOT_MESSAGES_BY_NODE, + STATE_KEY_PENDING_AGENT_NODE_HITL, STATE_KEY_PENDING_INTERRUPT, STATE_KEY_PENDING_INTERRUPT_AUTHOR, STATE_KEY_PENDING_INTERRUPT_BRANCH, @@ -114,6 +115,7 @@ def test_interrupt_keys(self): assert STATE_KEY_PENDING_INTERRUPT_ID == "_trpc_graph_pending_interrupt_id" assert STATE_KEY_PENDING_INTERRUPT_AUTHOR == "_trpc_graph_pending_interrupt_author" assert STATE_KEY_PENDING_INTERRUPT_BRANCH == "_trpc_graph_pending_interrupt_branch" + assert STATE_KEY_PENDING_AGENT_NODE_HITL == "_trpc_graph_pending_agent_node_hitl" assert STATE_KEY_LONG_RUNNING_PREFIX == "__trpc_graph_long_running__" def test_role_values(self): @@ -145,6 +147,7 @@ def test_unsafe_keys_contains_expected_members(self): STATE_KEY_PENDING_INTERRUPT_ID, STATE_KEY_PENDING_INTERRUPT_AUTHOR, STATE_KEY_PENDING_INTERRUPT_BRANCH, + STATE_KEY_PENDING_AGENT_NODE_HITL, } assert UNSAFE_STATE_KEYS == expected diff --git a/tests/trpc_agent_dsl/graph/test_events.py b/tests/trpc_agent_dsl/graph/test_events.py index fe3cbc8e..0c8ea666 100644 --- a/tests/trpc_agent_dsl/graph/test_events.py +++ b/tests/trpc_agent_dsl/graph/test_events.py @@ -54,6 +54,26 @@ def test_node_start_contains_structured_metadata_and_truncates_model_input(self) assert event.object == GraphEventType.GRAPH_NODE_START assert event.partial is True + def test_node_events_prefer_human_readable_description_in_text(self): + """User-visible node events should not leak internal node IDs.""" + start = self.builder.node_start( + node_id="tech-clarify", + node_description="技术澄清", + ) + complete = self.builder.node_complete( + node_id="tech-clarify", + node_description="技术澄清", + ) + error = self.builder.node_error( + node_id="tech-clarify", + node_description="技术澄清", + error="failed", + ) + + assert start.get_text() == "Starting node: 技术澄清" + assert complete.get_text().startswith("Completed node: 技术澄清") + assert error.get_text() == "Error in node 技术澄清: failed" + def test_model_complete_uses_error_phase_and_truncates_payloads(self): """Model completion should switch to error phase when error is provided.""" start_time = datetime.now() - timedelta(milliseconds=20) diff --git a/trpc_agent_sdk/agents/_llm_agent.py b/trpc_agent_sdk/agents/_llm_agent.py index 3546c1d6..e4e59472 100644 --- a/trpc_agent_sdk/agents/_llm_agent.py +++ b/trpc_agent_sdk/agents/_llm_agent.py @@ -12,52 +12,43 @@ from __future__ import annotations -from typing import Any -from typing import AsyncGenerator -from typing import Awaitable -from typing import Callable -from typing import List -from typing import Literal -from typing import Optional -from typing import TypeAlias -from typing import Union -from typing_extensions import override +from typing import Any, AsyncGenerator, Awaitable, Callable, List, Literal, Optional, TypeAlias, Union -from pydantic import BaseModel -from pydantic import Field -from pydantic import field_validator +from pydantic import BaseModel, Field, field_validator +from typing_extensions import override from trpc_agent_sdk.code_executors import BaseCodeExecutor from trpc_agent_sdk.context import InvocationContext from trpc_agent_sdk.events import Event from trpc_agent_sdk.log import logger -from trpc_agent_sdk.models import LLMModel -from trpc_agent_sdk.models import LlmRequest -from trpc_agent_sdk.models import ModelRegistry +from trpc_agent_sdk.models import LLMModel, LlmRequest, ModelRegistry from trpc_agent_sdk.planners import BasePlanner from trpc_agent_sdk.skills import BaseSkillRepository -from trpc_agent_sdk.tools import BaseTool -from trpc_agent_sdk.tools import BaseToolSet -from trpc_agent_sdk.tools import FunctionTool -from trpc_agent_sdk.tools import LongRunningFunctionTool -from trpc_agent_sdk.tools import transfer_to_agent -from trpc_agent_sdk.types import Content -from trpc_agent_sdk.types import GenerateContentConfig +from trpc_agent_sdk.tools import ( + BaseTool, + BaseToolSet, + FunctionTool, + LongRunningFunctionTool, + is_tool_execution_error, + transfer_to_agent, +) +from trpc_agent_sdk.types import Content, GenerateContentConfig from ..exceptions import RunCancelledException from ._base_agent import BaseAgent -from ._callback import ModelCallback -from ._callback import ToolCallback +from ._callback import ModelCallback, ToolCallback from ._constants import TRPC_AGENT_RUNNING_KEY -from .core import BranchFilterMode -from .core import CodeExecutionRequestProcessor -from .core import CodeExecutionResponseProcessor -from .core import LlmProcessor -from .core import TimelineFilterMode -from .core import ToolsProcessor -from .core import create_final_model_response_event -from .core import default_request_processor -from .core import get_structured_model_response +from .core import ( + BranchFilterMode, + CodeExecutionRequestProcessor, + CodeExecutionResponseProcessor, + LlmProcessor, + TimelineFilterMode, + ToolsProcessor, + create_final_model_response_event, + default_request_processor, + get_structured_model_response, +) # Type aliases for instruction providers InstructionProvider: TypeAlias = Callable[[InvocationContext], Union[str, Awaitable[str]]] @@ -118,7 +109,7 @@ class LlmAgent(BaseAgent): settings, etc. """ - include_contents: Literal['default', 'none'] = 'default' + include_contents: Literal["default", "none"] = "default" """Controls content inclusion in model requests. Options: @@ -368,8 +359,12 @@ def _should_enable_agent_transfer(self) -> bool: return True # Check if transfer to peers is possible - if (not self.disallow_transfer_to_peers and self.parent_agent and self._is_llm_agent(self.parent_agent) - and len(self.parent_agent.sub_agents) > 1): # Has siblings + if ( + not self.disallow_transfer_to_peers + and self.parent_agent + and self._is_llm_agent(self.parent_agent) + and len(self.parent_agent.sub_agents) > 1 + ): # Has siblings return True return False @@ -471,7 +466,9 @@ def accumulate_content(event: Event) -> None: await ctx.raise_if_cancelled() # Step 1: Build request using the request processor (includes conversation history) - request = LlmRequest(model=model_instance.name, ) + request = LlmRequest( + model=model_instance.name, + ) error_event = await default_request_processor.build_request( request, @@ -555,7 +552,6 @@ def accumulate_content(event: Event) -> None: # Step 3: Execute tools if any were collected if collected_tool_calls: logger.debug("Executing %s tool calls", len(collected_tool_calls)) - logger.debug("Executing %s tool calls", len(collected_tool_calls)) try: # Use extended tools processor that includes transfer tool if needed @@ -568,6 +564,14 @@ def accumulate_content(event: Event) -> None: if tool and isinstance(tool, LongRunningFunctionTool): long_running_tool_ids.add(tool_call.id) + if long_running_tool_ids and getattr(self, "parallel_tool_calls", False): + logger.warning( + "Long-running tools are mixed with parallel_tool_calls=True. " + "When a long-running tool suspends execution, other tools in the " + "same batch will be persisted but not summarized by the LLM until " + "the long-running tool resumes." + ) + # Execute tools and yield results (Runner will store them automatically) last_tool_event = None any_skip_summarization = False @@ -579,8 +583,7 @@ def accumulate_content(event: Event) -> None: # Check if this event contains responses from long-running tools if tool_event.content and tool_event.content.parts: for part in tool_event.content.parts: - if (part.function_response and part.function_response.id in long_running_tool_ids): - # This is a response from a long-running tool + if part.function_response and part.function_response.id in long_running_tool_ids: # Find the corresponding function call corresponding_call = None for call in collected_tool_calls: @@ -589,6 +592,16 @@ def accumulate_content(event: Event) -> None: break if corresponding_call: + if is_tool_execution_error(tool_event): + logger.warning( + "Long-running tool %s returned an invocation error; " + "continuing as a normal tool response", + corresponding_call.name, + ) + # The model can see the error and decide whether to + # correct and retry. Never suspend the graph here. + continue + # This is a response from a long-running tool # Import LongRunningEvent here to avoid circular imports from trpc_agent_sdk.events import LongRunningEvent @@ -607,8 +620,10 @@ def accumulate_content(event: Event) -> None: # Then yield the long-running event yield long_running_event - logger.debug("Long-running tool %s completed, yielding LongRunningEvent", - corresponding_call.name) + logger.debug( + "Long-running tool %s completed, yielding LongRunningEvent", + corresponding_call.name, + ) return # End agent execution after long-running event # Yield regular tool events @@ -701,33 +716,34 @@ def _save_output_to_state(self, ctx: InvocationContext, event: Event) -> None: """ if self.output_key: # Save output to state using the delta tracking system - result = ''.join([part.text if not part.thought else '' for part in event.content.parts]) + result = "".join([part.text if not part.thought else "" for part in event.content.parts]) ctx.state[self.output_key] = result event.actions.state_delta[self.output_key] = result logger.debug("Saved agent output to state key '%s': %s...", self.output_key, result[:100]) - @field_validator('generate_content_config', mode='after') + @field_validator("generate_content_config", mode="after") @classmethod def __validate_generate_content_config( - cls, generate_content_config: Optional[GenerateContentConfig]) -> GenerateContentConfig: + cls, generate_content_config: Optional[GenerateContentConfig] + ) -> GenerateContentConfig: if not generate_content_config: return GenerateContentConfig() if generate_content_config.thinking_config: - raise ValueError('Thinking config should be set via LlmAgent.planner.') + raise ValueError("Thinking config should be set via LlmAgent.planner.") if generate_content_config.tools: - raise ValueError('All tools must be set via LlmAgent.tools.') + raise ValueError("All tools must be set via LlmAgent.tools.") if generate_content_config.system_instruction: - raise ValueError('System instruction must be set via LlmAgent.instruction.') + raise ValueError("System instruction must be set via LlmAgent.instruction.") if generate_content_config.response_schema: - raise ValueError('Response schema must be set via LlmAgent.output_schema.') + raise ValueError("Response schema must be set via LlmAgent.output_schema.") return generate_content_config - @field_validator('code_executor', mode='after') + @field_validator("code_executor", mode="after") @classmethod def __validate_code_executor(cls, code_executor: Optional[BaseCodeExecutor]) -> Optional[BaseCodeExecutor]: """Validate code executor configuration.""" if code_executor and not isinstance(code_executor, BaseCodeExecutor): - raise ValueError('Code executor must be an instance of BaseCodeExecutor.') + raise ValueError("Code executor must be an instance of BaseCodeExecutor.") return code_executor diff --git a/trpc_agent_sdk/agents/core/_tools_processor.py b/trpc_agent_sdk/agents/core/_tools_processor.py index e6d1a2e3..041a4540 100644 --- a/trpc_agent_sdk/agents/core/_tools_processor.py +++ b/trpc_agent_sdk/agents/core/_tools_processor.py @@ -35,6 +35,10 @@ from trpc_agent_sdk.tools import BaseTool from trpc_agent_sdk.tools import BaseToolSet from trpc_agent_sdk.tools import convert_toolunion_to_tool_list +from trpc_agent_sdk.tools._function_tool import ToolArgumentErrorResponse +from trpc_agent_sdk.tools._long_running_tool import TOOL_ERROR_CODE_ARGUMENT_ERROR +from trpc_agent_sdk.tools._long_running_tool import TOOL_ERROR_CODE_EXECUTION_ERROR +from trpc_agent_sdk.tools._long_running_tool import TOOL_ERROR_CODE_NOT_FOUND from trpc_agent_sdk.types import Content from trpc_agent_sdk.types import FunctionCall from trpc_agent_sdk.types import Part @@ -141,7 +145,7 @@ async def __invoke_tools( logger.warning("No tool found for tool call: %s", tool_call.name) result_event = self._create_error_event( context, - "tool_not_found", + TOOL_ERROR_CODE_NOT_FOUND, f"Tool '{tool_call.name}' not found", tool_call.id, tool_call.name, @@ -153,7 +157,7 @@ async def __invoke_tools( logger.error("Error executing tool %s: %s", tool_call.name, ex, exc_info=True) result_event = self._create_error_event( context, - "tool_execution_error", + TOOL_ERROR_CODE_EXECUTION_ERROR, str(ex), tool_call.id, tool_call.name, @@ -256,7 +260,7 @@ async def execute_tools_async( if tool is None: yield self._create_error_event( context, - "tool_not_found", + TOOL_ERROR_CODE_NOT_FOUND, f"Tool '{tool_call.name}' not found", tool_call.id, tool_call.name, @@ -414,6 +418,10 @@ async def _execute_tool(self, tool_call: FunctionCall, tool: BaseTool, context: invocation_id=context.invocation_id, author=context.agent.name, content=content, + error_code=( + TOOL_ERROR_CODE_ARGUMENT_ERROR if isinstance(result, ToolArgumentErrorResponse) else None + ), + error_message=result.get("error") if isinstance(result, ToolArgumentErrorResponse) else None, custom_metadata={"execution_time": execution_time} if execution_time is not None else {}, branch=context.branch, ) @@ -454,7 +462,7 @@ async def _execute_tool(self, tool_call: FunctionCall, tool: BaseTool, context: error_type=type(ex).__name__, ) - error_event = self._create_error_event(context, "tool_execution_error", str(ex), tool_call.id, + error_event = self._create_error_event(context, TOOL_ERROR_CODE_EXECUTION_ERROR, str(ex), tool_call.id, tool_call.name) # Compute state after failed tool execution @@ -605,7 +613,7 @@ async def _execute_progress_streaming_tool( ) error_event = self._create_error_event( context, - "tool_execution_error", + TOOL_ERROR_CODE_EXECUTION_ERROR, str(ex), tool_call.id, tool_call.name, diff --git a/trpc_agent_sdk/dsl/graph/_constants.py b/trpc_agent_sdk/dsl/graph/_constants.py index c2d04072..ca70c358 100644 --- a/trpc_agent_sdk/dsl/graph/_constants.py +++ b/trpc_agent_sdk/dsl/graph/_constants.py @@ -148,6 +148,11 @@ STATE_KEY_PENDING_INTERRUPT_AUTHOR = "_trpc_graph_pending_interrupt_author" STATE_KEY_PENDING_INTERRUPT_BRANCH = "_trpc_graph_pending_interrupt_branch" +# Pending child-agent HITL context owned by AgentNodeAction. This bridges a +# child LongRunningEvent into a parent GraphAgent interrupt and supports +# replaying multiple clarification rounds in the same graph node. +STATE_KEY_PENDING_AGENT_NODE_HITL = "_trpc_graph_pending_agent_node_hitl" + # Prefix for synthetic function call IDs used by graph interrupt bridge STATE_KEY_LONG_RUNNING_PREFIX = "__trpc_graph_long_running__" @@ -164,22 +169,25 @@ # Unsafe State Keys (not serializable or should not be exposed) # ============================================================================= -UNSAFE_STATE_KEYS = frozenset({ - STATE_KEY_SESSION, - STATE_KEY_EXEC_CONTEXT, - STATE_KEY_CURRENT_NODE_ID, - STATE_KEY_TOOL_CALLBACKS, - STATE_KEY_MODEL_CALLBACKS, - STATE_KEY_AGENT_CALLBACKS, - STATE_KEY_NODE_CALLBACKS, - STATE_KEY_CHECKPOINTS, - STATE_KEY_CHECKPOINT_WRITES, - STATE_KEY_CHECKPOINT_BLOBS, - STATE_KEY_PENDING_INTERRUPT, - STATE_KEY_PENDING_INTERRUPT_ID, - STATE_KEY_PENDING_INTERRUPT_AUTHOR, - STATE_KEY_PENDING_INTERRUPT_BRANCH, -}) +UNSAFE_STATE_KEYS = frozenset( + { + STATE_KEY_SESSION, + STATE_KEY_EXEC_CONTEXT, + STATE_KEY_CURRENT_NODE_ID, + STATE_KEY_TOOL_CALLBACKS, + STATE_KEY_MODEL_CALLBACKS, + STATE_KEY_AGENT_CALLBACKS, + STATE_KEY_NODE_CALLBACKS, + STATE_KEY_CHECKPOINTS, + STATE_KEY_CHECKPOINT_WRITES, + STATE_KEY_CHECKPOINT_BLOBS, + STATE_KEY_PENDING_INTERRUPT, + STATE_KEY_PENDING_INTERRUPT_ID, + STATE_KEY_PENDING_INTERRUPT_AUTHOR, + STATE_KEY_PENDING_INTERRUPT_BRANCH, + STATE_KEY_PENDING_AGENT_NODE_HITL, + } +) def is_unsafe_state_key(key: str) -> bool: diff --git a/trpc_agent_sdk/dsl/graph/_events/_builder.py b/trpc_agent_sdk/dsl/graph/_events/_builder.py index 9ecd1955..83f45934 100644 --- a/trpc_agent_sdk/dsl/graph/_events/_builder.py +++ b/trpc_agent_sdk/dsl/graph/_events/_builder.py @@ -185,9 +185,10 @@ def node_start( # Store metadata in state_delta state_delta: dict[str, Any] = {} _store_metadata(state_delta, METADATA_KEY_NODE, metadata) + display_name = node_description.strip() if node_description and node_description.strip() else node_id return self._build_event( - text=f"Starting node: {node_id}", + text=f"Starting node: {display_name}", state_delta=state_delta, partial=True, author_override=self._author or node_id, @@ -227,9 +228,10 @@ def node_complete( # Store metadata in state_delta state_delta: dict[str, Any] = {} _store_metadata(state_delta, METADATA_KEY_NODE, metadata) + display_name = node_description.strip() if node_description and node_description.strip() else node_id return self._build_event( - text=f"Completed node: {node_id} ({duration_ms:.1f}ms)", + text=f"Completed node: {display_name} ({duration_ms:.1f}ms)", state_delta=state_delta, author_override=self._author or node_id, object_type=GraphEventType.GRAPH_NODE_COMPLETE, @@ -266,7 +268,8 @@ def node_error( _store_metadata(state_delta, METADATA_KEY_NODE, metadata) # Build error message - text = f"Error in node {node_id}: {error}" + display_name = node_description.strip() if node_description and node_description.strip() else node_id + text = f"Error in node {display_name}: {error}" return self._build_event( text=text, diff --git a/trpc_agent_sdk/dsl/graph/_graph_agent.py b/trpc_agent_sdk/dsl/graph/_graph_agent.py index 12e05c1d..3d484720 100644 --- a/trpc_agent_sdk/dsl/graph/_graph_agent.py +++ b/trpc_agent_sdk/dsl/graph/_graph_agent.py @@ -113,8 +113,10 @@ class GraphAgent(BaseAgent): @classmethod def _reject_sub_agents(cls, data: Any) -> Any: if isinstance(data, dict) and "sub_agents" in data: - raise ValueError("GraphAgent does not accept `sub_agents`; compose nested agents with " - "`StateGraph.add_agent_node(...)`.") + raise ValueError( + "GraphAgent does not accept `sub_agents`; compose nested agents with " + "`StateGraph.add_agent_node(...)`." + ) return data @override @@ -178,9 +180,9 @@ async def _run_async_impl(self, ctx: InvocationContext) -> AsyncGenerator[Event, # "updates" is required for state to propagate between nodes # "custom" enables EventEmitter streaming async for mode, chunk in self.graph.astream( - graph_input, - runnable_config, - stream_mode=["updates", "custom"], + graph_input, + runnable_config, + stream_mode=["updates", "custom"], ): # Cancellation checkpoint at each iteration await ctx.raise_if_cancelled() @@ -193,7 +195,8 @@ async def _run_async_impl(self, ctx: InvocationContext) -> AsyncGenerator[Event, interrupted = True for interrupt in interrupts: function_call_event, function_response_event, long_running_event = ( - self._create_interrupt_events(ctx, interrupt)) + self._create_interrupt_events(ctx, interrupt) + ) yield function_call_event await ctx.raise_if_cancelled() yield function_response_event @@ -216,7 +219,8 @@ async def _run_async_impl(self, ctx: InvocationContext) -> AsyncGenerator[Event, if isinstance(key, str): final_state[key] = value logger.debug( - f"[{self.name}] Node '{node_name}' updated keys: {list(node_output.keys())}") + f"[{self.name}] Node '{node_name}' updated keys: {list(node_output.keys())}" + ) # Emit state update event if updated_keys: @@ -293,7 +297,7 @@ def _extract_resume_command(self, ctx: InvocationContext) -> Optional[Command]: if not function_response_id.startswith(STATE_KEY_LONG_RUNNING_PREFIX): return None - interrupt_id = function_response_id[len(STATE_KEY_LONG_RUNNING_PREFIX):] + interrupt_id = function_response_id.removeprefix(STATE_KEY_LONG_RUNNING_PREFIX) if not interrupt_id: return None @@ -390,7 +394,11 @@ def _build_interrupt_function(self, interrupt: Interrupt) -> tuple[str, str, dic function_call_id = f"{STATE_KEY_LONG_RUNNING_PREFIX}{interrupt_id}" raw_args = interrupt.value - if isinstance(raw_args, dict): + if isinstance(raw_args, dict) and raw_args.get("_trpc_agent_node_hitl") is True: + function_name = str(raw_args.get("toolName") or function_name) + visible_args = raw_args.get("arguments") + function_args = visible_args if isinstance(visible_args, dict) else {} + elif isinstance(raw_args, dict): function_args = {str(key): value for key, value in raw_args.items()} else: function_args = {"value": raw_args} @@ -447,15 +455,17 @@ def _build_initial_state(self, ctx: InvocationContext) -> GraphState: else: initial_state[STATE_KEY_MESSAGES] = messages - initial_state.update({ - STATE_KEY_USER_INPUT: user_input, - STATE_KEY_METADATA: { - METADATA_KEY_INVOCATION_ID: ctx.invocation_id, - METADATA_KEY_BRANCH: ctx.branch or self.name, - METADATA_KEY_SESSION_ID: ctx.session.id, - METADATA_KEY_AGENT_NAME: self.name, - }, - }) + initial_state.update( + { + STATE_KEY_USER_INPUT: user_input, + STATE_KEY_METADATA: { + METADATA_KEY_INVOCATION_ID: ctx.invocation_id, + METADATA_KEY_BRANCH: ctx.branch or self.name, + METADATA_KEY_SESSION_ID: ctx.session.id, + METADATA_KEY_AGENT_NAME: self.name, + }, + } + ) # Built-ins should always exist for node logic consistency. initial_state.setdefault(STATE_KEY_LAST_RESPONSE, None) diff --git a/trpc_agent_sdk/dsl/graph/_node_action/_agent.py b/trpc_agent_sdk/dsl/graph/_node_action/_agent.py index 85d064dc..5bc7ed9b 100644 --- a/trpc_agent_sdk/dsl/graph/_node_action/_agent.py +++ b/trpc_agent_sdk/dsl/graph/_node_action/_agent.py @@ -10,12 +10,15 @@ from typing import Callable from typing import Optional +from langgraph.errors import GraphInterrupt from trpc_agent_sdk.agents import BaseAgent from trpc_agent_sdk.agents import LlmAgent from trpc_agent_sdk.context import InvocationContext from trpc_agent_sdk.events import Event +from trpc_agent_sdk.events import LongRunningEvent from trpc_agent_sdk.types import Content from trpc_agent_sdk.types import EventActions +from trpc_agent_sdk.types import FunctionResponse from trpc_agent_sdk.types import Part from .._callbacks import NodeCallbackContext @@ -23,9 +26,11 @@ from .._constants import STATE_KEY_LAST_RESPONSE from .._constants import STATE_KEY_MESSAGES from .._constants import STATE_KEY_NODE_RESPONSES +from .._constants import STATE_KEY_PENDING_AGENT_NODE_HITL from .._constants import STATE_KEY_USER_INPUT from .._event_writer import AsyncEventWriter from .._event_writer import EventWriter +from .._interrupt import interrupt from .._node_config import NodeConfig from .._state import State from .._state_mapper import SubgraphResult @@ -117,22 +122,45 @@ async def execute(self, state: State) -> dict[str, Any]: if parent_ctx is None: raise RuntimeError( f"Agent node '{self.name}' requires InvocationContext but none was set. " - "Pass context via config['configurable']['invocation_context'] when executing the graph.") + "Pass context via config['configurable']['invocation_context'] when executing the graph." + ) child_scope = self.event_scope or self.agent.name child_branch = f"{parent_ctx.branch}.{child_scope}" if parent_ctx.branch else child_scope child_user_input = child_state.get(STATE_KEY_USER_INPUT, "") + pending_hitl = self._get_pending_hitl(parent_ctx) + resume_content: Optional[Content] = None + if pending_hitl is not None: + completed_rounds = pending_hitl.get("completed", []) + for completed in completed_rounds if isinstance(completed_rounds, list) else []: + if isinstance(completed, dict): + interrupt(self._interrupt_payload(completed)) + + current_round = pending_hitl.get("current") + if not isinstance(current_round, dict): + raise RuntimeError(f"Agent node '{self.node_id}' has invalid pending HITL state.") + human_response = interrupt(self._interrupt_payload(current_round)) + resume_content = self._resume_content(current_round, human_response) + saved_child_state = current_round.get("child_state") + if isinstance(saved_child_state, dict): + child_state = dict(saved_child_state) + child_session = parent_ctx.session.model_copy(deep=True) child_session.state = dict(child_state) - child_events = self._build_child_events(parent_ctx, child_user_input, child_branch) + child_events = self._build_child_events( + parent_ctx, + child_user_input, + child_branch, + resume_content=resume_content, + ) if hasattr(child_session, "events"): child_session.events = child_events if self.isolated_messages: child_session.state[STATE_KEY_MESSAGES] = [] - child_user_content = None - if isinstance(child_user_input, str) and child_user_input: + child_user_content = resume_content + if child_user_content is None and isinstance(child_user_input, str) and child_user_input: child_user_content = Content( role="user", parts=[Part.from_text(text=child_user_input)], @@ -164,11 +192,35 @@ async def execute(self, state: State) -> dict[str, Any]: transfer_requested = False child_ctx.agent = current_agent - async for event in current_agent.run_async(child_ctx): + agent_stream = current_agent.run_async(child_ctx) + async for event in agent_stream: await self._run_agent_event_callbacks(state, event) + if isinstance(event, LongRunningEvent): + current_round = self._pending_round(event, child_ctx, final_state) + completed_rounds: list[dict[str, Any]] = [] + if pending_hitl is not None: + previous = pending_hitl.get("completed", []) + if isinstance(previous, list): + completed_rounds.extend(item for item in previous if isinstance(item, dict)) + previous_current = pending_hitl.get("current") + if isinstance(previous_current, dict): + completed_rounds.append( + {key: value for key, value in previous_current.items() if key != "child_state"} + ) + parent_ctx.state[STATE_KEY_PENDING_AGENT_NODE_HITL] = { + "node_id": self.node_id, + "completed": completed_rounds, + "current": current_round, + } + await agent_stream.aclose() + interrupt(self._interrupt_payload(current_round)) + raise RuntimeError(f"Agent node '{self.node_id}' did not suspend after LongRunningEvent.") + if (not event.partial) and hasattr(child_session, "events"): - child_session.events.append(event.model_copy(deep=True)) + existing_ids = {item.id for item in child_session.events if getattr(item, "id", None)} + if not event.id or event.id not in existing_ids: + child_session.events.append(event.model_copy(deep=True)) if event.actions and event.actions.state_delta: delta = dict(event.actions.state_delta) @@ -244,7 +296,14 @@ async def execute(self, state: State) -> dict[str, Any]: if isinstance(candidate, str) and candidate: last_response = candidate + if pending_hitl is not None: + parent_ctx.state[STATE_KEY_PENDING_AGENT_NODE_HITL] = None + + except GraphInterrupt: + raise except Exception as e: + if pending_hitl is not None: + parent_ctx.state[STATE_KEY_PENDING_AGENT_NODE_HITL] = None raise RuntimeError(f"Agent node '{self.name}' execution failed: {e}") from e if last_response: @@ -265,9 +324,7 @@ async def execute(self, state: State) -> dict[str, Any]: default_result = { STATE_KEY_LAST_RESPONSE: last_response, - STATE_KEY_NODE_RESPONSES: { - self.node_id: node_response - }, + STATE_KEY_NODE_RESPONSES: {self.node_id: node_response}, STATE_KEY_USER_INPUT: "", } @@ -276,8 +333,9 @@ async def execute(self, state: State) -> dict[str, Any]: if mapped is None: return {} if not isinstance(mapped, dict): - raise TypeError(f"Output mapper for agent node '{self.node_id}' must return dict, " - f"got {type(mapped).__name__}.") + raise TypeError( + f"Output mapper for agent node '{self.node_id}' must return dict, " f"got {type(mapped).__name__}." + ) if STATE_KEY_USER_INPUT not in mapped: mapped = dict(mapped) mapped[STATE_KEY_USER_INPUT] = "" @@ -290,10 +348,19 @@ def _build_child_events( parent_ctx: InvocationContext, child_user_input: Any, child_branch: str, + *, + resume_content: Optional[Content] = None, ) -> list[Event]: parent_events = getattr(parent_ctx.session, "events", []) - if self.isolated_messages: - child_events: list[Event] = [] + pending_hitl = self._get_pending_hitl(parent_ctx) + if self.isolated_messages and pending_hitl is not None: + child_events = [ + event.model_copy(deep=True) + for event in parent_events + if event.branch == child_branch or str(event.branch or "").startswith(f"{child_branch}.") + ] + elif self.isolated_messages: + child_events = [] else: child_events = [event.model_copy(deep=True) for event in parent_events] @@ -307,9 +374,77 @@ def _build_child_events( role="user", parts=[Part.from_text(text=child_user_input)], ), - )) + ) + ) + if resume_content is not None: + child_events.append( + Event( + invocation_id=parent_ctx.invocation_id, + author="user", + branch=child_branch, + content=resume_content.model_copy(deep=True), + ) + ) return child_events + def _get_pending_hitl(self, parent_ctx: InvocationContext) -> Optional[dict[str, Any]]: + value = parent_ctx.session.state.get(STATE_KEY_PENDING_AGENT_NODE_HITL) + if isinstance(value, dict) and value.get("node_id") == self.node_id: + return value + return None + + def _pending_round( + self, + event: LongRunningEvent, + child_ctx: InvocationContext, + child_state: dict[str, Any], + ) -> dict[str, Any]: + return { + "agent_name": event.author or self.agent.name, + "branch": event.branch or child_ctx.branch, + "function_call": event.function_call.model_dump(mode="json"), + "function_response": event.function_response.model_dump(mode="json"), + "child_state": dict(child_state), + } + + def _interrupt_payload(self, pending_round: dict[str, Any]) -> dict[str, Any]: + function_call = pending_round.get("function_call") + function_call = function_call if isinstance(function_call, dict) else {} + arguments = function_call.get("args") + arguments = arguments if isinstance(arguments, dict) else {} + function_response = pending_round.get("function_response") + function_response = function_response if isinstance(function_response, dict) else {} + response = function_response.get("response") + response = response if isinstance(response, dict) else {} + return { + "_trpc_agent_node_hitl": True, + "nodeId": self.node_id, + "agentName": str(pending_round.get("agent_name") or self.agent.name), + "toolName": str(function_call.get("name") or "graph_interrupt"), + # Long-running tools return the UI interaction contract from the + # tool execution. Preserve the model-supplied call arguments and + # let the tool result add derived fields such as stable IDs. + "arguments": {**arguments, **response}, + } + + @staticmethod + def _resume_content(pending_round: dict[str, Any], human_response: Any) -> Content: + function_call = pending_round.get("function_call") + function_call = function_call if isinstance(function_call, dict) else {} + response = human_response if isinstance(human_response, dict) else {"value": human_response} + return Content( + role="user", + parts=[ + Part( + function_response=FunctionResponse( + id=str(function_call.get("id") or ""), + name=str(function_call.get("name") or ""), + response=response, + ) + ) + ], + ) + async def _run_agent_event_callbacks(self, state: State, event: Event) -> None: if not self.callbacks or not self.callbacks.agent_event: return diff --git a/trpc_agent_sdk/models/_openai_model.py b/trpc_agent_sdk/models/_openai_model.py index bd30ae33..c84d3740 100644 --- a/trpc_agent_sdk/models/_openai_model.py +++ b/trpc_agent_sdk/models/_openai_model.py @@ -11,6 +11,7 @@ """ import base64 +import inspect import json import uuid from enum import Enum @@ -110,6 +111,10 @@ class ApiParamsKey(str, Enum): PROMPT_CACHE_RETENTION = "prompt_cache_retention" +_RESPONSES_INPUT_ITEMS = "responses_input_items" +_RESPONSES_LOGPROBS_REQUEST = "_trpc_responses_logprobs_request" + + @register_model(model_name="OpenAIModel", supported_models=[r"gpt-.*", r"o1-.*", r"deepseek-.*", r"hy3-.*"]) class OpenAIModel(LLMModel): """OpenAI model implementation using the abstract model interface. @@ -131,6 +136,11 @@ class OpenAIModel(LLMModel): will be used as the base, with per-request configs overriding specific fields. Useful for maintaining consistent model behavior across multiple calls. + use_responses_api: Use ``client.responses.create`` instead of Chat Completions. + Defaults to False for backward compatibility. + responses_api_params: Optional Responses-only parameters such as ``store``, + ``reasoning``, ``include``, or ``truncation``. The model, + input, and stream parameters remain managed by this class. **kwargs: Additional arguments passed to parent LLMModel class (e.g., api_key, base_url, etc.) @@ -168,6 +178,8 @@ def __init__( tool_prompt: str = "xml", generate_content_config: Optional[GenerateContentConfig] = None, http_client_provider_factory: HttpClientProviderFactory = temporary_http_client_provider_factory, + use_responses_api: bool = False, + responses_api_params: Optional[Dict[str, Any]] = None, **kwargs, ): super().__init__(model_name, filters_name, **kwargs) @@ -176,6 +188,12 @@ def __init__( # Extract OpenAI-specific config self.organization: str = kwargs.get(const.ORGANIZATION, "") self.client_args = kwargs.get(const.CLIENT_ARGS, {}) + self.use_responses_api = use_responses_api + self.responses_api_params = dict(responses_api_params or {}) + reserved_response_params = {"model", "input", "stream"}.intersection(self.responses_api_params) + if reserved_response_params: + names = ", ".join(sorted(reserved_response_params)) + raise ValueError(f"responses_api_params cannot override managed parameters: {names}") # Allow callers to inject a tuned httpx client http_client_provider_factory = http_client_provider_factory or temporary_http_client_provider_factory self._http_client_provider: BaseHttpClientProvider = http_client_provider_factory() @@ -234,7 +252,7 @@ def _create_async_client(self) -> openai.AsyncOpenAI: logging.getLogger("httpx").setLevel(logging.WARNING) client_args = self.client_args.copy() - client_args['http_client'] = self._http_client_provider.create_http_client() + client_args["http_client"] = self._http_client_provider.create_http_client() return openai.AsyncOpenAI( api_key=self._api_key, @@ -320,8 +338,13 @@ def _format_messages(self, request: LlmRequest) -> List[Dict[str, Any]]: parts: list[Part] = content.parts # type: ignore conditions_iter = [ - len(parts) == 1, parts[0].text, parts[0].function_call, parts[0].function_response, - parts[0].code_execution_result, parts[0].executable_code, parts[0].inline_data + len(parts) == 1, + parts[0].text, + parts[0].function_call, + parts[0].function_response, + parts[0].code_execution_result, + parts[0].executable_code, + parts[0].inline_data, ] # Handle different content structures if all(conditions_iter): @@ -337,11 +360,21 @@ def _format_messages(self, request: LlmRequest) -> List[Dict[str, Any]]: text_parts = [] image_parts = [] tool_calls = [] + responses_input_items = [] for part in parts: # type: ignore + if part.thought: + if self.use_responses_api and part.thought_signature: + try: + raw = part.thought_signature + raw = raw.decode("utf-8") if isinstance(raw, bytes) else raw + reasoning_item = json.loads(raw) + if isinstance(reasoning_item, dict) and reasoning_item.get("type") == "reasoning": + responses_input_items.append(reasoning_item) + except (UnicodeDecodeError, json.JSONDecodeError, TypeError): + logger.warning("Ignoring invalid Responses reasoning item metadata") + continue if part.text: - if part.thought: - continue text_parts.append(part.text) elif part.inline_data and part.inline_data.mime_type: # Handle image data - convert to OpenAI vision format @@ -355,10 +388,12 @@ def _format_messages(self, request: LlmRequest) -> List[Dict[str, Any]]: "id": getattr(part.function_call, "id", None) or f"call_{uuid.uuid4().hex[:24]}", "type": "function", "function": { - "name": - part.function_call.name, - "arguments": (part.function_call.args if isinstance(part.function_call.args, str) - else json.dumps(part.function_call.args, ensure_ascii=False)), + "name": part.function_call.name, + "arguments": ( + part.function_call.args + if isinstance(part.function_call.args, str) + else json.dumps(part.function_call.args, ensure_ascii=False) + ), }, } if self._adapter.should_include_thought_signature(): @@ -405,8 +440,11 @@ def _format_messages(self, request: LlmRequest) -> List[Dict[str, Any]]: else: for func_response in function_responses: # Standard tool message format for OpenAI API - raw_text = (json.dumps(func_response.response, ensure_ascii=False) if isinstance( - func_response.response, dict) else str(func_response.response)) + raw_text = ( + json.dumps(func_response.response, ensure_ascii=False) + if isinstance(func_response.response, dict) + else str(func_response.response) + ) clipped_text = self._clip_tool_response_text( raw_text, getattr(func_response, "name", "tool"), @@ -448,6 +486,9 @@ def _format_messages(self, request: LlmRequest) -> List[Dict[str, Any]]: if self._adapter.should_backfill_reasoning_content(role, message): message[const.REASONING_CONTENT] = "" + if responses_input_items: + message[_RESPONSES_INPUT_ITEMS] = responses_input_items + formatted_messages.append(message) # Validate and fix message sequence for OpenAI compatibility @@ -487,10 +528,9 @@ def _validate_and_fix_openai_messages(self, messages: List[Dict[str, Any]]) -> L dummy_response = { const.ROLE: const.TOOL, const.TOOL_CALL_ID: pending_call["id"], - const.CONTENT: json.dumps({ - "status": "completed", - "note": "Tool call completed by system" - }), + const.CONTENT: json.dumps( + {"status": "completed", "note": "Tool call completed by system"} + ), } fixed_messages.append(dummy_response) @@ -503,16 +543,17 @@ def _validate_and_fix_openai_messages(self, messages: List[Dict[str, Any]]) -> L # Assistant message without tool calls if pending_tool_calls: # Need to add dummy responses for pending tool calls - logger.warning("Adding dummy tool responses for %s pending tool calls before assistant message", - len(pending_tool_calls)) + logger.warning( + "Adding dummy tool responses for %s pending tool calls before assistant message", + len(pending_tool_calls), + ) for pending_call in pending_tool_calls: dummy_response = { const.ROLE: const.TOOL, const.TOOL_CALL_ID: pending_call["id"], - const.CONTENT: json.dumps({ - "status": "completed", - "note": "Tool call completed by system" - }), + const.CONTENT: json.dumps( + {"status": "completed", "note": "Tool call completed by system"} + ), } fixed_messages.append(dummy_response) pending_tool_calls = [] @@ -533,16 +574,16 @@ def _validate_and_fix_openai_messages(self, messages: List[Dict[str, Any]]) -> L # User or system message if pending_tool_calls: # Add dummy responses for any pending tool calls before user/system message - logger.warning("Adding dummy tool responses for %s pending tool calls before %s message", - len(pending_tool_calls), role) + logger.warning( + "Adding dummy tool responses for %s pending tool calls before %s message", + len(pending_tool_calls), + role, + ) for pending_call in pending_tool_calls: dummy_response = { const.ROLE: const.TOOL, const.TOOL_CALL_ID: pending_call["id"], - const.CONTENT: json.dumps({ - "status": "completed", - "note": "Tool call completed by system" - }), + const.CONTENT: json.dumps({"status": "completed", "note": "Tool call completed by system"}), } fixed_messages.append(dummy_response) pending_tool_calls = [] @@ -556,10 +597,7 @@ def _validate_and_fix_openai_messages(self, messages: List[Dict[str, Any]]) -> L dummy_response = { const.ROLE: const.TOOL, const.TOOL_CALL_ID: pending_call["id"], - const.CONTENT: json.dumps({ - "status": "completed", - "note": "Tool call completed by system" - }), + const.CONTENT: json.dumps({"status": "completed", "note": "Tool call completed by system"}), } fixed_messages.append(dummy_response) @@ -609,8 +647,10 @@ def _is_thinking_event(self, response: dict) -> bool: Returns: `bool`: True if this is a thinking event """ - return (response.get("object", "") == "stream_server.event" - and response.get("event", {}).get("name", "") == "thinking") + return ( + response.get("object", "") == "stream_server.event" + and response.get("event", {}).get("name", "") == "thinking" + ) def _set_thinking(self, request: LlmRequest, http_options: dict): """Set thinking parameters from request config.""" @@ -654,11 +694,15 @@ def _set_thinking(self, request: LlmRequest, http_options: dict): if thinking_config.thinking_budget <= max_output_tokens: processed_extra_body[const.THINKING_TOKENS] = thinking_config.thinking_budget else: - raise ValueError(f"thinking_budget: {thinking_config.thinking_budget} " - f"must be between 1024 and {max_output_tokens}") + raise ValueError( + f"thinking_budget: {thinking_config.thinking_budget} " + f"must be between 1024 and {max_output_tokens}" + ) else: - raise ValueError(f"Invalid thinking_budget value: {thinking_config.thinking_budget}. " - "Must be 0 (disabled), -1 (automatic), or positive integer.") + raise ValueError( + f"Invalid thinking_budget value: {thinking_config.thinking_budget}. " + "Must be 0 (disabled), -1 (automatic), or positive integer." + ) def _get_thinking_state(self, response: dict) -> int: """Get the thinking state from a thinking event. @@ -687,15 +731,14 @@ def _process_tool_call_delta(self, tool_call_delta: dict, accumulated_tool_calls # Ensure we have enough slots in accumulated_tool_calls while len(accumulated_tool_calls) <= index: - accumulated_tool_calls.append({ - ToolKey.ID: "", - ToolKey.TYPE: ToolKey.FUNCTION, - ToolKey.FUNCTION: { - ToolKey.NAME: "", - ToolKey.ARGUMENTS: "" - }, - ToolKey.THOUGHT_SIGNATURE: "", - }) + accumulated_tool_calls.append( + { + ToolKey.ID: "", + ToolKey.TYPE: ToolKey.FUNCTION, + ToolKey.FUNCTION: {ToolKey.NAME: "", ToolKey.ARGUMENTS: ""}, + ToolKey.THOUGHT_SIGNATURE: "", + } + ) # Capture thought_signature from delta or provider_specific_fields for next-round pass-through thought_sig = tool_call_delta.get(ToolKey.THOUGHT_SIGNATURE) @@ -715,8 +758,11 @@ def _process_tool_call_delta(self, tool_call_delta: dict, accumulated_tool_calls if ToolKey.FUNCTION in tool_call_delta: function_delta = tool_call_delta[ToolKey.FUNCTION] - if (ToolKey.NAME in function_delta and function_delta[ToolKey.NAME] is not None - and function_delta[ToolKey.NAME] != ""): + if ( + ToolKey.NAME in function_delta + and function_delta[ToolKey.NAME] is not None + and function_delta[ToolKey.NAME] != "" + ): accumulated_tool_calls[index][ToolKey.FUNCTION][ToolKey.NAME] = function_delta[ToolKey.NAME] # If function_delta[ToolKey.NAME] is None or empty, keep the existing name value if ToolKey.ARGUMENTS in function_delta and function_delta[ToolKey.ARGUMENTS] is not None: @@ -729,14 +775,16 @@ def _build_usage_metadata(usage_data: dict) -> GenerateContentResponseUsageMetad ``cache_read_input_tokens`` prefers Anthropic/LiteLLM-style top-level fields; falls back to OpenAI-style ``prompt_tokens_details.cached_tokens``. """ - completion_details = usage_data.get("completion_tokens_details") or {} + completion_details = ( + usage_data.get("completion_tokens_details") or usage_data.get("output_tokens_details") or {} + ) cache_read = usage_data.get("cache_read_input_tokens") if cache_read is None: - details = usage_data.get("prompt_tokens_details") + details = usage_data.get("prompt_tokens_details") or usage_data.get("input_tokens_details") cache_read = details.get("cached_tokens") if isinstance(details, dict) else None return GenerateContentResponseUsageMetadata( - prompt_token_count=usage_data.get("prompt_tokens", 0), - candidates_token_count=usage_data.get("completion_tokens", 0), + prompt_token_count=usage_data.get("prompt_tokens", usage_data.get("input_tokens", 0)), + candidates_token_count=usage_data.get("completion_tokens", usage_data.get("output_tokens", 0)), thoughts_token_count=completion_details.get("reasoning_tokens"), total_token_count=usage_data.get("total_tokens", 0), cache_read_input_tokens=cache_read, @@ -849,15 +897,20 @@ def _create_complete_tool_calls(self, accumulated_tool_calls: list[dict]) -> Opt logger.warning("Generated fallback ID '%s' for tool call with missing ID", tool_call_id) thought_sig = tool_call_data.get(ToolKey.THOUGHT_SIGNATURE) or None - logger.debug("Creating tool call: id=%s, name=%s, arguments=%s", tool_call_id, - function_map[ToolKey.NAME], arguments) + logger.debug( + "Creating tool call: id=%s, name=%s, arguments=%s", + tool_call_id, + function_map[ToolKey.NAME], + arguments, + ) complete_tool_calls.append( ToolCall( id=tool_call_id, name=function_map[ToolKey.NAME], arguments=arguments, thought_signature=thought_sig, - )) + ) + ) except json.JSONDecodeError as ex: # Arguments not complete yet, skip this tool call logger.debug("JSON decode error for tool call %s: %s", i, ex) @@ -996,7 +1049,8 @@ def _process_tool_calls_from_message(self, message: dict) -> Optional[List[ToolC name=tool_call[ToolKey.FUNCTION][ToolKey.NAME], arguments=arguments, thought_signature=thought_sig, - )) + ) + ) except (KeyError, json.JSONDecodeError, TypeError) as ex: logger.warning("Failed to parse tool call: %s, error: %s", tool_call, ex) continue @@ -1051,7 +1105,8 @@ def _create_response_with_content(self, response_dict: dict) -> LlmResponse: tool_call = ToolCall( id=f"call_{uuid.uuid4().hex[:24]}", name=func_call.name, # type: ignore - arguments=func_call.args) # type: ignore + arguments=func_call.args, + ) # type: ignore tool_calls.append(tool_call) except Exception as ex: # pylint: disable=broad-except logger.warning("Failed to parse function calls from text content: %s", ex) @@ -1102,11 +1157,252 @@ def _create_response_with_content(self, response_dict: dict) -> LlmResponse: return LlmResponse(content=content, usage_metadata=usage, error_code=error_code, response_id=response_id) - async def _generate_single(self, - api_params: Dict, - request: LlmRequest, - http_options: Dict[str, Any] | None = None, - ctx: InvocationContext | None = None) -> LlmResponse: + @staticmethod + def _model_dump(value: Any) -> dict: + """Return a dictionary for OpenAI SDK models and test doubles.""" + if isinstance(value, dict): + return value + return value.model_dump() + + def _convert_messages_to_responses_input(self, messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """Convert Chat Completions messages to Responses input items.""" + input_items: List[Dict[str, Any]] = [] + for message in messages: + role = message.get(const.ROLE, const.USER) + if role == const.TOOL: + input_items.append( + { + "type": "function_call_output", + "call_id": message.get(const.TOOL_CALL_ID, "unknown"), + "output": str(message.get(const.CONTENT, "")), + } + ) + continue + + input_items.extend(message.get(_RESPONSES_INPUT_ITEMS) or []) + + content = message.get(const.CONTENT) + if content not in (None, "", []): + if isinstance(content, list): + converted_content = [] + for item in content: + item_type = item.get("type") + if item_type == "image_url": + image = item.get("image_url") or {} + converted_content.append( + { + "type": "input_image", + "image_url": image.get("url", ""), + "detail": image.get("detail", "auto"), + } + ) + elif item_type == "text": + converted_content.append( + { + "type": "output_text" if role == const.ASSISTANT else "input_text", + "text": item.get("text", ""), + } + ) + else: + converted_content.append(item) + content = converted_content + input_items.append({"role": role, "content": content}) + + for tool_call in message.get(const.TOOL_CALLS, []) or []: + function = tool_call.get(ToolKey.FUNCTION, {}) + input_items.append( + { + "type": "function_call", + "call_id": tool_call.get(ToolKey.ID) or f"call_{uuid.uuid4().hex[:24]}", + "name": function.get(ToolKey.NAME, ""), + "arguments": function.get(ToolKey.ARGUMENTS, "{}"), + } + ) + return input_items + + @staticmethod + def _convert_tools_to_responses_format(tools: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """Flatten Chat Completions function definitions for Responses.""" + converted = [] + for tool in tools: + if tool.get("type") != "function": + converted.append(tool) + continue + function = dict(tool.get("function") or {}) + function["type"] = "function" + converted.append(function) + return converted + + @staticmethod + def _convert_response_format_to_responses(response_format: Dict[str, Any]) -> Dict[str, Any]: + """Convert Chat Completions response_format to Responses text config.""" + if response_format.get("type") == "json_schema": + schema = dict(response_format.get("json_schema") or {}) + return {"format": {"type": "json_schema", **schema}} + return {"format": dict(response_format)} + + def _convert_api_params_to_responses(self, api_params: Dict[str, Any]) -> Dict[str, Any]: + """Translate shared generation parameters to ``responses.create``.""" + responses_params: Dict[str, Any] = { + "model": api_params[ApiParamsKey.MODEL], + "input": self._convert_messages_to_responses_input(api_params[ApiParamsKey.MESSAGES]), + "stream": api_params[ApiParamsKey.STREAM], + } + parameter_map = { + ApiParamsKey.TEMPERATURE: "temperature", + ApiParamsKey.TOP_P: "top_p", + ApiParamsKey.TOOL_CHOICE: "tool_choice", + ApiParamsKey.PARALLEL_TOOL_CALLS: "parallel_tool_calls", + ApiParamsKey.PROMPT_CACHE_KEY: "prompt_cache_key", + ApiParamsKey.PROMPT_CACHE_RETENTION: "prompt_cache_retention", + } + for source, target in parameter_map.items(): + if source in api_params: + responses_params[target] = api_params[source] + + # The OpenAI Python SDK has used more than one Responses logprobs + # shape. Keep this framework-level request private until we can inspect + # the installed client's supported parameters immediately before send. + response_logprobs = api_params.get(ApiParamsKey.LOGPROBS) + top_logprobs = api_params.get(ApiParamsKey.TOP_LOGPROBS) + if response_logprobs is not None or top_logprobs is not None: + responses_params[_RESPONSES_LOGPROBS_REQUEST] = { + "enabled": bool(response_logprobs) if response_logprobs is not None else bool(top_logprobs), + "top_logprobs": top_logprobs, + } + + max_output_tokens = api_params.get(ApiParamsKey.MAX_COMPLETION_TOKENS) or api_params.get( + ApiParamsKey.MAX_TOKENS + ) + if max_output_tokens: + responses_params["max_output_tokens"] = max_output_tokens + if ApiParamsKey.TOOLS in api_params: + responses_params["tools"] = self._convert_tools_to_responses_format(api_params[ApiParamsKey.TOOLS]) + if ApiParamsKey.RESPONSE_FORMAT in api_params: + responses_params["text"] = self._convert_response_format_to_responses( + api_params[ApiParamsKey.RESPONSE_FORMAT] + ) + + responses_params.update(self.responses_api_params) + if responses_params.get("store") is False: + include = list(responses_params.get("include") or []) + if "reasoning.encrypted_content" not in include: + include.append("reasoning.encrypted_content") + responses_params["include"] = include + return responses_params + + @staticmethod + def _prepare_responses_api_params(client: Any, api_params: Dict[str, Any]) -> Dict[str, Any]: + """Adapt optional logprobs to the installed OpenAI client's capability. + + ``openai>=1.66`` exposes Responses but accepts neither of the later + logprobs parameter shapes. Newer clients may expose either a top-level + ``top_logprobs`` integer or a structured ``logprobs`` object. + """ + prepared = dict(api_params) + requested = prepared.pop(_RESPONSES_LOGPROBS_REQUEST, None) + if not requested or not requested["enabled"]: + return prepared + + try: + supported = inspect.signature(client.responses.create).parameters + except (TypeError, ValueError) as exc: + raise ValueError("Unable to determine Responses logprobs support") from exc + + if "logprobs" in supported: + prepared["logprobs"] = requested + return prepared + + top_logprobs = requested["top_logprobs"] + if "top_logprobs" in supported and top_logprobs: + prepared["top_logprobs"] = top_logprobs + return prepared + + raise ValueError( + "Responses logprobs requires an OpenAI SDK that supports the Responses " + "logprobs parameters; upgrade openai or disable logprobs." + ) + + @staticmethod + def _responses_error(response_dict: dict) -> tuple[Optional[str], Optional[str]]: + status = response_dict.get("status") + if status not in {"failed", "incomplete", "cancelled"}: + return None, None + error = response_dict.get("error") or response_dict.get("incomplete_details") or {} + if isinstance(error, dict): + return str(error.get("code") or status), str(error.get("message") or error.get("reason") or status) + return str(status), str(error) + + def _create_responses_response(self, response_dict: dict) -> LlmResponse: + """Convert a completed Responses payload into ``LlmResponse``.""" + parts = [] + for item in response_dict.get("output") or []: + item_type = item.get("type") + if item_type == "reasoning": + summaries = item.get("summary") or [] + if not summaries: + part = Part.from_text(text="") + part.thought = True + part.thought_signature = json.dumps(item, ensure_ascii=False).encode("utf-8") + parts.append(part) + for index, summary in enumerate(summaries): + text = summary.get("text") if isinstance(summary, dict) else None + if text: + part = Part.from_text(text=text) + part.thought = True + if index == 0: + part.thought_signature = json.dumps(item, ensure_ascii=False).encode("utf-8") + parts.append(part) + elif item_type == "message": + for output_part in item.get("content") or []: + text = output_part.get("text") or output_part.get("refusal") + if text: + part = Part.from_text(text=text) + part.thought = False + parts.append(part) + elif item_type == "function_call": + arguments = json_loads_repair(item.get("arguments") or "{}") + if not isinstance(arguments, dict): + logger.warning("Skipping Responses function call with non-dict arguments: %r", arguments) + continue + part = Part.from_function_call(name=item.get("name", ""), args=arguments) + part.function_call.id = item.get("call_id") or item.get("id") # type: ignore + parts.append(part) + + content = Content(parts=parts, role=const.MODEL) if parts else None + usage = self._process_usage_from_response(response_dict) + error_code, error_message = self._responses_error(response_dict) + return LlmResponse( + content=content, + usage_metadata=usage, + error_code=error_code, + error_message=error_message, + response_id=response_dict.get("id"), + ) + + async def _generate_responses_single( + self, + api_params: Dict[str, Any], + http_options: Optional[Dict[str, Any]] = None, + ) -> LlmResponse: + """Generate one non-streaming response through the Responses API.""" + client = self._create_async_client() + try: + response = await client.responses.create( + **self._prepare_responses_api_params(client, api_params), + **(http_options or {}), + ) + return self._create_responses_response(self._model_dump(response)) + finally: + await self._http_client_provider.close_http_client(client) + + async def _generate_single( + self, + api_params: Dict, + request: LlmRequest, + http_options: Dict[str, Any] | None = None, + ctx: InvocationContext | None = None, + ) -> LlmResponse: """Generate a single response (non-streaming).""" if http_options is None: http_options = {} @@ -1157,7 +1453,8 @@ def _convert_tools_to_openai_format(self, tools: List[Tool]) -> List[Dict[str, A # Convert parameters schema - always include parameters field if func_decl.parameters: openai_tool["function"]["parameters"] = self._convert_schema_to_openai_format( - func_decl.parameters) + func_decl.parameters + ) else: # When parameters are empty, provide the proper OpenAI format structure openai_tool["function"]["parameters"] = {"type": "object", "properties": {}} @@ -1295,22 +1592,14 @@ def _build_response_format(self, config: GenerateContentConfig) -> Optional[Dict openai_schema = self._ensure_additional_properties_false(openai_schema) return { "type": "json_schema", - "json_schema": { - "name": "response_schema", - "schema": openai_schema, - "strict": True - }, + "json_schema": {"name": "response_schema", "schema": openai_schema, "strict": True}, } elif config.response_json_schema: # Use provided JSON schema directly processed_schema = self._ensure_additional_properties_false(config.response_json_schema) return { "type": "json_schema", - "json_schema": { - "name": "response_schema", - "schema": processed_schema, - "strict": True - }, + "json_schema": {"name": "response_schema", "schema": processed_schema, "strict": True}, } else: # Basic JSON mode @@ -1428,13 +1717,13 @@ def _log_unsupported_config_options(self, config: GenerateContentConfig) -> None if unsupported_options: logger.warning( "The following configuration options are not supported in OpenAI models and will be ignored: %s", - ', '.join(unsupported_options)) + ", ".join(unsupported_options), + ) @override - async def _generate_async_impl(self, - request: LlmRequest, - stream: bool = False, - ctx: InvocationContext | None = None) -> AsyncGenerator[LlmResponse, None]: + async def _generate_async_impl( + self, request: LlmRequest, stream: bool = False, ctx: InvocationContext | None = None + ) -> AsyncGenerator[LlmResponse, None]: """Generate content asynchronously.""" self.validate_request(request) @@ -1443,8 +1732,12 @@ async def _generate_async_impl(self, # Update request with merged config request.config = merged_config - if (request.config and request.config.tools and self._adapter.requires_add_tools_to_prompt() - and not self.add_tools_to_prompt): + if ( + request.config + and request.config.tools + and self._adapter.requires_add_tools_to_prompt() + and not self.add_tools_to_prompt + ): raise ValueError(f"{self._model_name} requires add_tools_to_prompt=True when tools are used.") # Prepare OpenAI API parameters @@ -1481,18 +1774,23 @@ async def _generate_async_impl(self, api_params[ApiParamsKey.STOP] = request.config.stop_sequences # Additional OpenAI-specific parameters - if (request.config.frequency_penalty is not None - and not self._adapter.should_skip_config_param("frequency_penalty")): + if request.config.frequency_penalty is not None and not self._adapter.should_skip_config_param( + "frequency_penalty" + ): api_params[ApiParamsKey.FREQUENCY_PENALTY] = request.config.frequency_penalty - if (request.config.presence_penalty is not None - and not self._adapter.should_skip_config_param("presence_penalty")): + if request.config.presence_penalty is not None and not self._adapter.should_skip_config_param( + "presence_penalty" + ): api_params[ApiParamsKey.PRESENCE_PENALTY] = request.config.presence_penalty if request.config.seed is not None and not self._adapter.should_skip_config_param("seed"): api_params[ApiParamsKey.SEED] = request.config.seed # Handle candidate count (maps to OpenAI's 'n' parameter) - if (request.config.candidate_count is not None and request.config.candidate_count > 0 - and not self._adapter.should_skip_config_param("candidate_count")): + if ( + request.config.candidate_count is not None + and request.config.candidate_count > 0 + and not self._adapter.should_skip_config_param("candidate_count") + ): api_params[ApiParamsKey.N] = request.config.candidate_count # Handle logprobs configuration @@ -1528,21 +1826,210 @@ async def _generate_async_impl(self, http_options = {} if request.config: http_options = self._extract_http_options(request.config) - # set thinking params - self._set_thinking(request, http_options) + if self.use_responses_api: + api_params = self._convert_api_params_to_responses(api_params) + if ( + request.config + and request.config.thinking_config + and request.config.thinking_config.include_thoughts + and request.config.thinking_config.thinking_budget != 0 + ): + reasoning = dict(api_params.get("reasoning") or {}) + reasoning.setdefault("summary", "auto") + api_params["reasoning"] = reasoning + else: + # Chat Completions provider-specific thinking params. + self._set_thinking(request, http_options) if stream: - async for response in self._generate_stream(api_params, request, http_options, ctx): + generator = ( + self._generate_responses_stream(api_params, request, http_options) + if self.use_responses_api + else self._generate_stream(api_params, request, http_options, ctx) + ) + async for response in generator: yield response else: - response = await self._generate_single(api_params, request, http_options, ctx) + response = ( + await self._generate_responses_single(api_params, http_options) + if self.use_responses_api + else await self._generate_single(api_params, request, http_options, ctx) + ) yield response - async def _generate_stream(self, - api_params: Dict, - request: LlmRequest, - http_options: Dict[str, Any] | None = None, - ctx: InvocationContext | None = None) -> AsyncGenerator[LlmResponse, None]: + async def _generate_responses_stream( + self, + api_params: Dict[str, Any], + request: LlmRequest, + http_options: Optional[Dict[str, Any]] = None, + ) -> AsyncGenerator[LlmResponse, None]: + """Generate streaming responses through ``responses.create``.""" + client = self._create_async_client() + response_id: Optional[str] = None + completed_response: Optional[dict] = None + accumulated_text = "" + accumulated_reasoning = "" + function_calls: Dict[str, Dict[str, Any]] = {} + function_order: List[str] = [] + streaming_tool_names = getattr(request, "streaming_tool_names", None) or set() + + def upsert_function(item: dict) -> tuple[str, Dict[str, Any]]: + item_id = str(item.get("id") or item.get("call_id") or f"fc_{len(function_order)}") + if item_id not in function_calls: + function_calls[item_id] = { + "type": "function_call", + "id": item.get("id"), + "call_id": item.get("call_id") or item.get("id"), + "name": item.get("name", ""), + "arguments": item.get("arguments") or "", + } + function_order.append(item_id) + else: + current = function_calls[item_id] + for key in ("id", "call_id", "name"): + if item.get(key): + current[key] = item[key] + if item.get("arguments"): + current["arguments"] = item["arguments"] + return item_id, function_calls[item_id] + + try: + response = await client.responses.create( + **self._prepare_responses_api_params(client, api_params), + **(http_options or {}), + ) + if response is None: + raise ValueError("Empty response from Responses API") + + async for event in response: + event_dict = self._model_dump(event) + event_type = event_dict.get("type", "") + logger.debug("OpenAI Responses event: %s", json.dumps(event_dict, ensure_ascii=False)) + + response_data = event_dict.get("response") or {} + if response_id is None and response_data.get("id"): + response_id = response_data["id"] + if response_id is None and event_dict.get("response_id"): + response_id = event_dict["response_id"] + + if event_type == "error": + completed_response = { + "id": response_id, + "status": "failed", + "error": { + "code": event_dict.get("code") or "responses_stream_error", + "message": event_dict.get("message") or "Responses stream failed", + }, + "output": [], + } + continue + + if event_type in {"response.output_text.delta", "response.refusal.delta"}: + delta = event_dict.get("delta") or "" + if delta: + accumulated_text += delta + part = Part.from_text(text=delta) + part.thought = False + yield LlmResponse( + content=Content(parts=[part], role=const.MODEL), + partial=True, + response_id=response_id, + custom_metadata={const.CHUNK: event_dict}, + ) + continue + + if event_type in {"response.reasoning_summary_text.delta", "response.reasoning_text.delta"}: + delta = event_dict.get("delta") or "" + if delta: + accumulated_reasoning += delta + part = Part.from_text(text=delta) + part.thought = True + yield LlmResponse( + content=Content(parts=[part], role=const.MODEL), + partial=True, + response_id=response_id, + custom_metadata={const.CHUNK: event_dict}, + ) + continue + + if event_type in {"response.output_item.added", "response.output_item.done"}: + item = event_dict.get("item") or {} + if item.get("type") == "function_call": + upsert_function(item) + continue + + if event_type == "response.function_call_arguments.delta": + item_id = str(event_dict.get("item_id") or f"fc_{event_dict.get('output_index', 0)}") + if item_id not in function_calls: + function_order.append(item_id) + function_calls[item_id] = { + "type": "function_call", + "id": event_dict.get("item_id"), + "call_id": event_dict.get("call_id") or event_dict.get("item_id"), + "name": event_dict.get("name", ""), + "arguments": "", + } + function_call = function_calls[item_id] + delta = event_dict.get("delta") or "" + function_call["arguments"] += delta + name = function_call.get("name", "") + if delta and name and name in streaming_tool_names: + part = Part.from_function_call(name=name, args={const.TOOL_STREAMING_ARGS: delta}) + part.function_call.id = function_call.get("call_id") # type: ignore + yield LlmResponse( + content=Content(parts=[part], role=const.MODEL), + partial=True, + response_id=response_id, + custom_metadata={const.CHUNK: event_dict, const.TOOL_STREAMING: True}, + ) + continue + + if event_type == "response.function_call_arguments.done": + item_id = str(event_dict.get("item_id") or f"fc_{event_dict.get('output_index', 0)}") + item = event_dict.get("item") or {} + if item.get("type") == "function_call": + upsert_function(item) + elif item_id in function_calls: + function_calls[item_id]["arguments"] = event_dict.get("arguments") or "" + continue + + if event_type in {"response.completed", "response.failed", "response.incomplete"}: + completed_response = response_data + + if completed_response is None: + output = [] + if accumulated_reasoning: + output.append( + { + "type": "reasoning", + "summary": [{"type": "summary_text", "text": accumulated_reasoning}], + } + ) + if accumulated_text: + output.append( + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": accumulated_text}], + } + ) + output.extend(function_calls[item_id] for item_id in function_order) + completed_response = {"id": response_id, "status": "completed", "output": output} + + final_response = self._create_responses_response(completed_response) + final_response.partial = False + final_response.custom_metadata = {"stream_complete": True} + yield final_response + finally: + await self._http_client_provider.close_http_client(client) + + async def _generate_stream( + self, + api_params: Dict, + request: LlmRequest, + http_options: Dict[str, Any] | None = None, + ctx: InvocationContext | None = None, + ) -> AsyncGenerator[LlmResponse, None]: """Generate streaming responses.""" if http_options is None: http_options = {} @@ -1555,7 +2042,7 @@ async def _generate_stream(self, response_id: str | None = None # Track response ID from API # For streaming tool call arguments - get the set of tool names that should stream - streaming_tool_names = getattr(request, 'streaming_tool_names', None) or set() + streaming_tool_names = getattr(request, "streaming_tool_names", None) or set() # Create tool prompt instance for streaming if needed tool_prompt = None @@ -1613,9 +2100,9 @@ async def _generate_stream(self, if streaming_tool_names and delta_arguments and accumulated_tool_calls: # Yield streaming tool call event with delta arguments # Only for tools in streaming_tool_names - streaming_event = self._create_streaming_tool_call_response(accumulated_tool_calls, - delta_arguments, - streaming_tool_names) + streaming_event = self._create_streaming_tool_call_response( + accumulated_tool_calls, delta_arguments, streaming_tool_names + ) if streaming_event: yield streaming_event @@ -1633,11 +2120,15 @@ async def _generate_stream(self, reasoning_content = delta.get(const.REASONING_CONTENT) if reasoning_content is not None: partial_text = reasoning_content - if (tool_prompt and streaming_text_filter_state is not None - and self._adapter.should_filter_reasoning_text()): + if ( + tool_prompt + and streaming_text_filter_state is not None + and self._adapter.should_filter_reasoning_text() + ): reasoning_filter_state = streaming_text_filter_state["reasoning"] - partial_text = self._adapter.filter_streaming_text(reasoning_content, - reasoning_filter_state) + partial_text = self._adapter.filter_streaming_text( + reasoning_content, reasoning_filter_state + ) if not partial_text: continue @@ -1649,10 +2140,12 @@ async def _generate_stream(self, content_part.thought = True partial_content = Content(parts=[content_part], role=const.MODEL) - yield LlmResponse(content=partial_content, - partial=True, - response_id=response_id, - custom_metadata={const.CHUNK: chunk_dict}) + yield LlmResponse( + content=partial_content, + partial=True, + response_id=response_id, + custom_metadata={const.CHUNK: chunk_dict}, + ) # Handle regular content if delta.get(const.CONTENT): @@ -1677,10 +2170,12 @@ async def _generate_stream(self, content_part.thought = is_thinking partial_content = Content(parts=[content_part], role=const.MODEL) - yield LlmResponse(content=partial_content, - partial=True, - response_id=response_id, - custom_metadata={const.CHUNK: chunk_dict}) + yield LlmResponse( + content=partial_content, + partial=True, + response_id=response_id, + custom_metadata={const.CHUNK: chunk_dict}, + ) # Handle usage usage = self._process_usage(chunk_dict) @@ -1690,16 +2185,19 @@ async def _generate_stream(self, if tool_prompt and streaming_text_filter_state is not None: if self._adapter.should_filter_reasoning_text(): flushed_reasoning_text = self._adapter.flush_streaming_text( - streaming_text_filter_state["reasoning"]) + streaming_text_filter_state["reasoning"] + ) if flushed_reasoning_text: thought_content += flushed_reasoning_text content_part = Part.from_text(text=flushed_reasoning_text) content_part.thought = True partial_content = Content(parts=[content_part], role=const.MODEL) - yield LlmResponse(content=partial_content, - partial=True, - response_id=response_id, - custom_metadata={"stream_filter_flushed": "reasoning"}) + yield LlmResponse( + content=partial_content, + partial=True, + response_id=response_id, + custom_metadata={"stream_filter_flushed": "reasoning"}, + ) flushed_content_text = self._adapter.flush_streaming_text(streaming_text_filter_state["content"]) if flushed_content_text: @@ -1708,10 +2206,12 @@ async def _generate_stream(self, content_part = Part.from_text(text=flushed_content_text) content_part.thought = is_thinking partial_content = Content(parts=[content_part], role=const.MODEL) - yield LlmResponse(content=partial_content, - partial=True, - response_id=response_id, - custom_metadata={"stream_filter_flushed": "content"}) + yield LlmResponse( + content=partial_content, + partial=True, + response_id=response_id, + custom_metadata={"stream_filter_flushed": "content"}, + ) # Yield final complete response final_content = None @@ -1729,17 +2229,19 @@ async def _generate_stream(self, if tool_prompt and accumulated_content and not complete_tool_calls: try: parsed_function_calls = self._adapter.parse_tool_prompt_function_calls( - accumulated_content, tool_prompt) + accumulated_content, tool_prompt + ) if parsed_function_calls: # Convert FunctionCall objects to ToolCall objects complete_tool_calls = [] for func_call in parsed_function_calls: - tool_call = ToolCall(id=f"call_{uuid.uuid4().hex[:24]}", - name=func_call.name, - arguments=func_call.args) + tool_call = ToolCall( + id=f"call_{uuid.uuid4().hex[:24]}", name=func_call.name, arguments=func_call.args + ) complete_tool_calls.append(tool_call) - logger.debug("Parsed %s function calls from final accumulated content", - len(complete_tool_calls)) + logger.debug( + "Parsed %s function calls from final accumulated content", len(complete_tool_calls) + ) except Exception as ex: # pylint: disable=broad-except logger.warning("Failed to parse function calls from final accumulated content: %s", ex) diff --git a/trpc_agent_sdk/server/ag_ui/_core/_agui_agent.py b/trpc_agent_sdk/server/ag_ui/_core/_agui_agent.py index 26748674..534dc249 100644 --- a/trpc_agent_sdk/server/ag_ui/_core/_agui_agent.py +++ b/trpc_agent_sdk/server/ag_ui/_core/_agui_agent.py @@ -64,6 +64,12 @@ from ._http_req import set_agui_http_req from ._session_manager import SessionManager +# Prefix for LangGraph checkpoint state keys stored in Session.state +# (see trpc_agent_sdk/dsl/graph/_constants.py: STATE_KEY_CHECKPOINTS etc.). +# Used to detect GraphAgent resume scenarios where the session state must +# be preserved to avoid overwriting the LangGraph checkpoint. +_GRAPH_CHECKPOINT_STATE_KEY_PREFIX = "_trpc_graph_checkpoint" + class AgUiAgent: """Middleware to bridge AG-UI Protocol with TRPC agents. @@ -1100,13 +1106,22 @@ async def _run_trpc_in_background(self, set_agui_http_req(run_config, http_request) # Ensure session exists - await self._ensure_session_exists(app_name, user_id, input.thread_id, input.state) - - # this will always update the backend states with the frontend states - # Recipe Demo Example: if there is a state "salt" in the ingredients state and in frontend user - # remove this salt state using UI from the ingredients list then our backend should also update - # these state changes as well to sync both the states - await self._session_manager.update_session_state(input.thread_id, app_name, user_id, input.state) + session = await self._ensure_session_exists(app_name, user_id, input.thread_id, input.state) + + # A GraphAgent tool result is a continuation token, not a fresh + # state snapshot. Its LangGraph checkpoint lives in Session.state; + # applying the browser state can replace that checkpoint and make + # Command(resume=...) start a new graph from START. Keep the + # established state synchronisation for normal turns and for + # non-graph agents that intentionally submit a state patch with a + # tool result. + if not self._is_graph_checkpoint_resume(input, session=session): + await self._session_manager.update_session_state( + input.thread_id, + app_name, + user_id, + input.state, + ) # Convert messages # only use this new_message if there is no tool response from the user @@ -1285,6 +1300,26 @@ async def _run_trpc_in_background(self, # since toolset is now embedded in the agent's tools pass + def _is_graph_checkpoint_resume( + self, + input: RunAgentInput, + *, + session: Any = None, + ) -> bool: + """Whether this tool result must preserve a GraphAgent checkpoint. + + When *session* is supplied (e.g. from ``_ensure_session_exists``) the + check is performed without an extra database round-trip. Otherwise + the method falls back to fetching the session via the session service. + """ + if not self._is_tool_result_submission(input): + return False + state = getattr(session, "state", None) if session is not None else None + return isinstance(state, dict) and any( + key.startswith(_GRAPH_CHECKPOINT_STATE_KEY_PREFIX) + for key in state + ) + async def _execute_user_feedback_handler(self, tool_name: str, tool_message: str, thread_id: str, app_name: str, user_id: str) -> str: """Execute the user feedback handler if configured. diff --git a/trpc_agent_sdk/server/ag_ui/_core/_session_manager.py b/trpc_agent_sdk/server/ag_ui/_core/_session_manager.py index 848e7590..cce266d7 100644 --- a/trpc_agent_sdk/server/ag_ui/_core/_session_manager.py +++ b/trpc_agent_sdk/server/ag_ui/_core/_session_manager.py @@ -96,6 +96,15 @@ def __init__( session_timeout_seconds, cleanup_interval_seconds, max_sessions_per_user or 'unlimited', 'enabled' if memory_service else 'disabled') + @property + def session_service(self): + """Public accessor for the underlying session service. + + Provides clean access to the session service without reaching into + the private ``_session_service`` attribute. + """ + return self._session_service + @classmethod def get_instance(cls, **kwargs): """Get the singleton instance.""" diff --git a/trpc_agent_sdk/tools/__init__.py b/trpc_agent_sdk/tools/__init__.py index 715da600..d13d67e0 100644 --- a/trpc_agent_sdk/tools/__init__.py +++ b/trpc_agent_sdk/tools/__init__.py @@ -15,93 +15,86 @@ from trpc_agent_sdk.agents.sub_agent import DynamicSubAgentTool as DynamicSubAgentTool # noqa: F401 from trpc_agent_sdk.agents.sub_agent import SpawnSubAgentTool as SpawnSubAgentTool # noqa: F401 -from ._agent_tool import AGENT_TOOL_APP_NAME_SUFFIX -from ._agent_tool import AgentTool +from ._agent_tool import AGENT_TOOL_APP_NAME_SUFFIX, AgentTool from ._base_tool import BaseTool from ._constants import TOOL_NAME -from ._context_var import get_tool_var -from ._context_var import reset_tool_var -from ._context_var import set_tool_var +from ._context_var import get_tool_var, reset_tool_var, set_tool_var from ._default_toolset import DefaultToolSet from ._function_tool import FunctionTool -from ._load_memory_tool import LoadMemoryResponse -from ._load_memory_tool import LoadMemoryTool -from ._load_memory_tool import load_memory -from ._load_memory_tool import load_memory_tool -from ._long_running_tool import LongRunningFunctionTool -from ._preload_memory_tool import PreloadMemoryTool -from ._preload_memory_tool import preload_memory_tool -from ._registry import ToolRegistry -from ._registry import ToolSetRegistry -from ._registry import ToolType -from ._registry import get_tool -from ._registry import get_tool_set -from ._registry import register_tool -from ._registry import register_tool_set +from ._load_memory_tool import LoadMemoryResponse, LoadMemoryTool, load_memory, load_memory_tool +from ._long_running_tool import ( + FRAMEWORK_TOOL_ERROR_CODES, + LongRunningFunctionTool, + TOOL_ERROR_CODE_ARGUMENT_ERROR, + TOOL_ERROR_CODE_EXECUTION_ERROR, + TOOL_ERROR_CODE_NOT_FOUND, + is_tool_execution_error, +) +from ._preload_memory_tool import PreloadMemoryTool, preload_memory_tool +from ._registry import ToolRegistry, ToolSetRegistry, ToolType, get_tool, get_tool_set, register_tool, register_tool_set from ._set_model_response_tool import SetModelResponseTool from ._streaming_function_tool import StreamingFunctionTool from ._streaming_progress_tool import StreamingProgressTool -from ._todo_tool import DEFAULT_NUDGE_MESSAGE -from ._todo_tool import DEFAULT_STATE_KEY_PREFIX -from ._todo_tool import DEFAULT_TODO_DESCRIPTION -from ._todo_tool import DEFAULT_TODO_PROMPT -from ._todo_tool import TodoItem -from ._todo_tool import TodoStatus -from ._todo_tool import TodoWriteTool -from ._todo_tool import get_todos -from ._todo_tool import render_todos -from ._todo_tool import state_key -from ._todo_tool import validate_todos -from ._tool_adapter import convert_toolunion_to_tool_list -from ._tool_adapter import create_tool -from ._tool_adapter import create_toolset -from .task_tools import DEFAULT_TASK_PROMPT -from .task_tools import TaskCreateTool -from .task_tools import TaskGetTool -from .task_tools import TaskListSummary -from .task_tools import TaskListTool -from .task_tools import TaskRecord -from .task_tools import TaskStatus -from .task_tools import TaskStore -from .task_tools import TaskToolSet -from .task_tools import TaskUpdateTool -from .task_tools import get_task_store -from .task_tools import render_task_list -from .goal_tools import GoalOptions -from .goal_tools import GoalRecord -from .goal_tools import GoalStatus -from .goal_tools import GoalToolSet -from .goal_tools import OnRetry -from .goal_tools import RetryEvent -from .goal_tools import get_goal_record -from .goal_tools import setup_goal -from .goal_tools import render_goal +from ._todo_tool import ( + DEFAULT_NUDGE_MESSAGE, + DEFAULT_STATE_KEY_PREFIX, + DEFAULT_TODO_DESCRIPTION, + DEFAULT_TODO_PROMPT, + TodoItem, + TodoStatus, + TodoWriteTool, + get_todos, + render_todos, + state_key, + validate_todos, +) +from ._tool_adapter import convert_toolunion_to_tool_list, create_tool, create_toolset from ._transfer_to_agent_tool import transfer_to_agent -from ._webfetch_tool import FetchResult -from ._webfetch_tool import WebFetchTool -from ._websearch_tool import SearchHit -from ._websearch_tool import WebSearchResult -from ._websearch_tool import WebSearchTool -from .file_tools import BashTool -from .file_tools import EditTool -from .file_tools import FileToolSet -from .file_tools import GlobTool -from .file_tools import GrepTool -from .file_tools import ReadTool -from .file_tools import WriteTool -from .mcp_tool import MCPTool -from .mcp_tool import MCPToolset -from .mcp_tool import McpConnectionParamsType -from .mcp_tool import McpStdioServerParameters -from .mcp_tool import SseConnectionParams -from .mcp_tool import StdioConnectionParams -from .mcp_tool import StreamableHTTPConnectionParams -from .mcp_tool import patch_mcp_cancel_scope_exit_issue -from .utils import build_function_declaration -from .utils import from_function_with_options -from .utils import get_required_fields -from .utils import parse_schema_from_parameter -from .utils import register_checker +from ._webfetch_tool import FetchResult, WebFetchTool +from ._websearch_tool import SearchHit, WebSearchResult, WebSearchTool +from .file_tools import BashTool, EditTool, FileToolSet, GlobTool, GrepTool, ReadTool, WriteTool +from .goal_tools import ( + GoalOptions, + GoalRecord, + GoalStatus, + GoalToolSet, + OnRetry, + RetryEvent, + get_goal_record, + render_goal, + setup_goal, +) +from .mcp_tool import ( + McpConnectionParamsType, + McpStdioServerParameters, + MCPTool, + MCPToolset, + SseConnectionParams, + StdioConnectionParams, + StreamableHTTPConnectionParams, + patch_mcp_cancel_scope_exit_issue, +) +from .task_tools import ( + DEFAULT_TASK_PROMPT, + TaskCreateTool, + TaskGetTool, + TaskListSummary, + TaskListTool, + TaskRecord, + TaskStatus, + TaskStore, + TaskToolSet, + TaskUpdateTool, + get_task_store, + render_task_list, +) +from .utils import ( + build_function_declaration, + from_function_with_options, + get_required_fields, + parse_schema_from_parameter, + register_checker, +) __all__ = [ "ToolPredicate", @@ -119,6 +112,11 @@ "load_memory", "load_memory_tool", "LongRunningFunctionTool", + "is_tool_execution_error", + "TOOL_ERROR_CODE_NOT_FOUND", + "TOOL_ERROR_CODE_ARGUMENT_ERROR", + "TOOL_ERROR_CODE_EXECUTION_ERROR", + "FRAMEWORK_TOOL_ERROR_CODES", "PreloadMemoryTool", "preload_memory_tool", "ToolRegistry", @@ -207,6 +205,7 @@ def __getattr__(name): if name in _LAZY_REEXPORTS: import importlib + module_name, attr = _LAZY_REEXPORTS[name] obj = getattr(importlib.import_module(module_name), attr) globals()[name] = obj # cache: subsequent accesses skip __getattr__ diff --git a/trpc_agent_sdk/tools/_function_tool.py b/trpc_agent_sdk/tools/_function_tool.py index 0533f395..7e0fcb24 100644 --- a/trpc_agent_sdk/tools/_function_tool.py +++ b/trpc_agent_sdk/tools/_function_tool.py @@ -78,6 +78,15 @@ def my_function(param1: str, tool_context: InvocationContext): from .utils import get_mandatory_args +class ToolArgumentErrorResponse(dict): + """Private marker for framework-detected argument validation failures. + + It deliberately remains a dictionary so direct ``FunctionTool`` callers + keep the established ``{\"error\": ...}`` result contract. The tools + processor uses the type to attach an out-of-band Event error code. + """ + + class FunctionTool(BaseTool): """A tool that wraps a user-defined Python function. @@ -170,7 +179,7 @@ async def _run_async_impl(self, *, tool_context: InvocationContext, args: dict[s error_str = f"""Invoking `{self.name}()` failed as the following mandatory input parameters are not present: {missing_mandatory_args_str} You could retry calling this tool, but it is IMPORTANT for you to provide all the mandatory parameters.""" - return {'error': error_str} + return ToolArgumentErrorResponse(error=error_str) # Functions are callable objects, but not all callable objects are functions # checking coroutine function is not enough. We also need to check whether diff --git a/trpc_agent_sdk/tools/_long_running_tool.py b/trpc_agent_sdk/tools/_long_running_tool.py index 4e70954f..36d70332 100644 --- a/trpc_agent_sdk/tools/_long_running_tool.py +++ b/trpc_agent_sdk/tools/_long_running_tool.py @@ -7,15 +7,37 @@ from __future__ import annotations -from typing import Callable -from typing import Optional +from typing import Callable, Optional + from typing_extensions import override from trpc_agent_sdk.filter import BaseFilter +from trpc_agent_sdk.events import Event from trpc_agent_sdk.types import FunctionDeclaration from ._function_tool import FunctionTool +# Framework-level tool execution error codes shared between the tools +# processor (which emits them) and the long-running tool helper (which +# checks them). Keeping the literals in one place prevents silent drift +# when new codes are added. +TOOL_ERROR_CODE_NOT_FOUND = "tool_not_found" +TOOL_ERROR_CODE_ARGUMENT_ERROR = "tool_argument_error" +TOOL_ERROR_CODE_EXECUTION_ERROR = "tool_execution_error" + +FRAMEWORK_TOOL_ERROR_CODES = frozenset( + { + TOOL_ERROR_CODE_NOT_FOUND, + TOOL_ERROR_CODE_ARGUMENT_ERROR, + TOOL_ERROR_CODE_EXECUTION_ERROR, + } +) + + +def is_tool_execution_error(event: Event) -> bool: + """Return whether the framework, not a tool payload, reported a failure.""" + return event.error_code in FRAMEWORK_TOOL_ERROR_CODES + class LongRunningFunctionTool(FunctionTool): """A function tool that returns the result asynchronously. @@ -34,10 +56,9 @@ class LongRunningFunctionTool(FunctionTool): is_long_running: Whether the tool is a long running operation. """ - def __init__(self, - func: Callable, - filters_name: Optional[list[str]] = None, - filters: Optional[list[BaseFilter]] = None): + def __init__( + self, func: Callable, filters_name: Optional[list[str]] = None, filters: Optional[list[BaseFilter]] = None + ): """Initialize the long running function tool. Args: @@ -58,9 +79,11 @@ def _get_declaration(self) -> Optional[FunctionDeclaration]: """ declaration = super()._get_declaration() if declaration: - instruction = ("\n\nNOTE: This is a long-running operation. Do not call this tool" - " again if it has already returned some intermediate or pending" - " status.") + instruction = ( + "\n\nNOTE: This is a long-running operation. Do not call this tool" + " again if it has already returned some intermediate or pending" + " status." + ) if declaration.description: declaration.description += instruction else: