Skip to content

Integrate ChARGe RSA algorithm - #194

Open
shashankk42 wants to merge 3 commits into
FLASK-LLNL:mainfrom
shashankk42:feature/use-charge-rsa
Open

Integrate ChARGe RSA algorithm#194
shashankk42 wants to merge 3 commits into
FLASK-LLNL:mainfrom
shashankk42:feature/use-charge-rsa

Conversation

@shashankk42

@shashankk42 shashankk42 commented Apr 12, 2026

Copy link
Copy Markdown

Summary

Integrate ChARGe RSA into FLASK-Copilot retrosynthesis, with consolidated flag handling, extracted RSA driver, and shared root-node creation. Addresses all PR review comments.

Changes

Backend

  • New file charge_backend/retrosynthesis/ai_rsa.py — chemistry-specific RSA driver wrapping ChARGe's RSATask via a
    RetroRSATask subclass (overrides format_candidates, validate_proposal)
  • ai_based_retrosynthesis gains root_smiles: Optional[str] = None — creates the root Node itself when called from compute;
    db_then_ai_retrosynthesis path unchanged
  • New helper build_root_node(smiles, run_settings) -> Node in backend_helper_funcs.py, reused by
    template_based_retrosynthesis (dedupe)
  • Drop FlaskRunSettings.use_ai_based; everything reads data.get("aiOnly", True) at the top level of the WS payload
  • Drop redundant os.path.exists(config_file) checks in ai.py and backend_manager.py; the consolidated check now lives only
    in template.py:run_retro_planner
  • Surface a clogger.warning(...) when the model returns no reactants (refusal indicator)
  • Unhighlight reactant node when template expansion is skipped (avoid stuck-yellow UI)

Frontend (flask-app)

  • State rename useAiBasedaiOnly; sent at top level of compute / custom-query / compute-reaction-from payloads (always
    passed)
  • Updated combined_customization_modal.tsx, retrosynthesis_customization_content.tsx, types.ts accordingly
  • retrosynthesis_customization_content.tsx imports RsaSettingsPanel from lc-conductor (single source of truth)
  • Revert package-lock.json and .gitignore to upstream main (per review)

Prompts (charge_backend/retrosynthesis/prompts/)

  • 3-part chemistry prompts for both RAG and standalone modes: rsa_{mode}_system.txt, rsa_{mode}_proposal.txt,
    rsa_{mode}_aggregation.txt

Modes

  • Standalone: pure chemistry reasoning, query_reaction_database removed from tools
  • RAG: queries the reaction DB once, injects results into all proposal prompts; tool removed afterward

Testing

  • Server boots; HTTP 200; all import sanity passes
  • AI-only, RSA-standalone, and RSA-RAG verified end-to-end on caffeine
  • Template mode verified against local test_reactions.db
  • pytest charge_backend/tests/test_websocket_callbacks.py 2/2

Dependencies

Related

@shashankk42
shashankk42 marked this pull request as draft April 16, 2026 22:27
@shashankk42
shashankk42 force-pushed the feature/use-charge-rsa branch from db3ade3 to 2e99d29 Compare April 24, 2026 00:08
@shashankk42
shashankk42 marked this pull request as ready for review April 24, 2026 01:51
@shashankk42
shashankk42 marked this pull request as draft April 24, 2026 04:51
@shashankk42
shashankk42 marked this pull request as ready for review April 24, 2026 05:25
@shashankk42
shashankk42 marked this pull request as draft April 24, 2026 06:31
@shashankk42
shashankk42 marked this pull request as ready for review April 24, 2026 06:34
@bvanessen
bvanessen force-pushed the feature/use-charge-rsa branch from 34c7019 to 3d87dfc Compare May 1, 2026 20:47

@tbennun tbennun left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Overall this looks great! Prompts are well-written etc. The code itself should be refactored to a separate top-level function or file that is called from within ai_based_retrosynthesis since it is quite lengthy. Additionally, some code duplication exists in this PR, as well as a potential merge conflict.

