src/chemgraph/graphs/single_agent_xanes.py

import os

from langgraph.graph import StateGraph, START, END
from langchain_openai import ChatOpenAI
from langgraph.checkpoint.memory import MemorySaver
from langgraph.prebuilt import ToolNode
from chemgraph.tools.cheminformatics_tools import (
    molecule_name_to_smiles,
    smiles_to_coordinate_file,
)
from chemgraph.tools.ase_tools import run_ase
from chemgraph.tools.xanes_tools import (
    run_xanes, 
    fetch_xanes_data, 
    plot_xanes_data,)
from chemgraph.schemas.agent_response import ResponseFormatter
from chemgraph.prompt.xanes_prompt import (
    xanes_single_agent_prompt,
    xanes_formatter_prompt,
)
from chemgraph.utils.logging_config import setup_logger
from chemgraph.state.state import State

logger = setup_logger(__name__)


def _tool_call_signature(tool_calls) -> tuple:
    """Create a comparable signature for a list of tool calls.

    Parameters
    ----------
    tool_calls : list
        Tool-call dictionaries from an AI message.

    Returns
    -------
    tuple
        Deterministic signature of tool names and arguments.
    """
    signature = []
    for call in tool_calls or []:
        name = call.get("name") if isinstance(call, dict) else None
        args = call.get("args", {}) if isinstance(call, dict) else {}
        if isinstance(args, dict):
            args_sig = tuple(sorted(args.items()))
        else:
            args_sig = str(args)
        signature.append((name, args_sig))
    return tuple(signature)


def _is_repeated_tool_cycle(messages) -> bool:
    """Detect if the most recent AI tool-call set repeats the previous one.

    Parameters
    ----------
    messages : list
        Message history to inspect.

    Returns
    -------
    bool
        ``True`` when the last two AI tool-call sets are identical.
    """
    ai_with_calls = []
    for message in messages:
        if hasattr(message, "tool_calls") and getattr(message, "tool_calls", None):
            ai_with_calls.append(message)

    if len(ai_with_calls) < 2:
        return False

    last_calls = _tool_call_signature(ai_with_calls[-1].tool_calls)
    prev_calls = _tool_call_signature(ai_with_calls[-2].tool_calls)
    return bool(last_calls) and last_calls == prev_calls


def route_tools(state: State):
    """Route to the 'tools' node if the last message has tool calls; otherwise, route to 'done'.

    Parameters
    ----------
    state : State
        The current state containing messages and remaining steps

    Returns
    -------
    str
        Either 'tools' or 'done' based on the state conditions
    """
    if isinstance(state, list):
        ai_message = state[-1]
    elif messages := state.get("messages", []):
        ai_message = messages[-1]
    else:
        raise ValueError(f"No messages found in input state to tool_edge: {state}")
    if hasattr(ai_message, "tool_calls") and len(ai_message.tool_calls) > 0:
        if not isinstance(state, list) and _is_repeated_tool_cycle(messages):
            return "done"
        return "tools"
    return "done"


def XANESAgent(state: State, llm: ChatOpenAI, system_prompt: str, tools=None):
    """LLM node for XANES workflows that processes messages and decides next actions.

    Parameters
    ----------
    state : State
        The current state containing messages and remaining steps
    llm : ChatOpenAI
        The language model to use for processing
    system_prompt : str
        The system prompt to guide the LLM's behavior
    tools : list, optional
        List of tools available to the agent, by default None

    Returns
    -------
    dict
        Updated state containing the LLM's response
    """
    if tools is None:
        tools = [
            molecule_name_to_smiles,
            smiles_to_coordinate_file,
            run_ase,
            run_xanes,
            fetch_xanes_data,
            plot_xanes_data,
        ]
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": f"{state['messages']}"},
    ]
    llm_with_tools = llm.bind_tools(tools=tools)
    return {"messages": [llm_with_tools.invoke(messages)]}


