src/chemgraph/graphs/single_agent.py

import json

from langgraph.graph import StateGraph, START, END
from langchain_openai import ChatOpenAI
from langchain_core.messages import AIMessage, HumanMessage
from langgraph.checkpoint.memory import MemorySaver
from langgraph.prebuilt import ToolNode
from chemgraph.tools.ase_tools import (
    run_ase,
    extract_output_json,
)
from chemgraph.tools.cheminformatics_tools import (
    molecule_name_to_smiles,
    smiles_to_coordinate_file,
)
from chemgraph.tools.report_tools import generate_html
from chemgraph.tools.generic_tools import calculator, ask_human
from chemgraph.prompt.single_agent_prompt import (
    single_agent_prompt,
    formatter_prompt,
    report_prompt,
)
from chemgraph.utils.tool_mapping import (
    build_workflow_state,
    build_validation_failure_output,
    extract_query_from_messages,
    format_validation_failure_message,
    infer_requirement,
    parse_reaction_from_context,
    validate_completion,
    validate_tool_call_batch,
)
from chemgraph.utils.logging_config import setup_logger
from chemgraph.utils.parsing import parse_response_formatter
from chemgraph.state.state import State

logger = setup_logger(__name__)


def _normalize_tool_call(call):
    """Normalize LangChain/OpenAI/dict tool-call payloads."""
    if not isinstance(call, dict):
        return None
    if "function" in call:
        function = call.get("function") or {}
        raw_args = function.get("arguments") or {}
        if isinstance(raw_args, str):
            try:
                args = json.loads(raw_args)
            except json.JSONDecodeError:
                args = {}
        else:
            args = raw_args
        return {
            "id": call.get("id"),
            "name": function.get("name"),
            "args": args if isinstance(args, dict) else {},
        }
    return {
        "id": call.get("id"),
        "name": call.get("name"),
        "args": call.get("args") if isinstance(call.get("args"), dict) else {},
    }


