main.py

"""
Entry point — interactive CLI for the multi-agent RAG research system.

Usage:
    python main.py

Or import and use programmatically:
    from main import run_query
    result = run_query("Summarise arXiv:2310.06825 and generate a minimal training loop")
"""

from __future__ import annotations

import os
from dotenv import load_dotenv

load_dotenv()

from langchain_core.messages import HumanMessage

from supervisor import build_supervisor_graph


def run_query(user_input: str, verbose: bool = True) -> str:
    """
    Run a single query through the multi-agent system.
    Returns the final AI response as a string.
    """
    graph = build_supervisor_graph()

    initial_state = {
        "messages": [HumanMessage(content=user_input)],
        "next_agent": "",
        "iteration": 0,
    }

    if verbose:
        print("\n" + "=" * 60)
        print("USER:", user_input)
        print("=" * 60)

    final_state = graph.invoke(initial_state)

    # The last message in state is the final answer
    from langchain_core.messages import AIMessage
    last_ai = next(
        (m for m in reversed(final_state["messages"]) if isinstance(m, AIMessage)),
        None,
    )
    answer = last_ai.content if last_ai else "(No response generated)"

    if verbose:
        print("\nFINAL ANSWER:\n")
        print(answer)

    return answer


# ---------------------------------------------------------------------------
# Interactive CLI
# ---------------------------------------------------------------------------

EXAMPLES = [
    "Summarise the paper https://arxiv.org/abs/2310.06825",
    "Summarise https://arxiv.org/abs/1706.03762 (Attention is All You Need)",
    "Fetch https://github.com/karpathy/minGPT and generate a minimal GPT training script",
    "Summarise arXiv:2310.06825 and then generate Python code that implements the core idea",
]


def main() -> None:
    if not os.getenv("OPENAI_API_KEY"):
        raise EnvironmentError(
            "OPENAI_API_KEY is not set. Copy .env.example to .env and add your key."
        )

    print("\n" + "=" * 60)
    print("  Multi-Agent RAG Research System")
    print("  (paper summariser + code generator)")
    print("=" * 60)
    print("\nExample queries:")
    for i, ex in enumerate(EXAMPLES, 1):
        print(f"  {i}. {ex}")
    print("\nType 'quit' to exit.\n")

    while True:
        try:
            user_input = input("You: ").strip()
        except (EOFError, KeyboardInterrupt):
            print("\nGoodbye!")
            break

        if not user_input:
            continue
        if user_input.lower() in {"quit", "exit", "q"}:
            print("Goodbye!")
            break

        # Allow picking an example by number
        if user_input.isdigit() and 1 <= int(user_input) <= len(EXAMPLES):
            user_input = EXAMPLES[int(user_input) - 1]
            print(f"Using example: {user_input}\n")

        run_query(user_input, verbose=True)
        print()


if __name__ == "__main__":
    main()