Comment thread charge_backend/retrosynthesis/ai.py
Comment thread charge_backend/retrosynthesis/ai.py Outdated
Comment thread charge_backend/backend_helper_funcs.py Outdated
Comment thread charge_backend/backend_manager.py Outdated
Comment thread charge_backend/backend_manager.py Outdated
Comment thread charge_backend/backend_manager.py Outdated
Comment thread charge_backend/charge_server.py
Comment thread flask-app/src/components/combined_customization_modal.tsx Outdated
Comment thread flask-app/package-lock.json
@bvanessen
bvanessen force-pushed the feature/use-charge-rsa branch 2 times, most recently from 2c54c24 to 29864d7 Compare May 12, 2026 21:09
@shashankk42
shashankk42 force-pushed the feature/use-charge-rsa branch from 799b869 to 8921db6 Compare May 19, 2026 07:57
@shashankk42
shashankk42 force-pushed the feature/use-charge-rsa branch from 0ffd898 to 8921db6 Compare May 19, 2026 22:34

@tbennun tbennun left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

mostly minor comments

Comment thread charge_backend/retrosynthesis/ai.py Outdated
Comment thread charge_backend/retrosynthesis/ai_rsa.py Outdated
Comment on lines +268 to +273
task=None,
agent_name=f"retrosynth_{node_id}_proposal_{proposal_counter[0]}",
callback=independent_callback,
)

output, final_result = await rsa_task.run_rsa(

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

it's a bit strange that we extend Task but then end up not using the built in API to run a task with a created Agent. If we only extend Task by name, there is no reason to extend it. I think we should use it and run RSA somehow with the Task. Maybe Task can define a generic run function that in the superclass is just agent.run(instructions=self.system_prompt, task=self.user_prompt) etc. and then RSATask's run will execute run_rsa with the agents?

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Addressed in 835e21e (here) + FLASK-LLNL/ChARGe#124 (deacaa0). Task now has a generic run(agent, ...); RSATask.run overrides to delegate to run_rsa. RSA is now invoked via the Task API in ai_rsa.py. The N-K-T loop and run_rsa are unchanged for backward compatibility. ChARGe#124 to be approved first.

Comment thread charge_backend/backend_helper_funcs.py Outdated
Comment thread charge_backend/charge_server.py Outdated
Comment on lines +326 to +344
uvicorn.run(
app,
host=host,
port=args.port,
ws_ping_timeout=60.0, # Increase websocket ping timeout to 60 seconds
timeout_keep_alive=75, # Keep connections alive for 75 seconds
)
uvicorn.run(app, host=host, port=args.port)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why were these kwargs removed?

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Re-added: Previous commit by BVE.

shashankk42 pushed a commit to shashankk42/flask-copilot that referenced this pull request May 23, 2026
…docstring

- ai.py: remove stale RSAAggregationTask comment
- backend_helper_funcs.py: drop PR FLASK-LLNL#194 quote from build_root_node docstring
Comment thread charge_backend/charge_server.py Outdated
Comment thread charge_backend/charge_server.py Outdated
shashankk42 pushed a commit to shashankk42/flask-copilot that referenced this pull request May 23, 2026
- Restore backend default to livai (PR FLASK-LLNL#193 default)
- Restore upstream MCP endpoint handlers (check-mcp-servers and
  registered-mcp-servers) and manual_mcp_servers strip/filter,
  preserving wormhole bearer-token forwarding from PR FLASK-LLNL#193
- Drop the local CheckServersRequest model that shadowed the upstream
  version in lc_conductor.tool_registration

File is now byte-identical to upstream main.
@shashankk42
shashankk42 force-pushed the feature/use-charge-rsa branch from 3c10687 to 94e95af Compare June 1, 2026 02:55
@shashankk42
shashankk42 requested review from bvanessen and tbennun June 1, 2026 20:19
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants