diff --git a/.gitignore b/.gitignore index 04ce48ac..6c3ed36a 100644 --- a/.gitignore +++ b/.gitignore @@ -143,6 +143,10 @@ vite.config.ts.timestamp-* .venv/ venv/ __pycache__/ +*.egg-info/ +*.egg +dist/ +build/ # Visual Studio Code .vscode/ diff --git a/ChARGe b/ChARGe new file mode 160000 index 00000000..6bcfc14f --- /dev/null +++ b/ChARGe @@ -0,0 +1 @@ +Subproject commit 6bcfc14fc4c467c03bd4d40421e25aea45109398 diff --git a/charge_backend/backend_helper_funcs.py b/charge_backend/backend_helper_funcs.py index e42c77bb..8d4f175c 100644 --- a/charge_backend/backend_helper_funcs.py +++ b/charge_backend/backend_helper_funcs.py @@ -108,6 +108,7 @@ def json(self): @dataclass class FlaskRunSettings(RunSettings): molecule_name_format: MolNameFormat = Field(alias="moleculeName", default="brand") + # Inherit use_ai_based from RunSettings (defined in LC-Conductor) @dataclass(frozen=True) diff --git a/charge_backend/backend_manager.py b/charge_backend/backend_manager.py index f9def612..6b9d2f1f 100644 --- a/charge_backend/backend_manager.py +++ b/charge_backend/backend_manager.py @@ -50,6 +50,7 @@ def __init__( self.websocket = task_manager.websocket self.builtin_tool_definitions = builtin_tool_definitions or [] self.task_manager.available_builtin_tool_ids = None + self.retro_synth_context = None def _selected_mcp_tools(self) -> list[str]: if self.task_manager.available_tools is None: @@ -225,14 +226,74 @@ async def _handle_optimization( async def _handle_retrosynthesis(self, data: dict) -> None: """Handle retrosynthesis problem type.""" - run_func = partial( - template_based_retrosynthesis, - data["smiles"], - self.args.config_file, - self.get_retro_synth_context(), - self.task_manager.websocket, - self.run_settings, - ) + # Check if AI-based retrosynthesis is requested + # use_ai_based flag controls AI vs template approach + # If no config exists, fall back to AI-based + use_ai_based = self.run_settings.use_ai_based or not os.path.exists(self.args.config_file) + + if use_ai_based: + # Use AI-based retrosynthesis (supports both standard and RSA modes) + # Create root node like template_based_retrosynthesis does + from backend_helper_funcs import Node + from charge_backend.moleculedb.molecule_naming import smiles_to_html + from charge_backend.moleculedb.purchasable import is_purchasable + + context = self.get_retro_synth_context() + context.reset() # Clear context + + start_smiles = data["smiles"] + mol_sources = is_purchasable(start_smiles) + if mol_sources: + purchasable_str = f"Yes (via {', '.join(mol_sources)})" + else: + purchasable_str = "No" + + root = Node( + id="node_0", + smiles=start_smiles, + label=smiles_to_html(start_smiles, self.run_settings.molecule_name_format), + hoverInfo=f"""# Root molecule +**SMILES:** {start_smiles} + +**Purchasable**? {purchasable_str}""", + level=0, + parentId=None, + cost=None, + bandgap=None, + yield_=None, + purchasable=(len(mol_sources) > 0), + highlight="yellow", + x=100, + y=100, + ) + + await context.add_node(root, websocket=self.task_manager.websocket) + root_node_id = root.id + + run_func = partial( + ai_based_retrosynthesis, + root_node_id, + context, + data.get("query", None), + None, # Unconstrained + self.task_manager.websocket, + self.experiment, + self.args.config_file, + self.run_settings, + self._selected_mcp_tools(), + self._selected_builtin_tools(), + self.log_progress, + ) + else: + # Use template-based retrosynthesis + run_func = partial( + template_based_retrosynthesis, + data["smiles"], + self.args.config_file, + self.get_retro_synth_context(), + self.task_manager.websocket, + self.run_settings, + ) await self.task_manager.run_task(run_func()) diff --git a/charge_backend/charge_server.py b/charge_backend/charge_server.py index 13d519ae..49e37654 100644 --- a/charge_backend/charge_server.py +++ b/charge_backend/charge_server.py @@ -165,12 +165,21 @@ async def root(request: Request): with open(os.path.join(DIST_PATH, "index.html"), "r") as fp: html = fp.read() + # Use dynamic WebSocket URL based on request host + ws_server = os.getenv("WS_SERVER") + if not ws_server: + # Auto-detect from request host and protocol + host = request.headers.get("host", f"localhost:{args.port}") + # Use wss:// for HTTPS, ws:// for HTTP + scheme = "wss" if request.url.scheme == "https" else "ws" + ws_server = f"{scheme}://{host}/ws" + html = html.replace( "", f""" """, diff --git a/charge_backend/retrosynthesis/ai.py b/charge_backend/retrosynthesis/ai.py index 58dabe0f..ffce6944 100644 --- a/charge_backend/retrosynthesis/ai.py +++ b/charge_backend/retrosynthesis/ai.py @@ -1,5 +1,9 @@ import os import asyncio +import random +import json +import datetime +from pathlib import Path from fastapi import WebSocket from lc_conductor.callback_logger import CallbackLogger from typing import Any, Callable, Optional, Union @@ -28,6 +32,7 @@ from retrosynthesis.retrosynthesis_task import ( TemplateFreeRetrosynthesisTask as RetrosynthesisTask, TemplateFreeReactionOutputSchema as ReactionOutputSchema, + RSAAggregationTask, ) from charge.experiments.experiment import Experiment @@ -130,13 +135,185 @@ async def ai_based_retrosynthesis( f"Finding synthesis routes for {current_node.smiles} using available tools: {available_tools}." ) - # Run task + # Run task (with optional RSA mode) await highlight_node(current_node, websocket, True) if run_settings.prompt_debugging: await debug_prompt(runner, websocket) - output = await runner.run(log_progress) - if isinstance(callback_handler, CallbackHandler): - await callback_handler.drain() + + # Check if RSA mode is enabled + if run_settings.use_rsa: + try: + # RSA Mode: Recursive Self-Aggregation + from rsa_algorithm import run_rsa_loop + + rsa_n = run_settings.rsa_n if hasattr(run_settings, 'rsa_n') else 8 + rsa_k = run_settings.rsa_k if hasattr(run_settings, 'rsa_k') else 4 + rsa_t = run_settings.rsa_t if hasattr(run_settings, 'rsa_t') else 3 + rsa_mode = run_settings.rsa_mode if hasattr(run_settings, 'rsa_mode') else "standalone" + + await clogger.info( + f"Running RSA mode: {rsa_mode} with N={rsa_n}, K={rsa_k}, T={rsa_t}" + ) + + # Create directory for RSA execution logs + import datetime + timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") + rsa_log_dir = f"/tmp/rsa_execution_{timestamp}" + os.makedirs(rsa_log_dir, exist_ok=True) + await clogger.info(f"RSA execution logs will be saved to: {rsa_log_dir}") + + # For RAG mode: Query database once and inject into prompts + # For standalone mode: Remove database query tool + user_prompt_with_rag = user_prompt + builtin_tools_filtered = builtin_tools or [] + + if rsa_mode == "rag": + await clogger.info("RAG mode: Querying reaction database once...") + try: + from retrosynthesis.database import query_reaction_database + db_results = query_reaction_database(current_node.smiles, top_k=10) + + if db_results and not any("error" in r for r in db_results): + await clogger.info(f"Found {len(db_results)} similar reactions in database") + + # Log summary of database results to UI + summary_lines = [f"**Database Query Results ({len(db_results)} reactions found):**"] + for idx, reaction in enumerate(db_results[:5], 1): # Show first 5 in UI + name = reaction.get('name', f'Reaction {idx}') + summary_lines.append(f" {idx}. {name}") + if 'components' in reaction and reaction['components']: + # Extract reactants and products + reactants = [c.get('name', c.get('smiles', '?')) for c in reaction['components'] + if c.get('role') in ['Reactant', 'Reagent']] + products = [c.get('name', c.get('smiles', '?')) for c in reaction['components'] + if c.get('role') == 'Product'] + if reactants: + summary_lines.append(f" Reactants: {', '.join(reactants[:3])}") + if products: + summary_lines.append(f" Products: {', '.join(products[:2])}") + if len(db_results) > 5: + summary_lines.append(f" ... and {len(db_results) - 5} more reactions") + summary_lines.append("These reactions will be injected into all proposal prompts.") + await clogger.info("\n".join(summary_lines)) + + # Format database results with clear context + rag_context = "\n\n--- REACTION DATABASE RESULTS ---\n" + rag_context += f"These are similar reactions retrieved by comparing structural similarity to the target product ({current_node.smiles}):\n\n" + + for idx, reaction in enumerate(db_results[:10], 1): + rag_context += f"Reaction {idx}:\n" + if "reactants" in reaction: + rag_context += f" Reactants: {reaction.get('reactants', 'N/A')}\n" + if "products" in reaction: + rag_context += f" Products: {reaction.get('products', 'N/A')}\n" + if "text" in reaction and reaction.get("text"): + rag_context += f" Description: {reaction['text']}\n" + rag_context += "\n" + + rag_context += "Use these reactions as supporting evidence for your retrosynthesis proposal.\n" + rag_context += "--- END DATABASE RESULTS ---\n" + + user_prompt_with_rag = user_prompt + rag_context + + # Save database results to log + with open(f"{rsa_log_dir}/database_query_results.json", "w") as f: + json.dump(db_results, f, indent=2) + else: + await clogger.info("No reactions found in database") + user_prompt_with_rag = user_prompt + "\n\nNo similar reactions found in the database for this target molecule.\n" + except Exception as e: + await clogger.warning(f"Database query failed: {str(e)}") + user_prompt_with_rag = user_prompt + "\n\nDatabase query failed. Proceed using chemistry knowledge only.\n" + + # Filter out query_reaction_database from builtin tools (already queried once) + builtin_tools_filtered = [ + tool for tool in builtin_tools_filtered + if getattr(tool, '__name__', '') != 'query_reaction_database' + ] + await clogger.info("RAG mode: Removed query_reaction_database from tools (already queried)") + + elif rsa_mode == "standalone": + # Standalone mode: Remove database query tool entirely (no retrieval) + builtin_tools_filtered = [ + tool for tool in builtin_tools_filtered + if getattr(tool, '__name__', '') != 'query_reaction_database' + ] + await clogger.info("Standalone mode: Removed query_reaction_database from tools (no retrieval)") + + # Define task factories for retrosynthesis + def create_proposal_task(): + return RetrosynthesisTask( + user_prompt=user_prompt_with_rag, + server_urls=available_tools, + builtin_tools=builtin_tools_filtered, + ) + + def create_aggregation_task(candidates_text, subset, step, total_steps): + return RSAAggregationTask( + original_user_prompt=user_prompt, + candidates_text=candidates_text, + step=step, + total_steps=total_steps, + mode=rsa_mode, + server_urls=available_tools, + builtin_tools=builtin_tools_filtered, + ) + + def format_candidates(subset): + """Format retrosynthesis proposals for aggregation""" + candidates_text = "" + for idx, prop in enumerate(subset, 1): + prop_result = prop["result"] + candidates_text += f"\n---- Candidate {idx} ----\n" + candidates_text += f"Reasoning: {prop_result.reasoning_summary}\n" + candidates_text += f"Reactants: {', '.join(prop_result.reactants_smiles_list)}\n" + candidates_text += f"Products: {', '.join(prop_result.products_smiles_list)}\n" + return candidates_text + + # Runner factory for parallel execution + proposal_counter = [0] # Mutable counter for unique agent names + def create_runner(): + """Create independent runner for parallel proposals""" + proposal_counter[0] += 1 + return experiment.create_agent_with_experiment_state( + task=None, + agent_name=f"retrosynth_{node_id}_proposal_{proposal_counter[0]}", + callback=callback_handler if isinstance(callback_handler, CallbackHandler) else None, + ) + + # Run generic RSA loop + output, final_result = await run_rsa_loop( + n=rsa_n, + k=rsa_k, + t=rsa_t, + create_proposal_task=create_proposal_task, + create_aggregation_task=create_aggregation_task, + format_candidates=format_candidates, + runner=runner, + log_progress=log_progress, + clogger=clogger, + log_dir=rsa_log_dir, + output_schema=ReactionOutputSchema, + callback_handler=callback_handler if isinstance(callback_handler, CallbackHandler) else None, + parallel=True, + runner_factory=create_runner, + ) + + await clogger.info(f"RSA mode completed successfully. Logs saved to: {rsa_log_dir}") + + except Exception as e: + # Fallback to standard mode if RSA fails + await clogger.error(f"RSA mode failed: {str(e)}, falling back to standard retrosynthesis") + # Reset task to original + runner.task = retro_task + output = await runner.run(log_progress) + if isinstance(callback_handler, CallbackHandler): + await callback_handler.drain() + else: + # Standard mode (no RSA) + output = await runner.run(log_progress) + if isinstance(callback_handler, CallbackHandler): + await callback_handler.drain() if os.getenv("CHARGE_DISABLE_OUTPUT_VALIDATION", "0") == "1": await clogger.warning( @@ -239,48 +416,60 @@ async def ai_based_retrosynthesis( await context.add_node(node, websocket) await asyncio.sleep(0) - for node, purch in zip(new_nodes, purchasable): - if purch: # Skip purchasable nodes unless explicitly asked for - continue - - # Highlight node because we are looking for templates - await highlight_node(node, websocket, True) - - # Find paths for the leaf nodes - reaction, routes = await run_retro_planner( - config_file, node.smiles, clogger, run_settings - ) - if reaction is None: - await clogger.warning(f"No routes found for {node.smiles}. Skipping...") - continue - node.reaction = reaction - - # Use child nodes and edges of the first route - await generate_nodes_for_molecular_graph( - routes[0].nodes, - context, - websocket, - start_level=level, - include_root_node=False, - root_node_id=node.id, + # Template-based expansion of reactants (optional) + if not os.path.exists(config_file): + await clogger.info( + f"Template-based retrosynthesis config not found at {config_file}. " + "Skipping template expansion. AI-based retrosynthesis completed successfully." ) - - # Attach mapped reaction for immediate reactants -> product. - # This is needed so hover-highlighting works for the first template step - # discovered after an AI-generated step. - if node.reaction is not None: - child_smiles = [ - n.smiles - for nid, n in context.node_ids.items() - if context.parents.get(nid) == node.id - ] - node.reaction.mappedReaction = build_mapped_reaction_dict_or_none( - reactants=child_smiles, - products=[node.smiles], - log_msg="Failed to build rdkitjs mapped reaction for template node_id={node_id} smiles={smiles}", - node_id=node.id, - smiles=node.smiles, - ) - await context.update_node(node, websocket) # Also disables highlight + else: + for node, purch in zip(new_nodes, purchasable): + if purch: # Skip purchasable nodes unless explicitly asked for + continue + + # Highlight node because we are looking for templates + await highlight_node(node, websocket, True) + + # Find paths for the leaf nodes + try: + reaction, routes = await run_retro_planner( + config_file, node.smiles, clogger, run_settings + ) + if reaction is None: + await clogger.warning(f"No routes found for {node.smiles}. Skipping...") + continue + node.reaction = reaction + + # Use child nodes and edges of the first route + await generate_nodes_for_molecular_graph( + routes[0].nodes, + context, + websocket, + start_level=level, + include_root_node=False, + root_node_id=node.id, + ) + except Exception as e: + await clogger.warning( + f"Template-based expansion failed for {node.smiles}: {str(e)}. Continuing..." + ) + + # Attach mapped reaction for immediate reactants -> product. + # This is needed so hover-highlighting works for the first template step + # discovered after an AI-generated step. + if node.reaction is not None: + child_smiles = [ + n.smiles + for nid, n in context.node_ids.items() + if context.parents.get(nid) == node.id + ] + node.reaction.mappedReaction = build_mapped_reaction_dict_or_none( + reactants=child_smiles, + products=[node.smiles], + log_msg="Failed to build rdkitjs mapped reaction for template node_id={node_id} smiles={smiles}", + node_id=node.id, + smiles=node.smiles, + ) + await context.update_node(node, websocket) # Also disables highlight await websocket.send_json({"type": "complete"}) diff --git a/charge_backend/retrosynthesis/prompts/rsa_rag_aggregation.txt b/charge_backend/retrosynthesis/prompts/rsa_rag_aggregation.txt new file mode 100644 index 00000000..3c1b730c --- /dev/null +++ b/charge_backend/retrosynthesis/prompts/rsa_rag_aggregation.txt @@ -0,0 +1,27 @@ +You are given a chemistry task prompt and several candidate retrosynthesis solutions that may include insights from a reaction database. + +Original task: +{original_prompt} + +Candidate solutions (Step {step} of {total_steps}): +{candidates} + +Your task: +- Review all candidate solutions carefully, including any database-informed insights +- Identify portions of the reasoning that are chemically sound +- Look for common patterns or complementary insights across candidates +- Consider database evidence as supporting information, not absolute rules +- Synthesize this information to produce a single, high-quality, chemically correct retrosynthesis + +Requirements: +- Prefer candidates with more plausible chemistry over those with errors +- Use database patterns as evidence, but validate with chemistry principles +- Reject patterns that violate valence, require impossible reagents, or imply implausible bond changes +- Ensure your output is a valid single-step retrosynthesis +- Make sure all SMILES strings are valid +- Verify that the reactants can plausibly produce the target molecule + +Output a single improved retrosynthesis solution with: +- A reasoning summary explaining your aggregated approach +- A list of reactant SMILES strings +- A list of product SMILES strings (the target molecule) diff --git a/charge_backend/retrosynthesis/prompts/rsa_rag_system.txt b/charge_backend/retrosynthesis/prompts/rsa_rag_system.txt new file mode 100644 index 00000000..19875280 --- /dev/null +++ b/charge_backend/retrosynthesis/prompts/rsa_rag_system.txt @@ -0,0 +1,20 @@ +You are an expert chemist specializing in retrosynthesis with access to a reaction database. + +Your task is to provide a retrosynthetic pathway for the target molecule. Database results showing similar reactions are already provided in the user prompt below - do NOT use the query_reaction_database tool as the query has already been performed. + +Reason using the provided database results and retrosynthesis principles: +- Use database results as supporting evidence, but do NOT assume the closest-looking example is correct +- Propose a chemically plausible disconnection based on the database patterns and your chemistry knowledge +- Use functional group interconversion (FGI) only when justified +- Ensure synthons/synthetic equivalents are realistic +- Respect protecting group logic when needed +- Avoid implausible bond disconnections (wrong polarity, impossible leaving groups, severe strain) +- Prefer chemically plausible outcomes over superficial similarity +- Reject patterns that would violate valence, require impossible reagents, or imply implausible bond changes + +Perform only single step retrosynthesis. Make sure the SMILES strings are valid. Use tools to verify the SMILES strings and diagnose any issues that arise. + +Your output should include: +- A reasoning summary explaining your retrosynthetic approach (including insights from the database) +- A list of reactant SMILES strings +- A list of product SMILES strings (the target molecule) diff --git a/charge_backend/retrosynthesis/prompts/rsa_standalone_aggregation.txt b/charge_backend/retrosynthesis/prompts/rsa_standalone_aggregation.txt new file mode 100644 index 00000000..30036195 --- /dev/null +++ b/charge_backend/retrosynthesis/prompts/rsa_standalone_aggregation.txt @@ -0,0 +1,24 @@ +You are given a chemistry task prompt and several candidate retrosynthesis solutions. + +Original task: +{original_prompt} + +Candidate solutions (Step {step} of {total_steps}): +{candidates} + +Your task: +- Review all candidate solutions carefully +- Identify portions of the reasoning that are chemically sound +- Look for common patterns or complementary insights across candidates +- Synthesize this information to produce a single, high-quality, chemically correct retrosynthesis + +Requirements: +- Prefer candidates with more plausible chemistry over those with errors +- Ensure your output is a valid single-step retrosynthesis +- Make sure all SMILES strings are valid +- Verify that the reactants can plausibly produce the target molecule + +Output a single improved retrosynthesis solution with: +- A reasoning summary explaining your aggregated approach +- A list of reactant SMILES strings +- A list of product SMILES strings (the target molecule) diff --git a/charge_backend/retrosynthesis/prompts/rsa_standalone_system.txt b/charge_backend/retrosynthesis/prompts/rsa_standalone_system.txt new file mode 100644 index 00000000..36d4c2f7 --- /dev/null +++ b/charge_backend/retrosynthesis/prompts/rsa_standalone_system.txt @@ -0,0 +1,17 @@ +You are an expert chemist specializing in retrosynthesis. + +Your task is to provide a retrosynthetic pathway for the target molecule. Reason using retrosynthesis principles: +- Propose a chemically plausible disconnection +- Use functional group interconversion (FGI) only when justified +- Ensure synthons/synthetic equivalents are realistic +- Respect protecting group logic when needed +- Avoid implausible bond disconnections (wrong polarity, impossible leaving groups, severe strain) +- Avoid violating valence/charges +- Prefer chemically plausible precursor sets over superficial similarity + +Perform only single step retrosynthesis. Make sure the SMILES strings are valid. Use tools to verify the SMILES strings and diagnose any issues that arise. + +Your output should include: +- A reasoning summary explaining your retrosynthetic approach +- A list of reactant SMILES strings +- A list of product SMILES strings (the target molecule) diff --git a/charge_backend/retrosynthesis/retrosynthesis_task.py b/charge_backend/retrosynthesis/retrosynthesis_task.py index fa2b5a9a..e0fd9b52 100644 --- a/charge_backend/retrosynthesis/retrosynthesis_task.py +++ b/charge_backend/retrosynthesis/retrosynthesis_task.py @@ -3,6 +3,7 @@ from pydantic import BaseModel, field_validator from flask_tools.chemistry.smarts_reactions_utils import verify_reaction_SMARTS from flask_tools.chemistry.smiles_utils import verify_smiles +from pathlib import Path class ReactionOutputSchema(BaseModel): @@ -180,3 +181,77 @@ def __init__( + f"\n{self.user_prompt}" + f"\n{TEMPLATE_FREE_SYSTEM_PROMPT}" ) + + +class RSAAggregationTask(Task): + """ + Task for RSA (Recursive Self-Aggregation) that aggregates multiple candidate + retrosynthesis solutions into a single improved solution. + """ + + def __init__( + self, + original_user_prompt: str, + candidates_text: str, + step: int, + total_steps: int, + mode: str = "standalone", + system_prompt: Optional[str] = None, + **kwargs, + ): + """ + Initialize RSA Aggregation Task. + + Args: + original_user_prompt: The original retrosynthesis prompt + candidates_text: Formatted text containing candidate solutions + step: Current aggregation step (2..T) + total_steps: Total number of RSA steps (T) + mode: "standalone" or "rag" + system_prompt: Optional override for system prompt + **kwargs: Additional arguments passed to Task + """ + # Determine prompt file paths + prompts_dir = Path(__file__).parent / "prompts" + system_file = prompts_dir / f"rsa_{mode}_system.txt" + aggregation_file = prompts_dir / f"rsa_{mode}_aggregation.txt" + + # Load system prompt + if system_prompt is None: + if system_file.exists(): + system_prompt = system_file.read_text() + else: + raise FileNotFoundError( + f"RSA system prompt not found: {system_file}" + ) + + # Load and format aggregation template + if aggregation_file.exists(): + aggregation_template = aggregation_file.read_text() + user_prompt = aggregation_template.format( + original_prompt=original_user_prompt, + candidates=candidates_text, + step=step, + total_steps=total_steps, + ) + else: + raise FileNotFoundError( + f"RSA aggregation template not found: {aggregation_file}" + ) + + super().__init__( + system_prompt=system_prompt, + user_prompt=user_prompt, + **kwargs, + ) + self.system_prompt = system_prompt + self.user_prompt = user_prompt + self.original_user_prompt = original_user_prompt + self.step = step + self.total_steps = total_steps + self.mode = mode + self.set_structured_output_schema(TemplateFreeReactionOutputSchema) + + print( + f"RSAAggregationTask initialized (mode={mode}, step={step}/{total_steps})" + ) diff --git a/charge_backend/retrosynthesis/template.py b/charge_backend/retrosynthesis/template.py index 76d6443f..6b1cee33 100644 --- a/charge_backend/retrosynthesis/template.py +++ b/charge_backend/retrosynthesis/template.py @@ -1,3 +1,4 @@ +import os import asyncio from fastapi import WebSocket from charge_backend.retrosynthesis import aizynth_tools as azf @@ -205,6 +206,14 @@ async def run_retro_planner( :return: A 2-tuple of (Reaction object, list of routes) if routes found, or ``(None, [])`` if nothing was discovered. """ + # Check if config file exists + if not os.path.exists(config_file): + await clogger.info( + f"Template-based retrosynthesis config not found at {config_file}. " + "Template search unavailable." + ) + return None, [] + await clogger.info(f"Running RetroPlanner for SMILES: {smiles}") report_init = False if azf.RetroPlanner.finder is None: diff --git a/charge_backend/rsa_algorithm.py b/charge_backend/rsa_algorithm.py new file mode 100644 index 00000000..16454ba0 --- /dev/null +++ b/charge_backend/rsa_algorithm.py @@ -0,0 +1,335 @@ +""" +Generic Recursive Self-Aggregation (RSA) Algorithm + +This module provides the core N-K-T RSA loop that can be used by any task type +(retrosynthesis, LMO, etc.). It uses the existing ChARGe Agent/Task framework +without introducing new orchestration layers. + +RSA Algorithm: + Stage 1: Generate N diverse proposals + Stages 2-T: Recursively aggregate K-subset proposals + Output: Single best proposal from final stage +""" + +import random +import json +import os +import asyncio +from typing import Any, Callable, Optional +from pathlib import Path + + +async def run_rsa_loop( + n: int, + k: int, + t: int, + create_proposal_task: Callable[[], Any], + create_aggregation_task: Callable[[str, list[dict], int, int], Any], + format_candidates: Callable[[list[dict]], str], + runner: Any, + log_progress: Callable, + clogger: Any, + log_dir: str, + output_schema: Any, + callback_handler: Optional[Any] = None, + parallel: bool = True, + runner_factory: Optional[Callable[[], Any]] = None, +) -> tuple[str, Any]: + """ + Execute the generic N-K-T RSA algorithm. + + Args: + n: Number of initial proposals to generate + k: Size of subsets for aggregation (K <= N) + t: Total number of stages (including initial proposals) + create_proposal_task: Factory function that returns a Task for proposals + create_aggregation_task: Factory function that returns aggregation Task + Takes (candidates_text, subset, step, total_steps) + format_candidates: Function to format proposals into text for aggregation + Takes list of proposals, returns formatted string + runner: ChARGe agent runner with .task and .run() interface + log_progress: Callback for reasoning progress + clogger: Callback logger for UI messages + log_dir: Directory to save execution logs + output_schema: Pydantic schema for validating outputs + callback_handler: Optional callback handler to drain after each task + parallel: If True, generate initial proposals in parallel (default: True) + runner_factory: Factory to create independent runner instances for parallel execution + If None and parallel=True, uses the provided runner (not truly parallel) + + Returns: + tuple: (final_output_json, final_result_object) + + Raises: + ValueError: If all proposals fail or invalid parameters + """ + + # Validate parameters + if n < 1 or k < 1 or t < 1: + raise ValueError(f"Invalid RSA parameters: N={n}, K={k}, T={t} (all must be >= 1)") + if k > n: + await clogger.warning(f"K ({k}) > N ({n}), adjusting K to N") + k = n + + # Helper function to run a single proposal + async def run_single_proposal(proposal_index: int, proposal_runner: Any): + """Run a single proposal and return result or None if failed""" + try: + await clogger.info(f"Generating proposal {proposal_index+1}/{n}") + + # Create proposal task + proposal_task = create_proposal_task() + proposal_runner.task = proposal_task + + # Disable validation if requested + if os.getenv("CHARGE_DISABLE_OUTPUT_VALIDATION", "0") == "1": + proposal_task.structured_output_schema = None + + # Save proposal prompt + proposer_log = { + "proposal_index": proposal_index + 1, + "system_prompt": proposal_task.get_system_prompt(), + "user_prompt": proposal_task.get_user_prompt(), + } + with open(f"{log_dir}/proposer_{proposal_index+1:02d}_prompt.json", "w") as f: + json.dump(proposer_log, f, indent=2) + + # Run proposal + proposal_output = await proposal_runner.run(log_progress) + if callback_handler: + await callback_handler.drain() + + # Validate output + proposal_result = output_schema.model_validate_json(proposal_output) + + # Check if proposal contains actual chemical information + if not hasattr(proposal_result, 'reactants_smiles_list') or len(proposal_result.reactants_smiles_list) == 0: + await clogger.warning(f"Proposal {proposal_index+1} has empty reactants (model refused or failed), skipping") + return None + + # Save proposal output + proposer_output_log = { + "proposal_index": proposal_index + 1, + "result": proposal_result.model_dump(), + "full_output": json.loads(proposal_output) + } + with open(f"{log_dir}/proposer_{proposal_index+1:02d}_output.json", "w") as f: + json.dump(proposer_output_log, f, indent=2) + + await clogger.info(f"Proposal {proposal_index+1} completed successfully") + + return { + "output": proposal_output, + "result": proposal_result, + "index": proposal_index + } + + except Exception as e: + await clogger.warning(f"Proposal {proposal_index+1} failed: {str(e)}") + return None + + # Stage 1: Generate N initial proposals + await clogger.info(f"RSA Step 1/{t}: Generating {n} initial proposals" + + (" (parallel mode)" if parallel else " (sequential mode)")) + + if parallel and runner_factory: + # Parallel mode: generate all proposals concurrently + proposal_tasks = [] + for i in range(n): + # Create independent runner for each proposal + proposal_runner = runner_factory() + task = run_single_proposal(i, proposal_runner) + proposal_tasks.append(task) + + # Run all proposals in parallel + proposal_results = await asyncio.gather(*proposal_tasks, return_exceptions=True) + + # Filter valid proposals and handle exceptions + proposals = [] + for i, result in enumerate(proposal_results): + if isinstance(result, Exception): + await clogger.warning(f"Proposal {i+1} failed with exception: {str(result)}") + elif result is not None: + proposals.append(result) + else: + # Sequential mode: generate proposals one by one + if parallel and not runner_factory: + await clogger.warning("Parallel mode requested but no runner_factory provided, falling back to sequential") + + proposals = [] + for i in range(n): + result = await run_single_proposal(i, runner) + if result is not None: + proposals.append(result) + + if not proposals: + raise ValueError("All RSA proposals failed") + + await clogger.info(f"Generated {len(proposals)} valid proposals") + + # Helper function to run a single aggregation + async def run_single_aggregation(agg_index: int, step: int, current_proposals: list, agg_runner: Any): + """Run a single aggregation and return result or None if failed""" + try: + await clogger.info(f"Aggregation {agg_index+1}/{num_aggregations} (Step {step})") + + # Adjust K if needed + current_k = k + if len(current_proposals) < k: + current_k = len(current_proposals) + + # Select K random proposals + if len(current_proposals) <= current_k: + subset = current_proposals + else: + subset = random.sample(current_proposals, current_k) + + # Format candidates using task-specific formatter + candidates_text = format_candidates(subset) + subset_indices = [prop["index"] + 1 for prop in subset] + + # Create aggregation task + agg_task = create_aggregation_task( + candidates_text, + subset, + step, + t + ) + agg_runner.task = agg_task + + if os.getenv("CHARGE_DISABLE_OUTPUT_VALIDATION", "0") == "1": + agg_task.structured_output_schema = None + + # Save aggregation prompt + aggregator_log = { + "step": step, + "aggregation_index": agg_index + 1, + "k_subset_indices": subset_indices, + "system_prompt": agg_task.get_system_prompt(), + "user_prompt": agg_task.get_user_prompt(), + "candidates_text": candidates_text, + } + with open(f"{log_dir}/aggregator_step{step}_{agg_index+1:02d}_prompt.json", "w") as f: + json.dump(aggregator_log, f, indent=2) + + # Run aggregation + agg_output = await agg_runner.run(log_progress) + if callback_handler: + await callback_handler.drain() + + # Validate output + agg_result = output_schema.model_validate_json(agg_output) + + # Check if aggregation contains actual chemical information + if not hasattr(agg_result, 'reactants_smiles_list') or len(agg_result.reactants_smiles_list) == 0: + await clogger.warning(f"Aggregation {agg_index+1} has empty reactants (model refused or failed), skipping") + return None + + # Save aggregation output + aggregator_output_log = { + "step": step, + "aggregation_index": agg_index + 1, + "k_subset_indices": subset_indices, + "result": agg_result.model_dump(), + "full_output": json.loads(agg_output) + } + with open(f"{log_dir}/aggregator_step{step}_{agg_index+1:02d}_output.json", "w") as f: + json.dump(aggregator_output_log, f, indent=2) + + await clogger.info(f"Aggregation {agg_index+1} completed successfully") + + return { + "output": agg_output, + "result": agg_result, + "index": agg_index, + "step": step + } + + except Exception as e: + await clogger.warning(f"Aggregation {agg_index+1} failed: {str(e)}") + return None + + # Stages 2-T: Recursive aggregation + # BARRIER: Wait for all Stage 1 proposals to complete before starting Stage 2 + current_proposals = proposals + await clogger.info(f"Stage 1 complete. Generated {len(proposals)} valid proposals.") + + for step in range(2, t + 1): + await clogger.info( + f"RSA Step {step}/{t}: Aggregating {len(current_proposals)} proposals into {k}-subsets" + + (" (parallel mode)" if parallel else " (sequential mode)") + ) + + # Adjust K if needed + current_k = k + if len(current_proposals) < k: + await clogger.warning( + f"Not enough proposals ({len(current_proposals)}) for K={k}, using all available" + ) + current_k = len(current_proposals) + + # Generate aggregations + num_aggregations = max(n, len(current_proposals)) + + if parallel and runner_factory: + # Parallel mode: run all aggregations in this stage concurrently + agg_tasks = [] + for i in range(num_aggregations): + # Create independent runner for each aggregation + agg_runner = runner_factory() + task = run_single_aggregation(i, step, current_proposals, agg_runner) + agg_tasks.append(task) + + # Run all aggregations in parallel and wait for all to complete + # BARRIER: asyncio.gather waits for all aggregations in this stage + agg_results = await asyncio.gather(*agg_tasks, return_exceptions=True) + + # Filter valid aggregations and handle exceptions + next_proposals = [] + for i, result in enumerate(agg_results): + if isinstance(result, Exception): + await clogger.warning(f"Aggregation {i+1} failed with exception: {str(result)}") + elif result is not None: + next_proposals.append(result) + + else: + # Sequential mode: run aggregations one by one + if parallel and not runner_factory: + await clogger.warning("Parallel mode requested but no runner_factory provided, falling back to sequential") + + next_proposals = [] + for i in range(num_aggregations): + result = await run_single_aggregation(i, step, current_proposals, runner) + if result is not None: + next_proposals.append(result) + + if not next_proposals: + await clogger.warning(f"No successful aggregations in step {step}, using previous proposals") + break + + # BARRIER: All aggregations in current stage complete before moving to next stage + current_proposals = next_proposals + await clogger.info(f"Stage {step} complete. Generated {len(current_proposals)} valid aggregations.") + + # Select final proposal (first one from final stage) + if not current_proposals: + raise ValueError("RSA failed to produce any valid proposals") + + final_proposal = current_proposals[0] + final_output = final_proposal["output"] + final_result = final_proposal["result"] + + # Save final output + final_log = { + "final_step": step if step <= t else t, + "n_proposals": n, + "k_subset_size": k, + "t_stages": t, + "final_result": final_result.model_dump(), + } + with open(f"{log_dir}/FINAL_OUTPUT.json", "w") as f: + json.dump(final_log, f, indent=2) + + await clogger.info(f"RSA completed! Final output saved to {log_dir}") + + return final_output, final_result diff --git a/externals/lc_conductor b/externals/lc_conductor index 380c2cfd..afcf0c80 160000 --- a/externals/lc_conductor +++ b/externals/lc_conductor @@ -1 +1 @@ -Subproject commit 380c2cfd071b2ebb268b7a36e38b9ab1929b992b +Subproject commit afcf0c80ff356a1fd4037cfb6d8c001b617348f9 diff --git a/flask-app/package-lock.json b/flask-app/package-lock.json index d89bf5c3..9673cb14 100644 --- a/flask-app/package-lock.json +++ b/flask-app/package-lock.json @@ -372,6 +372,7 @@ "cpu": [ "ppc64" ], + "dev": true, "license": "MIT", "optional": true, "os": [ @@ -388,6 +389,7 @@ "cpu": [ "arm" ], + "dev": true, "license": "MIT", "optional": true, "os": [ @@ -404,6 +406,7 @@ "cpu": [ "arm64" ], + "dev": true, "license": "MIT", "optional": true, "os": [ @@ -420,6 +423,7 @@ "cpu": [ "x64" ], + "dev": true, "license": "MIT", "optional": true, "os": [ @@ -436,6 +440,7 @@ "cpu": [ "arm64" ], + "dev": true, "license": "MIT", "optional": true, "os": [ @@ -452,6 +457,7 @@ "cpu": [ "x64" ], + "dev": true, "license": "MIT", "optional": true, "os": [ @@ -468,6 +474,7 @@ "cpu": [ "arm64" ], + "dev": true, "license": "MIT", "optional": true, "os": [ @@ -484,6 +491,7 @@ "cpu": [ "x64" ], + "dev": true, "license": "MIT", "optional": true, "os": [ @@ -500,6 +508,7 @@ "cpu": [ "arm" ], + "dev": true, "license": "MIT", "optional": true, "os": [ @@ -516,6 +525,7 @@ "cpu": [ "arm64" ], + "dev": true, "license": "MIT", "optional": true, "os": [ @@ -532,6 +542,7 @@ "cpu": [ "ia32" ], + "dev": true, "license": "MIT", "optional": true, "os": [ @@ -548,6 +559,7 @@ "cpu": [ "loong64" ], + "dev": true, "license": "MIT", "optional": true, "os": [ @@ -564,6 +576,7 @@ "cpu": [ "mips64el" ], + "dev": true, "license": "MIT", "optional": true, "os": [ @@ -580,6 +593,7 @@ "cpu": [ "ppc64" ], + "dev": true, "license": "MIT", "optional": true, "os": [ @@ -596,6 +610,7 @@ "cpu": [ "riscv64" ], + "dev": true, "license": "MIT", "optional": true, "os": [ @@ -612,6 +627,7 @@ "cpu": [ "s390x" ], + "dev": true, "license": "MIT", "optional": true, "os": [ @@ -628,6 +644,7 @@ "cpu": [ "x64" ], + "dev": true, "license": "MIT", "optional": true, "os": [ @@ -644,6 +661,7 @@ "cpu": [ "arm64" ], + "dev": true, "license": "MIT", "optional": true, "os": [ @@ -660,6 +678,7 @@ "cpu": [ "x64" ], + "dev": true, "license": "MIT", "optional": true, "os": [ @@ -676,6 +695,7 @@ "cpu": [ "arm64" ], + "dev": true, "license": "MIT", "optional": true, "os": [ @@ -692,6 +712,7 @@ "cpu": [ "x64" ], + "dev": true, "license": "MIT", "optional": true, "os": [ @@ -708,6 +729,7 @@ "cpu": [ "arm64" ], + "dev": true, "license": "MIT", "optional": true, "os": [ @@ -724,6 +746,7 @@ "cpu": [ "x64" ], + "dev": true, "license": "MIT", "optional": true, "os": [ @@ -740,6 +763,7 @@ "cpu": [ "arm64" ], + "dev": true, "license": "MIT", "optional": true, "os": [ @@ -756,6 +780,7 @@ "cpu": [ "ia32" ], + "dev": true, "license": "MIT", "optional": true, "os": [ @@ -772,6 +797,7 @@ "cpu": [ "x64" ], + "dev": true, "license": "MIT", "optional": true, "os": [ @@ -1149,6 +1175,7 @@ "cpu": [ "arm" ], + "dev": true, "license": "MIT", "optional": true, "os": [ @@ -1162,6 +1189,7 @@ "cpu": [ "arm64" ], + "dev": true, "license": "MIT", "optional": true, "os": [ @@ -1175,6 +1203,7 @@ "cpu": [ "arm64" ], + "dev": true, "license": "MIT", "optional": true, "os": [ @@ -1188,6 +1217,7 @@ "cpu": [ "x64" ], + "dev": true, "license": "MIT", "optional": true, "os": [ @@ -1201,6 +1231,7 @@ "cpu": [ "arm64" ], + "dev": true, "license": "MIT", "optional": true, "os": [ @@ -1214,6 +1245,7 @@ "cpu": [ "x64" ], + "dev": true, "license": "MIT", "optional": true, "os": [ @@ -1227,6 +1259,7 @@ "cpu": [ "arm" ], + "dev": true, "license": "MIT", "optional": true, "os": [ @@ -1240,6 +1273,7 @@ "cpu": [ "arm" ], + "dev": true, "license": "MIT", "optional": true, "os": [ @@ -1253,6 +1287,7 @@ "cpu": [ "arm64" ], + "dev": true, "license": "MIT", "optional": true, "os": [ @@ -1266,6 +1301,7 @@ "cpu": [ "arm64" ], + "dev": true, "license": "MIT", "optional": true, "os": [ @@ -1279,6 +1315,7 @@ "cpu": [ "loong64" ], + "dev": true, "license": "MIT", "optional": true, "os": [ @@ -1292,6 +1329,7 @@ "cpu": [ "loong64" ], + "dev": true, "license": "MIT", "optional": true, "os": [ @@ -1305,6 +1343,7 @@ "cpu": [ "ppc64" ], + "dev": true, "license": "MIT", "optional": true, "os": [ @@ -1318,6 +1357,7 @@ "cpu": [ "ppc64" ], + "dev": true, "license": "MIT", "optional": true, "os": [ @@ -1331,6 +1371,7 @@ "cpu": [ "riscv64" ], + "dev": true, "license": "MIT", "optional": true, "os": [ @@ -1344,6 +1385,7 @@ "cpu": [ "riscv64" ], + "dev": true, "license": "MIT", "optional": true, "os": [ @@ -1357,6 +1399,7 @@ "cpu": [ "s390x" ], + "dev": true, "license": "MIT", "optional": true, "os": [ @@ -1370,6 +1413,7 @@ "cpu": [ "x64" ], + "dev": true, "license": "MIT", "optional": true, "os": [ @@ -1383,6 +1427,7 @@ "cpu": [ "x64" ], + "dev": true, "license": "MIT", "optional": true, "os": [ @@ -1396,6 +1441,7 @@ "cpu": [ "x64" ], + "dev": true, "license": "MIT", "optional": true, "os": [ @@ -1409,6 +1455,7 @@ "cpu": [ "arm64" ], + "dev": true, "license": "MIT", "optional": true, "os": [ @@ -1422,6 +1469,7 @@ "cpu": [ "arm64" ], + "dev": true, "license": "MIT", "optional": true, "os": [ @@ -1435,6 +1483,7 @@ "cpu": [ "ia32" ], + "dev": true, "license": "MIT", "optional": true, "os": [ @@ -1448,6 +1497,7 @@ "cpu": [ "x64" ], + "dev": true, "license": "MIT", "optional": true, "os": [ @@ -1461,6 +1511,7 @@ "cpu": [ "x64" ], + "dev": true, "license": "MIT", "optional": true, "os": [ @@ -2010,12 +2061,14 @@ "version": "15.7.15", "resolved": "https://registry.npmjs.org/@types/prop-types/-/prop-types-15.7.15.tgz", "integrity": "sha512-F6bEyamV9jKGAFBEmlQnesRPGOQqS2+Uwi0Em15xenOxHaf2hv6L8YCVn3rPdPJOiJfPiCnLIRyvwVaqMY3MIw==", + "dev": true, "license": "MIT" }, "node_modules/@types/react": { "version": "18.3.28", "resolved": "https://registry.npmjs.org/@types/react/-/react-18.3.28.tgz", "integrity": "sha512-z9VXpC7MWrhfWipitjNdgCauoMLRdIILQsAEV+ZesIzBq/oUlxk0m3ApZuMFCXdnS4U7KrI+l3WRUEGQ8K1QKw==", + "dev": true, "license": "MIT", "dependencies": { "@types/prop-types": "*", @@ -2026,7 +2079,7 @@ "version": "18.3.7", "resolved": "https://registry.npmjs.org/@types/react-dom/-/react-dom-18.3.7.tgz", "integrity": "sha512-MEe3UeoENYVFXzoXEWsvcpg6ZvlrFNlOQ7EOsvhI3CfAXwzPfO8Qwuxd40nepsYKqyyVQnTdEfv68q91yLcKrQ==", - "devOptional": true, + "dev": true, "license": "MIT", "peerDependencies": { "@types/react": "^18.0.0" @@ -2899,6 +2952,7 @@ "version": "0.27.3", "resolved": "https://registry.npmjs.org/esbuild/-/esbuild-0.27.3.tgz", "integrity": "sha512-8VwMnyGCONIs6cWue2IdpHxHnAjzxnw2Zr7MkVxB2vjmQ2ivqGFb4LEG3SMnv0Gb2F/G/2yA8zUaiL1gywDCCg==", + "dev": true, "hasInstallScript": true, "license": "MIT", "bin": { @@ -3016,6 +3070,7 @@ "version": "6.5.0", "resolved": "https://registry.npmjs.org/fdir/-/fdir-6.5.0.tgz", "integrity": "sha512-tIbYtZbucOs0BRGqPJkshJUYdL+SDH7dVM8gjy+ERp3WAUjLEFJE+02kanyHtwjWOnwrKYBiwAmM0p4kLJAnXg==", + "dev": true, "license": "MIT", "engines": { "node": ">=12.0.0" @@ -3070,6 +3125,7 @@ "version": "2.3.3", "resolved": "https://registry.npmjs.org/fsevents/-/fsevents-2.3.3.tgz", "integrity": "sha512-5xoDfX+fL7faATnagmWPpbFtwh/R77WmMMqqHGS65C3vvB0YHrgF+B1YmZ3441tMj5n63k0212XNoJwzlhffQw==", + "dev": true, "hasInstallScript": true, "license": "MIT", "optional": true, @@ -4769,6 +4825,7 @@ "version": "3.3.11", "resolved": "https://registry.npmjs.org/nanoid/-/nanoid-3.3.11.tgz", "integrity": "sha512-N8SpfPUnUp1bK+PMYW8qSWdl9U+wwNWI4QKxOYDy9JAro3WMX7p2OeVRF9v+347pnakNevPmiHhNmZ2HbFA76w==", + "dev": true, "funding": [ { "type": "github", @@ -4894,6 +4951,7 @@ "version": "8.5.6", "resolved": "https://registry.npmjs.org/postcss/-/postcss-8.5.6.tgz", "integrity": "sha512-3Ybi1tAuwAP9s0r1UQ2J4n5Y0G05bJkpUIO0/bI9MhwmD70S5aTWbXGBwxHrelT+XM1k6dM0pk+SwNkpTRN7Pg==", + "dev": true, "funding": [ { "type": "opencollective", @@ -5325,6 +5383,7 @@ "version": "4.57.1", "resolved": "https://registry.npmjs.org/rollup/-/rollup-4.57.1.tgz", "integrity": "sha512-oQL6lgK3e2QZeQ7gcgIkS2YZPg5slw37hYufJ3edKlfQSGGm8ICoxswK15ntSzF/a8+h7ekRy7k7oWc3BQ7y8A==", + "dev": true, "license": "MIT", "dependencies": { "@types/estree": "1.0.8" @@ -5548,6 +5607,7 @@ "version": "0.2.15", "resolved": "https://registry.npmjs.org/tinyglobby/-/tinyglobby-0.2.15.tgz", "integrity": "sha512-j2Zq4NyQYG5XMST4cbs02Ak8iJUdxRM0XI5QyxXuZOzKOINmWurp3smXu3y5wDcJrptwpSjgXHzIQxR0omXljQ==", + "dev": true, "license": "MIT", "dependencies": { "fdir": "^6.5.0", @@ -5811,6 +5871,7 @@ "version": "7.3.1", "resolved": "https://registry.npmjs.org/vite/-/vite-7.3.1.tgz", "integrity": "sha512-w+N7Hifpc3gRjZ63vYBXA56dvvRlNWRczTdmCBBa+CotUzAPf5b7YMdMR/8CQoeYE5LX3W4wj6RYTgonm1b9DA==", + "dev": true, "license": "MIT", "dependencies": { "esbuild": "^0.27.0", diff --git a/flask-app/src/App.tsx b/flask-app/src/App.tsx index 338f9c79..f65f19e0 100644 --- a/flask-app/src/App.tsx +++ b/flask-app/src/App.tsx @@ -130,6 +130,16 @@ const ChemistryTool: React.FC = () => { const [editedPrompt, setEditedPrompt] = useState(''); const [debugModalMinimized, setDebugModalMinimized] = useState(false); + // Retrosynthesis approach + const [useAiBased, setUseAiBased] = useState(true); + + // RSA settings + const [useRsa, setUseRsa] = useState(false); + const [rsaMode, setRsaMode] = useState<'standalone' | 'rag'>('standalone'); + const [rsaN, setRsaN] = useState(8); + const [rsaK, setRsaK] = useState(4); + const [rsaT, setRsaT] = useState(3); + // Function to refresh tools list from backend const refreshToolsList = useCallback(() => { if (wsRef.current && wsRef.current.readyState === WebSocket.OPEN) { @@ -486,6 +496,12 @@ const ChemistryTool: React.FC = () => { runSettings: { promptDebugging: debugMode, moleculeName: orchestratorSettings.moleculeName || 'brand', + useAiBased, + useRsa, + rsaMode, + rsaN, + rsaK, + rsaT, }, customization, }; @@ -1015,13 +1031,18 @@ const ChemistryTool: React.FC = () => { runSettings: { promptDebugging: debugMode, moleculeName: orchestratorSettings.moleculeName || 'brand', + useRsa, + rsaMode, + rsaN, + rsaK, + rsaT, }, ...data, }; wsRef.current.send(JSON.stringify(msg)); setContextMenu({ node: null, isReaction: false, x: 0, y: 0 }); }, - [debugMode, orchestratorSettings] + [debugMode, orchestratorSettings, useAiBased, useRsa, rsaMode, rsaN, rsaK, rsaT] ); const handleReactionCardClick = useCallback( @@ -1095,6 +1116,11 @@ const ChemistryTool: React.FC = () => { runSettings: { promptDebugging: debugMode, moleculeName: orchestratorSettings.moleculeName || 'brand', + useRsa, + rsaMode, + rsaN, + rsaK, + rsaT, }, }) ); @@ -1102,8 +1128,8 @@ const ChemistryTool: React.FC = () => { } setIsComputing(true); }, - [selectedReactionNode?.id, debugMode, orchestratorSettings] - ); // Only depend on the ID + [selectedReactionNode?.id, debugMode, orchestratorSettings, useRsa, rsaMode, rsaN, rsaK, rsaT] + ); const stableAlternatives = useMemo(() => { return selectedReactionNode?.reaction?.alternatives || []; @@ -1145,6 +1171,12 @@ const ChemistryTool: React.FC = () => { runSettings: { promptDebugging: debugMode, moleculeName: orchestratorSettings.moleculeName || 'brand', + useAiBased, + useRsa, + rsaMode, + rsaN, + rsaK, + rsaT, }, ...propertyDetails, }; @@ -1515,6 +1547,105 @@ const ChemistryTool: React.FC = () => { AI Debug Mode + + {/* Retrosynthesis Settings */} + {problemType === 'retrosynthesis' && ( +
+
+ +
+ + {useAiBased && ( +
+ +
+ )} + + {useAiBased && useRsa && ( +
+
+ + +
+ +
+
+ + setRsaN(parseInt(e.target.value) || 8)} + disabled={isComputing} + min="1" + max="20" + className="form-input text-sm w-full" + title="Number of initial proposals to generate" + /> +
+
+ + setRsaK(parseInt(e.target.value) || 4)} + disabled={isComputing} + min="1" + max={rsaN} + className="form-input text-sm w-full" + title="Subset size for aggregation" + /> +
+
+ + setRsaT(parseInt(e.target.value) || 3)} + disabled={isComputing} + min="1" + max="10" + className="form-input text-sm w-full" + title="Total number of aggregation steps" + /> +
+
+ +
+ Estimated runtime: ~{rsaN * rsaT} inferences +
+
+ )} +
+ )} diff --git a/flask-app/src/types.ts b/flask-app/src/types.ts index 2dfe2316..2b9c8d71 100644 --- a/flask-app/src/types.ts +++ b/flask-app/src/types.ts @@ -81,6 +81,12 @@ export type MoleculeNameFormat = 'brand' | 'iupac' | 'formula' | 'smiles'; export interface FlaskRunSettings { moleculeName: MoleculeNameFormat; promptDebugging: boolean; + useAiBased?: boolean; + useRsa?: boolean; + rsaMode?: 'standalone' | 'rag'; + rsaN?: number; + rsaK?: number; + rsaT?: number; } import type { OrchestratorSettings, SidebarMessage, SidebarState } from 'lc-conductor'; diff --git a/flask-tools b/flask-tools new file mode 160000 index 00000000..28bad724 --- /dev/null +++ b/flask-tools @@ -0,0 +1 @@ +Subproject commit 28bad7248d7ec2944cd52b9fa095c66cf87d57c7