src/chemgraph/agent/llm_agent.py

import asyncio
import datetime
import dataclasses
import os
from typing import Callable, List, Optional
import uuid

from chemgraph.memory.store import SessionStore
from chemgraph.memory.schemas import SessionMessage
from chemgraph.models.openai import load_openai_model
from chemgraph.models.alcf_endpoints import load_alcf_model
from chemgraph.models.local_model import load_ollama_model
from chemgraph.models.anthropic import load_anthropic_model
from chemgraph.models.gemini import load_gemini_model
from chemgraph.models.groq import load_groq_model
from chemgraph.models.supported_models import (
    supported_openai_models,
    supported_ollama_models,
    supported_anthropic_models,
    supported_alcf_models,
    supported_argo_models,
    supported_gemini_models,

)
from chemgraph.schemas.ase_input import (
    get_available_calculator_names,
    get_calculator_selection_context,
    get_default_calculator_name,
)

from chemgraph.prompt.single_agent_prompt import (
    single_agent_prompt,
    get_single_agent_prompt,
    formatter_prompt as default_formatter_prompt,
    report_prompt as default_report_prompt,
)
from chemgraph.prompt.multi_agent_prompt import (
    executor_prompt as default_executor_prompt,
    formatter_multi_prompt as default_formatter_multi_prompt,
    aggregator_prompt as default_aggregator_prompt,
    planner_prompt as default_planner_prompt,
)
from langgraph.types import Command
from langgraph.errors import GraphInterrupt

from chemgraph.graphs.single_agent import construct_single_agent_graph


from chemgraph.graphs.python_relp_agent import construct_relp_graph
from chemgraph.graphs.multi_agent import construct_multi_agent_graph
from chemgraph.graphs.graspa_agent import construct_graspa_graph
from chemgraph.graphs.mock_agent import construct_mock_agent_graph
from chemgraph.graphs.single_agent_mcp import construct_single_agent_mcp_graph
from chemgraph.graphs.graspa_mcp import construct_graspa_mcp_graph
from chemgraph.graphs.rag_agent import construct_rag_agent_graph
from chemgraph.graphs.single_agent_xanes import construct_single_agent_xanes_graph
from chemgraph.prompt.rag_prompt import rag_agent_prompt
from chemgraph.prompt.xanes_prompt import (
    xanes_single_agent_prompt as default_xanes_single_agent_prompt,
    xanes_formatter_prompt as default_xanes_formatter_prompt,
)

import logging

logger = logging.getLogger(__name__)


def _is_mock_object(value) -> bool:
    """Return True for unittest.mock objects without importing test-only APIs.

    Parameters
    ----------
    value : Any
        Object to inspect.

    Returns
    -------
    bool
        ``True`` when the object comes from ``unittest.mock``.
    """
    return value.__class__.__module__.startswith("unittest.mock")


def serialize_state(state, *, max_depth: int = 50, _seen: set[int] | None = None):
    """Convert non-serializable objects in state to a JSON-friendly format.

    Parameters
    ----------
    state : Any
        The state object to be serialized. Can be a list, dict, or object with __dict__
    max_depth : int, optional
        Maximum object nesting depth to serialize before falling back to a
        placeholder. This prevents runaway recursion for complex graph objects.

    Returns
    -------
    Any
        A JSON-serializable version of the input state
    """
    if _seen is None:
        _seen = set()

    if max_depth < 0:
        return f"<max depth exceeded: {type(state).__name__}>"

    if isinstance(state, (str, int, float, bool)) or state is None:
        return state

    if isinstance(state, (datetime.datetime, datetime.date)):
        return state.isoformat()

    if _is_mock_object(state):
        return str(state)

    state_id = id(state)
    if state_id in _seen:
        return f"<circular reference: {type(state).__name__}>"

    if isinstance(state, dict):
        _seen.add(state_id)
        try:
            return {
                str(key): serialize_state(
                    value, max_depth=max_depth - 1, _seen=_seen
                )
                for key, value in state.items()
            }
        finally:
            _seen.remove(state_id)

    if isinstance(state, (list, tuple, set, frozenset)):
        _seen.add(state_id)
        try:
            return [
                serialize_state(item, max_depth=max_depth - 1, _seen=_seen)
                for item in state
            ]
        finally:
            _seen.remove(state_id)

    model_dump = getattr(state, "model_dump", None)
    if callable(model_dump):
        _seen.add(state_id)
        try:
            try:
                dumped = model_dump(mode="json")
            except TypeError:
                dumped = model_dump()
            return serialize_state(dumped, max_depth=max_depth - 1, _seen=_seen)
        except Exception:
            return str(state)
        finally:
            _seen.remove(state_id)

    if dataclasses.is_dataclass(state) and not isinstance(state, type):
        _seen.add(state_id)
        try:
            return {
                field.name: serialize_state(
                    getattr(state, field.name),
                    max_depth=max_depth - 1,
                    _seen=_seen,
                )
                for field in dataclasses.fields(state)
            }
        finally:
            _seen.remove(state_id)

    if hasattr(state, "__dict__"):
        _seen.add(state_id)
        try:
            return {
                str(key): serialize_state(
                    value, max_depth=max_depth - 1, _seen=_seen
                )
                for key, value in vars(state).items()
            }
        finally:
            _seen.remove(state_id)

    return str(state)


