Add generic RSA (Recursive Self-Aggregation) algorithm framework by shashankk42 · Pull Request #124 · FLASK-LLNL/ChARGe

shashankk42 · 2026-04-12T08:59:17Z

Add generic Recursive Self-Aggregation (RSA) algorithm as a Task subclass. Domain-agnostic by default; subclass to customize
for any prompt-and-aggregate problem.

Changes

Add charge/algorithms/rsa.py with RSATask(Task) exposing the N-K-T loop via an async run_rsa() method
Hook methods on RSATask for domain customization (defaults work without subclassing):
- format_candidates(subset) : how candidates are rendered for aggregation
- validate_proposal(result) : accept/reject a proposal
- build_proposal_task() / build_aggregation_task(text, step) — per-stage Task construction
3-part prompt structure carried on RSATask:
- system_prompt: domain expert (proposals and aggregations)
- proposal_prompt: instructions for generating a single solution
- aggregation_prompt: template for synthesizing K-subsets ({original_prompt}, {candidates}, {step}, {total_steps})
Module-level helpers exported for one-off Task construction without the loop:
- create_default_proposal_task(...), create_default_aggregation_task(...), default_format_candidates(...),
  GenericRSAOutput
Default prompts in charge/algorithms/prompts/ (proposal system, aggregation template)
README rewritten with a minimal example + cross-domain example (non-retrosynthesis) demonstrating that no subclassing is
required

API

task = RSATask(n=4, k=2, t=2, user_prompt="...")
output, result = await task.run_rsa(runner, runner_factory=...)

Subclass for domain customization:
class MyDomainRSATask(RSATask):
    def format_candidates(self, subset): ...
    def validate_proposal(self, result): ...

Testing

N/K/T validation rejects values < 2 and K > N with messages naming the violator(s)
Default hooks introspect any Pydantic schema (verified with GenericRSAOutput and a custom schema)
Bearer token / server_urls / builtin_tools correctly propagated to every per-proposal and per-aggregation Task

Related PRs

Integrate ChARGe RSA algorithm flask-copilot#194 : consumes RSATask via RetroRSATask subclass
Add RSA settings LC-Conductor#15 : provides the RSA settings UI panel

tbennun

Some comments on the code. I think this PR needs to be refactored to the Task interface as we discussed. Maybe even have FLASK-Copilot use that interface directly with its prompts.

tbennun · 2026-05-12T18:49:41Z

+_DEFAULT_AGGREGATION_TEMPLATE = _PROMPTS_DIR / "default_aggregation_template.txt"
+
+
+# Generic default output schema


This code and all the classes seem quite excessive. Why do we need a class for RSAPrompts? If RSA extends a Task we should add those fields there?

tbennun · 2026-05-12T18:49:58Z

+                )
+
+
+class RSACallbacks:


unnecessary to put in a separate class

tbennun · 2026-05-12T18:51:16Z

+    return task
+
+
+class RSATaskFactories(Generic[T]):


Why make a generic RSA task factory? What is the need for the factory design pattern? I thought that there should just be an RSATask that extends Task

tbennun · 2026-05-12T18:52:02Z

+    k = config.k
+
+    # Helper function to run a single proposal
+    async def run_single_proposal(proposal_index: int, proposal_runner: Any):


Long nested functions... Hard to understand the flow.

tbennun · 2026-05-12T18:52:13Z

