diff --git a/docs/mkdocs/en/llm_agent.md b/docs/mkdocs/en/llm_agent.md index bcfff1892..0bc41d9b0 100644 --- a/docs/mkdocs/en/llm_agent.md +++ b/docs/mkdocs/en/llm_agent.md @@ -384,6 +384,102 @@ for query in demo_queries: ## Advanced Configuration and Control +### Limiting Work per Agent Invocation + +Use `RunConfig` to limit LLM calls, loop iterations, and tool calls for each Agent invocation. These limits prevent an unexpected execution path from consuming resources indefinitely. + +| Setting | Default | Meaning | +| --- | ---: | --- | +| `max_llm_calls` | `500` | Number of LLM calls that may be executed | +| `max_iterations` | `0` | Number of Agent loop iterations that may be executed; an iteration normally contains one LLM call and any tool execution requested by that call | +| `max_tool_calls` | `0` | Total number of tool calls that may be executed | + +A value of `0` disables the corresponding limit. A limit allows the configured number of operations and raises `RunLimitException` when the Agent attempts the next operation. For example, `max_iterations=1` allows the first iteration to complete and raises before the second iteration starts, so the exception contains `configured_value=1` and `observed_value=2`. + +`max_tool_calls` accumulates the number of tool calls returned by the LLM. If a batch would make the total exceed the limit, the framework raises before executing the batch, and none of the tools in that batch are executed. + +The following configuration applies to every Agent involved in this `Runner.run_async()` call: + +```python +from trpc_agent_sdk.configs import RunConfig + +run_config = RunConfig( + max_llm_calls=10, + max_iterations=5, + max_tool_calls=5, +) +``` + +In a multi-Agent application, use `agent_limits` to configure different limits by `agent.name`: + +```python +from trpc_agent_sdk.configs import AgentRunLimits, RunConfig + +run_config = RunConfig( + max_llm_calls=10, + max_iterations=5, + max_tool_calls=5, + agent_limits={ + "weather_agent": AgentRunLimits( + max_llm_calls=2, + max_iterations=2, + max_tool_calls=1, + ), + "summary_agent": AgentRunLimits( + max_llm_calls=1, + max_tool_calls=0, + ), + }, +) +``` + +Each `agent_limits` key must exactly match the target `agent.name`; otherwise, that override has no effect. Fields omitted from `AgentRunLimits` inherit the top-level `RunConfig` value. Explicitly setting a field to `0` disables the inherited limit for that Agent. + +Every call to an Agent's `run_async()` uses independent counters. Consequently, Agents involved in the same `Runner.run_async()` call are counted independently, and counters start over when a later invocation uses the same session. + +When a limit is exceeded, the framework terminates the current invocation and raises `RunLimitException` to the Python caller: + +```python +from trpc_agent_sdk.exceptions import RunLimitException + +try: + async for event in runner.run_async( + user_id=user_id, + session_id=session_id, + new_message=user_content, + run_config=run_config, + ): + ... +except RunLimitException as exc: + print(exc.error_code) + print(exc.agent_name) + print(exc.limit_type) + print(exc.configured_value, exc.observed_value) +``` + +The exception terminates only the current invocation; it does not close the session. The caller can pass a new `RunConfig` to a later `Runner.run_async()` call and continue using the same session. See [examples/llmagent_with_limit/run_agent.py](../../../examples/llmagent_with_limit/run_agent.py) for a complete, low-cost demonstration. + +`RunConfig` only limits work performed inside the Agent loop. The caller should decide when a timeout starts and how to handle it. In addition, `Runner.run_async()` returns an asynchronous event stream, and the run is not complete until that stream has been consumed. Therefore, the framework does not add a time limit to `RunConfig`; instead, the caller should apply a timeout to the complete event-consumption coroutine: + +```python +import asyncio + +async def run_once() -> None: + async for event in runner.run_async( + user_id=user_id, + session_id=session_id, + new_message=user_content, + run_config=run_config, + ): + ... + +try: + await asyncio.wait_for(run_once(), timeout=120.0) +except TimeoutError: + # Handle the invocation timeout. + ... +``` + ### GenerateContentConfig Used to adjust LLM generation behavior, such as temperature, top-p, and other parameters: diff --git a/docs/mkdocs/zh/llm_agent.md b/docs/mkdocs/zh/llm_agent.md index a3227d461..3a051ef38 100644 --- a/docs/mkdocs/zh/llm_agent.md +++ b/docs/mkdocs/zh/llm_agent.md @@ -385,6 +385,90 @@ for query in demo_queries: ## 高级配置与控制 +### 配置 Agent 的运行次数 + +使用 `RunConfig` 可以限制一次 Agent 调用中的 LLM 调用次数、循环次数和工具调用次数。 + +| 配置项 | 默认值 | 含义 | +| --- | ---: | --- | +| `max_llm_calls` | `500` | LLM 最多调用多少次 | +| `max_iterations` | `0` | Agent 最多循环多少次 | +| `max_tool_calls` | `0` | 工具最多调用多少次 | + +值为 `0` 表示不限制。配置的次数可以正常执行，下一次执行时才会抛出 `RunLimitException`。例如，`max_iterations=1` 时第一次循环可以正常执行，开始第二次循环时抛出异常，因此异常中的 `configured_value` 是 `1`，`observed_value` 是 `2`。 + +如果 LLM 一次返回多个工具调用，框架会一起计算这些工具调用。总数超过 `max_tool_calls` 时，这批工具都不会执行。 + +下面的配置会应用到本次调用中的所有 Agent： + +```python +from trpc_agent_sdk.configs import AgentRunLimits, RunConfig + +run_config = RunConfig( + max_llm_calls=10, + max_iterations=5, + max_tool_calls=5, + agent_limits={ + "weather_agent": AgentRunLimits( + max_llm_calls=2, + max_iterations=2, + max_tool_calls=1, + ), + "summary_agent": AgentRunLimits( + max_llm_calls=1, + max_tool_calls=0, + ), + }, +) +``` + +`agent_limits` 用于给不同 Agent 设置单独的限制。字典中的名称必须与 `agent.name` 完全一致，否则配置不会生效。`AgentRunLimits` 中没有填写的配置会沿用 `RunConfig` 中的值；设置为 `0` 表示该 Agent 不受这项限制。 + +每个 Agent 单独计数。再次调用 Agent 时会重新计数，即使继续使用同一个会话也不会沿用上一次的次数。 + +超过限制时，当前调用会停止，并抛出 `RunLimitException`： + +```python +from trpc_agent_sdk.exceptions import RunLimitException + +try: + async for event in runner.run_async( + user_id=user_id, + session_id=session_id, + new_message=user_content, + run_config=run_config, + ): + ... +except RunLimitException as exc: + print(exc.error_code) + print(exc.agent_name) + print(exc.limit_type) + print(exc.configured_value, exc.observed_value) +``` + +这个异常不会关闭会话。后续调用可以传入新的 `RunConfig`，继续使用原来的会话。完整示例请参考 [examples/llmagent_with_limit/run_agent.py](../../../examples/llmagent_with_limit/run_agent.py)。 + +`RunConfig` 只负责限制 Agent 内部的运行次数。超时从何时开始计算、超时后如何处理，应由调用方决定；而且 `Runner.run_async()` 返回的是异步事件流，需要消费事件后才算运行结束。因此，框架没有在 `RunConfig` 中增加时间限制，而是由调用方为完整的事件消费过程设置超时： + +```python +import asyncio + +async def run_once() -> None: + async for event in runner.run_async( + user_id=user_id, + session_id=session_id, + new_message=user_content, + run_config=run_config, + ): + ... + +try: + await asyncio.wait_for(run_once(), timeout=120.0) +except TimeoutError: + # 处理超时。 + ... +``` + ### GenerateContentConfig 用于调整LLM的生成行为，如temperature、top-p等参数： diff --git a/examples/llmagent_with_limit/.env b/examples/llmagent_with_limit/.env new file mode 100644 index 000000000..399f6375c --- /dev/null +++ b/examples/llmagent_with_limit/.env @@ -0,0 +1,4 @@ +# Set TRPC_AGENT_API_KEY, TRPC_AGENT_BASE_URL, and TRPC_AGENT_MODEL_NAME +TRPC_AGENT_API_KEY=your-api-key +TRPC_AGENT_BASE_URL=your-base-url +TRPC_AGENT_MODEL_NAME=your-model-name diff --git a/examples/llmagent_with_limit/README.md b/examples/llmagent_with_limit/README.md new file mode 100644 index 000000000..763705d5c --- /dev/null +++ b/examples/llmagent_with_limit/README.md @@ -0,0 +1,59 @@ +# LLM Agent 运行次数限制示例 + +本示例演示如何通过 `RunConfig` 限制一次 Agent 调用中的 LLM 调用次数、循环次数和工具调用次数。 + +## 验证内容 + +示例包含三个独立场景： + +- `max_llm_calls=1`：允许一次 LLM 调用，在第二次调用前抛出异常 +- `max_iterations=1`：允许 Agent 执行一次循环，在第二次循环开始前抛出异常 +- `max_tool_calls=1`：让 LLM 一次请求两个工具，超过限制后两个工具都不执行 + +每个场景使用一个单独的会话，并连续调用两次： + +1. 第一次调用触发 `RunLimitException` +2. 第二次调用关闭限制，在同一个会话中询问 `What did we do previously?` + +第二次调用可以正常完成，说明异常只会停止当前调用，不会关闭会话。 + +## 关键配置 + +```python +run_config = RunConfig( + agent_limits={ + root_agent.name: AgentRunLimits( + max_llm_calls=1, + max_iterations=0, + max_tool_calls=0, + ), + }, +) +``` + +`agent_limits` 中的名称需要与 `agent.name` 完全一致。值为 `0` 表示不限制。 + +## 运行示例 + +先在 `.env` 中配置以下环境变量： + +- `TRPC_AGENT_API_KEY` +- `TRPC_AGENT_BASE_URL` +- `TRPC_AGENT_MODEL_NAME` + +然后运行： + +```bash +cd examples/llmagent_with_limit +python3 run_agent.py +``` + +每个场景的预期结果如下： + +```text +⛔ [max_llm_calls_exceeded: Agent 'weather_agent' reached max_llm_calls=1.] +✅ Invocation 1 raised the expected limit: configured=1, observed=2 +✅ Invocation 2 continued and completed normally +``` + +另外两个场景会分别输出 `max_iterations_exceeded` 和 `max_tool_calls_exceeded`。 diff --git a/examples/llmagent_with_limit/agent/__init__.py b/examples/llmagent_with_limit/agent/__init__.py new file mode 100644 index 000000000..bc6e483f9 --- /dev/null +++ b/examples/llmagent_with_limit/agent/__init__.py @@ -0,0 +1,5 @@ +# 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. diff --git a/examples/llmagent_with_limit/agent/agent.py b/examples/llmagent_with_limit/agent/agent.py new file mode 100644 index 000000000..b6469e43f --- /dev/null +++ b/examples/llmagent_with_limit/agent/agent.py @@ -0,0 +1,43 @@ +# 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. +"""Agent used by the run-limit example.""" + +from trpc_agent_sdk.agents import LlmAgent +from trpc_agent_sdk.models import LLMModel +from trpc_agent_sdk.models import OpenAIModel +from trpc_agent_sdk.tools import FunctionTool + +from .config import get_model_config +from .prompts import INSTRUCTION +from .tools import get_weather_forecast +from .tools import get_weather_report + + +def _create_model() -> LLMModel: + """Create the model configured for this example.""" + api_key, base_url, model_name = get_model_config() + return OpenAIModel( + model_name=model_name, + api_key=api_key, + base_url=base_url, + ) + + +def create_agent() -> LlmAgent: + """Create the weather Agent used to trigger the run limits.""" + return LlmAgent( + name="weather_agent", + description="A weather assistant used to demonstrate run limits.", + model=_create_model(), + instruction=INSTRUCTION, + tools=[ + FunctionTool(get_weather_report), + FunctionTool(get_weather_forecast), + ], + ) + + +root_agent = create_agent() diff --git a/examples/llmagent_with_limit/agent/config.py b/examples/llmagent_with_limit/agent/config.py new file mode 100644 index 000000000..d44ac34c3 --- /dev/null +++ b/examples/llmagent_with_limit/agent/config.py @@ -0,0 +1,19 @@ +# 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. +"""Model configuration for the run-limit example.""" + +import os + + +def get_model_config() -> tuple[str, str, str]: + """Read the model configuration from environment variables.""" + api_key = os.getenv("TRPC_AGENT_API_KEY", "") + base_url = os.getenv("TRPC_AGENT_BASE_URL", "") + model_name = os.getenv("TRPC_AGENT_MODEL_NAME", "") + if not api_key or not base_url or not model_name: + raise ValueError("TRPC_AGENT_API_KEY, TRPC_AGENT_BASE_URL, and " + "TRPC_AGENT_MODEL_NAME must be set.") + return api_key, base_url, model_name diff --git a/examples/llmagent_with_limit/agent/prompts.py b/examples/llmagent_with_limit/agent/prompts.py new file mode 100644 index 000000000..b8dd579ad --- /dev/null +++ b/examples/llmagent_with_limit/agent/prompts.py @@ -0,0 +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. +"""Prompt for the run-limit example.""" + +INSTRUCTION = """ +You are a weather assistant for {user_name}, whose default city is {user_city}. + +Use `get_weather_report` for current weather and `get_weather_forecast` for a +multi-day forecast. Follow the user's tool-call instructions exactly. + +When asked what happened previously, answer from the conversation history +without calling a tool. +""" diff --git a/examples/llmagent_with_limit/agent/tools.py b/examples/llmagent_with_limit/agent/tools.py new file mode 100644 index 000000000..8e7b9943c --- /dev/null +++ b/examples/llmagent_with_limit/agent/tools.py @@ -0,0 +1,55 @@ +# 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. +"""Weather tools for the run-limit example.""" + + +def get_weather_report(city: str) -> dict[str, str]: + """Return simulated current weather for a city.""" + weather_data = { + "Beijing": { + "temperature": "25°C", + "condition": "Sunny", + "humidity": "60%", + }, + "Shanghai": { + "temperature": "28°C", + "condition": "Cloudy", + "humidity": "70%", + }, + } + return weather_data.get( + city, + { + "temperature": "Unknown", + "condition": "Data not available", + "humidity": "Unknown", + }, + ) + + +def get_weather_forecast(city: str, days: int = 3) -> list[dict[str, str]]: + """Return a simulated multi-day weather forecast for a city.""" + forecast = [ + { + "date": "2024-01-01", + "city": city, + "temperature": "25°C", + "condition": "Sunny", + }, + { + "date": "2024-01-02", + "city": city, + "temperature": "23°C", + "condition": "Cloudy", + }, + { + "date": "2024-01-03", + "city": city, + "temperature": "20°C", + "condition": "Light rain", + }, + ] + return forecast[:days] diff --git a/examples/llmagent_with_limit/run_agent.py b/examples/llmagent_with_limit/run_agent.py new file mode 100644 index 000000000..e15e1e24f --- /dev/null +++ b/examples/llmagent_with_limit/run_agent.py @@ -0,0 +1,211 @@ +# 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. +"""Run low-cost checks for each Agent run limit.""" + +import asyncio +import uuid +from dataclasses import dataclass + +from dotenv import load_dotenv +from trpc_agent_sdk.configs import AgentRunLimits +from trpc_agent_sdk.configs import RunConfig +from trpc_agent_sdk.events import Event +from trpc_agent_sdk.exceptions import RunLimitException +from trpc_agent_sdk.exceptions import RunLimitType +from trpc_agent_sdk.runners import Runner +from trpc_agent_sdk.sessions import InMemorySessionService +from trpc_agent_sdk.types import Content +from trpc_agent_sdk.types import Part + +load_dotenv() + +_CONTINUATION_QUERY = "What did we do previously?" + + +@dataclass(frozen=True) +class LimitScenario: + """Configuration for one run-limit check.""" + + name: str + trigger_query: str + limits: AgentRunLimits + expected_limit: RunLimitType + + +def _create_limit_scenarios() -> list[LimitScenario]: + """Create one low-cost scenario for each supported limit.""" + return [ + LimitScenario( + name="max_llm_calls", + trigger_query=("Use get_weather_report to check the current weather in Beijing."), + limits=AgentRunLimits( + max_llm_calls=1, + max_iterations=0, + max_tool_calls=0, + ), + expected_limit=RunLimitType.MAX_LLM_CALLS, + ), + LimitScenario( + name="max_iterations", + trigger_query=("Use get_weather_report to check the current weather in Beijing."), + limits=AgentRunLimits( + max_llm_calls=0, + max_iterations=1, + max_tool_calls=0, + ), + expected_limit=RunLimitType.MAX_ITERATIONS, + ), + LimitScenario( + name="max_tool_calls", + trigger_query=("For this message only, call both tools in the same response: " + "get_weather_report for Beijing and get_weather_forecast for " + "Beijing with days=1. Do not retry these calls in a later message."), + limits=AgentRunLimits( + max_llm_calls=0, + max_iterations=0, + max_tool_calls=1, + ), + expected_limit=RunLimitType.MAX_TOOL_CALLS, + ), + ] + + +def _print_event(event: Event) -> None: + """Print visible content from an Agent event.""" + if not event.content or not event.content.parts: + return + + if event.partial: + for part in event.content.parts: + if part.text: + print(part.text, end="", flush=True) + return + + for part in event.content.parts: + if part.thought: + continue + if part.function_call: + print(f"\n🔧 [Invoke Tool:: " + f"{part.function_call.name}({part.function_call.args})]") + elif part.function_response: + print(f"📊 [Tool Result: {part.function_response.response}]") + + +async def _run_invocation( + runner: Runner, + user_id: str, + session_id: str, + query: str, + run_config: RunConfig, +) -> bool: + """Run one prompt and return whether it produced a final response.""" + final_response_received = False + user_content = Content(parts=[Part.from_text(text=query)]) + async for event in runner.run_async( + user_id=user_id, + session_id=session_id, + new_message=user_content, + run_config=run_config, + ): + _print_event(event) + if event.is_final_response(): + final_response_received = True + return final_response_received + + +async def run_weather_agent() -> None: + """Run one isolated check for each supported run limit.""" + from agent.agent import root_agent + + app_name = "weather_agent_limit_demo" + user_id = "demo_user" + session_service = InMemorySessionService() + runner = Runner( + app_name=app_name, + agent=root_agent, + session_service=session_service, + ) + + scenarios = _create_limit_scenarios() + for index, scenario in enumerate(scenarios, 1): + session_id = str(uuid.uuid4()) + run_config = RunConfig(agent_limits={ + root_agent.name: scenario.limits, + }, ) + continuation_run_config = RunConfig(agent_limits={ + root_agent.name: + AgentRunLimits( + max_llm_calls=0, + max_iterations=0, + max_tool_calls=0, + ), + }, ) + + await session_service.create_session( + app_name=app_name, + user_id=user_id, + session_id=session_id, + state={ + "user_name": user_id, + "user_city": "Beijing", + }, + ) + + print(f"\n=== Scenario {index}/{len(scenarios)}: {scenario.name} ===") + print(f"⚙️ Limits: {scenario.limits.model_dump()}") + print(f"🆔 Session ID: {session_id[:8]}...") + + print("\n--- Invocation 1/2: trigger the limit ---") + print(f"📝 User: {scenario.trigger_query}") + print("🤖 Assistant: ", end="", flush=True) + try: + await _run_invocation( + runner, + user_id, + session_id, + scenario.trigger_query, + run_config, + ) + except RunLimitException as exc: + if exc.limit_type != scenario.expected_limit: + raise RuntimeError(f"Scenario '{scenario.name}' expected " + f"{scenario.expected_limit.value}, but received " + f"{exc.limit_type.value}.") from exc + if exc.configured_value != 1 or exc.observed_value != 2: + raise RuntimeError("Limit counters were unexpected: " + f"configured={exc.configured_value}, " + f"observed={exc.observed_value}.") from exc + print(f"\n⛔ [{exc.error_code}: {exc}]") + print("✅ Invocation 1 raised the expected limit: " + f"configured={exc.configured_value}, " + f"observed={exc.observed_value}") + else: + raise RuntimeError(f"Scenario '{scenario.name}' completed without raising " + f"{scenario.expected_limit.value}.") + + print("\n--- Invocation 2/2: continue with the same session ---") + print("⚙️ Limits disabled for the continuation invocation") + print(f"📝 User: {_CONTINUATION_QUERY}") + print("🤖 Assistant: ", end="", flush=True) + try: + final_response_received = await _run_invocation( + runner, + user_id, + session_id, + _CONTINUATION_QUERY, + continuation_run_config, + ) + except RunLimitException as exc: + raise RuntimeError(f"Invocation 2 unexpectedly raised {exc.error_code} even though " + "its limits were disabled.") from exc + if not final_response_received: + raise RuntimeError("Invocation 2 completed without a final response.") + print("\n✅ Invocation 2 continued and completed normally") + print("\n" + "-" * 40) + + +if __name__ == "__main__": + asyncio.run(run_weather_agent()) diff --git a/tests/agents/core/test_llm_processor.py b/tests/agents/core/test_llm_processor.py index e69170163..9f91036fc 100644 --- a/tests/agents/core/test_llm_processor.py +++ b/tests/agents/core/test_llm_processor.py @@ -242,6 +242,8 @@ async def run(): assert event.error_code == "STREAMING_ERROR" assert mock_trace.call_args.args[3].error_message == "rate limit exceeded" + assert mock_trace.call_args.kwargs["error_type"] is None + assert mock_trace.call_args.kwargs["error_message"] is None def test_partial_stream_close_traces_accumulated_text_and_error(self, invocation_context): m = MockLLMModel(model_name="test-llmproc-model") @@ -274,6 +276,8 @@ async def run(): assert response.content.role == "model" assert response.content.parts[0].text == "part1part2" assert response.custom_metadata is None + assert mock_trace.call_args.kwargs["error_type"] == "LlmCallGeneratorExit" + assert mock_trace.call_args.kwargs["error_message"] == "LLM call stopped with GeneratorExit." assert mock_report.call_args.args[2] is response assert mock_report.call_args.kwargs["error_type"] == "LlmCallGeneratorExit" @@ -364,4 +368,6 @@ async def run(): assert response.partial is True assert response.content.parts[0].text == "part1part2" assert response.custom_metadata == {"error_type": "RuntimeError"} + assert mock_trace.call_args.kwargs["error_type"] == "RuntimeError" + assert mock_trace.call_args.kwargs["error_message"] == "upstream failed" assert mock_report.call_args.kwargs["error_type"] == "RuntimeError" diff --git a/tests/agents/test_base_agent.py b/tests/agents/test_base_agent.py index a716111b0..fbd6e8190 100644 --- a/tests/agents/test_base_agent.py +++ b/tests/agents/test_base_agent.py @@ -231,16 +231,12 @@ async def run(): await anext(stream) await stream.aclose() - with patch("trpc_agent_sdk.agents._base_agent.mark_span_error") as mock_span_error, \ - patch("trpc_agent_sdk.agents._base_agent.report_invoke_agent") as mock_report, \ - patch("trpc_agent_sdk.agents._base_agent.trace_agent"), \ + with patch("trpc_agent_sdk.agents._base_agent.report_invoke_agent") as mock_report, \ + patch("trpc_agent_sdk.agents._base_agent.trace_agent") as mock_trace_agent, \ patch("trpc_agent_sdk.agents._base_agent.tracer") as mock_tracer: asyncio.run(run()) - agent_span = mock_tracer.start_as_current_span.return_value.__enter__.return_value - mock_span_error.assert_called_once_with( - agent_span, - error_type="AgentGeneratorExit", - description="Agent execution stopped with GeneratorExit.", - ) + assert mock_tracer.start_as_current_span.called + assert mock_trace_agent.call_args.kwargs["error_type"] == "AgentGeneratorExit" + assert mock_trace_agent.call_args.kwargs["error_message"] == "Agent execution stopped with GeneratorExit." assert mock_report.call_args.kwargs["error_type"] == "AgentGeneratorExit" diff --git a/tests/configs/test_run_config.py b/tests/configs/test_run_config.py index 9e9d4def3..c92baa9f8 100644 --- a/tests/configs/test_run_config.py +++ b/tests/configs/test_run_config.py @@ -25,7 +25,6 @@ from trpc_agent_sdk.configs import RunConfig from trpc_agent_sdk.configs._run_config import RunConfig as RunConfigDirect - # --------------------------------------------------------------------------- # Default values # --------------------------------------------------------------------------- @@ -266,6 +265,9 @@ def test_model_dump_returns_all_fields(self): d = cfg.model_dump() expected_keys = { "max_llm_calls", + "max_iterations", + "max_tool_calls", + "agent_limits", "streaming", "agent_run_config", "custom_data", @@ -285,6 +287,9 @@ def test_model_json_schema_has_all_fields(self): schema = RunConfig.model_json_schema() props = schema.get("properties", {}) assert "max_llm_calls" in props + assert "max_iterations" in props + assert "max_tool_calls" in props + assert "agent_limits" in props assert "streaming" in props assert "agent_run_config" in props assert "custom_data" in props diff --git a/tests/exceptions/test_exceptions.py b/tests/exceptions/test_exceptions.py index 991077e93..a6121645b 100644 --- a/tests/exceptions/test_exceptions.py +++ b/tests/exceptions/test_exceptions.py @@ -21,7 +21,6 @@ TrpcAgentException, ) - # --------------------------------------------------------------------------- # ErrorCode # --------------------------------------------------------------------------- @@ -42,6 +41,7 @@ def test_is_int_enum(self): (ErrorCode.ARTIFACT_SERVICE_NOT_FOUND, 603), (ErrorCode.LLM_AGENT_MODEL_NOT_FOUND, 604), (ErrorCode.RUN_CANCELLED, 605), + (ErrorCode.RUN_LIMIT_EXCEEDED, 606), ], ) def test_member_values(self, member: ErrorCode, expected_value: int): @@ -57,6 +57,7 @@ def test_member_values(self, member: ErrorCode, expected_value: int): (ErrorCode.ARTIFACT_SERVICE_NOT_FOUND, "artifact_service not found"), (ErrorCode.LLM_AGENT_MODEL_NOT_FOUND, "model not found"), (ErrorCode.RUN_CANCELLED, "run cancelled"), + (ErrorCode.RUN_LIMIT_EXCEEDED, "run limit exceeded"), ], ) def test_member_phrases(self, member: ErrorCode, expected_phrase: str): @@ -71,13 +72,14 @@ def test_member_phrases(self, member: ErrorCode, expected_phrase: str): (ErrorCode.ARTIFACT_SERVICE_NOT_FOUND, "the artifact_service maybe is none"), (ErrorCode.LLM_AGENT_MODEL_NOT_FOUND, "the artifact not found"), (ErrorCode.RUN_CANCELLED, "the run was cancelled by user request"), + (ErrorCode.RUN_LIMIT_EXCEEDED, "the agent invocation reached a configured run limit"), ], ) def test_member_descriptions(self, member: ErrorCode, expected_description: str): assert member.description == expected_description def test_total_member_count(self): - assert len(ErrorCode) == 6 + assert len(ErrorCode) == 7 def test_can_be_used_as_int(self): assert ErrorCode.OK + 1 == 1 @@ -87,10 +89,12 @@ def test_lookup_by_value(self): assert ErrorCode(0) is ErrorCode.OK assert ErrorCode(601) is ErrorCode.PARENT_AGENT_NOT_FOUND assert ErrorCode(605) is ErrorCode.RUN_CANCELLED + assert ErrorCode(606) is ErrorCode.RUN_LIMIT_EXCEEDED def test_lookup_by_name(self): assert ErrorCode["OK"] is ErrorCode.OK assert ErrorCode["RUN_CANCELLED"] is ErrorCode.RUN_CANCELLED + assert ErrorCode["RUN_LIMIT_EXCEEDED"] is ErrorCode.RUN_LIMIT_EXCEEDED def test_invalid_value_raises(self): with pytest.raises(ValueError): @@ -118,7 +122,7 @@ def test_init_stores_code(self): def test_init_sets_args_to_phrase(self): exc = TrpcAgentException(ErrorCode.OK) - assert exc.args == ("OK",) + assert exc.args == ("OK", ) def test_str_format(self): exc = TrpcAgentException(ErrorCode.PARENT_AGENT_NOT_FOUND) @@ -224,19 +228,19 @@ def test_llm_agent_model_not_found_code(self): def test_all_are_trpc_agent_exception_instances(self): for instance in ( - ParentAgentNotFound, - AgentFilterError, - ArtifactServiceNotFound, - LLMAgentModelNotFound, + ParentAgentNotFound, + AgentFilterError, + ArtifactServiceNotFound, + LLMAgentModelNotFound, ): assert isinstance(instance, TrpcAgentException) def test_all_are_exception_instances(self): for instance in ( - ParentAgentNotFound, - AgentFilterError, - ArtifactServiceNotFound, - LLMAgentModelNotFound, + ParentAgentNotFound, + AgentFilterError, + ArtifactServiceNotFound, + LLMAgentModelNotFound, ): assert isinstance(instance, Exception) @@ -262,10 +266,10 @@ def test_can_be_raised_and_caught(self, instance, code): def test_instances_are_not_run_cancelled(self): for instance in ( - ParentAgentNotFound, - AgentFilterError, - ArtifactServiceNotFound, - LLMAgentModelNotFound, + ParentAgentNotFound, + AgentFilterError, + ArtifactServiceNotFound, + LLMAgentModelNotFound, ): assert not isinstance(instance, RunCancelledException) @@ -288,6 +292,8 @@ def test_all_names_importable(self): "LLMAgentModelNotFound", "ParentAgentNotFound", "RunCancelledException", + "RunLimitException", + "RunLimitType", "TrpcAgentException", ] for name in expected: @@ -303,6 +309,8 @@ def test_all_attribute(self): "LLMAgentModelNotFound", "ParentAgentNotFound", "RunCancelledException", + "RunLimitException", + "RunLimitType", "TrpcAgentException", } assert set(mod.__all__) == expected_all diff --git a/tests/telemetry/test_trace.py b/tests/telemetry/test_trace.py index 6b8453b43..0a9fbf477 100644 --- a/tests/telemetry/test_trace.py +++ b/tests/telemetry/test_trace.py @@ -29,7 +29,6 @@ _build_llm_request_for_trace, _safe_json_serialize, get_trpc_agent_span_name, - mark_span_error, set_trpc_agent_span_name, trace_agent, trace_call_llm, @@ -181,37 +180,34 @@ def test_serialize_empty_dict(self): # --------------------------------------------------------------------------- -# Tests: mark_span_error +# Tests: trace_runner # --------------------------------------------------------------------------- -class TestMarkSpanError: +class TestTraceRunner: + + def setup_method(self): + set_trpc_agent_span_name("trpc.python.agent") - def test_marks_interruption_with_operation_specific_error(self): + @patch("trpc_agent_sdk.telemetry._trace.trace.get_current_span") + def test_error_sets_span_status_and_type(self, mock_get_span): span = _mock_span() + mock_get_span.return_value = span - mark_span_error( - span, + trace_runner( + "app", + "user", + "session", + _make_invocation_context(), error_type="RunnerGeneratorExit", - description="Runner invocation stopped with GeneratorExit.", + error_message="Runner invocation stopped with GeneratorExit.", ) span.set_status.assert_called_once_with( trace.StatusCode.ERROR, "Runner invocation stopped with GeneratorExit.", ) - span.set_attribute.assert_called_once_with("error.type", "RunnerGeneratorExit") - - -# --------------------------------------------------------------------------- -# Tests: trace_runner -# --------------------------------------------------------------------------- - - -class TestTraceRunner: - - def setup_method(self): - set_trpc_agent_span_name("trpc.python.agent") + span.set_attribute.assert_any_call("error.type", "RunnerGeneratorExit") @patch("trpc_agent_sdk.telemetry._trace.trace.get_current_span") def test_basic_attributes(self, mock_get_span): @@ -937,7 +933,7 @@ def test_basic_llm_trace(self, mock_get_span): span.set_attribute.assert_any_call("trpc.python.agent.event_id", "e-1") @patch("trpc_agent_sdk.telemetry._trace.trace.get_current_span") - def test_error_response_sets_status_message_and_keeps_llm_response_output(self, mock_get_span): + def test_explicit_error_sets_status_message_and_keeps_llm_response_output(self, mock_get_span): span = _mock_span() mock_get_span.return_value = span ctx = _make_invocation_context() @@ -949,20 +945,19 @@ def test_error_response_sets_status_message_and_keeps_llm_response_output(self, custom_metadata={"error_type": "RateLimitError"}, ) - trace_call_llm(ctx, event_id="e-1", llm_request=req, llm_response=resp) + trace_call_llm( + ctx, + event_id="e-1", + llm_request=req, + llm_response=resp, + error_type="RateLimitError", + error_message="rate limit exceeded", + ) # Output remains the LlmResponse JSON; status carries the error message. span.set_attribute.assert_any_call("trpc.python.agent.llm_response", '{"content": "response"}') span.set_status.assert_called_once_with(trace.StatusCode.ERROR, "rate limit exceeded") span.set_attribute.assert_any_call("error.type", "RateLimitError") - span.set_attribute.assert_any_call( - "trpc.python.agent.llm.error_code", - "STREAMING_ERROR", - ) - span.set_attribute.assert_any_call( - "trpc.python.agent.llm.error_message", - "rate limit exceeded", - ) span.add_event.assert_not_called() @patch("trpc_agent_sdk.telemetry._trace.trace.get_current_span") diff --git a/tests/test_runner.py b/tests/test_runner.py index 3756a962b..305189943 100644 --- a/tests/test_runner.py +++ b/tests/test_runner.py @@ -294,8 +294,7 @@ async def mock_agent_run(ctx): mock_agent.run_async = mock_agent_run - with patch("trpc_agent_sdk.runners.mark_span_error") as mock_span_error, \ - patch("trpc_agent_sdk.runners.trace_runner"), \ + with patch("trpc_agent_sdk.runners.trace_runner") as mock_trace_runner, \ patch("trpc_agent_sdk.runners.tracer") as mock_tracer: stream = runner.run_async( user_id="test_user", @@ -307,12 +306,9 @@ async def mock_agent_run(ctx): await stream.aclose() assert event.partial is True - invocation_span = mock_tracer.start_as_current_span.return_value.__enter__.return_value - mock_span_error.assert_called_once_with( - invocation_span, - error_type="RunnerGeneratorExit", - description="Runner invocation stopped with GeneratorExit.", - ) + assert mock_tracer.start_as_current_span.called + assert mock_trace_runner.call_args.kwargs["error_type"] == "RunnerGeneratorExit" + assert mock_trace_runner.call_args.kwargs["error_message"] == "Runner invocation stopped with GeneratorExit." @pytest.mark.asyncio async def test_run_async_non_streaming_mode(self, runner, mock_session_service, mock_agent, mock_session): diff --git a/tests/tools/test_agent_tool.py b/tests/tools/test_agent_tool.py index 467080769..59b54c68f 100644 --- a/tests/tools/test_agent_tool.py +++ b/tests/tools/test_agent_tool.py @@ -114,6 +114,7 @@ def mock_context(self): ctx.state = MagicMock() ctx.state.to_dict.return_value = {} ctx.event_actions = MagicMock() + ctx.run_config = None ctx.save_artifact = AsyncMock() return ctx @@ -268,6 +269,7 @@ def mock_context(self): ctx.state = MagicMock() ctx.state.to_dict.return_value = {} ctx.event_actions = MagicMock() + ctx.run_config = None ctx.save_artifact = AsyncMock() return ctx diff --git a/trpc_agent_sdk/agents/_base_agent.py b/trpc_agent_sdk/agents/_base_agent.py index ff39076c6..322268d36 100644 --- a/trpc_agent_sdk/agents/_base_agent.py +++ b/trpc_agent_sdk/agents/_base_agent.py @@ -43,9 +43,9 @@ from trpc_agent_sdk.context import reset_invocation_ctx from trpc_agent_sdk.context import set_invocation_ctx from trpc_agent_sdk.events import Event +from trpc_agent_sdk.exceptions import RunLimitException from trpc_agent_sdk.filter import get_filter from trpc_agent_sdk.filter import run_stream_filters -from trpc_agent_sdk.telemetry import mark_span_error from trpc_agent_sdk.telemetry import report_invoke_agent from trpc_agent_sdk.telemetry import tracer from trpc_agent_sdk.telemetry import trace_agent @@ -219,6 +219,7 @@ def model_post_init(self, __context: Any) -> None: def _create_invocation_context(self, parent_context: InvocationContext) -> InvocationContext: """Creates a new invocation context for this agent.""" invocation_context = parent_context.model_copy(update={"agent": self}) + invocation_context._reset_run_limit_observed() # Handle branch assignment: # - If parent_context.agent is the same as self, we're being called from runner @@ -284,7 +285,8 @@ async def run_async( mono_start = time.monotonic() t_first_visible: Optional[float] = None - metrics_error_type: Optional[str] = None + error_type: Optional[str] = None + error_message: Optional[str] = None try: gen_co = run_stream_filters(ctx.agent_context, None, self.filters, handle) # type: ignore @@ -297,15 +299,16 @@ async def run_async( non_partial_events.append(event) yield event # type: ignore except GeneratorExit: - metrics_error_type = "AgentGeneratorExit" - mark_span_error( - agent_span, - error_type=metrics_error_type, - description="Agent execution stopped with GeneratorExit.", - ) + error_type = "AgentGeneratorExit" + error_message = "Agent execution stopped with GeneratorExit." + raise + except RunLimitException as ex: + error_type = ex.error_code + error_message = str(ex) raise except Exception as ex: - metrics_error_type = type(ex).__name__ + error_type = type(ex).__name__ + error_message = str(ex) raise finally: # Compute state after agent run @@ -321,6 +324,8 @@ async def run_async( agent_action=agent_action, state_begin=state_begin, state_end=state_end, + error_type=error_type, + error_message=error_message, ) duration_s = time.monotonic() - mono_start @@ -334,7 +339,7 @@ async def run_async( input_tokens=input_tokens, output_tokens=output_tokens, is_stream=is_stream, - error_type=metrics_error_type, + error_type=error_type, ) # avoid memory leak diff --git a/trpc_agent_sdk/agents/_llm_agent.py b/trpc_agent_sdk/agents/_llm_agent.py index 3546c1d6f..b31844e07 100644 --- a/trpc_agent_sdk/agents/_llm_agent.py +++ b/trpc_agent_sdk/agents/_llm_agent.py @@ -30,6 +30,8 @@ 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.exceptions import RunLimitException +from trpc_agent_sdk.exceptions import RunLimitType from trpc_agent_sdk.log import logger from trpc_agent_sdk.models import LLMModel from trpc_agent_sdk.models import LlmRequest @@ -467,9 +469,11 @@ def accumulate_content(event: Event) -> None: running = agent_context.get_metadata(TRPC_AGENT_RUNNING_KEY, True) # Multi-turn conversation loop - continue until no more tool calls or code execution while running: - # CHECKPOINT 1: At start of each conversation turn + # CHECKPOINT 1: At start of each loop iteration await ctx.raise_if_cancelled() + ctx.raise_if_limit(RunLimitType.MAX_ITERATIONS) + # Step 1: Build request using the request processor (includes conversation history) request = LlmRequest(model=model_instance.name, ) @@ -496,6 +500,7 @@ def accumulate_content(event: Event) -> None: collected_tool_calls = [] code_was_executed = False + ctx.raise_if_limit(RunLimitType.MAX_LLM_CALLS) logger.debug("Starting LLM call for agent: %s", self.name) # Use LlmProcessor to get unified events @@ -561,6 +566,11 @@ def accumulate_content(event: Event) -> None: # Use extended tools processor that includes transfer tool if needed extended_tools_processor = self._get_extended_tools_processor(ctx) + ctx.raise_if_limit( + RunLimitType.MAX_TOOL_CALLS, + increment=len(collected_tool_calls), + ) + # Check if any of the tool calls are for long-running tools long_running_tool_ids = set() for tool_call in collected_tool_calls: @@ -659,6 +669,9 @@ def accumulate_content(event: Event) -> None: # raise to runner to handle raise + except RunLimitException: + raise + except Exception as ex: # pylint: disable=broad-except logger.error("Error executing tools for agent %s: %s", self.name, ex, exc_info=True) @@ -682,6 +695,8 @@ def accumulate_content(event: Event) -> None: except RunCancelledException: # raise to runner to handle raise + except RunLimitException: + raise except Exception as ex: # pylint: disable=broad-except logger.error("Unexpected error in LLM agent %s: %s", self.name, ex, exc_info=True) diff --git a/trpc_agent_sdk/agents/core/_llm_processor.py b/trpc_agent_sdk/agents/core/_llm_processor.py index ecbfb87ac..39163f257 100644 --- a/trpc_agent_sdk/agents/core/_llm_processor.py +++ b/trpc_agent_sdk/agents/core/_llm_processor.py @@ -130,7 +130,8 @@ def _build_interrupted_content() -> Optional[Content]: t_start = time.monotonic() t_first_token: Optional[float] = None - metrics_error_type: Optional[str] = None + error_type: Optional[str] = None + error_message: Optional[str] = None try: async for llm_response in self.model.generate_async(request, stream=stream, ctx=context): latest_llm_response = llm_response @@ -171,23 +172,25 @@ def _build_interrupted_content() -> Optional[Content]: yield event except GeneratorExit: - metrics_error_type = "LlmCallGeneratorExit" + error_type = "LlmCallGeneratorExit" + error_message = "LLM call stopped with GeneratorExit." final_llm_response = LlmResponse( content=_build_interrupted_content(), partial=True, - error_code=metrics_error_type, - error_message="LLM call stopped with GeneratorExit.", + error_code=error_type, + error_message=error_message, interrupted=True, ) raise except Exception as ex: - metrics_error_type = type(ex).__name__ + error_type = type(ex).__name__ + error_message = str(ex) final_llm_response = LlmResponse( content=_build_interrupted_content(), partial=True, error_code="LLM_CALL_ERROR", - error_message=str(ex), - custom_metadata={"error_type": type(ex).__name__}, + error_message=error_message, + custom_metadata={"error_type": error_type}, ) raise finally: @@ -201,6 +204,8 @@ def _build_interrupted_content() -> Optional[Content]: instruction_metadata=instruction_metadata, stream_function_calls_raw=aggregated_raw_function_calls, stream_function_calls_post_planner=aggregated_event_function_calls, + error_type=error_type, + error_message=error_message, ) duration_s = time.monotonic() - t_start @@ -212,7 +217,7 @@ def _build_interrupted_content() -> Optional[Content]: duration_s=duration_s, ttft_s=ttft_s, is_stream=stream, - error_type=metrics_error_type, + error_type=error_type, ) if terminal_event is not None: diff --git a/trpc_agent_sdk/agents/core/_tools_processor.py b/trpc_agent_sdk/agents/core/_tools_processor.py index abbce309d..7ec9bfb18 100644 --- a/trpc_agent_sdk/agents/core/_tools_processor.py +++ b/trpc_agent_sdk/agents/core/_tools_processor.py @@ -26,6 +26,7 @@ from trpc_agent_sdk.context import InvocationContext from trpc_agent_sdk.events import Event from trpc_agent_sdk.events import EventActions +from trpc_agent_sdk.exceptions import RunLimitException from trpc_agent_sdk.log import logger from trpc_agent_sdk.models import LlmRequest from trpc_agent_sdk.telemetry import report_execute_tool @@ -148,6 +149,8 @@ async def __invoke_tools( else: try: result_event = await self._execute_tool(tool_call, tool, context) + except RunLimitException: + raise except Exception as ex: # pylint: disable=broad-except logger.error("Error executing tool %s: %s", tool_call.name, ex, exc_info=True) result_event = self._create_error_event( @@ -211,10 +214,23 @@ async def execute_tools_async( if parallel_tool_calls: # Parallel execution: collect all events and merge them function_response_events: list[Event] = [] - async with asyncio.TaskGroup() as tg: - for tool_call in non_streaming_calls: - tg.create_task(self.__invoke_tools(context, resolved_tools, tool_call, - function_response_events)) + tasks = [ + asyncio.create_task( + self.__invoke_tools( + context, + resolved_tools, + tool_call, + function_response_events, + )) for tool_call in non_streaming_calls + ] + try: + await asyncio.gather(*tasks) + except RunLimitException: + for task in tasks: + if not task.done(): + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + raise # Handle merging and tracing based on number of events if function_response_events: @@ -421,6 +437,8 @@ async def _execute_tool(self, tool_call: FunctionCall, tool: BaseTool, context: return event + except RunLimitException: + raise except Exception as ex: # pylint: disable=broad-except report_execute_tool( context, @@ -442,6 +460,8 @@ async def _execute_tool(self, tool_call: FunctionCall, tool: BaseTool, context: function_response_event=error_event, state_begin=state_begin, state_end=state_end, + error_type=type(ex).__name__, + error_message=str(ex), ) return error_event @@ -579,6 +599,8 @@ async def _execute_progress_streaming_tool( yield final_event + except RunLimitException: + raise except Exception as ex: # pylint: disable=broad-except report_execute_tool( context, @@ -600,6 +622,8 @@ async def _execute_progress_streaming_tool( function_response_event=error_event, state_begin=state_begin, state_end=state_end, + error_type=type(ex).__name__, + error_message=str(ex), ) logger.error("Error executing streaming tool %s: %s", tool_call.name, ex, exc_info=True) yield error_event diff --git a/trpc_agent_sdk/agents/sub_agent/_runner.py b/trpc_agent_sdk/agents/sub_agent/_runner.py index d35bd2152..ed955d7b6 100644 --- a/trpc_agent_sdk/agents/sub_agent/_runner.py +++ b/trpc_agent_sdk/agents/sub_agent/_runner.py @@ -26,8 +26,10 @@ from trpc_agent_sdk.abc import ArtifactId from trpc_agent_sdk.agents._llm_agent import LlmAgent +from trpc_agent_sdk.configs import RunConfig from trpc_agent_sdk.context import InvocationContext from trpc_agent_sdk.exceptions import RunCancelledException +from trpc_agent_sdk.exceptions import RunLimitException from trpc_agent_sdk.log import logger from trpc_agent_sdk.memory import InMemoryMemoryService from trpc_agent_sdk.sessions import InMemorySessionService @@ -388,6 +390,7 @@ async def run_subagent_streaming( user_id=sub_session.user_id, session_id=sub_session.id, new_message=content, + run_config=parent_ctx.run_config or RunConfig(), ): last_event = event # Forward this sub-agent event to the parent consumer as a progress @@ -409,6 +412,8 @@ async def run_subagent_streaming( await _forward_artifacts(sub_runner, sub_session, parent_ctx) except RunCancelledException: final_value = "[sub-agent cancelled]" + except RunLimitException: + raise except Exception as ex: # noqa: BLE001 logger.error("sub-agent run failed: %s", ex, exc_info=True) final_value = {"status": "error", "message": str(ex)} diff --git a/trpc_agent_sdk/cancel/_session_utils.py b/trpc_agent_sdk/cancel/_session_utils.py index 7b5a235b0..d54f8d804 100644 --- a/trpc_agent_sdk/cancel/_session_utils.py +++ b/trpc_agent_sdk/cancel/_session_utils.py @@ -21,7 +21,10 @@ _CANCELING_SUFFIX = "Detect user cancel the agent execution." -async def cleanup_incomplete_function_calls(session: SessionABC) -> None: +async def cleanup_incomplete_function_calls( + session: SessionABC, + invocation_id: Optional[str] = None, +) -> None: """Remove function_calls from session.events that have no matching function_response. When cancellation occurs after tool execution, some function_calls may not @@ -30,6 +33,8 @@ async def cleanup_incomplete_function_calls(session: SessionABC) -> None: Args: session: The session to clean up + invocation_id: When provided, only remove unmatched function calls + created by this invocation. """ # Step 1: Collect all function_response IDs in the session response_ids: set[str] = set() @@ -40,6 +45,8 @@ async def cleanup_incomplete_function_calls(session: SessionABC) -> None: # Step 2: Find and remove incomplete function_calls for event in session.events: + if invocation_id is not None and event.invocation_id != invocation_id: + continue func_calls = event.get_function_calls() if not func_calls: continue diff --git a/trpc_agent_sdk/configs/__init__.py b/trpc_agent_sdk/configs/__init__.py index c505fbb5f..d6cb30fd6 100644 --- a/trpc_agent_sdk/configs/__init__.py +++ b/trpc_agent_sdk/configs/__init__.py @@ -5,12 +5,14 @@ # tRPC-Agent-Python is licensed under Apache-2.0. """Configs for TRPC Agent framework.""" +from ._agent_run_limits import AgentRunLimits from ._model_retry_config import ExponentialBackoffConfig from ._model_retry_config import ModelRetryConfig from ._prompt_cache_config import PromptCacheConfig from ._run_config import RunConfig __all__ = [ + "AgentRunLimits", "ExponentialBackoffConfig", "ModelRetryConfig", "PromptCacheConfig", diff --git a/trpc_agent_sdk/configs/_agent_run_limits.py b/trpc_agent_sdk/configs/_agent_run_limits.py new file mode 100644 index 000000000..ebc591edb --- /dev/null +++ b/trpc_agent_sdk/configs/_agent_run_limits.py @@ -0,0 +1,36 @@ +# 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. +"""Per-agent run-limit configuration.""" + +from __future__ import annotations + +import sys +from typing import Optional + +from pydantic import BaseModel +from pydantic import ConfigDict +from pydantic import Field + + +class AgentRunLimits(BaseModel): + """Per-agent overrides for limits configured on :class:`RunConfig`. + + A field set to ``None`` inherits the corresponding top-level value from + :class:`RunConfig`. A value of ``0`` disables that limit for the selected + agent, while a positive value enables it. + """ + + model_config = ConfigDict(extra="forbid") + """The Pydantic model configuration.""" + + max_llm_calls: Optional[int] = Field(default=None, ge=0, lt=sys.maxsize) + """Maximum logical LLM calls for each invocation of the selected agent.""" + + max_iterations: Optional[int] = Field(default=None, ge=0) + """Maximum loop iterations for each invocation of the selected agent.""" + + max_tool_calls: Optional[int] = Field(default=None, ge=0) + """Maximum tool calls for each invocation of the selected agent.""" diff --git a/trpc_agent_sdk/configs/_run_config.py b/trpc_agent_sdk/configs/_run_config.py index 50f16d581..1ba662f7f 100644 --- a/trpc_agent_sdk/configs/_run_config.py +++ b/trpc_agent_sdk/configs/_run_config.py @@ -18,6 +18,7 @@ from trpc_agent_sdk.log import logger +from ._agent_run_limits import AgentRunLimits from ._prompt_cache_config import PromptCacheConfig @@ -29,7 +30,7 @@ class RunConfig(BaseModel): max_llm_calls: int = 500 """ - A limit on the total number of llm calls for a given run. + A limit on the total number of llm calls for each agent invocation. Valid Values: - More than 0 and less than sys.maxsize: The bound on the number of llm @@ -37,6 +38,28 @@ class RunConfig(BaseModel): - Less than or equal to 0: This allows for unbounded number of llm calls. """ + max_iterations: int = Field(default=0, ge=0) + """Maximum loop iterations for each agent invocation. + + A value of ``0`` disables this limit. The counter is local to each call to + an agent's ``run_async`` method. + """ + + max_tool_calls: int = Field(default=0, ge=0) + """Maximum tool calls for each agent invocation. + + A value of ``0`` disables this limit. A batch that would exceed the limit + is rejected before any tool in that batch is executed. + """ + + agent_limits: dict[str, AgentRunLimits] = Field(default_factory=dict) + """Per-agent limit overrides keyed by the exact ``agent.name``. + + Unset fields inherit their top-level values from this ``RunConfig``. A + per-agent value of ``0`` explicitly disables the corresponding inherited + limit for that agent. + """ + streaming: bool = True """Whether to enable streaming mode. Default is True.""" diff --git a/trpc_agent_sdk/context/_invocation_context.py b/trpc_agent_sdk/context/_invocation_context.py index 4f00be0d8..b20a995ca 100644 --- a/trpc_agent_sdk/context/_invocation_context.py +++ b/trpc_agent_sdk/context/_invocation_context.py @@ -41,6 +41,7 @@ from pydantic import BaseModel from pydantic import ConfigDict from pydantic import Field +from pydantic import PrivateAttr from trpc_agent_sdk.abc import AgentABC from trpc_agent_sdk.abc import ArtifactEntry @@ -49,7 +50,10 @@ from trpc_agent_sdk.abc import MemoryServiceABC from trpc_agent_sdk.abc import SessionABC from trpc_agent_sdk.abc import SessionServiceABC +from trpc_agent_sdk.configs import AgentRunLimits from trpc_agent_sdk.configs import RunConfig +from trpc_agent_sdk.exceptions import RunLimitException +from trpc_agent_sdk.exceptions import RunLimitType from trpc_agent_sdk.types import ActiveStreamingTool from trpc_agent_sdk.types import Content from trpc_agent_sdk.types import EventActions @@ -134,6 +138,8 @@ class InvocationContext(BaseModel): run_config: Optional[RunConfig] = None """Configuration for agent execution under this invocation.""" + _observed_run_limits: AgentRunLimits = PrivateAttr(default_factory=AgentRunLimits) + event_actions: EventActions = Field(default=EventActions(), init=True) callback_state: Optional[State] = None @@ -252,7 +258,7 @@ async def save_artifact(self, filename: str, artifact: Part) -> int: The version of the artifact. """ if self.artifact_service is None: - raise ValueError("Artifact service is not initialized.") + raise ValueError('Artifact service is not initialized.') version = await self.artifact_service.save_artifact( artifact_id=ArtifactId( app_name=self.app_name, @@ -279,7 +285,7 @@ async def list_artifacts(self) -> list[str]: - Returns empty list if no artifacts exist """ if self.artifact_service is None: - raise ValueError('Artifact service is not initialized.') + raise ValueError("Artifact service is not initialized.") return await self.artifact_service.list_artifact_keys(artifact_id=ArtifactId( app_name=self.app_name, user_id=self.user_id, @@ -310,6 +316,68 @@ async def search_memory(self, query: str) -> SearchMemoryResponse: agent_context=self.agent_context, ) + def _reset_run_limit_observed(self) -> None: + """Reset observed run-limit values for a new agent invocation.""" + self._observed_run_limits = AgentRunLimits() + + def _get_run_limit(self, limit_type: RunLimitType) -> int: + """Return the effective limit for the current agent.""" + if self.run_config is None: + return 0 + + agent_limits = self.run_config.agent_limits.get(self.agent.name) + if limit_type == RunLimitType.MAX_LLM_CALLS: + if agent_limits is not None and agent_limits.max_llm_calls is not None: + return agent_limits.max_llm_calls + return self.run_config.max_llm_calls + if limit_type == RunLimitType.MAX_ITERATIONS: + if agent_limits is not None and agent_limits.max_iterations is not None: + return agent_limits.max_iterations + return self.run_config.max_iterations + if agent_limits is not None and agent_limits.max_tool_calls is not None: + return agent_limits.max_tool_calls + return self.run_config.max_tool_calls + + def raise_if_limit( + self, + limit_type: RunLimitType, + increment: int = 1, + ) -> None: + """Record observed work and raise if its configured limit is exceeded. + + Args: + limit_type: Limit whose observed value should be incremented. + increment: Number of newly observed operations. + + Raises: + RunLimitException: If the updated observed value exceeds + the effective limit for the current agent. + ValueError: If ``increment`` is negative. + """ + if increment < 0: + raise ValueError("Run-limit increment must not be negative.") + + configured_value = self._get_run_limit(limit_type) + if limit_type == RunLimitType.MAX_LLM_CALLS: + observed_value = (self._observed_run_limits.max_llm_calls or 0) + increment + self._observed_run_limits.max_llm_calls = observed_value + elif limit_type == RunLimitType.MAX_ITERATIONS: + observed_value = (self._observed_run_limits.max_iterations or 0) + increment + self._observed_run_limits.max_iterations = observed_value + else: + observed_value = (self._observed_run_limits.max_tool_calls or 0) + increment + self._observed_run_limits.max_tool_calls = observed_value + + if configured_value <= 0: + return + if observed_value > configured_value: + raise RunLimitException( + agent_name=self.agent.name, + limit_type=limit_type, + configured_value=configured_value, + observed_value=observed_value, + ) + async def raise_if_cancelled(self) -> None: """Raise RunCancelledException if this run is cancelled. diff --git a/trpc_agent_sdk/dsl/graph/_graph_agent.py b/trpc_agent_sdk/dsl/graph/_graph_agent.py index 12e05c1de..a3e1eb0c8 100644 --- a/trpc_agent_sdk/dsl/graph/_graph_agent.py +++ b/trpc_agent_sdk/dsl/graph/_graph_agent.py @@ -32,6 +32,7 @@ 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.exceptions import RunLimitException from trpc_agent_sdk.log import logger from trpc_agent_sdk.types import Content from trpc_agent_sdk.types import FunctionCall @@ -239,6 +240,8 @@ async def _run_async_impl(self, ctx: InvocationContext) -> AsyncGenerator[Event, yield chunk_event # Cancellation checkpoint after yielding events await ctx.raise_if_cancelled() + except RunLimitException: + raise except Exception as e: error_message = str(e) logger.error(f"[{self.name}] Graph execution failed: {e}", exc_info=True) diff --git a/trpc_agent_sdk/dsl/graph/_node_action/_agent.py b/trpc_agent_sdk/dsl/graph/_node_action/_agent.py index 85d064dc5..6ce6378d3 100644 --- a/trpc_agent_sdk/dsl/graph/_node_action/_agent.py +++ b/trpc_agent_sdk/dsl/graph/_node_action/_agent.py @@ -14,6 +14,7 @@ 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.exceptions import RunLimitException from trpc_agent_sdk.types import Content from trpc_agent_sdk.types import EventActions from trpc_agent_sdk.types import Part @@ -244,6 +245,8 @@ async def execute(self, state: State) -> dict[str, Any]: if isinstance(candidate, str) and candidate: last_response = candidate + except RunLimitException: + raise except Exception as e: raise RuntimeError(f"Agent node '{self.name}' execution failed: {e}") from e diff --git a/trpc_agent_sdk/exceptions/__init__.py b/trpc_agent_sdk/exceptions/__init__.py index c4feb4c58..e3941ec0f 100644 --- a/trpc_agent_sdk/exceptions/__init__.py +++ b/trpc_agent_sdk/exceptions/__init__.py @@ -11,6 +11,8 @@ from ._exceptions import LLMAgentModelNotFound from ._exceptions import ParentAgentNotFound from ._exceptions import RunCancelledException +from ._exceptions import RunLimitException +from ._exceptions import RunLimitType from ._exceptions import TrpcAgentException __all__ = [ @@ -20,5 +22,7 @@ "LLMAgentModelNotFound", "ParentAgentNotFound", "RunCancelledException", + "RunLimitException", + "RunLimitType", "TrpcAgentException", ] diff --git a/trpc_agent_sdk/exceptions/_exceptions.py b/trpc_agent_sdk/exceptions/_exceptions.py index 8070a128a..f6bf665c6 100644 --- a/trpc_agent_sdk/exceptions/_exceptions.py +++ b/trpc_agent_sdk/exceptions/_exceptions.py @@ -5,7 +5,17 @@ # tRPC-Agent-Python is licensed under Apache-2.0. """Exceptions for TRPC Agent framework.""" +from enum import Enum from enum import IntEnum +from typing import Union + + +class RunLimitType(str, Enum): + """Types of invocation-local count limits.""" + + MAX_LLM_CALLS = "max_llm_calls" + MAX_ITERATIONS = "max_iterations" + MAX_TOOL_CALLS = "max_tool_calls" class ErrorCode(IntEnum): @@ -25,6 +35,7 @@ def __new__(cls, value, phrase, description=''): ARTIFACT_SERVICE_NOT_FOUND = 603, 'artifact_service not found', 'the artifact_service maybe is none' LLM_AGENT_MODEL_NOT_FOUND = 604, 'model not found', 'the artifact not found' RUN_CANCELLED = 605, 'run cancelled', 'the run was cancelled by user request' + RUN_LIMIT_EXCEEDED = 606, 'run limit exceeded', 'the agent invocation reached a configured run limit' class TrpcAgentException(Exception): @@ -39,6 +50,50 @@ def __str__(self) -> str: return f'code: {self.code}, msg: {self.code.phrase}, reason: {self.code.description}' +class RunLimitException(TrpcAgentException): + """Exception raised when an agent invocation exceeds a configured limit. + + Attributes: + agent_name: Name of the agent whose invocation reached the limit. + limit_type: Type of the configured limit. + configured_value: Maximum value configured for the limit. + observed_value: Value observed when the limit was detected. + """ + + def __init__( + self, + *, + agent_name: str, + limit_type: RunLimitType, + configured_value: int, + observed_value: int, + ) -> None: + super().__init__(ErrorCode.RUN_LIMIT_EXCEEDED) + self.agent_name = agent_name + self.limit_type = limit_type + self.configured_value = configured_value + self.observed_value = observed_value + self.message = f"Agent '{agent_name}' reached {limit_type.value}={configured_value}." + + def __str__(self) -> str: + """Return the limit-specific error message.""" + return self.message + + @property + def error_code(self) -> str: + """Return the stable error code for protocol and telemetry adapters.""" + return f"{self.limit_type.value}_exceeded" + + def get_custom_metadata(self) -> dict[str, Union[str, int]]: + """Return JSON-serializable details for protocol adapters.""" + return { + "limit_type": self.limit_type.value, + "configured_value": self.configured_value, + "observed_value": self.observed_value, + "agent_name": self.agent_name, + } + + class RunCancelledException(TrpcAgentException): """Exception raised when a run is cancelled. diff --git a/trpc_agent_sdk/filter/_base_filter.py b/trpc_agent_sdk/filter/_base_filter.py index 8585a441c..04b5dd38d 100644 --- a/trpc_agent_sdk/filter/_base_filter.py +++ b/trpc_agent_sdk/filter/_base_filter.py @@ -53,10 +53,10 @@ async def _after(self, ctx, req): from trpc_agent_sdk.abc import FilterHandleType from trpc_agent_sdk.abc import FilterResult from trpc_agent_sdk.context import AgentContext +from trpc_agent_sdk.exceptions import RunCancelledException +from trpc_agent_sdk.exceptions import RunLimitException from trpc_agent_sdk.log import logger -from ..exceptions import RunCancelledException - class BaseFilter(FilterABC): """Abstract base class defining the filter interface. @@ -146,8 +146,8 @@ async def _handle_co(self, logger.debug(self._create_err_str(f"{msg} error: {result.error}")) if not result.is_continue: return - except RunCancelledException: - # raise to runner to handle + except (RunCancelledException, RunLimitException): + # Preserve framework control-flow exceptions for the execution boundary. raise except Exception as ex: # pylint: disable=broad-except logger.error("filter type: %s, name: %s run %s error: %s", diff --git a/trpc_agent_sdk/runners.py b/trpc_agent_sdk/runners.py index 7ee3f00b4..7c26a2302 100644 --- a/trpc_agent_sdk/runners.py +++ b/trpc_agent_sdk/runners.py @@ -32,11 +32,11 @@ from trpc_agent_sdk.events import AgentCancelledEvent from trpc_agent_sdk.events import Event from trpc_agent_sdk.exceptions import RunCancelledException +from trpc_agent_sdk.exceptions import RunLimitException from trpc_agent_sdk.log import logger from trpc_agent_sdk.memory import BaseMemoryService from trpc_agent_sdk.sessions import BaseSessionService from trpc_agent_sdk.sessions import Session -from trpc_agent_sdk.telemetry import mark_span_error from trpc_agent_sdk.telemetry import tracer from trpc_agent_sdk.telemetry import trace_cancellation from trpc_agent_sdk.telemetry import trace_runner @@ -382,6 +382,10 @@ async def run_async( Yields: The events generated by the agent. + + Raises: + RunLimitException: If an agent exceeds one of its configured run + limits. """ # Manually propagate span context using attach/detach instead of # start_as_current_span. This ensures child spans (agent_run, call_llm, @@ -458,6 +462,8 @@ async def run_async( # Track accumulated partial text for cancellation handling temp_text_parts: list[str] = [] runner_trace_recorded = False + trace_error_type: Optional[str] = None + trace_error_message: Optional[str] = None try: # Support multiple levels of agent transfers @@ -564,12 +570,21 @@ async def run_async( await self._schedule_post_turn_processing(invocation_context=invocation_context, ) except GeneratorExit: - mark_span_error( - invocation_span, - error_type="RunnerGeneratorExit", - description="Runner invocation stopped with GeneratorExit.", + trace_error_type = "RunnerGeneratorExit" + trace_error_message = "Runner invocation stopped with GeneratorExit." + raise + + except RunLimitException as ex: + trace_error_type = ex.error_code + trace_error_message = str(ex) + logger.warning("Run for session %s exceeded a configured limit: %s", session_id, ex) + await cancel.cleanup_incomplete_function_calls( + session, + invocation_id=invocation_context.invocation_id, ) + await self.session_service.update_session(session=session) raise + except RunCancelledException as ex: logger.info("Run for session %s was cancelled", session_id) logger.debug("Cancellation details: %s", ex, exc_info=True) @@ -613,6 +628,11 @@ async def run_async( branch=invocation_context.branch, ) + except Exception as ex: + trace_error_type = type(ex).__name__ + trace_error_message = str(ex) + raise + finally: if not runner_trace_recorded: state_end = dict(session.state) @@ -630,6 +650,8 @@ async def run_async( last_event=last_non_streaming_event, state_begin=state_begin, state_end=state_end, + error_type=trace_error_type, + error_message=trace_error_message, ) # Always cleanup cancellation tracking diff --git a/trpc_agent_sdk/server/a2a/_remote_a2a_agent.py b/trpc_agent_sdk/server/a2a/_remote_a2a_agent.py index e76540c50..ac09e69d3 100644 --- a/trpc_agent_sdk/server/a2a/_remote_a2a_agent.py +++ b/trpc_agent_sdk/server/a2a/_remote_a2a_agent.py @@ -412,6 +412,10 @@ def _events_from_response(self, result: Any, event_count: int, ctx: InvocationCo if state not in (TaskState.submitted, TaskState.working, TaskState.completed): partial = self._resolve_partial(result.metadata) ev = convert_a2a_message_to_event(msg, author=self.name, invocation_context=ctx, partial=partial) + if state == TaskState.failed: + error_code = get_metadata(result.metadata, "error_code") or get_metadata(msg.metadata, "error_code") + ev.error_code = error_code or "a2a_task_failed" + ev.error_message = ev.get_text() or "Remote A2A task failed" events.append(ev) return events diff --git a/trpc_agent_sdk/server/a2a/converters/_event_converter.py b/trpc_agent_sdk/server/a2a/converters/_event_converter.py index 04429329e..7e8874640 100644 --- a/trpc_agent_sdk/server/a2a/converters/_event_converter.py +++ b/trpc_agent_sdk/server/a2a/converters/_event_converter.py @@ -499,7 +499,9 @@ def create_exception_status_event( context_id: str, message_text: str, final: bool = True, + metadata: Optional[Dict[str, Any]] = None, ) -> TaskStatusUpdateEvent: + """Create a terminal failed task status for an execution exception.""" return TaskStatusUpdateEvent( task_id=task_id, status=TaskStatus( @@ -509,10 +511,12 @@ def create_exception_status_event( message_id=str(uuid.uuid4()), role=Role.agent, parts=[TextPart(text=message_text)], + metadata=metadata, ), ), context_id=context_id, final=final, + metadata=metadata, ) diff --git a/trpc_agent_sdk/server/a2a/executor/_a2a_agent_executor.py b/trpc_agent_sdk/server/a2a/executor/_a2a_agent_executor.py index 17159bf24..dafb5cbfa 100644 --- a/trpc_agent_sdk/server/a2a/executor/_a2a_agent_executor.py +++ b/trpc_agent_sdk/server/a2a/executor/_a2a_agent_executor.py @@ -41,9 +41,11 @@ from pydantic import BaseModel from trpc_agent_sdk.cancel import SessionKey from trpc_agent_sdk.cancel import is_run_cancelled +from trpc_agent_sdk.configs import RunConfig from trpc_agent_sdk.context import new_agent_context from trpc_agent_sdk.events import AgentCancelledEvent from trpc_agent_sdk.events import Event +from trpc_agent_sdk.exceptions import RunLimitException from trpc_agent_sdk.log import logger from trpc_agent_sdk.runners import Runner @@ -63,6 +65,8 @@ EventCallback = Callable[[Event, RequestContext], Union[Optional[Event], Awaitable[Optional[Event]]]] +RunConfigFactory = Callable[[RequestContext], Union[RunConfig, Awaitable[RunConfig]]] + class TrpcA2aAgentExecutorConfig(BaseModel): """Configuration for TrpcA2aAgentExecutor. @@ -75,6 +79,11 @@ class TrpcA2aAgentExecutorConfig(BaseModel): Return the event to continue, a modified event to alter behavior, or None to skip this event entirely. Useful for filtering, logging, or augmenting events (e.g. detecting streaming tool calls via event.is_streaming_tool_call()). + run_config: Optional server-owned configuration applied to every agent + invocation. Request metadata is preserved in ``agent_run_config``. + run_config_factory: Optional callback that creates a server-owned + configuration for each request. When set, it takes precedence over + ``run_config``. """ model_config = {"arbitrary_types_allowed": True} @@ -82,6 +91,8 @@ class TrpcA2aAgentExecutorConfig(BaseModel): cancel_wait_timeout: float = 1.0 user_id_extractor: Optional[UserIdExtractor] = None event_callback: Optional[EventCallback] = None + run_config: Optional[RunConfig] = None + run_config_factory: Optional[RunConfigFactory] = None class TrpcA2aAgentExecutor(AgentExecutor): @@ -116,6 +127,15 @@ async def _resolve_runner(self) -> Runner: return resolved raise TypeError(f"Runner must be a Runner instance or callable, got {type(self._runner)}") + async def _resolve_run_config(self, context: RequestContext) -> Optional[RunConfig]: + """Resolve the server-owned run configuration for one request.""" + if self._config.run_config_factory is None: + return self._config.run_config + result = self._config.run_config_factory(context) + if inspect.isawaitable(result): + return await result + return result + def _get_user_session_from_task_metadata( self, context: RequestContext, @@ -215,6 +235,20 @@ async def execute(self, context: RequestContext, event_queue: EventQueue): return await self._handle_request(context, event_queue) + except RunLimitException as ex: + logger.warning("A2A task %s exceeded a configured run limit: %s", context.task_id, ex) + metadata = ex.get_custom_metadata() + metadata["error_code"] = ex.error_code + try: + await event_queue.enqueue_event( + create_exception_status_event( + task_id=context.task_id, + context_id=context.context_id, + message_text=str(ex), + metadata=metadata, + )) + except Exception as enqueue_error: # pylint: disable=broad-except + logger.error("Failed to publish run-limit failure event: %s", enqueue_error, exc_info=True) except Exception as ex: # pylint: disable=broad-except logger.error("Error handling A2A request: %s", ex, exc_info=True) try: @@ -238,6 +272,19 @@ async def execute(self, context: RequestContext, event_queue: EventQueue): async def _handle_request(self, context: RequestContext, event_queue: EventQueue): runner = await self._resolve_runner() run_args = await convert_a2a_request_to_trpc_agent_run_args(context, self._user_id_extractor) + configured_run_config = await self._resolve_run_config(context) + if configured_run_config is not None: + request_agent_run_config = run_args["run_config"].agent_run_config + configured_agent_run_config = configured_run_config.agent_run_config + run_args["run_config"] = configured_run_config.model_copy( + update={ + "agent_run_config": { + **configured_agent_run_config, + **request_agent_run_config, + }, + }, + deep=True, + ) session = await self._prepare_session(run_args, runner) agent_context = new_agent_context() 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 267486747..6478aa8a8 100644 --- a/trpc_agent_sdk/server/ag_ui/_core/_agui_agent.py +++ b/trpc_agent_sdk/server/ag_ui/_core/_agui_agent.py @@ -46,6 +46,7 @@ from trpc_agent_sdk.configs import RunConfig as TRPCRunConfig from trpc_agent_sdk.events import EventTranslatorBase from trpc_agent_sdk.events import LongRunningEvent +from trpc_agent_sdk.exceptions import RunLimitException from trpc_agent_sdk.log import logger from trpc_agent_sdk.abc import ToolSetABC from trpc_agent_sdk.memory import BaseMemoryService @@ -918,7 +919,7 @@ async def _start_new_execution(self, logger.debug("Finished iterating over _stream_events for execution %s", execution.thread_id) # If we found tool calls, add them to session state BEFORE cleanup - if has_tool_calls: + if has_tool_calls and not has_error: app_name = self.get_app_name(input) user_id = self.get_user_id(input) for tool_call_id in tool_call_ids: @@ -1273,6 +1274,16 @@ async def _run_trpc_in_background(self, await event_queue.put(None) logger.debug("Background task completion signal sent for thread %s", input.thread_id) + except RunLimitException as ex: + logger.warning("AG-UI run %s exceeded a configured run limit: %s", input.run_id, ex) + async for ag_ui_event in event_translator.force_close_streaming_message(): + await event_queue.put(ag_ui_event) + await event_queue.put(RunErrorEvent( + type=EventType.RUN_ERROR, + message=str(ex), + code=ex.error_code, + )) + await event_queue.put(None) except Exception as ex: # pylint: disable=broad-except logger.error("Background execution error: %s", ex, exc_info=True) # Put error in queue diff --git a/trpc_agent_sdk/telemetry/__init__.py b/trpc_agent_sdk/telemetry/__init__.py index e50d18a5e..03bacbd44 100644 --- a/trpc_agent_sdk/telemetry/__init__.py +++ b/trpc_agent_sdk/telemetry/__init__.py @@ -11,7 +11,6 @@ from ._metrics import report_execute_tool from ._metrics import report_invoke_agent from ._trace import get_trpc_agent_span_name -from ._trace import mark_span_error from ._trace import set_trpc_agent_span_name from ._trace import trace_agent from ._trace import trace_call_llm @@ -27,7 +26,6 @@ "report_call_llm", "report_execute_tool", "report_invoke_agent", - "mark_span_error", "trace_agent", "trace_call_llm", "trace_cancellation", diff --git a/trpc_agent_sdk/telemetry/_trace.py b/trpc_agent_sdk/telemetry/_trace.py index 2a1af505b..9596c9084 100644 --- a/trpc_agent_sdk/telemetry/_trace.py +++ b/trpc_agent_sdk/telemetry/_trace.py @@ -62,21 +62,6 @@ def get_trpc_agent_span_name() -> str: return _trpc_agent_span_name -def mark_span_error(span: trace.Span, error_type: str, description: str) -> None: - """Mark a span as failed with an operation-specific error. - - The caller supplies an error type and description that identify the failed - operation. - - Args: - span: The failed operation's span. - error_type: The operation-specific error type. - description: The human-readable error description. - """ - span.set_status(trace.StatusCode.ERROR, description) - span.set_attribute("error.type", error_type) - - def _join_parts_with_thought_tag(parts) -> str: """Join part texts, wrapping thought parts in tags. @@ -127,6 +112,8 @@ def trace_runner( last_event: Optional[Event] = None, state_begin: Optional[dict[str, Any]] = None, state_end: Optional[dict[str, Any]] = None, + error_type: Optional[str] = None, + error_message: Optional[str] = None, ): """Traces runner execution. @@ -142,6 +129,8 @@ def trace_runner( last_event: The last non-streaming event from the agent execution. state_begin: The state before the runner execution. state_end: The state after the runner execution. + error_type: The error type when the runner does not complete normally. + error_message: The error message when the runner does not complete normally. """ span = trace.get_current_span() span.set_attribute("gen_ai.system", _trpc_agent_span_name) @@ -169,6 +158,10 @@ def trace_runner( if state_end is not None: span.set_attribute(f"{_trpc_agent_span_name}.state.end", _safe_json_serialize(state_end)) + if error_type: + span.set_status(trace.StatusCode.ERROR, error_message or error_type) + span.set_attribute("error.type", error_type) + def trace_cancellation( app_name: str, @@ -249,6 +242,8 @@ def trace_agent( agent_action: str = "", state_begin: Optional[dict[str, Any]] = None, state_end: Optional[dict[str, Any]] = None, + error_type: Optional[str] = None, + error_message: Optional[str] = None, ): """Traces agent execution. @@ -261,6 +256,8 @@ def trace_agent( (text, function calls, function responses). state_begin: The state before the agent run. state_end: The state after the agent run. + error_type: The error type when the agent does not complete normally. + error_message: The error message when the agent does not complete normally. """ span = trace.get_current_span() span.set_attribute("gen_ai.system", _trpc_agent_span_name) @@ -303,6 +300,10 @@ def trace_agent( if state_end is not None: span.set_attribute(f"{_trpc_agent_span_name}.state.end", _safe_json_serialize(state_end)) + if error_type: + span.set_status(trace.StatusCode.ERROR, error_message or error_type) + span.set_attribute("error.type", error_type) + def trace_tool_call( tool: BaseTool, @@ -310,6 +311,8 @@ def trace_tool_call( function_response_event: Event, state_begin: Optional[dict[str, Any]] = None, state_end: Optional[dict[str, Any]] = None, + error_type: Optional[str] = None, + error_message: Optional[str] = None, ): """Traces tool call. @@ -319,6 +322,8 @@ def trace_tool_call( function_response_event: The event with the function response details. state_begin: The state before the tool execution. state_end: The state after the tool execution. + error_type: The error type when the tool call does not complete normally. + error_message: The error message when the tool call does not complete normally. """ span = trace.get_current_span() span.set_attribute("gen_ai.system", _trpc_agent_span_name) @@ -367,6 +372,10 @@ def trace_tool_call( if state_end is not None: span.set_attribute(f"{_trpc_agent_span_name}.state.end", _safe_json_serialize(state_end)) + if error_type: + span.set_status(trace.StatusCode.ERROR, error_message or error_type) + span.set_attribute("error.type", error_type) + def trace_merged_tool_calls( response_event_id: str, @@ -427,6 +436,8 @@ def trace_call_llm( instruction_metadata: Optional[InstructionMetadata] = None, stream_function_calls_raw: Optional[list[dict[str, Any]]] = None, stream_function_calls_post_planner: Optional[list[dict[str, Any]]] = None, + error_type: Optional[str] = None, + error_message: Optional[str] = None, ): """Traces a call to the LLM. @@ -446,6 +457,8 @@ def trace_call_llm( raw LLM stream chunks. stream_function_calls_post_planner: Optional function calls collected from post-planner events emitted during stream processing. + error_type: The error type when the LLM call does not complete normally. + error_message: The error message when the LLM call does not complete normally. """ span = trace.get_current_span() # Special standard Open Telemetry GenaI attributes that indicate @@ -472,20 +485,9 @@ def trace_call_llm( llm_response_json, ) - # call_llm observation output is always the LlmResponse JSON above. On - # error, also set ERROR status (with description) and error attributes, but - # skip exception events so exporters keep using llm_response as output. - error_code = getattr(llm_response, "error_code", None) - if error_code: - error_message = getattr(llm_response, "error_message", None) - custom_metadata = getattr(llm_response, "custom_metadata", None) - error_type = custom_metadata.get("error_type") if isinstance(custom_metadata, dict) else None - error_type = str(error_type or error_code) - status_description = str(error_message or error_code) - mark_span_error(span, error_type, status_description) - span.set_attribute(f"{_trpc_agent_span_name}.llm.error_code", str(error_code)) - if error_message: - span.set_attribute(f"{_trpc_agent_span_name}.llm.error_message", str(error_message)) + if error_type: + span.set_status(trace.StatusCode.ERROR, error_message or error_type) + span.set_attribute("error.type", error_type) if stream_function_calls_raw: span.set_attribute( diff --git a/trpc_agent_sdk/tools/_agent_tool.py b/trpc_agent_sdk/tools/_agent_tool.py index 78ba4ad1c..fd2854262 100644 --- a/trpc_agent_sdk/tools/_agent_tool.py +++ b/trpc_agent_sdk/tools/_agent_tool.py @@ -49,6 +49,7 @@ from trpc_agent_sdk.abc import AgentABC from trpc_agent_sdk.abc import ArtifactId +from trpc_agent_sdk.configs import RunConfig from trpc_agent_sdk.context import InvocationContext from trpc_agent_sdk.events import Event from trpc_agent_sdk.filter import BaseFilter @@ -177,7 +178,12 @@ async def _run_async_impl( ) last_event = None - async for event in runner.run_async(user_id=session.user_id, session_id=session.id, new_message=content): + async for event in runner.run_async( + user_id=session.user_id, + session_id=session.id, + new_message=content, + run_config=tool_context.run_config or RunConfig(), + ): # Forward state delta to parent session. assert isinstance(event, Event) if event.actions.state_delta: