Build hybrid search with PostgreSQL

CREATE EXTENSION IF NOT EXISTS vector;

CREATE TABLE IF NOT EXISTS documents (
    source text PRIMARY KEY,
    sha256 text NOT NULL
);

CREATE TABLE IF NOT EXISTS chunks (
    id bigint GENERATED ALWAYS AS IDENTITY PRIMARY KEY,
    source text NOT NULL REFERENCES documents ON DELETE CASCADE,
    content text NOT NULL,
    embedding vector(1536) NOT NULL,
    tsv tsvector GENERATED ALWAYS AS (to_tsvector('english', content)) STORED
);

CREATE INDEX IF NOT EXISTS chunks_embedding_idx
    ON chunks USING hnsw (embedding vector_cosine_ops);

CREATE INDEX IF NOT EXISTS chunks_tsv_idx
    ON chunks USING gin (tsv);

import os
from pathlib import Path

import psycopg

schema = Path(__file__).with_name("schema.sql").read_text()
with psycopg.connect(os.environ["DB_URL"], autocommit=True) as conn:
    conn.execute(schema)
print("The schema is up to date")

from openai import AsyncOpenAI

# 1,536 dimensions, the size of the vector column in schema.sql
MODEL = "text-embedding-3-small"

client = AsyncOpenAI()


async def embed(texts: list[str]) -> list[str]:
    """Return one vector per text, in pgvector's text format: '[0.012, -0.034, ...]'."""
    response = await client.embeddings.create(model=MODEL, input=texts)
    return [str(item.embedding) for item in response.data]

import asyncio
import hashlib
import os
import sys
from pathlib import Path

import psycopg

from app.embeddings import embed

CHUNK_SIZE = 1500  # characters, about 375 tokens
OVERLAP = 200  # characters that each chunk repeats from the one before
BATCH_SIZE = 100  # chunks per embeddings request


