diff --git a/.env.example b/.env.example index f731735..a4ea9f0 100644 --- a/.env.example +++ b/.env.example @@ -1,2 +1,12 @@ ANTHROPIC_API_KEY= ANTHROPIC_MODEL=claude-haiku-4-5 + +# How chunks are retrieved: tfidf, semantic or hybrid (default: hybrid). +# tfidf needs no model download. semantic and hybrid download the embedding +# model (about 220 MB) the first time the server starts. +RETRIEVAL_MODE=hybrid + +# Qdrant server used by semantic and hybrid retrieval, for example +# http://localhost:6333 after `docker compose up -d`. +# Leave empty to run Qdrant in memory, with nothing to start. +QDRANT_URL= diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 12a13fb..ad05af2 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -6,6 +6,9 @@ on: pull_request: branches: [main] +env: + FASTEMBED_CACHE_PATH: ${{ github.workspace }}/.cache/fastembed + jobs: test: runs-on: ubuntu-latest @@ -20,3 +23,24 @@ jobs: run: ruff check . && ruff format --check . - name: Run tests run: pytest -v + + retrieval-evaluation: + runs-on: ubuntu-latest + timeout-minutes: 15 + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + - name: Install dependencies + run: pip install -r requirements.txt + - name: Cache the embedding model + uses: actions/cache@v4 + with: + path: ${{ env.FASTEMBED_CACHE_PATH }} + key: embedding-model-${{ hashFiles('rag/embeddings.py') }} + - name: Evaluate retrieval modes + shell: bash + env: + HF_HUB_DISABLE_PROGRESS_BARS: "1" + run: python -m evaluation.run --enforce-gate | tee -a "$GITHUB_STEP_SUMMARY" diff --git a/README.md b/README.md index bafe0e1..0f39380 100644 --- a/README.md +++ b/README.md @@ -1,29 +1,94 @@ # RAG Knowledge Assistant -A small, complete retrieval-augmented generation pipeline: ingest text, retrieve the passages relevant to a question with TF-IDF, and ask Claude to answer using only what was retrieved — with every answer citing the sources it came from. +A small, complete retrieval-augmented generation pipeline: ingest text, retrieve the passages relevant to a question with TF-IDF, embeddings, or both, and ask Claude to answer using only what was retrieved — with every answer citing the sources it came from. -## Why TF-IDF instead of a dense embedding model +## Retrieval modes -This is a deliberate choice, not a shortcut. TF-IDF needs no model download and no API key, so the entire pipeline — ingestion, retrieval, and its tests — runs offline and deterministically. The vector store's public interface (`add_documents`, `search`) says nothing about TF-IDF specifically, so swapping in a dense embedding model later is a change to one file, not a redesign. +The mode is chosen with `RETRIEVAL_MODE` (default `hybrid`, documented in `.env.example`): + +| Mode | How chunks are ranked | Needs | +| --- | --- | --- | +| `tfidf` | Cosine similarity over TF-IDF vectors | Nothing: no model download, runs fully offline | +| `semantic` | Cosine similarity between embeddings, searched in Qdrant | An embedding model (about 220 MB, downloaded on first start) | +| `hybrid` | The TF-IDF and semantic rankings merged with Reciprocal Rank Fusion | The same as `semantic` | + +All three implement the same `add_documents` / `search` interface, so the pipeline and the API do not change with the mode. TF-IDF stays a first-class option: it is the baseline the other modes are measured against, and the one to pick when nothing should be downloaded. + +- **Embeddings** come from [fastembed](https://github.com/qdrant/fastembed) running `sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2` as an ONNX model, so the install has no PyTorch. The model sits behind a small `Embedder` interface that the store never looks behind. +- **Vector database** is Qdrant. By default it runs in memory inside the process, so there is nothing to start. For a real server, run `docker compose up -d` and set `QDRANT_URL=http://localhost:6333`. The collection is rebuilt every time the app starts, so the semantic index always matches the lexical one. +- **Reciprocal Rank Fusion** scores a chunk as the sum of `1 / (60 + rank)` over the rankings it appears in. Only ranks are used, so the two scorers never need their scores normalized against each other. +- **`score`** in a response belongs to the mode that produced it (a TF-IDF cosine, an embedding cosine, or a fused rank score), so it is only comparable within one response. + +Semantic and hybrid retrieval always return the closest chunks, even for a question the corpus cannot answer. The prompt is what tells the model to say when the sources do not contain the answer. + +## Evaluation + +Each mode is measured on a golden set of 40 questions over the 18 sample notes (`evaluation/golden_set.json`). Every question is labeled with the note, or notes, that answer it. A question is a hit when one of the top 3 retrieved chunks comes from a labeled note, and MRR is the mean of 1 / the position of the first such chunk (0 if none is in the top 10). + +The questions come in three kinds: 14 that reuse many of the words of the answering note, 14 paraphrases that share few of them, and 12 exact identifiers or terms such as `B-tree` or `(user_id, created_at)`. A test checks that each kind keeps that definition and that no question repeats a phrase of the notes. The questions were written once, before any mode was run against them, and have not been edited since. + +Results with `sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2`, as printed by the `retrieval-evaluation` CI job: + +| Mode | Hit rate@3 | MRR | +| --- | --- | --- | +| tfidf | 0.875 (35/40) | 0.852 | +| semantic | 0.900 (36/40) | 0.877 | +| hybrid | 0.975 (39/40) | 0.941 | + +Hit rate@3 by question type: + +| Mode | Lexical overlap (14) | Paraphrase (14) | Exact term (12) | +| --- | --- | --- | --- | +| tfidf | 1.000 (14/14) | 0.643 (9/14) | 1.000 (12/12) | +| semantic | 1.000 (14/14) | 0.857 (12/14) | 0.833 (10/12) | +| hybrid | 1.000 (14/14) | 0.929 (13/14) | 1.000 (12/12) | + +MRR by question type: + +| Mode | Lexical overlap (14) | Paraphrase (14) | Exact term (12) | +| --- | --- | --- | --- | +| tfidf | 1.000 | 0.577 | 1.000 | +| semantic | 0.964 | 0.818 | 0.845 | +| hybrid | 1.000 | 0.830 | 1.000 | + +What the numbers say, including where the new modes do not win: + +- Hybrid has the highest hit rate overall and is never below TF-IDF on any question type. It matches TF-IDF on word-overlap and exact-term questions, which TF-IDF answers perfectly, and gains on paraphrases (13/14 against 9/14). +- Semantic retrieval on its own does not beat TF-IDF everywhere. It is better on paraphrases (12/14 against 9/14), equal on word-overlap questions, and worse on exact terms (10/12 against 12/12): it missed `B-tree` and `(user_id, created_at)`. Its MRR is also a little lower on word-overlap questions (0.964 against 1.000). +- Fusion is not free. Hybrid missed one question that semantic retrieval alone answered, `paraphrase-05` (a service receiving the same work item twice): TF-IDF ranked two chunks of the OAuth note first, and fusion let them push the right note out of the top 3. +- Hybrid answered two paraphrases (`paraphrase-03`, `paraphrase-04`) that neither TF-IDF nor semantic retrieval had in its own top 3, so combining the two rankings pulled the right note up. + +Forty questions over 18 short notes is a small sample. One question is 2.5 points of hit rate, so the one-question gap between semantic and TF-IDF overall says little by itself. The gains on paraphrases and the weakness of semantic retrieval on exact terms are larger. The notes and the questions share an author, which favors questions the notes answer cleanly; a real corpus is messier. + +To reproduce it: + +```bash +pip install -r requirements.txt +python -m evaluation.run +``` + +The first run downloads the embedding model. The command prints the tables above and logs every question each mode missed in its top 3. With `--enforce-gate` it exits with status 1 when hybrid's hit rate@3 is below the TF-IDF baseline, which is what CI does on every pull request, with the model cached between runs. The dependency versions are pinned, but a different model or library version can move a result by a question. ## How a question gets answered 1. `POST /documents` splits the text into overlapping word chunks and indexes them. -2. `POST /query` ranks the indexed chunks against the question with cosine similarity over TF-IDF vectors. +2. `POST /query` ranks the indexed chunks against the question using the configured retrieval mode. 3. The top chunks are placed in a prompt that instructs the model to answer only from what it was given and to cite each source it uses. 4. The response returns the answer alongside the exact chunks it was grounded in, so an answer can always be checked against its sources. -The app ships with a small sample knowledge base in `docs_sample/` (original short notes on REST APIs, database indexing, the testing pyramid, and caching) so it is queryable immediately, in addition to whatever is ingested through `/documents`. +The app ships with a small sample knowledge base in `docs_sample/` (18 original short notes, from REST APIs and database transactions to sourdough baking and Roman aqueducts) so it is queryable immediately, in addition to whatever is ingested through `/documents`. The same notes are the corpus of the retrieval evaluation. ## Running it ```bash python -m venv .venv && source .venv/bin/activate pip install -r requirements.txt -cp .env.example .env # set ANTHROPIC_API_KEY -uvicorn app.main:app --reload +cp .env.example .env # set ANTHROPIC_API_KEY, and RETRIEVAL_MODE if you want another mode +uvicorn app.main:app --reload --env-file .env ``` +With the default `hybrid` mode, the first start downloads the embedding model (about 220 MB) and indexes the sample notes. Set `RETRIEVAL_MODE=tfidf` to skip the download entirely. + ```bash curl -X POST localhost:8000/documents -d '{"document_id": "notes", "text": "..."}' curl -X POST localhost:8000/query -d '{"question": "..."}' @@ -38,7 +103,7 @@ pip install -r requirements-dev.txt pytest ``` -The suite covers chunking (including that no word from the source text is lost), retrieval ranking, prompt construction, and the API end to end — all with the Claude call mocked, so no API key or network access is needed to run it. +The suite covers chunking (including that no word from the source text is lost), the three stores, the fusion arithmetic, mode selection, lazy startup, the API contract in every mode, prompt construction, the evaluation metrics and the integrity of the golden set. It runs offline: the Claude call is mocked and a deterministic stand-in replaces the embedding model, so no API key, model download or network access is needed. The real model is only exercised by the evaluation. ## License diff --git a/app/main.py b/app/main.py index 5249b6a..0087e22 100644 --- a/app/main.py +++ b/app/main.py @@ -1,27 +1,31 @@ """FastAPI application exposing the RAG pipeline: ingest documents, then ask questions.""" -import os +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager from fastapi import FastAPI +from app.pipeline_factory import build_default_pipeline from app.schemas import IngestRequest, IngestResponse, QueryRequest, QueryResponse, SourceOut -from rag.generation import ClaudeAnswerGenerator from rag.pipeline import RagPipeline -from rag.sample_corpus import load_sample_documents +APP_VERSION = "1.1.0" -def create_app(pipeline: RagPipeline | None = None) -> FastAPI: - app = FastAPI(title="RAG Knowledge Assistant", version="1.0.0") - if pipeline is None: - generator = ClaudeAnswerGenerator( - api_key=os.environ.get("ANTHROPIC_API_KEY", ""), - model=os.environ.get("ANTHROPIC_MODEL", "claude-haiku-4-5"), - ) - pipeline = RagPipeline(generator=generator) - for document_id, text in load_sample_documents().items(): - pipeline.ingest(document_id, text) +@asynccontextmanager +async def _build_pipeline_on_startup(app: FastAPI) -> AsyncIterator[None]: + # Building the pipeline loads the embedding model and connects to Qdrant, so + # it happens here and never at import time. + if app.state.pipeline is None: + app.state.pipeline = build_default_pipeline() + yield + +def create_app(pipeline: RagPipeline | None = None) -> FastAPI: + """Create the API around `pipeline`, or around the environment-configured one at startup.""" + app = FastAPI( + title="RAG Knowledge Assistant", version=APP_VERSION, lifespan=_build_pipeline_on_startup + ) app.state.pipeline = pipeline @app.post("/documents", response_model=IngestResponse) diff --git a/app/pipeline_factory.py b/app/pipeline_factory.py new file mode 100644 index 0000000..06af9a5 --- /dev/null +++ b/app/pipeline_factory.py @@ -0,0 +1,30 @@ +"""Builds the default pipeline from environment variables when the server starts.""" + +import logging +import os +from collections.abc import Mapping + +from rag.generation import ClaudeAnswerGenerator +from rag.pipeline import RagPipeline +from rag.retrieval_settings import load_retrieval_settings +from rag.sample_corpus import load_sample_documents +from rag.store_factory import build_vector_store + +DEFAULT_ANTHROPIC_MODEL = "claude-haiku-4-5" + +logger = logging.getLogger(__name__) + + +def build_default_pipeline(environ: Mapping[str, str] = os.environ) -> RagPipeline: + """Create the pipeline, load the retrieval store it needs, and ingest the sample corpus.""" + retrieval_settings = load_retrieval_settings(environ) + logger.info("Starting with retrieval mode %s", retrieval_settings.mode) + + generator = ClaudeAnswerGenerator( + api_key=environ.get("ANTHROPIC_API_KEY", ""), + model=environ.get("ANTHROPIC_MODEL", DEFAULT_ANTHROPIC_MODEL), + ) + pipeline = RagPipeline(generator=generator, store=build_vector_store(retrieval_settings)) + for document_id, text in load_sample_documents().items(): + pipeline.ingest(document_id, text) + return pipeline diff --git a/app/schemas.py b/app/schemas.py index 6b3f90e..bcf25e2 100644 --- a/app/schemas.py +++ b/app/schemas.py @@ -2,6 +2,8 @@ from pydantic import BaseModel +from rag.vector_store import DEFAULT_TOP_K + class IngestRequest(BaseModel): document_id: str @@ -15,7 +17,7 @@ class IngestResponse(BaseModel): class QueryRequest(BaseModel): question: str - top_k: int = 3 + top_k: int = DEFAULT_TOP_K class SourceOut(BaseModel): diff --git a/docker-compose.yml b/docker-compose.yml new file mode 100644 index 0000000..3142014 --- /dev/null +++ b/docker-compose.yml @@ -0,0 +1,10 @@ +services: + qdrant: + image: qdrant/qdrant:v1.19.1 + ports: + - "6333:6333" + volumes: + - qdrant_storage:/qdrant/storage + +volumes: + qdrant_storage: diff --git a/docs/architecture.md b/docs/architecture.md index 4cff0ef..96ad80a 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -3,17 +3,45 @@ ```mermaid flowchart LR Ingest[POST /documents] --> Chunk[split_into_chunks] - Chunk --> Store[TfidfVectorStore] - Query[POST /query] --> Search[store.search] - Store --> Search + Chunk --> Add[VectorStore.add_documents] + Query[POST /query] --> Search[VectorStore.search] Search --> Prompt[build_grounded_prompt] Prompt --> Claude[ClaudeAnswerGenerator] Claude --> Response[answer + cited sources] + + Add --> Mode{RETRIEVAL_MODE} + Search --> Mode + Mode -->|tfidf| Tfidf[TfidfVectorStore] + Mode -->|semantic| Semantic[QdrantVectorStore] + Mode -->|hybrid| Hybrid[HybridVectorStore] + Hybrid --> Tfidf + Hybrid --> Semantic + Hybrid --> Fusion[reciprocal_rank_fusion] + Semantic --> Embedder[Embedder: fastembed, ONNX] + Semantic --> Qdrant[(Qdrant)] ``` ## Why generation is separate from retrieval -`rag/vector_store.py` and `rag/generation.py` know nothing about each other. `RagPipeline` is the only piece that wires them together, which is what makes each one testable on its own: the store is tested with real TF-IDF math and no network calls, and the generator is tested with a mocked HTTP client and no real store. +`rag/vector_store.py` and `rag/generation.py` know nothing about each other beyond the `SearchResult` they exchange. `RagPipeline` is the only piece that wires them together, which is what makes each one testable on its own: the stores are tested without a model or a network, and the generator is tested with a mocked HTTP client and no real store. + +## One retrieval contract, three modes + +Every store implements the same two methods, `add_documents(chunks)` and `search(query, top_k)`, declared as the `VectorStore` protocol in `rag/vector_store.py`. `RagPipeline` depends on that protocol and nothing else, and `rag/store_factory.py` is the only place that knows which concrete class serves which mode: + +- **`tfidf`** is `TfidfVectorStore`: cosine similarity over TF-IDF vectors. It is the lexical baseline, runs entirely in process and needs no model download. +- **`semantic`** is `QdrantVectorStore`: chunks are embedded and indexed in Qdrant, and a query is embedded the same way and matched by cosine similarity. The embedding model sits behind the `Embedder` protocol in `rag/embeddings.py`, so the store never depends on a specific model; the real implementation uses fastembed, which runs ONNX models and keeps PyTorch out of the install. The Qdrant client is injected: in memory for tests and quick starts, a server through `QDRANT_URL` for real runs. +- **`hybrid`** is `HybridVectorStore`: it asks the lexical and the semantic store for their own rankings and merges them with Reciprocal Rank Fusion. + +The mode comes from `RETRIEVAL_MODE` (default `hybrid`). The score on each returned source is mode specific: a TF-IDF cosine, an embedding cosine, or a fused rank score. Scores are only comparable within one response. + +## Why Reciprocal Rank Fusion + +TF-IDF cosine and embedding cosine live on different scales, so adding or averaging them needs a normalization step that quietly depends on the data. Reciprocal Rank Fusion sidesteps this: a chunk's fused score is the sum of `1 / (60 + rank)` over the rankings it appears in, so only the order matters. A chunk both stores agree on beats a chunk only one of them likes, and a store that finds nothing (TF-IDF on a pure paraphrase, for example) simply contributes nothing. The constant 60 is the value commonly used for RRF and has not been tuned on this project's data. Each store contributes its top 20 candidates, so a chunk ranked fifth by both can outrank one ranked first by only one. + +## Why initialization is lazy + +Building the semantic store downloads the embedding model on first use and connects to Qdrant, neither of which belongs in an import. `app.main` therefore builds the default pipeline in the application's startup hook, and `create_app(pipeline=...)` accepts a ready-made pipeline, which is how the tests run without any model. Importing `app.main` stays fast and offline. ## Why every answer carries its sources @@ -21,4 +49,10 @@ A RAG answer that cannot be traced back to what it was grounded in is not meanin ## Where this would need to change for production scale -`TfidfVectorStore` refits its vectorizer on every ingest call and keeps every chunk in memory, which is the right tradeoff for a small, testable demo and the wrong one for a large, frequently updated corpus. A production version would swap in a dense embedding model behind the same `add_documents` / `search` interface and move the index into a real vector database — nothing above that interface would need to know. +The semantic half already runs on a vector database; the lexical half and the surrounding bookkeeping do not. + +- `TfidfVectorStore` refits its vectorizer over every chunk on each ingest call and keeps everything in memory, which is fine for a small demo and wrong for a large, frequently updated corpus. Qdrant can hold sparse vectors next to dense ones and fuse them on the server, which would retire the in-process lexical index. +- `QdrantVectorStore` recreates its collection when the app starts, so the lexical and semantic indexes always describe the same chunks. That treats Qdrant as an index rebuilt from the sample corpus, not as the system of record. A production version would keep chunks in a durable store and re-index from it, which also covers re-embedding the corpus when the model changes. +- Embedding runs inside the request that ingests a document and inside the query path. At scale, ingestion belongs in a background worker with batched embedding calls. +- Semantic and hybrid retrieval always return the closest chunks, even for a question the corpus cannot answer; there is no similarity threshold, because a sensible one depends on the model and the corpus. The grounded prompt is what tells the model to say when the sources do not contain the answer. +- The evaluation set is 40 questions over 18 short notes. It is enough to compare the modes here and to catch a regression, not to predict quality on a real corpus. diff --git a/docs_sample/coffee_extraction.md b/docs_sample/coffee_extraction.md new file mode 100644 index 0000000..ae6ddb5 --- /dev/null +++ b/docs_sample/coffee_extraction.md @@ -0,0 +1,7 @@ +# Coffee Extraction + +Brewing coffee is the act of dissolving flavor compounds out of ground beans into hot water. Acids dissolve first, then sugars, and bitter compounds last, which explains the two classic failures: under-extracted coffee tastes sour and thin, while over-extracted coffee tastes bitter and harsh. + +Four variables control the result. Grind size sets how much surface area the water can reach: a finer grind extracts faster, a coarser one more slowly. Water temperature, usually between 92 and 96 degrees Celsius for filter brewing, sets how aggressively compounds dissolve. Brew time is how long water and grounds stay in contact. The brew ratio is the weight of coffee relative to water, and 1:16, or 1 gram of coffee for every 16 grams of water, is a common starting point. + +When a cup tastes sour, grind finer or brew longer; when it tastes bitter, go coarser or shorten the brew. Weighing the coffee and the water makes these adjustments repeatable. diff --git a/docs_sample/compound_interest.md b/docs_sample/compound_interest.md new file mode 100644 index 0000000..d31f696 --- /dev/null +++ b/docs_sample/compound_interest.md @@ -0,0 +1,7 @@ +# Compound Interest + +Compound interest is interest earned on both the original amount and on the interest already added. A deposit of 1,000 at 7 percent a year grows to about 1,967 after ten years, whereas simple interest, which only ever pays on the original amount, would give 1,700. The gap widens every year because the growing balance is what earns the next payment. + +The rule of 72 gives a quick estimate of how long money takes to double: divide 72 by the annual rate in percent. At 6 percent the answer is about twelve years, and at 9 percent about eight. The rule is an approximation that works best for rates between roughly 6 and 10 percent. + +How often interest is compounded, monthly rather than yearly, matters far less than the rate and the time. Time is the strongest lever, so starting early beats contributing more later. The effect works against borrowers too: a credit card balance that compounds at a high rate can grow faster than the minimum payments shrink it. diff --git a/docs_sample/database_transactions.md b/docs_sample/database_transactions.md new file mode 100644 index 0000000..33da337 --- /dev/null +++ b/docs_sample/database_transactions.md @@ -0,0 +1,7 @@ +# Database Transactions and Isolation + +A transaction groups several statements into one unit of work that either completes entirely or leaves no trace. The guarantees are usually summarized as ACID: atomicity, consistency, isolation, and durability. Atomicity is what protects a bank transfer: if the system crashes after taking money out of one account but before putting it into the other, the database rolls the whole transfer back instead of leaving the money in limbo. Durability means that once a commit is acknowledged, the change survives a crash or a power loss. + +Isolation decides how much concurrent transactions can see of each other's unfinished work. Weak levels allow anomalies such as dirty reads, where a transaction reads data another one has not committed yet, and phantom rows, where a repeated query suddenly returns extra matches. The strongest level, SERIALIZABLE, makes concurrent transactions behave as if they ran one after another, at the price of more conflicts and lower throughput. + +When two transactions each hold a lock the other needs, the result is a deadlock. Databases detect it and abort one of them, so application code should be ready to retry, and transactions should stay short to keep the window for conflicts small. diff --git a/docs_sample/docker_images.md b/docs_sample/docker_images.md new file mode 100644 index 0000000..2d6da10 --- /dev/null +++ b/docs_sample/docker_images.md @@ -0,0 +1,7 @@ +# Docker Images and Layers + +A Docker image is a read-only template built from a Dockerfile, and a container is a running instance of it. Every instruction in the Dockerfile produces a layer, and layers are cached: if an instruction and everything before it are unchanged, the builder reuses the stored layer instead of running it again. + +That makes instruction order a performance decision. Steps that rarely change, such as installing dependencies, should come first, and steps that change on every commit, such as copying the application source, should come last. Copying only requirements.txt, installing the packages, and then copying the rest of the code means that editing a source file does not invalidate the slow install layer. + +Two habits keep images small. A multi-stage build compiles or installs tools in one stage and copies only the finished artifacts into a slim final image. A .dockerignore file keeps things like .git folders and virtual environments out of the build context so they never enter any layer. diff --git a/docs_sample/git_branching.md b/docs_sample/git_branching.md new file mode 100644 index 0000000..b2be38c --- /dev/null +++ b/docs_sample/git_branching.md @@ -0,0 +1,7 @@ +# Git Branching Strategies + +A branching strategy decides how work moves from an idea into the main line of a repository. In trunk-based development, everyone integrates into main at least once a day through short-lived branches, and unfinished features hide behind feature flags. The opposite habit, keeping a feature on a long-lived branch for weeks, lets the branch drift away from main and turns the eventual merge into a painful afternoon of conflicts. + +There are two ways to bring a branch up to date. A merge adds a merge commit that ties the two histories together, preserving exactly what happened. A git rebase instead replays the branch's commits on top of the latest main, producing a straight, linear history, but it rewrites commit hashes. For that reason, commits that other people have already pulled should never be rebased. + +A squash merge is a third option: it collapses all the commits of a branch into one before landing it on main, trading detailed history for a tidy log. diff --git a/docs_sample/marathon_training.md b/docs_sample/marathon_training.md new file mode 100644 index 0000000..b78b60f --- /dev/null +++ b/docs_sample/marathon_training.md @@ -0,0 +1,7 @@ +# Marathon Training Basics + +Most marathon plans run between sixteen and twenty weeks and raise weekly mileage gradually, commonly by no more than about ten percent from one week to the next. The cornerstone of the week is the long run, which grows by a few kilometers at a time and peaks at around thirty to thirty-two kilometers (roughly twenty miles) about three weeks before race day. + +Most of the running should be easy, at a pace where you could hold a conversation, with only a minority of the weekly volume spent on tempo runs and intervals. This split, often called the 80/20 approach, builds endurance while limiting injury risk. + +In the final two to three weeks, the taper cuts volume substantially while keeping some short, quicker running, so the legs arrive at the start line fresh rather than fatigued. Fueling is trainable as well: long runs are the time to practice taking in thirty to sixty grams of carbohydrate per hour, which delays the sudden exhaustion runners call hitting the wall. diff --git a/docs_sample/message_queues.md b/docs_sample/message_queues.md new file mode 100644 index 0000000..b39fe9d --- /dev/null +++ b/docs_sample/message_queues.md @@ -0,0 +1,7 @@ +# Message Queues and Delivery Guarantees + +A message queue lets one part of a system hand work to another without waiting for it. The producer drops a message on the queue and moves on, and consumers pick messages up at their own pace, which smooths out sudden traffic spikes that would otherwise overwhelm the downstream service. + +Brokers offer different delivery guarantees. At-most-once delivery never duplicates a message but may lose one; at-least-once delivery never loses a message but may deliver it more than once; exactly-once delivery is expensive and is usually approximated by combining at-least-once delivery with deduplication. Because at-least-once is the common default, consumers should be idempotent: handling the same message twice must have the same effect as handling it once, which is often achieved by storing an idempotency key for every processed message. + +Messages that keep failing should not block the queue forever. After a limited number of retries, a broker can move them to a dead-letter queue, where they wait for a person or a repair job to inspect them. diff --git a/docs_sample/oauth_authorization.md b/docs_sample/oauth_authorization.md new file mode 100644 index 0000000..7ad66e5 --- /dev/null +++ b/docs_sample/oauth_authorization.md @@ -0,0 +1,7 @@ +# OAuth 2.0 in Plain Terms + +OAuth 2.0 lets a person give an application limited access to their data on another service without ever handing the application their password. In the authorization code flow, the application sends the person to the provider, the person approves a list of scopes such as "read contacts", and the provider returns a short-lived authorization code that the application exchanges for an access token. + +Mobile and single-page apps cannot keep a client secret safe, so they add PKCE. Before starting, the app generates a random code verifier and sends only a hash of it; when it later redeems the code, it must present the original verifier, which makes a stolen code useless to anyone else. + +Access tokens expire quickly on purpose. When an API receives an expired token it answers 401, and the client can use its refresh token to obtain a new access token without sending the person through the login screen again. Refresh tokens are long-lived, so they deserve stricter storage than access tokens. diff --git a/docs_sample/photosynthesis.md b/docs_sample/photosynthesis.md new file mode 100644 index 0000000..25af2f2 --- /dev/null +++ b/docs_sample/photosynthesis.md @@ -0,0 +1,7 @@ +# Photosynthesis + +Photosynthesis is how plants, algae, and some bacteria turn light into chemical energy stored in sugar. In plants it happens inside chloroplasts, organelles filled with the pigment chlorophyll. Chlorophyll absorbs mostly red and blue light and reflects green, which is why leaves look green. + +The process has two stages. In the light-dependent reactions, which take place in the thylakoid membranes, light energy splits water molecules, releasing oxygen as a by-product and producing the energy carriers ATP and NADPH. In the Calvin cycle, which takes place in the fluid-filled stroma, those carriers power the fixation of carbon dioxide into sugar. Summed up, carbon dioxide and water plus light energy become glucose and oxygen. + +Leaves take in carbon dioxide through tiny pores called stomata. Opening them also lets water vapor escape, so a plant constantly balances the need for carbon dioxide against the risk of drying out. diff --git a/docs_sample/rate_limiting.md b/docs_sample/rate_limiting.md new file mode 100644 index 0000000..dd2fa5e --- /dev/null +++ b/docs_sample/rate_limiting.md @@ -0,0 +1,7 @@ +# Rate Limiting + +A rate limiter caps how many requests a client may send in a given period. It protects a service from overload and stops one noisy client from starving everyone else. + +The token bucket is a popular algorithm: a bucket refills with tokens at a steady rate, every request spends one, and a request that finds the bucket empty is rejected. Because the bucket can hold a reserve, short bursts are allowed while the long-term average stays capped. A fixed-window counter is simpler but lets a client send a double burst across the boundary between two windows. + +A rejected request should receive HTTP 429 Too Many Requests, ideally with a Retry-After header telling the client how many seconds to wait. Well-behaved clients respond by backing off exponentially and adding random jitter, so that thousands of them do not all retry at the same instant and recreate the spike. diff --git a/docs_sample/roman_aqueducts.md b/docs_sample/roman_aqueducts.md new file mode 100644 index 0000000..a4ee83b --- /dev/null +++ b/docs_sample/roman_aqueducts.md @@ -0,0 +1,7 @@ +# Roman Aqueducts + +Roman engineers carried water to cities over long distances using nothing but gravity. The channel was laid with a gentle, steady downward slope, in some cases only a few tenths of a percent, so the water kept flowing without ever rushing. Most of the route ran underground or along the surface, where it was cheaper and easier to maintain. + +Where the land dipped, arches lifted the channel across valleys. The Pont du Gard in southern France, about 49 meters tall and built in the first century, is among the best preserved examples. Settling tanks along the way let sediment drop out so the water arriving in the city was clearer. + +The first aqueduct serving Rome, the Aqua Appia, was completed in 312 BC, and by the third century AD the city was supplied by eleven of them. The water fed public fountains, baths, and some wealthy households, and the surplus flushed the sewers. diff --git a/docs_sample/sleep_hygiene.md b/docs_sample/sleep_hygiene.md new file mode 100644 index 0000000..a732117 --- /dev/null +++ b/docs_sample/sleep_hygiene.md @@ -0,0 +1,7 @@ +# Sleep Hygiene + +The circadian rhythm is the body's internal clock, which runs on a cycle of roughly twenty-four hours and is set mainly by light. Bright light soon after waking anchors the clock, while bright screens late at night delay the release of melatonin and push bedtime later. A consistent wake-up time does more for sleep quality than a consistent bedtime. + +Caffeine has a half-life of about five to six hours, which means a cup consumed late in the afternoon can still be half as strong when you go to bed. Alcohol makes it easier to fall asleep but fragments the second half of the night and reduces REM sleep, the stage most associated with vivid dreaming and memory consolidation. + +A cool, dark bedroom, around 18 degrees Celsius, helps the body lower its core temperature. If you are still awake after about twenty minutes, it is better to get up and do something calm in dim light than to lie there watching the clock. diff --git a/docs_sample/sourdough_baking.md b/docs_sample/sourdough_baking.md new file mode 100644 index 0000000..3394590 --- /dev/null +++ b/docs_sample/sourdough_baking.md @@ -0,0 +1,7 @@ +# Sourdough Basics + +A sourdough starter is a living mixture of flour and water inhabited by wild yeast and lactic acid bacteria. To keep it healthy, a baker feeds it regularly: discard most of it, then add fresh flour and water. A starter is ready to use when it has roughly doubled in volume and is full of bubbles. + +Hydration is the weight of water divided by the weight of flour. A dough at 75 percent hydration is wetter than one at 65 percent, which gives a more open crumb but is stickier and harder to shape. After mixing, the dough goes through bulk fermentation, its first long rise, which can take four to eight hours at room temperature. Warmth speeds fermentation up and cold slows it down, which is why many bakers shape the loaf and then proof it overnight in the refrigerator. + +Baking in a preheated Dutch oven traps steam around the loaf, letting it expand fully before the crust hardens. diff --git a/docs_sample/volcano_types.md b/docs_sample/volcano_types.md new file mode 100644 index 0000000..287ff6c --- /dev/null +++ b/docs_sample/volcano_types.md @@ -0,0 +1,7 @@ +# Types of Volcanoes + +Volcanoes are usually grouped by the shape they build, and that shape reflects the kind of magma behind it. Shield volcanoes, such as Mauna Loa in Hawaii, are broad and gently sloping because they are built from thin, runny basaltic lava that flows a long way before it cools. Their eruptions tend to be effusive, with lava pouring out rather than exploding. + +Stratovolcanoes, also called composite volcanoes, are steep, symmetrical cones made of alternating layers of lava and ash. Their magma is thick and sticky and traps gas, so pressure builds until the eruption becomes explosive. Mount Fuji and Mount St. Helens, which erupted violently in 1980, are stratovolcanoes. + +Cinder cones are the smallest type: steep piles of loose fragments thrown out around a single vent. They often erupt only once and stop within a few years. A caldera forms when a very large eruption empties the magma chamber and the ground above collapses into a wide basin. diff --git a/evaluation/__init__.py b/evaluation/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/evaluation/golden_set.json b/evaluation/golden_set.json new file mode 100644 index 0000000..5aec24e --- /dev/null +++ b/evaluation/golden_set.json @@ -0,0 +1,45 @@ +{ + "version": 1, + "questions": [ + {"id": "lexical-01", "kind": "lexical_overlap", "question": "Why does the column order matter in a composite index?", "relevant_documents": ["database_indexing"]}, + {"id": "lexical-02", "kind": "lexical_overlap", "question": "What is the difference between cache-aside and write-through caching?", "relevant_documents": ["caching_strategies"]}, + {"id": "lexical-03", "kind": "lexical_overlap", "question": "How many end-to-end tests should a healthy test suite have compared to unit tests?", "relevant_documents": ["testing_pyramid"]}, + {"id": "lexical-04", "kind": "lexical_overlap", "question": "Which HTTP methods does a REST API use to read, create, update and delete a resource?", "relevant_documents": ["rest_apis"]}, + {"id": "lexical-05", "kind": "lexical_overlap", "question": "What does isolation mean in a database transaction?", "relevant_documents": ["database_transactions"]}, + {"id": "lexical-06", "kind": "lexical_overlap", "question": "What is a dead-letter queue used for?", "relevant_documents": ["message_queues"]}, + {"id": "lexical-07", "kind": "lexical_overlap", "question": "What is the difference between a merge and a rebase in git?", "relevant_documents": ["git_branching"]}, + {"id": "lexical-08", "kind": "lexical_overlap", "question": "What does a refresh token do in OAuth 2.0?", "relevant_documents": ["oauth_authorization"]}, + {"id": "lexical-09", "kind": "lexical_overlap", "question": "Why does the order of instructions in a Dockerfile affect layer caching?", "relevant_documents": ["docker_images"]}, + {"id": "lexical-10", "kind": "lexical_overlap", "question": "What does hydration mean in sourdough baking?", "relevant_documents": ["sourdough_baking"]}, + {"id": "lexical-11", "kind": "lexical_overlap", "question": "Where in the plant cell does the Calvin cycle take place?", "relevant_documents": ["photosynthesis"]}, + {"id": "lexical-12", "kind": "lexical_overlap", "question": "How long should a marathon taper last?", "relevant_documents": ["marathon_training"]}, + {"id": "lexical-13", "kind": "lexical_overlap", "question": "How does grind size affect coffee extraction?", "relevant_documents": ["coffee_extraction"]}, + {"id": "lexical-14", "kind": "lexical_overlap", "question": "What kind of lava builds a shield volcano?", "relevant_documents": ["volcano_types"]}, + {"id": "paraphrase-01", "kind": "paraphrase", "question": "Why shouldn't I add a lookup structure to every field of my table?", "relevant_documents": ["database_indexing"]}, + {"id": "paraphrase-02", "kind": "paraphrase", "question": "Why is it risky to rely mostly on slow checks that drive the whole application like a real user?", "relevant_documents": ["testing_pyramid"]}, + {"id": "paraphrase-03", "kind": "paraphrase", "question": "Why is it bad practice for a web service to report success while hiding the failure inside the payload?", "relevant_documents": ["rest_apis"]}, + {"id": "paraphrase-04", "kind": "paraphrase", "question": "How can a payment system guarantee that a half-finished operation never leaves customer balances inconsistent?", "relevant_documents": ["database_transactions"]}, + {"id": "paraphrase-05", "kind": "paraphrase", "question": "How can an application stay safe if its work items get delivered to it again by mistake?", "relevant_documents": ["message_queues"]}, + {"id": "paraphrase-06", "kind": "paraphrase", "question": "What goes wrong when a team develops a big change separately for a long time before combining it with everyone else's work?", "relevant_documents": ["git_branching"]}, + {"id": "paraphrase-07", "kind": "paraphrase", "question": "How can an app on my phone be allowed to act on my behalf at a third-party service without learning my credentials?", "relevant_documents": ["oauth_authorization"]}, + {"id": "paraphrase-08", "kind": "paraphrase", "question": "What should a client do when a server tells it to slow down because it is sending too many calls?", "relevant_documents": ["rate_limiting"]}, + {"id": "paraphrase-09", "kind": "paraphrase", "question": "Why does an espresso after lunch sometimes keep me awake until midnight?", "relevant_documents": ["sleep_hygiene"]}, + {"id": "paraphrase-10", "kind": "paraphrase", "question": "Is there a mental shortcut to work out the time needed for an investment to become twice as large?", "relevant_documents": ["compound_interest"]}, + {"id": "paraphrase-11", "kind": "paraphrase", "question": "How did ancient builders deliver fresh supplies to their cities without any machinery?", "relevant_documents": ["roman_aqueducts"]}, + {"id": "paraphrase-12", "kind": "paraphrase", "question": "What makes the foliage of trees appear the way it does, rather than red or blue?", "relevant_documents": ["photosynthesis"]}, + {"id": "paraphrase-13", "kind": "paraphrase", "question": "Why does my homemade bread take hours longer to be ready in winter?", "relevant_documents": ["sourdough_baking"]}, + {"id": "paraphrase-14", "kind": "paraphrase", "question": "Which techniques stop a system from being swamped when a huge number of requests arrives at once?", "relevant_documents": ["rate_limiting", "message_queues"]}, + {"id": "exact-01", "kind": "exact_term", "question": "HTTP 429 Retry-After", "relevant_documents": ["rate_limiting"]}, + {"id": "exact-02", "kind": "exact_term", "question": "SERIALIZABLE", "relevant_documents": ["database_transactions"]}, + {"id": "exact-03", "kind": "exact_term", "question": "PKCE code verifier", "relevant_documents": ["oauth_authorization"]}, + {"id": "exact-04", "kind": "exact_term", "question": "Pont du Gard", "relevant_documents": ["roman_aqueducts"]}, + {"id": "exact-05", "kind": "exact_term", "question": "Rule of 72", "relevant_documents": ["compound_interest"]}, + {"id": "exact-06", "kind": "exact_term", "question": "TTL", "relevant_documents": ["caching_strategies"]}, + {"id": "exact-07", "kind": "exact_term", "question": "B-tree", "relevant_documents": ["database_indexing"]}, + {"id": "exact-08", "kind": "exact_term", "question": "(user_id, created_at)", "relevant_documents": ["database_indexing"]}, + {"id": "exact-09", "kind": "exact_term", "question": ".dockerignore multi-stage build", "relevant_documents": ["docker_images"]}, + {"id": "exact-10", "kind": "exact_term", "question": "401 status code", "relevant_documents": ["rest_apis", "oauth_authorization"]}, + {"id": "exact-11", "kind": "exact_term", "question": "REM sleep", "relevant_documents": ["sleep_hygiene"]}, + {"id": "exact-12", "kind": "exact_term", "question": "Mount St. Helens", "relevant_documents": ["volcano_types"]} + ] +} diff --git a/evaluation/golden_set.py b/evaluation/golden_set.py new file mode 100644 index 0000000..732075f --- /dev/null +++ b/evaluation/golden_set.py @@ -0,0 +1,85 @@ +"""The retrieval golden set: questions labeled with the documents that answer them.""" + +import json +from collections import Counter +from collections.abc import Collection, Mapping +from dataclasses import dataclass +from enum import StrEnum +from pathlib import Path + +GOLDEN_SET_PATH = Path(__file__).resolve().parent / "golden_set.json" +REQUIRED_FIELDS = ("id", "kind", "question", "relevant_documents") + + +class QuestionKind(StrEnum): + LEXICAL_OVERLAP = "lexical_overlap" + PARAPHRASE = "paraphrase" + EXACT_TERM = "exact_term" + + +@dataclass(frozen=True) +class GoldenQuestion: + question_id: str + kind: QuestionKind + text: str + relevant_documents: frozenset[str] + + +@dataclass(frozen=True) +class GoldenSet: + version: int + questions: list[GoldenQuestion] + + +class GoldenSetError(ValueError): + """Raised when the golden set file is malformed or labels a document that does not exist.""" + + +def load_golden_set(path: Path = GOLDEN_SET_PATH) -> GoldenSet: + contents = json.loads(path.read_text(encoding="utf-8")) + if "version" not in contents or "questions" not in contents: + raise GoldenSetError(f"{path.name} needs a 'version' and a 'questions' list") + questions = [_parse_question(entry) for entry in contents["questions"]] + _reject_duplicate_ids(questions) + return GoldenSet(version=contents["version"], questions=questions) + + +def ensure_labels_exist(golden_set: GoldenSet, document_ids: Collection[str]) -> None: + """Fail loudly when a question is labeled with a document that is not in the corpus.""" + for question in golden_set.questions: + unknown_documents = sorted(question.relevant_documents - set(document_ids)) + if unknown_documents: + raise GoldenSetError( + f"Question {question.question_id} is labeled with unknown documents: " + f"{unknown_documents}" + ) + + +def _parse_question(entry: Mapping[str, object]) -> GoldenQuestion: + missing_fields = [name for name in REQUIRED_FIELDS if name not in entry] + if missing_fields: + raise GoldenSetError(f"Golden set entry {dict(entry)} is missing {missing_fields}") + + question_id = str(entry["id"]) + try: + kind = QuestionKind(entry["kind"]) + except ValueError: + raise GoldenSetError(f"Question {question_id} has unknown kind {entry['kind']!r}") from None + + relevant_documents = frozenset(entry["relevant_documents"]) + if not relevant_documents: + raise GoldenSetError(f"Question {question_id} has no relevant document") + + return GoldenQuestion( + question_id=question_id, + kind=kind, + text=str(entry["question"]), + relevant_documents=relevant_documents, + ) + + +def _reject_duplicate_ids(questions: list[GoldenQuestion]) -> None: + id_counts = Counter(question.question_id for question in questions) + repeated_ids = sorted(question_id for question_id, count in id_counts.items() if count > 1) + if repeated_ids: + raise GoldenSetError(f"Duplicate question ids: {repeated_ids}") diff --git a/evaluation/metrics.py b/evaluation/metrics.py new file mode 100644 index 0000000..46d087f --- /dev/null +++ b/evaluation/metrics.py @@ -0,0 +1,61 @@ +"""Retrieval metrics computed from ranked document ids and the labeled relevant documents.""" + +from collections.abc import Sequence +from dataclasses import dataclass + +HIT_RATE_CUTOFF = 3 +RANKING_DEPTH = 10 + + +@dataclass(frozen=True) +class RetrievalOutcome: + """What one store returned for one question, next to the documents that answer it. + + `ranked_document_ids` holds one entry per retrieved chunk, best first, so a + document can appear more than once. + """ + + ranked_document_ids: tuple[str, ...] + relevant_documents: frozenset[str] + + +@dataclass(frozen=True) +class RetrievalScores: + hits: int + questions: int + mean_reciprocal_rank: float + + @property + def hit_rate(self) -> float: + return self.hits / self.questions + + +def hit_at_k(outcome: RetrievalOutcome, k: int) -> bool: + """Whether a relevant document appears among the first `k` retrieved chunks.""" + return any( + document_id in outcome.relevant_documents for document_id in outcome.ranked_document_ids[:k] + ) + + +def reciprocal_rank(outcome: RetrievalOutcome) -> float: + """1 divided by the position of the first relevant chunk, or 0 if none was retrieved. + + Rankings are cut at `RANKING_DEPTH` chunks, so a relevant document ranked + deeper than that scores 0. + """ + for position, document_id in enumerate(outcome.ranked_document_ids, start=1): + if document_id in outcome.relevant_documents: + return 1.0 / position + return 0.0 + + +def score_outcomes( + outcomes: Sequence[RetrievalOutcome], k: int = HIT_RATE_CUTOFF +) -> RetrievalScores: + if not outcomes: + raise ValueError("Cannot score zero questions") + return RetrievalScores( + hits=sum(hit_at_k(outcome, k) for outcome in outcomes), + questions=len(outcomes), + mean_reciprocal_rank=sum(reciprocal_rank(outcome) for outcome in outcomes) / len(outcomes), + ) diff --git a/evaluation/report.py b/evaluation/report.py new file mode 100644 index 0000000..67f5156 --- /dev/null +++ b/evaluation/report.py @@ -0,0 +1,92 @@ +"""The shape of an evaluation report and its Markdown rendering.""" + +from collections.abc import Sequence +from dataclasses import dataclass + +from evaluation.golden_set import QuestionKind +from evaluation.metrics import HIT_RATE_CUTOFF, RetrievalScores +from rag.retrieval_settings import RetrievalMode + + +@dataclass(frozen=True) +class ModeReport: + mode: RetrievalMode + overall: RetrievalScores + by_kind: dict[QuestionKind, RetrievalScores] + + +@dataclass(frozen=True) +class EvaluationReport: + golden_set_version: int + document_count: int + chunk_count: int + embedding_model: str + modes: list[ModeReport] + + +def render_markdown(report: EvaluationReport) -> str: + question_count = report.modes[0].overall.questions + description = ( + f"Golden set v{report.golden_set_version}: {question_count} questions over " + f"{report.document_count} notes ({report.chunk_count} chunks). " + f"Embedding model: {report.embedding_model}." + ) + hit_rate_label = f"Hit rate@{HIT_RATE_CUTOFF}" + kinds = [kind for kind in QuestionKind if kind in report.modes[0].by_kind] + kind_headers = [ + f"{_kind_label(kind)} ({report.modes[0].by_kind[kind].questions})" for kind in kinds + ] + + overall_table = _markdown_table( + ["Mode", hit_rate_label, "MRR"], + [ + [ + mode_report.mode.value, + _format_hit_rate(mode_report.overall), + f"{mode_report.overall.mean_reciprocal_rank:.3f}", + ] + for mode_report in report.modes + ], + ) + hit_rate_by_kind_table = _markdown_table( + ["Mode", *kind_headers], + [ + [mode_report.mode.value, *(_format_hit_rate(mode_report.by_kind[k]) for k in kinds)] + for mode_report in report.modes + ], + ) + mrr_by_kind_table = _markdown_table( + ["Mode", *kind_headers], + [ + [ + mode_report.mode.value, + *(f"{mode_report.by_kind[k].mean_reciprocal_rank:.3f}" for k in kinds), + ] + for mode_report in report.modes + ], + ) + return "\n\n".join( + [ + description, + overall_table, + f"{hit_rate_label} by question type:", + hit_rate_by_kind_table, + "MRR by question type:", + mrr_by_kind_table, + ] + ) + + +def _format_hit_rate(scores: RetrievalScores) -> str: + return f"{scores.hit_rate:.3f} ({scores.hits}/{scores.questions})" + + +def _kind_label(kind: QuestionKind) -> str: + return kind.value.replace("_", " ").capitalize() + + +def _markdown_table(header: Sequence[str], rows: Sequence[Sequence[str]]) -> str: + separator = ["---"] * len(header) + return "\n".join( + "| " + " | ".join(cells) + " |" for cells in [list(header), separator, *map(list, rows)] + ) diff --git a/evaluation/run.py b/evaluation/run.py new file mode 100644 index 0000000..1b212d4 --- /dev/null +++ b/evaluation/run.py @@ -0,0 +1,158 @@ +"""Evaluates every retrieval mode on the golden set. + +Run `python -m evaluation.run` to print the results. With `--enforce-gate` the +process also exits with status 1 when hybrid retrieval scores a lower hit rate +than the TF-IDF baseline. +""" + +import argparse +import logging +import sys +from collections import defaultdict +from collections.abc import Callable, Mapping, Sequence +from functools import cache + +from evaluation.golden_set import ( + GoldenQuestion, + GoldenSet, + QuestionKind, + ensure_labels_exist, + load_golden_set, +) +from evaluation.metrics import ( + HIT_RATE_CUTOFF, + RANKING_DEPTH, + RetrievalOutcome, + RetrievalScores, + hit_at_k, + score_outcomes, +) +from evaluation.report import EvaluationReport, ModeReport, render_markdown +from rag.chunking import Chunk, split_into_chunks +from rag.embeddings import Embedder, FastEmbedEmbedder +from rag.retrieval_settings import RetrievalMode, RetrievalSettings +from rag.sample_corpus import load_sample_documents +from rag.store_factory import build_vector_store +from rag.vector_store import VectorStore + +EXIT_OK = 0 +EXIT_QUALITY_GATE_FAILED = 1 + +logger = logging.getLogger(__name__) + + +def run_evaluation( + golden_set: GoldenSet, + documents: Mapping[str, str], + embedder_factory: Callable[[], Embedder], +) -> EvaluationReport: + ensure_labels_exist(golden_set, documents.keys()) + chunks = _chunk_corpus(documents) + # The semantic and hybrid modes share one embedder so the model loads once. + shared_embedder_factory = cache(embedder_factory) + mode_reports = [ + _evaluate_mode(mode, chunks, golden_set.questions, shared_embedder_factory) + for mode in RetrievalMode + ] + return EvaluationReport( + golden_set_version=golden_set.version, + document_count=len(documents), + chunk_count=len(chunks), + embedding_model=shared_embedder_factory().model_name, + modes=mode_reports, + ) + + +def hybrid_meets_tfidf_baseline(report: EvaluationReport) -> bool: + hit_rates = {mode_report.mode: mode_report.overall.hit_rate for mode_report in report.modes} + return hit_rates[RetrievalMode.HYBRID] >= hit_rates[RetrievalMode.TFIDF] + + +def main( + argv: Sequence[str] | None = None, + embedder_factory: Callable[[], Embedder] = FastEmbedEmbedder, +) -> int: + logging.basicConfig(level=logging.INFO, format="%(message)s") + arguments = _parse_arguments(argv) + + report = run_evaluation(load_golden_set(), load_sample_documents(), embedder_factory) + sys.stdout.write(render_markdown(report) + "\n") + + if arguments.enforce_gate and not hybrid_meets_tfidf_baseline(report): + logger.error( + "Quality gate failed: hybrid hit rate@%d is below the TF-IDF baseline", HIT_RATE_CUTOFF + ) + return EXIT_QUALITY_GATE_FAILED + return EXIT_OK + + +def _parse_arguments(argv: Sequence[str] | None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Evaluate retrieval modes on the golden set.") + parser.add_argument( + "--enforce-gate", + action="store_true", + help="exit with status 1 if hybrid hit rate falls below the TF-IDF baseline", + ) + return parser.parse_args(argv) + + +def _chunk_corpus(documents: Mapping[str, str]) -> list[Chunk]: + return [ + chunk + for document_id, text in documents.items() + for chunk in split_into_chunks(document_id, text) + ] + + +def _evaluate_mode( + mode: RetrievalMode, + chunks: list[Chunk], + questions: list[GoldenQuestion], + embedder_factory: Callable[[], Embedder], +) -> ModeReport: + store = build_vector_store(RetrievalSettings(mode=mode), embedder_factory) + store.add_documents(chunks) + + outcomes = [ + RetrievalOutcome(_rank_documents(store, question.text), question.relevant_documents) + for question in questions + ] + _log_misses(mode, questions, outcomes) + return ModeReport( + mode=mode, + overall=score_outcomes(outcomes), + by_kind=_score_by_kind(questions, outcomes), + ) + + +def _rank_documents(store: VectorStore, question_text: str) -> tuple[str, ...]: + matches = store.search(question_text, top_k=RANKING_DEPTH) + return tuple(match.chunk.document_id for match in matches) + + +def _score_by_kind( + questions: list[GoldenQuestion], outcomes: list[RetrievalOutcome] +) -> dict[QuestionKind, RetrievalScores]: + outcomes_by_kind: dict[QuestionKind, list[RetrievalOutcome]] = defaultdict(list) + for question, outcome in zip(questions, outcomes, strict=True): + outcomes_by_kind[question.kind].append(outcome) + return {kind: score_outcomes(kind_outcomes) for kind, kind_outcomes in outcomes_by_kind.items()} + + +def _log_misses( + mode: RetrievalMode, questions: list[GoldenQuestion], outcomes: list[RetrievalOutcome] +) -> None: + for question, outcome in zip(questions, outcomes, strict=True): + if not hit_at_k(outcome, HIT_RATE_CUTOFF): + logger.info( + "[%s] miss at %d: %s %r -> %s", + mode.value, + HIT_RATE_CUTOFF, + question.question_id, + question.text, + outcome.ranked_document_ids[:HIT_RATE_CUTOFF], + ) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/rag/embeddings.py b/rag/embeddings.py new file mode 100644 index 0000000..e4e6c57 --- /dev/null +++ b/rag/embeddings.py @@ -0,0 +1,49 @@ +"""Turns text into vectors behind a small interface, so no store depends on one model.""" + +from collections.abc import Sequence +from typing import Protocol + +from fastembed import TextEmbedding + +DEFAULT_EMBEDDING_MODEL = "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2" + + +class Embedder(Protocol): + """Maps text to fixed-size vectors in which closer vectors mean closer meaning.""" + + @property + def model_name(self) -> str: ... + + @property + def dimension(self) -> int: ... + + def embed_passages(self, texts: Sequence[str]) -> list[list[float]]: ... + + def embed_query(self, text: str) -> list[float]: ... + + +class FastEmbedEmbedder: + """ONNX embeddings through fastembed, which keeps PyTorch out of the install. + + Constructing it downloads the model on first use and loads it into memory, so + it must be built at application startup and never at import time. + """ + + def __init__(self, model_name: str = DEFAULT_EMBEDDING_MODEL, cache_dir: str | None = None): + self._model_name = model_name + self._model = TextEmbedding(model_name=model_name, cache_dir=cache_dir) + self._dimension = TextEmbedding.get_embedding_size(model_name) + + @property + def model_name(self) -> str: + return self._model_name + + @property + def dimension(self) -> int: + return self._dimension + + def embed_passages(self, texts: Sequence[str]) -> list[list[float]]: + return [vector.tolist() for vector in self._model.passage_embed(texts)] + + def embed_query(self, text: str) -> list[float]: + return next(iter(self._model.query_embed(text))).tolist() diff --git a/rag/fusion.py b/rag/fusion.py new file mode 100644 index 0000000..287900e --- /dev/null +++ b/rag/fusion.py @@ -0,0 +1,51 @@ +"""Reciprocal Rank Fusion: merges several rankings without comparing their scores.""" + +from collections.abc import Sequence + +from rag.chunking import Chunk +from rag.vector_store import SearchResult + +RRF_RANK_CONSTANT = 60 + + +def reciprocal_rank_fusion( + rankings: Sequence[Sequence[SearchResult]], top_k: int +) -> list[SearchResult]: + """Fuse `rankings` into one list of at most `top_k` results, best first. + + A chunk's fused score is the sum of 1 / (RRF_RANK_CONSTANT + rank) over every + ranking it appears in. Only ranks are used, never the raw scores, which is + what makes rankings from different scorers (TF-IDF cosine, embedding cosine) + safe to combine. A chunk listed twice in one ranking counts once, at its best + position. Equal fused scores keep the order in which the chunks were first + seen, so earlier rankings win ties. + """ + fused_scores: dict[tuple[str, int], float] = {} + chunks_by_key: dict[tuple[str, int], Chunk] = {} + + for ranking in rankings: + for rank, scored_chunk in enumerate(_first_occurrences(ranking), start=1): + key = _chunk_key(scored_chunk.chunk) + fused_scores[key] = fused_scores.get(key, 0.0) + 1.0 / (RRF_RANK_CONSTANT + rank) + chunks_by_key.setdefault(key, scored_chunk.chunk) + + best_first = sorted(fused_scores, key=fused_scores.__getitem__, reverse=True) + return [ + SearchResult(chunk=chunks_by_key[key], score=fused_scores[key]) + for key in best_first[:top_k] + ] + + +def _first_occurrences(ranking: Sequence[SearchResult]) -> list[SearchResult]: + seen: set[tuple[str, int]] = set() + distinct: list[SearchResult] = [] + for scored_chunk in ranking: + key = _chunk_key(scored_chunk.chunk) + if key not in seen: + seen.add(key) + distinct.append(scored_chunk) + return distinct + + +def _chunk_key(chunk: Chunk) -> tuple[str, int]: + return (chunk.document_id, chunk.index) diff --git a/rag/hybrid_store.py b/rag/hybrid_store.py new file mode 100644 index 0000000..ecbf780 --- /dev/null +++ b/rag/hybrid_store.py @@ -0,0 +1,31 @@ +"""Hybrid retrieval: fuses a lexical and a semantic ranking with Reciprocal Rank Fusion.""" + +from rag.chunking import Chunk +from rag.fusion import reciprocal_rank_fusion +from rag.vector_store import DEFAULT_TOP_K, SearchResult, VectorStore + +CANDIDATES_PER_RANKER = 20 + + +class HybridVectorStore: + """Combines two stores behind the common `VectorStore` contract. + + Each store ranks the query on its own and only the ranks are fused, so the + two scoring scales never have to be reconciled. + """ + + def __init__(self, lexical_store: VectorStore, semantic_store: VectorStore): + self._lexical_store = lexical_store + self._semantic_store = semantic_store + + def add_documents(self, chunks: list[Chunk]) -> None: + self._lexical_store.add_documents(chunks) + self._semantic_store.add_documents(chunks) + + def search(self, query: str, top_k: int = DEFAULT_TOP_K) -> list[SearchResult]: + # Each store is asked for more than `top_k` candidates: a chunk ranked + # fifth by both stores can beat one ranked first by only one of them. + candidate_count = max(top_k, CANDIDATES_PER_RANKER) + lexical_ranking = self._lexical_store.search(query, top_k=candidate_count) + semantic_ranking = self._semantic_store.search(query, top_k=candidate_count) + return reciprocal_rank_fusion([lexical_ranking, semantic_ranking], top_k=top_k) diff --git a/rag/pipeline.py b/rag/pipeline.py index 4086837..b648f59 100644 --- a/rag/pipeline.py +++ b/rag/pipeline.py @@ -1,10 +1,10 @@ -"""Wires chunking, the vector store, and answer generation into one pipeline.""" +"""Wires chunking, a vector store, and answer generation into one pipeline.""" from dataclasses import dataclass from rag.chunking import split_into_chunks from rag.generation import ClaudeAnswerGenerator -from rag.vector_store import TfidfVectorStore +from rag.vector_store import DEFAULT_TOP_K, VectorStore @dataclass(frozen=True) @@ -22,8 +22,8 @@ class Answer: class RagPipeline: - def __init__(self, generator: ClaudeAnswerGenerator, store: TfidfVectorStore | None = None): - self._store = store or TfidfVectorStore() + def __init__(self, generator: ClaudeAnswerGenerator, store: VectorStore): + self._store = store self._generator = generator def ingest(self, document_id: str, text: str) -> int: @@ -32,7 +32,7 @@ def ingest(self, document_id: str, text: str) -> int: self._store.add_documents(chunks) return len(chunks) - async def ask(self, question: str, top_k: int = 3) -> Answer: + async def ask(self, question: str, top_k: int = DEFAULT_TOP_K) -> Answer: results = self._store.search(question, top_k=top_k) if not results: return Answer(text="No indexed documents match this question.", sources=[]) diff --git a/rag/qdrant_store.py b/rag/qdrant_store.py new file mode 100644 index 0000000..8df601d --- /dev/null +++ b/rag/qdrant_store.py @@ -0,0 +1,79 @@ +"""Semantic vector store: chunk embeddings indexed and searched in Qdrant.""" + +import uuid +from dataclasses import asdict + +from qdrant_client import QdrantClient +from qdrant_client.models import Distance, PointStruct, ScoredPoint, VectorParams + +from rag.chunking import Chunk +from rag.embeddings import Embedder +from rag.vector_store import DEFAULT_TOP_K, SearchResult + +COLLECTION_NAME = "knowledge_chunks" +UPSERT_BATCH_SIZE = 64 +CHUNK_ID_NAMESPACE = uuid.uuid5(uuid.NAMESPACE_DNS, "rag-knowledge-assistant") + + +class QdrantVectorStore: + """Ranks chunks by cosine similarity between their embeddings and the query's. + + The client decides where Qdrant runs: `QdrantClient(location=":memory:")` for + tests and quick starts, `QdrantClient(url=...)` for a real server. + """ + + def __init__( + self, + embedder: Embedder, + client: QdrantClient, + collection_name: str = COLLECTION_NAME, + ): + self._embedder = embedder + self._client = client + self._collection_name = collection_name + self._recreate_collection() + + def add_documents(self, chunks: list[Chunk]) -> None: + for start in range(0, len(chunks), UPSERT_BATCH_SIZE): + self._upsert_batch(chunks[start : start + UPSERT_BATCH_SIZE]) + + def search(self, query: str, top_k: int = DEFAULT_TOP_K) -> list[SearchResult]: + # Qdrant rejects a limit below 1 while the TF-IDF store answers with an + # empty list, so the stores stay interchangeable only if this one does too. + if top_k < 1: + return [] + response = self._client.query_points( + collection_name=self._collection_name, + query=self._embedder.embed_query(query), + limit=top_k, + ) + return [_to_search_result(point) for point in response.points] + + def _recreate_collection(self) -> None: + # Starting empty keeps the lexical and semantic indexes describing the + # same chunks after a restart; the collection is an index, not the + # system of record. + if self._client.collection_exists(self._collection_name): + self._client.delete_collection(self._collection_name) + self._client.create_collection( + collection_name=self._collection_name, + vectors_config=VectorParams(size=self._embedder.dimension, distance=Distance.COSINE), + ) + + def _upsert_batch(self, chunks: list[Chunk]) -> None: + vectors = self._embedder.embed_passages([chunk.text for chunk in chunks]) + points = [ + PointStruct(id=_point_id(chunk), vector=vector, payload=asdict(chunk)) + for chunk, vector in zip(chunks, vectors, strict=True) + ] + self._client.upsert(collection_name=self._collection_name, points=points) + + +def _point_id(chunk: Chunk) -> str: + # Deterministic ids make re-ingesting the same chunk overwrite it instead of + # storing a duplicate. + return str(uuid.uuid5(CHUNK_ID_NAMESPACE, f"{chunk.document_id}#{chunk.index}")) + + +def _to_search_result(point: ScoredPoint) -> SearchResult: + return SearchResult(chunk=Chunk(**point.payload), score=point.score) diff --git a/rag/retrieval_settings.py b/rag/retrieval_settings.py new file mode 100644 index 0000000..80f2c48 --- /dev/null +++ b/rag/retrieval_settings.py @@ -0,0 +1,47 @@ +"""Retrieval configuration read from environment variables.""" + +from collections.abc import Mapping +from dataclasses import dataclass +from enum import StrEnum + +RETRIEVAL_MODE_VARIABLE = "RETRIEVAL_MODE" +QDRANT_URL_VARIABLE = "QDRANT_URL" + + +class RetrievalMode(StrEnum): + TFIDF = "tfidf" + SEMANTIC = "semantic" + HYBRID = "hybrid" + + +DEFAULT_RETRIEVAL_MODE = RetrievalMode.HYBRID + + +class InvalidRetrievalModeError(ValueError): + """Raised when RETRIEVAL_MODE names a mode that does not exist.""" + + +@dataclass(frozen=True) +class RetrievalSettings: + mode: RetrievalMode = DEFAULT_RETRIEVAL_MODE + qdrant_url: str | None = None + + +def load_retrieval_settings(environ: Mapping[str, str]) -> RetrievalSettings: + return RetrievalSettings( + mode=_parse_mode(environ.get(RETRIEVAL_MODE_VARIABLE, "")), + qdrant_url=environ.get(QDRANT_URL_VARIABLE, "").strip() or None, + ) + + +def _parse_mode(raw_mode: str) -> RetrievalMode: + normalized_mode = raw_mode.strip().lower() + if not normalized_mode: + return DEFAULT_RETRIEVAL_MODE + try: + return RetrievalMode(normalized_mode) + except ValueError: + valid_modes = ", ".join(mode.value for mode in RetrievalMode) + raise InvalidRetrievalModeError( + f"{RETRIEVAL_MODE_VARIABLE} must be one of: {valid_modes} (got {raw_mode!r})" + ) from None diff --git a/rag/store_factory.py b/rag/store_factory.py new file mode 100644 index 0000000..15b5d38 --- /dev/null +++ b/rag/store_factory.py @@ -0,0 +1,39 @@ +"""Assembles the vector store that matches the configured retrieval mode.""" + +from collections.abc import Callable + +from qdrant_client import QdrantClient + +from rag.embeddings import Embedder, FastEmbedEmbedder +from rag.hybrid_store import HybridVectorStore +from rag.qdrant_store import QdrantVectorStore +from rag.retrieval_settings import RetrievalMode, RetrievalSettings +from rag.vector_store import TfidfVectorStore, VectorStore + +IN_MEMORY_QDRANT_LOCATION = ":memory:" + + +def build_vector_store( + settings: RetrievalSettings, + embedder_factory: Callable[[], Embedder] = FastEmbedEmbedder, +) -> VectorStore: + """Build the store for `settings.mode`. + + `embedder_factory` is only called for the semantic and hybrid modes, so the + `tfidf` mode never loads, or even downloads, an embedding model. + """ + if settings.mode is RetrievalMode.TFIDF: + return TfidfVectorStore() + + semantic_store = QdrantVectorStore( + embedder=embedder_factory(), client=_connect_to_qdrant(settings) + ) + if settings.mode is RetrievalMode.SEMANTIC: + return semantic_store + return HybridVectorStore(lexical_store=TfidfVectorStore(), semantic_store=semantic_store) + + +def _connect_to_qdrant(settings: RetrievalSettings) -> QdrantClient: + if settings.qdrant_url: + return QdrantClient(url=settings.qdrant_url) + return QdrantClient(location=IN_MEMORY_QDRANT_LOCATION) diff --git a/rag/vector_store.py b/rag/vector_store.py index 1ada130..d8d9c41 100644 --- a/rag/vector_store.py +++ b/rag/vector_store.py @@ -1,19 +1,22 @@ -"""A small TF-IDF backed vector store for similarity search over chunks. +"""The retrieval contract shared by every store, plus the TF-IDF store. -TF-IDF is the deliberate baseline here: it needs no model download and no API -key, so the whole pipeline runs and is fully tested offline. The public -interface below (`add_documents`, `search`) does not expose anything -TF-IDF-specific, so swapping in a dense embedding model later only touches -this file. +TF-IDF is the lexical baseline and a first-class retrieval mode: it needs no +model download and no API key, so it runs fully offline. The semantic store +(`rag.qdrant_store`) and the hybrid combiner (`rag.hybrid_store`) implement the +same `VectorStore` contract, which is why the pipeline never depends on one +concrete store. """ from dataclasses import dataclass +from typing import Protocol from sklearn.feature_extraction.text import TfidfVectorizer from sklearn.metrics.pairwise import cosine_similarity from rag.chunking import Chunk +DEFAULT_TOP_K = 3 + @dataclass(frozen=True) class SearchResult: @@ -21,6 +24,18 @@ class SearchResult: score: float +class VectorStore(Protocol): + """Indexes chunks and ranks them against a query, best match first. + + Every store scores on its own scale, so scores are only comparable within + the results of a single `search` call. + """ + + def add_documents(self, chunks: list[Chunk]) -> None: ... + + def search(self, query: str, top_k: int = DEFAULT_TOP_K) -> list[SearchResult]: ... + + class TfidfVectorStore: def __init__(self): self._vectorizer = TfidfVectorizer() @@ -34,7 +49,7 @@ def add_documents(self, chunks: list[Chunk]) -> None: self._chunks.extend(chunks) self._matrix = self._vectorizer.fit_transform(chunk.text for chunk in self._chunks) - def search(self, query: str, top_k: int = 3) -> list[SearchResult]: + def search(self, query: str, top_k: int = DEFAULT_TOP_K) -> list[SearchResult]: if not self._chunks: return [] query_vector = self._vectorizer.transform([query]) diff --git a/requirements.txt b/requirements.txt index d165da9..d2e8506 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,3 +3,5 @@ uvicorn[standard]==0.54.0 scikit-learn==1.9.1 httpx2==2.13.1 python-dotenv==1.2.3 +qdrant-client==1.19.1 +fastembed==0.9.0 diff --git a/tests/fakes.py b/tests/fakes.py new file mode 100644 index 0000000..643c28e --- /dev/null +++ b/tests/fakes.py @@ -0,0 +1,58 @@ +"""Deterministic test doubles shared by the retrieval tests.""" + +import hashlib +import re +from collections.abc import Sequence + +from rag.chunking import Chunk +from rag.vector_store import DEFAULT_TOP_K, SearchResult + +HASHING_EMBEDDER_DIMENSION = 256 +WORD_HASH_BYTES = 4 + + +class HashingEmbedder: + """Bag-of-words embedder: every word hashes into one bucket of the vector. + + Texts that share words end up close together, which is enough to exercise the + stores without a model, and the vectors are identical on every run. + """ + + model_name = "hashing-bag-of-words" + dimension = HASHING_EMBEDDER_DIMENSION + + def embed_passages(self, texts: Sequence[str]) -> list[list[float]]: + return [self._embed(text) for text in texts] + + def embed_query(self, text: str) -> list[float]: + return self._embed(text) + + def _embed(self, text: str) -> list[float]: + vector = [0.0] * self.dimension + for word in re.findall(r"\w+", text.lower()): + digest = hashlib.blake2b(word.encode(), digest_size=WORD_HASH_BYTES).digest() + vector[int.from_bytes(digest, "big") % self.dimension] += 1.0 + return vector + + +class PresetRankingStore: + """A store that ignores the query and answers with a fixed ranking. + + It also remembers what it was given and how many results it was asked for. + """ + + def __init__(self, ranking: list[SearchResult] | None = None): + self._ranking = ranking or [] + self.added_chunks: list[Chunk] = [] + self.requested_top_k: list[int] = [] + + def add_documents(self, chunks: list[Chunk]) -> None: + self.added_chunks.extend(chunks) + + def search(self, query: str, top_k: int = DEFAULT_TOP_K) -> list[SearchResult]: + self.requested_top_k.append(top_k) + return self._ranking[:top_k] + + +def scored_chunk(document_id: str, index: int = 0, score: float = 1.0) -> SearchResult: + return SearchResult(chunk=Chunk(document_id=document_id, index=index, text="..."), score=score) diff --git a/tests/test_api.py b/tests/test_api.py index f726866..6fbc83b 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -7,6 +7,7 @@ from app.main import create_app from rag.generation import ClaudeAnswerGenerator from rag.pipeline import RagPipeline +from rag.vector_store import TfidfVectorStore @pytest.fixture @@ -22,7 +23,7 @@ def handler(request: httpx2.Request) -> httpx2.Response: model="test-model", client=mock_http_client, ) - pipeline = RagPipeline(generator=generator) + pipeline = RagPipeline(generator=generator, store=TfidfVectorStore()) app = create_app(pipeline=pipeline) return TestClient(app) diff --git a/tests/test_api_retrieval_modes.py b/tests/test_api_retrieval_modes.py new file mode 100644 index 0000000..6309c1f --- /dev/null +++ b/tests/test_api_retrieval_modes.py @@ -0,0 +1,98 @@ +"""The HTTP contract must not change with the retrieval mode: same routes, same response shapes.""" + +import httpx2 +import pytest +from fastapi.testclient import TestClient + +from app.main import create_app +from rag.generation import ClaudeAnswerGenerator +from rag.pipeline import RagPipeline +from rag.retrieval_settings import RetrievalMode, RetrievalSettings +from rag.store_factory import build_vector_store +from tests.fakes import HashingEmbedder + +SOURCE_FIELDS = {"document_id", "chunk_index", "text", "score"} + + +def build_client(mode: RetrievalMode) -> TestClient: + def answer_with_a_fixed_reply(request: httpx2.Request) -> httpx2.Response: + return httpx2.Response( + 200, json={"content": [{"type": "text", "text": "Cats are mammals [source 1]."}]} + ) + + http_client = httpx2.AsyncClient(transport=httpx2.MockTransport(answer_with_a_fixed_reply)) + generator = ClaudeAnswerGenerator(api_key="test-key", model="test-model", client=http_client) + store = build_vector_store(RetrievalSettings(mode=mode), HashingEmbedder) + return TestClient(create_app(pipeline=RagPipeline(generator=generator, store=store))) + + +@pytest.fixture(params=list(RetrievalMode), ids=lambda mode: mode.value) +def client(request) -> TestClient: + return build_client(request.param) + + +def test_ingesting_reports_the_document_and_chunk_count(client): + response = client.post( + "/documents", json={"document_id": "animals", "text": "Cats are small mammals."} + ) + + assert response.status_code == 200 + assert response.json() == {"document_id": "animals", "chunks_created": 1} + + +def test_querying_an_empty_index_returns_the_same_no_match_answer(client): + response = client.post("/query", json={"question": "anything"}) + + assert response.status_code == 200 + assert response.json() == { + "answer": "No indexed documents match this question.", + "sources": [], + } + + +def test_a_query_returns_the_answer_with_sources_in_the_original_shape(client): + client.post( + "/documents", + json={"document_id": "animals", "text": "Cats are small domesticated mammals."}, + ) + + payload = client.post("/query", json={"question": "Are cats mammals?"}).json() + + assert set(payload) == {"answer", "sources"} + assert payload["answer"] == "Cats are mammals [source 1]." + assert payload["sources"][0]["document_id"] == "animals" + assert set(payload["sources"][0]) == SOURCE_FIELDS + assert isinstance(payload["sources"][0]["score"], float) + + +def test_top_k_still_limits_the_number_of_sources(client): + for position in range(4): + client.post( + "/documents", + json={"document_id": f"note-{position}", "text": "Cats are small mammals."}, + ) + + payload = client.post("/query", json={"question": "cats mammals", "top_k": 2}).json() + + assert len(payload["sources"]) == 2 + + +def test_a_top_k_of_zero_still_answers_instead_of_failing(client): + client.post("/documents", json={"document_id": "animals", "text": "Cats are mammals."}) + + response = client.post("/query", json={"question": "cats", "top_k": 0}) + + assert response.status_code == 200 + assert response.json()["sources"] == [] + + +def test_the_published_routes_and_request_fields_are_unchanged(): + schema = build_client(RetrievalMode.TFIDF).get("/openapi.json").json() + + assert set(schema["paths"]) == {"/documents", "/query"} + schemas = schema["components"]["schemas"] + assert set(schemas["IngestRequest"]["properties"]) == {"document_id", "text"} + assert set(schemas["QueryRequest"]["properties"]) == {"question", "top_k"} + assert schemas["QueryRequest"]["properties"]["top_k"]["default"] == 3 + assert set(schemas["QueryResponse"]["properties"]) == {"answer", "sources"} + assert set(schemas["SourceOut"]["properties"]) == SOURCE_FIELDS diff --git a/tests/test_app_startup.py b/tests/test_app_startup.py new file mode 100644 index 0000000..75a8ebe --- /dev/null +++ b/tests/test_app_startup.py @@ -0,0 +1,69 @@ +"""Tests that the default pipeline is built at server startup, never at import time.""" + +import httpx2 +from fastapi.testclient import TestClient + +from app import main +from app.main import create_app +from rag.generation import ClaudeAnswerGenerator +from rag.pipeline import RagPipeline +from rag.vector_store import TfidfVectorStore + + +class PipelineBuilderSpy: + def __init__(self, pipeline: RagPipeline): + self._pipeline = pipeline + self.calls = 0 + + def __call__(self) -> RagPipeline: + self.calls += 1 + return self._pipeline + + +def build_offline_pipeline() -> RagPipeline: + client = httpx2.AsyncClient( + transport=httpx2.MockTransport(lambda request: httpx2.Response(200)) + ) + generator = ClaudeAnswerGenerator(api_key="test-key", model="test-model", client=client) + return RagPipeline(generator=generator, store=TfidfVectorStore()) + + +def test_creating_the_app_builds_no_pipeline(monkeypatch): + spy = PipelineBuilderSpy(build_offline_pipeline()) + monkeypatch.setattr(main, "build_default_pipeline", spy) + + create_app() + + assert spy.calls == 0 + + +def test_starting_the_server_builds_the_pipeline_once(monkeypatch): + spy = PipelineBuilderSpy(build_offline_pipeline()) + monkeypatch.setattr(main, "build_default_pipeline", spy) + + with TestClient(create_app()) as client: + client.post("/documents", json={"document_id": "notes", "text": "Cats purr."}) + client.post("/documents", json={"document_id": "more", "text": "Dogs bark."}) + + assert spy.calls == 1 + + +def test_the_pipeline_built_at_startup_serves_requests(monkeypatch): + monkeypatch.setattr(main, "build_default_pipeline", build_offline_pipeline) + + with TestClient(create_app()) as client: + response = client.post("/documents", json={"document_id": "notes", "text": "Cats purr."}) + + assert response.json() == {"document_id": "notes", "chunks_created": 1} + + +def test_an_injected_pipeline_is_used_and_nothing_is_built_at_startup(monkeypatch): + spy = PipelineBuilderSpy(build_offline_pipeline()) + monkeypatch.setattr(main, "build_default_pipeline", spy) + injected_pipeline = build_offline_pipeline() + + with TestClient(create_app(pipeline=injected_pipeline)) as client: + client.post("/documents", json={"document_id": "notes", "text": "Cats purr."}) + + assert spy.calls == 0 + assert injected_pipeline._store.search("cats")[0].chunk.document_id == "notes" diff --git a/tests/test_embeddings.py b/tests/test_embeddings.py new file mode 100644 index 0000000..c66dca6 --- /dev/null +++ b/tests/test_embeddings.py @@ -0,0 +1,78 @@ +"""Tests for the fastembed adapter, using a stand-in for the model so nothing is downloaded.""" + +import numpy as np +import pytest + +from rag import embeddings +from rag.embeddings import DEFAULT_EMBEDDING_MODEL, FastEmbedEmbedder + +STUB_DIMENSION = 3 + + +class StubTextEmbedding: + """Mimics the parts of fastembed's TextEmbedding that the adapter uses.""" + + def __init__(self, model_name: str, cache_dir: str | None = None): + self.model_name = model_name + self.cache_dir = cache_dir + self.passage_calls: list[list[str]] = [] + self.query_calls: list[str] = [] + + @classmethod + def get_embedding_size(cls, model_name: str) -> int: + return STUB_DIMENSION + + def passage_embed(self, texts): + self.passage_calls.append(list(texts)) + return ( + np.full(STUB_DIMENSION, float(index), dtype=np.float32) for index, _ in enumerate(texts) + ) + + def query_embed(self, query): + self.query_calls.append(query) + return iter([np.array([0.5, 0.25, 0.125], dtype=np.float32)]) + + +@pytest.fixture +def embedder(monkeypatch) -> FastEmbedEmbedder: + monkeypatch.setattr(embeddings, "TextEmbedding", StubTextEmbedding) + return FastEmbedEmbedder() + + +def test_the_default_model_is_a_multilingual_one(embedder): + assert embedder.model_name == DEFAULT_EMBEDDING_MODEL + assert "multilingual" in embedder.model_name + + +def test_the_dimension_comes_from_the_model_registry(embedder): + assert embedder.dimension == STUB_DIMENSION + + +def test_passages_become_plain_python_float_lists(embedder): + vectors = embedder.embed_passages(["first", "second"]) + + assert vectors == [[0.0, 0.0, 0.0], [1.0, 1.0, 1.0]] + assert all(isinstance(number, float) for vector in vectors for number in vector) + + +def test_passages_use_the_passage_side_of_the_model(embedder): + embedder.embed_passages(["first", "second"]) + + assert embedder._model.passage_calls == [["first", "second"]] + assert embedder._model.query_calls == [] + + +def test_a_query_becomes_a_single_float_list_through_the_query_side(embedder): + vector = embedder.embed_query("what is a cache?") + + assert vector == [0.5, 0.25, 0.125] + assert embedder._model.query_calls == ["what is a cache?"] + + +def test_the_model_name_and_cache_directory_reach_the_model(monkeypatch): + monkeypatch.setattr(embeddings, "TextEmbedding", StubTextEmbedding) + + embedder = FastEmbedEmbedder(model_name="some/model", cache_dir="/tmp/models") + + assert embedder.model_name == "some/model" + assert embedder._model.cache_dir == "/tmp/models" diff --git a/tests/test_evaluation_report.py b/tests/test_evaluation_report.py new file mode 100644 index 0000000..9a75130 --- /dev/null +++ b/tests/test_evaluation_report.py @@ -0,0 +1,72 @@ +"""Tests for the Markdown rendering of an evaluation report.""" + +from evaluation.golden_set import QuestionKind +from evaluation.metrics import RetrievalScores +from evaluation.report import EvaluationReport, ModeReport, render_markdown +from rag.retrieval_settings import RetrievalMode + + +def build_report() -> EvaluationReport: + def mode_report(mode: RetrievalMode, hits: int, mrr: float) -> ModeReport: + return ModeReport( + mode=mode, + overall=RetrievalScores(hits=hits, questions=4, mean_reciprocal_rank=mrr), + by_kind={ + QuestionKind.LEXICAL_OVERLAP: RetrievalScores( + hits=hits - 1, questions=3, mean_reciprocal_rank=mrr + ), + QuestionKind.EXACT_TERM: RetrievalScores( + hits=1, questions=1, mean_reciprocal_rank=1.0 + ), + }, + ) + + return EvaluationReport( + golden_set_version=2, + document_count=18, + chunk_count=36, + embedding_model="some/model", + modes=[ + mode_report(RetrievalMode.TFIDF, 3, 0.75), + mode_report(RetrievalMode.SEMANTIC, 2, 0.5), + mode_report(RetrievalMode.HYBRID, 4, 0.875), + ], + ) + + +def test_the_report_opens_with_the_corpus_and_model_it_describes(): + lines = render_markdown(build_report()).splitlines() + + assert lines[0] == ( + "Golden set v2: 4 questions over 18 notes (36 chunks). Embedding model: some/model." + ) + + +def test_the_overall_table_shows_hit_rate_with_counts_and_mrr_per_mode(): + rendered = render_markdown(build_report()) + + assert "| Mode | Hit rate@3 | MRR |" in rendered + assert "| tfidf | 0.750 (3/4) | 0.750 |" in rendered + assert "| semantic | 0.500 (2/4) | 0.500 |" in rendered + assert "| hybrid | 1.000 (4/4) | 0.875 |" in rendered + + +def test_the_per_type_tables_list_only_the_types_that_were_asked(): + rendered = render_markdown(build_report()) + + assert "| Mode | Lexical overlap (3) | Exact term (1) |" in rendered + assert "Paraphrase" not in rendered + + +def test_the_per_type_tables_hold_hit_rate_and_mrr_for_each_mode(): + rendered = render_markdown(build_report()) + + assert "| tfidf | 0.667 (2/3) | 1.000 (1/1) |" in rendered + assert "| hybrid | 1.000 (3/3) | 1.000 (1/1) |" in rendered + assert "| tfidf | 0.750 | 1.000 |" in rendered + + +def test_every_table_has_a_header_separator(): + rendered = render_markdown(build_report()) + + assert rendered.count("| --- | --- | --- |") == 3 diff --git a/tests/test_evaluation_run.py b/tests/test_evaluation_run.py new file mode 100644 index 0000000..c5cd868 --- /dev/null +++ b/tests/test_evaluation_run.py @@ -0,0 +1,170 @@ +"""Tests for running every retrieval mode against a golden set and gating on the result.""" + +import logging + +import pytest + +from evaluation import run +from evaluation.golden_set import GoldenQuestion, GoldenSet, GoldenSetError, QuestionKind +from evaluation.metrics import RetrievalScores +from evaluation.report import EvaluationReport, ModeReport +from evaluation.run import ( + EXIT_OK, + EXIT_QUALITY_GATE_FAILED, + hybrid_meets_tfidf_baseline, + main, + run_evaluation, +) +from rag.retrieval_settings import RetrievalMode +from tests.fakes import HashingEmbedder + +DOCUMENTS = { + "cats": "Cats are small domesticated feline animals that purr.", + "rockets": "Rockets use combustion to reach orbit around the planet.", + "bread": "Sourdough bread rises slowly during a long cold fermentation.", + "bridges": "Suspension bridges hang their deck from cables anchored at both ends.", + "glaciers": "Glaciers are slow rivers of ice that carve valleys over centuries.", + "tea": "Green tea is steeped briefly in water that is hot but not boiling.", +} + +GOLDEN_SET = GoldenSet( + version=7, + questions=[ + GoldenQuestion( + "q1", QuestionKind.LEXICAL_OVERLAP, "Which animals purr?", frozenset({"cats"}) + ), + GoldenQuestion("q2", QuestionKind.EXACT_TERM, "combustion", frozenset({"rockets"})), + GoldenQuestion("q3", QuestionKind.PARAPHRASE, "zzz qqq", frozenset({"bread"})), + ], +) + + +class CountingEmbedderFactory: + def __init__(self): + self.calls = 0 + + def __call__(self) -> HashingEmbedder: + self.calls += 1 + return HashingEmbedder() + + +def scores(hits: int, questions: int = 10) -> RetrievalScores: + return RetrievalScores(hits=hits, questions=questions, mean_reciprocal_rank=0.5) + + +def report_with_hits(tfidf_hits: int, hybrid_hits: int) -> EvaluationReport: + def mode_report(mode: RetrievalMode, hits: int) -> ModeReport: + return ModeReport( + mode=mode, overall=scores(hits), by_kind={QuestionKind.PARAPHRASE: scores(hits)} + ) + + return EvaluationReport( + golden_set_version=1, + document_count=1, + chunk_count=1, + embedding_model="stub", + modes=[ + mode_report(RetrievalMode.TFIDF, tfidf_hits), + mode_report(RetrievalMode.SEMANTIC, 0), + mode_report(RetrievalMode.HYBRID, hybrid_hits), + ], + ) + + +def test_every_retrieval_mode_is_evaluated_in_a_fixed_order(): + report = run_evaluation(GOLDEN_SET, DOCUMENTS, HashingEmbedder) + + assert [mode_report.mode for mode_report in report.modes] == list(RetrievalMode) + + +def test_the_report_describes_the_corpus_and_the_model_it_ran_on(): + report = run_evaluation(GOLDEN_SET, DOCUMENTS, HashingEmbedder) + + assert report.golden_set_version == 7 + assert report.document_count == len(DOCUMENTS) + assert report.chunk_count == len(DOCUMENTS) + assert report.embedding_model == "hashing-bag-of-words" + + +def test_tfidf_is_scored_overall_and_per_question_type(): + report = run_evaluation(GOLDEN_SET, DOCUMENTS, HashingEmbedder) + tfidf = report.modes[0] + + assert (tfidf.overall.hits, tfidf.overall.questions) == (2, 3) + assert tfidf.by_kind[QuestionKind.LEXICAL_OVERLAP].hits == 1 + assert tfidf.by_kind[QuestionKind.EXACT_TERM].hits == 1 + assert tfidf.by_kind[QuestionKind.PARAPHRASE].hits == 0 + assert tfidf.overall.mean_reciprocal_rank == pytest.approx(2 / 3) + + +def test_the_embedding_model_is_loaded_once_for_all_modes(): + embedder_factory = CountingEmbedderFactory() + + run_evaluation(GOLDEN_SET, DOCUMENTS, embedder_factory) + + assert embedder_factory.calls == 1 + + +def test_a_label_for_a_missing_note_stops_the_run_before_any_model_loads(): + mislabeled = GoldenSet( + version=1, + questions=[GoldenQuestion("q1", QuestionKind.PARAPHRASE, "?", frozenset({"ghost"}))], + ) + embedder_factory = CountingEmbedderFactory() + + with pytest.raises(GoldenSetError, match="ghost"): + run_evaluation(mislabeled, DOCUMENTS, embedder_factory) + + assert embedder_factory.calls == 0 + + +def test_missed_questions_are_logged_with_the_documents_that_were_returned(caplog): + with caplog.at_level(logging.INFO, logger="evaluation.run"): + run_evaluation(GOLDEN_SET, DOCUMENTS, HashingEmbedder) + + tfidf_misses = [ + record.getMessage() for record in caplog.records if "[tfidf]" in record.getMessage() + ] + assert len(tfidf_misses) == 1 + assert "q3" in tfidf_misses[0] + + +@pytest.mark.parametrize( + ("tfidf_hits", "hybrid_hits", "passes"), + [(8, 9, True), (8, 8, True), (8, 7, False)], +) +def test_the_gate_fails_only_when_hybrid_hits_fewer_questions_than_tfidf( + tfidf_hits, hybrid_hits, passes +): + assert hybrid_meets_tfidf_baseline(report_with_hits(tfidf_hits, hybrid_hits)) is passes + + +def test_the_command_prints_the_results_table_for_every_mode(capsys): + exit_status = main([], embedder_factory=HashingEmbedder) + + output = capsys.readouterr().out + assert exit_status == EXIT_OK + for mode in RetrievalMode: + assert f"| {mode.value} |" in output + + +def test_the_gate_flag_returns_a_failure_status_when_hybrid_regresses(monkeypatch, caplog): + monkeypatch.setattr(run, "run_evaluation", lambda *arguments: report_with_hits(8, 5)) + + with caplog.at_level(logging.ERROR, logger="evaluation.run"): + exit_status = main(["--enforce-gate"], embedder_factory=HashingEmbedder) + + assert exit_status == EXIT_QUALITY_GATE_FAILED + assert "Quality gate failed" in caplog.text + + +def test_without_the_gate_flag_a_regression_does_not_fail_the_command(monkeypatch): + monkeypatch.setattr(run, "run_evaluation", lambda *arguments: report_with_hits(8, 5)) + + assert main([], embedder_factory=HashingEmbedder) == EXIT_OK + + +def test_the_gate_flag_returns_success_when_hybrid_holds_the_baseline(monkeypatch): + monkeypatch.setattr(run, "run_evaluation", lambda *arguments: report_with_hits(8, 8)) + + assert main(["--enforce-gate"], embedder_factory=HashingEmbedder) == EXIT_OK diff --git a/tests/test_fakes.py b/tests/test_fakes.py new file mode 100644 index 0000000..7ea6372 --- /dev/null +++ b/tests/test_fakes.py @@ -0,0 +1,38 @@ +"""The retrieval tests lean on the fake embedder, so its own guarantees are checked here.""" + +from tests.fakes import HashingEmbedder + + +def cosine_similarity(left: list[float], right: list[float]) -> float: + dot_product = sum(a * b for a, b in zip(left, right, strict=True)) + left_norm = sum(a * a for a in left) ** 0.5 + right_norm = sum(b * b for b in right) ** 0.5 + return dot_product / (left_norm * right_norm) + + +def test_the_same_text_always_produces_the_same_vector(): + embedder = HashingEmbedder() + + assert embedder.embed_query("cats purr") == embedder.embed_query("cats purr") + + +def test_vectors_have_the_declared_dimension(): + embedder = HashingEmbedder() + + assert len(embedder.embed_query("cats purr")) == embedder.dimension + + +def test_passages_and_queries_share_one_vector_space(): + embedder = HashingEmbedder() + + assert embedder.embed_passages(["cats purr"])[0] == embedder.embed_query("cats purr") + + +def test_texts_sharing_words_are_closer_than_unrelated_texts(): + embedder = HashingEmbedder() + query = embedder.embed_query("cats purr loudly") + + related = cosine_similarity(query, embedder.embed_query("cats purr when happy")) + unrelated = cosine_similarity(query, embedder.embed_query("rockets reach orbit")) + + assert related > unrelated diff --git a/tests/test_fusion.py b/tests/test_fusion.py new file mode 100644 index 0000000..9977dbf --- /dev/null +++ b/tests/test_fusion.py @@ -0,0 +1,93 @@ +"""Tests for Reciprocal Rank Fusion, including the arithmetic of the fused scores.""" + +import pytest + +from rag.fusion import RRF_RANK_CONSTANT, reciprocal_rank_fusion +from tests.fakes import scored_chunk + + +def document_order(results) -> list[str]: + return [result.chunk.document_id for result in results] + + +def test_a_chunk_ranked_first_in_two_rankings_scores_twice_the_single_rank_score(): + results = reciprocal_rank_fusion([[scored_chunk("a")], [scored_chunk("a")]], top_k=1) + + assert results[0].score == pytest.approx(2 / (RRF_RANK_CONSTANT + 1)) + + +def test_each_rank_contributes_one_over_the_rank_constant_plus_the_rank(): + ranking = [scored_chunk("a"), scored_chunk("b"), scored_chunk("c")] + + results = reciprocal_rank_fusion([ranking], top_k=3) + + assert [result.score for result in results] == pytest.approx( + [1 / (RRF_RANK_CONSTANT + rank) for rank in (1, 2, 3)] + ) + + +def test_a_chunk_found_by_both_rankings_beats_chunks_found_by_only_one(): + lexical = [scored_chunk("only-lexical"), scored_chunk("shared")] + semantic = [scored_chunk("only-semantic"), scored_chunk("shared")] + + results = reciprocal_rank_fusion([lexical, semantic], top_k=3) + + assert document_order(results)[0] == "shared" + + +def test_the_raw_scores_of_the_rankings_are_ignored(): + small_scale = [scored_chunk("a", score=0.9), scored_chunk("b", score=0.1)] + huge_scale = [scored_chunk("b", score=9000.0), scored_chunk("a", score=1.0)] + + results = reciprocal_rank_fusion([small_scale, huge_scale], top_k=2) + + assert [result.score for result in results] == pytest.approx( + [1 / (RRF_RANK_CONSTANT + 1) + 1 / (RRF_RANK_CONSTANT + 2)] * 2 + ) + + +def test_top_k_limits_the_number_of_fused_results(): + ranking = [scored_chunk("a"), scored_chunk("b"), scored_chunk("c")] + + assert document_order(reciprocal_rank_fusion([ranking], top_k=2)) == ["a", "b"] + + +def test_equal_scores_keep_the_order_in_which_chunks_were_first_seen(): + results = reciprocal_rank_fusion([[scored_chunk("first")], [scored_chunk("second")]], top_k=2) + + assert document_order(results) == ["first", "second"] + + +def test_chunks_of_the_same_document_are_told_apart_by_their_index(): + ranking = [scored_chunk("doc", index=0), scored_chunk("doc", index=1)] + + results = reciprocal_rank_fusion([ranking], top_k=2) + + assert [result.chunk.index for result in results] == [0, 1] + + +def test_a_chunk_repeated_within_one_ranking_counts_once_at_its_best_position(): + duplicated = [scored_chunk("a"), scored_chunk("a"), scored_chunk("b")] + + results = reciprocal_rank_fusion([duplicated], top_k=2) + + assert document_order(results) == ["a", "b"] + assert results[0].score == pytest.approx(1 / (RRF_RANK_CONSTANT + 1)) + assert results[1].score == pytest.approx(1 / (RRF_RANK_CONSTANT + 2)) + + +def test_fusing_a_single_ranking_keeps_its_order(): + ranking = [scored_chunk("a"), scored_chunk("b"), scored_chunk("c")] + + assert document_order(reciprocal_rank_fusion([ranking], top_k=3)) == ["a", "b", "c"] + + +def test_empty_rankings_fuse_to_nothing(): + assert reciprocal_rank_fusion([[], []], top_k=3) == [] + assert reciprocal_rank_fusion([], top_k=3) == [] + + +def test_an_empty_ranking_does_not_disturb_the_other_one(): + ranking = [scored_chunk("a"), scored_chunk("b")] + + assert document_order(reciprocal_rank_fusion([[], ranking], top_k=2)) == ["a", "b"] diff --git a/tests/test_golden_set.py b/tests/test_golden_set.py new file mode 100644 index 0000000..f0870c9 --- /dev/null +++ b/tests/test_golden_set.py @@ -0,0 +1,188 @@ +"""Tests that the golden set file is well formed and that its question types mean what they say.""" + +import json +import re + +import pytest +from sklearn.feature_extraction.text import ENGLISH_STOP_WORDS + +from evaluation.golden_set import ( + GoldenQuestion, + GoldenSetError, + QuestionKind, + ensure_labels_exist, + load_golden_set, +) +from rag.sample_corpus import load_sample_documents + +MIN_QUESTIONS = 40 +MIN_QUESTIONS_PER_KIND = 10 +MIN_LEXICAL_OVERLAP = 0.5 +MAX_PARAPHRASE_OVERLAP = 0.45 +COPIED_PHRASE_WORDS = 5 + + +@pytest.fixture(scope="module") +def golden_set(): + return load_golden_set() + + +@pytest.fixture(scope="module") +def documents() -> dict[str, str]: + return load_sample_documents() + + +def words_of(text: str) -> list[str]: + return re.findall(r"[a-z0-9_]+", text.lower()) + + +def content_words(text: str) -> set[str]: + return {word for word in words_of(text) if word not in ENGLISH_STOP_WORDS and len(word) > 1} + + +def overlap_with_notes(question: GoldenQuestion, documents: dict[str, str]) -> float: + question_words = content_words(question.text) + note_words = set().union( + *(content_words(documents[document_id]) for document_id in question.relevant_documents) + ) + return len(question_words & note_words) / len(question_words) + + +def phrases_of(text: str, length: int) -> set[tuple[str, ...]]: + words = words_of(text) + return {tuple(words[start : start + length]) for start in range(len(words) - length + 1)} + + +def write_golden_set(path, questions: list[dict], version: int = 1): + path.write_text(json.dumps({"version": version, "questions": questions})) + return path + + +def valid_entry(**overrides) -> dict: + entry = { + "id": "q-1", + "kind": "paraphrase", + "question": "Why?", + "relevant_documents": ["notes"], + } + return entry | overrides + + +def test_the_committed_golden_set_is_versioned_and_large_enough(golden_set): + assert golden_set.version >= 1 + assert len(golden_set.questions) >= MIN_QUESTIONS + + +def test_every_question_type_has_enough_questions(golden_set): + for kind in QuestionKind: + count = sum(question.kind is kind for question in golden_set.questions) + assert count >= MIN_QUESTIONS_PER_KIND, kind + + +def test_every_label_points_at_a_note_that_exists(golden_set, documents): + ensure_labels_exist(golden_set, documents.keys()) + + +def test_every_note_is_the_answer_to_at_least_one_question(golden_set, documents): + labeled_documents = { + document_id + for question in golden_set.questions + for document_id in question.relevant_documents + } + + assert labeled_documents == set(documents) + + +def test_no_question_repeats_a_phrase_of_the_notes_verbatim(golden_set, documents): + note_phrases = set().union( + *(phrases_of(text, COPIED_PHRASE_WORDS) for text in documents.values()) + ) + + for question in golden_set.questions: + copied = phrases_of(question.text, COPIED_PHRASE_WORDS) & note_phrases + assert not copied, (question.question_id, copied) + + +def test_lexical_overlap_questions_reuse_the_words_of_their_notes(golden_set, documents): + for question in golden_set.questions: + if question.kind is QuestionKind.LEXICAL_OVERLAP: + overlap = overlap_with_notes(question, documents) + assert overlap >= MIN_LEXICAL_OVERLAP, (question.question_id, overlap) + + +def test_paraphrase_questions_share_few_words_with_their_notes(golden_set, documents): + for question in golden_set.questions: + if question.kind is QuestionKind.PARAPHRASE: + overlap = overlap_with_notes(question, documents) + assert overlap <= MAX_PARAPHRASE_OVERLAP, (question.question_id, overlap) + + +def test_exact_term_questions_are_short_terms_found_in_their_notes(golden_set, documents): + for question in golden_set.questions: + if question.kind is QuestionKind.EXACT_TERM: + assert len(question.text.split()) <= 4, question.question_id + assert overlap_with_notes(question, documents) == 1.0, question.question_id + + +def test_a_valid_file_loads_with_its_labels(tmp_path): + path = write_golden_set(tmp_path / "golden.json", [valid_entry()], version=3) + + loaded = load_golden_set(path) + + assert loaded.version == 3 + assert loaded.questions == [ + GoldenQuestion( + question_id="q-1", + kind=QuestionKind.PARAPHRASE, + text="Why?", + relevant_documents=frozenset({"notes"}), + ) + ] + + +@pytest.mark.parametrize("missing_field", ["id", "kind", "question", "relevant_documents"]) +def test_an_entry_missing_a_field_is_rejected(tmp_path, missing_field): + entry = valid_entry() + del entry[missing_field] + path = write_golden_set(tmp_path / "golden.json", [entry]) + + with pytest.raises(GoldenSetError, match=missing_field): + load_golden_set(path) + + +def test_an_unknown_question_kind_is_rejected(tmp_path): + path = write_golden_set(tmp_path / "golden.json", [valid_entry(kind="trick_question")]) + + with pytest.raises(GoldenSetError, match="trick_question"): + load_golden_set(path) + + +def test_a_question_without_relevant_documents_is_rejected(tmp_path): + path = write_golden_set(tmp_path / "golden.json", [valid_entry(relevant_documents=[])]) + + with pytest.raises(GoldenSetError, match="no relevant document"): + load_golden_set(path) + + +def test_duplicate_question_ids_are_rejected(tmp_path): + path = write_golden_set(tmp_path / "golden.json", [valid_entry(), valid_entry()]) + + with pytest.raises(GoldenSetError, match="q-1"): + load_golden_set(path) + + +def test_a_file_without_a_version_is_rejected(tmp_path): + path = tmp_path / "golden.json" + path.write_text(json.dumps({"questions": [valid_entry()]})) + + with pytest.raises(GoldenSetError, match="version"): + load_golden_set(path) + + +def test_a_label_for_a_missing_note_is_reported_with_the_question(tmp_path): + path = write_golden_set( + tmp_path / "golden.json", [valid_entry(relevant_documents=["notes", "ghost"])] + ) + + with pytest.raises(GoldenSetError, match=r"q-1.*ghost"): + ensure_labels_exist(load_golden_set(path), {"notes"}) diff --git a/tests/test_hybrid_store.py b/tests/test_hybrid_store.py new file mode 100644 index 0000000..0a46a4f --- /dev/null +++ b/tests/test_hybrid_store.py @@ -0,0 +1,117 @@ +"""Tests for the hybrid store that fuses a lexical and a semantic ranking.""" + +import pytest +from qdrant_client import QdrantClient + +from rag.chunking import Chunk +from rag.fusion import RRF_RANK_CONSTANT +from rag.hybrid_store import CANDIDATES_PER_RANKER, HybridVectorStore +from rag.qdrant_store import QdrantVectorStore +from rag.vector_store import TfidfVectorStore +from tests.fakes import HashingEmbedder, PresetRankingStore, scored_chunk + + +def document_order(results) -> list[str]: + return [result.chunk.document_id for result in results] + + +def test_added_chunks_reach_both_stores(): + lexical_store, semantic_store = PresetRankingStore(), PresetRankingStore() + hybrid_store = HybridVectorStore(lexical_store, semantic_store) + chunks = [Chunk(document_id="doc", index=0, text="some text")] + + hybrid_store.add_documents(chunks) + + assert lexical_store.added_chunks == chunks + assert semantic_store.added_chunks == chunks + + +def test_a_chunk_both_stores_agree_on_is_ranked_first(): + lexical_store = PresetRankingStore([scored_chunk("lexical-pick"), scored_chunk("agreed")]) + semantic_store = PresetRankingStore([scored_chunk("semantic-pick"), scored_chunk("agreed")]) + + results = HybridVectorStore(lexical_store, semantic_store).search("question") + + assert document_order(results)[0] == "agreed" + + +def test_the_semantic_ranking_alone_is_used_when_the_lexical_store_finds_nothing(): + lexical_store = PresetRankingStore([]) + semantic_store = PresetRankingStore([scored_chunk("paraphrased-match")]) + + results = HybridVectorStore(lexical_store, semantic_store).search("question") + + assert document_order(results) == ["paraphrased-match"] + + +def test_the_lexical_ranking_alone_is_used_when_the_semantic_store_finds_nothing(): + lexical_store = PresetRankingStore([scored_chunk("exact-term-match")]) + semantic_store = PresetRankingStore([]) + + results = HybridVectorStore(lexical_store, semantic_store).search("question") + + assert document_order(results) == ["exact-term-match"] + + +def test_results_carry_the_fused_score_not_a_store_score(): + lexical_store = PresetRankingStore([scored_chunk("a", score=0.87)]) + semantic_store = PresetRankingStore([scored_chunk("a", score=0.41)]) + + results = HybridVectorStore(lexical_store, semantic_store).search("question") + + assert results[0].score == pytest.approx(2 / (RRF_RANK_CONSTANT + 1)) + + +def test_top_k_limits_the_fused_results(): + ranking = [scored_chunk(f"doc-{position}") for position in range(6)] + hybrid_store = HybridVectorStore(PresetRankingStore(ranking), PresetRankingStore(ranking)) + + assert len(hybrid_store.search("question", top_k=2)) == 2 + + +def test_each_store_is_asked_for_more_candidates_than_top_k(): + lexical_store, semantic_store = PresetRankingStore(), PresetRankingStore() + + HybridVectorStore(lexical_store, semantic_store).search("question", top_k=3) + + assert lexical_store.requested_top_k == [CANDIDATES_PER_RANKER] + assert semantic_store.requested_top_k == [CANDIDATES_PER_RANKER] + + +def test_a_top_k_larger_than_the_candidate_pool_is_still_honored(): + lexical_store, semantic_store = PresetRankingStore(), PresetRankingStore() + + HybridVectorStore(lexical_store, semantic_store).search( + "question", top_k=CANDIDATES_PER_RANKER + 5 + ) + + assert lexical_store.requested_top_k == [CANDIDATES_PER_RANKER + 5] + + +def test_a_chunk_ranked_deep_by_both_stores_can_beat_one_ranked_first_by_one(): + deep_ranking = [scored_chunk(f"filler-{position}") for position in range(4)] + lexical_store = PresetRankingStore([scored_chunk("lexical-only"), *deep_ranking]) + semantic_store = PresetRankingStore([scored_chunk("semantic-only"), *deep_ranking]) + + results = HybridVectorStore(lexical_store, semantic_store).search("question", top_k=3) + + assert "lexical-only" not in document_order(results)[:1] + assert document_order(results)[0] == "filler-0" + + +def test_hybrid_over_real_stores_finds_documents_the_words_point_to(): + hybrid_store = HybridVectorStore( + lexical_store=TfidfVectorStore(), + semantic_store=QdrantVectorStore( + embedder=HashingEmbedder(), client=QdrantClient(location=":memory:") + ), + ) + hybrid_store.add_documents( + [ + Chunk(document_id="cats", index=0, text="Cats are small domesticated feline animals."), + Chunk(document_id="rockets", index=0, text="Rockets use combustion to reach orbit."), + ] + ) + + assert hybrid_store.search("feline animals")[0].chunk.document_id == "cats" + assert hybrid_store.search("orbit combustion")[0].chunk.document_id == "rockets" diff --git a/tests/test_metrics.py b/tests/test_metrics.py new file mode 100644 index 0000000..b0948c8 --- /dev/null +++ b/tests/test_metrics.py @@ -0,0 +1,89 @@ +"""Tests for the retrieval metrics, checked against values worked out by hand.""" + +import pytest + +from evaluation.metrics import RetrievalOutcome, hit_at_k, reciprocal_rank, score_outcomes + + +def outcome(ranked_document_ids: list[str], relevant_documents: set[str]) -> RetrievalOutcome: + return RetrievalOutcome( + ranked_document_ids=tuple(ranked_document_ids), + relevant_documents=frozenset(relevant_documents), + ) + + +def test_a_relevant_document_at_the_cutoff_position_is_a_hit(): + assert hit_at_k(outcome(["x", "y", "a"], {"a"}), k=3) + + +def test_a_relevant_document_just_past_the_cutoff_is_a_miss(): + assert not hit_at_k(outcome(["x", "y", "z", "a"], {"a"}), k=3) + + +def test_any_of_several_relevant_documents_counts_as_a_hit(): + assert hit_at_k(outcome(["x", "b", "y"], {"a", "b"}), k=3) + + +def test_chunks_of_one_document_fill_the_cutoff_like_any_other_chunks(): + ranking = ["x", "x", "x", "a"] + + assert not hit_at_k(outcome(ranking, {"a"}), k=3) + assert hit_at_k(outcome(ranking, {"a"}), k=4) + + +def test_an_empty_ranking_is_a_miss(): + assert not hit_at_k(outcome([], {"a"}), k=3) + + +@pytest.mark.parametrize( + ("ranking", "expected"), + [ + (["a", "x", "y"], 1.0), + (["x", "a", "y"], 1 / 2), + (["x", "y", "a"], 1 / 3), + (["x", "y", "z"], 0.0), + ([], 0.0), + ], +) +def test_reciprocal_rank_is_one_over_the_position_of_the_first_relevant_chunk(ranking, expected): + assert reciprocal_rank(outcome(ranking, {"a"})) == pytest.approx(expected) + + +def test_reciprocal_rank_uses_the_earliest_of_several_relevant_documents(): + assert reciprocal_rank(outcome(["x", "b", "a"], {"a", "b"})) == pytest.approx(1 / 2) + + +def test_reciprocal_rank_counts_every_chunk_that_comes_before_the_answer(): + assert reciprocal_rank(outcome(["x", "x", "a"], {"a"})) == pytest.approx(1 / 3) + + +def test_scores_aggregate_hits_and_the_mean_reciprocal_rank(): + outcomes = [ + outcome(["a", "x"], {"a"}), + outcome(["x", "y", "a"], {"a"}), + outcome(["x", "y", "z"], {"a"}), + ] + + scores = score_outcomes(outcomes, k=3) + + assert scores.hits == 2 + assert scores.questions == 3 + assert scores.hit_rate == pytest.approx(2 / 3) + assert scores.mean_reciprocal_rank == pytest.approx((1 + 1 / 3 + 0) / 3) + + +def test_scores_honor_the_requested_cutoff(): + outcomes = [outcome(["x", "a"], {"a"})] + + assert score_outcomes(outcomes, k=1).hits == 0 + assert score_outcomes(outcomes, k=2).hits == 1 + + +def test_the_default_cutoff_is_three(): + assert score_outcomes([outcome(["x", "y", "a"], {"a"})]).hits == 1 + assert score_outcomes([outcome(["x", "y", "z", "a"], {"a"})]).hits == 0 + + +def test_scoring_zero_questions_is_an_error(): + with pytest.raises(ValueError, match="zero questions"): + score_outcomes([]) diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index b27bc92..6606382 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -4,12 +4,13 @@ from rag.generation import ClaudeAnswerGenerator from rag.pipeline import RagPipeline +from rag.vector_store import TfidfVectorStore def _pipeline_with_generator(handler) -> RagPipeline: client = httpx2.AsyncClient(transport=httpx2.MockTransport(handler)) generator = ClaudeAnswerGenerator(api_key="test-key", model="test-model", client=client) - return RagPipeline(generator=generator) + return RagPipeline(generator=generator, store=TfidfVectorStore()) async def test_asking_before_ingesting_anything_returns_no_results(): diff --git a/tests/test_qdrant_store.py b/tests/test_qdrant_store.py new file mode 100644 index 0000000..0a8bfb1 --- /dev/null +++ b/tests/test_qdrant_store.py @@ -0,0 +1,105 @@ +"""Tests for the Qdrant-backed semantic store, running Qdrant in memory with a fake embedder.""" + +import pytest +from qdrant_client import QdrantClient + +from rag.chunking import Chunk +from rag.qdrant_store import COLLECTION_NAME, UPSERT_BATCH_SIZE, QdrantVectorStore +from tests.fakes import HashingEmbedder + + +@pytest.fixture +def client() -> QdrantClient: + return QdrantClient(location=":memory:") + + +@pytest.fixture +def store(client) -> QdrantVectorStore: + return QdrantVectorStore(embedder=HashingEmbedder(), client=client) + + +def test_search_on_an_empty_store_returns_nothing(store): + assert store.search("anything") == [] + + +def test_search_ranks_the_more_relevant_chunk_first(store): + store.add_documents( + [ + Chunk(document_id="cats", index=0, text="Cats are small domesticated feline animals."), + Chunk(document_id="rockets", index=0, text="Rockets use combustion to reach orbit."), + ] + ) + + results = store.search("feline animals") + + assert results[0].chunk.document_id == "cats" + assert results[0].score > results[1].score + + +def test_search_returns_the_stored_chunk_unchanged(store): + stored_chunk = Chunk(document_id="cats", index=3, text="Cats purr when they are content.") + store.add_documents([stored_chunk]) + + results = store.search("cats purr") + + assert results[0].chunk == stored_chunk + + +def test_search_respects_top_k(store): + store.add_documents( + [ + Chunk(document_id="doc", index=position, text="python programming") + for position in range(5) + ] + ) + + assert len(store.search("python programming", top_k=2)) == 2 + + +def test_a_top_k_below_one_returns_nothing_instead_of_failing(store): + store.add_documents([Chunk(document_id="doc", index=0, text="python programming")]) + + assert store.search("python programming", top_k=0) == [] + + +def test_documents_added_later_are_still_searchable(store): + store.add_documents([Chunk(document_id="a", index=0, text="apples and oranges")]) + store.add_documents([Chunk(document_id="b", index=0, text="bicycles and cars")]) + + assert store.search("bicycles")[0].chunk.document_id == "b" + + +def test_adding_the_same_chunk_twice_stores_it_once(store, client): + chunk = Chunk(document_id="a", index=0, text="apples and oranges") + + store.add_documents([chunk]) + store.add_documents([chunk]) + + assert client.count(COLLECTION_NAME).count == 1 + + +def test_more_chunks_than_one_batch_are_all_stored(store, client): + chunk_count = UPSERT_BATCH_SIZE * 2 + 5 + store.add_documents( + [ + Chunk(document_id="doc", index=position, text=f"note {position}") + for position in range(chunk_count) + ] + ) + + assert client.count(COLLECTION_NAME).count == chunk_count + + +def test_adding_no_chunks_changes_nothing(store, client): + store.add_documents([]) + + assert client.count(COLLECTION_NAME).count == 0 + + +def test_a_new_store_starts_from_an_empty_collection(client): + first_store = QdrantVectorStore(embedder=HashingEmbedder(), client=client) + first_store.add_documents([Chunk(document_id="old", index=0, text="stale content")]) + + second_store = QdrantVectorStore(embedder=HashingEmbedder(), client=client) + + assert second_store.search("stale content") == [] diff --git a/tests/test_retrieval_settings.py b/tests/test_retrieval_settings.py new file mode 100644 index 0000000..6f764ea --- /dev/null +++ b/tests/test_retrieval_settings.py @@ -0,0 +1,50 @@ +"""Tests for reading the retrieval mode and the Qdrant location from the environment.""" + +import pytest + +from rag.retrieval_settings import ( + DEFAULT_RETRIEVAL_MODE, + InvalidRetrievalModeError, + RetrievalMode, + load_retrieval_settings, +) + + +def test_hybrid_is_the_default_mode(): + assert DEFAULT_RETRIEVAL_MODE is RetrievalMode.HYBRID + assert load_retrieval_settings({}).mode is RetrievalMode.HYBRID + + +@pytest.mark.parametrize("name", ["tfidf", "semantic", "hybrid"]) +def test_every_documented_mode_can_be_selected(name): + settings = load_retrieval_settings({"RETRIEVAL_MODE": name}) + + assert settings.mode is RetrievalMode(name) + + +def test_the_mode_name_ignores_case_and_surrounding_spaces(): + assert load_retrieval_settings({"RETRIEVAL_MODE": " TfIdf "}).mode is RetrievalMode.TFIDF + + +def test_a_blank_mode_falls_back_to_the_default(): + assert load_retrieval_settings({"RETRIEVAL_MODE": " "}).mode is DEFAULT_RETRIEVAL_MODE + + +def test_an_unknown_mode_is_rejected_and_the_valid_ones_are_listed(): + with pytest.raises(InvalidRetrievalModeError) as error: + load_retrieval_settings({"RETRIEVAL_MODE": "bm25"}) + + message = str(error.value) + assert "bm25" in message + assert all(mode.value in message for mode in RetrievalMode) + + +def test_without_a_qdrant_url_the_in_memory_mode_is_implied(): + assert load_retrieval_settings({}).qdrant_url is None + assert load_retrieval_settings({"QDRANT_URL": " "}).qdrant_url is None + + +def test_a_qdrant_url_is_passed_through(): + settings = load_retrieval_settings({"QDRANT_URL": "http://localhost:6333"}) + + assert settings.qdrant_url == "http://localhost:6333" diff --git a/tests/test_store_factory.py b/tests/test_store_factory.py new file mode 100644 index 0000000..208ecb5 --- /dev/null +++ b/tests/test_store_factory.py @@ -0,0 +1,101 @@ +"""Tests for choosing the vector store that matches the configured retrieval mode.""" + +import pytest +from qdrant_client import QdrantClient + +from rag import store_factory +from rag.chunking import Chunk +from rag.hybrid_store import HybridVectorStore +from rag.qdrant_store import QdrantVectorStore +from rag.retrieval_settings import RetrievalMode, RetrievalSettings +from rag.store_factory import build_vector_store +from rag.vector_store import TfidfVectorStore +from tests.fakes import HashingEmbedder + + +class CountingEmbedderFactory: + def __init__(self): + self.calls = 0 + + def __call__(self) -> HashingEmbedder: + self.calls += 1 + return HashingEmbedder() + + +@pytest.fixture +def qdrant_connections(monkeypatch) -> list[dict]: + connections: list[dict] = [] + + def connect(**connection_arguments) -> QdrantClient: + connections.append(connection_arguments) + return QdrantClient(location=":memory:") + + monkeypatch.setattr(store_factory, "QdrantClient", connect) + return connections + + +def test_tfidf_mode_builds_the_lexical_store(): + settings = RetrievalSettings(mode=RetrievalMode.TFIDF) + + assert isinstance(build_vector_store(settings, CountingEmbedderFactory()), TfidfVectorStore) + + +def test_tfidf_mode_never_asks_for_an_embedding_model(): + embedder_factory = CountingEmbedderFactory() + + build_vector_store(RetrievalSettings(mode=RetrievalMode.TFIDF), embedder_factory) + + assert embedder_factory.calls == 0 + + +def test_semantic_mode_builds_the_qdrant_store(): + embedder_factory = CountingEmbedderFactory() + + store = build_vector_store(RetrievalSettings(mode=RetrievalMode.SEMANTIC), embedder_factory) + + assert isinstance(store, QdrantVectorStore) + assert embedder_factory.calls == 1 + + +def test_hybrid_mode_builds_the_hybrid_combiner(): + embedder_factory = CountingEmbedderFactory() + + store = build_vector_store(RetrievalSettings(mode=RetrievalMode.HYBRID), embedder_factory) + + assert isinstance(store, HybridVectorStore) + assert embedder_factory.calls == 1 + + +@pytest.mark.parametrize("mode", list(RetrievalMode)) +def test_every_mode_indexes_and_finds_a_document_through_the_common_interface(mode): + store = build_vector_store(RetrievalSettings(mode=mode), HashingEmbedder) + store.add_documents( + [ + Chunk(document_id="cats", index=0, text="Cats are small feline animals."), + Chunk(document_id="rockets", index=0, text="Rockets reach orbit."), + ] + ) + + assert store.search("feline animals")[0].chunk.document_id == "cats" + + +def test_a_configured_qdrant_url_is_used_to_connect(qdrant_connections): + settings = RetrievalSettings( + mode=RetrievalMode.SEMANTIC, qdrant_url="http://qdrant.internal:6333" + ) + + build_vector_store(settings, HashingEmbedder) + + assert qdrant_connections == [{"url": "http://qdrant.internal:6333"}] + + +def test_without_a_qdrant_url_the_in_memory_mode_is_used(qdrant_connections): + build_vector_store(RetrievalSettings(mode=RetrievalMode.SEMANTIC), HashingEmbedder) + + assert qdrant_connections == [{"location": ":memory:"}] + + +def test_tfidf_mode_never_connects_to_qdrant(qdrant_connections): + build_vector_store(RetrievalSettings(mode=RetrievalMode.TFIDF), HashingEmbedder) + + assert qdrant_connections == []