README.md

Paper · PDF · Project page · Code

ThinkRetrieve makes a reasoning model recall worked examples mid-thought. Instead of only "thinking longer" (which drifts and compounds errors), it pauses at each reasoning step, retrieves the most similar solved example from a bank, and injects it into the thinking trace — guidance on how to reason, not just what facts to use.

It works over any chat API (OpenAI-compatible, Anthropic, Amazon Bedrock, or a local model via Ollama / LM Studio / MLX / llama.cpp) — no GPU or vLLM required.

pip install "thinkretrieve[faiss]"

The local example below requires a running model. With Ollama, download the tutorial model first:

ollama pull qwen3:4b

To verify a cloned repository without downloading a model, dataset, or embedding encoder, install the core package and run examples/offline_dummy.py:

git clone https://github.com/itsvaibhav01/ThinkRetrieve.git
cd ThinkRetrieve/thinkretrieve
python -m pip install -e .
python examples/offline_dummy.py

The full tutorial explains local and hosted models, required downloads, small smoke runs, full dataset runs, and troubleshooting.

Quickstart

from thinkretrieve import ThinkRetrieve, FaissRetriever, OpenAICompatBackend

# 1. A bank of solved (question, worked-solution) pairs — a dataset, your docs, past tickets…
bank = FaissRetriever.from_examples([
    ("A jacket costs \$120 and is discounted 25%. Final price?",
     "Discount = 0.25*120 = 30. Final = 120-30 = 90. Answer: \$90"),
    ("What is 15% of 80?", "0.15*80 = 12. Answer: 12"),
])

# 2. Any chat model — local (shown) or hosted.
backend = OpenAICompatBackend(model="qwen3:4b", base_url="http://localhost:11434/v1")

# 3. Reason with mid-thought retrieval.
result = ThinkRetrieve(backend, bank).run(
    "A phone costs \$250 after a 20% discount. What was the original price?")
print(result.answer)          # the final answer
print(result.think_trace)     # full reasoning, with injected examples visible
print(result.retrievals)      # what was retrieved, when, and why

How it works

At each reasoning boundary the model elicits an interim answer, retrieves the most similar solved example (E5 + FAISS by default), injects it into the trace, and continues — repeating until the thinking budget is spent. Passing retriever=None gives you the plain "think longer" baseline for A/B comparison.

Results

Across 5 reasoning models × 4 benchmarks (paper), ThinkRetrieve beats standard thinking, sequential test-time scaling (TTS), static in-context ICL, and random per-step retrieval on every cell — and, because injected tokens count against the budget but aren't generated, it produces fewer model tokens than TTS at the same budget (~6% wall-clock overhead).

Sequential TTS degrades or plateaus as the budget grows; ThinkRetrieve keeps improving — most dramatically on the hardest benchmark:

Use cases

python examples/sciq_example.py --limit 30

result = agent.run(task)
if verified(result):            # tests pass / human approves
    memory.add(task, result.answer)   # the agent now remembers HOW
    memory.save("agent_memory")

from thinkretrieve.types import Example

class MyStore:                              # pgvector / Pinecone / Chroma / …
    def retrieve(self, question, interim_answer="", k=1, exclude_ids=frozenset()):
        hits = my_index.search(f"{question}\n{interim_answer}", k + len(exclude_ids))
        return [Example(id=h.id, question=h.q, solution=h.body)
                for h in hits if h.id not in exclude_ids][:k]

ThinkRetrieve(backend, MyStore()).run(question)

python examples/compare_tts_vs_thinkretrieve.py --backend openai --model qwen3:4b

Backends

Or subclass ChatTranscriptBackend and implement one _chat() method.

Install options

pip install "thinkretrieve[faiss]"            # core + FAISS retrieval (recommended)
pip install "thinkretrieve[faiss,anthropic]"  # + Anthropic
pip install "thinkretrieve[all]"              # everything

No GPU needed: retrieval runs on CPU / Apple Silicon; generation runs wherever your model lives. Full walkthrough in TUTORIAL.md.

Citation

@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},
  url    = {https://arxiv.org/abs/2608.10928}
}

License

MIT