llm-rsa 0.0.2__tar.gz → 0.0.3__tar.gz

llm_rsa-0.0.3/PKG-INFO

@@ -0,0 +1,287 @@
+Metadata-Version: 2.4
+Name: llm_rsa
+Version: 0.0.3
+Summary: Recursive self aggregation
+Home-page: https://github.com/risheekkumarb/llm_rsa
+Author: Risheek kumar B
+Author-email: b.risheekkumar@gmail.com
+License: Apache-2.0
+Keywords: nbdev jupyter notebook python
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: Natural Language :: English
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: fastcore
+Requires-Dist: litellm
+Requires-Dist: fastprogress
+Provides-Extra: dev
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: keywords
+Dynamic: license
+Dynamic: license-file
+Dynamic: provides-extra
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
+
+# RSA - Recursive Self-Aggregation
+
+
+<!-- WARNING: THIS FILE WAS AUTOGENERATED! DO NOT EDIT! -->
+
+RSA implements Recursive Self-Aggregation, a technique for improving LLM
+responses by generating multiple candidate answers and iteratively
+aggregating them. The algorithm samples k candidates from a pool of M
+responses, asks the LLM to synthesize an improved answer, and repeats
+this process across multiple loops to converge on higher-quality
+outputs.
+
+## Developer Guide
+
+If you are new to using `nbdev` here are some useful pointers to get you
+started.
+
+### Install in Development mode
+
+``` sh
+# make sure  package is installed in development mode
+$ pip install -e .
+
+# make changes under nbs/ directory
+# ...
+
+# compile to have changes apply to 
+$ nbdev_prepare
+```
+
+## Usage
+
+### Installation
+
+Install latest from the GitHub
+[repository](https://github.com/risheekkumarb/llm_rsa):
+
+``` sh
+$ pip install git+https://github.com//.git
+```
+
+or from [conda](https://anaconda.org/risheekkumarb/llm_rsa)
+
+``` sh
+$ conda install -c  
+```
+
+or from [pypi](https://pypi.org/project/llm_rsa)
+
+``` sh
+$ pip install 
+```
+
+### Documentation
+
+Documentation can be found hosted on this GitHub
+[repository](https://github.com/risheekkumarb/llm_rsa)’s
+[pages](https://risheekkumarb.github.io/llm_rsa/). Additionally you can
+find package manager specific guidelines on
+[conda](https://anaconda.org/risheekkumarb/llm_rsa) and
+[pypi](https://pypi.org/project/llm_rsa) respectively.
+
+## How to use
+
+### Basic Usage
+
+Create an RSA instance with your task prompt and call it to run the
+aggregation:
+
+``` python
+task_prompt = '''Three people check into a hotel room that costs $30. They each contribute $10. 
+Later, the manager realizes the room only costs $25 and gives $5 to the bellboy to return. 
+The bellboy keeps $2 and gives $1 back to each person. 
+So each person paid $9 (total $27), plus the bellboy has $2, which equals $29. 
+Where did the extra dollar go?'''
+```
+
+``` python
+agg_prompt = """Below is a reasoning problem followed by several candidate solutions. 
+Your job is to:
+1. Carefully analyze each candidate's reasoning step-by-step
+2. Identify which candidates make logical errors or arithmetic mistakes  
+3. Note which approaches lead to correct reasoning
+4. Synthesize the best reasoning into a single, clear, correct solution
+
+Show your work step-by-step, then state your final answer clearly."""
+```
+
+``` python
+from llm_rsa.core import RSA
+
+# Create RSA instance with a reasoning task
+rsa = RSA(
+    task_prompt=task_prompt,
+    agg_prompt=agg_prompt, 
+    N=4,
+    K=2,
+    loops=2
+)
+
+# Run the aggregation
+results = rsa.run()
+print(f"Generated {len(rsa.history)} total candidates across {rsa.loops} loops")
+print('llm response: \n', results[-1].response)
+```
+
+<style>
+    progress { appearance: none; border: none; border-radius: 4px; width: 300px;
+        height: 20px; vertical-align: middle; background: #e0e0e0; }
+&#10;    progress::-webkit-progress-bar { background: #e0e0e0; border-radius: 4px; }
+    progress::-webkit-progress-value { background: #2196F3; border-radius: 4px; }
+    progress::-moz-progress-bar { background: #2196F3; border-radius: 4px; }
+&#10;    progress:not([value]) {
+        background: repeating-linear-gradient(45deg, #7e7e7e, #7e7e7e 10px, #5c5c5c 10px, #5c5c5c 20px); }
+&#10;    progress.progress-bar-interrupted::-webkit-progress-value { background: #F44336; }
+    progress.progress-bar-interrupted::-moz-progress-value { background: #F44336; }
+    progress.progress-bar-interrupted::-webkit-progress-bar { background: #F44336; }
+    progress.progress-bar-interrupted::-moz-progress-bar { background: #F44336; }
+    progress.progress-bar-interrupted { background: #F44336; }    
+&#10;    table.fastprogress { border-collapse: collapse; margin: 1em 0; font-size: 0.9em; }
+    table.fastprogress th, table.fastprogress td { padding: 8px 12px; border: 1px solid #ddd; text-align: left; }
+    table.fastprogress thead tr { background: #f8f9fa; font-weight: bold; }
+    table.fastprogress tbody tr:nth-of-type(even) { background: #f8f9fa; }
+</style>
+
+``` html
+<div>
+<progress max="2" value="2"></progress> 100.00% [2/2 00:20&lt;00:00... Loop 2]</div>
+```
+
+    Generated 8 total candidates across 2 loops
+    llm response: 
+     ### Analysis of Candidate Reasoning
+
+    Both **Candidate 1** and **Candidate 2** provide excellent, accurate explanations of the "missing dollar" riddle. 
+
+    1.  **Candidate 1 Analysis:** This candidate correctly identifies the logical fallacy of adding the bellboy's kept money to the amount spent by the guests. They provide a clear "Follow the Money" breakdown showing that the $30 is distributed as $25 (hotel), $2 (bellboy), and $3 (guests). They correctly state that the $27 spent by the guests *already includes* the $2 held by the bellboy.
+    2.  **Candidate 2 Analysis:** This candidate identifies the "false premise" and "incorrect logic" of the riddle. Like Candidate 1, they break down the $30 correctly and explain that the riddle incorrectly adds a "cost" ($27) to a "profit" ($2) instead of adding the "cost" ($27) to the "refund" ($3).
+
+    Both candidates conclude that no money is actually missing and that the riddle is based on an arithmetic trick.
+
+    ---
+
+    ### Synthetic Correct Solution
+
+    The "missing dollar" is a result of a misleading calculation. To resolve the mystery, we must track the money accurately using two different perspectives: **where the money is now** and **the total amount spent vs. kept.**
+
+    #### 1. Where is the money now? (The $30 Breakdown)
+    The original $30 can be accounted for by looking at who currently holds the cash:
+    *   **$25:** Held by the hotel (the actual cost of the room).
+    *   **$2:** Held by the bellboy (the amount he kept).
+    *   **$3:** Held by the three guests ($1 each in their pockets).
+    *   **Total: $25 + $2 + $3 = $30.** 
+    Nothing is missing.
+
+    #### 2. The Fallacy in the Riddle
+    The riddle states: *"Each person paid $9 (total $27), plus the bellboy has $2, which equals $29."*
+    This is logically incorrect because it **double-counts** the bellboy's money. 
+
+    *   **The Net Payment:** The guests spent a total of **$27**. 
+    *   **The Destination of that payment:** Of that $27, **$25** went to the hotel and **$2** went to the bellboy.
+    *   **The Calculation:** Adding the $2 to the $27 is nonsensical because the $2 is *already part* of the $27. 
+
+    To reach the original $30, you must add the money the guests **kept** to the money they **spent**:
+    **$27 (Spent) + $3 (Returned to them) = $30.**
+
+    ### Final Answer
+    The dollar did not go anywhere. The riddle creates an illusion by adding the bellboy's $2 to the $27 spent, when it should be subtracting the $2 from the $27 to find the hotel’s $25, or adding the $3 refund to the $27 to find the original $30.
+
+``` python
+from pydantic import BaseModel
+class Answer(BaseModel):
+    answer: str
+    confidence: float
+
+prompt, response = rsa.aggregate(response_model=Answer)
+print(response)
+```
+
+    {"answer":"The mystery of the missing dollar is caused by a logical fallacy known as misdirection. The riddle incorrectly adds the bellboy's $2 to the $27 paid by the guests, creating a mathematically irrelevant number ($29). To solve the puzzle, we simply need to track the original $30 using two balance methods:\n\n1. The Distribution Method (Where is the money now?):\n- $25 is in the hotel's register.\n- $2 is in the bellboy's pocket.\n- $3 is in the guests' pockets ($1 each).\n- Total: $25 + $2 + $3 = $30. \nEverything is accounted for.\n\n2. The Net Expenditure Method (What did the guests pay?):\nThe guests paid $30 and got $3 back, meaning they spent exactly $27. \n- $25 of that $27 went to the hotel room cost.\n- $2 of that $27 went to the bellboy as a tip.\n- Total: $25 + $2 = $27.\n\nThe error in the riddle is adding the $2 to the $27. Because the $2 is already part of the $27, adding them together double-counts the bellboy's tip. To get back to the original $30, you must add the $27 spent to the $3 refund ($27 + $3 = $30). There is no missing dollar.","confidence":1.0}
+
+``` python
+from litellm import completion
+
+# Single direct call (baseline)
+response = completion(
+    model='openrouter/google/gemini-3-flash-preview',
+    messages=[{"role": "user", "content": task_prompt}],
+    temperature=1.0
+)
+baseline_answer = response.choices[0].message.content
+print("=== BASELINE (single call) ===")
+print(baseline_answer)
+```
+
+    === BASELINE (single call) ===
+    This is a classic riddle that relies on a **logical fallacy**—specifically, an error in how the numbers are added together at the end.
+
+    The "lost" dollar doesn't exist; it only appears to be missing because the math at the end of the story adds two numbers that should actually be **subtracted**.
+
+    Here is the correct breakdown of the money:
+
+    ### 1. Follow the Money
+    Instead of adding the bellboy's tip to the guests' payment, look at where the original $30 is at the very end:
+    *   **$25** is in the hotel cash register.
+    *   **$2** is in the bellboy's pocket.
+    *   **$3** was returned to the guests ($1 each).
+    *   **Total: $25 + $2 + $3 = $30.** (The math is perfect).
+
+    ### 2. The Flaw in the Riddle
+    The riddle says: *"Each person paid $9 (total $27), plus the bellboy has $2, which equals $29."* 
+
+    **The error is adding the $2 to the $27.** 
+    The $27 that the guests spent **already includes** the $2 that the bellboy took.
+
+    Think of it this way:
+    *   The guests paid **$27**.
+    *   Where did that $27 go? **$25** went to the hotel and **$2** went to the bellboy.
+    *   To reach the original $30, you should add the **$3** they got back, not the $2 the bellboy kept.
+
+    **The correct equation is:**
+    $27 (Paid) + $3 (Refund) = $30. 
+    *OR*
+    $27 (Paid) - $2 (Bellboy's Tip) = $25 (Room Cost).
+
+### Configuration Options
+
+| Parameter | Default | Description |
+|----|----|----|
+| `task_prompt` | (required) | The main task/question to solve |
+| `model` | `'openrouter/google/gemini-3-flash-preview'` | LLM model to use (any litellm-compatible model) |
+| `N` | 4 | Population size (candidates per loop) |
+| `K` | 3 | Number of candidates to aggregate |
+| `loops` | 2 | Number of aggregation iterations |
+| `temperature` | 1.0 | LLM sampling temperature |
+| `n_workers` | 4 | Parallel workers for LLM calls |
+| `agg_prompt` | (auto) | Custom aggregation prompt (optional) |
+
+### How RSA Works
+
+1.  **Loop 0**: Generate N independent responses to the task prompt
+2.  **Loop 1+**: For each of N new candidates, randomly sample K
+    previous candidates and ask the LLM to aggregate them into an
+    improved answer
+3.  **Repeat** for the specified number of loops
+4.  **Return** the final pool of aggregated candidates
+
+The `history` attribute stores all candidates across all loops, allowing
+you to trace the aggregation process.

llm_rsa-0.0.3/README.md

@@ -0,0 +1,249 @@
+# RSA - Recursive Self-Aggregation
+
+
+<!-- WARNING: THIS FILE WAS AUTOGENERATED! DO NOT EDIT! -->
+
+RSA implements Recursive Self-Aggregation, a technique for improving LLM
+responses by generating multiple candidate answers and iteratively
+aggregating them. The algorithm samples k candidates from a pool of M
+responses, asks the LLM to synthesize an improved answer, and repeats
+this process across multiple loops to converge on higher-quality
+outputs.
+
+## Developer Guide
+
+If you are new to using `nbdev` here are some useful pointers to get you
+started.
+
+### Install in Development mode
+
+``` sh
+# make sure  package is installed in development mode
+$ pip install -e .
+
+# make changes under nbs/ directory
+# ...
+
+# compile to have changes apply to 
+$ nbdev_prepare
+```
+
+## Usage
+
+### Installation
+
+Install latest from the GitHub
+[repository](https://github.com/risheekkumarb/llm_rsa):
+
+``` sh
+$ pip install git+https://github.com//.git
+```
+
+or from [conda](https://anaconda.org/risheekkumarb/llm_rsa)
+
+``` sh
+$ conda install -c  
+```
+
+or from [pypi](https://pypi.org/project/llm_rsa)
+
+``` sh
+$ pip install 
+```
+
+### Documentation
+
+Documentation can be found hosted on this GitHub
+[repository](https://github.com/risheekkumarb/llm_rsa)’s
+[pages](https://risheekkumarb.github.io/llm_rsa/). Additionally you can
+find package manager specific guidelines on
+[conda](https://anaconda.org/risheekkumarb/llm_rsa) and
+[pypi](https://pypi.org/project/llm_rsa) respectively.
+
+## How to use
+
+### Basic Usage
+
+Create an RSA instance with your task prompt and call it to run the
+aggregation:
+
+``` python
+task_prompt = '''Three people check into a hotel room that costs $30. They each contribute $10. 
+Later, the manager realizes the room only costs $25 and gives $5 to the bellboy to return. 
+The bellboy keeps $2 and gives $1 back to each person. 
+So each person paid $9 (total $27), plus the bellboy has $2, which equals $29. 
+Where did the extra dollar go?'''
+```
+
+``` python
+agg_prompt = """Below is a reasoning problem followed by several candidate solutions. 
+Your job is to:
+1. Carefully analyze each candidate's reasoning step-by-step
+2. Identify which candidates make logical errors or arithmetic mistakes  
+3. Note which approaches lead to correct reasoning
+4. Synthesize the best reasoning into a single, clear, correct solution
+
+Show your work step-by-step, then state your final answer clearly."""
+```
+
+``` python
+from llm_rsa.core import RSA
+
+# Create RSA instance with a reasoning task
+rsa = RSA(
+    task_prompt=task_prompt,
+    agg_prompt=agg_prompt, 
+    N=4,
+    K=2,
+    loops=2
+)
+
+# Run the aggregation
+results = rsa.run()
+print(f"Generated {len(rsa.history)} total candidates across {rsa.loops} loops")
+print('llm response: \n', results[-1].response)
+```
+
+<style>
+    progress { appearance: none; border: none; border-radius: 4px; width: 300px;
+        height: 20px; vertical-align: middle; background: #e0e0e0; }
+&#10;    progress::-webkit-progress-bar { background: #e0e0e0; border-radius: 4px; }
+    progress::-webkit-progress-value { background: #2196F3; border-radius: 4px; }
+    progress::-moz-progress-bar { background: #2196F3; border-radius: 4px; }
+&#10;    progress:not([value]) {
+        background: repeating-linear-gradient(45deg, #7e7e7e, #7e7e7e 10px, #5c5c5c 10px, #5c5c5c 20px); }
+&#10;    progress.progress-bar-interrupted::-webkit-progress-value { background: #F44336; }
+    progress.progress-bar-interrupted::-moz-progress-value { background: #F44336; }
+    progress.progress-bar-interrupted::-webkit-progress-bar { background: #F44336; }
+    progress.progress-bar-interrupted::-moz-progress-bar { background: #F44336; }
+    progress.progress-bar-interrupted { background: #F44336; }    
+&#10;    table.fastprogress { border-collapse: collapse; margin: 1em 0; font-size: 0.9em; }
+    table.fastprogress th, table.fastprogress td { padding: 8px 12px; border: 1px solid #ddd; text-align: left; }
+    table.fastprogress thead tr { background: #f8f9fa; font-weight: bold; }
+    table.fastprogress tbody tr:nth-of-type(even) { background: #f8f9fa; }
+</style>
+
+``` html
+<div>
+<progress max="2" value="2"></progress> 100.00% [2/2 00:20&lt;00:00... Loop 2]</div>
+```
+
+    Generated 8 total candidates across 2 loops
+    llm response: 
+     ### Analysis of Candidate Reasoning
+
+    Both **Candidate 1** and **Candidate 2** provide excellent, accurate explanations of the "missing dollar" riddle. 
+
+    1.  **Candidate 1 Analysis:** This candidate correctly identifies the logical fallacy of adding the bellboy's kept money to the amount spent by the guests. They provide a clear "Follow the Money" breakdown showing that the $30 is distributed as $25 (hotel), $2 (bellboy), and $3 (guests). They correctly state that the $27 spent by the guests *already includes* the $2 held by the bellboy.
+    2.  **Candidate 2 Analysis:** This candidate identifies the "false premise" and "incorrect logic" of the riddle. Like Candidate 1, they break down the $30 correctly and explain that the riddle incorrectly adds a "cost" ($27) to a "profit" ($2) instead of adding the "cost" ($27) to the "refund" ($3).
+
+    Both candidates conclude that no money is actually missing and that the riddle is based on an arithmetic trick.
+
+    ---
+
+    ### Synthetic Correct Solution
+
+    The "missing dollar" is a result of a misleading calculation. To resolve the mystery, we must track the money accurately using two different perspectives: **where the money is now** and **the total amount spent vs. kept.**
+
+    #### 1. Where is the money now? (The $30 Breakdown)
+    The original $30 can be accounted for by looking at who currently holds the cash:
+    *   **$25:** Held by the hotel (the actual cost of the room).
+    *   **$2:** Held by the bellboy (the amount he kept).
+    *   **$3:** Held by the three guests ($1 each in their pockets).
+    *   **Total: $25 + $2 + $3 = $30.** 
+    Nothing is missing.
+
+    #### 2. The Fallacy in the Riddle
+    The riddle states: *"Each person paid $9 (total $27), plus the bellboy has $2, which equals $29."*
+    This is logically incorrect because it **double-counts** the bellboy's money. 
+
+    *   **The Net Payment:** The guests spent a total of **$27**. 
+    *   **The Destination of that payment:** Of that $27, **$25** went to the hotel and **$2** went to the bellboy.
+    *   **The Calculation:** Adding the $2 to the $27 is nonsensical because the $2 is *already part* of the $27. 
+
+    To reach the original $30, you must add the money the guests **kept** to the money they **spent**:
+    **$27 (Spent) + $3 (Returned to them) = $30.**
+
+    ### Final Answer
+    The dollar did not go anywhere. The riddle creates an illusion by adding the bellboy's $2 to the $27 spent, when it should be subtracting the $2 from the $27 to find the hotel’s $25, or adding the $3 refund to the $27 to find the original $30.
+
+``` python
+from pydantic import BaseModel
+class Answer(BaseModel):
+    answer: str
+    confidence: float
+
+prompt, response = rsa.aggregate(response_model=Answer)
+print(response)
+```
+
+    {"answer":"The mystery of the missing dollar is caused by a logical fallacy known as misdirection. The riddle incorrectly adds the bellboy's $2 to the $27 paid by the guests, creating a mathematically irrelevant number ($29). To solve the puzzle, we simply need to track the original $30 using two balance methods:\n\n1. The Distribution Method (Where is the money now?):\n- $25 is in the hotel's register.\n- $2 is in the bellboy's pocket.\n- $3 is in the guests' pockets ($1 each).\n- Total: $25 + $2 + $3 = $30. \nEverything is accounted for.\n\n2. The Net Expenditure Method (What did the guests pay?):\nThe guests paid $30 and got $3 back, meaning they spent exactly $27. \n- $25 of that $27 went to the hotel room cost.\n- $2 of that $27 went to the bellboy as a tip.\n- Total: $25 + $2 = $27.\n\nThe error in the riddle is adding the $2 to the $27. Because the $2 is already part of the $27, adding them together double-counts the bellboy's tip. To get back to the original $30, you must add the $27 spent to the $3 refund ($27 + $3 = $30). There is no missing dollar.","confidence":1.0}
+
+``` python
+from litellm import completion
+
+# Single direct call (baseline)
+response = completion(
+    model='openrouter/google/gemini-3-flash-preview',
+    messages=[{"role": "user", "content": task_prompt}],
+    temperature=1.0
+)
+baseline_answer = response.choices[0].message.content
+print("=== BASELINE (single call) ===")
+print(baseline_answer)
+```
+
+    === BASELINE (single call) ===
+    This is a classic riddle that relies on a **logical fallacy**—specifically, an error in how the numbers are added together at the end.
+
+    The "lost" dollar doesn't exist; it only appears to be missing because the math at the end of the story adds two numbers that should actually be **subtracted**.
+
+    Here is the correct breakdown of the money:
+
+    ### 1. Follow the Money
+    Instead of adding the bellboy's tip to the guests' payment, look at where the original $30 is at the very end:
+    *   **$25** is in the hotel cash register.
+    *   **$2** is in the bellboy's pocket.
+    *   **$3** was returned to the guests ($1 each).
+    *   **Total: $25 + $2 + $3 = $30.** (The math is perfect).
+
+    ### 2. The Flaw in the Riddle
+    The riddle says: *"Each person paid $9 (total $27), plus the bellboy has $2, which equals $29."* 
+
+    **The error is adding the $2 to the $27.** 
+    The $27 that the guests spent **already includes** the $2 that the bellboy took.
+
+    Think of it this way:
+    *   The guests paid **$27**.
+    *   Where did that $27 go? **$25** went to the hotel and **$2** went to the bellboy.
+    *   To reach the original $30, you should add the **$3** they got back, not the $2 the bellboy kept.
+
+    **The correct equation is:**
+    $27 (Paid) + $3 (Refund) = $30. 
+    *OR*
+    $27 (Paid) - $2 (Bellboy's Tip) = $25 (Room Cost).
+
+### Configuration Options
+
+| Parameter | Default | Description |
+|----|----|----|
+| `task_prompt` | (required) | The main task/question to solve |
+| `model` | `'openrouter/google/gemini-3-flash-preview'` | LLM model to use (any litellm-compatible model) |
+| `N` | 4 | Population size (candidates per loop) |
+| `K` | 3 | Number of candidates to aggregate |
+| `loops` | 2 | Number of aggregation iterations |
+| `temperature` | 1.0 | LLM sampling temperature |
+| `n_workers` | 4 | Parallel workers for LLM calls |
+| `agg_prompt` | (auto) | Custom aggregation prompt (optional) |
+
+### How RSA Works
+
+1.  **Loop 0**: Generate N independent responses to the task prompt
+2.  **Loop 1+**: For each of N new candidates, randomly sample K
+    previous candidates and ask the LLM to aggregate them into an
+    improved answer
+3.  **Repeat** for the specified number of loops
+4.  **Return** the final pool of aggregated candidates
+
+The `history` attribute stores all candidates across all loops, allowing
+you to trace the aggregation process.

{llm_rsa-0.0.2 → llm_rsa-0.0.3}/llm_rsa/core.py

@@ -43,12 +43,14 @@ class RSA:
 
 # %% ../nbs/00_core.ipynb #3c39e9e6
 @patch
-def _call_llm(self:RSA, prompt):
+def _call_llm(self:RSA, prompt, **kwargs):
     "Call the LLM with the given prompt and return the response content"
     response = completion(
         model=self.model,
         messages=[{"role": "user", "content": prompt}],
         temperature=self.temperature,
+        num_retries=3,
+        **kwargs
     )
     return response.choices[0].message.content
 
@@ -97,11 +99,13 @@ def run(self:RSA):
         self.history.extend(pool)
     return pool
 
-# %% ../nbs/00_core.ipynb #b4edf71c
+# %% ../nbs/00_core.ipynb #72cb9b61
 @patch
-def aggregate(self:RSA):
-    "Final aggregation one LLM call to aggregate all the final loop candidates"
+def aggregate(self:RSA, agg_prompt=None, response_model=None):
+    "Final aggregation: one LLM call to aggregate all final loop candidates, with optional structured output"
+    agg_prompt = agg_prompt or self.agg_prompt
     candidates = self.history.filter(lambda x: x.loop_id==(self.loops-1))
     responses = '\n'.join(f"---- Candidate {i+1} ----\n{c.response}" for i, c in enumerate(candidates))
-    prompt = f"{self.agg_prompt}\n{self.task_prompt}\n\nCANDIDATE ANSWERS:\n{responses}\n\nProvide the best aggregated answer:"
-    return prompt, self._call_llm(prompt)
+    prompt = f"{agg_prompt}\n{self.task_prompt}\n\nCANDIDATE ANSWERS:\n{responses}\n\nProvide the best aggregated answer:"
+    result = self._call_llm(prompt, **({'response_format': response_model} if response_model else {}))
+    return prompt, result