swiftrag 0.1.0

Instant, optimized Retrieval-Augmented Generation. Pass your text + a model, get a RAG-powered LLM in one line.

swiftrag

Instant, optimized Retrieval-Augmented Generation. Pass your text and a model — get a RAG-powered LLM in one line.

from swiftrag import RAG

rag = RAG(
    documents="The Eiffel Tower is 330 metres tall and located in Paris.",
    embedding_model="openai:text-embedding-3-small",
    llm_model="openai:gpt-4o-mini",
)

print(rag.query("How tall is the Eiffel Tower?"))
# -> "The Eiffel Tower is 330 metres tall."

That's the whole API. You bring documents (a string, a list of strings, or dicts) and a model spec; swiftrag handles chunking, embedding, vector indexing, retrieval, and prompt construction.

---

Why swiftrag

• One line to a working RAG. No glue code, no framework to learn.
• Optimized core. L2-normalized embeddings + a single BLAS matmul for search, argpartition top-k (no full sort), batched & concurrent embedding requests, token-aware chunking, and an optional FAISS backend for large corpora.
• Tiny footprint. The core depends only on numpy. It installs in seconds.
• Runs offline out of the box. With no API key it uses a built-in hashing embedder and an extractive answerer, so the full pipeline works in tests/CI/demos.
• Provider-agnostic. OpenAI, Anthropic, local sentence-transformers, or any custom callable/object you plug in.
• MIT licensed.

Install

pip install swiftrag                 # core (numpy only) — works offline
pip install "swiftrag[openai]"       # OpenAI embeddings + LLM
pip install "swiftrag[anthropic]"    # Claude LLM
pip install "swiftrag[local]"        # local sentence-transformers embeddings
pip install "swiftrag[faiss]"        # FAISS backend for big corpora
pip install "swiftrag[all]"          # everything

Usage

Pick your models with a simple "provider:model" string

# OpenAI (needs OPENAI_API_KEY)
RAG(documents=text, embedding_model="openai:text-embedding-3-small", llm_model="openai:gpt-4o-mini")

# Anthropic for generation, local embeddings (no embedding API calls)
RAG(documents=text, embedding_model="local:all-MiniLM-L6-v2", llm_model="anthropic:claude-3-5-sonnet-latest")

# Fully offline (default) — no keys required
RAG(documents=text)

Build straight from files or a folder

rag = RAG.from_files("docs/", embedding_model="openai:text-embedding-3-small")
# each file becomes a document tagged with metadata={"source": <path>}

resp = rag.query("What's our deployment process?")
print(resp.answer)
print(resp.format_sources())   # numbered, human-readable citations

Documents can be a string, list, or dicts with metadata

rag = RAG(documents=[
    "Plain string document.",
    {"text": "Document with metadata.", "metadata": {"source": "handbook", "page": 12}},
])

Query, stream, or just retrieve

resp = rag.query("What does the handbook say about refunds?")
print(resp.answer)
for s in resp.sources:
    print(s.score, s.metadata, s.text[:80])

# Token streaming
for token in rag.stream("Summarize the refund policy."):
    print(token, end="", flush=True)

# Retrieval only (no LLM call)
chunks = rag.retrieve("refunds", top_k=5)

Filter by metadata and score

# Only consider chunks whose metadata matches, and drop weak matches.
resp = rag.query(
    "What is the refund window?",
    where={"source": "handbook"},     # exact metadata match (or pass a Chunk -> bool callable)
    min_score=0.25,                    # cosine threshold; weaker chunks are ignored
    top_k=3,
)

Repeated queries reuse a cached query embedding (LRU, configurable via query_cache_size), so re-asking the same question skips the embedding call.

Add documents incrementally, save, and reload

rag = RAG().add("first batch").add("second batch")
rag.save("index.pkl")

rag = RAG.load("index.pkl", embedding_model="openai:text-embedding-3-small",
               llm_model="openai:gpt-4o-mini")

Batch and async

# Answer many questions at once — embeddings are batched, generation is parallelized.
responses = rag.query_many(["q1?", "q2?", "q3?"], max_workers=8)

# Async API (non-blocking, great for web servers):
resp = await rag.aquery("your question")
async for token in rag.astream("your question"):
    print(token, end="", flush=True)

Bring your own provider

# Any callable fn(prompt) -> str works as an LLM:
rag = RAG(documents=text, llm_model=lambda prompt: my_model.generate(prompt))

# Any object with embed_documents(list[str]) and embed_query(str) works as an embedder.

Configuration

Argument	Default	Description
documents	None	str / list[str] / list[dict] / Document(s) to index.
embedding_model	"hash"	"provider:model" string or a custom provider.
llm_model	None (offline)	"provider:model" string, callable, or provider.
chunk_size	512	Target chunk size in tokens.
chunk_overlap	64	Token overlap between chunks.
top_k	4	Chunks retrieved per query.
use_mmr	False	Maximal Marginal Relevance re-ranking for diverse results.
use_faiss	False	Use FAISS index (install swiftrag[faiss]).
min_score	None	Default cosine threshold for dropping weak matches.
max_context_tokens	None	Cap the tokens of retrieved context packed into the prompt.
dedup	True	Skip exact-duplicate chunks on ingest.
query_cache_size	128	LRU size for cached query embeddings (0 disables).
system_prompt	grounded default	System prompt for the LLM.

How it works

documents ─▶ chunk (token-aware) ─▶ embed (batched) ─▶ normalize ─▶ vector store
                                                                         │
query ─▶ embed ─▶ cosine (BLAS matmul) ─▶ top-k (argpartition) ─▶ context ─▶ LLM ─▶ answer

Development

pip install -e ".[dev]"
pytest
ruff check .

License

MIT — see LICENSE.