<p align="center">
  <img src="assets/banner.png" width="900" alt="ThinkRetrieve — retrieval-augmented reasoning traces">
</p>

<p><a href="https://itsvaibhav01.github.io/ThinkRetrieve/article/retrieval-timing.html">New: a short, reproducible walkthrough of retrieval timing, including the failure cases.</a></p>

<h1 align="center">Your RAG Agent May Be Searching Too Early</h1>

<p align="center">
  <strong>A small local model went from describing symptoms to recalling the right private procedure—by retrieving after it started reasoning.</strong>
</p>

<p align="center">
  Science reasoning · Procedural memory · Local models · Tool dispatch
</p>

---

> **The result in one sentence:** On a tiny six-case science-agent screen,
> ThinkRetrieve selected the correct private tool route in **4/6 cases with a
> local 4B model and 5/6 with a 9B model**; plain reasoning, prompt RAG, and
> simply thinking longer scored between 0/6 and 2/6.

Most retrieval systems search before the model has understood the problem.

That is often exactly the wrong moment.

A user reports observations: cells grow slowly, their morphology changes, the
medium remains clear, and ordinary sterility checks are negative. The useful
memory, however, may be indexed under the *cause*—“occult mycoplasma
contamination”—rather than those surface symptoms.

Raw RAG searches with the user's words. A reasoning model can form a much
better query after it has begun diagnosing the incident.

