From 11da87991525de3cee26dc7378c6848b80d16af0 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 | 191 ++- tests/models/test_openai_responses_model.py | 1089 +++++++++++++++++ tests/server/ag_ui/_core/test_agui_agent.py | 305 +++-- tests/tools/test_function_tool.py | 21 +- tests/tools/test_long_running_tool.py | 61 +- .../graph/test_agent_node_hitl.py | 490 ++++++++ tests/trpc_agent_dsl/graph/test_constants.py | 16 +- tests/trpc_agent_dsl/graph/test_events.py | 20 + trpc_agent_sdk/agents/_llm_agent.py | 103 +- .../agents/core/_tools_processor.py | 17 +- trpc_agent_sdk/dsl/graph/_constants.py | 6 + trpc_agent_sdk/dsl/graph/_events/_builder.py | 9 +- trpc_agent_sdk/dsl/graph/_graph_agent.py | 10 +- .../dsl/graph/_node_action/_agent.py | 148 ++- trpc_agent_sdk/models/_openai_model.py | 592 ++++++++- .../server/ag_ui/_core/_agui_agent.py | 66 +- .../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 | 24 +- 25 files changed, 3050 insertions(+), 356 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..f550f1f5 100644 --- a/tests/agents/test_llm_agent_ext.py +++ b/tests/agents/test_llm_agent_ext.py @@ -23,13 +23,13 @@ from trpc_agent_sdk.sessions import InMemorySessionService from trpc_agent_sdk.types import Content, GenerateContentConfig, Part - # --------------------------------------------------------------------------- # Stubs / helpers # --------------------------------------------------------------------------- class MockLLMModel(LLMModel): + @classmethod def supported_models(cls) -> List[str]: return [r"test-llm-ext-.*"] @@ -58,9 +58,7 @@ def _agent(**kwargs): @pytest.fixture def invocation_context(): service = InMemorySessionService() - session = asyncio.run( - service.create_session(app_name="test", user_id="u1", session_id="s1") - ) + session = asyncio.run(service.create_session(app_name="test", user_id="u1", session_id="s1")) agent = _agent() ctx = InvocationContext( session_service=service, @@ -79,6 +77,7 @@ def invocation_context(): class TestToolsProcessorProperty: + def test_returns_processor_with_no_tools(self): """Returns ToolsProcessor even when agent has no tools.""" agent = _agent(tools=[]) @@ -105,6 +104,7 @@ def my_func(x: str) -> str: class TestGetExtendedToolsProcessor: + def test_no_transfer_when_no_sub_agents(self, invocation_context): """No transfer tool added when agent has no sub_agents.""" agent = _agent() @@ -139,6 +139,7 @@ def my_func(x: str) -> str: class TestShouldEnableAgentTransferExtended: + def test_transfer_to_peers_enabled(self): """Transfer enabled when agent has siblings.""" child1 = _agent(name="child1") @@ -158,7 +159,9 @@ def test_disallow_both_disables(self): def test_non_llm_parent_disables_parent_transfer(self): """Non-LlmAgent parent disables parent transfer.""" + class SimpleAgent(BaseAgent): + async def _run_async_impl(self, ctx): yield @@ -173,6 +176,7 @@ async def _run_async_impl(self, ctx): class TestResolveModelExtended: + def test_factory_receives_custom_data(self, invocation_context): """Factory callback receives custom_data from run_config.""" invocation_context.run_config.custom_data["key"] = "value" @@ -213,6 +217,7 @@ async def run(): class TestCreateErrorEventExtended: + def test_branch_is_included(self, invocation_context): """Error event includes the context's branch.""" invocation_context.branch = "some.branch" @@ -233,6 +238,7 @@ def test_author_is_agent_name(self, invocation_context): class TestSaveOutputToStateExtended: + def test_filters_thought_parts(self, invocation_context): """Thought parts are excluded from saved output.""" agent = _agent(output_key="result") @@ -266,6 +272,7 @@ def test_state_delta_set(self, invocation_context): class TestBranchFilterModeExtended: + def test_prefix_overrides_true_include_previous(self): """PREFIX mode takes precedence over include_previous_history=True.""" agent = _agent( @@ -281,6 +288,7 @@ def test_prefix_overrides_true_include_previous(self): class TestValidatorsExtended: + def test_response_schema_in_config_raises(self): """Response schema in generate_content_config raises.""" from pydantic import BaseModel as PydanticModel @@ -289,17 +297,11 @@ class TestSchema(PydanticModel): name: str with pytest.raises(ValueError, match="Response schema"): - _agent( - generate_content_config=GenerateContentConfig( - response_schema=TestSchema, - ), - ) + _agent(generate_content_config=GenerateContentConfig(response_schema=TestSchema, ), ) def test_valid_config_passes(self): """Valid generate_content_config passes validation.""" - agent = _agent( - generate_content_config=GenerateContentConfig(temperature=0.5), - ) + agent = _agent(generate_content_config=GenerateContentConfig(temperature=0.5), ) assert agent.generate_content_config.temperature == 0.5 def test_none_config_produces_default(self): @@ -314,6 +316,7 @@ def test_none_config_produces_default(self): class TestModelPostInit: + def test_string_model_resolved_on_init(self): """String model is resolved to LLMModel on init.""" agent = _agent(model="test-llm-ext-model") @@ -321,6 +324,7 @@ def test_string_model_resolved_on_init(self): def test_callable_model_not_resolved_on_init(self): """Callable model factory is not resolved on init.""" + async def factory(custom_data): return MockLLMModel(model_name="test-llm-ext-model") @@ -340,6 +344,7 @@ def test_model_instance_used_directly(self): class TestAgentFieldDefaults: + def test_default_include_contents(self): """Default include_contents is 'default'.""" agent = _agent() @@ -382,6 +387,7 @@ def test_max_history_messages_default(self): class TestRunAsyncImplErrorPath: + def test_yields_error_on_request_build_failure(self, invocation_context): """Error event yielded when request building fails.""" agent = _agent() @@ -396,9 +402,7 @@ def test_yields_error_on_request_build_failure(self, invocation_context): async def run(): events = [] - with patch( - "trpc_agent_sdk.agents._llm_agent.default_request_processor" - ) as mock_rp: + with patch("trpc_agent_sdk.agents._llm_agent.default_request_processor") as mock_rp: mock_rp.build_request = AsyncMock(return_value=mock_event) async for event in agent._run_async_impl(invocation_context): events.append(event) @@ -409,5 +413,162 @@ 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..c95d2459 --- /dev/null +++ b/tests/models/test_openai_responses_model.py @@ -0,0 +1,1089 @@ +# 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. + +import json +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"] + + def test_pure_reasoning_turn_preserves_responses_input_items(self): + """A turn with only thought parts (no text/image/tool calls) must + still produce a message carrying ``responses_input_items`` so that + reasoning signatures are not silently dropped from history. + """ + from trpc_agent_sdk.models._openai_model import _RESPONSES_INPUT_ITEMS + model = _model(use_responses_api=True) + + reasoning_item = { + "type": "reasoning", + "id": "rs_pure", + "encrypted_content": "encrypted-pure", + "summary": [{ + "type": "summary_text", + "text": "Pure thought" + }], + } + + # Assistant content with only a thought part — no text, no tool calls. + # thought_signature is set via attribute assignment (not construction) + # to bypass Part's base64 validation, matching _create_responses_response. + thought_part = Part(thought=True, text="") + thought_part.thought_signature = json.dumps(reasoning_item, ensure_ascii=False).encode("utf-8") + assistant_content = Content( + role="model", + parts=[thought_part], + ) + user_content = Content(role="user", parts=[Part.from_text(text="follow up")]) + + messages = model._format_messages(_request([assistant_content, user_content])) + + # The assistant message must be present and carry the reasoning items. + assistant_msgs = [m for m in messages if m.get("role") == "assistant"] + assert len(assistant_msgs) == 1 + assert _RESPONSES_INPUT_ITEMS in assistant_msgs[0] + assert assistant_msgs[0][_RESPONSES_INPUT_ITEMS] == [reasoning_item] + + @pytest.mark.asyncio + async def test_responses_maps_thinking_budget_to_reasoning_effort(self): + """The Responses path should map thinking_budget to reasoning.effort + so that reasoning depth can be controlled.""" + from trpc_agent_sdk.types import ThinkingConfig + + model = _model(use_responses_api=True) + request = _request( + [Content(parts=[Part.from_text(text="Think hard")], role="user")], + GenerateContentConfig( + max_output_tokens=4096, + thinking_config=ThinkingConfig( + include_thoughts=True, + thinking_budget=20000, + ), + ), + ) + response = Mock() + response.model_dump.return_value = { + "id": "resp_effort", + "status": "completed", + "output": [ + { + "type": "message", + "role": "assistant", + "content": [{ + "type": "output_text", + "text": "Result" + }], + }, + ], + "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 + [item async for item in model.generate_async(request, stream=False)] + + reasoning = captured.get("reasoning") + assert reasoning is not None + assert reasoning["summary"] == "auto" + assert reasoning["effort"] == "high" + + @pytest.mark.asyncio + async def test_responses_thinking_budget_negative_one_skips_effort(self): + """thinking_budget=-1 (automatic) should not set reasoning.effort, + letting the model decide.""" + from trpc_agent_sdk.types import ThinkingConfig + + model = _model(use_responses_api=True) + request = _request( + [Content(parts=[Part.from_text(text="Think")], role="user")], + GenerateContentConfig( + max_output_tokens=4096, + thinking_config=ThinkingConfig( + include_thoughts=True, + thinking_budget=-1, + ), + ), + ) + response = Mock() + response.model_dump.return_value = { + "id": "resp_auto", + "status": "completed", + "output": [ + { + "type": "message", + "role": "assistant", + "content": [{ + "type": "output_text", + "text": "OK" + }], + }, + ], + "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 + [item async for item in model.generate_async(request, stream=False)] + + reasoning = captured.get("reasoning") + assert reasoning is not None + assert reasoning["summary"] == "auto" + assert "effort" not in reasoning + + @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..269926d2 100644 --- a/tests/server/ag_ui/_core/test_agui_agent.py +++ b/tests/server/ag_ui/_core/test_agui_agent.py @@ -27,7 +27,6 @@ from trpc_agent_sdk.server.ag_ui._core._feed_back_content import AgUiUserFeedBack from trpc_agent_sdk.server.ag_ui._core._session_manager import SessionManager - # --------------------------------------------------------------------------- # Fixtures # --------------------------------------------------------------------------- @@ -116,6 +115,7 @@ def _make_tool_call(tc_id="tc-1", name="my_tool"): class TestAgUiAgentInit: + def test_init_with_required_params_only(self, mock_agent): agent = AgUiAgent(trpc_agent=mock_agent, auto_cleanup=False) assert agent._trpc_agent is mock_agent @@ -173,6 +173,7 @@ def test_init_stores_timeout_values(self, mock_agent): class TestGetAppName: + def test_returns_static_app_name(self, mock_agent): agent = AgUiAgent(trpc_agent=mock_agent, app_name="my_app", auto_cleanup=False) result = agent.get_app_name(_make_input()) @@ -202,6 +203,7 @@ def test_defaults_to_agent_name(self, mock_agent): class TestGetUserId: + def test_returns_static_user_id(self, mock_agent): agent = AgUiAgent(trpc_agent=mock_agent, user_id="uid-42", auto_cleanup=False) result = agent.get_user_id(_make_input()) @@ -231,6 +233,7 @@ def test_defaults_to_thread_user_prefix(self, mock_agent): class TestIsToolResultSubmission: + def test_true_when_last_message_is_tool(self, agui_agent): inp = _make_input(messages=[_make_user_message(), _make_tool_message()]) assert agui_agent._is_tool_result_submission(inp) is True @@ -250,6 +253,7 @@ def test_false_when_last_message_is_user(self, agui_agent): class TestExtractToolResults: + async def test_extracts_most_recent_tool_message(self, agui_agent): tool_msg = _make_tool_message(content='{"result": "ok"}', tool_call_id="tc-5") inp = _make_input(messages=[_make_user_message(), tool_msg]) @@ -289,6 +293,7 @@ async def test_returns_empty_list_when_messages_empty(self, agui_agent): class TestConvertLatestMessage: + @patch("trpc_agent_sdk.server.ag_ui._core._agui_agent.convert_message_content_to_parts") async def test_converts_user_message(self, mock_convert, agui_agent): mock_convert.return_value = [_make_text_part("hello")] @@ -338,6 +343,7 @@ async def test_returns_none_when_convert_returns_empty(self, mock_convert, agui_ class TestExtractLongRunningToolNames: + def test_extracts_from_long_running_function_tool(self, agui_agent): lrt = Mock(spec=LongRunningFunctionTool) lrt.name = "slow_tool" @@ -421,6 +427,7 @@ async def test_expands_nested_toolset_long_running_tools(self, agui_agent): class TestResolveToolNameFromSession: + @pytest.mark.asyncio async def test_resolves_tool_name_from_session_events(self, agui_agent): from trpc_agent_sdk import types @@ -437,12 +444,11 @@ async def test_resolves_tool_name_from_session_events(self, agui_agent): content=types.Content( role="model", parts=[ - types.Part( - function_call=types.FunctionCall( - id="call_abc", - name="ask_user_question", - args={"question": "Pick one"}, - )) + types.Part(function_call=types.FunctionCall( + id="call_abc", + name="ask_user_question", + args={"question": "Pick one"}, + )) ], ), ) @@ -473,12 +479,11 @@ async def test_extract_tool_results_falls_back_to_session(self, agui_agent): content=types.Content( role="model", parts=[ - types.Part( - function_call=types.FunctionCall( - id="call_xyz", - name="ask_user_question", - args={"question": "Pick one"}, - )) + types.Part(function_call=types.FunctionCall( + id="call_xyz", + name="ask_user_question", + args={"question": "Pick one"}, + )) ], ), ) @@ -517,6 +522,7 @@ async def test_extract_tool_results_falls_back_to_session(self, agui_agent): class TestDefaultRunConfig: + def test_returns_run_config_with_streaming(self, agui_agent): inp = _make_input() config = agui_agent._default_run_config(inp) @@ -531,6 +537,7 @@ def test_returns_run_config_with_streaming(self, agui_agent): class TestGetSessionMetadata: + def test_returns_from_cache(self, agui_agent): agui_agent._session_lookup_cache["sess-1"] = { "app_name": "cached_app", @@ -560,9 +567,7 @@ def test_returns_none_when_not_found(self, agui_agent): assert result is None def test_returns_none_on_exception(self, agui_agent): - agui_agent._session_manager._user_sessions = Mock( - items=Mock(side_effect=RuntimeError("boom")) - ) + agui_agent._session_manager._user_sessions = Mock(items=Mock(side_effect=RuntimeError("boom"))) result = agui_agent._get_session_metadata("sess-err") assert result is None @@ -573,6 +578,7 @@ def test_returns_none_on_exception(self, agui_agent): class TestCancelRun: + @patch("trpc_agent_sdk.server.ag_ui._core._agui_agent.cancel") async def test_cancels_active_run(self, mock_cancel_module, agui_agent): cleanup_event = asyncio.Event() @@ -653,15 +659,19 @@ async def test_handles_timeout_during_cancel_wait(self, mock_cancel_module, agui class TestExecuteUserFeedbackHandler: + async def test_no_handler_returns_original_message(self, agui_agent): agui_agent._user_feedback_handler = None - result = await agui_agent._execute_user_feedback_handler( - tool_name="t", tool_message="original", thread_id="th", app_name="app", user_id="u" - ) + result = await agui_agent._execute_user_feedback_handler(tool_name="t", + tool_message="original", + thread_id="th", + app_name="app", + user_id="u") assert result == "original" async def test_handler_modifies_tool_message(self, mock_agent): + async def handler(feedback: AgUiUserFeedBack): feedback.tool_message = "modified" @@ -676,12 +686,15 @@ async def handler(feedback: AgUiUserFeedBack): real_session = Session(id="th-1", app_name="test_app", user_id="test_user", save_key="k", state={}) agent._session_manager._session_service.get_session = AsyncMock(return_value=real_session) - result = await agent._execute_user_feedback_handler( - tool_name="tool_x", tool_message="original", thread_id="th-1", app_name="test_app", user_id="test_user" - ) + result = await agent._execute_user_feedback_handler(tool_name="tool_x", + tool_message="original", + thread_id="th-1", + app_name="test_app", + user_id="test_user") assert result == "modified" async def test_handler_marks_session_modified_triggers_update(self, mock_agent): + async def handler(feedback: AgUiUserFeedBack): feedback.mark_session_modified() @@ -697,13 +710,16 @@ async def handler(feedback: AgUiUserFeedBack): agent._session_manager._session_service.get_session = AsyncMock(return_value=real_session) agent._session_manager._session_service.update_session = AsyncMock() - await agent._execute_user_feedback_handler( - tool_name="t", tool_message="msg", thread_id="th-1", app_name="test_app", user_id="test_user" - ) + await agent._execute_user_feedback_handler(tool_name="t", + tool_message="msg", + thread_id="th-1", + app_name="test_app", + user_id="test_user") agent._session_manager._session_service.update_session.assert_awaited_once_with(real_session) async def test_handler_exception_returns_original_message(self, mock_agent): + async def handler(feedback: AgUiUserFeedBack): raise RuntimeError("handler broke") @@ -718,12 +734,15 @@ async def handler(feedback: AgUiUserFeedBack): real_session = Session(id="th-1", app_name="test_app", user_id="test_user", save_key="k", state={}) agent._session_manager._session_service.get_session = AsyncMock(return_value=real_session) - result = await agent._execute_user_feedback_handler( - tool_name="t", tool_message="original", thread_id="th-1", app_name="test_app", user_id="test_user" - ) + result = await agent._execute_user_feedback_handler(tool_name="t", + tool_message="original", + thread_id="th-1", + app_name="test_app", + user_id="test_user") assert result == "original" async def test_session_not_found_returns_original_message(self, mock_agent): + async def handler(feedback: AgUiUserFeedBack): feedback.tool_message = "should not reach" @@ -737,9 +756,11 @@ async def handler(feedback: AgUiUserFeedBack): agent._session_manager._session_service.get_session = AsyncMock(return_value=None) - result = await agent._execute_user_feedback_handler( - tool_name="t", tool_message="original", thread_id="th-1", app_name="test_app", user_id="test_user" - ) + result = await agent._execute_user_feedback_handler(tool_name="t", + tool_message="original", + thread_id="th-1", + app_name="test_app", + user_id="test_user") assert result == "original" @@ -749,6 +770,7 @@ async def handler(feedback: AgUiUserFeedBack): class TestCleanupStaleExecutions: + async def test_removes_stale_executions(self, agui_agent): stale_exec = AsyncMock(spec=ExecutionState) stale_exec.is_stale = Mock(return_value=True) @@ -795,6 +817,7 @@ async def test_mixed_stale_and_fresh(self, agui_agent): class TestClose: + async def test_cancels_all_active_executions(self, agui_agent): exec_1 = AsyncMock(spec=ExecutionState) exec_1.cancel = AsyncMock() @@ -834,6 +857,7 @@ async def test_clears_session_lookup_cache(self, agui_agent): class TestEnsureSessionExists: + async def test_creates_session_and_populates_cache(self, agui_agent): mock_session = Mock() agui_agent._session_manager.get_or_create_session = AsyncMock(return_value=mock_session) @@ -850,9 +874,7 @@ async def test_creates_session_and_populates_cache(self, agui_agent): assert agui_agent._session_lookup_cache["sess-1"] == {"app_name": "app", "user_id": "user"} async def test_propagates_exception(self, agui_agent): - agui_agent._session_manager.get_or_create_session = AsyncMock( - side_effect=RuntimeError("db error") - ) + agui_agent._session_manager.get_or_create_session = AsyncMock(side_effect=RuntimeError("db error")) with pytest.raises(RuntimeError, match="db error"): await agui_agent._ensure_session_exists("app", "user", "sess-err", {}) @@ -864,6 +886,7 @@ async def test_propagates_exception(self, agui_agent): class TestCreateRunner: + @patch("trpc_agent_sdk.server.ag_ui._core._agui_agent.Runner") def test_creates_runner_with_correct_params(self, mock_runner_cls, agui_agent, mock_agent): agui_agent._create_runner(mock_agent, "user-1", "app-1") @@ -882,6 +905,7 @@ def test_creates_runner_with_correct_params(self, mock_runner_cls, agui_agent, m class TestDefaultAppExtractor: + def test_returns_agent_name(self, mock_agent): mock_agent.name = "my_special_agent" agent = AgUiAgent(trpc_agent=mock_agent, auto_cleanup=False) @@ -902,6 +926,7 @@ def test_returns_fallback_on_exception(self, mock_agent): class TestAddPendingToolCallWithContext: + async def test_adds_tool_call_to_pending_list(self, agui_agent): agui_agent._session_manager.get_state_value = AsyncMock(return_value=[]) agui_agent._session_manager.set_state_value = AsyncMock(return_value=True) @@ -909,8 +934,11 @@ async def test_adds_tool_call_to_pending_list(self, agui_agent): await agui_agent._add_pending_tool_call_with_context("sess-1", "tc-1", "app", "user") agui_agent._session_manager.set_state_value.assert_awaited_once_with( - session_id="sess-1", app_name="app", user_id="user", - key="pending_tool_calls", value=["tc-1"], + session_id="sess-1", + app_name="app", + user_id="user", + key="pending_tool_calls", + value=["tc-1"], ) async def test_does_not_add_duplicate(self, agui_agent): @@ -934,6 +962,7 @@ async def test_handles_exception(self, agui_agent): class TestRemovePendingToolCall: + async def test_removes_tool_call(self, agui_agent): agui_agent._session_lookup_cache["sess-1"] = {"app_name": "app", "user_id": "user"} agui_agent._session_manager.get_state_value = AsyncMock(return_value=["tc-1", "tc-2"]) @@ -942,8 +971,11 @@ async def test_removes_tool_call(self, agui_agent): await agui_agent._remove_pending_tool_call("sess-1", "tc-1") agui_agent._session_manager.set_state_value.assert_awaited_once_with( - session_id="sess-1", app_name="app", user_id="user", - key="pending_tool_calls", value=["tc-2"], + session_id="sess-1", + app_name="app", + user_id="user", + key="pending_tool_calls", + value=["tc-2"], ) async def test_no_metadata_found(self, agui_agent): @@ -975,6 +1007,7 @@ async def test_handles_exception(self, agui_agent): class TestHasPendingToolCalls: + async def test_returns_true_when_pending(self, agui_agent): agui_agent._session_lookup_cache["sess-1"] = {"app_name": "app", "user_id": "user"} agui_agent._session_manager.get_state_value = AsyncMock(return_value=["tc-1"]) @@ -1010,6 +1043,7 @@ async def test_returns_false_on_exception(self, agui_agent): class TestRun: + async def test_delegates_to_tool_result_submission(self, agui_agent): tool_msg = _make_tool_message() inp = _make_input(messages=[_make_user_message(), tool_msg]) @@ -1052,6 +1086,7 @@ async def fake_execution(*args, **kwargs): class TestHandleToolResultSubmission: + async def test_no_tool_results_yields_error(self, agui_agent): inp = _make_input(messages=[_make_user_message()]) @@ -1111,6 +1146,7 @@ async def test_exception_yields_error_event(self, agui_agent): class TestStreamEvents: + async def test_streams_events_until_none(self, agui_agent): from ag_ui.core import TextMessageStartEvent queue = asyncio.Queue() @@ -1178,18 +1214,21 @@ async def quick_task(): class TestIsHitlTextScenario: + async def test_detects_hitl_pattern(self, agui_agent): from trpc_agent_sdk import types from trpc_agent_sdk.events import Event func_call = types.FunctionCall(id="fc-1", name="ask_user", args={}) second_last = Event( - invocation_id="inv-1", author="agent", + invocation_id="inv-1", + author="agent", content=types.Content(role="model", parts=[types.Part(function_call=func_call)]), ) func_resp = types.FunctionResponse(id="fc-1", name="ask_user", response={"text": "hello"}) last = Event( - invocation_id="inv-2", author="user", + invocation_id="inv-2", + author="user", content=types.Content(role="function", parts=[types.Part(function_response=func_resp)]), ) @@ -1221,12 +1260,14 @@ async def test_returns_none_when_ids_dont_match(self, agui_agent): func_call = types.FunctionCall(id="fc-1", name="ask_user", args={}) second_last = Event( - invocation_id="inv-1", author="agent", + invocation_id="inv-1", + author="agent", content=types.Content(role="model", parts=[types.Part(function_call=func_call)]), ) func_resp = types.FunctionResponse(id="fc-DIFFERENT", name="ask_user", response={"text": "hello"}) last = Event( - invocation_id="inv-2", author="user", + invocation_id="inv-2", + author="user", content=types.Content(role="function", parts=[types.Part(function_response=func_resp)]), ) @@ -1238,8 +1279,7 @@ async def test_returns_none_when_ids_dont_match(self, agui_agent): assert result is None async def test_returns_none_on_exception(self, agui_agent): - agui_agent._session_manager._session_service.get_session = AsyncMock( - side_effect=RuntimeError("fail")) + agui_agent._session_manager._session_service.get_session = AsyncMock(side_effect=RuntimeError("fail")) result = await agui_agent._is_hitl_text_scenario("t1", "app", "user") assert result is None @@ -1248,12 +1288,14 @@ async def test_returns_none_when_no_function_call_in_second_last(self, agui_agen from trpc_agent_sdk.events import Event second_last = Event( - invocation_id="inv-1", author="agent", + invocation_id="inv-1", + author="agent", content=types.Content(role="model", parts=[types.Part(text="hello")]), ) func_resp = types.FunctionResponse(id="fc-1", name="ask_user", response={"text": "hi"}) last = Event( - invocation_id="inv-2", author="user", + invocation_id="inv-2", + author="user", content=types.Content(role="function", parts=[types.Part(function_response=func_resp)]), ) @@ -1271,6 +1313,7 @@ async def test_returns_none_when_no_function_call_in_second_last(self, agui_agen class TestStartNewExecution: + async def test_emits_run_started_and_run_finished(self, agui_agent): from ag_ui.core import RunStartedEvent, RunFinishedEvent @@ -1310,6 +1353,7 @@ async def test_max_concurrent_executions_error(self, agui_agent): assert any(isinstance(e, RunErrorEvent) for e in events) async def test_cleans_up_execution_on_completion(self, agui_agent): + async def fake_bg_execution(input, http_request=None): queue = asyncio.Queue() await queue.put(None) @@ -1338,6 +1382,7 @@ async def noop(): class TestStartBackgroundExecution: + async def test_returns_execution_state(self, agui_agent): from ag_ui.core import SystemMessage as AGUISystemMessage @@ -1357,7 +1402,7 @@ async def empty_run(*args, **kwargs): mock_runner.run_async = empty_run MockRunner.return_value = mock_runner with patch("trpc_agent_sdk.server.ag_ui._core._agui_agent.convert_message_content_to_parts", - return_value=[_make_text_part("hello")]): + return_value=[_make_text_part("hello")]): exec_state = await agui_agent._start_background_execution(inp) assert isinstance(exec_state, ExecutionState) @@ -1394,7 +1439,7 @@ async def empty_run(*args, **kwargs): mock_runner.run_async = empty_run MockRunner.return_value = mock_runner with patch("trpc_agent_sdk.server.ag_ui._core._agui_agent.convert_message_content_to_parts", - return_value=[_make_text_part("hello")]): + return_value=[_make_text_part("hello")]): exec_state = await agui_agent._start_background_execution(inp) # model_copy should have been called with tools update @@ -1432,7 +1477,7 @@ async def empty_run(*args, **kwargs): mock_runner.run_async = empty_run MockRunner.return_value = mock_runner with patch("trpc_agent_sdk.server.ag_ui._core._agui_agent.convert_message_content_to_parts", - return_value=[_make_text_part("hello")]): + return_value=[_make_text_part("hello")]): # Patch isinstance to detect our mock as SystemMessage with patch("trpc_agent_sdk.server.ag_ui._core._agui_agent.SystemMessage", type(sys_msg)): exec_state = await agui_agent._start_background_execution(inp) @@ -1450,6 +1495,7 @@ async def empty_run(*args, **kwargs): class TestRunTrpcInBackground: + async def test_runs_agent_and_puts_events_in_queue(self, agui_agent, mock_agent): from trpc_agent_sdk.events import Event from trpc_agent_sdk import types @@ -1478,10 +1524,13 @@ async def mock_run_async(**kwargs): mock_runner.run_async = mock_run_async MockRunner.return_value = mock_runner with patch("trpc_agent_sdk.server.ag_ui._core._agui_agent.convert_message_content_to_parts", - return_value=[_make_text_part("hello")]): + return_value=[_make_text_part("hello")]): await agui_agent._run_trpc_in_background( - input=inp, agent=mock_agent, user_id="test_user", - app_name="test_app", event_queue=queue, + input=inp, + agent=mock_agent, + user_id="test_user", + app_name="test_app", + event_queue=queue, ) # Should have put events + None sentinel @@ -1490,12 +1539,17 @@ 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() - agui_agent._session_manager.get_or_create_session = AsyncMock( - side_effect=RuntimeError("session error")) + agui_agent._session_manager.get_or_create_session = AsyncMock(side_effect=RuntimeError("session error")) agui_agent._session_manager.update_session_state = AsyncMock() inp = _make_input(messages=[_make_user_message("hello")]) @@ -1503,10 +1557,13 @@ async def test_handles_error_and_puts_error_event(self, agui_agent, mock_agent): with patch("trpc_agent_sdk.server.ag_ui._core._agui_agent.Runner") as MockRunner: MockRunner.return_value = Mock() with patch("trpc_agent_sdk.server.ag_ui._core._agui_agent.convert_message_content_to_parts", - return_value=[_make_text_part("hello")]): + return_value=[_make_text_part("hello")]): await agui_agent._run_trpc_in_background( - input=inp, agent=mock_agent, user_id="test_user", - app_name="test_app", event_queue=queue, + input=inp, + agent=mock_agent, + user_id="test_user", + app_name="test_app", + event_queue=queue, ) events = [] @@ -1528,8 +1585,11 @@ async def test_no_message_yields_error(self, agui_agent, mock_agent): with patch("trpc_agent_sdk.server.ag_ui._core._agui_agent.Runner") as MockRunner: MockRunner.return_value = Mock() await agui_agent._run_trpc_in_background( - input=inp, agent=mock_agent, user_id="test_user", - app_name="test_app", event_queue=queue, + input=inp, + agent=mock_agent, + user_id="test_user", + app_name="test_app", + event_queue=queue, ) events = [] @@ -1546,9 +1606,11 @@ async def test_handles_tool_result_submission_in_background(self, agui_agent, mo queue = asyncio.Queue() trpc_event = Event( - invocation_id="inv-1", author="agent", + invocation_id="inv-1", + author="agent", content=types.Content(role="model", parts=[types.Part(text="done")]), - partial=False, timestamp=1000.0, + partial=False, + timestamp=1000.0, ) agui_agent._session_manager.get_or_create_session = AsyncMock(return_value=Mock()) @@ -1569,16 +1631,107 @@ async def mock_run_async(**kwargs): 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, + input=inp, + agent=mock_agent, + user_id="test_user", + app_name="test_app", + event_queue=queue, ) events = [] 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): + async def test_preserves_graph_state_on_pending_interrupt_without_checkpoint(self, agui_agent, mock_agent): + """When auto_persist=False, the first interrupt does not persist + checkpoint keys to session.state — only ``_trpc_graph_pending_interrupt`` + is set. The resume detection must still skip state sync. + """ + 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, + ) + # Session has pending interrupt but NO checkpoint keys — this is the + # common case for the first interrupt when auto_persist=False. + agui_agent._session_manager.get_or_create_session = AsyncMock( + return_value=Mock(state={"_trpc_graph_pending_interrupt": True}), ) + 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() from trpc_agent_sdk.events import LongRunningEvent from trpc_agent_sdk import types @@ -1602,10 +1755,13 @@ async def mock_run_async(**kwargs): mock_runner.run_async = mock_run_async MockRunner.return_value = mock_runner with patch("trpc_agent_sdk.server.ag_ui._core._agui_agent.convert_message_content_to_parts", - return_value=[_make_text_part("hello")]): + return_value=[_make_text_part("hello")]): await agui_agent._run_trpc_in_background( - input=inp, agent=mock_agent, user_id="test_user", - app_name="test_app", event_queue=queue, + input=inp, + agent=mock_agent, + user_id="test_user", + app_name="test_app", + event_queue=queue, ) events = [] @@ -1621,9 +1777,11 @@ async def test_hitl_text_scenario_converts_to_function_response(self, agui_agent queue = asyncio.Queue() trpc_event = Event( - invocation_id="inv-1", author="agent", + invocation_id="inv-1", + author="agent", content=types.Content(role="model", parts=[types.Part(text="result")]), - partial=False, timestamp=1000.0, + partial=False, + timestamp=1000.0, ) func_call_obj = types.FunctionCall(id="fc-1", name="ask_user", args={}) @@ -1646,10 +1804,13 @@ async def mock_run_async(**kwargs): mock_runner.run_async = mock_run_async MockRunner.return_value = mock_runner with patch("trpc_agent_sdk.server.ag_ui._core._agui_agent.convert_message_content_to_parts", - return_value=[_make_text_part("my answer")]): + return_value=[_make_text_part("my answer")]): await agui_agent._run_trpc_in_background( - input=inp, agent=mock_agent, user_id="test_user", - app_name="test_app", event_queue=queue, + input=inp, + agent=mock_agent, + user_id="test_user", + app_name="test_app", + event_queue=queue, ) events = [] diff --git a/tests/tools/test_function_tool.py b/tests/tools/test_function_tool.py index c550ba98..d67402f8 100644 --- a/tests/tools/test_function_tool.py +++ b/tests/tools/test_function_tool.py @@ -6,19 +6,18 @@ 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 --- + def sync_func(param1: str, param2: int = 10) -> str: """A sync test function.""" return f"{param1}-{param2}" @@ -128,7 +127,10 @@ async def test_run_sync_function(self, mock_context): tool = FunctionTool(sync_func) result = await tool._run_async_impl( tool_context=mock_context, - args={"param1": "hello", "param2": 5}, + args={ + "param1": "hello", + "param2": 5 + }, ) assert result == "hello-5" @@ -158,11 +160,13 @@ 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"] @pytest.mark.asyncio async def test_returns_empty_dict_for_none(self, mock_context): + def returns_none(): return None @@ -178,7 +182,10 @@ async def test_pydantic_model_return(self, mock_context): tool = FunctionTool(func_returns_model) result = await tool._run_async_impl( tool_context=mock_context, - args={"name": "test", "value": 42}, + args={ + "name": "test", + "value": 42 + }, ) assert isinstance(result, str) assert "test" in result @@ -200,7 +207,9 @@ def slow_func(x: str) -> str: @pytest.mark.asyncio async def test_async_callable_object(self, mock_context): + class AsyncCallable: + async def __call__(self, value: str) -> str: """Async callable.""" return f"async-{value}" 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..97ae3503 --- /dev/null +++ b/tests/trpc_agent_dsl/graph/test_agent_node_hitl.py @@ -0,0 +1,490 @@ +# 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..8e8be7ee 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 @@ -169,10 +172,15 @@ def test_is_unsafe_state_key_returns_false_for_safe_keys(self): def test_serializable_keys_are_not_unsafe(self): """Keys that represent user-visible data should never appear in UNSAFE_STATE_KEYS.""" serializable = { - STATE_KEY_USER_INPUT, STATE_KEY_MESSAGES, STATE_KEY_LAST_RESPONSE, - STATE_KEY_LAST_RESPONSE_ID, STATE_KEY_LAST_TOOL_RESPONSE, - STATE_KEY_NODE_RESPONSES, STATE_KEY_ONE_SHOT_MESSAGES, - STATE_KEY_ONE_SHOT_MESSAGES_BY_NODE, STATE_KEY_METADATA, + STATE_KEY_USER_INPUT, + STATE_KEY_MESSAGES, + STATE_KEY_LAST_RESPONSE, + STATE_KEY_LAST_RESPONSE_ID, + STATE_KEY_LAST_TOOL_RESPONSE, + STATE_KEY_NODE_RESPONSES, + STATE_KEY_ONE_SHOT_MESSAGES, + STATE_KEY_ONE_SHOT_MESSAGES_BY_NODE, + STATE_KEY_METADATA, STATE_KEY_STEP_NUMBER, } assert serializable.isdisjoint(UNSAFE_STATE_KEYS) 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..959aebba 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: @@ -555,7 +546,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 +558,12 @@ 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 +575,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 +584,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 +612,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 +708,33 @@ 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: 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..074205d6 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,9 @@ 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 +461,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 +612,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..95fa6c25 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__" @@ -179,6 +184,7 @@ STATE_KEY_PENDING_INTERRUPT_ID, STATE_KEY_PENDING_INTERRUPT_AUTHOR, STATE_KEY_PENDING_INTERRUPT_BRANCH, + STATE_KEY_PENDING_AGENT_NODE_HITL, }) 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..3d61036d 100644 --- a/trpc_agent_sdk/dsl/graph/_graph_agent.py +++ b/trpc_agent_sdk/dsl/graph/_graph_agent.py @@ -293,7 +293,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 @@ -356,7 +356,7 @@ def _create_interrupt_events( if isinstance(interrupt.value, dict): interrupt_response = interrupt.value else: - interrupt_response = {"desicion": interrupt.value} + interrupt_response = {"decision": interrupt.value} function_response = FunctionResponse( id=function_call.id, @@ -390,7 +390,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} diff --git a/trpc_agent_sdk/dsl/graph/_node_action/_agent.py b/trpc_agent_sdk/dsl/graph/_node_action/_agent.py index 85d064dc..0f2d6553 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 @@ -123,16 +128,38 @@ async def execute(self, state: State) -> dict[str, Any]: 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 +191,36 @@ 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: @@ -290,10 +349,18 @@ 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] @@ -308,8 +375,75 @@ def _build_child_events( 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..5abaea1e 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 @@ -448,6 +481,16 @@ 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) + elif responses_input_items: + # Pure-reasoning turn: the assistant emitted only thought + # parts with no text, image, or tool call. We must still + # emit a message so that the collected Responses reasoning + # items are not silently dropped from conversation history. + message = {const.ROLE: role, const.CONTENT: "", _RESPONSES_INPUT_ITEMS: responses_input_items} formatted_messages.append(message) # Validate and fix message sequence for OpenAI compatibility @@ -503,8 +546,10 @@ 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, @@ -533,8 +578,11 @@ 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, @@ -729,14 +777,15 @@ 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,8 +898,12 @@ 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, @@ -1051,7 +1104,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 +1156,240 @@ 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 = {} @@ -1428,7 +1711,8 @@ 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, @@ -1481,11 +1765,11 @@ 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 @@ -1528,21 +1812,217 @@ 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") + # Map thinking_budget to reasoning.effort so the Responses + # API can also control reasoning depth. -1 means "automatic" + # and is left to the model; positive budgets are bucketed + # into low / medium / high. + budget = request.config.thinking_config.thinking_budget + if isinstance(budget, int) and budget > 0: + if budget <= 5000: + reasoning.setdefault("effort", "low") + elif budget <= 15000: + reasoning.setdefault("effort", "medium") + else: + reasoning.setdefault("effort", "high") + 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 +2035,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 @@ -1649,10 +2129,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 +2159,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) @@ -1696,10 +2180,12 @@ async def _generate_stream(self, 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 +2194,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 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..44fe129b 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,21 @@ 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" + +# Pending interrupt marker key set in Session.state when a GraphAgent +# interrupts for HITL (see trpc_agent_sdk/dsl/graph/_constants.py). +# When auto_persist=False (the runner default), checkpoint data is only +# persisted via the completion event's state_delta — which is NOT emitted +# on interrupt. The pending-interrupt key, however, IS persisted via the +# interrupt bridge event's state_delta. Checking this key as a fallback +# ensures we detect graph resumes even when checkpoint keys are absent. +_GRAPH_PENDING_INTERRUPT_STATE_KEY = "_trpc_graph_pending_interrupt" + class AgUiAgent: """Middleware to bridge AG-UI Protocol with TRPC agents. @@ -1100,13 +1115,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 +1309,34 @@ 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. + + Two signals are checked: + 1. ``_trpc_graph_checkpoint*`` keys — present when a previous + completion persisted checkpoint data to session.state. + 2. ``_trpc_graph_pending_interrupt`` — set to True by the interrupt + bridge event and always persisted, even when no completion event + was emitted (the common case for the first and intermediate + interrupts when ``auto_persist=False``). + """ + if not self._is_tool_result_submission(input): + return False + state = getattr(session, "state", None) if session is not None else None + if not isinstance(state, dict): + return False + has_checkpoint = any(key.startswith(_GRAPH_CHECKPOINT_STATE_KEY_PREFIX) for key in state) + has_pending_interrupt = state.get(_GRAPH_PENDING_INTERRUPT_STATE_KEY) is True + return has_checkpoint or has_pending_interrupt + 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..9d246e06 100644 --- a/trpc_agent_sdk/tools/_long_running_tool.py +++ b/trpc_agent_sdk/tools/_long_running_tool.py @@ -7,15 +7,35 @@ 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.