Build and update a hybrid PDF search index

import argparse
from io import BytesIO
from pathlib import Path

from pypdf import PdfReader
from pypdf.errors import PdfReadError


def extract_pages(data):
    reader = PdfReader(BytesIO(data))
    if reader.is_encrypted:
        raise ValueError("Use a PDF you can open without a password for this exercise.")
    return [(number, page.extract_text().strip())
            for number, page in enumerate(reader.pages, start=1)]


def main():
    parser = argparse.ArgumentParser(description="Inspect a PDF's embedded text by page.")
    parser.add_argument("pdf", type=Path)
    parser.add_argument("--page", type=int)
    args = parser.parse_args()

    try:
        pages = extract_pages(args.pdf.read_bytes())
        if args.page is not None:
            if not 1 <= args.page <= len(pages):
                raise ValueError(f"Choose a page from 1 to {len(pages)}.")
            number, text = pages[args.page - 1]
            print(f"Page {number}\n{text or '(no embedded text)'}")
        else:
            for number, text in pages:
                message = f"{len(text)} characters" if text else "no embedded text"
                print(f"Page {number}: {message}")
    except (OSError, ValueError, PdfReadError) as error:
        raise SystemExit(str(error)) from None


if __name__ == "__main__":
    main()

import argparse
from contextlib import closing
from dataclasses import dataclass
from hashlib import sha256
import json
from pathlib import Path
import re
import sqlite3
import textwrap

from pypdf.errors import PdfReadError
from pdf_pages import extract_pages
from semantic_notes import cosine, embed, model_digest

DEFAULT_DB = Path("handbook.sqlite3")


@dataclass(frozen=True)
class Passage:
    id: int
    source: str
    page: int
    text: str


def load_pdf_files(path):
    data = path.read_bytes()
    ocr_files = {}
    pattern = re.compile(re.escape(path.stem) + r"\.page-([1-9]\d*)\.txt")
    for sibling in sorted(path.parent.iterdir()):
        match = pattern.fullmatch(sibling.name)
        if match and not sibling.is_symlink() and sibling.is_file():
            ocr_files[int(match[1])] = sibling.read_bytes()
    checksum = sha256()
    checksum.update(sha256(data).digest())
    for page, text in sorted(ocr_files.items()):
        checksum.update(f"{page}:".encode("ascii") + sha256(text).digest())
    return data, ocr_files, checksum.hexdigest()


def replace_pdf(db, source, data, ocr_files, checksum):
    pages = extract_pages(data)
    if ocr_files.keys() - {number for number, _ in pages}:
        raise ValueError("An OCR text filename refers to a page outside the PDF.")
    pages = [
        (page, text or ocr_files.get(page, b"").decode("utf-8").strip())
        for page, text in pages
    ]
    chunks = [
        (page, chunk)
        for page, text in pages
        for chunk in textwrap.wrap(" ".join(text.split()), width=480)
    ]
    db.execute("INSERT INTO files VALUES (?, ?)", (source, checksum))
    for start in range(0, len(chunks), 16):
        batch = chunks[start:start + 16]
        vectors = embed([text for _, text in batch])
        for (page, text), vector in zip(batch, vectors):
            row = db.execute(
                "INSERT INTO passages(source, page, text, vector) VALUES (?, ?, ?, ?)",
                (source, page, text, json.dumps(vector)),
            )
            db.execute("INSERT INTO words(rowid, text) VALUES (?, ?)", (row.lastrowid, text))
    return [page for page, text in pages if not text]


def remove_pdf(db, source):
    db.execute("DELETE FROM words WHERE rowid IN (SELECT id FROM passages WHERE source = ?)", (source,))
    db.execute("DELETE FROM passages WHERE source = ?", (source,))
    db.execute("DELETE FROM files WHERE source = ?", (source,))


