src/chemgraph/graphs/graspa_mcp.py

from typing import Union
from functools import partial

from langgraph.prebuilt import ToolNode
from langgraph.graph import StateGraph, END
from langchain_openai import ChatOpenAI
from langgraph.checkpoint.memory import MemorySaver
from langgraph.types import Send

from chemgraph.utils.logging_config import setup_logger
from chemgraph.state.graspa_state import (
    ExecutorState,
    PlannerState,
    PlannerResponse,
)
from chemgraph.prompt.graspa_prompt import (
    planner_prompt,
    executor_prompt,
    analyst_prompt,
)

logger = setup_logger(__name__)


def planner_agent(
    state: PlannerState,
    llm: ChatOpenAI,
    system_prompt: str,
):
    """Plan the next gRASPA MCP workflow step.

    Parameters
    ----------
    state : PlannerState
        Current planner state containing messages and executor results.
    llm : ChatOpenAI
        Chat model used for planning.
    system_prompt : str
        Planner system prompt.

    Returns
    -------
    dict
        Planner state update containing messages, next step, and tasks.
    """
    executor_outputs = state.get("executor_results", [])
    content_block = f"Current Conversation History: {state['messages']}"
    if executor_outputs:
        results_text = "\n".join(executor_outputs)
        content_block += (
            f"\n\n### UPDATED: Results from Executor Tasks ###\n{results_text}"
        )

    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "assistant", "content": content_block},
    ]

    structured_llm = llm.with_structured_output(PlannerResponse)
    response_obj = structured_llm.invoke(messages)
    print(f"PLANNER: {response_obj.model_dump_json()}")
    return {
        "messages": [response_obj.thought_process],
        "next_step": response_obj.next_step,
        "tasks": response_obj.tasks if response_obj.tasks else [],
    }


def unified_planner_router(state: PlannerState) -> Union[str, list[Send]]:
    """Route based on the planner's structured ``next_step``.

    Parameters
    ----------
    state : PlannerState
        Current planner state.

    Returns
    -------
    str or list[Send]
        Next node name, ``END``, or fan-out executor sends.
    """
    next_step = state.get("next_step")

    if next_step == "executor_subgraph":
        tasks = state.get("tasks", [])
        return [
            Send(
                "executor_subgraph",
                {
                    "executor_id": f"worker_{getattr(t, 'task_index', i + 1)}",
                    "messages": [getattr(t, 'prompt')],
                },
            )
            for i, t in enumerate(tasks)
        ]

    elif next_step == "insight_analyst":
        return "insight_analyst"

    return END


async def executor_model_node(
    state: ExecutorState,
    llm: ChatOpenAI,
    system_prompt: str,
    tools: list,
):
    """Run the reasoning step for a single gRASPA executor.

    Parameters
    ----------
    state : ExecutorState
        Local executor state.
    llm : ChatOpenAI
        Chat model used by the executor.
    system_prompt : str
        Executor system prompt.
    tools : list
        Tools available to the executor.

    Returns
    -------
    dict
        Executor state update containing the model response.
    """
    messages = [{"role": "system", "content": system_prompt}] + state["messages"]

    # Flatten MCP/LangChain content blocks to plain text before ChatOpenAI
    for m in messages:
        content = m.get("content") if isinstance(m, dict) else getattr(m, "content", None)
        if isinstance(content, list):
            text = "\n".join(
                block.get("text", str(block)) if isinstance(block, dict) else str(block)
                for block in content
            )
            if isinstance(m, dict):
                m["content"] = text
            else:
                m.content = text

    llm_with_tools = llm.bind_tools(tools)
    response = await llm_with_tools.ainvoke(messages)

    return {"messages": [response]}

def route_executor(state: ExecutorState):
    """Route executor output to tools or completion.

    Parameters
    ----------
    state : ExecutorState
        Local executor state.

    Returns
    -------
    str
        ``"tools"`` when tool calls are present, otherwise ``"done"``.
    """
    messages = state["messages"]
    last_message = messages[-1]
    if hasattr(last_message, "tool_calls") and last_message.tool_calls:
        return "tools"
    return "done"


def format_executor_output(state: ExecutorState) -> PlannerState:
    """Convert local executor state into a global planner update.

    Parameters
    ----------
    state : ExecutorState
        Local executor state at subgraph completion.

    Returns
    -------
    PlannerState
        Planner update containing executor results and logs.
    """
    executor_id = state["executor_id"]
    final_message = state["messages"][-1].content
    full_history = state["messages"]

    return {
        "executor_results": [f"[{executor_id}] Result: {final_message}"],
        "executor_logs": {executor_id: full_history},
    }


