Metadata-Version: 2.4
Name: thinkretrieve
Version: 0.1.0
Summary: Retrieval-augmented test-time scaling: inject retrieved worked examples into a reasoning model's thinking trace, mid-thought. Works over any chat API (OpenAI-compatible, Anthropic, Bedrock) — no GPU required.
Project-URL: Homepage, https://thinkretrieve.github.io
Project-URL: Paper, https://arxiv.org/abs/XXXX.XXXXX
Project-URL: Repository, https://github.com/itsvaibhav01/RAG_thinker
Project-URL: Documentation, https://github.com/itsvaibhav01/RAG_thinker/blob/main/thinkretrieve/TUTORIAL.md
Author: Vaibhav Singh
License: MIT License
        
        Copyright (c) 2026 Vaibhav Singh
        
        Permission is hereby granted, free of charge, to any person obtaining a copy
        of this software and associated documentation files (the "Software"), to deal
        in the Software without restriction, including without limitation the rights
        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
        copies of the Software, and to permit persons to whom the Software is
        furnished to do so, subject to the following conditions:
        
        The above copyright notice and this permission notice shall be included in all
        copies or substantial portions of the Software.
        
        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
        SOFTWARE.
License-File: LICENSE
Keywords: faiss,in-context-learning,llm,rag,reasoning,retrieval,test-time-scaling
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.9
Requires-Dist: httpx>=0.24
Provides-Extra: all
Requires-Dist: anthropic>=0.40; extra == 'all'
Requires-Dist: boto3>=1.34; extra == 'all'
Requires-Dist: faiss-cpu>=1.7; extra == 'all'
Requires-Dist: numpy; extra == 'all'
Requires-Dist: sentence-transformers>=2.2; extra == 'all'
Provides-Extra: anthropic
Requires-Dist: anthropic>=0.40; extra == 'anthropic'
Provides-Extra: bedrock
Requires-Dist: boto3>=1.34; extra == 'bedrock'
Provides-Extra: faiss
Requires-Dist: faiss-cpu>=1.7; extra == 'faiss'
Requires-Dist: numpy; extra == 'faiss'
Requires-Dist: sentence-transformers>=2.2; extra == 'faiss'
Description-Content-Type: text/markdown

<!-- Images resolve on PyPI/GitHub only when the repo is public. Base path:
     https://raw.githubusercontent.com/itsvaibhav01/RAG_thinker/main/thinkretrieve/assets/ -->
<p align="center">
  <img src="https://raw.githubusercontent.com/itsvaibhav01/RAG_thinker/main/thinkretrieve/assets/banner.png" width="720" alt="ThinkRetrieve">
</p>

<p align="center">
  <a href="https://pypi.org/project/thinkretrieve/"><img src="https://img.shields.io/pypi/v/thinkretrieve.svg?color=2a78d6" alt="PyPI"></a>
  <a href="https://pypi.org/project/thinkretrieve/"><img src="https://img.shields.io/pypi/pyversions/thinkretrieve.svg" alt="Python"></a>
  <a href="https://github.com/itsvaibhav01/RAG_thinker/blob/main/thinkretrieve/LICENSE"><img src="https://img.shields.io/badge/license-MIT-green.svg" alt="License"></a>
  <a href="https://arxiv.org/abs/XXXX.XXXXX"><img src="https://img.shields.io/badge/arXiv-XXXX.XXXXX-b31b1b.svg" alt="arXiv"></a>
</p>

<p align="center">
  <b><a href="https://arxiv.org/abs/XXXX.XXXXX">Paper</a></b> ·
  <b><a href="https://thinkretrieve.github.io">Project page</a></b> ·
  <b><a href="https://github.com/itsvaibhav01/RAG_thinker">Code</a></b>
</p>

---

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

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

## Quickstart

```python
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

<p align="center">
  <img src="https://raw.githubusercontent.com/itsvaibhav01/RAG_thinker/main/thinkretrieve/assets/pipeline.png" width="860" alt="How ThinkRetrieve works">
</p>

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

<p align="center">
  <img src="https://raw.githubusercontent.com/itsvaibhav01/RAG_thinker/main/thinkretrieve/assets/wins_every_cell.png" width="620" alt="Wins every cell">
</p>

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

<p align="center">
  <img src="https://raw.githubusercontent.com/itsvaibhav01/RAG_thinker/main/thinkretrieve/assets/scaling.png" width="720" alt="Scaling behavior">
</p>

| Best accuracy (%) | GSM-8K | MATH-500 | AIME 2025 |
|---|---|---|---|
| | TTS / **Ours** | TTS / **Ours** | TTS / **Ours** |
| Qwen3-1.7B | 90.3 / **92.1** | 91.0 / **92.5** | 22.2 / **35.6** |
| Qwen3-4B | 95.1 / **96.8** | 93.7 / **96.1** | 64.4 / **66.7** |
| Qwen3-8B | 96.4 / **97.2** | 94.0 / **94.8** | 68.9 / **71.1** |

## Use cases

<details open>
<summary><b>1 · Reason over a dataset (math, science QA, …)</b></summary>

Index any dataset with worked solutions and go. See
[`examples/sciq_example.py`](examples/sciq_example.py) for a full SciQ run
(build bank from the train split → answer test questions → score):

```bash
python examples/sciq_example.py --limit 30
```
Ready-made banks: NuminaMath, MetaMathQA, GSM8K, MATH, SciQ — anything with
`(question, step-by-step solution)` pairs.
</details>

<details>
<summary><b>2 · Procedural memory for agents</b></summary>

Most agent memory stores *facts*. A ThinkRetrieve bank stores *procedures* —
solved tasks — and recalls them **mid-reasoning**. The bank grows as the agent
works ([`examples/agent_memory.py`](examples/agent_memory.py)):

```python
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")
```
</details>

<details>
<summary><b>3 · Upgrade an existing RAG stack</b></summary>

Prompt-level RAG is ~neutral on reasoning tasks; injecting the *same* content
*inside the trace* is what wins (paper §5). Keep your index — wrap it in a 5-line
retriever ([`examples/rag_integration.py`](examples/rag_integration.py)):

```python
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)
```
</details>

<details>
<summary><b>Compare it yourself (one command)</b></summary>

```bash
python examples/compare_tts_vs_thinkretrieve.py --backend openai --model qwen3:4b
```
Runs plain / RAG / long-thinking / ThinkRetrieve at the same budget and prints
the table. Add `--backend bedrock --model qwen.qwen3-32b-v1:0` for Amazon Bedrock.
</details>

## Backends

| Backend | Import | Use for |
|---|---|---|
| OpenAI-compatible | `OpenAICompatBackend` | Ollama, LM Studio, MLX, llama.cpp, vLLM, OpenAI, Together, Groq, DeepSeek, OpenRouter |
| Anthropic | `AnthropicBackend` | Claude models (`pip install "thinkretrieve[anthropic]"`) |
| Amazon Bedrock | `BedrockConverseBackend` | Bedrock models (`pip install "thinkretrieve[bedrock]"`) |

Or subclass `ChatTranscriptBackend` and implement one `_chat()` method.

## Install options

```bash
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](TUTORIAL.md).

## Citation

```bibtex
@article{thinkretrieve2026,
  title  = {ThinkRetrieve: Retrieval-Augmented Reasoning Traces for Test-Time Scaling},
  author = {Singh, Vaibhav and Ghosal, Soumya Sourav and Gharat, Sarvesh and
            Pal, Soumyabrata and Narayanam, Ramasuri and Manocha, Dinesh},
  journal = {arXiv preprint arXiv:XXXX.XXXXX},
  year   = {2026}
}
```

## License

MIT
