ThinkRetrieve — retrieval-augmented reasoning traces

New: a short, reproducible walkthrough of retrieval timing, including the failure cases.

Your RAG Agent May Be Searching Too Early

A small local model went from describing symptoms to recalling the right private procedure—by retrieving after it started reasoning.

Science reasoning · Procedural memory · Local models · Tool dispatch


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: 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.

Raw prompt RAG follows surface wording to the wrong procedure, while ThinkRetrieve retrieves after diagnosing mycoplasma and selects the correct route

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:

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:

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:

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:

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

Accuracy comparison showing ThinkRetrieve at 67 percent for Qwen 3.5 4B and 83 percent for Qwen 3.5 9B

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

ThinkRetrieve pipeline: think, retrieve a solved example, inject it into the trace, and continue

The core loop is intentionally small:

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:

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, start it, and download a model:

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

Then run all four arms:

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, 4B raw records, and 9B raw records. 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:

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.

pip install "thinkretrieve[faiss]"

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 · Download BibTeX · Download RIS for Zotero / EndNote

@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}
}

Think → diagnose → retrieve → act.