agents/multi/research.py

"""Research agent subgraph."""

from typing import Annotated, TypedDict

from langchain_core.language_models import BaseChatModel
from langchain_core.messages import HumanMessage, RemoveMessage, SystemMessage
from langchain_core.tools import BaseTool
from langgraph.graph import END, START, StateGraph
from langgraph.graph.message import add_messages
from langgraph.graph.state import CompiledStateGraph
from langgraph.prebuilt import ToolNode
from pydantic import BaseModel

_SYSTEM = (
    "You are a music researcher. Search the web to find chart information for a song. "
    "Make sure it's actually the song by the artist in question, not a song of the same name "
    "by another artist. Form queries like: '\"Let It Be\" Beatles UK Singles Chart peak position'."
)


class ResearchState(TypedDict):
    songs_to_research: list[dict]      # received from parent
    current_song_index: int            # internal loop counter
    researched_neighbours: list[dict]  # returned to parent (merged by parent reducer)
    messages: Annotated[list, add_messages]  # private, reset per song


class _ResearchedSong(BaseModel):
    title: str
    artist: str
    chart_peak: int | None = None
    chart_name: str | None = None
    is_famous_artist: bool | None = None
    notes: str | None = None


def build_research_subgraph(
    research_llm: BaseChatModel,
    search_tools: list[BaseTool],
) -> CompiledStateGraph:
    """Build the music research agent subgraph.

    :param research_llm: LLM for internet research (must support tool calling).
    :param search_tools: Web search tools for research.
    """
    research_llm_with_tools = research_llm.bind_tools(search_tools)

    def start(state: ResearchState) -> dict:
        return {"researched_neighbours": [], "current_song_index": 0}

    def prepare_song(state: ResearchState) -> dict:
        song = state["songs_to_research"][state["current_song_index"]]
        clear = [RemoveMessage(id=m.id) for m in state["messages"]]
        return {
            "messages": clear + [
                SystemMessage(content=_SYSTEM),
                HumanMessage(content=f"Research \"{song['title']}\" by {song['artist']}."),
            ]
        }

    def agent(state: ResearchState) -> dict:
        return {"messages": [research_llm_with_tools.invoke(state["messages"])]}

    def router(state: ResearchState) -> str:
        last = state["messages"][-1]
        return "tools" if getattr(last, "tool_calls", None) else "extract_song"

    def extract_song(state: ResearchState) -> dict:
        song = state["songs_to_research"][state["current_song_index"]]
        structured = research_llm.with_structured_output(_ResearchedSong)
        prompt = HumanMessage(content=(
            f"Summarise your findings for \"{song['title']}\" by {song['artist']}. "
            "If it charted, give the peak position and chart name. "
            "If not, note whether the artist is well-known."
        ))
        result: _ResearchedSong = structured.invoke(state["messages"] + [prompt])
        return {
            "researched_neighbours": state["researched_neighbours"] + [result.model_dump()],
            "current_song_index": state["current_song_index"] + 1,
        }

    def route_next(state: ResearchState) -> str:
        return "prepare_song" if state["current_song_index"] < len(state["songs_to_research"]) else END

    graph = StateGraph(ResearchState)
    graph.add_node("start", start)
    graph.add_node("prepare_song", prepare_song)
    graph.add_node("agent", agent)
    graph.add_node("tools", ToolNode(search_tools))
    graph.add_node("extract_song", extract_song)

    graph.add_edge(START, "start")
    graph.add_edge("start", "prepare_song")
    graph.add_edge("prepare_song", "agent")
    graph.add_conditional_edges("agent", router, ["tools", "extract_song"])
    graph.add_edge("tools", "agent")
    graph.add_conditional_edges("extract_song", route_next, ["prepare_song", END])

    return graph.compile()