wehub-resource-sync
136 lines
7.1 KiB
Python
136 lines
7.1 KiB
Python
import asyncio
|
|
import tempfile
|
|
import unittest
|
|
|
|
from agent_runtime.context_compactor import ContextCompactor
|
|
from agent_runtime.llm import AssistantMessage
|
|
from agent_runtime.loop import AgentLoop
|
|
from agent_runtime.models import ToolCall, ToolResult
|
|
from agent_runtime.prompts import PromptBuilder
|
|
from agent_runtime.session_index import SessionIndex
|
|
from agent_runtime.tool_executor import ToolExecutor
|
|
from agent_runtime.tools import ToolArgumentSchema, ToolRegistry, ToolSpec
|
|
|
|
|
|
class FakeLLM:
|
|
def __init__(self, replies):
|
|
self.replies = list(replies)
|
|
|
|
async def complete(self, messages, tools):
|
|
return self.replies.pop(0)
|
|
|
|
|
|
class FailingLLM:
|
|
async def complete(self, messages, tools):
|
|
raise RuntimeError("provider returned invalid response shape")
|
|
|
|
|
|
class AgentLoopTests(unittest.IsolatedAsyncioTestCase):
|
|
async def test_no_tool_call_finishes(self):
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
index = SessionIndex(tmp)
|
|
registry = ToolRegistry([])
|
|
loop = AgentLoop(index, PromptBuilder(f"{tmp}/prompts", index, registry), registry, ToolExecutor(registry, index), FakeLLM([AssistantMessage(text="done")]))
|
|
events = [event async for event in loop.stream_events("hi")]
|
|
self.assertEqual(events[-2]["type"], "done")
|
|
turn_id = events[0]["turn_id"]
|
|
self.assertTrue(all(event.get("turn_id") == turn_id for event in events))
|
|
log_text = (index.logs_dir / "loop_history.jsonl").read_text(encoding="utf-8")
|
|
self.assertIn("assistant_finished_without_tools", log_text)
|
|
|
|
|
|
async def test_turn_record_follows_session_created_by_tool(self):
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
index = SessionIndex(tmp)
|
|
old = index.create(idea="old")
|
|
|
|
def create_actual(args):
|
|
record = index.create(idea="actual")
|
|
return ToolResult("create_actual", True, record["session_id"])
|
|
|
|
registry = ToolRegistry([ToolSpec("create_actual", "Create actual session", create_actual, schema={})])
|
|
llm = FakeLLM([AssistantMessage(tool_calls=[ToolCall(name="create_actual", arguments={})]), AssistantMessage(text="finished")])
|
|
loop = AgentLoop(index, PromptBuilder(f"{tmp}/prompts", index, registry), registry, ToolExecutor(registry, index), llm)
|
|
events = [event async for event in loop.stream_events("start new project")]
|
|
active = index.active()
|
|
self.assertNotEqual(active["session_id"], old["session_id"])
|
|
self.assertEqual(len(index.get(active["session_id"])["recent_turn_records"]), 1)
|
|
self.assertEqual(index.get(old["session_id"])["recent_turn_records"], [])
|
|
self.assertEqual(events[-1]["session"]["active_session_id"], active["session_id"])
|
|
|
|
|
|
async def test_tool_progress_streams_before_tool_result(self):
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
index = SessionIndex(tmp)
|
|
release = asyncio.Event()
|
|
|
|
async def slow_tool(args, runtime):
|
|
runtime.emit_progress("started", stage="running")
|
|
await release.wait()
|
|
return ToolResult("slow_tool", True, "done")
|
|
|
|
registry = ToolRegistry([ToolSpec("slow_tool", "Slow tool", slow_tool, schema={})])
|
|
llm = FakeLLM([AssistantMessage(tool_calls=[ToolCall(name="slow_tool", arguments={})]), AssistantMessage(text="finished")])
|
|
loop = AgentLoop(index, PromptBuilder(f"{tmp}/prompts", index, registry), registry, ToolExecutor(registry, index), llm)
|
|
agen = loop.stream_events("start")
|
|
seen = []
|
|
while True:
|
|
event = await asyncio.wait_for(anext(agen), timeout=1)
|
|
seen.append(event["type"])
|
|
if event["type"] == "tool_progress":
|
|
self.assertFalse(release.is_set())
|
|
break
|
|
release.set()
|
|
async for event in agen:
|
|
seen.append(event["type"])
|
|
self.assertLess(seen.index("tool_progress"), seen.index("tool_result"))
|
|
|
|
|
|
async def test_preflight_compact_summarizes_old_history(self):
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
index = SessionIndex(tmp)
|
|
registry = ToolRegistry([])
|
|
compactor = ContextCompactor(None, token_threshold=200, buffer_tokens=0, preserve_last_n=2, summary_max_chars=2000)
|
|
loop = AgentLoop(index, PromptBuilder(f"{tmp}/prompts", index, registry), registry, ToolExecutor(registry, index), FakeLLM([AssistantMessage(text="after compact")]), compactor)
|
|
loop.history = [
|
|
{"role": "user", "content": "old request " + "x" * 1200},
|
|
{"role": "assistant", "content": "old answer " + "y" * 1200},
|
|
{"role": "user", "content": "recent request"},
|
|
{"role": "assistant", "content": "recent answer"},
|
|
]
|
|
events = [event async for event in loop.stream_events("continue")]
|
|
self.assertIn("compact", [event.get("phase") for event in events if event["type"] == "status"])
|
|
session = index.active()
|
|
self.assertIn("Reference Context Only", session["compacted_summary"])
|
|
self.assertGreaterEqual(session["compacted_turns"], 1)
|
|
self.assertTrue(session["compaction_snapshots"])
|
|
self.assertEqual(loop.history[0]["role"], "system")
|
|
self.assertIn("after compact", loop.history[-1]["content"])
|
|
self.assertNotIn("old request", index.memory_text())
|
|
|
|
|
|
async def test_llm_sampling_error_yields_error_without_crashing_loop(self):
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
index = SessionIndex(tmp)
|
|
registry = ToolRegistry([])
|
|
loop = AgentLoop(index, PromptBuilder(f"{tmp}/prompts", index, registry), registry, ToolExecutor(registry, index), FailingLLM())
|
|
events = [event async for event in loop.stream_events("start")]
|
|
self.assertTrue(any(event["type"] == "error" and event.get("metadata", {}).get("error_type") == "llm_sampling_failed" for event in events))
|
|
self.assertEqual(events[-2]["type"], "done")
|
|
self.assertEqual(events[-1]["type"], "session")
|
|
self.assertEqual(index.active()["recent_turn_records"][-1]["status"], "failed")
|
|
|
|
async def test_tool_call_continues_then_finishes(self):
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
index = SessionIndex(tmp)
|
|
|
|
def hello(args):
|
|
return ToolResult("hello", True, "hello result")
|
|
|
|
registry = ToolRegistry([ToolSpec("hello", "Say hello", hello, schema={"name": ToolArgumentSchema(str, False, "x")})])
|
|
llm = FakeLLM([AssistantMessage(tool_calls=[ToolCall(name="hello", arguments={})]), AssistantMessage(text="finished")])
|
|
loop = AgentLoop(index, PromptBuilder(f"{tmp}/prompts", index, registry), registry, ToolExecutor(registry, index), llm)
|
|
events = [event async for event in loop.stream_events("start")]
|
|
self.assertTrue(any(event["type"] == "tool_result" for event in events))
|
|
self.assertEqual(events[-2]["assistant"], "finished")
|