test.py

"""Quick smoke test for the multi-agent system.

Checks:
  1. GOOGLE_API_KEY is set and reachable
  2. Core imports work (state, tools, orchestrator)
  3. WorkflowState initialisation is correct
  4. Routing logic behaves as expected

Run:
    .venv/Scripts/python test.py
"""
from __future__ import annotations

import asyncio
import os
import sys

try:
    sys.stdout.reconfigure(encoding="utf-8", errors="replace")
except Exception:
    pass

from dotenv import load_dotenv

load_dotenv()

SEP = "=" * 60


def check_env() -> bool:
    key = os.getenv("GOOGLE_API_KEY", "")
    print(SEP)
    print("GOOGLE_API_KEY set :", bool(key), f"(...{key[-4:]})" if key else "")
    print("PLANNER_MODEL      :", os.getenv("PLANNER_MODEL", "gemini-2.5-flash"))
    print("DATABASE_URL       :", os.getenv("DATABASE_URL", "(not set)"))
    print("REDIS_URL          :", os.getenv("REDIS_URL", "(not set)"))
    print(SEP)
    if not key:
        print("[X] No GOOGLE_API_KEY found — set it in .env")
        return False
    return True


async def check_gemini(key: str) -> bool:
    import httpx

    model = os.getenv("PLANNER_MODEL", "gemini-2.5-flash")
    url = (
        f"https://generativelanguage.googleapis.com/v1beta/models/"
        f"{model}:generateContent?key={key}"
    )
    body = {"contents": [{"parts": [{"text": "Reply with exactly: API key works"}]}]}
    async with httpx.AsyncClient(timeout=15) as c:
        r = await c.post(url, json=body)

    print(f"\n[1] Gemini live call -> HTTP {r.status_code}")
    if r.status_code >= 300:
        print("    error:", r.text[:400])
        print("\n[X] Common causes:")
        print("    400 API_KEY_INVALID        -> wrong/expired key")
        print("    403 PERMISSION_DENIED      -> 'Generative Language API' not enabled,")
        print("                                  OR Google project access denied —")
        print("                                  create a new key at https://aistudio.google.com/apikey")
        print("    404 model not found        -> check PLANNER_MODEL in .env")
        print("    429                        -> rate/quota limit, try again shortly")
        return False

    try:
        reply = r.json()["candidates"][0]["content"]["parts"][0]["text"].strip()
    except Exception:
        reply = r.text[:200]
    print("    model reply:", reply)
    print("[OK] Gemini API key is working.\n")
    return True


def check_imports() -> bool:
    print("[2] Checking imports ...")
    try:
        from backend.state.graph_state import (  # noqa: F401
            WorkflowState, TaskStatus, StepStatus, AgentRole,
            create_initial_state, make_plan_step, make_agent_event,
        )
        from backend.agents.orchestrator import (  # noqa: F401
            build_workflow, route_after_executor,
            route_after_critic, route_after_planner,
        )
        print("    backend.state        OK")
        print("    backend.orchestrator OK")
    except ImportError as e:
        print(f"    [X] Import failed: {e}")
        return False

    try:
        from backend.tools.registry import execute_tool, calculate  # noqa: F401
        print("    backend.tools        OK")
    except ImportError as e:
        print(f"    [X] tools import failed: {e}")
        return False

    print("[OK] All imports succeeded.\n")
    return True


def check_state() -> bool:
    print("[3] WorkflowState smoke test ...")
    from backend.state.graph_state import (
        create_initial_state, TaskStatus, make_plan_step, StepStatus,
    )

    state = create_initial_state("Write a hello-world script")
    assert state["task"] == "Write a hello-world script"
    assert state["status"] == TaskStatus.PENDING
    assert state["plan"] == []
    assert state["iteration"] == 0
    assert state["total_tokens"] == 0
    assert len(state["events"]) == 1

    step = make_plan_step("s1", "Write code", "Create hello.py", tool="run_python")
    assert step["status"] == StepStatus.PENDING
    assert step["attempts"] == 0

    print("    initial state    OK")
    print("    make_plan_step   OK")
    print("[OK] State checks passed.\n")
    return True


def check_routing() -> bool:
    print("[4] Routing logic smoke test ...")
    from backend.state.graph_state import create_initial_state, TaskStatus, StepStatus
    from backend.agents.orchestrator import (
        route_after_executor, route_after_critic, route_after_planner,
    )

    def _state(status, plan=None, needs_replanning=False):
        s = create_initial_state("test")
        s["status"] = status
        s["plan"] = plan or []
        s["needs_replanning"] = needs_replanning
        return s

    # executor → critic when REFLECTING
    assert route_after_executor(_state(TaskStatus.REFLECTING)) == "critic"
    # executor → self when steps pending
    assert route_after_executor(
        _state(TaskStatus.EXECUTING, [{"step_id": "s1", "status": StepStatus.PENDING}])
    ) == "executor"
    # executor → end on FAILED
    assert route_after_executor(_state(TaskStatus.FAILED)) == "end"
    # critic → planner on needs_replanning
    assert route_after_critic(_state(TaskStatus.PLANNING, needs_replanning=True)) == "planner"
    # critic → memory_store when approved
    assert route_after_critic(_state(TaskStatus.COMPLETED)) == "memory_store"
    # planner → end with empty plan
    assert route_after_planner(_state(TaskStatus.EXECUTING, plan=[])) == "end"
    # planner → executor with valid plan
    assert route_after_planner(
        _state(TaskStatus.EXECUTING, [{"step_id": "s1", "status": StepStatus.PENDING}])
    ) == "executor"

    print("    route_after_executor OK")
    print("    route_after_critic   OK")
    print("    route_after_planner  OK")
    print("[OK] Routing checks passed.\n")
    return True


async def main() -> None:
    ok = check_env()
    if not ok:
        sys.exit(1)

    key = os.getenv("GOOGLE_API_KEY", "")
    gemini_ok = await check_gemini(key)

    imports_ok = check_imports()
    state_ok = check_state() if imports_ok else False
    routing_ok = check_routing() if imports_ok else False

    print(SEP)
    print("Summary")
    print(SEP)
    print("  Gemini API   :", "[OK]" if gemini_ok else "[FAIL]")
    print("  Imports      :", "[OK]" if imports_ok else "[FAIL]")
    print("  State        :", "[OK]" if state_ok else "[FAIL]")
    print("  Routing      :", "[OK]" if routing_ok else "[FAIL]")
    print(SEP)

    system_ok = imports_ok and state_ok and routing_ok
    if system_ok and gemini_ok:
        print("\nAll checks passed. The system is ready.")
    elif system_ok:
        print("\nSystem checks passed. Fix the Gemini API key to enable LLM calls.")
    else:
        print("\nSystem checks failed. See above for details.")
        sys.exit(1)


if __name__ == "__main__":
    asyncio.run(main())