From 6ec6f568393ad1e37c7c18d8604fa5b72ba6dfc5 Mon Sep 17 00:00:00 2001 From: Shashank Kushwaha Date: Wed, 1 Apr 2026 21:59:17 -0700 Subject: [PATCH 01/24] feat: Add RSA prompts and RSAAggregationTask class --- .../prompts/rsa_rag_aggregation.txt | 27 +++++++ .../retrosynthesis/prompts/rsa_rag_system.txt | 20 +++++ .../prompts/rsa_standalone_aggregation.txt | 24 ++++++ .../prompts/rsa_standalone_system.txt | 17 +++++ .../retrosynthesis/retrosynthesis_task.py | 75 +++++++++++++++++++ 5 files changed, 163 insertions(+) create mode 100644 charge_backend/retrosynthesis/prompts/rsa_rag_aggregation.txt create mode 100644 charge_backend/retrosynthesis/prompts/rsa_rag_system.txt create mode 100644 charge_backend/retrosynthesis/prompts/rsa_standalone_aggregation.txt create mode 100644 charge_backend/retrosynthesis/prompts/rsa_standalone_system.txt 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..8dc78b7a --- /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. First, query the reaction database using the query_reaction_database tool to find known reactions that produce the target molecule. + +After querying the database, reason using 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})" + ) From ff57eff9e0360636a526eaa7638ce14d6fe0f69b Mon Sep 17 00:00:00 2001 From: Shashank Kushwaha Date: Wed, 1 Apr 2026 22:02:17 -0700 Subject: [PATCH 02/24] feat: Add RSA orchestration logic to ai_based_retrosynthesis --- charge_backend/retrosynthesis/ai.py | 159 +++++++++++++++++++++++++++- 1 file changed, 155 insertions(+), 4 deletions(-) diff --git a/charge_backend/retrosynthesis/ai.py b/charge_backend/retrosynthesis/ai.py index 58dabe0f..abf069a7 100644 --- a/charge_backend/retrosynthesis/ai.py +++ b/charge_backend/retrosynthesis/ai.py @@ -1,5 +1,7 @@ import os import asyncio +import random +from pathlib import Path from fastapi import WebSocket from lc_conductor.callback_logger import CallbackLogger from typing import Any, Callable, Optional, Union @@ -28,6 +30,7 @@ from retrosynthesis.retrosynthesis_task import ( TemplateFreeRetrosynthesisTask as RetrosynthesisTask, TemplateFreeReactionOutputSchema as ReactionOutputSchema, + RSAAggregationTask, ) from charge.experiments.experiment import Experiment @@ -130,13 +133,161 @@ 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 + 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}" + ) + + # Step 1: Generate N initial proposals + await clogger.info(f"RSA Step 1/{rsa_t}: Generating {rsa_n} initial proposals") + proposals = [] + for i in range(rsa_n): + await clogger.info(f"Generating proposal {i+1}/{rsa_n}") + try: + # Create a fresh task for each proposal + proposal_task = RetrosynthesisTask( + user_prompt=user_prompt, + server_urls=available_tools, + builtin_tools=builtin_tools or [], + ) + runner.task = proposal_task + + if os.getenv("CHARGE_DISABLE_OUTPUT_VALIDATION", "0") == "1": + proposal_task.structured_output_schema = None + + # Run proposal + proposal_output = await runner.run(log_progress) + if isinstance(callback_handler, CallbackHandler): + await callback_handler.drain() + + # Validate and store + proposal_result = ReactionOutputSchema.model_validate_json(proposal_output) + proposals.append({ + "output": proposal_output, + "result": proposal_result, + "index": i + }) + await clogger.info(f"Proposal {i+1} completed successfully") + except Exception as e: + await clogger.warning(f"Proposal {i+1} failed: {str(e)}") + continue + + if not proposals: + raise ValueError("All RSA proposals failed, falling back to standard mode") + + await clogger.info(f"Generated {len(proposals)} valid proposals") + + # Steps 2..T: Recursive aggregation + current_proposals = proposals + for step in range(2, rsa_t + 1): + await clogger.info( + f"RSA Step {step}/{rsa_t}: Aggregating {len(current_proposals)} proposals into {rsa_k}-subsets" + ) + + if len(current_proposals) < rsa_k: + await clogger.warning( + f"Not enough proposals ({len(current_proposals)}) for K={rsa_k}, using all available" + ) + rsa_k = len(current_proposals) + + # Generate new proposals by aggregating K-subsets + next_proposals = [] + num_aggregations = max(rsa_n, len(current_proposals)) + + for i in range(num_aggregations): + await clogger.info(f"Aggregation {i+1}/{num_aggregations}") + try: + # Select K random proposals + if len(current_proposals) <= rsa_k: + subset = current_proposals + else: + subset = random.sample(current_proposals, rsa_k) + + # Format candidates text + 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" + + # Create aggregation task + agg_task = RSAAggregationTask( + original_user_prompt=user_prompt, + candidates_text=candidates_text, + step=step, + total_steps=rsa_t, + mode=rsa_mode, + server_urls=available_tools, + builtin_tools=builtin_tools or [], + ) + runner.task = agg_task + + if os.getenv("CHARGE_DISABLE_OUTPUT_VALIDATION", "0") == "1": + agg_task.structured_output_schema = None + + # Run aggregation + agg_output = await runner.run(log_progress) + if isinstance(callback_handler, CallbackHandler): + await callback_handler.drain() + + # Validate and store + agg_result = ReactionOutputSchema.model_validate_json(agg_output) + next_proposals.append({ + "output": agg_output, + "result": agg_result, + "index": i, + "step": step + }) + await clogger.info(f"Aggregation {i+1} completed successfully") + except Exception as e: + await clogger.warning(f"Aggregation {i+1} failed: {str(e)}") + continue + + if not next_proposals: + await clogger.warning( + f"All aggregations at step {step} failed, using previous step results" + ) + break + + current_proposals = next_proposals + await clogger.info(f"Step {step} produced {len(current_proposals)} aggregated proposals") + + # Select final output (use first/best from final step) + if current_proposals: + final_proposal = current_proposals[0] + output = final_proposal["output"] + await clogger.info("RSA mode completed successfully") + else: + raise ValueError("RSA aggregation produced no valid results") + + 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( From 6d23afacfe550e355b9ecd309cef38b00e3fb508 Mon Sep 17 00:00:00 2001 From: Shashank Kushwaha Date: Wed, 1 Apr 2026 22:02:49 -0700 Subject: [PATCH 03/24] feat: Add RSA fields to FlaskRunSettings interface --- flask-app/src/types.ts | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/flask-app/src/types.ts b/flask-app/src/types.ts index 2dfe2316..ab144309 100644 --- a/flask-app/src/types.ts +++ b/flask-app/src/types.ts @@ -81,6 +81,11 @@ export type MoleculeNameFormat = 'brand' | 'iupac' | 'formula' | 'smiles'; export interface FlaskRunSettings { moleculeName: MoleculeNameFormat; promptDebugging: boolean; + useRsa?: boolean; + rsaMode?: 'standalone' | 'rag'; + rsaN?: number; + rsaK?: number; + rsaT?: number; } import type { OrchestratorSettings, SidebarMessage, SidebarState } from 'lc-conductor'; From 85e156325673af87b700eb1ec74784e694b58371 Mon Sep 17 00:00:00 2001 From: Shashank Kushwaha Date: Wed, 1 Apr 2026 22:04:50 -0700 Subject: [PATCH 04/24] feat: Add RSA UI controls to App.tsx --- flask-app/src/App.tsx | 100 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 100 insertions(+) diff --git a/flask-app/src/App.tsx b/flask-app/src/App.tsx index 338f9c79..bb4b285e 100644 --- a/flask-app/src/App.tsx +++ b/flask-app/src/App.tsx @@ -130,6 +130,13 @@ const ChemistryTool: React.FC = () => { const [editedPrompt, setEditedPrompt] = useState(''); const [debugModalMinimized, setDebugModalMinimized] = useState(false); + // 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 +493,11 @@ const ChemistryTool: React.FC = () => { runSettings: { promptDebugging: debugMode, moleculeName: orchestratorSettings.moleculeName || 'brand', + useRsa, + rsaMode, + rsaN, + rsaK, + rsaT, }, customization, }; @@ -1145,6 +1157,11 @@ const ChemistryTool: React.FC = () => { runSettings: { promptDebugging: debugMode, moleculeName: orchestratorSettings.moleculeName || 'brand', + useRsa, + rsaMode, + rsaN, + rsaK, + rsaT, }, ...propertyDetails, }; @@ -1515,6 +1532,89 @@ const ChemistryTool: React.FC = () => { AI Debug Mode + + {/* RSA Settings - Only for Retrosynthesis */} + {problemType === 'retrosynthesis' && ( +
+
+ +
+ + {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 +
+
+ )} +
+ )} From 2046f3499ceef36278942a626ad46cdcbadbeb82 Mon Sep 17 00:00:00 2001 From: Shashank Kushwaha Date: Thu, 2 Apr 2026 00:04:30 -0700 Subject: [PATCH 05/24] Fix AttributeError by initializing retro_synth_context in __init__ --- charge_backend/backend_manager.py | 1 + 1 file changed, 1 insertion(+) diff --git a/charge_backend/backend_manager.py b/charge_backend/backend_manager.py index f9def612..974b5a94 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: From 181203d35b7b552ad96bdcb86c30bd54dfa375d1 Mon Sep 17 00:00:00 2001 From: Shashank Kushwaha Date: Thu, 2 Apr 2026 00:56:31 -0700 Subject: [PATCH 06/24] Update ChARGe submodule to restore OpenAIResponsesClient --- ChARGe | 1 + 1 file changed, 1 insertion(+) create mode 160000 ChARGe diff --git a/ChARGe b/ChARGe new file mode 160000 index 00000000..30209322 --- /dev/null +++ b/ChARGe @@ -0,0 +1 @@ +Subproject commit 30209322451c4ca307198431f4c576a0a06b568d From 0d1f6f627b427fa17b4cc35539bd2cd30624cf43 Mon Sep 17 00:00:00 2001 From: Shashank Kushwaha Date: Thu, 2 Apr 2026 00:59:06 -0700 Subject: [PATCH 07/24] Update ChARGe submodule with OpenAIResponsesClient parameter fix --- ChARGe | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ChARGe b/ChARGe index 30209322..c945255b 160000 --- a/ChARGe +++ b/ChARGe @@ -1 +1 @@ -Subproject commit 30209322451c4ca307198431f4c576a0a06b568d +Subproject commit c945255bfc0d7ab53cc01c8b9a8f438a53949219 From 287151a39abdc6004448711ab2481b3d2ab804a7 Mon Sep 17 00:00:00 2001 From: Shashank Kushwaha Date: Thu, 2 Apr 2026 01:12:43 -0700 Subject: [PATCH 08/24] Fix session reset between retries and add temperature=0.8 for RSA proposers --- ChARGe | 2 +- charge_backend/retrosynthesis/ai.py | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/ChARGe b/ChARGe index c945255b..6bcfc14f 160000 --- a/ChARGe +++ b/ChARGe @@ -1 +1 @@ -Subproject commit c945255bfc0d7ab53cc01c8b9a8f438a53949219 +Subproject commit 6bcfc14fc4c467c03bd4d40421e25aea45109398 diff --git a/charge_backend/retrosynthesis/ai.py b/charge_backend/retrosynthesis/ai.py index abf069a7..db5d7547 100644 --- a/charge_backend/retrosynthesis/ai.py +++ b/charge_backend/retrosynthesis/ai.py @@ -157,11 +157,12 @@ async def ai_based_retrosynthesis( for i in range(rsa_n): await clogger.info(f"Generating proposal {i+1}/{rsa_n}") try: - # Create a fresh task for each proposal + # Create a fresh task for each proposal with higher temperature for diversity proposal_task = RetrosynthesisTask( user_prompt=user_prompt, server_urls=available_tools, builtin_tools=builtin_tools or [], + temperature=0.8, # Higher temperature for diverse proposals ) runner.task = proposal_task From 892cde181f19ba2e28161a6d41255483d8ca1e9b Mon Sep 17 00:00:00 2001 From: Shashank Kushwaha Date: Thu, 2 Apr 2026 01:15:12 -0700 Subject: [PATCH 09/24] Remove temperature parameter for gpt-5.2 compatibility --- charge_backend/retrosynthesis/ai.py | 4 +- externals/lc_conductor | 2 +- flask-app/package-lock.json | 63 ++++++++++- flask-tools | 1 + flask_copilot.egg-info/PKG-INFO | 117 ++++++++++++++++++++ flask_copilot.egg-info/SOURCES.txt | 39 +++++++ flask_copilot.egg-info/dependency_links.txt | 1 + flask_copilot.egg-info/entry_points.txt | 2 + flask_copilot.egg-info/requires.txt | 36 ++++++ flask_copilot.egg-info/top_level.txt | 1 + 10 files changed, 262 insertions(+), 4 deletions(-) create mode 160000 flask-tools create mode 100644 flask_copilot.egg-info/PKG-INFO create mode 100644 flask_copilot.egg-info/SOURCES.txt create mode 100644 flask_copilot.egg-info/dependency_links.txt create mode 100644 flask_copilot.egg-info/entry_points.txt create mode 100644 flask_copilot.egg-info/requires.txt create mode 100644 flask_copilot.egg-info/top_level.txt diff --git a/charge_backend/retrosynthesis/ai.py b/charge_backend/retrosynthesis/ai.py index db5d7547..3dd9c3c3 100644 --- a/charge_backend/retrosynthesis/ai.py +++ b/charge_backend/retrosynthesis/ai.py @@ -157,12 +157,12 @@ async def ai_based_retrosynthesis( for i in range(rsa_n): await clogger.info(f"Generating proposal {i+1}/{rsa_n}") try: - # Create a fresh task for each proposal with higher temperature for diversity + # Create a fresh task for each proposal + # TODO: Add temperature=0.8 for models that support it (Claude, etc.) proposal_task = RetrosynthesisTask( user_prompt=user_prompt, server_urls=available_tools, builtin_tools=builtin_tools or [], - temperature=0.8, # Higher temperature for diverse proposals ) runner.task = proposal_task diff --git a/externals/lc_conductor b/externals/lc_conductor index 380c2cfd..2873f402 160000 --- a/externals/lc_conductor +++ b/externals/lc_conductor @@ -1 +1 @@ -Subproject commit 380c2cfd071b2ebb268b7a36e38b9ab1929b992b +Subproject commit 2873f40228693d71da1da1b232ea201ab9279ca1 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-tools b/flask-tools new file mode 160000 index 00000000..28bad724 --- /dev/null +++ b/flask-tools @@ -0,0 +1 @@ +Subproject commit 28bad7248d7ec2944cd52b9fa095c66cf87d57c7 diff --git a/flask_copilot.egg-info/PKG-INFO b/flask_copilot.egg-info/PKG-INFO new file mode 100644 index 00000000..1d1e4772 --- /dev/null +++ b/flask_copilot.egg-info/PKG-INFO @@ -0,0 +1,117 @@ +Metadata-Version: 2.4 +Name: flask-copilot +Version: 0.1.0 +Summary: FLASK-copilot: Interactive Agentic Framework for Molecular discovery +License: ################################################################################ + ## Copyright 2025 Lawrence Livermore National Security, LLC. + ## See the top-level LICENSE file for details. + ## + ## SPDX-License-Identifier: Apache-2.0 + ################################################################################ + Copyright (c) 2025, Lawrence Livermore National Security, LLC. + Produced at the Lawrence Livermore National Laboratory. + + LLNL-CODE-2006345. + All rights reserved. + + This file is part of FLASK Project: Foundation Learning AI for Synthesis Knowledge. + For details, see https://github.com/FLASK-LLNL + + Licensed under the Apache License, Version 2.0 (the "Licensee"); you + may not use this file except in compliance with the License. You may + obtain a copy of the License at: + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied. See the License for the specific language governing + permissions and limitations under the license. + +Requires-Python: <3.13,>=3.10 +Description-Content-Type: text/markdown +License-File: LICENSE +License-File: NOTICE +Requires-Dist: charge>=0.1.0 +Requires-Dist: lc_conductor>=0.1.0 +Requires-Dist: flask_tools>=0.1.0 +Requires-Dist: mcp>=1.10.0 +Requires-Dist: jsonschema>=4.17.3 +Requires-Dist: loguru>=0.7.0 +Requires-Dist: pydantic>=2.10.7 +Requires-Dist: click +Requires-Dist: pre-commit +Requires-Dist: black +Requires-Dist: fastapi +Requires-Dist: uvicorn +Requires-Dist: python-multipart +Requires-Dist: websockets +Requires-Dist: requests +Provides-Extra: aizynthfinder +Requires-Dist: paretoset; extra == "aizynthfinder" +Requires-Dist: rdchiral; extra == "aizynthfinder" +Requires-Dist: wrapt_timeout_decorator; extra == "aizynthfinder" +Requires-Dist: swifter; extra == "aizynthfinder" +Requires-Dist: apted; extra == "aizynthfinder" +Requires-Dist: scipy; extra == "aizynthfinder" +Requires-Dist: onnxruntime; extra == "aizynthfinder" +Requires-Dist: dask[dataframe]>=2025.9.0; extra == "aizynthfinder" +Requires-Dist: tables; extra == "aizynthfinder" +Requires-Dist: networkx; extra == "aizynthfinder" +Requires-Dist: xxhash; extra == "aizynthfinder" +Requires-Dist: jinja2; extra == "aizynthfinder" +Requires-Dist: markupsafe; extra == "aizynthfinder" +Provides-Extra: pds +Requires-Dist: flask-copilot[aizynthfinder]; extra == "pds" +Provides-Extra: all +Requires-Dist: flask-copilot[aizynthfinder]; extra == "all" +Dynamic: license-file + +# FLASK Copilot Web UI + +This is a Web UI for the FLASK Copilot, which presents computed molecules and +their properties. It can be used for reaction prediction, lead molecule +optimization, and other custom prompts. + +The FLASK Copilot consists of a React application as a frontend, and a Python +WebSocket-powered server as the backend. + +## Installing + +- Backend: + + - Make sure you have Python installed with a virtualenv. + - Go to the main folder and run `pip install -r requirements.txt` + - Alternativey: To install the package, clone the repository and run: + + ```bash + pip install -e .[all] + flask-copilot-install --extras all + ``` + +- Frontend: + - Install `npm` + - `cd` into the `flask-app` folder and run `npm install` + - Go into the `flask-app` directory and run `npm start dev` for development + work or `npm run build` for a production build of the app. + +## Running + +To run FLASK Copilot, both the frontend and the backend need to run. The backend +will also serve the frontend web UI on the same port, if `npm run build` was +run. If this is not the case (e.g., with `npm start dev`), the backend still +needs to run. A server that creates mock data will run with +`python mock_server.py`. + +Note: if the server was not running when the web UI started, click the blinking +red dot on the top right side to reconnect. + +## License + +Copyright (c) 2025, Lawrence Livermore National Security, LLC. +Produced at the Lawrence Livermore National Laboratory. + +SPDX-License-Identifier: Apache-2.0 + +LLNL-CODE-2006345 diff --git a/flask_copilot.egg-info/SOURCES.txt b/flask_copilot.egg-info/SOURCES.txt new file mode 100644 index 00000000..67659691 --- /dev/null +++ b/flask_copilot.egg-info/SOURCES.txt @@ -0,0 +1,39 @@ +LICENSE +NOTICE +README.md +pyproject.toml +charge_backend/__init__.py +charge_backend/backend_helper_funcs.py +charge_backend/backend_manager.py +charge_backend/builtin_tools.py +charge_backend/charge_backend_custom.py +charge_backend/charge_server.py +charge_backend/install.py +charge_backend/prompt_debugger.py +charge_backend/rdkit_mol_differ.py +charge_backend/rdkitjs_payload.py +charge_backend/lmo/lmo_charge_backend_funcs.py +charge_backend/lmo/lmo_task.py +charge_backend/moleculedb/__init__.py +charge_backend/moleculedb/dynamic_import.py +charge_backend/moleculedb/molecule_naming.py +charge_backend/moleculedb/purchasable.py +charge_backend/moleculedb/reactiondb_query.py +charge_backend/retrosynthesis/__init__.py +charge_backend/retrosynthesis/ai.py +charge_backend/retrosynthesis/aizynth_tools.py +charge_backend/retrosynthesis/alternatives.py +charge_backend/retrosynthesis/context.py +charge_backend/retrosynthesis/database.py +charge_backend/retrosynthesis/mapping.py +charge_backend/retrosynthesis/reaction_task.py +charge_backend/retrosynthesis/retrosynthesis_task.py +charge_backend/retrosynthesis/template.py +charge_backend/tests/test_rdkit_mol_differ.py +charge_backend/tests/test_websocket_callbacks.py +flask_copilot.egg-info/PKG-INFO +flask_copilot.egg-info/SOURCES.txt +flask_copilot.egg-info/dependency_links.txt +flask_copilot.egg-info/entry_points.txt +flask_copilot.egg-info/requires.txt +flask_copilot.egg-info/top_level.txt \ No newline at end of file diff --git a/flask_copilot.egg-info/dependency_links.txt b/flask_copilot.egg-info/dependency_links.txt new file mode 100644 index 00000000..8b137891 --- /dev/null +++ b/flask_copilot.egg-info/dependency_links.txt @@ -0,0 +1 @@ + diff --git a/flask_copilot.egg-info/entry_points.txt b/flask_copilot.egg-info/entry_points.txt new file mode 100644 index 00000000..c06fda36 --- /dev/null +++ b/flask_copilot.egg-info/entry_points.txt @@ -0,0 +1,2 @@ +[console_scripts] +flask-copilot-install = charge_backend.install:main diff --git a/flask_copilot.egg-info/requires.txt b/flask_copilot.egg-info/requires.txt new file mode 100644 index 00000000..def61c16 --- /dev/null +++ b/flask_copilot.egg-info/requires.txt @@ -0,0 +1,36 @@ +charge>=0.1.0 +lc_conductor>=0.1.0 +flask_tools>=0.1.0 +mcp>=1.10.0 +jsonschema>=4.17.3 +loguru>=0.7.0 +pydantic>=2.10.7 +click +pre-commit +black +fastapi +uvicorn +python-multipart +websockets +requests + +[aizynthfinder] +paretoset +rdchiral +wrapt_timeout_decorator +swifter +apted +scipy +onnxruntime +dask[dataframe]>=2025.9.0 +tables +networkx +xxhash +jinja2 +markupsafe + +[all] +flask-copilot[aizynthfinder] + +[pds] +flask-copilot[aizynthfinder] diff --git a/flask_copilot.egg-info/top_level.txt b/flask_copilot.egg-info/top_level.txt new file mode 100644 index 00000000..17602508 --- /dev/null +++ b/flask_copilot.egg-info/top_level.txt @@ -0,0 +1 @@ +charge_backend From f7f9c2e3a51d4c5e4bd5c4f39fb33c213b1c6749 Mon Sep 17 00:00:00 2001 From: Shashank Kushwaha Date: Thu, 2 Apr 2026 01:38:41 -0700 Subject: [PATCH 10/24] Add comprehensive RSA execution logging to capture prompts and outputs --- charge_backend/retrosynthesis/ai.py | 80 ++++++++++++++++++++++++++++- 1 file changed, 79 insertions(+), 1 deletion(-) diff --git a/charge_backend/retrosynthesis/ai.py b/charge_backend/retrosynthesis/ai.py index 3dd9c3c3..5ca338fb 100644 --- a/charge_backend/retrosynthesis/ai.py +++ b/charge_backend/retrosynthesis/ai.py @@ -1,6 +1,8 @@ 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 @@ -151,6 +153,13 @@ async def ai_based_retrosynthesis( 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}") + # Step 1: Generate N initial proposals await clogger.info(f"RSA Step 1/{rsa_t}: Generating {rsa_n} initial proposals") proposals = [] @@ -169,6 +178,16 @@ async def ai_based_retrosynthesis( if os.getenv("CHARGE_DISABLE_OUTPUT_VALIDATION", "0") == "1": proposal_task.structured_output_schema = None + # Save proposer prompt for examination + proposer_log = { + "proposal_index": i + 1, + "system_prompt": proposal_task.get_system_prompt(), + "user_prompt": user_prompt, + "mode": rsa_mode, + } + with open(f"{rsa_log_dir}/proposer_{i+1:02d}_prompt.json", "w") as f: + json.dump(proposer_log, f, indent=2) + # Run proposal proposal_output = await runner.run(log_progress) if isinstance(callback_handler, CallbackHandler): @@ -176,6 +195,18 @@ async def ai_based_retrosynthesis( # Validate and store proposal_result = ReactionOutputSchema.model_validate_json(proposal_output) + + # Save proposer output for examination + proposer_output_log = { + "proposal_index": i + 1, + "reasoning_summary": proposal_result.reasoning_summary, + "reactants_smiles": proposal_result.reactants_smiles_list, + "products_smiles": proposal_result.products_smiles_list, + "full_output": json.loads(proposal_output) + } + with open(f"{rsa_log_dir}/proposer_{i+1:02d}_output.json", "w") as f: + json.dump(proposer_output_log, f, indent=2) + proposals.append({ "output": proposal_output, "result": proposal_result, @@ -219,8 +250,10 @@ async def ai_based_retrosynthesis( # Format candidates text candidates_text = "" + subset_indices = [] for idx, prop in enumerate(subset, 1): prop_result = prop["result"] + subset_indices.append(prop["index"] + 1) # Convert to 1-indexed 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" @@ -241,6 +274,20 @@ async def ai_based_retrosynthesis( if os.getenv("CHARGE_DISABLE_OUTPUT_VALIDATION", "0") == "1": agg_task.structured_output_schema = None + # Save aggregator prompt for examination + aggregator_log = { + "step": step, + "aggregation_index": i + 1, + "k_subset_indices": subset_indices, # Which proposals were selected + "system_prompt": agg_task.get_system_prompt(), + "user_prompt": agg_task.get_user_prompt(), + "original_user_prompt": user_prompt, + "candidates_text": candidates_text, + "mode": rsa_mode, + } + with open(f"{rsa_log_dir}/aggregator_step{step}_{i+1:02d}_prompt.json", "w") as f: + json.dump(aggregator_log, f, indent=2) + # Run aggregation agg_output = await runner.run(log_progress) if isinstance(callback_handler, CallbackHandler): @@ -248,6 +295,20 @@ async def ai_based_retrosynthesis( # Validate and store agg_result = ReactionOutputSchema.model_validate_json(agg_output) + + # Save aggregator output for examination + aggregator_output_log = { + "step": step, + "aggregation_index": i + 1, + "k_subset_indices": subset_indices, + "reasoning_summary": agg_result.reasoning_summary, + "reactants_smiles": agg_result.reactants_smiles_list, + "products_smiles": agg_result.products_smiles_list, + "full_output": json.loads(agg_output) + } + with open(f"{rsa_log_dir}/aggregator_step{step}_{i+1:02d}_output.json", "w") as f: + json.dump(aggregator_output_log, f, indent=2) + next_proposals.append({ "output": agg_output, "result": agg_result, @@ -272,7 +333,24 @@ async def ai_based_retrosynthesis( if current_proposals: final_proposal = current_proposals[0] output = final_proposal["output"] - await clogger.info("RSA mode completed successfully") + + # Save final output for examination + final_result = json.loads(output) + final_output_log = { + "final_step": rsa_t, + "mode": rsa_mode, + "n_proposals": rsa_n, + "k_subset_size": rsa_k, + "t_stages": rsa_t, + "final_reasoning": final_result.get("reasoning_summary", ""), + "final_reactants_smiles": final_result.get("reactants_smiles_list", []), + "final_products_smiles": final_result.get("products_smiles_list", []), + "full_output": final_result + } + with open(f"{rsa_log_dir}/FINAL_OUTPUT.json", "w") as f: + json.dump(final_output_log, f, indent=2) + + await clogger.info(f"RSA mode completed successfully. Logs saved to: {rsa_log_dir}") else: raise ValueError("RSA aggregation produced no valid results") From 276feeb9fdbf3c1fcd07352586a0455559564aff Mon Sep 17 00:00:00 2001 From: Shashank Kushwaha Date: Thu, 2 Apr 2026 04:27:49 -0700 Subject: [PATCH 11/24] Fix RSA settings propagation to backend - Add RSA parameters (useRsa, rsaMode, rsaN, rsaK, rsaT) to sendMessageToServer - Add RSA parameters to handleComputeFlaskAI for compute-reaction-from action - Include RSA settings in dependency arrays for React hooks - Ensures RSA mode is sent when clicking 'How do I make this?' on nodes --- flask-app/src/App.tsx | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/flask-app/src/App.tsx b/flask-app/src/App.tsx index bb4b285e..c505df9c 100644 --- a/flask-app/src/App.tsx +++ b/flask-app/src/App.tsx @@ -1027,13 +1027,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, useRsa, rsaMode, rsaN, rsaK, rsaT] ); const handleReactionCardClick = useCallback( @@ -1107,6 +1112,11 @@ const ChemistryTool: React.FC = () => { runSettings: { promptDebugging: debugMode, moleculeName: orchestratorSettings.moleculeName || 'brand', + useRsa, + rsaMode, + rsaN, + rsaK, + rsaT, }, }) ); @@ -1114,8 +1124,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 || []; From b8809a4d33650ba7105a0243d3aeb4656f21c90e Mon Sep 17 00:00:00 2001 From: Shashank Kushwaha Date: Thu, 2 Apr 2026 04:28:17 -0700 Subject: [PATCH 12/24] Add intelligent routing for AI-based vs template-based retrosynthesis - Route to AI-based retrosynthesis when RSA enabled or config file missing - Create root node for AI-based path (matching template-based behavior) - Support both standard one-shot and RSA modes from initial compute action - Maintain backward compatibility with template-based when config exists - Fix import for smiles_to_html from correct module (molecule_naming) --- charge_backend/backend_manager.py | 75 +++++++++++++++++++++++++++---- 1 file changed, 67 insertions(+), 8 deletions(-) diff --git a/charge_backend/backend_manager.py b/charge_backend/backend_manager.py index 974b5a94..2f90e894 100644 --- a/charge_backend/backend_manager.py +++ b/charge_backend/backend_manager.py @@ -226,14 +226,73 @@ 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 (RSA or standard AI) + # AI-based is the default; template-based only runs if config exists and AI not requested + use_ai_based = self.run_settings.use_rsa 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()) From 58834b7a279981c7d3f0e7b1550dddb5978f768b Mon Sep 17 00:00:00 2001 From: Shashank Kushwaha Date: Thu, 2 Apr 2026 04:29:03 -0700 Subject: [PATCH 13/24] Make template-based retrosynthesis optional when config missing - Add os import at module level in both ai.py and template.py - Check if config file exists before attempting template expansion - Gracefully skip template search with info message when config unavailable - Prevent FileNotFoundError when /data/config.yml doesn't exist - Allow AI-based retrosynthesis to complete without template dependency - Production-ready: works with or without AiZynthFinder config --- charge_backend/retrosynthesis/ai.py | 92 +++++++++++++---------- charge_backend/retrosynthesis/template.py | 9 +++ 2 files changed, 61 insertions(+), 40 deletions(-) diff --git a/charge_backend/retrosynthesis/ai.py b/charge_backend/retrosynthesis/ai.py index 5ca338fb..e0484802 100644 --- a/charge_backend/retrosynthesis/ai.py +++ b/charge_backend/retrosynthesis/ai.py @@ -469,48 +469,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 + # 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." + ) + 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) + # 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, - ) + # 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 + # 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/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: From dd360ca8bbc6e3464fe3bf56b70f564952671718 Mon Sep 17 00:00:00 2001 From: Shashank Kushwaha Date: Thu, 2 Apr 2026 04:35:58 -0700 Subject: [PATCH 14/24] Add dynamic WebSocket URL detection with HTTPS support - Auto-detect WebSocket URL from request host header and protocol - Use wss:// for HTTPS requests, ws:// for HTTP requests - Support VS Code port forwarding (e.g., localhost:8003 -> ws://localhost:8003/ws) - Remove hardcoded localhost:8001 default - Maintain backward compatibility with WS_SERVER environment variable - Production-ready: works for local dev, port forwarding, and HTTPS hosting --- charge_backend/charge_server.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) 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""" """, From af1608fe05cf6e1a032b3bd11eb0144541fdb6e7 Mon Sep 17 00:00:00 2001 From: Shashank Kushwaha Date: Thu, 2 Apr 2026 14:19:39 -0700 Subject: [PATCH 15/24] Reduce RAG database queries from N(T+1) to 1 - Query once before proposals, inject into all prompts - Update system prompt to prevent redundant database queries --- charge_backend/retrosynthesis/ai.py | 40 +++++++++++++++++++ .../retrosynthesis/prompts/rsa_rag_system.txt | 4 +- 2 files changed, 42 insertions(+), 2 deletions(-) diff --git a/charge_backend/retrosynthesis/ai.py b/charge_backend/retrosynthesis/ai.py index e0484802..77f55cac 100644 --- a/charge_backend/retrosynthesis/ai.py +++ b/charge_backend/retrosynthesis/ai.py @@ -160,6 +160,46 @@ async def ai_based_retrosynthesis( 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 + user_prompt_with_rag = user_prompt + 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") + + # 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" + # Step 1: Generate N initial proposals await clogger.info(f"RSA Step 1/{rsa_t}: Generating {rsa_n} initial proposals") proposals = [] diff --git a/charge_backend/retrosynthesis/prompts/rsa_rag_system.txt b/charge_backend/retrosynthesis/prompts/rsa_rag_system.txt index 8dc78b7a..19875280 100644 --- a/charge_backend/retrosynthesis/prompts/rsa_rag_system.txt +++ b/charge_backend/retrosynthesis/prompts/rsa_rag_system.txt @@ -1,8 +1,8 @@ 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. First, query the reaction database using the query_reaction_database tool to find known reactions that produce the target molecule. +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. -After querying the database, reason using retrosynthesis principles: +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 From 0cd6ac8133eadb49802706b85e121f0cd042840f Mon Sep 17 00:00:00 2001 From: Shashank Kushwaha Date: Thu, 2 Apr 2026 14:46:01 -0700 Subject: [PATCH 16/24] Fix RAG mode to use user_prompt_with_rag in proposals Bug: Proposals were using user_prompt instead of user_prompt_with_rag Result: Database queried N times instead of 1 (once per proposal) Fix: Use user_prompt_with_rag for proposal tasks (safe for standalone/RAG/non-RSA) --- charge_backend/retrosynthesis/ai.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/charge_backend/retrosynthesis/ai.py b/charge_backend/retrosynthesis/ai.py index 77f55cac..267a77fc 100644 --- a/charge_backend/retrosynthesis/ai.py +++ b/charge_backend/retrosynthesis/ai.py @@ -209,7 +209,7 @@ async def ai_based_retrosynthesis( # Create a fresh task for each proposal # TODO: Add temperature=0.8 for models that support it (Claude, etc.) proposal_task = RetrosynthesisTask( - user_prompt=user_prompt, + user_prompt=user_prompt_with_rag, server_urls=available_tools, builtin_tools=builtin_tools or [], ) @@ -222,7 +222,7 @@ async def ai_based_retrosynthesis( proposer_log = { "proposal_index": i + 1, "system_prompt": proposal_task.get_system_prompt(), - "user_prompt": user_prompt, + "user_prompt": user_prompt_with_rag, "mode": rsa_mode, } with open(f"{rsa_log_dir}/proposer_{i+1:02d}_prompt.json", "w") as f: From e901a0d646d14d5f5d9f68d17f23e293d00f2866 Mon Sep 17 00:00:00 2001 From: Shashank Kushwaha Date: Thu, 2 Apr 2026 15:03:18 -0700 Subject: [PATCH 17/24] Filter query_reaction_database tool in RSA modes RAG mode: Query once, inject results, then remove tool (1 query total) Standalone mode: Remove tool entirely (0 queries) Non-RSA: Unaffected (tool remains available) Changes only within RSA block (lines 144-404) --- charge_backend/retrosynthesis/ai.py | 22 ++++++++++++++++++++-- 1 file changed, 20 insertions(+), 2 deletions(-) diff --git a/charge_backend/retrosynthesis/ai.py b/charge_backend/retrosynthesis/ai.py index 267a77fc..f8543788 100644 --- a/charge_backend/retrosynthesis/ai.py +++ b/charge_backend/retrosynthesis/ai.py @@ -161,7 +161,10 @@ async def ai_based_retrosynthesis( 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: @@ -200,6 +203,21 @@ async def ai_based_retrosynthesis( 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)") + # Step 1: Generate N initial proposals await clogger.info(f"RSA Step 1/{rsa_t}: Generating {rsa_n} initial proposals") proposals = [] @@ -211,7 +229,7 @@ async def ai_based_retrosynthesis( proposal_task = RetrosynthesisTask( user_prompt=user_prompt_with_rag, server_urls=available_tools, - builtin_tools=builtin_tools or [], + builtin_tools=builtin_tools_filtered, ) runner.task = proposal_task @@ -307,7 +325,7 @@ async def ai_based_retrosynthesis( total_steps=rsa_t, mode=rsa_mode, server_urls=available_tools, - builtin_tools=builtin_tools or [], + builtin_tools=builtin_tools_filtered, ) runner.task = agg_task From eb00bb7c298e0c25f8e0c7357d79cb4b395857ca Mon Sep 17 00:00:00 2001 From: Shashank Kushwaha Date: Thu, 2 Apr 2026 15:19:24 -0700 Subject: [PATCH 18/24] Add UI display for RAG database query results Show summary of retrieved reactions in reasoning panel: - First 5 reactions with names and components - Reactants and products for each - Confirmation that results are injected into prompts --- charge_backend/retrosynthesis/ai.py | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/charge_backend/retrosynthesis/ai.py b/charge_backend/retrosynthesis/ai.py index f8543788..964ac716 100644 --- a/charge_backend/retrosynthesis/ai.py +++ b/charge_backend/retrosynthesis/ai.py @@ -174,6 +174,26 @@ async def ai_based_retrosynthesis( 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" From 1830ed16135143c700f62491c8cec6d60b79b236 Mon Sep 17 00:00:00 2001 From: Shashank Kushwaha Date: Thu, 2 Apr 2026 20:48:22 -0700 Subject: [PATCH 19/24] Extract generic RSA algorithm into reusable helper function Created rsa_algorithm.py with run_rsa_loop() and refactored retrosynthesis to use it. --- charge_backend/retrosynthesis/ai.py | 236 ++++++--------------------- charge_backend/rsa_algorithm.py | 240 ++++++++++++++++++++++++++++ 2 files changed, 286 insertions(+), 190 deletions(-) create mode 100644 charge_backend/rsa_algorithm.py diff --git a/charge_backend/retrosynthesis/ai.py b/charge_backend/retrosynthesis/ai.py index 964ac716..0508dd1f 100644 --- a/charge_backend/retrosynthesis/ai.py +++ b/charge_backend/retrosynthesis/ai.py @@ -144,6 +144,8 @@ async def ai_based_retrosynthesis( 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 @@ -238,199 +240,53 @@ async def ai_based_retrosynthesis( ] await clogger.info("Standalone mode: Removed query_reaction_database from tools (no retrieval)") - # Step 1: Generate N initial proposals - await clogger.info(f"RSA Step 1/{rsa_t}: Generating {rsa_n} initial proposals") - proposals = [] - for i in range(rsa_n): - await clogger.info(f"Generating proposal {i+1}/{rsa_n}") - try: - # Create a fresh task for each proposal - # TODO: Add temperature=0.8 for models that support it (Claude, etc.) - proposal_task = RetrosynthesisTask( - user_prompt=user_prompt_with_rag, - server_urls=available_tools, - builtin_tools=builtin_tools_filtered, - ) - runner.task = proposal_task - - if os.getenv("CHARGE_DISABLE_OUTPUT_VALIDATION", "0") == "1": - proposal_task.structured_output_schema = None - - # Save proposer prompt for examination - proposer_log = { - "proposal_index": i + 1, - "system_prompt": proposal_task.get_system_prompt(), - "user_prompt": user_prompt_with_rag, - "mode": rsa_mode, - } - with open(f"{rsa_log_dir}/proposer_{i+1:02d}_prompt.json", "w") as f: - json.dump(proposer_log, f, indent=2) - - # Run proposal - proposal_output = await runner.run(log_progress) - if isinstance(callback_handler, CallbackHandler): - await callback_handler.drain() - - # Validate and store - proposal_result = ReactionOutputSchema.model_validate_json(proposal_output) - - # Save proposer output for examination - proposer_output_log = { - "proposal_index": i + 1, - "reasoning_summary": proposal_result.reasoning_summary, - "reactants_smiles": proposal_result.reactants_smiles_list, - "products_smiles": proposal_result.products_smiles_list, - "full_output": json.loads(proposal_output) - } - with open(f"{rsa_log_dir}/proposer_{i+1:02d}_output.json", "w") as f: - json.dump(proposer_output_log, f, indent=2) - - proposals.append({ - "output": proposal_output, - "result": proposal_result, - "index": i - }) - await clogger.info(f"Proposal {i+1} completed successfully") - except Exception as e: - await clogger.warning(f"Proposal {i+1} failed: {str(e)}") - continue - - if not proposals: - raise ValueError("All RSA proposals failed, falling back to standard mode") - - await clogger.info(f"Generated {len(proposals)} valid proposals") + # 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, + ) - # Steps 2..T: Recursive aggregation - current_proposals = proposals - for step in range(2, rsa_t + 1): - await clogger.info( - f"RSA Step {step}/{rsa_t}: Aggregating {len(current_proposals)} proposals into {rsa_k}-subsets" + 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, ) - if len(current_proposals) < rsa_k: - await clogger.warning( - f"Not enough proposals ({len(current_proposals)}) for K={rsa_k}, using all available" - ) - rsa_k = len(current_proposals) - - # Generate new proposals by aggregating K-subsets - next_proposals = [] - num_aggregations = max(rsa_n, len(current_proposals)) - - for i in range(num_aggregations): - await clogger.info(f"Aggregation {i+1}/{num_aggregations}") - try: - # Select K random proposals - if len(current_proposals) <= rsa_k: - subset = current_proposals - else: - subset = random.sample(current_proposals, rsa_k) - - # Format candidates text - candidates_text = "" - subset_indices = [] - for idx, prop in enumerate(subset, 1): - prop_result = prop["result"] - subset_indices.append(prop["index"] + 1) # Convert to 1-indexed - 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" - - # Create aggregation task - agg_task = RSAAggregationTask( - original_user_prompt=user_prompt, - candidates_text=candidates_text, - step=step, - total_steps=rsa_t, - mode=rsa_mode, - server_urls=available_tools, - builtin_tools=builtin_tools_filtered, - ) - runner.task = agg_task - - if os.getenv("CHARGE_DISABLE_OUTPUT_VALIDATION", "0") == "1": - agg_task.structured_output_schema = None - - # Save aggregator prompt for examination - aggregator_log = { - "step": step, - "aggregation_index": i + 1, - "k_subset_indices": subset_indices, # Which proposals were selected - "system_prompt": agg_task.get_system_prompt(), - "user_prompt": agg_task.get_user_prompt(), - "original_user_prompt": user_prompt, - "candidates_text": candidates_text, - "mode": rsa_mode, - } - with open(f"{rsa_log_dir}/aggregator_step{step}_{i+1:02d}_prompt.json", "w") as f: - json.dump(aggregator_log, f, indent=2) - - # Run aggregation - agg_output = await runner.run(log_progress) - if isinstance(callback_handler, CallbackHandler): - await callback_handler.drain() - - # Validate and store - agg_result = ReactionOutputSchema.model_validate_json(agg_output) - - # Save aggregator output for examination - aggregator_output_log = { - "step": step, - "aggregation_index": i + 1, - "k_subset_indices": subset_indices, - "reasoning_summary": agg_result.reasoning_summary, - "reactants_smiles": agg_result.reactants_smiles_list, - "products_smiles": agg_result.products_smiles_list, - "full_output": json.loads(agg_output) - } - with open(f"{rsa_log_dir}/aggregator_step{step}_{i+1:02d}_output.json", "w") as f: - json.dump(aggregator_output_log, f, indent=2) - - next_proposals.append({ - "output": agg_output, - "result": agg_result, - "index": i, - "step": step - }) - await clogger.info(f"Aggregation {i+1} completed successfully") - except Exception as e: - await clogger.warning(f"Aggregation {i+1} failed: {str(e)}") - continue - - if not next_proposals: - await clogger.warning( - f"All aggregations at step {step} failed, using previous step results" - ) - break - - current_proposals = next_proposals - await clogger.info(f"Step {step} produced {len(current_proposals)} aggregated proposals") - - # Select final output (use first/best from final step) - if current_proposals: - final_proposal = current_proposals[0] - output = final_proposal["output"] - - # Save final output for examination - final_result = json.loads(output) - final_output_log = { - "final_step": rsa_t, - "mode": rsa_mode, - "n_proposals": rsa_n, - "k_subset_size": rsa_k, - "t_stages": rsa_t, - "final_reasoning": final_result.get("reasoning_summary", ""), - "final_reactants_smiles": final_result.get("reactants_smiles_list", []), - "final_products_smiles": final_result.get("products_smiles_list", []), - "full_output": final_result - } - with open(f"{rsa_log_dir}/FINAL_OUTPUT.json", "w") as f: - json.dump(final_output_log, f, indent=2) - - await clogger.info(f"RSA mode completed successfully. Logs saved to: {rsa_log_dir}") - else: - raise ValueError("RSA aggregation produced no valid results") + 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 + + # 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, + ) + + 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 diff --git a/charge_backend/rsa_algorithm.py b/charge_backend/rsa_algorithm.py new file mode 100644 index 00000000..dd33a70e --- /dev/null +++ b/charge_backend/rsa_algorithm.py @@ -0,0 +1,240 @@ +""" +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 +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, +) -> 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 + + 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 + + # Stage 1: Generate N initial proposals + await clogger.info(f"RSA Step 1/{t}: Generating {n} initial proposals") + proposals = [] + + for i in range(n): + await clogger.info(f"Generating proposal {i+1}/{n}") + try: + # Create proposal task + proposal_task = create_proposal_task() + 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": i + 1, + "system_prompt": proposal_task.get_system_prompt(), + "user_prompt": proposal_task.get_user_prompt(), + } + with open(f"{log_dir}/proposer_{i+1:02d}_prompt.json", "w") as f: + json.dump(proposer_log, f, indent=2) + + # Run proposal + proposal_output = await runner.run(log_progress) + if callback_handler: + await callback_handler.drain() + + # Validate output + proposal_result = output_schema.model_validate_json(proposal_output) + + # Save proposal output + proposer_output_log = { + "proposal_index": i + 1, + "result": proposal_result.model_dump(), + "full_output": json.loads(proposal_output) + } + with open(f"{log_dir}/proposer_{i+1:02d}_output.json", "w") as f: + json.dump(proposer_output_log, f, indent=2) + + proposals.append({ + "output": proposal_output, + "result": proposal_result, + "index": i + }) + await clogger.info(f"Proposal {i+1} completed successfully") + + except Exception as e: + await clogger.warning(f"Proposal {i+1} failed: {str(e)}") + continue + + if not proposals: + raise ValueError("All RSA proposals failed") + + await clogger.info(f"Generated {len(proposals)} valid proposals") + + # Stages 2-T: Recursive aggregation + current_proposals = proposals + + for step in range(2, t + 1): + await clogger.info( + f"RSA Step {step}/{t}: Aggregating {len(current_proposals)} proposals into {k}-subsets" + ) + + # 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 + next_proposals = [] + num_aggregations = max(n, len(current_proposals)) + + for i in range(num_aggregations): + await clogger.info(f"Aggregation {i+1}/{num_aggregations}") + try: + # 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 + ) + 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": i + 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}_{i+1:02d}_prompt.json", "w") as f: + json.dump(aggregator_log, f, indent=2) + + # Run aggregation + agg_output = await runner.run(log_progress) + if callback_handler: + await callback_handler.drain() + + # Validate output + agg_result = output_schema.model_validate_json(agg_output) + + # Save aggregation output + aggregator_output_log = { + "step": step, + "aggregation_index": i + 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}_{i+1:02d}_output.json", "w") as f: + json.dump(aggregator_output_log, f, indent=2) + + next_proposals.append({ + "output": agg_output, + "result": agg_result, + "index": i, + "step": step + }) + await clogger.info(f"Aggregation {i+1} completed successfully") + + except Exception as e: + await clogger.warning(f"Aggregation {i+1} failed: {str(e)}") + continue + + if not next_proposals: + await clogger.warning(f"No successful aggregations in step {step}, using previous proposals") + break + + current_proposals = next_proposals + + # 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 From 5f9fcdd0b7ba2657790043487604992201c1af9c Mon Sep 17 00:00:00 2001 From: Shashank Kushwaha Date: Thu, 2 Apr 2026 21:29:28 -0700 Subject: [PATCH 20/24] Add parallel proposal generation and fix empty reactants validation - Implement parallel execution for Stage 1 proposals using asyncio.gather() - Add runner_factory parameter to create independent runner instances - Add validation to filter proposals/aggregations with empty reactants - Default to parallel mode (parallel=True) with fallback to sequential - Retrosynthesis now creates independent runners for parallel execution --- charge_backend/retrosynthesis/ai.py | 13 +++++ charge_backend/rsa_algorithm.py | 86 +++++++++++++++++++++++------ 2 files changed, 81 insertions(+), 18 deletions(-) diff --git a/charge_backend/retrosynthesis/ai.py b/charge_backend/retrosynthesis/ai.py index 0508dd1f..ffce6944 100644 --- a/charge_backend/retrosynthesis/ai.py +++ b/charge_backend/retrosynthesis/ai.py @@ -270,6 +270,17 @@ def format_candidates(subset): 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, @@ -284,6 +295,8 @@ def format_candidates(subset): 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}") diff --git a/charge_backend/rsa_algorithm.py b/charge_backend/rsa_algorithm.py index dd33a70e..2078a8c4 100644 --- a/charge_backend/rsa_algorithm.py +++ b/charge_backend/rsa_algorithm.py @@ -14,6 +14,7 @@ import random import json import os +import asyncio from typing import Any, Callable, Optional from pathlib import Path @@ -31,6 +32,8 @@ async def run_rsa_loop( 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. @@ -50,6 +53,9 @@ async def run_rsa_loop( 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) @@ -65,16 +71,15 @@ async def run_rsa_loop( await clogger.warning(f"K ({k}) > N ({n}), adjusting K to N") k = n - # Stage 1: Generate N initial proposals - await clogger.info(f"RSA Step 1/{t}: Generating {n} initial proposals") - proposals = [] - - for i in range(n): - await clogger.info(f"Generating proposal {i+1}/{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() - runner.task = proposal_task + proposal_runner.task = proposal_task # Disable validation if requested if os.getenv("CHARGE_DISABLE_OUTPUT_VALIDATION", "0") == "1": @@ -82,40 +87,80 @@ async def run_rsa_loop( # Save proposal prompt proposer_log = { - "proposal_index": i + 1, + "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_{i+1:02d}_prompt.json", "w") as f: + 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 runner.run(log_progress) + 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": i + 1, + "proposal_index": proposal_index + 1, "result": proposal_result.model_dump(), "full_output": json.loads(proposal_output) } - with open(f"{log_dir}/proposer_{i+1:02d}_output.json", "w") as f: + with open(f"{log_dir}/proposer_{proposal_index+1:02d}_output.json", "w") as f: json.dump(proposer_output_log, f, indent=2) - proposals.append({ + await clogger.info(f"Proposal {proposal_index+1} completed successfully") + + return { "output": proposal_output, "result": proposal_result, - "index": i - }) - await clogger.info(f"Proposal {i+1} completed successfully") + "index": proposal_index + } except Exception as e: - await clogger.warning(f"Proposal {i+1} failed: {str(e)}") - continue + 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") @@ -187,6 +232,11 @@ async def run_rsa_loop( # 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 {i+1} has empty reactants (model refused or failed), skipping") + continue + # Save aggregation output aggregator_output_log = { "step": step, From e97d7a160bf772cf99bd420b4d8e8fba997b5dd8 Mon Sep 17 00:00:00 2001 From: Shashank Kushwaha Date: Thu, 2 Apr 2026 21:51:32 -0700 Subject: [PATCH 21/24] Parallelize aggregation stages with proper RSA barriers Aggregations within each stage run in parallel. Barriers enforce sequential stage progression. --- charge_backend/rsa_algorithm.py | 193 ++++++++++++++++++++------------ 1 file changed, 119 insertions(+), 74 deletions(-) diff --git a/charge_backend/rsa_algorithm.py b/charge_backend/rsa_algorithm.py index 2078a8c4..16454ba0 100644 --- a/charge_backend/rsa_algorithm.py +++ b/charge_backend/rsa_algorithm.py @@ -167,12 +167,97 @@ async def run_single_proposal(proposal_index: int, proposal_runner: Any): 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" + f"RSA Step {step}/{t}: Aggregating {len(current_proposals)} proposals into {k}-subsets" + + (" (parallel mode)" if parallel else " (sequential mode)") ) # Adjust K if needed @@ -184,87 +269,47 @@ async def run_single_proposal(proposal_index: int, proposal_runner: Any): current_k = len(current_proposals) # Generate aggregations - next_proposals = [] num_aggregations = max(n, len(current_proposals)) - for i in range(num_aggregations): - await clogger.info(f"Aggregation {i+1}/{num_aggregations}") - try: - # 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 - ) - 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": i + 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}_{i+1:02d}_prompt.json", "w") as f: - json.dump(aggregator_log, f, indent=2) - - # Run aggregation - agg_output = await 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 {i+1} has empty reactants (model refused or failed), skipping") - continue - - # Save aggregation output - aggregator_output_log = { - "step": step, - "aggregation_index": i + 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}_{i+1:02d}_output.json", "w") as f: - json.dump(aggregator_output_log, f, indent=2) - - next_proposals.append({ - "output": agg_output, - "result": agg_result, - "index": i, - "step": step - }) - await clogger.info(f"Aggregation {i+1} completed successfully") - - except Exception as e: - await clogger.warning(f"Aggregation {i+1} failed: {str(e)}") - continue + 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: From 1381b77700012e088fb88d9ac7d0fc184d599abe Mon Sep 17 00:00:00 2001 From: Shashank Kushwaha Date: Thu, 2 Apr 2026 23:09:16 -0700 Subject: [PATCH 22/24] Remove egg-info build artifacts and update gitignore --- .gitignore | 4 + flask_copilot.egg-info/PKG-INFO | 117 -------------------- flask_copilot.egg-info/SOURCES.txt | 39 ------- flask_copilot.egg-info/dependency_links.txt | 1 - flask_copilot.egg-info/entry_points.txt | 2 - flask_copilot.egg-info/requires.txt | 36 ------ flask_copilot.egg-info/top_level.txt | 1 - 7 files changed, 4 insertions(+), 196 deletions(-) delete mode 100644 flask_copilot.egg-info/PKG-INFO delete mode 100644 flask_copilot.egg-info/SOURCES.txt delete mode 100644 flask_copilot.egg-info/dependency_links.txt delete mode 100644 flask_copilot.egg-info/entry_points.txt delete mode 100644 flask_copilot.egg-info/requires.txt delete mode 100644 flask_copilot.egg-info/top_level.txt 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/flask_copilot.egg-info/PKG-INFO b/flask_copilot.egg-info/PKG-INFO deleted file mode 100644 index 1d1e4772..00000000 --- a/flask_copilot.egg-info/PKG-INFO +++ /dev/null @@ -1,117 +0,0 @@ -Metadata-Version: 2.4 -Name: flask-copilot -Version: 0.1.0 -Summary: FLASK-copilot: Interactive Agentic Framework for Molecular discovery -License: ################################################################################ - ## Copyright 2025 Lawrence Livermore National Security, LLC. - ## See the top-level LICENSE file for details. - ## - ## SPDX-License-Identifier: Apache-2.0 - ################################################################################ - Copyright (c) 2025, Lawrence Livermore National Security, LLC. - Produced at the Lawrence Livermore National Laboratory. - - LLNL-CODE-2006345. - All rights reserved. - - This file is part of FLASK Project: Foundation Learning AI for Synthesis Knowledge. - For details, see https://github.com/FLASK-LLNL - - Licensed under the Apache License, Version 2.0 (the "Licensee"); you - may not use this file except in compliance with the License. You may - obtain a copy of the License at: - - http://www.apache.org/licenses/LICENSE-2.0 - - Unless required by applicable law or agreed to in writing, software - distributed under the License is distributed on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or - implied. See the License for the specific language governing - permissions and limitations under the license. - -Requires-Python: <3.13,>=3.10 -Description-Content-Type: text/markdown -License-File: LICENSE -License-File: NOTICE -Requires-Dist: charge>=0.1.0 -Requires-Dist: lc_conductor>=0.1.0 -Requires-Dist: flask_tools>=0.1.0 -Requires-Dist: mcp>=1.10.0 -Requires-Dist: jsonschema>=4.17.3 -Requires-Dist: loguru>=0.7.0 -Requires-Dist: pydantic>=2.10.7 -Requires-Dist: click -Requires-Dist: pre-commit -Requires-Dist: black -Requires-Dist: fastapi -Requires-Dist: uvicorn -Requires-Dist: python-multipart -Requires-Dist: websockets -Requires-Dist: requests -Provides-Extra: aizynthfinder -Requires-Dist: paretoset; extra == "aizynthfinder" -Requires-Dist: rdchiral; extra == "aizynthfinder" -Requires-Dist: wrapt_timeout_decorator; extra == "aizynthfinder" -Requires-Dist: swifter; extra == "aizynthfinder" -Requires-Dist: apted; extra == "aizynthfinder" -Requires-Dist: scipy; extra == "aizynthfinder" -Requires-Dist: onnxruntime; extra == "aizynthfinder" -Requires-Dist: dask[dataframe]>=2025.9.0; extra == "aizynthfinder" -Requires-Dist: tables; extra == "aizynthfinder" -Requires-Dist: networkx; extra == "aizynthfinder" -Requires-Dist: xxhash; extra == "aizynthfinder" -Requires-Dist: jinja2; extra == "aizynthfinder" -Requires-Dist: markupsafe; extra == "aizynthfinder" -Provides-Extra: pds -Requires-Dist: flask-copilot[aizynthfinder]; extra == "pds" -Provides-Extra: all -Requires-Dist: flask-copilot[aizynthfinder]; extra == "all" -Dynamic: license-file - -# FLASK Copilot Web UI - -This is a Web UI for the FLASK Copilot, which presents computed molecules and -their properties. It can be used for reaction prediction, lead molecule -optimization, and other custom prompts. - -The FLASK Copilot consists of a React application as a frontend, and a Python -WebSocket-powered server as the backend. - -## Installing - -- Backend: - - - Make sure you have Python installed with a virtualenv. - - Go to the main folder and run `pip install -r requirements.txt` - - Alternativey: To install the package, clone the repository and run: - - ```bash - pip install -e .[all] - flask-copilot-install --extras all - ``` - -- Frontend: - - Install `npm` - - `cd` into the `flask-app` folder and run `npm install` - - Go into the `flask-app` directory and run `npm start dev` for development - work or `npm run build` for a production build of the app. - -## Running - -To run FLASK Copilot, both the frontend and the backend need to run. The backend -will also serve the frontend web UI on the same port, if `npm run build` was -run. If this is not the case (e.g., with `npm start dev`), the backend still -needs to run. A server that creates mock data will run with -`python mock_server.py`. - -Note: if the server was not running when the web UI started, click the blinking -red dot on the top right side to reconnect. - -## License - -Copyright (c) 2025, Lawrence Livermore National Security, LLC. -Produced at the Lawrence Livermore National Laboratory. - -SPDX-License-Identifier: Apache-2.0 - -LLNL-CODE-2006345 diff --git a/flask_copilot.egg-info/SOURCES.txt b/flask_copilot.egg-info/SOURCES.txt deleted file mode 100644 index 67659691..00000000 --- a/flask_copilot.egg-info/SOURCES.txt +++ /dev/null @@ -1,39 +0,0 @@ -LICENSE -NOTICE -README.md -pyproject.toml -charge_backend/__init__.py -charge_backend/backend_helper_funcs.py -charge_backend/backend_manager.py -charge_backend/builtin_tools.py -charge_backend/charge_backend_custom.py -charge_backend/charge_server.py -charge_backend/install.py -charge_backend/prompt_debugger.py -charge_backend/rdkit_mol_differ.py -charge_backend/rdkitjs_payload.py -charge_backend/lmo/lmo_charge_backend_funcs.py -charge_backend/lmo/lmo_task.py -charge_backend/moleculedb/__init__.py -charge_backend/moleculedb/dynamic_import.py -charge_backend/moleculedb/molecule_naming.py -charge_backend/moleculedb/purchasable.py -charge_backend/moleculedb/reactiondb_query.py -charge_backend/retrosynthesis/__init__.py -charge_backend/retrosynthesis/ai.py -charge_backend/retrosynthesis/aizynth_tools.py -charge_backend/retrosynthesis/alternatives.py -charge_backend/retrosynthesis/context.py -charge_backend/retrosynthesis/database.py -charge_backend/retrosynthesis/mapping.py -charge_backend/retrosynthesis/reaction_task.py -charge_backend/retrosynthesis/retrosynthesis_task.py -charge_backend/retrosynthesis/template.py -charge_backend/tests/test_rdkit_mol_differ.py -charge_backend/tests/test_websocket_callbacks.py -flask_copilot.egg-info/PKG-INFO -flask_copilot.egg-info/SOURCES.txt -flask_copilot.egg-info/dependency_links.txt -flask_copilot.egg-info/entry_points.txt -flask_copilot.egg-info/requires.txt -flask_copilot.egg-info/top_level.txt \ No newline at end of file diff --git a/flask_copilot.egg-info/dependency_links.txt b/flask_copilot.egg-info/dependency_links.txt deleted file mode 100644 index 8b137891..00000000 --- a/flask_copilot.egg-info/dependency_links.txt +++ /dev/null @@ -1 +0,0 @@ - diff --git a/flask_copilot.egg-info/entry_points.txt b/flask_copilot.egg-info/entry_points.txt deleted file mode 100644 index c06fda36..00000000 --- a/flask_copilot.egg-info/entry_points.txt +++ /dev/null @@ -1,2 +0,0 @@ -[console_scripts] -flask-copilot-install = charge_backend.install:main diff --git a/flask_copilot.egg-info/requires.txt b/flask_copilot.egg-info/requires.txt deleted file mode 100644 index def61c16..00000000 --- a/flask_copilot.egg-info/requires.txt +++ /dev/null @@ -1,36 +0,0 @@ -charge>=0.1.0 -lc_conductor>=0.1.0 -flask_tools>=0.1.0 -mcp>=1.10.0 -jsonschema>=4.17.3 -loguru>=0.7.0 -pydantic>=2.10.7 -click -pre-commit -black -fastapi -uvicorn -python-multipart -websockets -requests - -[aizynthfinder] -paretoset -rdchiral -wrapt_timeout_decorator -swifter -apted -scipy -onnxruntime -dask[dataframe]>=2025.9.0 -tables -networkx -xxhash -jinja2 -markupsafe - -[all] -flask-copilot[aizynthfinder] - -[pds] -flask-copilot[aizynthfinder] diff --git a/flask_copilot.egg-info/top_level.txt b/flask_copilot.egg-info/top_level.txt deleted file mode 100644 index 17602508..00000000 --- a/flask_copilot.egg-info/top_level.txt +++ /dev/null @@ -1 +0,0 @@ -charge_backend From 2a76705de1bb0f4cf5ff60abef1a09f3602cc698 Mon Sep 17 00:00:00 2001 From: Shashank Kushwaha Date: Thu, 2 Apr 2026 23:29:25 -0700 Subject: [PATCH 23/24] Separate AI-based flag from RSA flag for explicit retrosynthesis approach control --- charge_backend/backend_manager.py | 7 +++--- externals/lc_conductor | 2 +- flask-app/src/App.tsx | 39 ++++++++++++++++++++++++------- flask-app/src/types.ts | 1 + 4 files changed, 36 insertions(+), 13 deletions(-) diff --git a/charge_backend/backend_manager.py b/charge_backend/backend_manager.py index 2f90e894..6b9d2f1f 100644 --- a/charge_backend/backend_manager.py +++ b/charge_backend/backend_manager.py @@ -226,9 +226,10 @@ async def _handle_optimization( async def _handle_retrosynthesis(self, data: dict) -> None: """Handle retrosynthesis problem type.""" - # Check if AI-based retrosynthesis is requested (RSA or standard AI) - # AI-based is the default; template-based only runs if config exists and AI not requested - use_ai_based = self.run_settings.use_rsa or not os.path.exists(self.args.config_file) + # 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) diff --git a/externals/lc_conductor b/externals/lc_conductor index 2873f402..afcf0c80 160000 --- a/externals/lc_conductor +++ b/externals/lc_conductor @@ -1 +1 @@ -Subproject commit 2873f40228693d71da1da1b232ea201ab9279ca1 +Subproject commit afcf0c80ff356a1fd4037cfb6d8c001b617348f9 diff --git a/flask-app/src/App.tsx b/flask-app/src/App.tsx index c505df9c..f65f19e0 100644 --- a/flask-app/src/App.tsx +++ b/flask-app/src/App.tsx @@ -130,6 +130,9 @@ 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'); @@ -493,6 +496,7 @@ const ChemistryTool: React.FC = () => { runSettings: { promptDebugging: debugMode, moleculeName: orchestratorSettings.moleculeName || 'brand', + useAiBased, useRsa, rsaMode, rsaN, @@ -1038,7 +1042,7 @@ const ChemistryTool: React.FC = () => { wsRef.current.send(JSON.stringify(msg)); setContextMenu({ node: null, isReaction: false, x: 0, y: 0 }); }, - [debugMode, orchestratorSettings, useRsa, rsaMode, rsaN, rsaK, rsaT] + [debugMode, orchestratorSettings, useAiBased, useRsa, rsaMode, rsaN, rsaK, rsaT] ); const handleReactionCardClick = useCallback( @@ -1167,6 +1171,7 @@ const ChemistryTool: React.FC = () => { runSettings: { promptDebugging: debugMode, moleculeName: orchestratorSettings.moleculeName || 'brand', + useAiBased, useRsa, rsaMode, rsaN, @@ -1543,26 +1548,42 @@ const ChemistryTool: React.FC = () => { - {/* RSA Settings - Only for Retrosynthesis */} + {/* Retrosynthesis Settings */} {problemType === 'retrosynthesis' && (
- {useRsa && ( -
+ {useAiBased && ( +
+ +
+ )} + + {useAiBased && useRsa && ( +