Case-Based Plan Reuse for AI Agents

from dataclasses import dataclass
from math import exp


@dataclass
class PlanTrace:
    trace_id: str
    goal_text: str
    plan: list[str]
    avg_score: float
    created_at_seconds: float
    reuse_count: int = 0
    attempt_count: int = 0


class ScorableRanker:
    def __init__(self, embedding_store, weights=None):
        self.embedding_store = embedding_store
        self.weights = weights or {
            "similarity": 0.4,
            "reward": 0.3,
            "recency": 0.2,
            "adaptability": 0.1,
        }

    def rank(self, goal_text, traces, now_seconds):
        goal_emb = self.embedding_store.get_or_create(goal_text)
        ranked = []
        for trace in traces:
            trace_emb = self.embedding_store.get_or_create(trace.goal_text + "\n" + "\n".join(trace.plan))
            components = {
                "similarity": self.embedding_store.cosine(goal_emb, trace_emb),
                "reward": trace.avg_score,
                "recency": exp(-(now_seconds - trace.created_at_seconds) / (30 * 24 * 60 * 60)),
                "adaptability": trace.reuse_count / max(trace.attempt_count, 1),
            }
            score = sum(components[k] * self.weights[k] for k in self.weights)
            ranked.append((score, components, trace))
        return sorted(ranked, key=lambda item: item[0], reverse=True)

class PlannerReuseAgent:
    def __init__(self, memory, ranker, llm, top_k=3):
        self.memory = memory
        self.ranker = ranker
        self.llm = llm
        self.top_k = top_k

    async def run(self, context):
        goal_text = context["goal"]["goal_text"]
        all_traces = self.memory.plan_traces.get_all(limit=500)
        ranked = self.ranker.rank(goal_text, all_traces, context["now_seconds"])
        selected = [trace for _score, _parts, trace in ranked[: self.top_k]]

        prompt = self._adaptation_prompt(goal_text, selected)
        new_plan = self.llm.generate_plan(prompt)
        new_trace_id = self.memory.plan_traces.create(goal_text=goal_text, plan=new_plan)

        for parent in selected:
            self.memory.plan_traces.add_reuse_link(
                parent_trace_id=parent.trace_id,
                child_trace_id=new_trace_id,
            )

        context["plan_trace_id"] = new_trace_id
        context["plan"] = new_plan
        context["reused_trace_ids"] = [trace.trace_id for trace in selected]
        return context

    def _adaptation_prompt(self, goal_text, traces):
        cases = "\n\n".join(
            f"Case {i+1}: {trace.goal_text}\nScore: {trace.avg_score}\nPlan:\n"
            + "\n".join(f"- {step}" for step in trace.plan)
            for i, trace in enumerate(traces)
        )
        return f"Adapt the useful patterns from these prior cases to solve:\n{goal_text}\n\n{cases}"

agent = PlannerReuseAgent(memory=memory, ranker=ScorableRanker(memory.embedding), llm=llm)
context = await agent.run({
    "goal": {"goal_text": "Build a RAG evaluator for a new paper corpus"},
    "now_seconds": time.time(),
})