"""Recursive Self-Aggregation (RSA) - A general-purpose LLM aggregation algorithm using litellm based on the paper **https://rsa-llm.github.io/**"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/00_core.ipynb.

# %% auto #0
__all__ = ['RSACandidate', 'RSA']

# %% ../nbs/00_core.ipynb #ccd7cb56
from fastcore.all import *
from fastcore.test import *
from litellm import completion
import random, uuid
from itertools import combinations
from fastprogress import progress_bar

# %% ../nbs/00_core.ipynb #c8def011
class RSACandidate:
    "A candidate response in the RSA algorithm"
    def __init__(self, id:str, loop_id:int, prompt:str, response:str=None, parent_ids:list=None): store_attr()
    def __repr__(self): return f'id:{self.id}\nloop_id:{self.loop_id}\nprompt:\n{self.prompt}\nresponse:\n{self.response}\nparent_ids:\n{self.parent_ids}'

# %% ../nbs/00_core.ipynb #cc5af69e
class RSA:
    "Recursive Self-Aggregation algorithm for LLM response aggregation"
    def __init__(
        self,
        task_prompt:str,  # The main task/question to solve
        agg_prompt:str=None,  # Custom aggregation prompt
        model:str='openrouter/google/gemini-3-flash-preview',  # LLM model to use
        N:int=4,  # Population size (candidates per loop)
        K:int=3,  # Number of candidates to aggregate
        loops:int=2,  # Number of aggregation loops
        history:list=None,  # History of all candidates
        temperature:float=1.0,  # LLM temperature
        n_workers:int=4  # Parallel workers
    ): 
        if not task_prompt: raise ValueError("task_prompt is required")
        store_attr()
        if not history: self.history = L()
        if not self.agg_prompt: self.agg_prompt = """You are given question with training examples and a test input.\nYou are also provided several candidate solutions. Some candidates may be incorrect\nAggregate/consider all the candidates and use their help to produce the improved correct solution"""
    
    def __repr__(self): return f'RSA(model={self.model!r}, \nN={self.N}, \nK={self.K}, \nloops={self.loops}, \nhistory={len(self.history)} candidates, \ntask_prompt={self.task_prompt})'

# %% ../nbs/00_core.ipynb #3c39e9e6
@patch
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

# %% ../nbs/00_core.ipynb #ef493c51
@patch
def _agg_prompt(self:RSA, candidates: list[RSACandidate]) -> str:
    "Build an aggregation prompt combining the task prompt with candidate responses"
    parts = [
        self.agg_prompt,
        self.task_prompt,
        "\nCANDIDATE ANSWERS (may contain mistakes):",
    ]
    for i, cand in enumerate(candidates, 1):
        parts.append(f"---- Candidate {i} ----\n{cand.response}")
    parts.append("\nYour response:")
    return "\n".join(parts)

# %% ../nbs/00_core.ipynb #81deb527
@patch
def get_prompts(self:RSA, loop_id, cands=None):
    "Generate candidate prompts for a given loop: N initial candidates, or all C(n,K) combinations for aggregation"
    if not cands: return L(RSACandidate(id=str(uuid.uuid4()), loop_id=loop_id, prompt=self.task_prompt) for _ in range(self.N))
    if len(cands) < self.K: raise ValueError(f"Need at least {self.K} candidates, got {len(cands)}")
    sel_cands = L(combinations(cands, self.K))
    if len(sel_cands) < self.N: raise ValueError(f"C({len(cands)},{self.K})={len(sel_cands)} combinations is less than N={self.N}")
    sel_cands = sel_cands.shuffle()[:self.N]
    return sel_cands.map(lambda x: RSACandidate(id=str(uuid.uuid4()), loop_id=loop_id, prompt=self._agg_prompt(x), parent_ids=L(x).attrgot('id')))

# %% ../nbs/00_core.ipynb #80329d80
@patch
def _run_loop(self:RSA, loop_id, pool=None):
    "Execute one loop: generate prompts, call LLM in parallel, attach responses"
    prompts = self.get_prompts(loop_id, pool)
    responses = parallel(self._call_llm, prompts.attrgot('prompt'), n_workers=self.n_workers, progress=True)
    for p, r in zip(prompts, responses): p.response = r
    return prompts

# %% ../nbs/00_core.ipynb #02599e94
@patch
def run(self:RSA):
    "Run the full RSA algorithm for the configured number of loops and return the final candidate pool"
    pbar = progress_bar(range(self.loops))
    for i in pbar:
        pbar.comment = f"Loop {i+1}"
        pool = self._run_loop(i, pool if i > 0 else None)
        self.history.extend(pool)
    return pool

# %% ../nbs/00_core.ipynb #72cb9b61
@patch
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"{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
