Retrieval-Augmented Actions

Published:

Standard RAG retrieves knowledge — documents, facts, passages that answer a question. But agents have to act. And when the question is "what should I do next?", the most useful thing to retrieve is an example of someone (or something) successfully doing a similar task before.

This is retrieval-augmented action: instead of hard-coding a handful of few-shot examples in the system prompt, you dynamically retrieve relevant demonstrations from a library of past successful trajectories. The agent sees how a similar problem was solved, adapts the approach to the current situation, and acts accordingly. It is the difference between a cookbook with three recipes stapled to the fridge and a searchable recipe database that surfaces the most relevant one based on what ingredients you have.

Why Static Few-Shot Examples Break Down #

Few-shot prompting works. You show the model two or three examples of the input-output pattern you want, and it generalizes. For narrow tasks — "extract the date from this email," "classify this ticket as bug or feature" — a fixed set of examples in the system prompt is sufficient.

But agents face a combinatorial problem. A coding agent might encounter thousands of distinct task shapes: "add a REST endpoint," "fix a race condition," "migrate a database schema," "refactor a class hierarchy." A customer-support agent handles billing disputes, password resets, refund requests, escalation paths, and dozens of edge cases. You cannot fit representative examples of every task type into a single prompt.

Static few-shot examples fail in three ways:

  • Coverage gaps. The fixed examples do not resemble the current task, so the model gets no useful signal from them. Worse, irrelevant examples can mislead.
  • Context budget. Every example consumes tokens. With a library of 200 possible demonstrations, you cannot include them all. You need to pick the right ones.
  • Staleness. The environment changes — new tools get added, APIs evolve, best practices shift — but the examples baked into the prompt stay frozen at the moment they were written.

Retrieval solves all three. You maintain a large library of demonstrations offline. At runtime, you retrieve only the ones that match the current task, spending context budget where it matters.

The Demonstration Library #

A demonstration library is a structured collection of past agent trajectories — successful runs that show how to accomplish specific tasks. Each entry captures enough context for the model to understand the situation, the approach, and why it worked.

┌──────────────────────────────────────────────────────────────┐
│                    Demonstration Library                     │
│                                                              │
│  ┌────────────────────────────────────────────────────────┐  │
│  │  Entry                                                 │  │
│  │  ├── task_description: "Add pagination to /api/users"  │  │
│  │  ├── task_embedding: [0.12, -0.34, 0.56, ...]          │  │
│  │  ├── tags: ["api", "pagination", "rest"]               │  │
│  │  ├── trajectory:                                       │  │
│  │  │     Step 1: read_file("routes/users.py")            │  │
│  │  │     Step 2: read_file("models/user.py")             │  │
│  │  │     Step 3: edit_file("routes/users.py", ...)       │  │
│  │  │     Step 4: edit_file("tests/test_users.py", ...)   │  │
│  │  │     Step 5: run_tests() → all pass                  │  │
│  │  ├── outcome: success                                  │  │
│  │  ├── quality_score: 0.94                               │  │
│  │  └── metadata: {tools_used, duration, model, date}     │  │
│  └────────────────────────────────────────────────────────┘  │
│                                                              │
│  Index: vector store on task_embedding + tag filters         │
│  Size: hundreds to thousands of entries                      │
│  Update: append new successes, prune stale entries           │
└──────────────────────────────────────────────────────────────┘

Good Demonstration #

Not every successful run deserves a spot in the library. The best demonstrations share these properties:

  • Representative. They illustrate a common pattern, not a one-off edge case that will never recur.
  • Concise. The trajectory is the clean path — the approach you would take if you already knew the codebase.
  • Complete. The demonstration includes the full loop: understanding the problem, taking action, verifying the result. Truncated examples leave the model guessing about the end state.
  • Annotated. Brief rationale at key decision points explaining why the agent chose that action over alternatives.

Data Model #

from dataclasses import dataclass, field
from datetime import datetime


@dataclass
class TrajectoryStep:
    """A single action within a demonstration."""

    action: str  # Tool name or operation
    input: dict  # Arguments passed
    output: str  # Result received
    rationale: str = ""  # Why this step was chosen


@dataclass
class Demonstration:
    """A complete demonstration entry in the library."""

    demo_id: str
    task_description: str
    task_embedding: list[float] = field(default_factory=list)
    tags: list[str] = field(default_factory=list)
    trajectory: list[TrajectoryStep] = field(default_factory=list)
    outcome: str = "success"  # Only store successes
    quality_score: float = 0.0
    created_at: datetime = field(default_factory=datetime.now)
    model_id: str = ""
    tools_available: list[str] = field(default_factory=list)


def normalize_trajectory(
    trajectory: list[dict | TrajectoryStep],
) -> list[TrajectoryStep]:
    """Convert serialized trajectory steps into typed objects."""
    return [
        step if isinstance(step, TrajectoryStep) else TrajectoryStep(**step)
        for step in trajectory
    ]

Building the Library #

There are three primary sources for demonstrations.

Source 1: Human-Authored Seeds #

Write 20-50 demonstrations by hand for the most common task types. These are your highest-quality entries — they represent how you want the agent to behave, not how it happened to behave on a given run. They are expensive to produce but set the baseline.

Source 2: Curated Agent Runs #

As the agent operates in production, successful runs become candidates for the library. But you need quality gates:

def should_add_to_library(run: dict) -> bool:
    """Decide whether a successful run qualifies as a demonstration."""
    if run["outcome"] != "success":
        return False

    # Too many steps suggests fumbling, not expertise
    if len(run["trajectory"]) > 15:
        return False

    # Skip runs that used retry/backoff heavily
    retry_count = sum(
        1 for step in run["trajectory"]
        if step["action"] == "retry"
    )
    if retry_count > 2:
        return False

    # Require explicit verification (test pass, user confirmation)
    has_verification = any(
        step["action"] in ("run_tests", "verify_output", "user_confirmed")
        for step in run["trajectory"]
    )
    if not has_verification:
        return False

    return True

Source 3: Synthetic Generation #

For task types where you have specifications but no real runs yet, generate synthetic demonstrations. Give the agent a task with a known-good solution, let it solve it, then validate the output against the ground truth. If it passes, add the trajectory.

def generate_synthetic_demo(
    task: str,
    ground_truth_validator,
    max_attempts: int = 3,
) -> Demonstration | None:
    """Generate a demonstration by solving a task with known answers."""
    for attempt in range(max_attempts):
        result = run_agent(task)

        if not ground_truth_validator(result["output"]):
            continue

        if should_add_to_library(result):
            return Demonstration(
                demo_id=generate_id(),
                task_description=task,
                trajectory=normalize_trajectory(result["trajectory"]),
                outcome="success",
                quality_score=score_trajectory(result["trajectory"]),
            )

    return None

Retrieval at Runtime #

When a new task arrives, the agent retrieves the most relevant demonstrations before beginning its own planning. The retrieval pipeline looks like this:

  New Task
     │
     ▼
┌──────────────┐      ┌────────────────────┐
│  Embed task  │────▶ │  Vector similarity │
│  description │      │  search over demos │
└──────────────┘      └────────┬───────────┘
                               │
                               ▼
                       ┌───────────────────┐
                       │  Candidate demos  │
                       │  (top-k by sim)   │
                       └────────┬──────────┘
                                │
                                ▼
                       ┌───────────────────┐
                       │  Filter & rerank  │
                       │  (tool overlap,   │
                       │   recency, score) │
                       └────────┬──────────┘
                                │
                                ▼
                       ┌───────────────────┐
                       │  Format as        │
                       │  few-shot context │
                       └───────────────────┘

Embedding and Similarity #

The task description gets embedded using the same model used to embed the library entries. Cosine similarity identifies the closest matches. But raw semantic similarity is often not enough — you need domain-specific re-ranking.