def update(folder, database):
    if not folder.is_dir():
        raise ValueError(f"PDF folder not found: {folder}")
    files = {
        path.relative_to(folder).as_posix(): path for path in sorted(folder.rglob("*"))
        if path.is_file() and not path.is_symlink() and path.suffix.lower() == ".pdf"
    }
    digest = model_digest()
    counts = {"indexed": 0, "reused": 0, "removed": 0}
    empty_pages = []
    with closing(sqlite3.connect(database)) as db, db:
        db.execute("BEGIN IMMEDIATE")
        db.execute("CREATE TABLE IF NOT EXISTS settings (id INTEGER PRIMARY KEY, digest TEXT NOT NULL)")
        db.execute("CREATE TABLE IF NOT EXISTS files (source TEXT PRIMARY KEY, checksum TEXT NOT NULL)")
        db.execute(
            "CREATE TABLE IF NOT EXISTS passages (id INTEGER PRIMARY KEY, source TEXT NOT NULL, "
            "page INTEGER NOT NULL, text TEXT NOT NULL, vector TEXT NOT NULL)"
        )
        db.execute("CREATE VIRTUAL TABLE IF NOT EXISTS words USING fts5(text)")
        stored = dict(db.execute("SELECT source, checksum FROM files"))
        revision = db.execute("SELECT digest FROM settings WHERE id = 1").fetchone()
        for source in sorted(stored.keys() - files.keys()):
            remove_pdf(db, source)
            counts["removed"] += 1
        for source, path in files.items():
            data, ocr_files, checksum = load_pdf_files(path)
            if stored.get(source) == checksum and revision == (digest,):
                counts["reused"] += 1
                continue
            remove_pdf(db, source)
            try:
                skipped = replace_pdf(db, source, data, ocr_files, checksum)
            except (PdfReadError, ValueError) as error:
                raise ValueError(f"{source}: {error}") from error
            empty_pages.extend((source, page) for page in skipped)
            counts["indexed"] += 1
        if model_digest() != digest:
            raise ValueError("The embedding model changed during the update. Run it again.")
        db.execute("INSERT OR REPLACE INTO settings VALUES (1, ?)", (digest,))

    for source, page in empty_pages:
        print(f"Skipped {source}, page {page}: no embedded text or OCR text file.")
    return counts


def fuse_rankings(*rankings):
    scores = {}
    for ranking in rankings:
        for rank, row_id in enumerate(ranking, start=1):
            scores[row_id] = scores.get(row_id, 0) + 1 / (60 + rank)
    return sorted(scores, key=lambda row_id: (-scores[row_id], row_id))


def search(database, query, mode="hybrid", limit=3):
    query = query.strip()
    if not query:
        raise ValueError("Enter a search question.")
    if not database.is_file():
        raise ValueError("Build the handbook index first.")
    terms = " OR ".join(f'"{word}"' for word in re.findall(r"[^\W_]+", query))
    with closing(sqlite3.connect(database.resolve().as_uri() + "?mode=ro", uri=True)) as db:
        db.execute("BEGIN")
        revision = db.execute("SELECT digest FROM settings WHERE id = 1").fetchone()
        keyword = [
            row[0] for row in db.execute(
                "SELECT rowid FROM words WHERE words MATCH ? ORDER BY rank, rowid LIMIT 8", (terms,)
            )
        ] if terms and mode != "vector" else []
        if mode == "keyword":
            hits = [db.execute(
                "SELECT id, source, page, text FROM passages WHERE id = ?", (row_id,)
            ).fetchone() for row_id in keyword[:limit]]
            return [Passage(*row) for row in hits]
        rows = db.execute("SELECT id, source, page, text, vector FROM passages ORDER BY id").fetchall()

    vector = []
    if rows:
        digest = model_digest()
        if revision != (digest,):
            raise ValueError("The embedding model changed. Run update before searching.")
        query_vector, = embed([query])
        if model_digest() != digest:
            raise ValueError("The embedding model changed during search. Run update and retry.")
        ranked = sorted(rows, key=lambda row: cosine(json.loads(row[4]), query_vector), reverse=True)
        vector = [row[0] for row in ranked[:8]]

    ids = fuse_rankings(keyword, vector) if mode == "hybrid" else vector
    candidates = {row[0]: Passage(*row[:4]) for row in rows}
    return [candidates[row_id] for row_id in ids[:limit]]


def main():
    parser = argparse.ArgumentParser(description="Update and search a small collection of PDF files.")
    parser.add_argument("--db", type=Path, default=DEFAULT_DB, help="Dedicated handbook SQLite file")
    commands = parser.add_subparsers(dest="command", required=True)
    rebuild = commands.add_parser("update")
    rebuild.add_argument("folder", type=Path)
    find = commands.add_parser("find")
    find.add_argument("query")
    find.add_argument("--mode", choices=("keyword", "vector", "hybrid"), default="hybrid")
    args = parser.parse_args()
    try:
        if args.command == "update":
            counts = update(args.folder, args.db)
            print(f"Indexed {counts['indexed']}, reused {counts['reused']}, removed {counts['removed']} PDFs.")
        else:
            matches = search(args.db, args.query, args.mode)
            for passage in matches:
                print(f"{passage.source}, page {passage.page}\n{passage.text}\n")
            if not matches:
                print("No matching passages.")
    except (OSError, sqlite3.Error, ValueError) as error:
        raise SystemExit(str(error)) from None


if __name__ == "__main__":
    main()

