File size: 2,038 Bytes
999bb04 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 | 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"
|