def retrieve_demonstrations(
    task: str,
    library: "VectorStore",
    top_k: int = 5,
    final_k: int = 2,
    available_tools: list[str] | None = None,
) -> list[Demonstration]:
    """Retrieve and rerank demonstrations for a given task."""
    # Phase 1: Broad semantic search
    task_embedding = embed(task)
    candidates = library.search(task_embedding, top_k=top_k)

    # Phase 2: Rerank by multiple signals
    scored = []
    for demo in candidates:
        score = demo.similarity_score  # From vector search

        # Boost demos that use the same tools currently available
        if available_tools:
            tool_overlap = len(
                set(demo.tools_available) & set(available_tools)
            ) / max(len(demo.tools_available), 1)
            score += 0.2 * tool_overlap

        # Penalize old demonstrations (tools/APIs may have changed)
        age_days = max((datetime.now() - demo.created_at).days, 0)
        recency_penalty = 0.3 * min(age_days / 365, 1.0)
        score -= recency_penalty

        # Boost high-quality demonstrations
        score += 0.1 * demo.quality_score

        scored.append((score, demo))

    # Phase 3: Select top-k after reranking, with diversity
    scored.sort(key=lambda x: x[0], reverse=True)
    return deduplicate_and_select(scored, final_k)

Diversity in Retrieval #

If the top-5 candidates all demonstrate the same approach, the agent gets a narrow view. Sometimes you want diversity — retrieve demonstrations that solved the problem in different ways, so the model can choose the best strategy for the current context.

def deduplicate_and_select(
    scored: list[tuple[float, "Demonstration"]],
    k: int,
) -> list["Demonstration"]:
    """Select top-k demos while maintaining diversity."""
    selected = []
    selected_embeddings = []

    for score, demo in scored:
        if len(selected) >= k:
            break

        # Check similarity to already-selected demos
        if selected_embeddings:
            max_sim = max(
                cosine_similarity(demo.task_embedding, emb)
                for emb in selected_embeddings
            )
            if max_sim > 0.95:  # Too similar to one already picked
                continue

        selected.append(demo)
        selected_embeddings.append(demo.task_embedding)

    return selected

Formatting Demonstrations for the Prompt #

Retrieved demonstrations need to be formatted in a way the model can parse and generalize from. The format should clearly delineate the task, the steps taken, and the outcome — while being compact enough to fit within budget.

def format_demonstration(demo: Demonstration, max_steps: int = 8) -> str:
    """Format a demonstration as prompt context."""
    lines = [f"### Example: {demo.task_description}"]

    steps = demo.trajectory[:max_steps]
    for i, step in enumerate(steps, 1):
        action_line = f"Step {i}: {step.action}({_summarize_input(step.input)})"
        lines.append(action_line)

        if step.rationale:
            lines.append(f"  Rationale: {step.rationale}")

        # Show output briefly (truncate long outputs)
        if step.output:
            output_preview = step.output[:200]
            if len(step.output) > 200:
                output_preview += "..."
            lines.append(f"  → {output_preview}")

    lines.append(f"Outcome: {demo.outcome}")
    return "\n".join(lines)


def build_demonstration_context(
    demonstrations: list[Demonstration],
    max_tokens: int = 2000,
) -> str:
    """Build the full demonstration section for the prompt."""
    sections = []
    token_budget = max_tokens

    for demo in demonstrations:
        formatted = format_demonstration(demo)
        estimated_tokens = len(formatted.split()) * 1.3  # Rough estimate

        if estimated_tokens > token_budget:
            break

        sections.append(formatted)
        token_budget -= estimated_tokens

    if not sections:
        return ""

    header = (
        "Here are examples of similar tasks completed successfully. "
        "Use them as reference for your approach, but adapt to the "
        "current situation:\n\n"
    )
    return header + "\n\n".join(sections)

The Full Pipeline #

Putting it all together, the agent's execution flow becomes:

def run_agent_with_rag_actions(
    task: str,
    system_prompt: str,
    library: "VectorStore",
    tools: list[dict],
) -> dict:
    """Execute an agent loop with retrieval-augmented action selection."""
    # Retrieve relevant demonstrations
    available_tool_names = [t["name"] for t in tools]
    demonstrations = retrieve_demonstrations(
        task=task,
        library=library,
        available_tools=available_tool_names,
    )

    # Build the augmented prompt
    demo_context = build_demonstration_context(demonstrations)
    augmented_prompt = f"{system_prompt}\n\n{demo_context}"

    # Run the agent loop with the augmented context
    result = execute_agent_loop(
        system_prompt=augmented_prompt,
        user_message=task,
        tools=tools,
    )

    # If successful, consider adding this run to the library
    if result["outcome"] == "success" and should_add_to_library(result):
        new_demo = Demonstration(
            demo_id=generate_id(),
            task_description=task,
            task_embedding=embed(task),
            trajectory=normalize_trajectory(result["trajectory"]),
            outcome="success",
            quality_score=score_trajectory(result["trajectory"]),
            tools_available=available_tool_names,
        )
        library.add(new_demo)

    return result

