Agent Zero
c760b5961c
Check Cross-Plugin Imports / check (push) Has been cancelled
- app/ai/agent_loop.py: ReActStep + ReActResult dataclasses, run_react_loop() with LLM→Tool→Observe loop, max_steps/timeout graceful stop, ErrorCategory retry (TRANSIENT→retry, PERMANENT→stop, PARTIAL→continue), cost accumulation, agent.step hook, on_step callback - app/plugins/builtins/automation/models.py: AgentRunStep model - alembic/versions/0121_agent_run_steps.py: migration for agent run steps table - app/plugins/builtins/automation/agent_runner.py: refactored to use run_react_loop(), saves steps to DB, updates AgentRun with cost/status/duration - tests/test_agent_loop.py: 11 tests (all passing, mocked, no DB/LLM needed) - PROGRESS.md: Phase F started, F-LOOP marked done
482 lines
16 KiB
Python
482 lines
16 KiB
Python
"""Tests for the ReAct agent loop (app/ai/agent_loop.py).
|
|
|
|
All tests mock llm_complete and tool_registry — no real LLM or DB needed.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
import uuid
|
|
from dataclasses import dataclass
|
|
from typing import Any
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
from app.ai.agent_loop import ReActResult, ReActStep, run_react_loop
|
|
|
|
|
|
# ──────────────────────────────────────────────────────────────────────────
|
|
# Test helpers
|
|
# ──────────────────────────────────────────────────────────────────────────
|
|
|
|
|
|
@dataclass
|
|
class MockAgentDefinition:
|
|
"""Minimal stand-in for AgentDefinition used by run_react_loop."""
|
|
|
|
id: uuid.UUID
|
|
llm_model: str = "gpt-4o"
|
|
system_prompt: str = "You are a helpful assistant."
|
|
api_key: str | None = None
|
|
api_base: str | None = None
|
|
provider: str | None = None
|
|
max_tokens: int = 1000
|
|
|
|
|
|
def _make_tool_call(
|
|
call_id: str = "call_1",
|
|
name: str = "search",
|
|
arguments: dict[str, Any] | None = None,
|
|
) -> dict[str, Any]:
|
|
"""Build a mock LiteLLM tool-call object."""
|
|
args_str = json.dumps(arguments or {})
|
|
function_mock = MagicMock()
|
|
function_mock.name = name
|
|
function_mock.arguments = args_str
|
|
tc = MagicMock()
|
|
tc.id = call_id
|
|
tc.function = function_mock
|
|
return tc
|
|
|
|
|
|
def _make_llm_response(
|
|
content: str = "",
|
|
tool_calls: list[Any] | None = None,
|
|
cost_usd: float = 0.001,
|
|
) -> dict[str, Any]:
|
|
"""Build a mock llm_complete() return value."""
|
|
message = MagicMock()
|
|
message.content = content
|
|
message.tool_calls = tool_calls or None
|
|
|
|
choice = MagicMock()
|
|
choice.message = message
|
|
|
|
raw_response = MagicMock()
|
|
raw_response.choices = [choice]
|
|
|
|
return {
|
|
"content": content,
|
|
"usage": {"total_tokens": 100},
|
|
"cost_usd": cost_usd,
|
|
"model": "gpt-4o",
|
|
"raw_response": raw_response,
|
|
}
|
|
|
|
|
|
def _make_tool_registry(tools: dict[str, AsyncMock] | None = None) -> MagicMock:
|
|
"""Build a mock ToolRegistry."""
|
|
registry = MagicMock()
|
|
_tools = tools or {}
|
|
|
|
def _get(name: str) -> Any:
|
|
if name in _tools:
|
|
tool = MagicMock()
|
|
tool.handler = _tools[name]
|
|
return tool
|
|
return None
|
|
|
|
registry.get = _get
|
|
return registry
|
|
|
|
|
|
@pytest.fixture
|
|
def agent_def() -> MockAgentDefinition:
|
|
return MockAgentDefinition(id=uuid.uuid4())
|
|
|
|
|
|
@pytest.fixture
|
|
def tenant_id() -> uuid.UUID:
|
|
return uuid.uuid4()
|
|
|
|
|
|
@pytest.fixture
|
|
def user_id() -> uuid.UUID:
|
|
return uuid.uuid4()
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_db() -> AsyncMock:
|
|
return AsyncMock()
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _mock_hooks():
|
|
"""Patch do_action so the loop doesn't try to fire real hooks."""
|
|
with patch("app.core.hooks.do_action", new_callable=AsyncMock):
|
|
yield
|
|
|
|
|
|
# ──────────────────────────────────────────────────────────────────────────
|
|
# Tests
|
|
# ──────────────────────────────────────────────────────────────────────────
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_basic_react_loop_no_tools(agent_def, tenant_id, user_id, mock_db):
|
|
"""Single-step loop: LLM returns final response, no tool calls."""
|
|
registry = _make_tool_registry()
|
|
|
|
with patch("app.ai.agent_loop.llm_complete", new_callable=AsyncMock) as mock_llm:
|
|
mock_llm.return_value = _make_llm_response(
|
|
content="Hello! How can I help you?",
|
|
cost_usd=0.002,
|
|
)
|
|
|
|
result = await run_react_loop(
|
|
agent_definition=agent_def,
|
|
messages=[{"role": "user", "content": "Hi"}],
|
|
tools=[],
|
|
tool_registry=registry,
|
|
db=mock_db,
|
|
tenant_id=tenant_id,
|
|
user_id=user_id,
|
|
)
|
|
|
|
assert result.status == "completed"
|
|
assert result.final_content == "Hello! How can I help you?"
|
|
assert result.steps_taken == 1
|
|
assert len(result.steps) == 1
|
|
assert result.steps[0].action is None # No tool call
|
|
assert result.steps[0].thought == "Hello! How can I help you?"
|
|
assert result.total_cost_usd == pytest.approx(0.002)
|
|
mock_llm.assert_awaited_once()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_multi_step_loop_with_tool_calls(agent_def, tenant_id, user_id, mock_db):
|
|
"""Multi-step loop: LLM calls a tool, gets result, then gives final answer."""
|
|
search_handler = AsyncMock(return_value="Found 3 contacts: Alice, Bob, Charlie")
|
|
registry = _make_tool_registry({"search": search_handler})
|
|
|
|
# Step 1: LLM calls search tool
|
|
# Step 2: LLM gives final answer
|
|
responses = [
|
|
_make_llm_response(
|
|
content="I'll search for contacts.",
|
|
tool_calls=[_make_tool_call(name="search", arguments={"query": "contacts"})],
|
|
cost_usd=0.003,
|
|
),
|
|
_make_llm_response(
|
|
content="I found 3 contacts: Alice, Bob, and Charlie.",
|
|
cost_usd=0.004,
|
|
),
|
|
]
|
|
|
|
with patch("app.ai.agent_loop.llm_complete", new_callable=AsyncMock) as mock_llm:
|
|
mock_llm.side_effect = responses
|
|
|
|
result = await run_react_loop(
|
|
agent_definition=agent_def,
|
|
messages=[{"role": "user", "content": "Find contacts"}],
|
|
tools=[{"type": "function", "function": {"name": "search"}}],
|
|
tool_registry=registry,
|
|
db=mock_db,
|
|
tenant_id=tenant_id,
|
|
user_id=user_id,
|
|
)
|
|
|
|
assert result.status == "completed"
|
|
assert result.steps_taken == 2
|
|
assert len(result.steps) == 2
|
|
assert result.steps[0].action == "search"
|
|
assert result.steps[0].action_input == {"query": "contacts"}
|
|
assert result.steps[0].observation == "Found 3 contacts: Alice, Bob, Charlie"
|
|
assert result.steps[1].action is None # Final response
|
|
assert result.final_content == "I found 3 contacts: Alice, Bob, and Charlie."
|
|
assert result.total_cost_usd == pytest.approx(0.007)
|
|
assert mock_llm.await_count == 2
|
|
search_handler.assert_awaited_once()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_max_steps_limit(agent_def, tenant_id, user_id, mock_db):
|
|
"""Loop stops gracefully when max_steps is reached."""
|
|
search_handler = AsyncMock(return_value="result")
|
|
registry = _make_tool_registry({"search": search_handler})
|
|
|
|
# Every response has a tool call — never gives final answer
|
|
tool_call = _make_tool_call(name="search", arguments={})
|
|
response = _make_llm_response(
|
|
content="Searching...",
|
|
tool_calls=[tool_call],
|
|
cost_usd=0.001,
|
|
)
|
|
|
|
with patch("app.ai.agent_loop.llm_complete", new_callable=AsyncMock) as mock_llm:
|
|
mock_llm.return_value = response
|
|
|
|
result = await run_react_loop(
|
|
agent_definition=agent_def,
|
|
messages=[{"role": "user", "content": "Keep searching"}],
|
|
tools=[{"type": "function", "function": {"name": "search"}}],
|
|
tool_registry=registry,
|
|
db=mock_db,
|
|
tenant_id=tenant_id,
|
|
user_id=user_id,
|
|
max_steps=3,
|
|
)
|
|
|
|
assert result.status == "stopped_max_steps"
|
|
assert result.steps_taken == 3
|
|
assert len(result.steps) == 3
|
|
assert "max_steps" in (result.error or "")
|
|
assert mock_llm.await_count == 3
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_timeout_graceful_stop(agent_def, tenant_id, user_id, mock_db):
|
|
"""Loop stops gracefully when timeout is exceeded."""
|
|
registry = _make_tool_registry()
|
|
|
|
# Simulate slow LLM responses that eventually exceed timeout
|
|
async def slow_llm(**kwargs: Any) -> dict[str, Any]:
|
|
await asyncio.sleep(0.5)
|
|
return _make_llm_response(content="thinking...", tool_calls=None, cost_usd=0.001)
|
|
|
|
with patch("app.ai.agent_loop.llm_complete", new_callable=AsyncMock) as mock_llm:
|
|
mock_llm.side_effect = slow_llm
|
|
|
|
result = await run_react_loop(
|
|
agent_definition=agent_def,
|
|
messages=[{"role": "user", "content": "Hi"}],
|
|
tools=[],
|
|
tool_registry=registry,
|
|
db=mock_db,
|
|
tenant_id=tenant_id,
|
|
user_id=user_id,
|
|
timeout_seconds=0, # Immediate timeout on first check
|
|
)
|
|
|
|
assert result.status == "stopped_timeout"
|
|
assert "timeout" in (result.error or "").lower()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_transient_error_retry(agent_def, tenant_id, user_id, mock_db):
|
|
"""Transient errors are retried, then succeed."""
|
|
registry = _make_tool_registry()
|
|
|
|
# First call raises transient error, second succeeds
|
|
responses: list[Any] = [
|
|
Exception("rate limit exceeded"),
|
|
_make_llm_response(content="Success after retry", cost_usd=0.002),
|
|
]
|
|
|
|
with patch("app.ai.agent_loop.llm_complete", new_callable=AsyncMock) as mock_llm:
|
|
mock_llm.side_effect = responses
|
|
with patch("app.ai.agent_loop.asyncio.sleep", new_callable=AsyncMock):
|
|
result = await run_react_loop(
|
|
agent_definition=agent_def,
|
|
messages=[{"role": "user", "content": "Hi"}],
|
|
tools=[],
|
|
tool_registry=registry,
|
|
db=mock_db,
|
|
tenant_id=tenant_id,
|
|
user_id=user_id,
|
|
)
|
|
|
|
assert result.status == "completed"
|
|
assert result.final_content == "Success after retry"
|
|
assert mock_llm.await_count == 2 # First failed, second succeeded
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_permanent_error_stops(agent_def, tenant_id, user_id, mock_db):
|
|
"""Permanent errors stop the loop immediately."""
|
|
registry = _make_tool_registry()
|
|
|
|
with patch("app.ai.agent_loop.llm_complete", new_callable=AsyncMock) as mock_llm:
|
|
mock_llm.side_effect = Exception("authentication error: invalid api key")
|
|
|
|
result = await run_react_loop(
|
|
agent_definition=agent_def,
|
|
messages=[{"role": "user", "content": "Hi"}],
|
|
tools=[],
|
|
tool_registry=registry,
|
|
db=mock_db,
|
|
tenant_id=tenant_id,
|
|
user_id=user_id,
|
|
)
|
|
|
|
assert result.status == "stopped_error"
|
|
assert "Permanent error" in (result.error or "")
|
|
assert mock_llm.await_count == 1 # No retry for permanent errors
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cost_accumulation(agent_def, tenant_id, user_id, mock_db):
|
|
"""Cost is accumulated across multiple LLM calls."""
|
|
search_handler = AsyncMock(return_value="result")
|
|
registry = _make_tool_registry({"search": search_handler})
|
|
|
|
responses = [
|
|
_make_llm_response(
|
|
content="Step 1",
|
|
tool_calls=[_make_tool_call(name="search")],
|
|
cost_usd=0.01,
|
|
),
|
|
_make_llm_response(
|
|
content="Step 2",
|
|
tool_calls=[_make_tool_call(name="search")],
|
|
cost_usd=0.02,
|
|
),
|
|
_make_llm_response(
|
|
content="Final answer",
|
|
cost_usd=0.03,
|
|
),
|
|
]
|
|
|
|
with patch("app.ai.agent_loop.llm_complete", new_callable=AsyncMock) as mock_llm:
|
|
mock_llm.side_effect = responses
|
|
|
|
result = await run_react_loop(
|
|
agent_definition=agent_def,
|
|
messages=[{"role": "user", "content": "Search twice"}],
|
|
tools=[{"type": "function", "function": {"name": "search"}}],
|
|
tool_registry=registry,
|
|
db=mock_db,
|
|
tenant_id=tenant_id,
|
|
user_id=user_id,
|
|
)
|
|
|
|
assert result.status == "completed"
|
|
assert result.total_cost_usd == pytest.approx(0.06) # 0.01 + 0.02 + 0.03
|
|
assert result.steps_taken == 3
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_step_persistence_via_callback(agent_def, tenant_id, user_id, mock_db):
|
|
"""Steps are passed to on_step callback for persistence."""
|
|
registry = _make_tool_registry()
|
|
collected_steps: list[ReActStep] = []
|
|
|
|
async def on_step(step: ReActStep) -> None:
|
|
collected_steps.append(step)
|
|
|
|
with patch("app.ai.agent_loop.llm_complete", new_callable=AsyncMock) as mock_llm:
|
|
mock_llm.return_value = _make_llm_response(content="Done", cost_usd=0.001)
|
|
|
|
result = await run_react_loop(
|
|
agent_definition=agent_def,
|
|
messages=[{"role": "user", "content": "Hi"}],
|
|
tools=[],
|
|
tool_registry=registry,
|
|
db=mock_db,
|
|
tenant_id=tenant_id,
|
|
user_id=user_id,
|
|
on_step=on_step,
|
|
)
|
|
|
|
assert len(collected_steps) == 1
|
|
assert collected_steps[0].step_number == 1
|
|
assert collected_steps[0].thought == "Done"
|
|
assert result.status == "completed"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_tool_not_found(agent_def, tenant_id, user_id, mock_db):
|
|
"""When a tool is not found, the observation contains an error message."""
|
|
registry = _make_tool_registry() # No tools registered
|
|
|
|
with patch("app.ai.agent_loop.llm_complete", new_callable=AsyncMock) as mock_llm:
|
|
mock_llm.side_effect = [
|
|
_make_llm_response(
|
|
content="Calling unknown tool",
|
|
tool_calls=[_make_tool_call(name="nonexistent", arguments={})],
|
|
cost_usd=0.001,
|
|
),
|
|
_make_llm_response(content="OK", cost_usd=0.001),
|
|
]
|
|
|
|
result = await run_react_loop(
|
|
agent_definition=agent_def,
|
|
messages=[{"role": "user", "content": "Call unknown tool"}],
|
|
tools=[{"type": "function", "function": {"name": "nonexistent"}}],
|
|
tool_registry=registry,
|
|
db=mock_db,
|
|
tenant_id=tenant_id,
|
|
user_id=user_id,
|
|
)
|
|
|
|
assert result.status == "completed"
|
|
assert result.steps[0].action == "nonexistent"
|
|
assert "not found" in (result.steps[0].observation or "").lower()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_multiple_tool_calls_per_step(agent_def, tenant_id, user_id, mock_db):
|
|
"""Multiple tool calls in a single LLM response are all executed."""
|
|
handler_a = AsyncMock(return_value="Result A")
|
|
handler_b = AsyncMock(return_value="Result B")
|
|
registry = _make_tool_registry({"tool_a": handler_a, "tool_b": handler_b})
|
|
|
|
with patch("app.ai.agent_loop.llm_complete", new_callable=AsyncMock) as mock_llm:
|
|
mock_llm.side_effect = [
|
|
_make_llm_response(
|
|
content="Calling two tools",
|
|
tool_calls=[
|
|
_make_tool_call(call_id="c1", name="tool_a", arguments={}),
|
|
_make_tool_call(call_id="c2", name="tool_b", arguments={}),
|
|
],
|
|
cost_usd=0.005,
|
|
),
|
|
_make_llm_response(content="Done with both", cost_usd=0.002),
|
|
]
|
|
|
|
result = await run_react_loop(
|
|
agent_definition=agent_def,
|
|
messages=[{"role": "user", "content": "Call two tools"}],
|
|
tools=[
|
|
{"type": "function", "function": {"name": "tool_a"}},
|
|
{"type": "function", "function": {"name": "tool_b"}},
|
|
],
|
|
tool_registry=registry,
|
|
db=mock_db,
|
|
tenant_id=tenant_id,
|
|
user_id=user_id,
|
|
)
|
|
|
|
assert result.status == "completed"
|
|
assert result.steps_taken == 2
|
|
# Step 1 has two tool calls → two step entries
|
|
assert len(result.steps) == 3 # 2 tool steps + 1 final step
|
|
handler_a.assert_awaited_once()
|
|
handler_b.assert_awaited_once()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_react_result_dataclass_fields():
|
|
"""ReActResult and ReActStep have correct default values."""
|
|
step = ReActStep(
|
|
step_number=1,
|
|
thought="test",
|
|
action=None,
|
|
action_input=None,
|
|
observation=None,
|
|
cost_usd=0.0,
|
|
timestamp="2025-01-01T00:00:00+00:00",
|
|
)
|
|
assert step.step_number == 1
|
|
assert step.thought == "test"
|
|
|
|
result = ReActResult(final_content="hello")
|
|
assert result.final_content == "hello"
|
|
assert result.steps == []
|
|
assert result.total_cost_usd == 0.0
|
|
assert result.steps_taken == 0
|
|
assert result.status == "completed"
|
|
assert result.error is None
|