This is the idea behind
[ThinkRetrieve](https://github.com/itsvaibhav01/ThinkRetrieve): let the model
think, use its interim answer to retrieve a worked solution, inject that
solution into the live reasoning trace, and then continue.

No fine-tuning. No new model. No giant agent framework.

Just better timing for retrieval.

<p align="center">
  <img src="assets/reasoning-aware-retrieval.png" width="1100" alt="Raw prompt RAG follows surface wording to the wrong procedure, while ThinkRetrieve retrieves after diagnosing mycoplasma and selects the correct route">
</p>

## The use case: from laboratory symptoms to a tool call

We wanted a demonstration outside the paper's math benchmarks—something that
combined science, reasoning, retrieval, and an agentic decision in one small
workflow.

So we built a synthetic procedural-memory task:

1. The model receives a simulated cell-culture incident.
2. It must infer the likely underlying cause.
3. It searches a four-entry private “Atlas” procedure bank.
4. It selects an opaque `dispatch_route(...)` action.

The route names are deliberately meaningless without memory:

```text
A. dispatch_route("atlas_blue")
B. dispatch_route("atlas_red")
C. dispatch_route("atlas_green")
D. dispatch_route("atlas_amber")
```

Knowing general biology is not enough. The agent must connect its diagnosis to
the organization's private procedure mapping.

This is where procedural memory differs from ordinary factual RAG. The agent
does not merely need another fact. It needs to recall **what process worked for
this kind of problem**.

## Why prompt RAG failed

The private bank contained four plausible procedures:

- occult mycoplasma contamination;
- incubator CO₂ calibration drift;
- ordinary bacterial contamination;
- nutrient depletion after a media error.

The incident wording strongly resembled the CO₂ entry: slow growth, clear
medium, granular-looking cells. Raw-question retrieval followed those shared
words and retrieved the wrong procedure in all six recorded runs for both
models.

ThinkRetrieve first elicited an interim root-cause diagnosis. When the model
identified mycoplasma, that diagnosis became part of the retrieval query. The
correct procedure moved to the top, was inserted into the trace, and the agent
could resolve it to `atlas_blue`.

The useful information flow is:

```text
symptoms → reasoning → diagnosis → relevant procedure → tool route
```

Prompt RAG tries to jump directly from symptoms to procedure. Sometimes that
works. In this case, it repeatedly matched the wrong abstraction level.

## What we measured

We froze six paraphrases of the incident and moved the correct route to
different answer positions. Every arm used:

- the same local Ollama model;
- the same four-entry memory bank;
- temperature 0;
- the same 768-token thinking budget;
- the same action resolver;
- the same NVIDIA RTX 4080 Super.

We compared four methods:

| Method | Behavior |
|---|---|
| **Plain** | One uninterrupted reasoning pass, without private memory |
| **Prompt RAG** | Retrieve once from the raw incident and prepend the result |
| **More thinking (TTS)** | Split the same budget and force reconsideration, without retrieval |
| **ThinkRetrieve** | Retrieve from the incident plus interim diagnosis, inject, and continue |

<p align="center">
  <img src="assets/accuracy-by-method.png" width="1100" alt="Accuracy comparison showing ThinkRetrieve at 67 percent for Qwen 3.5 4B and 83 percent for Qwen 3.5 9B">
</p>

| Local model | Plain | Prompt RAG | More thinking | ThinkRetrieve |
|---|---:|---:|---:|---:|
| Qwen 3.5 4B | 2/6 | 0/6 | 2/6 | **4/6** |
| Qwen 3.5 9B | 0/6 | 0/6 | 0/6 | **5/6** |

The 9B run is the cleanest illustration: **0/6 with all three controls and 5/6
with reasoning-aware retrieval**.

The failure is useful too. In the missed 9B case, the model diagnosed CO₂ drift
and therefore retrieved the CO₂ procedure. ThinkRetrieve did not magically
override a wrong diagnosis; it faithfully amplified the model's current
reasoning state.

That tells us where the method helps and where it can still fail:

> Retrieval quality improves when the model's emerging diagnosis is more
> informative than the original wording. A wrong diagnosis can still produce a
> wrong memory.

## This is not “free accuracy”

Retrieval adds work. Mean wall-clock time per case was:

| Model | Plain | ThinkRetrieve | Difference |
|---|---:|---:|---:|
| Qwen 3.5 4B | 6.15 s | 6.06 s | −0.09 s |
| Qwen 3.5 9B | 7.65 s | 8.36 s | +0.71 s |

On this local setup, the 4B timing difference was negligible and the 9B method
added about 0.7 seconds per case. These are single-machine exploratory timings,
not a throughput benchmark, but they make the engineering trade-off visible.

The more important cost is operational: the procedure bank must contain
verified solutions. Bad procedural memory can make an agent confidently follow
the wrong process.

## What the library actually does

<p align="center">
  <img src="assets/pipeline.png" width="1100" alt="ThinkRetrieve pipeline: think, retrieve a solved example, inject it into the trace, and continue">
</p>

The core loop is intentionally small:

```python
from thinkretrieve import FaissRetriever, OpenAICompatBackend, ThinkRetrieve

memory = FaissRetriever.from_examples([
    (
        "Atlas procedure for occult mycoplasma contamination",
        "Dispatch atlas_blue: quarantine -> confirm -> dispose -> decontaminate",
    ),
    # Add verified runbooks, solved tickets, or successful tool traces here.
])

backend = OpenAICompatBackend(
    model="qwen3.5:4b",
    base_url="http://localhost:11434/v1",
)

result = ThinkRetrieve(backend, memory).run(incident)

print(result.answer)
print(result.retrievals)  # inspect what influenced the decision
```

During generation, ThinkRetrieve:

1. lets the model reason for a segment;
2. temporarily asks for its current best answer;
3. retrieves using `question + interim answer`;
4. injects a solved example into the trace;
5. lets the same model reconsider and finish.

Because it uses a chat-style backend, the pattern works with Ollama, LM Studio,
MLX, llama.cpp, vLLM, hosted OpenAI-compatible endpoints, Anthropic, and Amazon
Bedrock.

## Reproduce the exact experiment

Clone the repository and install the retrieval extra:

```bash
git clone https://github.com/itsvaibhav01/ThinkRetrieve.git
cd ThinkRetrieve

python -m venv .venv
source .venv/bin/activate
python -m pip install -e "./thinkretrieve[faiss]"
```

Install [Ollama](https://ollama.com/download), start it, and download a model:

```bash
ollama pull qwen3.5:4b
# or
ollama pull qwen3.5:9b
```

Then run all four arms:

```bash
python thinkretrieve/benchmarks/science_agent/run.py \
  --model qwen3.5:4b \
  --output science_agent_results.json
```

The first run also downloads `intfloat/e5-base-v2` for CPU retrieval.

The repository includes the
[exact runner](https://github.com/itsvaibhav01/ThinkRetrieve/blob/main/thinkretrieve/benchmarks/science_agent/run.py),
[4B raw records](https://github.com/itsvaibhav01/ThinkRetrieve/blob/main/thinkretrieve/benchmarks/science_agent/qwen35_4b.json),
and
[9B raw records](https://github.com/itsvaibhav01/ThinkRetrieve/blob/main/thinkretrieve/benchmarks/science_agent/qwen35_9b.json).
Each record preserves the interim diagnosis, raw retrieval, dynamic retrieval,
answer, token count, latency, and correctness.

## Where this pattern could matter

The experiment uses simulated laboratory incidents, but the underlying shape
appears in many agent workflows:

| User describes… | Model infers… | Private memory stores… |
|---|---|---|
| symptoms and instrument behavior | likely failure mode | troubleshooting SOP |
| logs and a production incident | probable root cause | remediation runbook |
| a customer complaint | policy category and risk | verified resolution sequence |
| test failures and stack traces | bug class | successful repair pattern |
| a business request | intent and constraints | approved multi-tool workflow |

In each case, the raw request and the useful memory may use different language.
The model's reasoning can bridge them before retrieval happens.

This is especially attractive for small local models. Instead of retraining the
model every time a team learns a better procedure, add the verified solution to
the memory bank:

```python
if verifier.approved(task, result):
    memory.add(task, result.answer)
    memory.save("agent_procedural_memory")
```

The model stays the same. The agent's procedural memory improves.

## What we are—and are not—claiming

This is a **small, synthetic, exploratory use-case demonstration**. It is not a
biological safety benchmark, and no real laboratory tool was executed.

The use case was selected after exploratory screening. The six cases are
development examples around one failure mode, not a preregistered or held-out
evaluation. Temperature 0 also does not guarantee bit-for-bit determinism on
every local GPU stack.

So the honest claim is narrow:

> In this simulated agent workflow, two small local models retrieved private
> procedures more successfully after forming a diagnosis than when retrieving
> directly from the raw incident.

The next serious evaluation should freeze a larger suite across qPCR,
microscopy, chromatography, instrument triage, software incidents, and
multi-step tool execution—and report multiple runs with confidence intervals.

## Try ThinkRetrieve

If your model can often explain the problem but still chooses the wrong process,
give it access to examples of successful reasoning **while it thinks**.

```bash
pip install "thinkretrieve[faiss]"
```

- [GitHub repository](https://github.com/itsvaibhav01/ThinkRetrieve)
- [10-minute tutorial](https://github.com/itsvaibhav01/ThinkRetrieve/blob/main/thinkretrieve/TUTORIAL.md)
- [Project website](https://itsvaibhav01.github.io/ThinkRetrieve/)
- [Paper](https://arxiv.org/abs/2608.10928)
- [Science-agent benchmark](https://github.com/itsvaibhav01/ThinkRetrieve/tree/main/thinkretrieve/benchmarks/science_agent)

If you test it on your own runbooks, solved tickets, scientific workflows, or
tool traces, share the full comparison—including the failures. Procedural
memory becomes much more useful when the community can see where it works and
where it does not.

---

## Cite the work

[Paper abstract and citation details](../paper/) ·
[Download BibTeX](../paper/thinkretrieve.bib) ·
[Download RIS for Zotero / EndNote](../paper/thinkretrieve.ris)

```bibtex
@article{thinkretrieve2026,
  title   = {ThinkRetrieve: Retrieval-Augmented Reasoning Traces for Test-Time Scaling},
  author  = {Singh, Vaibhav and Ghosal, Soumya Suvra and Gharat, Sarvesh and
             Pal, Soumyabrata and Narayanam, Ramasuri and Manocha, Dinesh},
  journal = {arXiv preprint arXiv:2608.10928},
  year    = {2026},
  doi     = {10.48550/arXiv.2608.10928},
  url     = {https://arxiv.org/abs/2608.10928}
}
```

<p align="center"><strong>Think → diagnose → retrieve → act.</strong></p>