import argparse
from contextlib import closing
from itertools import islice
import json
import math
from pathlib import Path
import sqlite3
from urllib.error import HTTPError
from urllib.request import HTTPRedirectHandler, ProxyHandler, Request, build_opener

BASE = "http://127.0.0.1:11434"
MODEL = "all-minilm:latest"
MAX_PASSAGES = 200


class NoRedirects(HTTPRedirectHandler):
    def redirect_request(self, req, fp, code, msg, headers, newurl):
        return None


def api(endpoint, payload=None):
    request = Request(BASE + endpoint)
    if payload is not None:
        request.data = json.dumps(payload).encode("utf-8")
        request.add_header("Content-Type", "application/json")
    opener = build_opener(ProxyHandler({}), NoRedirects())
    with opener.open(request, timeout=120) as response:
        return json.load(response)


def model_digest():
    for model in api("/api/tags")["models"]:
        if model["name"] == MODEL:
            return model["digest"]
    raise ValueError(f"Pull {MODEL} with Ollama first.")


def embed(texts):
    result = api("/api/embed", {
        "model": MODEL,
        "input": texts,
        "truncate": False,
    })["embeddings"]
    if len(result) != len(texts):
        raise ValueError("Ollama returned the wrong number of embeddings.")
    return result


def cosine(left, right):
    if len(left) != len(right):
        raise ValueError("Vector dimensions changed. Rebuild the index.")
    length = math.hypot(*left) * math.hypot(*right)
    return sum(x * y for x, y in zip(left, right)) / length if length else 0


def build_index(folder, database):
    from index_notes import passages

    rows = list(islice(passages(folder), MAX_PASSAGES + 1))
    if not rows or len(rows) > MAX_PASSAGES:
        raise ValueError(f"Use 1–{MAX_PASSAGES} nonempty passages for this example.")
    digest = model_digest()
    stored = []
    for start in range(0, len(rows), 16):
        batch = rows[start:start + 16]
        vectors = embed([text for _, _, text in batch])
        stored.extend((*row, json.dumps(vector)) for row, vector in zip(batch, vectors))
    if model_digest() != digest:
        raise ValueError("The embedding model changed during indexing. Build the index again.")
    with closing(sqlite3.connect(database)) as connection, connection:
        connection.execute("CREATE TABLE IF NOT EXISTS model (name TEXT, digest TEXT)")
        connection.execute(
            "CREATE TABLE IF NOT EXISTS passages "
            "(source TEXT, number INTEGER, text TEXT, vector TEXT)"
        )
        connection.execute("DELETE FROM passages")
        connection.execute("DELETE FROM model")
        connection.execute("INSERT INTO model VALUES (?, ?)", (MODEL, digest))
        connection.executemany("INSERT INTO passages VALUES (?, ?, ?, ?)", stored)
    print(f"Saved {len(stored)} embeddings in {database}.")


def find_notes(database, query):
    if not query.strip() or not database.is_file():
        raise ValueError("Build the index and enter a nonempty question.")
    with closing(sqlite3.connect(database.resolve().as_uri() + "?mode=ro", uri=True)) as connection:
        connection.execute("BEGIN")
        identity = connection.execute("SELECT name, digest FROM model").fetchone()
        rows = connection.execute("SELECT source, number, text, vector FROM passages").fetchall()
    digest = model_digest()
    if identity != (MODEL, digest):
        raise ValueError("The embedding model differs from the one in the index. Rebuild it.")
    vector, = embed([query])
    if model_digest() != digest:
        raise ValueError("The embedding model changed during the query. Rebuild the index.")
    matches = [
        (cosine(json.loads(embedding), vector), source, number, text)
        for source, number, text, embedding in rows
    ]
    ranked = sorted(matches, key=lambda match: match[0], reverse=True)
    for score, source, number, text in ranked[:3]:
        print(f"{score:.3f} | {source}, passage {number}\n{text}\n")


def main():
    parser = argparse.ArgumentParser(description="Store local embeddings and search them later.")
    parser.add_argument("--db", type=Path, default=Path("notes-vectors.sqlite3"))
    commands = parser.add_subparsers(dest="command", required=True)
    build = commands.add_parser("build")
    build.add_argument("folder", type=Path)
    find = commands.add_parser("find")
    find.add_argument("query")
    args = parser.parse_args()

    try:
        if args.command == "build":
            build_index(args.folder, args.db)
        else:
            find_notes(args.db, args.query)
    except HTTPError as error:
        raise SystemExit(f"Ollama returned HTTP {error.code}: {error.read().decode()}") from None
    except (OSError, sqlite3.Error, ValueError) as error:
        raise SystemExit(str(error)) from None


if __name__ == "__main__":
    main()