Skip to content

Add full-text keyword search with tsvector #32

Description

@AJaccP

🧠 Context

Retrieval today is embeddings-only: Repository.top_k_chunks in src/infrastructure/db/repository.py ranks chunks by cosine distance against the query embedding. That's strong for meaning ("how do I take fewer classes?") but weak for exact terms — a query like "when should I take COMP 1805" needs chunks that literally contain COMP 1805, and semantic search tends to blur exact course codes and acronyms into "some related course".

Postgres full-text search covers that gap. to_tsvector('english', content) normalizes a chunk's text into a sorted list of stemmed lexemes with stop-words removed; a query is normalized the same way and matched with the @@ operator. Reference: the Postgres Full Text Search chapter (see Searching a Table and ts_rank) and Generated Columns.

This ticket adds that capability at the storage + repository layer only:

  1. a generated content_tsv column on chunks plus a GIN index, and
  2. a new repository method that returns the top-k chunks for a text query, ranked by full-text relevance.

It does not change how the app retrieves today. Wiring keyword results together with the vector results (fusion) is a separate, later effort that will consume this method — so retrieval_service.py and the completion flow are out of scope here.

Files this ticket owns:

  • new src/infrastructure/db/versions/0002_*.py — the migration
  • src/infrastructure/db/models.py — declare the new column on ChunkRow
  • src/infrastructure/db/repository.py — add the keyword retrieval method
  • tests/infrastructure/db/ — new tests (repository behaviour + schema)

🛠 Implementation Plan

Note: Some functions/code snippets here may be inaccurate or not the best choices as you work through the ticket. Feel free to adjust as needed so long as the core functionality outlined here stays the same.

1. Migration — src/infrastructure/db/versions/0002_*.py

Clone the raw-SQL pattern from versions/0001_initial.py (op.execute(...), no autogenerate, no imports from models.py). Revision "0002", down_revision = "0001".

upgrade() adds a stored generated column and a GIN index:

  • content_tsv tsvector GENERATED ALWAYS AS (to_tsvector('english', content)) STORED NOT NULL — Postgres computes and maintains it on every insert/update, so there is no application code and no backfill; existing rows get populated automatically when the column is added. Add NOT NULL (valid on a generated column, and content is already NOT NULL so the tsvector never is) so the DB enforces what the model's nullable=False declares — otherwise the two drift.
  • CREATE INDEX ... USING GIN (content_tsv) — required for @@ lookups to be fast.

downgrade() must fully reverse this: drop the index, then drop the column. The test fixture runs downgrade base at setup and teardown, so a stubbed or incomplete downgrade breaks the entire test suite, not just one test.

Use the literal text-search config 'english' in the column definition, and use the same 'english' on the query side (step 3). Stemming and stop-words must match on both sides or nothing lines up.

2. Model — src/infrastructure/db/models.py

Add the column to ChunkRow so the ORM knows it exists and can reference it in queries. Declare it as a generated column (SQLAlchemy Computed(..., persisted=True)) with type TSVECTOR from sqlalchemy.dialects.postgresql. Marking it generated keeps the model honest with the migration and guarantees the existing upsert_chunk insert never tries to write it.

The column is only ever used inside query expressions — it is never read into the domain Chunk (which is built field-by-field in the repository). Load it deferred so a plain select(ChunkRow) doesn't pull the tsvector text over the wire.

3. Repository — src/infrastructure/db/repository.py

Add a static method alongside top_k_chunks, e.g. keyword_top_k_chunks(session, query: str, k: int) -> list[RetrievedChunk]:

  • Build the query with func.websearch_to_tsquery("english", query). Prefer websearch_to_tsquery over plainto_tsquery: it accepts free-form user input (quoted phrases, or, -term) and never raises on odd punctuation.
  • Filter to real matches with the @@ operator (ChunkRow.content_tsv.op("@@")(tsquery)) so only matching rows are returned and the GIN index is used.
  • Rank with func.ts_rank_cd(ChunkRow.content_tsv, tsquery) (cover density — rewards chunks where the query terms appear close together), order_by(rank.desc()), limit(k).
  • Build RetrievedChunk exactly like top_k_chunks does, but set score to the ts_rank_cd value. See the note below — this score is not on the same scale as the cosine score.

Return [] naturally when nothing matches.

4. Tests — tests/infrastructure/db/

The migration-based fixture (conftest.py) already applies migrations, so the new column and index exist in the test DB automatically.

  • Repository (extend test_repository.py): insert a few chunks with known content (chunks need an embedding — reuse whatever the existing repository tests use for that), then assert:
    • an exact-term query returns the chunk that contains that term and excludes ones that don't;
    • stemming works — a query using a different word form (e.g. reducing) matches a chunk containing another form (e.g. reduced);
    • a query with no shared lexemes returns an empty list;
    • results are capped at k.
  • Schema (extend test_schema.py): assert the migrated chunks table has a content_tsv column of type tsvector and that a GIN index exists on it.

📝 Notes

  • The score field is a lexical rank, not a similarity. ts_rank_cd returns a small, unbounded, corpus-relative number that is not comparable to the 1.0 - cosine_distance score used by the vector path. Don't try to normalize or blend the two here. A later ticket that combines keyword and vector results will fuse them by rank position, not by these raw numbers — so storing the raw ts_rank_cd in score is fine and intentional.
  • Keyword search is literal. to_tsvector splits COMP 1805 into two lexemes comp and 1805; a chunk that writes COMP1805 with no space becomes one token and won't match. That brittleness is expected and is exactly why this complements (doesn't replace) the embedding path — no need to work around it here.
  • Migration stays self-contained — raw SQL only, no import from models.py (keeps replay-from-scratch working). content_tsv is a generated column = pure DDL, no data backfill step.
  • Don't touch retrieval_service.py, completion_service.py, or settings unless necessary. Reuse the caller-supplied k; no new config.
  • make test still needs Docker up for the DB tests; run make lint and make test before pushing.

✅ Acceptance Criteria

  • A 0002 migration adds a NOT NULL content_tsv generated tsvector column (config 'english') and a GIN index on it, with a downgrade() that drops both; make migrate up and downgrade base both succeed.
  • ChunkRow declares content_tsv as a deferred generated column; upsert_chunk is unchanged and still works.
  • Repository.keyword_top_k_chunks(session, query, k) returns up to k RetrievedChunks ranked by ts_rank_cd, filtered by @@, using the 'english' config; empty list when nothing matches.
  • Tests cover exact-term match, stemming, no-match, and the k cap, plus a schema check for the column and GIN index.
  • No changes to retrieval_service.py/completion_service.py. make test and make lint pass.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Fields

    No fields configured for issues without a type.

    Projects

    Status
    Ready

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions