Skip to content

Latest commit

 

History

24 Commits

Folders and files

NameName
Last commit message
Last commit date
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

rag-prototype

A retrieval-augmented generation (RAG) system built from scratch, one small, typed, tested slice at a time — deliberately without frameworks until the lower-level mechanics justify them.

The goal is understanding, not shortcuts: every architectural decision here is explicit and defensible. Frameworks and infrastructure are introduced only once the primitive they abstract is understood first: persistent pgvector storage arrived only after a hand-written in-memory store proved the retrieval core, and heavier frameworks (FastAPI, Pydantic AI, LangChain) stay deferred until the mechanics they hide are built by hand.

The RAG pipeline

                 INDEXING TIME
  Document ──► chunk ──► embed ──► ┌─────────────────┐
                                   │  vector store   │
  Query ─────────────► embed ────► │ (EmbeddedChunks)│
                                   └────────┬────────┘
                 QUERY TIME                 │ cosine similarity
                                            ▼
                                   ranked RetrievedChunks ──► (LLM answer)

The core idea: text is split into chunks, each chunk is turned into a numeric embedding, and both live together in a vector store. At query time the question is embedded with the same model, compared against every stored embedding by cosine similarity, and the closest chunks' original text (not their vectors) is what would be handed to a chat model.

Modules

Module Responsibility
domain.py Document and TextChunk value objects, with invariants enforced in __post_init__.
chunking.py Deterministic fixed-character chunker with configurable size and overlap. Guarantees chunk.content == document.content[chunk.start_offset:chunk.end_offset].
embedding.py The Embedding type, the EmbeddedChunk value object, and the pure cosine_similarity function.
embedder.py The Embedder Protocol and HashingEmbedder — a deterministic, framework-free fake embedder for testing the pipeline end-to-end.
lmstudio.py LMStudioEmbedder — a real Embedder backed by LM Studio's OpenAI-compatible HTTP API, with an injectable httpx.Client and an EmbeddingError boundary.
store.py The VectorStore Protocol, the RetrievedChunk result type, and InMemoryVectorStore — a linear-scan store returning ranked top-k results by cosine similarity.
pgvector_store.py PgvectorStore — a persistent VectorStore backed by PostgreSQL + pgvector: ranking is pushed into SQL via the cosine-distance operator (<=>) and converted back to similarity. Ships a create_schema DDL helper and a validated, configurable embedding dimensions.

Design decisions worth knowing

These are the "why"s behind the code — the interesting part of the project.

  • Value objects are frozen dataclasses with guarded invariants. A TextChunk or EmbeddedChunk cannot exist in an invalid state; construction fails loudly rather than deferring a confusing error downstream.

  • Embedding is a tuple[float, ...], not a class. It's a pure value with no behavior of its own — an immutable, hashable sequence. A type alias documents intent at zero runtime cost.

  • The zero-vector guard lives in cosine_similarity, not only on EmbeddedChunk. The query vector never passes through EmbeddedChunk, so the similarity function must defend its own precondition — callers get a clear ValueError, never a raw ZeroDivisionError.

  • The embedder is a Protocol (structural typing). The fake HashingEmbedder and the real LMStudioEmbedder are interchangeable without inheritance — and mypy verifies the substitution statically. The rest of the pipeline never changed when the real adapter arrived; that is the payoff.

  • Adapters own their failure boundary. LMStudioEmbedder wraps malformed responses in an EmbeddingError (with raise ... from exc chaining) instead of leaking raw KeyError/IndexError, so callers depend on the Embedder contract, not on httpx internals. The HTTP client is injected so tests drive it with httpx.MockTransport and never touch the network.

  • dimensions is embedder instance state, not a per-call argument. This guarantees indexing and querying use the identical transform; you cannot accidentally embed a chunk and a query into differently-sized spaces.

  • The fake embedder uses feature hashing, not random noise. Hashing tokens into a fixed number of buckets makes lexical overlap produce higher cosine similarity, so retrieval tests assert something meaningful rather than merely "exact text finds itself."

  • hashlib.sha256, never the built-in hash(). Python's hash() for strings is randomized per process (PYTHONHASHSEED), which would silently make the embedder non-reproducible across runs. sha256 is stable everywhere.

  • The persistent store pushes ranking into SQL, then restores the contract. PgvectorStore ranks with pgvector's cosine-distance operator (<=>, ORDER BY ... ASC) and converts back to similarity (score = 1 - distance) so it agrees with InMemoryVectorStore about what "best" means. The math leaves Python entirely; because both satisfy the VectorStore Protocol, nothing upstream notices the swap.

  • The embedding dimension is a validated, interpolated integer — not a bound parameter. A column type like VECTOR(768) cannot be a %(param)s, so dimensions is checked (positive int, excluding bool) and interpolated into the DDL, while every value still flows through psycopg's %(name)s params. Store and schema share the width, enforcing "you cannot mix vectors from two different embedders in one store."

Testing philosophy

  • Behaviour is pinned by tests before moving on; the suite is the specification.
  • Value objects are constructed directly (via small factory helpers), never mocked — they're cheap and their invariants are part of what's under test.
  • The fake embedder is a fake, not a mock: real, deterministic behavior lets the whole retrieval pipeline be tested with no network and no model weights.
  • Integration tests earn the name. A SQL store's behaviour is the SQL, so PgvectorStore is tested against a real pgvector Postgres — marked @pytest.mark.integration, skipped when none is reachable, and isolated by transaction-per-test rollback — cross-checking its ranking against the in-memory store. A fake returning canned rows would prove nothing.

Getting started

Requires Python ≥ 3.11 and uv.

uv sync                 # install dependencies
uv run pytest           # run the test suite
uv run mypy src tests   # type-check
uv run ruff check       # lint

See it retrieve

examples/demo.py runs the whole pipeline end-to-end — index a small corpus, then ask questions and print the ranked chunks:

uv run python examples/demo.py
Indexed 6 chunks from 3 documents.

❓ how do I set a breakpoint in python?
  1. [0.217] (python.md) Python uses the built-in breakpoint function to drop into the pdb debugger...
  2. [0.186] (ruby.md) Ruby developers reach for the debugger gem or pry to set breakpoints...
  3. [0.079] (python.md) with the pdb flag to inspect a failing test at the point it broke.

Each question surfaces its own topic first; the cross-topic overlap (Ruby's debugging doc ranking second for a Python debugging question) is the feature hashing capturing shared vocabulary — not a coincidence.

examples/demo_pgvector.py is the fully-real counterpart: the same corpus and queries, but embedded by LM Studio and stored in a pgvector Postgres — both Embedder and VectorStore Protocols swapped for real backends at once, the pipeline in between unchanged. It probes the model's real embedding dimension at startup (so it fits any model), treats both services as hard requirements, and degrades gracefully with a hint if either is unreachable. Configure via DATABASE_URL, LM_STUDIO_URL, and LM_STUDIO_MODEL (see docs/STEP_06_pgvector_store.md):

LM_STUDIO_URL=http://<host>:1234/v1 LM_STUDIO_MODEL=<model> \
  uv run python examples/demo_pgvector.py

Roadmap

  • Domain models + deterministic chunker
  • Embedding value types + cosine similarity
  • Deterministic fake embedder (feature hashing)
  • In-memory vector store + retrieval tests
  • Real LM Studio embedding adapter
  • Persistent PostgreSQL / pgvector storage
  • FastAPI boundary with Pydantic request/response models
  • RAG answer generation and evaluation
  • Pydantic AI tools and agent workflows

Design briefs for each step live in docs/.

About

RAG Learning project

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages