supervisor.py

"""
Supervisor (orchestrator) agent built with LangGraph.

Routes requests to:
  • "summariser"     → PaperSummarizerAgent
  • "code_explainer" → CodeExplainerAgent
  • "FINISH"         → return final answer
"""

from __future__ import annotations

import operator
from typing import Annotated, Literal, TypedDict

from langchain_core.messages import AIMessage, BaseMessage, SystemMessage
from langchain_openai import ChatOpenAI
from langchain_groq import ChatGroq
from langgraph.graph import END, StateGraph

from agents import build_code_explainer, build_paper_summariser


class SupervisorState(TypedDict):
    messages: Annotated[list[BaseMessage], operator.add]
    next_agent: str
    iteration: int


_SUPERVISOR_SYSTEM = """
You are a research orchestrator that coordinates two specialist agents:

  • summariser     – fetches and summarises research papers (arXiv, PDF)
  • code_explainer – finds GitHub repos linked in the paper, explains algorithms,
                     and produces how-to-run guides
  • FINISH         – use this when the task is fully complete

Given the conversation history, decide which agent to call next, OR output
FINISH if the task is fully complete.

Reply with EXACTLY one of: summariser | code_explainer | FINISH
No other text.
""".strip()

# _supervisor_llm = ChatOpenAI(model="gpt-4o", temperature=0)

_supervisor_llm = ChatGroq(model="llama-3.3-70b-versatile", temperature=0)

def supervisor_node(state: SupervisorState) -> SupervisorState:
    messages = [SystemMessage(content=_SUPERVISOR_SYSTEM)] + state["messages"]
    response = _supervisor_llm.invoke(messages)
    next_agent = response.content.strip()

    iteration = state.get("iteration", 0) + 1
    if iteration > 6 or next_agent not in {"summariser", "code_explainer"}:
        next_agent = "FINISH"

    return {**state, "next_agent": next_agent, "iteration": iteration}


_summariser    = build_paper_summariser()
_code_explainer = build_code_explainer()


def _run_agent(agent, state: SupervisorState) -> SupervisorState:
    result = agent.invoke({"messages": state["messages"]})
    agent_messages = result.get("messages", [])
    last = next(
        (m for m in reversed(agent_messages) if isinstance(m, AIMessage)),
        AIMessage(content="(no response)"),
    )
    return {**state, "messages": state["messages"] + [last]}


def summariser_node(state: SupervisorState) -> SupervisorState:
    return _run_agent(_summariser, state)


def code_explainer_node(state: SupervisorState) -> SupervisorState:
    return _run_agent(_code_explainer, state)


def route(state: SupervisorState) -> Literal["summariser", "code_explainer", "__end__"]:
    nxt = state.get("next_agent", "FINISH")
    if nxt == "summariser":
        return "summariser"
    if nxt == "code_explainer":
        return "code_explainer"
    return END


def build_supervisor_graph() -> StateGraph:
    graph = StateGraph(SupervisorState)

    graph.add_node("supervisor",     supervisor_node)
    graph.add_node("summariser",     summariser_node)
    graph.add_node("code_explainer", code_explainer_node)

    graph.set_entry_point("supervisor")
    graph.add_conditional_edges("supervisor", route)
    graph.add_edge("summariser",     "supervisor")
    graph.add_edge("code_explainer", "supervisor")

    return graph.compile()