graphs/workflows/orchestrated_v2.py

"""Orchestrated V2 workflow with tool-calling retrieval autonomy."""
from datetime import datetime
from langgraph.graph import StateGraph, END
from langchain_core.language_models.chat_models import BaseChatModel

from graphs.state import AgentState
from graphs.agents.classifier_agent import classifier_node
from graphs.agents.chat_tools_agent import chat_with_tools_node
from graphs.agents.summarizer_agent import summarizer_llm_node, summarizer_export_node
from graphs.prompts_v2 import load_v2_prompt
from typing import Callable


def _build_v2_workflow(
    classify_runner: Callable[[AgentState], AgentState],
    tools_agent_runner: Callable[[AgentState], AgentState],
    summarizer_llm_runner: Callable[[AgentState], AgentState],
    summarizer_export_runner: Callable[[AgentState], AgentState],
):
    workflow = StateGraph(AgentState)

    workflow.add_node("classify", classify_runner)
    workflow.add_node("tools_agent", tools_agent_runner)
    workflow.add_node("summarizer_llm", summarizer_llm_runner)
    workflow.add_node("summarizer_export", summarizer_export_runner)

    workflow.set_entry_point("classify")

    workflow.add_conditional_edges(
        "classify",
        lambda s: getattr(s.get("classification"), "classification", "CLASSIC"),
        {
            "CLASSIC": "tools_agent",
            "SUMMARIZE": "summarizer_llm",
            "UNKNOWN": "tools_agent",
        },
    )

    workflow.add_edge("tools_agent", END)

    workflow.add_edge("summarizer_llm", "summarizer_export")
    workflow.add_edge("summarizer_export", END)
    return workflow


def create_orchestrated_graph_v2(llm: BaseChatModel, checkpointer=None):
    # Lazy imports keep module importable in notebook contexts that do not
    # need runtime PDF generation dependencies.
    from tools.pdf import markdown_to_pdf
    from tools.storage import upload_pdf_to_supabase

    chat_prompt_v2_template = load_v2_prompt("chat_system.md")
    today_date = datetime.now().strftime("%d/%m/%Y")
    chat_prompt_v2 = chat_prompt_v2_template.replace("{{TODAY_DATE}}", today_date)
    tools_policy_v2 = load_v2_prompt("tools_policy.md")
    classifier_prompt_v2 = load_v2_prompt("classifier_system.md")
    summarizer_prompt_v2 = load_v2_prompt("summarizer_system.md")

    workflow = _build_v2_workflow(
        classify_runner=classifier_node(llm, system_prompt=classifier_prompt_v2),
        tools_agent_runner=chat_with_tools_node(
            llm,
            base_system_prompt=chat_prompt_v2,
            tools_policy_prompt=tools_policy_v2,
        ),
        summarizer_llm_runner=summarizer_llm_node(
            llm,
            system_prompt=summarizer_prompt_v2,
        ),
        summarizer_export_runner=summarizer_export_node(
            markdown_to_pdf=markdown_to_pdf,
            upload_pdf=upload_pdf_to_supabase,
        ),
    )
    return workflow.compile(checkpointer=checkpointer)