def _message_tool_calls(message) -> list[dict]:
    """Return normalized tool calls from a message-like object."""
    if isinstance(message, dict):
        raw_calls = message.get("tool_calls")
        if not raw_calls:
            raw_calls = (message.get("additional_kwargs") or {}).get("tool_calls")
    else:
        raw_calls = getattr(message, "tool_calls", None)
        if not raw_calls:
            raw_calls = (
                getattr(message, "additional_kwargs", {}) or {}
            ).get("tool_calls")
    calls = [_normalize_tool_call(call) for call in raw_calls or []]
    return [call for call in calls if call and call.get("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 []:
        normalized = _normalize_tool_call(call) if isinstance(call, dict) else None
        name = normalized.get("name") if normalized else None
        args = normalized.get("args", {}) if normalized else {}
        # Normalize args for deterministic comparisons across repeated cycles.
        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 _message_tool_calls(message):
            ai_with_calls.append(message)

    if len(ai_with_calls) < 2:
        return False

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


def _tool_message_name(message):
    """Extract tool name from a message-like object.

    Parameters
    ----------
    message : Any
        Message dictionary or object.

    Returns
    -------
    str or None
        Tool name when present.
    """
    if isinstance(message, dict):
        return message.get("name")
    return getattr(message, "name", None)


def _tool_message_content(message):
    """Extract content text from a message-like object.

    Parameters
    ----------
    message : Any
        Message dictionary or object.

    Returns
    -------
    Any
        Message content, or an empty string when unavailable.
    """
    if isinstance(message, dict):
        return message.get("content", "")
    return getattr(message, "content", "")


def _is_successful_report_message(message) -> bool:
    """Return True when a message indicates successful report generation.

    Parameters
    ----------
    message : Any
        Tool message dictionary or object.

    Returns
    -------
    bool
        ``True`` for non-error ``generate_html`` tool output.
    """
    if _tool_message_name(message) != "generate_html":
        return False

    content = _tool_message_content(message)
    content_text = str(content).strip().lower() if content is not None else ""
    if not content_text:
        return False

    # ToolNode formats failures as "Error: ..."; treat only non-error output as success.
    return not content_text.startswith("error")


def _message_type(message):
    if isinstance(message, dict):
        return message.get("type") or message.get("role")
    return getattr(message, "type", None)


def _message_content(message):
    if isinstance(message, dict):
        return message.get("content", "")
    return getattr(message, "content", "")


def _message_tool_name(message):
    if isinstance(message, dict):
        return message.get("name")
    return getattr(message, "name", None)


def _compact_for_prompt(value, *, max_items: int = 16, max_text: int = 1200):
    """Recursively keep only prompt-useful context from large values."""
    if isinstance(value, dict):
        compact = {}
        for key, child in value.items():
            if key in {"final_structure", "simulation_input"}:
                continue
            compact[key] = _compact_for_prompt(
                child, max_items=max_items, max_text=max_text
            )
        return compact
    if isinstance(value, list):
        if len(value) > max_items:
            return {
                "count": len(value),
                "sample": [
                    _compact_for_prompt(
                        item, max_items=max_items, max_text=max_text
                    )
                    for item in value[:max_items]
                ],
            }
        return [
            _compact_for_prompt(item, max_items=max_items, max_text=max_text)
            for item in value
        ]
    if isinstance(value, str) and len(value) > max_text:
        return value[:max_text] + "...<truncated>"
    return value


def _recent_text_messages(messages, max_messages: int = 6) -> list[dict]:
    """Return a short text-only conversation tail for follow-up resolution."""
    recent = []
    for message in messages or []:
        msg_type = _message_type(message)
        if msg_type == "tool":
            continue
        content = _message_content(message)
        if not content:
            continue
        if hasattr(message, "tool_calls") and getattr(message, "tool_calls", None):
            recent.append(
                {
                    "role": "ai",
                    "content": "Tool calls requested.",
                    "tool_calls": [
                        call.get("name") for call in getattr(message, "tool_calls", [])
                    ],
                }
            )
            continue
        recent.append(
            {
                "role": str(msg_type or "message"),
                "content": _compact_for_prompt(str(content), max_text=500),
            }
        )
    return recent[-max_messages:]


def _compact_state_for_llm(state: State) -> str:
    """Build a bounded workflow transcript for LLM routing/composition.

    The raw LangGraph message list may contain large tool outputs such as IR
    spectra.  The model only needs the latest user intent, what tools ran,
    where artifacts were saved, and what output is still missing.
    """
    messages = state.get("messages", [])
    query = extract_query_from_messages(messages)
    validation = state.get("completion_validation")
    try:
        workflow_state = build_workflow_state(
            query,
            messages,
            validation=validation,
        )
    except Exception as exc:
        workflow_state = {
            "intent": {"task_type": "unknown"},
            "task_plan": {"next_action": "inspect_messages"},
            "tool_trace": [],
            "observations": {},
            "artifacts": {},
            "validation": {
                "complete": False,
                "reason": f"workflow_state_error: {exc}",
            },
        }

    payload = {
        "latest_user_query": query,
        "recent_text_messages": _recent_text_messages(messages),
        "intent": workflow_state.get("intent", {}),
        "task_plan": workflow_state.get("task_plan", {}),
        "action_plan": workflow_state.get("action_plan", {}),
        "identity_resolution": workflow_state.get("identity_resolution", {}),
        "clarification": workflow_state.get("clarification", {}),
        "memory_refs": workflow_state.get("memory_refs", {}),
        "tool_trace": workflow_state.get("tool_trace", []),
        "observations": workflow_state.get("observations", {}),
        "artifacts": workflow_state.get("artifacts", {}),
        "validation": workflow_state.get("validation", {}),
    }
    compact_payload = _compact_for_prompt(payload)
    return (
        "Compact ChemGraph workflow state. Use only this bounded context for "
        "tool routing and answer composition. Reuse listed artifacts and "
        "observations when the latest query is a follow-up. Do not repeat "
        "satisfied tools unless the latest query requires a different "
        "property, molecule, calculator, temperature, pressure, or driver. "
        "Full numerical outputs are saved in the artifact paths.\n\n"
        + json.dumps(compact_payload, indent=2, default=str)
    )


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}")
    tool_calls = _message_tool_calls(ai_message)
    if tool_calls:
        if not isinstance(state, list):
            query = extract_query_from_messages(messages)
            try:
                validation = validate_completion(query, messages)
            except Exception:
                validation = {}
            if validation.get("complete") and validation.get("task_type") != "unknown":
                return "done"
            try:
                tool_guard = validate_tool_call_batch(
                    query,
                    messages[:-1],
                    tool_calls,
                )
            except Exception:
                tool_guard = {"allowed": True}
            if not tool_guard.get("allowed", True):
                return "done"
        if not isinstance(state, list) and _is_repeated_tool_cycle(messages):
            return "done"
        return "tools"
    return "done"


def route_after_completion_validation(state: State):
    """Route after deterministic completion validation.

    The graph can only enter the response/composer step when the mapped
    tool-output requirements are satisfied.  A small retry cap prevents an
    endless loop when the LLM cannot repair an incomplete workflow.
    """
    validation = state.get("completion_validation") or {}
    if validation.get("complete"):
        return "response"
    if state.get("validation_attempts", 0) >= 3:
        return "response"
    if "balanced_reaction" in validation.get("missing", []):
        return "intention"
    return "agent"


def IntentionAgent(state: State, llm: ChatOpenAI):
    """Create a machine-readable scientific plan before tool execution.

    This node handles chemistry intent that should be decided by the agent,
    such as turning a plain-language reaction name into a balanced reaction.
    Deterministic validators then parse this explicit agent decision.
    """
    messages = state.get("messages", [])
    query = extract_query_from_messages(messages)
    requirement = infer_requirement(query)
    if requirement.task_type not in {"reaction_enthalpy", "reaction_gibbs_energy"}:
        return {}
    if parse_reaction_from_context(query, messages) is not None:
        return {}

    prompt_messages = [
        {
            "role": "system",
            "content": (
                "You are a computational chemistry planning node. For the "
                "user's reaction task, determine the balanced chemical reaction. "
                "Return exactly one line in this format and no extra text: "
                "Balanced reaction: Reactant + 2 Reactant -> Product + 2 Product. "
                "Use common molecule names, not formulas, so downstream tools can "
                "resolve each species.\n\n"
                "Examples:\n"
                "User: What is the reaction enthalpy of methane combustion?\n"
                "Balanced reaction: Methane + 2 Oxygen -> Carbon dioxide + 2 Water\n"
                "User: What is the Gibbs free energy for hydrogen combustion?\n"
                "Balanced reaction: 2 Hydrogen + Oxygen -> 2 Water"
            ),
        },
        {"role": "user", "content": query},
    ]
    response = llm.invoke(prompt_messages)
    content = str(getattr(response, "content", response)).strip()
    if not content:
        return {}
    if not content.lower().startswith("balanced reaction"):
        content = f"Balanced reaction: {content}"
    return {"messages": [AIMessage(content=content)]}


def CompletionValidator(state: State):
    """Check mapped tool requirements before final response composition."""
    messages = state.get("messages", [])
    query = extract_query_from_messages(messages)
    validation = validate_completion(query, messages)
    tool_guard = _latest_tool_guard(query, messages)
    workflow_state = build_workflow_state(query, messages, validation=validation)
    attempts = int(state.get("validation_attempts", 0)) + 1

    update = {
        **workflow_state,
        "completion_validation": validation,
        "tool_guard": tool_guard,
        "validation_attempts": attempts,
    }
    deterministic_output = validation.get("structured_output")
    if validation.get("complete") and deterministic_output is not None:
        update["final_output"] = deterministic_output
        update["messages"] = [AIMessage(content=json.dumps(deterministic_output))]
        return update
    if not validation.get("complete"):
        if attempts >= 3:
            failure_output = build_validation_failure_output(
                validation,
                workflow_state=workflow_state,
            )
            update["final_output"] = failure_output
            update["messages"] = [
                AIMessage(content=format_validation_failure_message(failure_output))
            ]
        else:
            missing = ", ".join(validation.get("missing", [])) or "unknown"
            repair_instruction = (
                tool_guard.get("repair_instruction")
                if not tool_guard.get("allowed", True)
                else validation.get("repair_instruction", "")
            )
            update["messages"] = [
                HumanMessage(
                    content=(
                        "Validation failed: the requested task is not complete.\n"
                        f"Task type: {validation.get('task_type')}\n"
                        f"Missing required output: {missing}\n"
                        f"Tool guard: {tool_guard.get('reason', 'not evaluated')}\n"
                        f"Next action: {repair_instruction}\n"
                        "Do not provide a final answer yet. Call the required tool(s)."
                    )
                )
            ]
    return update


def _latest_tool_guard(query: str, messages: list) -> dict:
    """Evaluate the current AI tool-call batch against workflow state."""
    if not messages:
        return {"allowed": True}
    latest = messages[-1]
    tool_calls = _message_tool_calls(latest)
    if not tool_calls:
        return {"allowed": True}
    try:
        return validate_tool_call_batch(query, messages[:-1], tool_calls)
    except Exception as exc:
        return {
            "allowed": True,
            "blocked_tool": None,
            "expected_next_action": None,
            "repair_instruction": "",
            "reason": f"Tool guard skipped after internal error: {exc}",
        }


def route_report_tools(state: State):
    """Route report tool execution and stop if a report was already generated.

    Parameters
    ----------
    state : State
        Current graph state or message list.

    Returns
    -------
    str
        ``"tools"`` when ``generate_html`` should run, otherwise ``"done"``.
    """
    if isinstance(state, list):
        messages = state
        ai_message = state[-1] if state else None
    elif messages := state.get("messages", []):
        ai_message = messages[-1]
    else:
        raise ValueError(f"No messages found in input state to tool_edge: {state}")

    tool_calls = _message_tool_calls(ai_message)
    if not tool_calls:
        return "done"

    # Only allow known report tool calls to reach ToolNode.
    valid_report_tools = {"generate_html"}
    requested_tools = {
        call.get("name")
        for call in tool_calls
        if isinstance(call, dict)
    }
    if not requested_tools or not requested_tools.issubset(valid_report_tools):
        return "done"

    report_generated = any(
        _is_successful_report_message(message) for message in messages
    )
    return "done" if report_generated else "tools"


def route_after_report_tools(state: State):
    """After report tool execution, stop on success or retry on failure.

    Parameters
    ----------
    state : State
        Current graph state or message list after report tool execution.

    Returns
    -------
    str
        ``"done"`` after a successful report message, otherwise ``"retry"``.
    """
    if isinstance(state, list):
        messages = state
    elif messages := state.get("messages", []):
        pass
    else:
        raise ValueError(f"No messages found in input state to tool_edge: {state}")

    return "done" if _is_successful_report_message(messages[-1]) else "retry"


def ChemGraphAgent(
    state: State,
    llm: ChatOpenAI,
    system_prompt: str,
    tools=None,
    human_supervised: bool = False,
):
    """LLM node 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
    human_supervised : bool, optional
        Whether to include the ``ask_human`` tool, by default False

    Returns
    -------
    dict
        Updated state containing the LLM's response
    """

    # Load default tools if no tool is specified.
    if tools is None:
        tools = [
            smiles_to_coordinate_file,
            run_ase,
            molecule_name_to_smiles,
            extract_output_json,
            calculator,
        ]
        if human_supervised:
            tools.append(ask_human)
    elif human_supervised and ask_human not in tools:
        # Ensure ask_human is available when custom tools are provided
        # and human supervision is enabled.
        tools = list(tools) + [ask_human]
    compact_state = _compact_state_for_llm(state)
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": compact_state},
    ]
    llm_with_tools = llm.bind_tools(tools=tools)
    return {"messages": [llm_with_tools.invoke(messages)]}


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

    When the LLM response cannot be parsed into a valid
    :class:`ResponseFormatter`, the agent retries the LLM call up to
    ``max_retries`` times, sending the parse error back to the model so
    it can correct its output.

    If all attempts fail, an empty ``ResponseFormatter`` is returned
    with a ``_parse_error`` key in the serialised JSON so that
    downstream evaluation can detect the failure.

    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 behaviour.
    max_retries : int, optional
        Maximum number of retry attempts on parse failure (default 1).

    Returns
    -------
    dict
        Updated state containing the formatted response.
    """
    validation = state.get("completion_validation") or {}
    deterministic_output = validation.get("structured_output")
    if validation.get("complete") and deterministic_output is not None:
        return {
            "messages": [json.dumps(deterministic_output)],
            "final_output": deterministic_output,
        }
    if validation and not validation.get("complete") and validation.get("task_type") != "unknown":
        query = extract_query_from_messages(state.get("messages", []))
        workflow_state = build_workflow_state(
            query,
            state.get("messages", []),
            validation=validation,
        )
        failure_output = build_validation_failure_output(
            validation,
            workflow_state=workflow_state,
        )
        return {
            "messages": [json.dumps(failure_output)],
            "final_output": failure_output,
        }

    compact_state = _compact_state_for_llm(state)
    messages = [
        {"role": "system", "content": formatter_prompt},
        {"role": "user", "content": compact_state},
    ]
    raw_response = llm.invoke(messages).content
    formatter, parse_error = parse_response_formatter(raw_response)

    # Retry loop: re-invoke the LLM with the error feedback.
    retries = 0
    while parse_error is not None and retries < max_retries:
        retries += 1
        logger.warning(
            "ResponseAgent: parse attempt %d failed (%s); retrying LLM.",
            retries,
            parse_error,
        )
        retry_messages = [
            {"role": "system", "content": formatter_prompt},
            {"role": "user", "content": compact_state},
            {
                "role": "assistant",
                "content": raw_response,
            },
            {
                "role": "user",
                "content": (
                    f"Error: {parse_error}\n\n"
                    "Your previous response could not be parsed. "
                    "Please output ONLY a valid JSON object matching the "
                    "ResponseFormatter schema. Do not include any text, "
                    "markdown fences, or explanation outside the JSON object."
                ),
            },
        ]
        raw_response = llm.invoke(retry_messages).content
        formatter, parse_error = parse_response_formatter(raw_response)

    # Serialise to JSON, injecting ``_parse_error`` when parsing failed.
    result = json.loads(formatter.model_dump_json())
    if parse_error is not None:
        logger.error(
            "ResponseAgent: all %d retries exhausted; returning empty "
            "ResponseFormatter with _parse_error.",
            max_retries,
        )
        result["_parse_error"] = parse_error
    response = json.dumps(result)
    return {"messages": [response], "final_output": result}


def ReportAgent(
    state: State, llm: ChatOpenAI, system_prompt: str, tools=[generate_html]
):
    """LLM node that generates a report from the messages.

    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 [generate_html]

    Returns
    -------
    dict
        Updated state containing the LLM's response
    """

    # Load default tools if no tool is specified.
    if tools is None:
        tools = [
            generate_html,
        ]
    compact_state = _compact_state_for_llm(state)
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": compact_state},
    ]
    llm_with_tools = llm.bind_tools(
        tools=tools,
        tool_choice="generate_html",
        parallel_tool_calls=False,
    )
    return {"messages": [llm_with_tools.invoke(messages)]}


def construct_single_agent_graph(
    llm: ChatOpenAI,
    system_prompt: str = single_agent_prompt,
    structured_output: bool = False,
    formatter_prompt: str = formatter_prompt,
    generate_report: bool = False,
    report_prompt: str = report_prompt,
    tools: list = None,
    max_retries: int = 1,
    human_supervised: bool = False,
):
    """Construct a geometry optimization graph.

    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 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 formatter_prompt
    generate_report: bool, optional
        Whether to generate a report, by default False
    report_prompt: str, optional
        The prompt to guide the LLM's report generation behavior, by default report_prompt
    tools : list, optional
        The list of tools for the main agent, by default None
    max_retries : int, optional
        Maximum number of LLM retry attempts when the ResponseAgent
        fails to parse the formatter output, by default 1
    human_supervised : bool, optional
        Whether to include the ``ask_human`` tool so the agent can
        pause and request human input, by default False

    Returns
    -------
    StateGraph
        The constructed single agent graph
    """
    try:
        logger.info("Constructing single agent graph")
        checkpointer = MemorySaver()
        if tools is None:
            tools = [
                smiles_to_coordinate_file,
                molecule_name_to_smiles,
                run_ase,
                extract_output_json,
                calculator,
            ]
            if human_supervised:
                tools.append(ask_human)
        elif human_supervised and ask_human not in tools:
            # Ensure ask_human is available when custom tools are provided
            # and human supervision is enabled.
            tools = list(tools) + [ask_human]
        tool_node = ToolNode(tools=tools)
        graph_builder = StateGraph(State)

        if not structured_output:
            graph_builder.add_node(
                "IntentionAgent",
                lambda state: IntentionAgent(state, llm),
            )
            graph_builder.add_node(
                "ChemGraphAgent",
                lambda state: ChemGraphAgent(
                    state,
                    llm,
                    system_prompt=system_prompt,
                    tools=tools,
                    human_supervised=human_supervised,
                ),
            )
            graph_builder.add_node("tools", tool_node)
            graph_builder.add_node("CompletionValidator", CompletionValidator)
            graph_builder.add_edge(START, "IntentionAgent")
            graph_builder.add_edge("IntentionAgent", "ChemGraphAgent")

            if generate_report:
                tool_node_report = ToolNode(tools=[generate_html])
                graph_builder.add_node("report_tools", tool_node_report)

                graph_builder.add_node(
                    "ReportAgent",
                    lambda state: ReportAgent(
                        state, llm, system_prompt=report_prompt, tools=[generate_html]
                    ),
                )
                graph_builder.add_conditional_edges(
                    "ChemGraphAgent",
                    route_tools,
                    {"tools": "tools", "done": "CompletionValidator"},
                )
                graph_builder.add_conditional_edges(
                    "CompletionValidator",
                    route_after_completion_validation,
                    {
                        "intention": "IntentionAgent",
                        "agent": "ChemGraphAgent",
                        "response": "ReportAgent",
                    },
                )
                graph_builder.add_edge("tools", "ChemGraphAgent")
                graph_builder.add_conditional_edges(
                    "ReportAgent",
                    route_report_tools,
                    {"tools": "report_tools", "done": END},
                )
                graph_builder.add_conditional_edges(
                    "report_tools",
                    route_after_report_tools,
                    {"retry": "ReportAgent", "done": END},
                )
            else:
                graph_builder.add_conditional_edges(
                    "ChemGraphAgent",
                    route_tools,
                    {"tools": "tools", "done": "CompletionValidator"},
                )
                graph_builder.add_conditional_edges(
                    "CompletionValidator",
                    route_after_completion_validation,
                    {
                        "intention": "IntentionAgent",
                        "agent": "ChemGraphAgent",
                        "response": END,
                    },
                )
                graph_builder.add_edge("tools", "ChemGraphAgent")

            graph = graph_builder.compile(checkpointer=checkpointer)
            logger.info("Graph construction completed")
            return graph
        else:
            graph_builder.add_node(
                "IntentionAgent",
                lambda state: IntentionAgent(state, llm),
            )
            graph_builder.add_node(
                "ChemGraphAgent",
                lambda state: ChemGraphAgent(
                    state,
                    llm,
                    system_prompt=system_prompt,
                    tools=tools,
                    human_supervised=human_supervised,
                ),
            )
            graph_builder.add_node("tools", tool_node)
            graph_builder.add_node("CompletionValidator", CompletionValidator)
            graph_builder.add_node(
                "ResponseAgent",
                lambda state: ResponseAgent(
                    state,
                    llm,
                    formatter_prompt=formatter_prompt,
                    max_retries=max_retries,
                ),
            )
            graph_builder.add_conditional_edges(
                "ChemGraphAgent",
                route_tools,
                {"tools": "tools", "done": "CompletionValidator"},
            )
            graph_builder.add_conditional_edges(
                "CompletionValidator",
                route_after_completion_validation,
                {
                    "intention": "IntentionAgent",
                    "agent": "ChemGraphAgent",
                    "response": "ResponseAgent",
                },
            )
            graph_builder.add_edge("tools", "ChemGraphAgent")
            graph_builder.add_edge(START, "IntentionAgent")
            graph_builder.add_edge("IntentionAgent", "ChemGraphAgent")
            graph_builder.add_edge("ResponseAgent", END)

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

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