This creates a flywheel: the agent retrieves demonstrations to perform well, and its successful runs feed back into the library for future retrieval.

Trajectory Compression #

Real agent trajectories are verbose. A coding agent might read 15 files, make 8 edits, run tests 4 times, and backtrack twice. Storing and retrieving the raw trace wastes tokens and obscures the important decisions. You need compression.

Approach 1: Keep Only Decision Points #

Strip out routine observations and keep only the steps where the agent made a meaningful choice — tool selections, strategy pivots, error handling.

def compress_trajectory(
    trajectory: list[TrajectoryStep],
) -> list[TrajectoryStep]:
    """Remove noise, keep decision points."""
    compressed = []
    skip_actions = {"wait", "sleep", "retry"}

    for i, step in enumerate(trajectory):
        # Skip retries and waits
        if step.action in skip_actions:
            continue

        # Skip consecutive reads of the same type without intervening action
        if (
            i > 0
            and step.action == "read_file"
            and trajectory[i - 1].action == "read_file"
            and not compressed[-1:] == []
            and compressed[-1].action == "read_file"
        ):
            # Merge into a summary: "Read files: a.py, b.py, c.py"
            compressed[-1].input["files"] = compressed[-1].input.get(
                "files", [compressed[-1].input.get("path", "")]
            )
            compressed[-1].input["files"].append(step.input.get("path", ""))
            continue

        compressed.append(step)

    return compressed

Approach 2: LLM-Summarized Trajectories #

For particularly long runs, use the model itself to summarize the trajectory into a compact narrative:

def summarize_trajectory(trajectory: list[TrajectoryStep]) -> str:
    """Use an LLM to compress a trajectory into a concise summary."""
    raw = "\n".join(
        f"{i+1}. {s.action}({s.input}) → {s.output[:100]}"
        for i, s in enumerate(trajectory)
    )

    prompt = f"""Summarize this agent trajectory into a concise step-by-step
guide (max 6 steps). Focus on the key decisions and strategy, not
routine operations. Write it as instructions another agent could follow.

Raw trajectory:
{raw}

Concise guide:
"""
    return call_model(prompt, temperature=0.0).text

The summarized version is what gets stored and later retrieved — compact enough to fit in the context window, clear enough to be actionable.

Retrieval vs Planning #

Retrieval-augmented actions are not always the right choice. The decision depends on how novel the current task is relative to the library.

Signal Action
High similarity match (> 0.85) and same tools available Retrieve and follow closely
Moderate similarity (0.6-0.85) Retrieve for inspiration, but plan independently
Low similarity (< 0.6) or no matches Plan from scratch; do not force irrelevant demos
Retrieved demo uses deprecated tools Retrieve for strategy, ignore specific tool calls
def decide_retrieval_strategy(
    task: str,
    library: "VectorStore",
    similarity_threshold: float = 0.6,
) -> str:
    """Decide how heavily to rely on retrieved demonstrations."""
    task_embedding = embed(task)
    top_match = library.search(task_embedding, top_k=1)

    if not top_match:
        return "plan_from_scratch"

    similarity = top_match[0].similarity_score

    if similarity > 0.85:
        return "follow_closely"
    elif similarity > similarity_threshold:
        return "use_as_inspiration"
    else:
        return "plan_from_scratch"

An agent that blindly follows a retrieved demonstration when the task has subtle differences will get into trouble. The retrieval should inform planning, not replace it.

Handling Stale and Contradictory Demonstrations #

Libraries accumulate cruft. Tools get deprecated. APIs change their response formats. A demonstration from six months ago might reference a function that no longer exists or follow a workflow that has been superseded.

