"""Multi-agent workflow using the LangGraph Send() (map-reduce) pattern. Architecture ------------ Main graph (``PlannerState``):: Planner --condition--> Send(executor_subgraph, task1..N) --> Planner |-> ResponseAgent --> END (when structured_output) |-> END (when FINISH, no formatting) Executor subgraph (``ExecutorState``):: executor_agent --> ToolNode --> executor_agent (ReAct loop) |-> finalize --> END (no more tool calls) """ import json from typing import Any, Union from functools import partial from pydantic import BaseModel from langgraph.prebuilt import ToolNode from langgraph.graph import StateGraph, END from langchain_core.messages import AIMessage, BaseMessage, ToolMessage from langchain_openai import ChatOpenAI from langgraph.checkpoint.memory import MemorySaver from langgraph.types import Send, interrupt from chemgraph.utils.logging_config import setup_logger from chemgraph.utils.parsing import extract_json_block, parse_response_formatter from chemgraph.state.multi_agent_state import ExecutorState, PlannerState from chemgraph.schemas.multi_agent_response import PlannerResponse from chemgraph.prompt.multi_agent_prompt import ( planner_prompt as default_planner_prompt, executor_prompt as default_executor_prompt, formatter_multi_prompt as default_formatter_prompt, ) logger = setup_logger(__name__) # --------------------------------------------------------------------------- # Serialization helpers # --------------------------------------------------------------------------- def _to_jsonable(obj: Any) -> Any: """Recursively convert Pydantic models to plain dictionaries. Parameters ---------- obj : Any Object to convert. Returns ------- Any JSON-compatible version of the object where possible. """ if isinstance(obj, BaseModel): return obj.model_dump() elif isinstance(obj, dict): return {k: _to_jsonable(v) for k, v in obj.items()} elif isinstance(obj, (list, tuple)): return [_to_jsonable(v) for v in obj] else: return obj def sanitize_tool_calls(messages: list[BaseMessage]) -> list[BaseMessage]: """Ensure tool_call['args'] contains only JSON-serializable data. After LangChain's ToolNode validates tool-call arguments against Pydantic schemas (e.g. ``ASEInputSchema``), nested calculator dicts may be replaced by live Pydantic objects (e.g. ``MaceCalc``). When these messages are later re-sent to the LLM, LangChain serialises ``tool_call['args']`` with ``json.dumps`` — which raises ``TypeError`` for Pydantic instances. This function walks every ``AIMessage.tool_calls`` entry and recursively converts Pydantic models back to plain dicts. Parameters ---------- messages : list[BaseMessage] LangChain messages that may contain tool-call arguments. Returns ------- list[BaseMessage] Messages with JSON-serializable tool-call arguments. """ for m in messages: if isinstance(m, AIMessage) and getattr(m, "tool_calls", None): new_tool_calls = [] for tc in m.tool_calls: tc = dict(tc) tc["args"] = _to_jsonable(tc.get("args")) new_tool_calls.append(tc) m.tool_calls = new_tool_calls return messages # --------------------------------------------------------------------------- # Planner helpers # --------------------------------------------------------------------------- def _parse_planner_response( raw_text: str, ) -> tuple[PlannerResponse | None, str | None]: """Parse raw LLM text into a :class:`PlannerResponse`. Returns ``(parsed_response, None)`` on success, or ``(None, error_msg)`` on failure. Parameters ---------- raw_text : str Raw planner model output. Returns ------- tuple[PlannerResponse | None, str | None] Parsed response and optional parse error. """ # 1. Direct validation try: return PlannerResponse.model_validate_json(raw_text.strip()), None except Exception: pass # 2. Extract JSON block (handles ```json ... ``` or bare {}) extracted = extract_json_block(raw_text) if extracted: try: return PlannerResponse.model_validate_json(extracted), None except Exception: pass try: return PlannerResponse.model_validate(json.loads(extracted)), None except Exception: pass # 3. All attempts failed return None, f"Could not parse planner response from: {raw_text[:200]}" # --------------------------------------------------------------------------- # Planner node # --------------------------------------------------------------------------- def planner_agent( state: PlannerState, llm: ChatOpenAI, system_prompt: str, max_retries: int = 1, ): """Planner that decomposes tasks and routes the workflow. On the first invocation it sees only the user query in ``messages``. On subsequent invocations it also sees ``executor_results`` from completed executor subgraphs and can decide to re-plan or finish. The LLM is prompted to return a JSON object matching the ``PlannerResponse`` schema. If parsing fails, the LLM is retried up to ``max_retries`` times with error feedback. Parameters ---------- state : PlannerState Current global planner state. llm : ChatOpenAI Chat model used for planning. system_prompt : str Planner system prompt. max_retries : int, optional Number of parse-retry attempts after invalid planner output. Returns ------- dict Planner state update containing messages, next step, tasks, and iteration count. """ executor_outputs = state.get("executor_results", []) failed_tasks = state.get("failed_tasks", []) content_block = f"Current Conversation History: {state['messages']}" if executor_outputs: results_text = "\n".join( m.content if hasattr(m, "content") else str(m) for m in executor_outputs ) content_block += ( f"\n\n### UPDATED: Results from Executor Tasks ###\n{results_text}" ) if failed_tasks: failure_lines = [] for ft in failed_tasks: failure_lines.append( f"- Task {ft.get('task_index', '?')} " f"(retry #{ft.get('retry_count', 0)}): " f"{ft.get('error', 'unknown error')}" ) content_block += ( "\n\n### FAILED TASKS (may be retried) ###\n" + "\n".join(failure_lines) + "\n\nYou may retry failed tasks by including them in your " "tasks list with the same task_index. Use the error information " "above to adjust the prompt if needed (e.g., fix molecule names, " "adjust parameters). If a task cannot succeed, set next_step " "to FINISH and explain the failure in thought_process." ) messages = [ {"role": "system", "content": system_prompt}, {"role": "user", "content": content_block}, ] raw_response = llm.invoke(messages).content response_obj, parse_error = _parse_planner_response(raw_response) retries = 0 while response_obj is None and retries < max_retries: retries += 1 logger.warning( "Planner: parse attempt %d failed (%s); retrying.", retries, parse_error, ) retry_messages = messages + [ {"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 " "required format. No markdown fences, no text outside " "the JSON." ), }, ] raw_response = llm.invoke(retry_messages).content response_obj, parse_error = _parse_planner_response(raw_response) if response_obj is None: raise ValueError( f"Planner failed to produce valid JSON after " f"{max_retries} retries: {parse_error}" ) logger.info("PLANNER: %s", response_obj.model_dump_json()) current_iterations = state.get("planner_iterations", 0) result = { "messages": [AIMessage(content=response_obj.thought_process)], "next_step": response_obj.next_step, "tasks": response_obj.tasks if response_obj.tasks else [], "planner_iterations": current_iterations + 1, } if response_obj.next_step == "ask_human" and response_obj.clarification: result["clarification"] = response_obj.clarification return result # --------------------------------------------------------------------------- # Human review node (interrupt for human-in-the-loop) # --------------------------------------------------------------------------- def human_review_node(state: PlannerState): """Pause the graph and ask the human for clarification. This node calls ``interrupt()`` with the planner's clarification question. Execution halts until a human provides a response via ``Command(resume=...)``. The human's answer is injected back into the conversation as an ``AIMessage`` summarising what was asked and what the human replied, then control returns to the Planner. Parameters ---------- state : PlannerState Current planner state containing the clarification question. Returns ------- dict Planner state update containing the human clarification message. """ question = state.get("clarification", "Could you please provide more details?") logger.info("HUMAN_REVIEW: interrupting with question: %s", question) human_response = interrupt({"question": question}) # Normalise the response to a plain string. if isinstance(human_response, dict): answer = human_response.get( "answer", human_response.get("response", str(human_response)) ) else: answer = str(human_response) logger.info("HUMAN_REVIEW: received response: %s", answer) return { "messages": [ AIMessage( content=( f"Human clarification received.\n" f"Question: {question}\n" f"Answer: {answer}" ) ) ], } # --------------------------------------------------------------------------- # Planner router (conditional edge) # --------------------------------------------------------------------------- def unified_planner_router( state: PlannerState, structured_output: bool = False, max_planner_iterations: int = 3, max_task_retries: int = 2, ) -> Union[str, list[Send]]: """Route based on the planner's ``next_step`` decision. * ``executor_subgraph`` -- fan-out tasks via ``Send()`` * ``ask_human`` -- pause for human clarification via ``human_review`` * ``FINISH`` -- go to ``ResponseAgent`` (if structured_output) or ``END`` A cycle guard forces ``FINISH`` when the planner has dispatched executors ``max_planner_iterations`` times to prevent infinite loops. For retried tasks, the ``retry_count`` from the ``WorkerTask`` is checked against ``max_task_retries``. Tasks that have exceeded the retry limit are skipped and logged as permanently failed. Parameters ---------- state : PlannerState Current planner state. structured_output : bool, optional Whether to route final output through the response formatter. max_planner_iterations : int, optional Maximum planner dispatch iterations before forcing completion. max_task_retries : int, optional Maximum retry count for individual executor tasks. Returns ------- str or list[Send] Next graph node name, ``END``, or fan-out ``Send`` instructions. """ next_step = state.get("next_step") iterations = state.get("planner_iterations", 0) if next_step == "ask_human": return "human_review" if next_step == "executor_subgraph": if iterations > max_planner_iterations: logger.warning( "Planner exceeded max iterations (%d); forcing FINISH.", max_planner_iterations, ) if structured_output: return "ResponseAgent" return END tasks = state.get("tasks", []) # Build a lookup of previous failure counts from state. # This covers cases where the planner emits a task without # explicitly setting retry_count — we infer it from history. failed_history: dict[int, int] = {} for ft in state.get("failed_tasks", []): tidx = ft.get("task_index", -1) prev = ft.get("retry_count", 0) # Track the highest retry_count seen for each task_index failed_history[tidx] = max(failed_history.get(tidx, 0), prev + 1) sends = [] for i, t in enumerate(tasks): task_index = getattr(t, "task_index", i + 1) # Determine retry_count: use whichever is larger — # the value from the task object or the inferred history. task_retry = getattr(t, "retry_count", 0) inferred_retry = failed_history.get(task_index, 0) retry_count = max(task_retry, inferred_retry) if retry_count >= max_task_retries: logger.warning( "Task %d exceeded max retries (%d); skipping.", task_index, max_task_retries, ) continue sends.append( Send( "executor_subgraph", { "executor_id": f"worker_{task_index}", "task_index": task_index, "retry_count": retry_count, "messages": [getattr(t, "prompt", str(t))], }, ) ) if not sends: # All tasks were skipped (max retries exceeded). logger.warning( "All dispatched tasks exceeded retry limits; forcing FINISH." ) if structured_output: return "ResponseAgent" return END return sends # FINISH if structured_output: return "ResponseAgent" return END # --------------------------------------------------------------------------- # Executor subgraph nodes # --------------------------------------------------------------------------- async def executor_model_node( state: ExecutorState, llm: ChatOpenAI, system_prompt: str, tools: list, ): """ReAct reasoning step inside an executor subgraph. Reads its own ``messages`` history, calls the LLM with bound tools, and returns the response. 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 new model message. """ sanitized = sanitize_tool_calls(list(state["messages"])) messages = [{"role": "system", "content": system_prompt}] + sanitized # Flatten MCP/LangChain content blocks to plain text for 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) logger.debug("Executor response: %s", response) 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 the last message has tool calls, otherwise ``"done"``. """ last_message = state["messages"][-1] if hasattr(last_message, "tool_calls") and last_message.tool_calls: return "tools" return "done" _ERROR_MARKERS = [ "Error:", "error:", "Exception:", "exception:", "Traceback", "failed", "FAILED", "could not", "Could not", "No PubChem compound found", "ValueError", "TypeError", "KeyError", "RuntimeError", ] def _detect_executor_failure(messages: list) -> tuple[bool, str | None]: """Scan executor message history for signs of failure. Checks for: 1. ``ToolMessage`` objects with ``status == "error"`` (produced by ``ToolNode(handle_tool_errors=True)``). 2. Error markers in the final assistant message content. Returns ``(is_failed, error_summary)``. Parameters ---------- messages : list Executor message history. Returns ------- tuple[bool, str | None] Failure flag and optional error summary. """ # Collect all tool-level errors tool_errors = [] for m in messages: if isinstance(m, ToolMessage): if getattr(m, "status", None) == "error": tool_errors.append(m.content) if tool_errors: return True, "; ".join(tool_errors) # Check the final message for error markers final = messages[-1] if messages else None if final is not None: content = getattr(final, "content", str(final)) if isinstance(content, str): for marker in _ERROR_MARKERS: if marker in content: # Only flag as failure if the executor itself reports failure, # not if it's merely describing a prior error it recovered from. # Heuristic: if the last message also contains "success" or # "result", treat it as a recovered scenario. lower = content.lower() if "success" not in lower and "result:" not in lower: return True, content[:500] return False, None def format_executor_output(state: ExecutorState) -> dict: """Bridge: convert local ``ExecutorState`` into a ``PlannerState`` update. Writes the executor's final answer into ``executor_results`` and its full message history into ``executor_logs`` so the planner can inspect them on the next iteration. Detects executor failures by scanning the message history for tool errors and error markers. When a failure is detected, populates ``failed_tasks`` so the planner can decide whether to retry. Parameters ---------- state : ExecutorState Local executor state at subgraph completion. Returns ------- dict Planner-state update with executor results, logs, and failure data. """ executor_id = state["executor_id"] task_index = state.get("task_index", -1) retry_count = state.get("retry_count", 0) final_message = state["messages"][-1].content full_history = state["messages"] is_failed, error_summary = _detect_executor_failure(list(state["messages"])) result: dict = { "executor_logs": {executor_id: full_history}, } if is_failed: logger.warning( "Executor %s (task_index=%d, retry=%d) FAILED: %s", executor_id, task_index, retry_count, error_summary, ) result["executor_results"] = [ f"[{executor_id}] FAILED (task_index={task_index}, " f"retry={retry_count}): {error_summary}" ] result["failed_tasks"] = [ { "task_index": task_index, "executor_id": executor_id, "error": error_summary, "retry_count": retry_count, } ] else: result["executor_results"] = [ f"[{executor_id}] Result (task_index={task_index}): {final_message}" ] result["failed_tasks"] = [] return result def construct_executor_subgraph( llm: ChatOpenAI, tools: list, system_prompt: str, ): """Build the reusable executor subgraph (Agent -> Tools -> Agent loop). The subgraph is compiled and used as a node in the main graph. Each ``Send()`` invocation creates an independent copy with its own ``ExecutorState``. 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, handle_tool_errors=True)) workflow.add_node("finalize", format_executor_output) 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() # --------------------------------------------------------------------------- # Response agent (prompt-based, same approach as single_agent.py) # --------------------------------------------------------------------------- def response_agent( state: PlannerState, llm: ChatOpenAI, formatter_prompt: str, max_retries: int = 1, ): """Format the final answer using a prompt (no ``with_structured_output``). Mirrors the ``ResponseAgent`` from ``single_agent.py``: invokes the LLM with a formatter prompt and manually parses the response into a ``ResponseFormatter`` with retry logic on parse failure. Parameters ---------- state : PlannerState Final planner state to summarize. llm : ChatOpenAI Chat model used for response formatting. formatter_prompt : str Prompt instructing the model how to format the final answer. max_retries : int, optional Number of parse-retry attempts after invalid formatter output. Returns ------- dict State update containing the formatted response message. """ messages = [ {"role": "system", "content": formatter_prompt}, {"role": "user", "content": f"{state['messages']}"}, ] raw_response = llm.invoke(messages).content formatter, parse_error = parse_response_formatter(raw_response) 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": f"{state['messages']}"}, {"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) 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]} # --------------------------------------------------------------------------- # Main graph constructor # --------------------------------------------------------------------------- def construct_multi_agent_graph( llm: ChatOpenAI, planner_prompt: str = default_planner_prompt, executor_prompt: str = default_executor_prompt, executor_tools: list = None, structured_output: bool = False, formatter_prompt: str = default_formatter_prompt, max_retries: int = 1, max_task_retries: int = 2, ): """Construct the planner-executor graph using the Send() pattern. Parameters ---------- llm : ChatOpenAI The language model shared by all agents. planner_prompt : str System prompt for the planner agent. executor_prompt : str System prompt for each executor subgraph. executor_tools : list Tools available to executor agents (LangChain tools or MCP tools). structured_output : bool If ``True``, route to ``ResponseAgent`` for structured formatting before ending. If ``False``, the workflow ends directly after the planner decides ``FINISH``. formatter_prompt : str System prompt for the ``ResponseAgent`` (used only when ``structured_output=True``). max_retries : int Number of LLM retry attempts when the planner or response agent fails to parse its output, by default 1. max_task_retries : int Maximum number of times a single executor task may be retried after failure. Once a task reaches this limit, the router skips it and the planner must finish without it, by default 2. Returns ------- CompiledGraph A compiled LangGraph state graph. """ if executor_tools is None: 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.generic_tools import calculator executor_tools = [ molecule_name_to_smiles, smiles_to_coordinate_file, run_ase, extract_output_json, calculator, ] checkpointer = MemorySaver() # Build the executor subgraph executor_subgraph = construct_executor_subgraph( llm, executor_tools, executor_prompt ) # Build the main graph graph_builder = StateGraph(PlannerState) # -- Nodes -- graph_builder.add_node( "Planner", lambda state: planner_agent( state, llm, planner_prompt, max_retries=max_retries ), ) graph_builder.add_node("executor_subgraph", executor_subgraph) graph_builder.add_node("human_review", human_review_node) # Conditional destinations list for the planner router conditional_targets = ["executor_subgraph", "human_review", END] if structured_output: graph_builder.add_node( "ResponseAgent", lambda state: response_agent( state, llm, formatter_prompt=formatter_prompt, max_retries=max_retries, ), ) conditional_targets.append("ResponseAgent") # -- Edges -- graph_builder.set_entry_point("Planner") graph_builder.add_conditional_edges( "Planner", partial( unified_planner_router, structured_output=structured_output, max_task_retries=max_task_retries, ), conditional_targets, ) # Executors feed results back to the planner graph_builder.add_edge("executor_subgraph", "Planner") # After human clarification, return to the planner for re-planning graph_builder.add_edge("human_review", "Planner") if structured_output: graph_builder.add_edge("ResponseAgent", END) graph = graph_builder.compile(checkpointer=checkpointer) logger.info("Multi-agent graph (Send pattern) constructed successfully") return graph