+    await callbacks.logger_info(f"Generated {len(proposals)} valid proposals")
+
+    # Helper function to run a single aggregation
+    async def run_single_aggregation(


same, please refactor out

shashankk42 · 2026-05-23T08:17:43Z

Pushed deacaa0 after the approval to address review feedback on the consuming PR (FLASK-LLNL/flask-copilot#194, comment from @tbennun): RSATask was extending Task without actually using its API. This adds a generic Task.run(agent, ...) to the base class and an override on RSATask that delegates to the existing run_rsa (preserved for backward compatibility). No behavior change. RSA still runs the same N-K-T loop. Flagging in case you want to take another look before merge.

Create charge/algorithms/ with RSAConfig, RSACallbacks, RSATaskFactories to make RSA reusable for any task type.

Add task-agnostic default prompts that users can swap out for domain-specific needs.

Document how to use default prompts and swap them for domain-specific needs.

Add GenericRSAOutput schema, default_format_candidates, and default task factories so RSA works immediately without customization.

Allows domain-specific prompts to be passed to aggregation tasks, enabling clean separation between generic RSA and domain customization.

- system_prompt: Domain expert definition (constant across proposal and aggregation) - proposal_prompt: Task instructions for generating proposals - aggregation_prompt: Task instructions for evaluating/aggregating This provides cleaner separation: expertise (system) vs task (user prompt).

- Add minimal working example using ChARGe only - Document 3-part prompt structure (system, proposal, aggregation) - Remove stale 2-part prompt references - Clear progression from basic to advanced usage - Remove flask-copilot dependencies from examples

- Enforce N, K, T >= 2 (T=1 would skip aggregation) - Enforce K <= N at construction time in RSAConfig.__post_init__ - Remove redundant K validation in run_rsa_loop - Improve error messages with actionable guidance - Add logging when K is adjusted due to proposal failures

…sed.

Task gains a default async run(agent, *args, **kwargs) that assigns self to agent.task and dispatches agent.run, so callers can invoke any Task through a uniform API. RSATask overrides run() to delegate to the existing run_rsa N-K-T orchestrator, preserving the public run_rsa API for external consumers.

tbennun · 2026-06-01T18:05:33Z


        self.constructor_args = {}

+    async def run(self, agent, *args, **kwargs):


@bvanessen FYI, this change means we need to use Task.run() everywhere instead of agent.run(task).

Let's leave that for a follow up PR

shashankk42 mentioned this pull request Apr 12, 2026

Integrate ChARGe RSA algorithm FLASK-LLNL/flask-copilot#194

Open

shashankk42 marked this pull request as draft April 16, 2026 22:28

shashankk42 force-pushed the feature/rsa-generic-algorithm branch from 83e7a55 to cdc9fa2 Compare April 24, 2026 00:04

shashankk42 marked this pull request as ready for review April 24, 2026 01:50

shashankk42 marked this pull request as draft April 24, 2026 04:51

shashankk42 marked this pull request as ready for review April 24, 2026 05:02

shashankk42 marked this pull request as draft April 24, 2026 06:30

shashankk42 marked this pull request as ready for review April 24, 2026 06:34

bvanessen force-pushed the feature/rsa-generic-algorithm branch from 9b73603 to ffbe801 Compare May 1, 2026 20:47

bvanessen reviewed May 12, 2026

View reviewed changes

Comment thread charge/algorithms/rsa.py

tbennun requested changes May 12, 2026

View reviewed changes

bvanessen force-pushed the feature/rsa-generic-algorithm branch from 79c870e to 4435b50 Compare May 13, 2026 14:19

This was referenced May 19, 2026

RSA Support: Fix OpenAI client compatibility and session handling for RSA retrosynthesis #120

Closed

Add RSA settings FLASK-LLNL/LC-Conductor#15

Open

tbennun approved these changes May 23, 2026

View reviewed changes

shashankk42 requested a review from tbennun May 23, 2026 08:18

Shashank Kushwaha and others added 13 commits May 30, 2026 03:21

Add generic RSA algorithm module to ChARGe

00be1c8

Create charge/algorithms/ with RSAConfig, RSACallbacks, RSATaskFactories to make RSA reusable for any task type.

Add default prompts and RSAPrompts for easy customization

db3eadc

Add task-agnostic default prompts that users can swap out for domain-specific needs.

Add README with usage examples for RSA prompts

0b47904

Document how to use default prompts and swap them for domain-specific needs.

Add production-ready defaults for out-of-box RSA usage

8f38d3e

Add GenericRSAOutput schema, default_format_candidates, and default task factories so RSA works immediately without customization.

Add system_prompt parameter to aggregation tasks

9d0563a

Allows domain-specific prompts to be passed to aggregation tasks, enabling clean separation between generic RSA and domain customization.

Added better error handling for when tasks are canceled.

742c274

Update doc string to indicate that a bearer_token can / should be pas…

46d75c0

…sed.

Apply suggestion from @bvanessen

8ca3d5a

Refactor RSA into RSATask(Task) per review

d15c5b1

shashankk42 force-pushed the feature/rsa-generic-algorithm branch from deacaa0 to 63f1fd0 Compare June 1, 2026 02:32

tbennun approved these changes Jun 1, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add generic RSA (Recursive Self-Aggregation) algorithm framework#124

Add generic RSA (Recursive Self-Aggregation) algorithm framework#124
shashankk42 wants to merge 13 commits into
FLASK-LLNL:mainfrom
shashankk42:feature/rsa-generic-algorithm

shashankk42 commented Apr 12, 2026 •

edited

Loading

Uh oh!

Uh oh!

tbennun left a comment •

edited

Loading

Uh oh!

tbennun May 12, 2026

Uh oh!

tbennun May 12, 2026

Uh oh!

tbennun May 12, 2026

Uh oh!

tbennun May 12, 2026

Uh oh!

tbennun May 12, 2026

Uh oh!

shashankk42 commented May 23, 2026

Uh oh!

tbennun Jun 1, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

		_DEFAULT_AGGREGATION_TEMPLATE = _PROMPTS_DIR / "default_aggregation_template.txt"


		# Generic default output schema


		self.constructor_args = {}

		async def run(self, agent, args, *kwargs):

Conversation

shashankk42 commented Apr 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Changes

API

Uh oh!

Uh oh!

tbennun left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tbennun May 12, 2026

Choose a reason for hiding this comment

Uh oh!

tbennun May 12, 2026

Choose a reason for hiding this comment

Uh oh!

tbennun May 12, 2026

Choose a reason for hiding this comment

Uh oh!

tbennun May 12, 2026

Choose a reason for hiding this comment

Uh oh!

tbennun May 12, 2026

Choose a reason for hiding this comment

Uh oh!

shashankk42 commented May 23, 2026

Uh oh!

tbennun Jun 1, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

shashankk42 commented Apr 12, 2026 •

edited

Loading

tbennun left a comment •

edited

Loading