Staleness Detection #

def detect_stale_demos(
    library: list[Demonstration],
    current_tools: list[str],
    max_age_days: int = 180,
) -> list[str]:
    """Identify demonstrations that may be outdated."""
    stale_ids = []

    for demo in library:
        # Age-based staleness
        age = (datetime.now() - demo.created_at).days
        if age > max_age_days:
            stale_ids.append(demo.demo_id)
            continue

        # Tool-based staleness: demo uses tools that no longer exist
        missing_tools = set(demo.tools_available) - set(current_tools)
        if missing_tools:
            stale_ids.append(demo.demo_id)

    return stale_ids

Contradiction Resolution #

When two demonstrations for similar tasks suggest different approaches, you have a contradiction. The resolution strategy: prefer the more recent demo (it reflects the current environment), the higher quality score, or — if scores are close — include both and let the model decide.

def resolve_contradictions(
    demos: list[Demonstration],
) -> list[Demonstration]:
    """When demos conflict, keep the most reliable one per cluster."""
    clusters = cluster_by_similarity(demos, threshold=0.9)
    resolved = []

    for cluster in clusters:
        if len(cluster) == 1:
            resolved.append(cluster[0])
        else:
            # Pick the best in the cluster
            best = max(
                cluster,
                key=lambda d: (
                    d.quality_score * 0.6
                    + recency_score(d.created_at) * 0.4
                ),
            )
            resolved.append(best)

    return resolved

Trade-Offs #

Like every architectural decision, retrieval-augmented actions come with costs.

Context window budget. Every demonstration you include displaces tokens that could be used for the current task's observations, tool outputs, or reasoning. If the agent is working with large codebases or long documents, the demonstration budget might need to be tiny — one example, compressed to its essence.

Retrieval latency. The embedding + search + rerank pipeline adds latency before the agent can begin acting. For interactive use cases (a user waiting for a response), this might add 100-500ms. For batch processing, it is negligible.

Library maintenance. Demonstrations are not a write-once asset. They require ongoing curation: pruning stale entries, revalidating against the current environment, adding coverage for new task types. Without maintenance, the library degrades and eventually misleads more than it helps.

Overfitting to demonstrations. If the agent always follows retrieved demonstrations literally, it cannot innovate or handle novel situations. The demonstrations should be suggestions, not mandates. The system prompt should make this explicit: "Use these examples as reference, but adapt your approach based on the current context."

Cold start. A new deployment has no demonstration library. You start with hand-written seeds and synthetic generation, then grow the library organically as the agent accumulates successful runs. The first few weeks of operation are the weakest.

Retrieval-Augmented Actions vs Fine-Tuning #

Both approaches aim to make the agent better at specific tasks. The trade-off is flexibility versus efficiency.

Dimension RAG Actions Fine-Tuning
Update speed Instant (add/remove demos) Hours-days (retrain)
Context cost Uses tokens at runtime Zero runtime cost
Flexibility Can adapt per-query Fixed after training
Coverage Scales to thousands of patterns Limited by training data
Interpretability Can inspect which demo was used Opaque weight changes
Cold-start Needs seed demos Needs training data

For most agent deployments, retrieval-augmented actions are the first approach to try. They are faster to iterate on, easier to debug (you can see exactly which demonstration influenced the agent), and do not require a training pipeline. Fine-tuning becomes attractive when the retrieval overhead is too expensive at scale or when the demonstrations can be baked in to avoid runtime token costs.

Conclusion #

Retrieval-augmented actions extend the RAG pattern from knowledge retrieval to behavior retrieval. Instead of asking "what information do I need?", the agent asks "how has a similar task been done before?" The demonstration library becomes a shared, evolving repository of institutional knowledge.

The key implementation decisions are: what qualifies as a good demonstration (quality gates), how to match tasks to demonstrations (embedding + reranking), how to fit demonstrations into the prompt (compression and formatting), and when to follow versus when to ignore them (similarity thresholds and novelty detection). Get these right, and the agent improves with every successful run — without retraining, without manual prompt engineering, and without unbounded context growth.