agents/multi/graph.py

"""Parent orchestration graph for the harmonic analysis multi-agent system."""

import json
from typing import Annotated, TypedDict

from langchain_core.language_models import BaseChatModel
from langchain_core.messages import HumanMessage
from langchain_core.tools import BaseTool
from langgraph.graph import END, START, StateGraph
from langgraph.graph.state import CompiledStateGraph
from pydantic import BaseModel

from .analysis import build_analysis_subgraph
from .research import build_research_subgraph

_MIN_SONGS = 10
_INITIAL_LIMIT = 10
_LIMIT_INCREMENT = 10
_MAX_LIMIT = 50


class HarmonicInput(BaseModel):
    user_input: str


def _merge_researched(existing: list[dict], new: list[dict]) -> list[dict]:
    seen = {(s["title"], s["artist"]) for s in existing}
    return existing + [s for s in new if (s["title"], s["artist"]) not in seen]


class HarmonicState(TypedDict):
    user_input: str
    limit: int
    originality_score: float | None
    neighbours: list[dict]
    researched_neighbours: Annotated[list[dict], _merge_researched]
    songs_sent_to_research: list[dict]  # deterministic parent-owned log of dispatched songs
    songs_to_research: list[dict]       # set by parent before calling research subgraph


class HarmonicOutput(TypedDict):
    response: str


def build_harmonic_graph(
    analysis_llm: BaseChatModel,
    research_llm: BaseChatModel,
    mcp_tools: list[BaseTool],
    search_tools: list[BaseTool],
    min_songs: int = _MIN_SONGS,
    presenter_llm: BaseChatModel | None = None,
) -> CompiledStateGraph:
    """Build the harmonic analysis → research multi-agent graph.

    :param analysis_llm: LLM for harmonic analysis (must support tool calling).
    :param research_llm: LLM for internet research (must support tool calling).
    :param mcp_tools: Harmonic analysis MCP tools.
    :param search_tools: Web search tools for research.
    :param min_songs: Minimum well-known songs the researcher must find before the graph ends.
    :param presenter_llm: LLM for generating the final summary. Defaults to analysis_llm.
    """
    analysis_subgraph = build_analysis_subgraph(analysis_llm, mcp_tools)
    research_subgraph = build_research_subgraph(research_llm, search_tools)
    presenter_llm = presenter_llm or analysis_llm

    def prepare_analysis(state: HarmonicState) -> dict:
        return {
            "limit": _INITIAL_LIMIT,
            "originality_score": None,
            "neighbours": [],
            "researched_neighbours": [],
            "songs_sent_to_research": [],
        }

    def prepare_research(state: HarmonicState) -> dict:
        already_sent = {(s["title"], s["artist"]) for s in state["songs_sent_to_research"]}
        to_research = [
            n for n in state["neighbours"]
            if (n["title"], n["artist"]) not in already_sent
        ]
        return {
            "songs_to_research": to_research,
            "songs_sent_to_research": state["songs_sent_to_research"] + to_research,
        }

    def widen(state: HarmonicState) -> dict:
        return {
            "limit": state["limit"] + _LIMIT_INCREMENT,
            "originality_score": None,
            "neighbours": [],
        }

    def check_well_known(state: HarmonicState) -> str:
        well_known = sum(
            1 for s in state["researched_neighbours"]
            if s.get("chart_peak") is not None or s.get("is_famous_artist")
        )
        if well_known >= min_songs or state["limit"] >= _MAX_LIMIT:
            return "present"
        return "widen"

    def present(state: HarmonicState) -> dict:
        similarity_by_song = {(n["title"], n["artist"]): n.get("similarity") for n in state["neighbours"]}
        well_known = []
        for s in state["researched_neighbours"]:
            if s.get("chart_peak") is None and not s.get("is_famous_artist"):
                continue
            similarity = similarity_by_song.get((s["title"], s["artist"]))
            if similarity is None:
                continue
            well_known.append({**s, "similarity": similarity})
        prompt = (
            f"The chord sequence has an originality score of {state['originality_score']:.4f} "
            f"(0 = many harmonic matches, 1 = unique).\n\n"
            f"The following well-known songs were found to be harmonically similar:\n"
            f"{json.dumps(well_known, indent=2)}\n\n"
            "Write a short, readable summary for a musician. List each song with its similarity score "
            "and any interesting facts from the research (chart position, chart name, artist notes). "
            "Do not reproduce raw JSON — write in natural prose."
        )
        response = presenter_llm.invoke([HumanMessage(content=prompt)])
        return {"response": response.content}

    graph = StateGraph(HarmonicState, input=HarmonicInput, output=HarmonicOutput)
    graph.add_node("prepare_analysis", prepare_analysis)
    graph.add_node("analysis", analysis_subgraph)
    graph.add_node("prepare_research", prepare_research)
    graph.add_node("research", research_subgraph)
    graph.add_node("widen", widen)
    graph.add_node("present", present)

    graph.add_edge(START, "prepare_analysis")
    graph.add_edge("prepare_analysis", "analysis")
    graph.add_edge("analysis", "prepare_research")
    graph.add_edge("prepare_research", "research")
    graph.add_conditional_edges("research", check_well_known, ["widen", "present"])
    graph.add_edge("widen", "analysis")
    graph.add_edge("present", END)

    return graph.compile()