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] + "..." 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