"""Command implementations for the ChemGraph CLI. Each public function corresponds to a CLI action: running a query, starting interactive mode, managing sessions, etc. """ from __future__ import annotations import os import time from concurrent.futures import ThreadPoolExecutor, TimeoutError as FuturesTimeoutError from typing import Any, Dict, Optional from rich.panel import Panel from rich.progress import Progress, SpinnerColumn, TextColumn from rich.prompt import Prompt from rich.table import Table from chemgraph.memory.store import SessionStore from chemgraph.models.supported_models import ( supported_alcf_models, supported_anthropic_models, supported_gemini_models, supported_ollama_models, supported_openai_models, supported_argo_models, ) from chemgraph.utils.async_utils import run_async_callable from chemgraph.cli.formatting import ( console, create_banner, format_response, ) # --------------------------------------------------------------------------- # Workflow helpers # --------------------------------------------------------------------------- # All workflow types registered in ChemGraph.workflow_map ALL_WORKFLOW_TYPES = [ "single_agent", "multi_agent", "python_relp", "graspa", "mock_agent", "single_agent_mcp", "graspa_mcp", "rag_agent", "single_agent_xanes", ] # Common aliases so users can type the "obvious" name. WORKFLOW_ALIASES: Dict[str, str] = { "python_repl": "python_relp", "graspa_agent": "graspa", } def resolve_workflow(name: str) -> str: """Resolve a workflow name, applying aliases. Parameters ---------- name : str Workflow name or supported alias. Returns ------- str Canonical workflow name. """ return WORKFLOW_ALIASES.get(name, name) # --------------------------------------------------------------------------- # API-key validation # --------------------------------------------------------------------------- def check_api_keys(model_name: str) -> tuple[bool, str]: """Check if required API keys are available for *model_name*. Parameters ---------- model_name : str Model identifier selected for a run. Returns ------- tuple[bool, str] ``(is_available, error_message)``. The message is empty when the required credentials are available or not required. """ model_lower = model_name.lower() # OpenAI models (including GPT family, o-series, and Argo OpenAI) if ( model_name in supported_openai_models or model_name in supported_argo_models or model_lower.startswith("gpt") or any(prefix in model_lower for prefix in ["o1", "o3", "o4"]) ): # Argo models use a different auth mechanism; skip key check. if model_name in supported_argo_models: pass elif not os.getenv("OPENAI_API_KEY"): return ( False, "OpenAI API key not found. Set the OPENAI_API_KEY environment variable.", ) # Anthropic models elif "claude" in model_lower or model_name in supported_anthropic_models: if not os.getenv("ANTHROPIC_API_KEY"): return ( False, "Anthropic API key not found. Set the ANTHROPIC_API_KEY environment variable.", ) # Google models elif "gemini" in model_lower or model_name in supported_gemini_models: if not os.getenv("GEMINI_API_KEY"): return ( False, "Gemini API key not found. Set the GEMINI_API_KEY environment variable.", ) # GROQ models (groq: prefix) elif model_name.startswith("groq:"): if not os.getenv("GROQ_API_KEY"): return ( False, "GROQ API key not found. Set the GROQ_API_KEY environment variable.", ) # ALCF models (Globus OAuth access token) elif model_name in supported_alcf_models: if not os.getenv("ALCF_ACCESS_TOKEN"): return ( False, "ALCF access token not found. To authenticate with ALCF:\n" " 1. pip install globus_sdk\n" " 2. wget https://raw.githubusercontent.com/argonne-lcf/" "inference-endpoints/refs/heads/main/inference_auth_token.py\n" " 3. python inference_auth_token.py authenticate\n" " 4. export ALCF_ACCESS_TOKEN=$(python inference_auth_token.py get_access_token)\n" "\n" " See: https://docs.alcf.anl.gov/services/inference-endpoints/#api-access", ) # Local models (no API key needed) elif model_name in supported_ollama_models or any( local in model_lower for local in ["llama", "qwen", "ollama"] ): pass return True, "" # --------------------------------------------------------------------------- # Agent initialization # --------------------------------------------------------------------------- _INIT_TIMEOUT_SECONDS = 30 def initialize_agent( model_name: str, workflow_type: str, structured_output: bool, return_option: str, generate_report: bool, recursion_limit: int, base_url: Optional[str] = None, argo_user: Optional[str] = None, verbose: bool = False, human_supervised: bool = False, tools: Optional[list] = None, ) -> Any: """Initialize a ChemGraph agent with progress indication. Uses a thread-pool executor for the timeout so it works on all platforms. Parameters ---------- model_name : str LLM model identifier. workflow_type : str ChemGraph workflow name or alias. structured_output : bool Whether to request structured final output. return_option : str Agent return mode, such as ``"state"`` or ``"last_message"``. generate_report : bool Whether the agent should generate an HTML report. recursion_limit : int LangGraph recursion limit for the run. base_url : str, optional Custom model endpoint URL. argo_user : str, optional Argo username for Argo-hosted models. verbose : bool, optional Whether to print initialization details. human_supervised : bool, optional Whether to enable human-interrupt tooling. tools : list, optional Custom tool list for MCP-backed workflows. Returns ------- Any Initialized ``ChemGraph`` instance, or ``None`` when initialization fails. """ # Resolve workflow alias before initializing. workflow_type = resolve_workflow(workflow_type) if verbose: console.print("[blue]Initializing agent with:[/blue]") console.print(f" Model: {model_name}") console.print(f" Workflow: {workflow_type}") console.print(f" Structured Output: {structured_output}") console.print(f" Return Option: {return_option}") console.print(f" Generate Report: {generate_report}") console.print(f" Human Supervised: {human_supervised}") console.print(f" Recursion Limit: {recursion_limit}") if base_url: console.print(f" Base URL: {base_url}") if argo_user: console.print(f" Argo User: {argo_user}") if tools: console.print(f" MCP Tools: {len(tools)} loaded") # Check API keys before attempting initialization api_key_available, error_msg = check_api_keys(model_name) if not api_key_available: console.print(f"[red]{error_msg}[/red]") console.print( "[dim]Tip: Set environment variables in your shell or .env file[/dim]" ) console.print( "[dim] Example: export OPENAI_API_KEY='your_api_key_here'[/dim]" ) return None # Resolve API key for providers that need one passed explicitly. api_key: Optional[str] = None if model_name in supported_alcf_models: api_key = os.getenv("ALCF_ACCESS_TOKEN") with Progress( SpinnerColumn(), TextColumn("[progress.description]{task.description}"), console=console, transient=True, ) as progress: task = progress.add_task("Initializing ChemGraph agent...", total=None) def _create_agent() -> Any: """Create the ChemGraph agent inside the initialization worker. Returns ------- Any Initialized ``ChemGraph`` instance. """ from chemgraph.agent.llm_agent import ChemGraph return ChemGraph( model_name=model_name, workflow_type=workflow_type, base_url=base_url, api_key=api_key, argo_user=argo_user, generate_report=generate_report, return_option=return_option, recursion_limit=recursion_limit, structured_output=structured_output, human_supervised=human_supervised, tools=tools, ) try: with ThreadPoolExecutor(max_workers=1) as pool: future = pool.submit(_create_agent) agent = future.result(timeout=_INIT_TIMEOUT_SECONDS) progress.update(task, description="[green]Agent initialized successfully!") time.sleep(0.5) return agent except FuturesTimeoutError: progress.update(task, description="[red]Agent initialization timed out!") console.print( f"[red]Agent initialization timed out after {_INIT_TIMEOUT_SECONDS}s[/red]" ) console.print( "[dim]This might indicate network issues or invalid API credentials[/dim]" ) return None except Exception as e: progress.update(task, description="[red]Agent initialization failed!") console.print(f"[red]Error initializing agent: {e}[/red]") err_str = str(e).lower() if "authentication" in err_str or "api" in err_str: console.print( "[dim]This looks like an API key issue. Check your credentials.[/dim]" ) elif "connection" in err_str or "network" in err_str: console.print( "[dim]This looks like a network connectivity issue.[/dim]" ) return None # --------------------------------------------------------------------------- # Query execution # --------------------------------------------------------------------------- # Thread-ID counter for interactive mode so each query gets unique state. _thread_counter: int = 0 def _next_thread_id() -> int: """Return the next interactive-mode thread ID. Returns ------- int Incremented thread ID. """ global _thread_counter _thread_counter += 1 return _thread_counter def run_query( agent: Any, query: str, thread_id: Optional[int] = None, verbose: bool = False, resume_from: Optional[str] = None, ) -> Any: """Execute a query with the agent. When the graph pauses for human input (``HumanInputRequired``), the spinner is stopped, the question is shown in a Rich panel, and the user is prompted for a response. The graph is then resumed with the user's answer and the spinner restarts. This loop repeats until the graph completes or a non-interrupt error occurs. Parameters ---------- agent : Any Initialized ChemGraph-like agent with ``run`` and ``workflow`` methods. query : str User query to execute. thread_id : int, optional LangGraph thread identifier. A new ID is allocated when omitted. verbose : bool, optional Whether to print execution details. resume_from : str, optional Previous ChemGraph session ID to load as context. Returns ------- Any Agent result, resumed graph result, or ``None`` on failure. """ from langgraph.types import Command from chemgraph.agent.llm_agent import HumanInputRequired if thread_id is None: thread_id = _next_thread_id() if verbose: console.print(f"[blue]Executing query:[/blue] {query}") console.print(f"[blue]Thread ID:[/blue] {thread_id}") if resume_from: console.print(f"[blue]Resuming from session:[/blue] {resume_from}") config = {"configurable": {"thread_id": thread_id}} max_interrupts = 10 # safety guard interrupt_count = 0 # --- First invocation: run the full agent.run() --- with Progress( SpinnerColumn(), TextColumn("[progress.description]{task.description}"), console=console, transient=True, ) as progress: task = progress.add_task("Processing query...", total=None) try: result = run_async_callable( lambda: agent.run(query, config=config, resume_from=resume_from) ) progress.update(task, description="[green]Query completed!") time.sleep(0.3) return result except HumanInputRequired as hir: progress.update(task, description="[yellow]Agent needs your input") time.sleep(0.2) question = hir.question except Exception as e: progress.update(task, description="[red]Query failed!") console.print(f"[red]Error processing query: {e}[/red]") return None # --- Interrupt-resume loop --- # The spinner's `with` block has exited, so the terminal is free # for interactive user input. while question is not None: interrupt_count += 1 if interrupt_count > max_interrupts: console.print( "[red]Exceeded maximum number of human interrupts. Aborting.[/red]" ) return None console.print( Panel( question, title="[bold yellow]Agent needs your input[/bold yellow]", style="yellow", ) ) human_answer = Prompt.ask("[bold cyan]Your response[/bold cyan]") # Resume the graph, streaming messages so tool-call parameters # are printed just like the initial invocation. resume_config = dict(config) resume_config["recursion_limit"] = agent.recursion_limit async def _resume_stream(): """Resume an interrupted graph and stream updates until completion. Returns ------- dict or None Final streamed graph state. """ prev_msgs: list = [] last_st = None async for s in agent.workflow.astream( Command(resume=human_answer), stream_mode="values", config=resume_config, ): if "messages" in s and s["messages"] != prev_msgs: new_message = s["messages"][-1] try: new_message.pretty_print() except Exception: pass prev_msgs = s["messages"] last_st = s return last_st try: result = run_async_callable(_resume_stream) if result is None: console.print("[red]Resume produced no output.[/red]") return None if agent.return_option == "last_message": return result["messages"][-1] if result else None elif agent.return_option == "state": from chemgraph.agent.llm_agent import serialize_state return serialize_state(agent.get_state(config=config)) return result except HumanInputRequired as hir: question = hir.question except Exception as e: console.print(f"[red]Error processing query: {e}[/red]") return None return None # --------------------------------------------------------------------------- # Session management # --------------------------------------------------------------------------- def list_sessions(limit: int = 20, db_path: Optional[str] = None) -> None: """Display recent sessions in a formatted table. Parameters ---------- limit : int, optional Maximum number of sessions to display. db_path : str, optional Path to the session SQLite database. """ store = SessionStore(db_path=db_path) sessions = store.list_sessions(limit=limit) if not sessions: console.print("[dim]No sessions found.[/dim]") return console.print(Panel(f"Recent Sessions ({len(sessions)})", style="bold cyan")) table = Table(show_header=True, header_style="bold magenta") table.add_column("Session ID", style="cyan", width=10) table.add_column("Title", style="white", width=40) table.add_column("Model", style="green", width=16) table.add_column("Workflow", style="yellow", width=14) table.add_column("Queries", style="white", justify="right", width=8) table.add_column("Messages", style="white", justify="right", width=9) table.add_column("Date", style="dim", width=16) for s in sessions: table.add_row( s.session_id, s.title or "[dim]Untitled[/dim]", s.model_name, s.workflow_type, str(s.query_count), str(s.message_count), s.updated_at.strftime("%Y-%m-%d %H:%M"), ) console.print(table) console.print( "\n[dim]Use 'chemgraph session show ' to view a session. " "Prefix matching is supported.[/dim]" ) def show_session( session_id: str, db_path: Optional[str] = None, max_content: int = 500, ) -> None: """Display a session's full conversation. Parameters ---------- session_id : str Session ID or unique session prefix. db_path : str, optional Path to the session SQLite database. max_content : int, optional Maximum number of characters displayed for each message. """ store = SessionStore(db_path=db_path) session = store.get_session(session_id) if session is None: console.print( f"[red]Session '{session_id}' not found. " f"The ID may be ambiguous or nonexistent.[/red]" ) console.print("[dim]Use 'chemgraph session list' to see available sessions.[/dim]") return # Session metadata header meta_table = Table(show_header=False, box=None, padding=(0, 2)) meta_table.add_column("Key", style="bold cyan") meta_table.add_column("Value") meta_table.add_row("Session ID", session.session_id) meta_table.add_row("Title", session.title or "Untitled") meta_table.add_row("Model", session.model_name) meta_table.add_row("Workflow", session.workflow_type) meta_table.add_row("Queries", str(session.query_count)) meta_table.add_row("Created", session.created_at.strftime("%Y-%m-%d %H:%M:%S")) meta_table.add_row("Updated", session.updated_at.strftime("%Y-%m-%d %H:%M:%S")) if session.log_dir: meta_table.add_row("Log Dir", session.log_dir) console.print(Panel(meta_table, title="Session Info", style="bold cyan")) if not session.messages: console.print("[dim]No messages in this session.[/dim]") return # Display conversation console.print(f"\n[bold]Conversation ({len(session.messages)} messages):[/bold]\n") for msg in session.messages: if msg.role == "human": label = "[bold cyan]User[/bold cyan]" elif msg.role == "ai": label = "[bold green]Assistant[/bold green]" elif msg.role == "tool": tool_label = f" ({msg.tool_name})" if msg.tool_name else "" label = f"[bold yellow]Tool{tool_label}[/bold yellow]" else: label = f"[dim]{msg.role}[/dim]" content = msg.content if max_content and len(content) > max_content: content = ( content[:max_content] + f"\n... [truncated, {len(msg.content)} chars total]" ) timestamp = msg.timestamp.strftime("%H:%M:%S") if msg.timestamp else "" console.print(f" {label} [dim]{timestamp}[/dim]") console.print(f" {content}\n") def delete_session_cmd(session_id: str, db_path: Optional[str] = None) -> None: """Delete a session from the database. Parameters ---------- session_id : str Session ID or unique session prefix to delete. db_path : str, optional Path to the session SQLite database. """ store = SessionStore(db_path=db_path) # Show session info before deleting session = store.get_session(session_id) if session is None: console.print(f"[red]Session '{session_id}' not found.[/red]") return console.print( f"[yellow]Deleting session: {session.session_id} " f"({session.title or 'Untitled'})[/yellow]" ) if store.delete_session(session_id): console.print("[green]Session deleted.[/green]") else: console.print("[red]Failed to delete session.[/red]") # --------------------------------------------------------------------------- # Output helpers # --------------------------------------------------------------------------- def save_output(content: str, output_file: str) -> None: """Save output to a file. Parameters ---------- content : str Text content to write. output_file : str Destination file path. """ try: with open(output_file, "w") as f: f.write(content) console.print(f"[green]Output saved to: {output_file}[/green]") except Exception as e: console.print(f"[red]Error saving output: {e}[/red]") # --------------------------------------------------------------------------- # Interactive REPL # --------------------------------------------------------------------------- def interactive_mode( model: str = "gpt-4o-mini", workflow: str = "single_agent", structured: bool = False, return_option: str = "state", generate_report: bool = True, human_supervised: bool = False, recursion_limit: int = 20, base_url: Optional[str] = None, argo_user: Optional[str] = None, verbose: bool = False, tools: Optional[list] = None, ) -> None: """Start interactive REPL mode for ChemGraph CLI. Accepts the same configuration parameters as a normal run so that ``--config`` and CLI flags are honoured when entering interactive mode. Parameters ---------- model : str, optional Initial model selection. workflow : str, optional Initial workflow selection. structured : bool, optional Whether structured output is requested. return_option : str, optional Agent return mode. generate_report : bool, optional Whether report generation is enabled. human_supervised : bool, optional Whether human supervision tools are enabled. recursion_limit : int, optional LangGraph recursion limit. base_url : str, optional Custom model endpoint URL. argo_user : str, optional Argo username for Argo-hosted models. verbose : bool, optional Whether to print diagnostic output. tools : list, optional Custom tool list for MCP-backed workflows. """ console.print(create_banner()) console.print("[bold green]Welcome to ChemGraph Interactive Mode![/bold green]") console.print( "Type your queries and get AI-powered computational chemistry insights." ) console.print( "[dim]Type 'quit', 'exit', or 'q' to exit. Type 'help' for commands.[/dim]\n" ) # Allow the user to override model/workflow at startup. model = Prompt.ask( "Select model (or type a custom model ID)", default=model ) workflow = Prompt.ask( "Select workflow", choices=ALL_WORKFLOW_TYPES, default=resolve_workflow(workflow), ) # Initialize agent with the full config context. agent = initialize_agent( model, workflow, structured, return_option, generate_report, recursion_limit, base_url=base_url, argo_user=argo_user, verbose=verbose, human_supervised=human_supervised, tools=tools, ) if not agent: return console.print( "[green]Ready! You can now ask computational chemistry questions.[/green]\n" ) while True: try: query = Prompt.ask("\n[bold cyan]ChemGraph[/bold cyan]") if query.lower() in ("quit", "exit", "q"): console.print("[yellow]Goodbye![/yellow]") break elif query.lower() == "help": console.print( Panel( """ Available commands: quit/exit/q Exit interactive mode help Show this help message clear Clear screen config Show current configuration model Change model workflow Change workflow type Session commands: history List recent sessions show Show a session's conversation resume Resume from a previous session Example queries: What is the SMILES string for water? Optimize the geometry of methane Calculate CO2 vibrational frequencies Show me the structure of caffeine """, title="Help", style="blue", ) ) continue elif query.lower() == "clear": console.clear() continue elif query.lower() == "config": console.print(f"Model: {model}") console.print(f"Workflow: {workflow}") if hasattr(agent, "session_id"): console.print(f"Session ID: {agent.session_id}") continue elif query.lower() == "history": list_sessions() continue elif query.lower().startswith("show "): sid = query[5:].strip() if sid: show_session(sid) else: console.print("[red]Usage: show [/red]") continue elif query.lower().startswith("resume "): sid = query[7:].strip() if not sid: console.print("[red]Usage: resume [/red]") continue resume_query = Prompt.ask( "[bold cyan]Enter query to continue with[/bold cyan]" ) if resume_query.strip(): result = run_query( agent, resume_query, verbose=verbose, resume_from=sid, ) if result: format_response(result, verbose=verbose) continue elif query.startswith("model "): new_model = query[6:].strip() model = new_model agent = initialize_agent( model, workflow, structured, return_option, generate_report, recursion_limit, base_url=base_url, argo_user=argo_user, human_supervised=human_supervised, tools=tools, ) if agent: console.print(f"[green]Model changed to: {model}[/green]") continue elif query.startswith("workflow "): new_workflow = resolve_workflow(query[9:].strip()) if new_workflow in ALL_WORKFLOW_TYPES: workflow = new_workflow agent = initialize_agent( model, workflow, structured, return_option, generate_report, recursion_limit, base_url=base_url, argo_user=argo_user, human_supervised=human_supervised, tools=tools, ) if agent: console.print( f"[green]Workflow changed to: {workflow}[/green]" ) else: console.print(f"[red]Invalid workflow: {new_workflow}[/red]") console.print( f"[dim]Available: {', '.join(ALL_WORKFLOW_TYPES)}[/dim]" ) continue # Execute query (each query gets a unique thread ID) result = run_query(agent, query, verbose=verbose) if result: format_response(result, verbose=verbose) if hasattr(agent, "session_id") and agent.session_id: console.print(f"[dim]Session: {agent.session_id}[/dim]") except KeyboardInterrupt: console.print( "\n[yellow]Interrupted. Type 'quit' to exit.[/yellow]" ) except Exception as e: console.print(f"[red]Error: {e}[/red]")