lambdagent/tests/test_agentruntime.py

"""Tests for lambdagent.agentruntime components."""
import pytest
from lambdagent.agentruntime.action_parser import ActionParser, Action, ParseError
from lambdagent.agentruntime.termination import TerminationOracle
from lambdagent.agentruntime.memory_backend import LocalMemory, SQLiteMemory, MemoryBackend
from lambdagent.agentruntime.trace_store import TraceStore, TraceRecord
from lambdagent.agentruntime.config import RuntimeConfig, LLMConfig, MemoryConfig


# ActionParser tests

def test_parser_json_block():
    p = ActionParser(["search", "terminate"])
    action = p.parse('Let me search. ```json\n{"action": "search", "input": {"q": "test"}}\n```')
    assert action.tool == "search"


def test_parser_inline_json():
    p = ActionParser(["calc", "terminate"])
    action = p.parse('I will calculate. {"action": "calc", "input": {"expr": "1+1"}}')
    assert action.tool == "calc"


def test_parser_implicit_terminate():
    p = ActionParser(["search", "terminate"])
    action = p.parse("Final Answer: The result is 42.")
    assert action.tool == "terminate"
    assert "42" in str(action.input)


def test_parser_unknown_tool():
    p = ActionParser(["search"])
    with pytest.raises(ParseError):
        p.parse('{"action": "unknown_tool"}')


# TerminationOracle tests

def test_termination_implicit():
    oracle = TerminationOracle()
    assert oracle.should_stop("Task complete. Here is the result.", None, 0) == True
    assert oracle.should_stop("I need to search more.", None, 0) == False


def test_termination_disabled():
    oracle = TerminationOracle(implicit_detection=False)
    assert oracle.should_stop("Final Answer: 42", None, 0) == False


# LocalMemory tests

def test_local_memory_write_read():
    mem = LocalMemory(size=5, ttl=3600)
    mem.write("k1", "v1")
    assert mem.read("k1") == "v1"


def test_local_memory_lru():
    mem = LocalMemory(size=3, ttl=3600)
    for i in range(5):
        mem.write(f"k{i}", f"v{i}")
    assert mem.read("k0") is None  # evicted
    assert mem.read("k4") == "v4"


def test_local_memory_recent():
    mem = LocalMemory(size=10, ttl=3600)
    mem.write("a", "1")
    mem.write("b", "2")
    recent = mem.read_recent(5)
    assert len(recent) == 2


def test_local_memory_clear():
    mem = LocalMemory()
    mem.write("k", "v")
    mem.clear()
    assert mem.read("k") is None


# SQLiteMemory tests

def test_sqlite_memory():
    mem = SQLiteMemory(size=5, ttl=3600, db_path=":memory:")
    mem.write("k1", "v1")
    assert mem.read("k1") == "v1"
    mem.clear()
    assert mem.read("k1") is None


# MemoryBackend factory

def test_memory_factory():
    mem = MemoryBackend.create({"strategy": "local", "size": 10, "ttl": 3600})
    assert isinstance(mem, LocalMemory)

    mem2 = MemoryBackend.create({"strategy": "sqlite", "size": 10, "ttl": 3600})
    assert isinstance(mem2, SQLiteMemory)


# TraceStore tests

def test_trace_store():
    ts = TraceStore()
    ts.append(TraceRecord(step=0, term_name="think", term_type="Lam", duration_ms=100))
    ts.append(TraceRecord(step=1, term_name="tool", term_type="Tool", duration_ms=50))
    assert len(ts.get_all()) == 2
    assert ts.get_step(0).term_name == "think"
    stats = ts.stats()
    assert stats.total_steps == 2
    assert stats.total_time_ms == 150


def test_trace_timeline():
    ts = TraceStore()
    ts.append(TraceRecord(step=0, term_name="think", term_type="Lam", duration_ms=100))
    timeline = ts.to_timeline()
    assert "think" in timeline


def test_trace_json():
    ts = TraceStore()
    ts.append(TraceRecord(step=0, term_name="t", term_type="Lam"))
    j = ts.to_json()
    assert '"term_name": "t"' in j


# RuntimeConfig tests

def test_config_defaults():
    cfg = RuntimeConfig()
    assert cfg.llm.model == "claude-sonnet-4-20250514"
    assert cfg.memory.strategy == "local"
    assert cfg.react.max_steps == 10