def split(text: str) -> list[str]:
    """Split text into overlapping chunks, cut at a paragraph, line or word break."""
    chunks, start = [], 0
    while True:
        end = start + CHUNK_SIZE
        if end >= len(text):
            chunks.append(text[start:])
            return [chunk.strip() for chunk in chunks if chunk.strip()]
        for separator in ("\n\n", "\n", " "):
            cut = text.rfind(separator, start + CHUNK_SIZE // 2, end)
            if cut != -1:
                end = cut + len(separator)
                break
        chunks.append(text[start:end])
        # Start the next chunk at the first line or word in the overlap.
        overlap = end - OVERLAP
        boundary = text.find("\n", overlap, end)
        if boundary == -1:
            boundary = text.find(" ", overlap, end)
        start = boundary + 1 if boundary != -1 else overlap


async def ingest(folder: Path) -> None:
    files = sorted(p for p in folder.rglob("*") if p.suffix in {".md", ".mdx", ".txt"})
    sources = [p.relative_to(folder).as_posix() for p in files]
    async with await psycopg.AsyncConnection.connect(os.environ["DB_URL"], autocommit=True) as conn:
        for path, source in zip(files, sources):
            text = path.read_text()
            sha256 = hashlib.sha256(text.encode()).hexdigest()
            cursor = await conn.execute("SELECT sha256 FROM documents WHERE source = %s", (source,))
            if await cursor.fetchone() == (sha256,):
                continue  # unchanged since the last run

            chunks = [f"{source}\n\n{chunk}" for chunk in split(text)]
            vectors = []
            for i in range(0, len(chunks), BATCH_SIZE):
                vectors += await embed(chunks[i : i + BATCH_SIZE])

            async with conn.transaction():
                await conn.execute("DELETE FROM documents WHERE source = %s", (source,))
                await conn.execute(
                    "INSERT INTO documents (source, sha256) VALUES (%s, %s)", (source, sha256)
                )
                async with conn.cursor() as cursor:
                    await cursor.executemany(
                        "INSERT INTO chunks (source, content, embedding) VALUES (%s, %s, %s)",
                        [(source, chunk, vector) for chunk, vector in zip(chunks, vectors)],
                    )
            print(f"{source}: {len(chunks)} chunks")

        # Remove the documents whose files are gone, with their chunks.
        await conn.execute("DELETE FROM documents WHERE source <> ALL(%s)", (sources,))
        cursor = await conn.execute("SELECT count(*), count(DISTINCT source) FROM chunks")
        chunk_count, document_count = await cursor.fetchone()
        print(f"Ingested {chunk_count} chunks from {document_count} documents")


if __name__ == "__main__":
    asyncio.run(ingest(Path(sys.argv[1] if len(sys.argv) > 1 else "docs")))

import os
import secrets
from contextlib import asynccontextmanager
from typing import Annotated

from fastapi import Depends, FastAPI, HTTPException, Query
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
from psycopg.rows import dict_row
from psycopg_pool import AsyncConnectionPool

from app.embeddings import embed

CANDIDATES = 20  # results each search contributes to the fusion
RRF_K = 60  # reciprocal rank fusion constant

SEARCH = """
WITH semantic AS (
    SELECT id, rank() OVER (ORDER BY embedding <=> %(vector)s::vector) AS rank
    FROM chunks
    ORDER BY embedding <=> %(vector)s::vector
    LIMIT %(candidates)s
),
keyword AS (
    SELECT id, rank() OVER (ORDER BY ts_rank_cd(tsv, query) DESC) AS rank
    FROM chunks, websearch_to_tsquery('english', %(q)s) AS query
    WHERE tsv @@ query
    ORDER BY ts_rank_cd(tsv, query) DESC
    LIMIT %(candidates)s
),
fused AS (
    SELECT id, semantic.rank AS semantic_rank, keyword.rank AS keyword_rank,
           coalesce(1.0 / (%(rrf_k)s + semantic.rank), 0)
         + coalesce(1.0 / (%(rrf_k)s + keyword.rank), 0) AS score
    FROM semantic
    FULL OUTER JOIN keyword USING (id)
)
SELECT source, content, semantic_rank, keyword_rank, round(score, 4)::float AS score
FROM fused
JOIN chunks USING (id)
ORDER BY fused.score DESC, keyword_rank
LIMIT %(limit)s
"""

pool = AsyncConnectionPool(
    os.environ["DB_URL"],
    open=False,
    check=AsyncConnectionPool.check_connection,
    kwargs={"row_factory": dict_row},
)


@asynccontextmanager
async def lifespan(app: FastAPI):
    await pool.open()
    yield
    await pool.close()


app = FastAPI(lifespan=lifespan)


def require_token(
    credentials: Annotated[HTTPAuthorizationCredentials, Depends(HTTPBearer())],
) -> None:
    if not secrets.compare_digest(credentials.credentials, os.environ["SEARCH_TOKEN"]):
        raise HTTPException(status_code=401, detail="Invalid token")


@app.get("/")
async def index():
    async with pool.connection() as conn:
        cursor = await conn.execute("SELECT count(*) AS chunks FROM chunks")
        return await cursor.fetchone()


@app.get("/search", dependencies=[Depends(require_token)])
async def search(
    q: Annotated[str, Query(min_length=1, max_length=500)],
    limit: Annotated[int, Query(ge=1, le=20)] = 5,
):
    [vector] = await embed([q])
    params = {"vector": vector, "q": q, "candidates": CANDIDATES, "rrf_k": RRF_K, "limit": limit}
    async with pool.connection() as conn:
        cursor = await conn.execute(SEARCH, params)
        return {"query": q, "results": await cursor.fetchall()}

DB_URL=postgres://postgres:docsearch@localhost:5432/postgres
OPENAI_API_KEY=<your OpenAI API key>
SEARCH_TOKEN=local-token

ALTER TABLE chunks ADD COLUMN IF NOT EXISTS
    tsv tsvector GENERATED ALWAYS AS (to_tsvector('english', content)) STORED;

CREATE INDEX IF NOT EXISTS chunks_tsv_idx ON chunks USING gin (tsv);

FROM python:3.14-slim-trixie AS build
COPY --from=ghcr.io/astral-sh/uv:0.13 /uv /usr/local/bin/
ENV UV_COMPILE_BYTECODE=1 \
    UV_PYTHON_DOWNLOADS=0
WORKDIR /app
COPY pyproject.toml uv.lock .python-version ./
RUN uv sync --locked --no-dev

FROM python:3.14-slim-trixie
RUN useradd --uid 1000 --create-home app
ENV PATH="/app/.venv/bin:$PATH" \
    PYTHONUNBUFFERED=1 \
    PORT=8080 \
    FORWARDED_ALLOW_IPS="*"
WORKDIR /app
COPY --from=build /app/.venv .venv
COPY . .
RUN python -m compileall -q app
USER 1000:1000
EXPOSE 8080
CMD ["sh", "-c", "exec uvicorn app.main:app --host 0.0.0.0 --port \"$PORT\""]

.git
.venv
**/__pycache__
.env*
Dockerfile
.dockerignore