def ResponseAgent(state: State, llm: ChatOpenAI, formatter_prompt: str):
    """An LLM agent responsible for formatting final message.

    Parameters
    ----------
    state : State
        The current state containing messages and remaining steps
    llm : ChatOpenAI
        The language model to use for formatting
    formatter_prompt : str
        The prompt to guide the LLM's formatting behavior

    Returns
    -------
    dict
        Updated state containing the formatted response
    """
    messages = [
        {"role": "system", "content": formatter_prompt},
        {"role": "user", "content": f"{state['messages']}"},
    ]
    llm_structured_output = llm.with_structured_output(ResponseFormatter)
    response = llm_structured_output.invoke(messages).model_dump_json()
    return {"messages": [response]}


def construct_single_agent_xanes_graph(
    llm: ChatOpenAI,
    system_prompt: str = xanes_single_agent_prompt,
    structured_output: bool = False,
    formatter_prompt: str = xanes_formatter_prompt,
    tools: list = None,
):
    """Construct a single-agent graph for XANES/FDMNES workflows.

    Parameters
    ----------
    llm : ChatOpenAI
        The language model to use for the graph
    system_prompt : str, optional
        The system prompt to guide the LLM's behavior,
        by default xanes_single_agent_prompt
    structured_output : bool, optional
        Whether to use structured output, by default False
    formatter_prompt : str, optional
        The prompt to guide the LLM's formatting behavior,
        by default xanes_formatter_prompt
    tools : list, optional
        The list of tools for the main agent, by default None

    Returns
    -------
    StateGraph
        The constructed single agent XANES graph
    """
    try:
        logger.info("Constructing single agent XANES graph")

        if not os.environ.get("MP_API_KEY"):
            logger.warning(
                "MP_API_KEY environment variable is not set. "
                "The fetch_xanes_data tool will require an API key "
                "to be passed explicitly."
            )
        if not os.environ.get("FDMNES_EXE"):
            logger.warning(
                "FDMNES_EXE environment variable is not set. "
                "The run_xanes tool will not work without the FDMNES executable."
            )

        checkpointer = MemorySaver()
        if tools is None:
            tools = [
                molecule_name_to_smiles,
                smiles_to_coordinate_file,
                run_ase,
                run_xanes,
                fetch_xanes_data,
                plot_xanes_data,
            ]
        tool_node = ToolNode(tools=tools)
        graph_builder = StateGraph(State)

        if not structured_output:
            graph_builder.add_node(
                "XANESAgent",
                lambda state: XANESAgent(
                    state, llm, system_prompt=system_prompt, tools=tools
                ),
            )
            graph_builder.add_node("tools", tool_node)
            graph_builder.add_edge(START, "XANESAgent")
            graph_builder.add_conditional_edges(
                "XANESAgent",
                route_tools,
                {"tools": "tools", "done": END},
            )
            graph_builder.add_edge("tools", "XANESAgent")
            graph_builder.add_edge("XANESAgent", END)

            graph = graph_builder.compile(checkpointer=checkpointer)
            logger.info("XANES graph construction completed")
            return graph
        else:
            graph_builder.add_node(
                "XANESAgent",
                lambda state: XANESAgent(
                    state, llm, system_prompt=system_prompt, tools=tools
                ),
            )
            graph_builder.add_node("tools", tool_node)
            graph_builder.add_node(
                "ResponseAgent",
                lambda state: ResponseAgent(
                    state, llm, formatter_prompt=formatter_prompt
                ),
            )
            graph_builder.add_conditional_edges(
                "XANESAgent",
                route_tools,
                {"tools": "tools", "done": "ResponseAgent"},
            )
            graph_builder.add_edge("tools", "XANESAgent")
            graph_builder.add_edge(START, "XANESAgent")
            graph_builder.add_edge("ResponseAgent", END)

            graph = graph_builder.compile(checkpointer=checkpointer)
            logger.info("XANES graph construction completed")
            return graph

    except Exception as e:
        logger.error(f"Error constructing XANES graph: {str(e)}")
        raise