| from __future__ import annotations |
|
|
| import pytest |
|
|
| from aura.core.agent_loop import ReActAgentLoop |
| from aura.core.event_bus import EventBus |
| from aura.core.llm_router import LLMResult |
| from aura.core.tools import ToolRegistry, ToolSpec |
|
|
|
|
| class FakeRouter: |
| def __init__(self, outputs): |
| self.outputs = outputs |
| self.calls = 0 |
|
|
| async def chat(self, messages): |
| output = self.outputs[self.calls] |
| self.calls += 1 |
| return LLMResult(ok=True, model="fake", content=output) |
|
|
|
|
| @pytest.mark.asyncio |
| async def test_agent_loop_executes_tool_then_finishes(): |
| registry = ToolRegistry() |
| registry.register( |
| ToolSpec( |
| name="echo", |
| description="Echo input", |
| tier=1, |
| arguments_schema={"type": "object"}, |
| return_schema={"type": "object"}, |
| handler=lambda args: args, |
| ) |
| ) |
| router = FakeRouter([ |
| '{"type":"tool","tool":"echo","arguments":{"value":42}}', |
| '{"type":"final","response":"finished"}', |
| ]) |
| loop = ReActAgentLoop(router=router, registry=registry, event_bus=EventBus()) |
| result = await loop.run("run echo") |
| assert result.ok is True |
| assert result.answer == "finished" |
| assert result.steps[0]["tool"] == "echo" |
|
|
|
|
| @pytest.mark.asyncio |
| async def test_agent_loop_routes_high_importance_through_ensemble(monkeypatch): |
| registry = ToolRegistry() |
| router = FakeRouter(['{"type":"final","response":"direct"}']) |
| loop = ReActAgentLoop(router=router, registry=registry, event_bus=EventBus()) |
|
|
| called = {"ensemble": False} |
|
|
| async def fake_ensemble_answer(task, importance_level=2, models=None, context=None): |
| called["ensemble"] = True |
| return type("R", (), {"synthesized_answer": "ensemble-final"})() |
|
|
| monkeypatch.setattr("aura.agents.ensemble.tools.ensemble_answer", fake_ensemble_answer) |
| result = await loop.run("Please research and decide the best option") |
| assert called["ensemble"] is True |
| assert result.answer == "ensemble-final" |
|
|