class ChemGraph:
    """A graph-based workflow for LLM-powered computational chemistry tasks.

    This class manages different types of workflows for computational chemistry tasks,
    supporting various LLM models and workflow types.

    Parameters
    ----------
    model_name : str, optional
        Name of the language model to use, by default "gpt-4o-mini"
    workflow_type : str, optional
        Type of workflow to use. Options:
        - "single_agent"
        - "multi_agent"
        - "python_relp"
        - "graspa_agent"
        by default "single_agent"
    base_url : str, optional
        Base URL for API calls, by default None
    api_key : str, optional
        API key for authentication, by default None
    system_prompt : str, optional
        System prompt for the language model, by default single_agent_prompt
    formatter_prompt : str, optional
        Prompt for formatting output, by default formatter_prompt
    structured_output : bool, optional
        Whether to use structured output, by default False
    return_option : str, optional
        What to return from the workflow. Options:
        - "last_message"
        - "state"
        by default "last_message"
    recursion_limit : int, optional
        Maximum number of recursive steps in the workflow, by default 50
    max_retries : int, optional
        Maximum number of LLM retry attempts when an agent
        fails to parse its output, by default 1
    human_input_handler : callable, optional
        A callback ``f(question: str) -> str`` invoked when the graph
        pauses for human input (via ``interrupt()``).  Receives the
        question text and must return the human's answer as a string.
        If ``None`` (default), interrupts will propagate as
        ``GraphInterrupt`` exceptions.  The handler may also be an
        ``async`` callable.
    human_supervised : bool, optional
        Whether to include the ``ask_human`` tool so the agent can
        pause and request human input.  When ``False`` the tool is
        excluded from the tool list and the corresponding instruction
        is removed from the default system prompt, by default False.

    Raises
    ------
    ValueError
        If the workflow_type is not supported
    Exception
        If there is an error loading the specified model
    """

    def __init__(
        self,
        model_name: str = "gpt-4o-mini",
        workflow_type: str = "single_agent",
        base_url: str = None,
        api_key: str = None,
        argo_user: str = None,
        system_prompt: str = single_agent_prompt,
        formatter_prompt: str = default_formatter_prompt,
        structured_output: bool = False,
        return_option: str = "last_message",
        recursion_limit: int = 50,
        planner_prompt: str = default_planner_prompt,
        executor_prompt: str = default_executor_prompt,
        aggregator_prompt: str = default_aggregator_prompt,
        formatter_multi_prompt: str = default_formatter_multi_prompt,
        generate_report: bool = False,
        report_prompt: str = default_report_prompt,
        support_structured_output: bool = True,
        tools: List = None,
        data_tools: List = None,
        session_store: Optional[SessionStore] = None,
        enable_memory: bool = True,
        memory_db_path: Optional[str] = None,
        log_dir: Optional[str] = None,
        max_retries: int = 1,
        human_input_handler: Optional[Callable[[str], str]] = None,
        human_supervised: bool = False,
    ):
        """Initialize a ChemGraph workflow instance.

        Parameters
        ----------
        model_name : str, optional
            LLM model identifier.
        workflow_type : str, optional
            Workflow constructor key.
        base_url : str, optional
            Custom provider endpoint URL.
        api_key : str, optional
            API key passed to compatible model loaders.
        argo_user : str, optional
            Argo username for Argo-hosted models.
        system_prompt : str, optional
            System prompt for single-agent-style workflows.
        formatter_prompt : str, optional
            Prompt used to format single-agent final output.
        structured_output : bool, optional
            Whether structured final output is requested.
        return_option : str, optional
            Return mode, such as ``"last_message"`` or ``"state"``.
        recursion_limit : int, optional
            LangGraph recursion limit.
        planner_prompt : str, optional
            Planner prompt for multi-agent workflows.
        executor_prompt : str, optional
            Executor prompt for multi-agent workflows.
        aggregator_prompt : str, optional
            Aggregator prompt retained for compatibility.
        formatter_multi_prompt : str, optional
            Formatter prompt for multi-agent workflows.
        generate_report : bool, optional
            Whether report generation is enabled.
        report_prompt : str, optional
            Prompt used by the report-generation workflow.
        support_structured_output : bool, optional
            Whether the selected model supports structured output.
        tools : list, optional
            Custom tool list for applicable workflows.
        data_tools : list, optional
            Additional data-analysis tools for MCP workflows.
        session_store : SessionStore, optional
            Existing session store instance.
        enable_memory : bool, optional
            Whether persistent session memory is enabled.
        memory_db_path : str, optional
            SQLite path for the session store.
        log_dir : str, optional
            Directory for run logs and artifacts.
        max_retries : int, optional
            LLM parse-retry limit for formatter/planner nodes.
        human_input_handler : Callable[[str], str], optional
            Callback used to answer graph human-interrupt prompts.
        human_supervised : bool, optional
            Whether to expose human-supervision tools to the agent.
        """
        # Always generate a unique identifier for this instance
        self.uuid = str(uuid.uuid4())[:8]

        # Initialize log directory.  Explicit ``log_dir`` argument takes
        # precedence over the ``CHEMGRAPH_LOG_DIR`` environment variable,
        # which in turn takes precedence over the auto-generated default.
        self.log_dir = log_dir or os.environ.get("CHEMGRAPH_LOG_DIR")
        if not self.log_dir:
            # Create a new session log directory under cg_logs/
            timestamp = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
            # Use abspath to ensure tools getting this env var have a full path
            self.log_dir = os.path.join(
                os.getcwd(), "cg_logs", f"session_{timestamp}_{self.uuid}"
            )
            os.makedirs(self.log_dir, exist_ok=True)
            # Set env var for tools to pick up
            os.environ["CHEMGRAPH_LOG_DIR"] = self.log_dir

        # Initialize session memory store
        if session_store is not None:
            self.session_store = session_store
        elif enable_memory:
            self.session_store = SessionStore(db_path=memory_db_path)
        else:
            self.session_store = None

        # Track whether session has been registered in the memory store
        self._session_created: bool = False
        self._session_title: Optional[str] = None

        try:
            # Use hardcoded optimal values for tool calling
            temperature = 0.0  # Deterministic responses
            max_tokens = 4000  # Sufficient for most tasks
            top_p = 1.0  # No nucleus sampling filtering
            frequency_penalty = 0.0  # No repetition penalty
            presence_penalty = 0.0  # No presence penalty

            if (
                model_name in supported_openai_models
                or model_name in supported_argo_models
            ):
                openai_load_kwargs = {
                    "model_name": model_name,
                    "temperature": temperature,
                    "base_url": base_url,
                }
                if argo_user is not None:
                    openai_load_kwargs["argo_user"] = argo_user
                llm = load_openai_model(
                    **openai_load_kwargs,
                )
            elif model_name in supported_ollama_models:
                llm = load_ollama_model(model_name=model_name, temperature=temperature)
            elif model_name in supported_alcf_models:
                llm = load_alcf_model(
                    model_name=model_name, base_url=base_url, api_key=api_key
                )
            elif model_name in supported_anthropic_models:
                llm = load_anthropic_model(
                    model_name=model_name, api_key=api_key, temperature=temperature
                )
            elif model_name in supported_gemini_models:
                llm = load_gemini_model(
                    model_name=model_name, api_key=api_key, temperature=temperature
                )
            elif model_name.startswith("groq:"):
                llm = load_groq_model(
                    model_name=model_name, api_key=api_key, temperature=temperature
                )

            else:  # Assume it might be a vLLM or other custom OpenAI-compatible endpoint
                # Use environment variables for vLLM base_url and a dummy api_key if not provided
                # These would be set by docker-compose for the jupyter_lab service
                vllm_base_url = os.getenv("VLLM_BASE_URL", base_url)
                # ChatOpenAI requires an api_key, even if the endpoint doesn't use it.
                vllm_api_key = os.getenv(
                    "OPENAI_API_KEY", api_key if api_key else "dummy_vllm_key"
                )

                if vllm_base_url:
                    logger.info(
                        f"Attempting to load model '{model_name}' from custom endpoint: {vllm_base_url}"
                    )
                    from langchain_openai import ChatOpenAI

                    llm = ChatOpenAI(
                        model=model_name,
                        temperature=temperature,
                        base_url=vllm_base_url,
                        api_key=vllm_api_key,
                        max_tokens=max_tokens,
                        top_p=top_p,
                        frequency_penalty=frequency_penalty,
                        presence_penalty=presence_penalty,
                    )
                    logger.info(
                        f"Successfully initialized ChatOpenAI for model '{model_name}' at {vllm_base_url}"
                    )
                else:
                    logger.error(
                        f"Model '{model_name}' is not in any supported list and no VLLM_BASE_URL/base_url provided."
                    )
                    raise ValueError(
                        f"Unsupported model or missing base URL for: {model_name}"
                    )

        except Exception as e:
            logger.error(f"Exception thrown when loading {model_name}: {str(e)}")
            raise e

        self.workflow_type = workflow_type
        self.model_name = model_name
        self.system_prompt = system_prompt
        self.formatter_prompt = formatter_prompt
        self.structured_output = structured_output
        self.generate_report = generate_report
        self.report_prompt = report_prompt
        self.return_option = return_option
        self.recursion_limit = recursion_limit
        self.planner_prompt = planner_prompt
        self.executor_prompt = executor_prompt
        self.aggregator_prompt = aggregator_prompt
        self.formatter_multi_prompt = formatter_multi_prompt
        self.tools = tools
        self.data_tools = data_tools
        self.max_retries = max_retries
        self.human_input_handler = human_input_handler
        self.human_supervised = human_supervised

        # When human supervision is disabled and the caller is using the
        # default system prompt, strip the ask_human instructions so the
        # LLM is not told to call a tool that is unavailable.
        if not self.human_supervised and self.system_prompt == single_agent_prompt:
            self.system_prompt = get_single_agent_prompt(human_supervised=False)

        self.available_calculators = get_available_calculator_names()
        self.default_calculator = get_default_calculator_name()
        self.calculator_selection_context = get_calculator_selection_context()

        def append_calculator_context(prompt: str) -> str:
            """Append calculator availability guidance to a prompt once.

            Parameters
            ----------
            prompt : str
                Prompt text to augment.

            Returns
            -------
            str
                Prompt with calculator-selection context appended.
            """
            if self.calculator_selection_context in prompt:
                return prompt
            return f"{prompt}{self.calculator_selection_context}"

        if self.workflow_type in {"single_agent", "mock_agent", "single_agent_mcp"}:
            self.system_prompt = append_calculator_context(self.system_prompt)
        elif self.workflow_type == "multi_agent":
            self.planner_prompt = append_calculator_context(self.planner_prompt)
            self.executor_prompt = append_calculator_context(self.executor_prompt)

        if model_name in supported_argo_models:
            self.support_structured_output = False
        else:
            self.support_structured_output = support_structured_output

        self.workflow_map = {
            "single_agent": {"constructor": construct_single_agent_graph},
            "multi_agent": {"constructor": construct_multi_agent_graph},
            "python_relp": {"constructor": construct_relp_graph},
            "graspa": {"constructor": construct_graspa_graph},
            "mock_agent": {"constructor": construct_mock_agent_graph},
            "single_agent_mcp": {"constructor": construct_single_agent_mcp_graph},
            "graspa_mcp": {"constructor": construct_graspa_mcp_graph},
            "rag_agent": {"constructor": construct_rag_agent_graph},
            "single_agent_xanes": {"constructor": construct_single_agent_xanes_graph},
        }

        if workflow_type not in self.workflow_map:
            raise ValueError(
                f"Unsupported workflow type: {workflow_type}. Available types: {list(self.workflow_map.keys())}"
            )

        if self.workflow_type == "single_agent":
            self.workflow = self.workflow_map[workflow_type]["constructor"](
                llm,
                self.system_prompt,
                self.structured_output,
                self.formatter_prompt,
                self.generate_report,
                self.report_prompt,
                self.tools,
                max_retries=self.max_retries,
                human_supervised=self.human_supervised,
            )
        elif self.workflow_type == "multi_agent":
            self.workflow = self.workflow_map[workflow_type]["constructor"](
                llm,
                planner_prompt=self.planner_prompt,
                executor_prompt=self.executor_prompt,
                executor_tools=self.tools,
                structured_output=self.structured_output,
                formatter_prompt=self.formatter_multi_prompt,
                max_retries=self.max_retries,
            )
        elif self.workflow_type == "python_relp":
            self.workflow = self.workflow_map[workflow_type]["constructor"](
                llm,
                self.system_prompt,
            )
        elif self.workflow_type == "graspa":
            self.workflow = self.workflow_map[workflow_type]["constructor"](
                llm,
                self.system_prompt,
                self.structured_output,
                self.formatter_prompt,
            )
        elif self.workflow_type == "mock_agent":
            self.workflow = self.workflow_map[workflow_type]["constructor"](
                llm=llm,
                system_prompt=self.system_prompt,
            )
        elif self.workflow_type == "single_agent_mcp":
            self.workflow = self.workflow_map[workflow_type]["constructor"](
                llm=llm,
                system_prompt=self.system_prompt,
                tools=self.tools,
            )
        elif self.workflow_type == "graspa_mcp":
            self.workflow = self.workflow_map[workflow_type]["constructor"](
                llm=llm,
                executor_tools=self.tools,
                analysis_tools=self.data_tools,
            )
        elif self.workflow_type == "rag_agent":
            self.workflow = self.workflow_map[workflow_type]["constructor"](
                llm=llm,
                system_prompt=self.system_prompt
                if self.system_prompt != single_agent_prompt
                else rag_agent_prompt,
                tools=self.tools,
            )
        elif self.workflow_type == "single_agent_xanes":
            self.workflow = self.workflow_map[workflow_type]["constructor"](
                llm,
                system_prompt=self.system_prompt
                if self.system_prompt != single_agent_prompt
                else default_xanes_single_agent_prompt,
                structured_output=self.structured_output,
                formatter_prompt=self.formatter_prompt
                if self.formatter_prompt != default_formatter_prompt
                else default_xanes_formatter_prompt,
                tools=self.tools,
            )

    def visualize(self, method: str = "ascii"):
        """Visualize the LangGraph graph structure.

        This method creates and displays a visual representation of the workflow graph
        using Mermaid diagrams. The visualization is shown in Jupyter notebooks.

        Parameters
        ----------
        method : str, optional
            Visualization backend. ``"ascii"`` returns an ASCII graph;
            any other value renders a Mermaid PNG in the active notebook.

        Returns
        -------
        str or None
            ASCII graph text when ``method`` is ``"ascii"``; otherwise
            displays an image and returns ``None``.

        Notes
        -----
        Requires IPython and nest_asyncio to be installed.
        The visualization uses Mermaid diagrams with custom styling.
        """
        import nest_asyncio
        from IPython.display import Image, display
        from langchain_core.runnables.graph import (
            CurveStyle,
            MermaidDrawMethod,
            NodeStyles,
        )

        if method == "ascii":
            return self.workflow.get_graph().draw_ascii()
        else:
            nest_asyncio.apply()  # Required for Jupyter Notebook to run async functions

            display(
                Image(
                    self.workflow.get_graph().draw_mermaid_png(
                        curve_style=CurveStyle.LINEAR,
                        node_colors=NodeStyles(
                            first="#ffdfba", last="#baffc9", default="#fad7de"
                        ),
                        wrap_label_n_words=9,
                        output_file_path=None,
                        draw_method=MermaidDrawMethod.PYPPETEER,
                        background_color="white",
                        padding=6,
                    )
                )
            )

    def get_state(self, config={"configurable": {"thread_id": "1"}}):
        """Get the current state of the workflow.

        Parameters
        ----------
        config : dict, optional
            Configuration dictionary containing thread information,
            by default {"configurable": {"thread_id": "1"}}

        Returns
        -------
        list
            List of messages in the current state
        """
        return self.workflow.get_state(config).values

    def write_state(
        self,
        config: dict = None,
        file_path: str = None,
        file_name: str = None,
    ):
        """Write log of ChemGraph run to a JSON file, including workflow-specific prompts.

        Parameters
        ----------
        config : dict, optional
            Workflow config, must include 'configurable.thread_id'
        file_path : str, optional
            Full path to output file. If not provided, writes to 'cg_logs/state_thread_<thread_id>_<timestamp>.json'
        file_name : str, optional
            Optional filename to use if file_path is not provided

        Returns
        -------
        dict or str
            Dictionary of metadata if successful, or "Error" if failed.
        """
        import json
        import subprocess

        try:
            if config is None:
                config = {"configurable": {"thread_id": "1"}}
            timestamp = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
            thread_id = config["configurable"]["thread_id"]
            if not file_path:
                log_dir = getattr(self, "log_dir", None) or os.environ.get(
                    "CHEMGRAPH_LOG_DIR", "cg_logs"
                )
                os.makedirs(log_dir, exist_ok=True)
                if not file_name:
                    file_name = f"state_thread_{thread_id}_{self.uuid}_{timestamp}.json"
                file_path = os.path.join(log_dir, file_name)

            state = self.get_state(config=config)
            serialized_state = serialize_state(state)

            try:
                git_commit = (
                    subprocess.check_output(
                        ["git", "rev-parse", "HEAD"], stderr=subprocess.DEVNULL
                    )
                    .decode("utf-8")
                    .strip()
                )
            except (subprocess.CalledProcessError, FileNotFoundError):
                git_commit = "unknown"

            # Base log info
            output_data = {
                "timestamp": datetime.datetime.now().isoformat(),
                "model_name": self.model_name,
                "thread_id": thread_id,
                "git_commit": git_commit,
                "state": serialized_state,
            }

            # Add prompts depending on workflow_type
            if self.workflow_type in {
                "single_agent",
                "single_agent_xanes",
                "graspa",
                "python_relp",
                "rag_agent",
            }:
                output_data.update(
                    {
                        "system_prompt": self.system_prompt,
                        "formatter_prompt": self.formatter_prompt,
                    }
                )

            elif self.workflow_type == "graspa_mcp":
                output_data.update(
                    {
                        "system_prompt": self.system_prompt,
                    }
                )

            elif self.workflow_type == "mock_agent":
                output_data.update(
                    {
                        "system_prompt": self.system_prompt,
                    }
                )
            elif self.workflow_type == "multi_agent":
                output_data.update(
                    {
                        "planner_prompt": self.planner_prompt,
                        "executor_prompt": self.executor_prompt,
                        "formatter_prompt": self.formatter_multi_prompt,
                    }
                )
            else:
                output_data.update(
                    {
                        "system_prompt": "unknown",
                        "formatter_prompt": "unknown",
                    }
                )

            with open(file_path, "w", encoding="utf-8") as json_file:
                json.dump(output_data, json_file, indent=4)
            return output_data

        except Exception as e:
            print("Error with write_state: ", str(e))
            return "Error"

    @property
    def session_id(self) -> str:
        """Current session ID (always available, derived from self.uuid)."""
        return self.uuid

    def _ensure_session(self, query: str) -> None:
        """Create a session record on first run if memory is enabled.

        Parameters
        ----------
        query : str
            User query used to generate the session title.
        """
        if self.session_store is None:
            return
        if self._session_created:
            return

        self._session_title = SessionStore.generate_title(query)
        self.session_store.create_session(
            session_id=self.uuid,
            model_name=self.model_name,
            workflow_type=self.workflow_type,
            title=self._session_title,
            log_dir=self.log_dir,
        )
        self._session_created = True
        logger.info(f"Created session {self.uuid}: {self._session_title}")

    def _save_messages_to_store(self, last_state: dict, query: str) -> None:
        """Extract messages from workflow state and persist to session store.

        Parameters
        ----------
        last_state : dict
            Latest LangGraph state containing a ``messages`` sequence.
        query : str
            Original user query associated with the saved messages.
        """
        if self.session_store is None or not self._session_created:
            return

        try:
            messages_to_save = []
            state_messages = last_state.get("messages", [])

            for msg in state_messages:
                role = None
                content = ""
                tool_name = None

                if hasattr(msg, "type"):
                    # LangChain message objects
                    if msg.type == "human":
                        role = "human"
                    elif msg.type == "ai":
                        role = "ai"
                    elif msg.type == "tool":
                        role = "tool"
                        tool_name = getattr(msg, "name", None)
                    content = getattr(msg, "content", str(msg))
                elif isinstance(msg, dict):
                    role = msg.get("type") or msg.get("role")
                    content = msg.get("content", "")
                    tool_name = msg.get("name")

                # MCP tool messages may return content as a list of
                # content blocks (e.g. [{'type': 'text', 'text': '...'}])
                # instead of a plain string. Normalize to str.
                if isinstance(content, list):
                    content = "\n".join(
                        block.get("text", str(block))
                        if isinstance(block, dict)
                        else str(block)
                        for block in content
                    )
                elif not isinstance(content, str):
                    content = str(content)

                if role and content:
                    messages_to_save.append(
                        SessionMessage(
                            role=role,
                            content=content,
                            tool_name=tool_name,
                        )
                    )

            self.session_store.save_messages(
                session_id=self.uuid,
                messages=messages_to_save,
                title=self._session_title,
            )
            logger.info(
                f"Saved {len(messages_to_save)} messages to session {self.uuid}"
            )
        except Exception as e:
            logger.warning(f"Failed to save messages to session store: {e}")

    def load_previous_context(
        self,
        session_id: str,
        max_messages: Optional[int] = None,
    ) -> str:
        """Load context from a previous session as a summary string.

        This can be injected into the conversation to give the agent
        awareness of prior work.

        Parameters
        ----------
        session_id : str
            Previous session ID (or unique prefix).
        max_messages : int, optional
            Limit the number of messages included.

        Returns
        -------
        str
            Formatted context summary, or empty string if not found.
        """
        if self.session_store is None:
            logger.warning("Memory is disabled; cannot load previous context.")
            return ""
        return self.session_store.build_context_summary(session_id)

    async def _call_human_input_handler(self, question: str) -> str:
        """Invoke the human_input_handler, supporting both sync and async callables.

        Raises :class:`HumanInputRequired` when no handler is configured,
        allowing external callers (CLI, UI) to catch it, prompt the user,
        and resume the graph.

        Parameters
        ----------
        question : str
            Prompt emitted by the graph for a human response.

        Returns
        -------
        str
            Human response returned by the configured handler.
        """
        handler = self.human_input_handler
        if handler is None:
            raise HumanInputRequired(question)
        if asyncio.iscoroutinefunction(handler):
            return await handler(question)
        return handler(question)

    async def run(self, query: str, config=None, resume_from: Optional[str] = None):
        """
        Async-only runner. Requires `self.workflow.astream(...)`.
        Streams values, logs new messages, writes state, and returns according to
        `self.return_option` ("last_message" or "state").

        When the graph pauses for human input (via ``interrupt()``), the
        ``human_input_handler`` callback is invoked to obtain the user's
        response, and the graph is automatically resumed.  If no handler
        is configured, the ``GraphInterrupt`` exception propagates to the
        caller.

        Parameters
        ----------
        query : str
            The user query to execute.
        config : dict, optional
            LangGraph config with thread_id, etc.
        resume_from : str, optional
            Session ID to load context from. The previous conversation
            summary is prepended to the query.
        """

        def _validate_config(cfg):
            """Normalize and validate the LangGraph run configuration.

            Parameters
            ----------
            cfg : dict or None
                User-provided configuration, optionally with top-level
                ``thread_id``.

            Returns
            -------
            dict
                Config with ``configurable.thread_id`` and recursion limit set.
            """
            if cfg is None:
                cfg = {}
            if not isinstance(cfg, dict):
                raise TypeError(
                    f"`config` must be a dictionary, got {type(cfg).__name__}"
                )

            # Support top-level thread_id for convenience
            if "thread_id" in cfg:
                if "configurable" not in cfg:
                    cfg["configurable"] = {}
                cfg["configurable"]["thread_id"] = str(cfg["thread_id"])

            cfg.setdefault("configurable", {}).setdefault("thread_id", "1")
            cfg["recursion_limit"] = self.recursion_limit
            return cfg

        def _save_state_and_select_return(last_state, cfg):
            """Persist the final state and apply the configured return option.

            Parameters
            ----------
            last_state : dict
                Final streamed graph state.
            cfg : dict
                LangGraph run configuration used to retrieve/write state.

            Returns
            -------
            Any
                Final message or serialized state, depending on
                ``self.return_option``.
            """
            log_dir = self.log_dir
            if not log_dir:
                log_dir = "cg_logs"

            os.makedirs(log_dir, exist_ok=True)
            log_path = None
            self.write_state(config=cfg, file_path=log_path)

            if self.return_option == "last_message":
                return last_state["messages"][-1]
            elif self.return_option == "state":
                return serialize_state(self.get_state(config=cfg))
            else:
                raise ValueError(
                    f"Unsupported return_option: {self.return_option}. Use 'last_message' or 'state'."
                )

        async def _stream_until_interrupt(stream_input, cfg):
            """Stream the workflow until completion or an interrupt.

            Parameters
            ----------
            stream_input : dict or Command
                Initial graph input or resume command to stream.
            cfg : dict
                LangGraph run configuration.

            Returns
            -------
            tuple
                ``(last_state, interrupt_value)`` where ``interrupt_value`` is
                ``None`` when the graph completed normally.

            LangGraph's ``astream(stream_mode="values")`` does **not**
            raise ``GraphInterrupt``.  Instead the stream emits a state
            containing an ``__interrupt__`` key and then ends.  We
            detect this in two ways:

            1. Check for the ``__interrupt__`` key in streamed states.
            2. After the stream ends, inspect the checkpoint snapshot
               for pending interrupt tasks.
            """
            prev_msgs: list = []
            last_st = None
            interrupt_val = None
            try:
                async for s in self.workflow.astream(
                    stream_input, stream_mode="values", config=cfg
                ):
                    # Detect inline interrupt marker emitted by astream.
                    if "__interrupt__" in s:
                        int_data = s["__interrupt__"]
                        if isinstance(int_data, (list, tuple)) and int_data:
                            interrupt_val = int_data[0].value
                        elif hasattr(int_data, "value"):
                            interrupt_val = int_data.value
                        else:
                            interrupt_val = {
                                "question": "The workflow needs your input."
                            }

                    if "messages" in s and s["messages"] != prev_msgs:
                        new_message = s["messages"][-1]
                        try:
                            new_message.pretty_print()
                        except Exception:
                            pass
                        logger.info(new_message)
                        prev_msgs = s["messages"]
                    last_st = s
            except GraphInterrupt as gi:
                # Fallback: some LangGraph versions may still raise.
                interrupts = gi.args[0] if gi.args else []
                if interrupts:
                    interrupt_val = interrupts[0].value
                else:
                    interrupt_val = {
                        "question": "The workflow needs your input."
                    }

            # Double-check the checkpoint for pending interrupts that
            # the stream may not have surfaced explicitly.
            if interrupt_val is None:
                try:
                    snapshot = self.workflow.get_state(cfg)
                    if snapshot and snapshot.tasks:
                        for t in snapshot.tasks:
                            t_interrupts = getattr(t, "interrupts", None)
                            if t_interrupts:
                                interrupt_val = t_interrupts[0].value
                                break
                except Exception:
                    pass

            if interrupt_val is not None:
                logger.info("Graph interrupted: %s", interrupt_val)
                # Refresh state from checkpoint for consistency.
                try:
                    snapshot = self.workflow.get_state(cfg)
                    if snapshot:
                        last_st = snapshot.values
                except Exception:
                    pass

            return last_st, interrupt_val

        logger.debug("run called with config=%s", config)
        config = _validate_config(config)
        logger.debug("validated config=%s", config)

        # Initialize logging directory before determining inputs or running workflow
        # Check if CHEMGRAPH_LOG_DIR is already set
        if not os.environ.get("CHEMGRAPH_LOG_DIR"):
            os.environ["CHEMGRAPH_LOG_DIR"] = self.log_dir

        # Ensure session exists in memory store
        self._ensure_session(query)

        # If resuming from a previous session, prepend context
        if resume_from and self.session_store:
            context = self.session_store.build_context_summary(resume_from)
            if context:
                query = (
                    f"{context}\n\n"
                    f"Now, continuing from the previous session above, "
                    f"please help with the following:\n\n{query}"
                )
                logger.info(f"Injected context from session {resume_from}")

        inputs = {"messages": query}

        try:
            last_state, interrupt_value = await _stream_until_interrupt(inputs, config)

            # --- Human-in-the-loop resume loop ---
            # When the graph pauses with an interrupt, ask the human and
            # resume.  This loop handles chains of multiple interrupts
            # (e.g., the agent asks a follow-up question after receiving
            # the first answer).
            max_interrupts = 10  # safety guard against infinite interrupt loops
            interrupt_count = 0
            while interrupt_value is not None:
                interrupt_count += 1
                if interrupt_count > max_interrupts:
                    logger.error(
                        "Exceeded maximum number of human interrupts (%d); "
                        "aborting workflow.",
                        max_interrupts,
                    )
                    raise RuntimeError(
                        f"Workflow exceeded maximum of {max_interrupts} "
                        f"human interrupts."
                    )

                # Extract the question text from the interrupt value.
                if isinstance(interrupt_value, dict):
                    question = interrupt_value.get(
                        "question",
                        interrupt_value.get("message", str(interrupt_value)),
                    )
                else:
                    question = str(interrupt_value)

                logger.info("Requesting human input: %s", question)
                human_answer = await self._call_human_input_handler(question)
                logger.info("Human responded: %s", human_answer)

                # Resume the graph from the checkpoint with the human's answer.
                resume_cmd = Command(resume=human_answer)
                last_state, interrupt_value = await _stream_until_interrupt(
                    resume_cmd, config
                )

            if last_state is None:
                raise RuntimeError("Workflow produced no states.")

            # Save messages to persistent session store
            self._save_messages_to_store(last_state, query)

            return _save_state_and_select_return(last_state, config)

        except HumanInputRequired:
            # No human_input_handler configured — propagate so the
            # caller (CLI / UI) can prompt the user and resume.
            raise
        except Exception as e:
            logger.error(f"Error running workflow {self.workflow_type}: {e}")
            raise

class HumanInputRequired(Exception):
    """Raised when the graph needs human input but no handler is configured.

    Carries the question text so that external callers (CLI, UI) can
    present it to the user and resume the graph with
    ``Command(resume=answer)``.
    """

    def __init__(self, question: str):
        """Initialize the exception with the pending human question.

        Parameters
        ----------
        question : str
            Question that should be presented to the user.
        """
        self.question = question
        super().__init__(question)