tests/test_memory_agents.py

"""Test persistent memory, goal memory, and agents."""

import sys
import os
import tempfile
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))

from splitbit_llm.memory.persistent import PersistentMemory
from splitbit_llm.memory.goal_memory import GoalMemory, Goal
from splitbit_llm.agents.agent_manager import AgentManager
from splitbit_llm.agents.planner_agent import PlannerAgent
from splitbit_llm.agents.coder_agent import CoderAgent
from splitbit_llm.agents.researcher_agent import ResearcherAgent
from splitbit_llm.agents.reviewer_agent import ReviewerAgent
from splitbit_llm.agents.executor_agent import ExecutorAgent


def test_persistent_memory():
    """Test persistent memory storage and recall."""
    with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as f:
        db_path = f.name

    mem = PersistentMemory(db_path=db_path)
    mem.set_session("test-session")

    # Store episodic
    mem.add_episodic("user", "What is Python?", importance=0.7)
    mem.add_episodic("assistant", "Python is a programming language.", importance=0.8)

    # Store semantic
    mem.add_semantic("Python is an interpreted programming language", confidence=0.9)

    # Recall
    episodic = mem.recall_episodic("Python", max_results=3)
    assert len(episodic) >= 1, "No episodic memories recalled"
    print(f"  Episodic recalled: {len(episodic)}")

    semantic = mem.recall_semantic("Python", max_results=3)
    assert len(semantic) >= 1, "No semantic memories recalled"
    print(f"  Semantic recalled: {len(semantic)}")

    # Context injection
    context = mem.get_context("Python programming")
    assert "Python" in context or "programming" in context, f"Context empty: {context}"
    print(f"  Context: {context[:80]}")

    # Persistence test
    mem2 = PersistentMemory(db_path=db_path)
    stats = mem2.get_stats()
    assert stats["episodic_total"] >= 2, f"Episodic not persisted: {stats}"
    assert stats["semantic_total"] >= 1, f"Semantic not persisted: {stats}"
    print(f"  Persisted: {stats['episodic_total']} episodic, {stats['semantic_total']} semantic")


def test_goal_memory():
    """Test goal creation, planning, and execution."""
    with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as f:
        db_path = f.name

    gm = GoalMemory(db_path=db_path)

    # Create goal
    goal = gm.create_goal("Build a web app", "Create a simple Flask web application", priority="high")
    assert goal.status == "pending"
    print(f"  Created goal: {goal.title} ({goal.id[:8]})")

    # Plan goal
    result = gm.plan_goal(goal.id, steps=[
        {"title": "Setup", "description": "Create project structure"},
        {"title": "Code", "description": "Write the Flask app"},
        {"title": "Test", "description": "Test the application"},
    ])
    assert result["success"]
    goal = gm.get_goal(goal.id)
    assert goal.status == "in_progress"
    assert len(goal.steps) == 3
    print(f"  Planned: {len(goal.steps)} steps, status: {goal.status}")

    # Execute steps
    r1 = gm.execute_step(goal.id, "Project structure created", success=True)
    assert r1["success"]
    r2 = gm.execute_step(goal.id, "Flask app written", success=True)
    assert r2["success"]
    r3 = gm.execute_step(goal.id, "Tests passed", success=True)
    assert r3["success"]

    goal = gm.get_goal(goal.id)
    assert goal.status == "completed"
    assert goal.progress() == 1.0
    print(f"  Completed: {goal.progress():.0%}")

    # Persistence
    gm2 = GoalMemory(db_path=db_path)
    stats = gm2.get_stats()
    assert stats["total"] == 1
    assert stats["completed"] == 1
    print(f"  Persisted: {stats}")


def test_agent_manager():
    """Test agent manager with 5 agents."""
    with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as f:
        goal_db = f.name
    with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as f:
        mem_db = f.name

    gm = GoalMemory(db_path=goal_db)
    pm = PersistentMemory(db_path=mem_db)

    # Mock generate function
    def mock_generate(prompt):
        if "planning" in prompt.lower():
            return '{"steps": [{"title": "Step 1", "description": "Do thing 1"}], "sub_goals": []}'
        return "Mock response for: " + prompt[:50]

    manager = AgentManager(gm, pm, generate_fn=mock_generate)
    assert len(manager.agents) == 5
    print(f"  Agents: {list(manager.agents.keys())}")

    # Create a project
    goal = manager.create_project("Test project", "A test project for validation", priority="high")
    print(f"  Created project: {goal.id[:8]}")

    # Get status
    status = manager.get_agent_status()
    assert len(status) == 5
    for name, info in status.items():
        assert info["name"] == name
        assert info["running"] == False  # not started yet
    print(f"  All 5 agents present: {list(status.keys())}")

    # Get stats
    stats = manager.get_stats()
    assert "agents" in stats
    assert "goals" in stats
    print(f"  Stats: {stats['goals']}")



def test_agent_specialization():
    """Test that agents have correct specializations."""
    with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as f:
        db_path = f.name

    gm = GoalMemory(db_path=db_path)

    planner = PlannerAgent(gm)
    coder = CoderAgent(gm)
    researcher = ResearcherAgent(gm)
    reviewer = ReviewerAgent(gm)
    executor = ExecutorAgent(gm)

    # Test coder can handle code goals
    code_goal = gm.create_goal("Write code", "Implement a Python function to sort data")
    assert coder._can_handle(code_goal), "Coder should handle code goals"
    assert not researcher._can_handle(code_goal), "Researcher should not handle code goals"
    print("  Coder handles code goals: OK")

    # Test researcher can handle research goals
    research_goal = gm.create_goal("Research topic", "Research the best sorting algorithms")
    assert researcher._can_handle(research_goal), "Researcher should handle research goals"
    print("  Researcher handles research goals: OK")

    # Test executor can handle execution goals
    exec_goal = gm.create_goal("Run tests", "Execute the test suite and deploy")
    assert executor._can_handle(exec_goal), "Executor should handle execution goals"
    print("  Executor handles execution goals: OK")



if __name__ == "__main__":
    print("Running memory & agents tests...")
    test_persistent_memory()
    print("  ✓ test_persistent_memory")
    test_goal_memory()
    print("  ✓ test_goal_memory")
    test_agent_manager()
    print("  ✓ test_agent_manager")
    test_agent_specialization()
    print("  ✓ test_agent_specialization")
    print("\nAll memory & agents tests passed!")