def construct_executor_subgraph(llm: ChatOpenAI, tools: list, system_prompt: str):
    """Build the reusable executor subgraph.

    Parameters
    ----------
    llm : ChatOpenAI
        Chat model used by executor agents.
    tools : list
        Tools available to executor agents.
    system_prompt : str
        Executor system prompt.

    Returns
    -------
    CompiledStateGraph
        Compiled executor subgraph.
    """
    workflow = StateGraph(ExecutorState)
    workflow.add_node(
        "executor_agent",
        partial(executor_model_node, llm=llm, system_prompt=system_prompt, tools=tools),
    )
    workflow.add_node("tools", ToolNode(tools))

    # Format the output back to the Main State schema
    workflow.add_node("finalize", format_executor_output)

    # Edges
    workflow.set_entry_point("executor_agent")
    workflow.add_conditional_edges(
        "executor_agent",
        route_executor,
        {"tools": "tools", "done": "finalize"},
    )
    workflow.add_edge("tools", "executor_agent")
    workflow.add_edge("finalize", END)

    return workflow.compile()


def insight_analyst_node(
    state: PlannerState,
    llm: ChatOpenAI,
    tools: list,
    system_prompt: str,
):
    """Analyze gathered executor results.

    Parameters
    ----------
    state : PlannerState
        Planner state containing executor results.
    llm : ChatOpenAI
        Chat model used by the analyst.
    tools : list
        Analysis tools available to the analyst.
    system_prompt : str
        Analyst system prompt.

    Returns
    -------
    dict
        Planner state update containing the analyst response.
    """
    results_text = "\n".join(state["executor_results"])
    messages = [
        {"role": "system", "content": system_prompt},
        {
            "role": "assistant",
            "content": f"Previous Messages: {state['messages']}\n\nExecution Results:\n{results_text}",
        },
    ]

    # Simple single-step analyst (can be expanded to loop like Executor)
    response = llm.bind_tools(tools).invoke(messages)
    return {"messages": [response]}


def route_analyst(state: PlannerState):
    """Route analyst output to tools or back to the planner.

    Parameters
    ----------
    state : PlannerState
        Planner state containing the analyst's latest message.

    Returns
    -------
    str
        ``"analyst_tools"`` when tool calls are present, otherwise
        ``"Planner"``.
    """
    last_msg = state["messages"][-1]

    # If the Analyst wants to use Python/Pandas -> Go to ToolNode
    if hasattr(last_msg, "tool_calls") and last_msg.tool_calls:
        return "analyst_tools"

    # If the Analyst is done -> Go back to Planner (who will then trigger FINISH)
    return "Planner"


def construct_graspa_mcp_graph(
    llm: ChatOpenAI,
    planner_prompt: str = planner_prompt,
    executor_prompt: str = executor_prompt,
    analyst_prompt: str = analyst_prompt,
    executor_tools: list = None,
    analysis_tools: list = None,
):
    """Construct the gRASPA MCP map-reduce graph.

    Parameters
    ----------
    llm : ChatOpenAI
        Chat model shared by planner, executors, and analyst.
    planner_prompt : str, optional
        Planner system prompt.
    executor_prompt : str, optional
        Executor system prompt.
    analyst_prompt : str, optional
        Analyst system prompt.
    executor_tools : list, optional
        Tools available to executor subgraphs.
    analysis_tools : list, optional
        Tools available to the analyst node.

    Returns
    -------
    CompiledStateGraph
        Compiled gRASPA MCP graph.
    """
    checkpointer = MemorySaver()

    # Create the Executor subgraph
    executor_subgraph = construct_executor_subgraph(
        llm,
        executor_tools,
        executor_prompt,
    )

    # Create the Main Manager Graph
    graph_builder = StateGraph(PlannerState)

    # Planner agent
    graph_builder.add_node(
        "Planner",
        lambda state: planner_agent(
            state,
            llm,
            planner_prompt,
        ),
    )

    # Executor agent
    graph_builder.add_node("executor_subgraph", executor_subgraph)

    # Analyst agent
    graph_builder.add_node(
        "insight_analyst",
        partial(
            insight_analyst_node,
            llm=llm,
            tools=analysis_tools,
            system_prompt=analyst_prompt,
        ),
    )
    # Tool nodes
    graph_builder.add_node("analyst_tools", ToolNode(analysis_tools))

    # -- Edges --
    graph_builder.set_entry_point("Planner")

    # Conditional Edge: Planner -> [Send(executor)...]
    graph_builder.add_conditional_edges(
        "Planner",
        unified_planner_router,
        [
            "insight_analyst",
            "executor_subgraph",
            END,
        ],
    )
    graph_builder.add_edge("executor_subgraph", "Planner")
    graph_builder.add_conditional_edges(
        "insight_analyst",
        route_analyst,
        {
            "analyst_tools": "analyst_tools",
            "Planner": END,
        },
    )
    graph_builder.add_edge("analyst_tools", "insight_analyst")

    return graph_builder.compile(checkpointer=checkpointer)