From 2e8c0570e409bd06a296b2a4da17322cce58ea49 Mon Sep 17 00:00:00 2001 From: Richie Cahill Date: Wed, 10 Jun 2026 20:16:59 -0400 Subject: [PATCH 01/28] adding embedding Models to jeeves --- systems/bob/llms.nix | 2 +- systems/jeeves/services/llms.nix | 5 ++++- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/systems/bob/llms.nix b/systems/bob/llms.nix index 9476ea4..6d30844 100644 --- a/systems/bob/llms.nix +++ b/systems/bob/llms.nix @@ -4,7 +4,7 @@ host = "0.0.0.0"; enable = true; - syncModels = true; + syncModels = false; loadModels = [ "codellama:7b" "deepscaler:1.5b" diff --git a/systems/jeeves/services/llms.nix b/systems/jeeves/services/llms.nix index cdcd0ec..3abe93d 100644 --- a/systems/jeeves/services/llms.nix +++ b/systems/jeeves/services/llms.nix @@ -6,7 +6,7 @@ in user = "ollama"; enable = true; host = "0.0.0.0"; - syncModels = true; + syncModels = false; loadModels = [ "codellama:7b" "deepscaler:1.5b" @@ -30,6 +30,9 @@ in "ministral-3:14b" "nemotron-3-nano:30b" "qwen3-coder:30b" + "qwen3-embedding:0.6b" + "qwen3-embedding:4b" + "qwen3-embedding:8b" "qwen3-vl:32b" "qwen3:14b" "qwen3.5:35b" -- 2.54.0 From 3ee884f6b419b58723a8a2f5233fe23654adcb6e Mon Sep 17 00:00:00 2001 From: Richie Cahill Date: Wed, 10 Jun 2026 20:18:16 -0400 Subject: [PATCH 02/28] removed hedgedoc --- systems/jeeves/services/postgress.nix | 11 ----------- 1 file changed, 11 deletions(-) diff --git a/systems/jeeves/services/postgress.nix b/systems/jeeves/services/postgress.nix index 546ffab..06a4747 100644 --- a/systems/jeeves/services/postgress.nix +++ b/systems/jeeves/services/postgress.nix @@ -38,9 +38,6 @@ in # signalbot local signalbot signalbot trust - # hedgedoc - local hedgedoc hedgedoc trust - # math local postgres math trust host postgres math 127.0.0.1/32 trust @@ -120,19 +117,11 @@ in login = true; }; } - { - name = "hedgedoc"; - ensureDBOwnership = true; - ensureClauses = { - login = true; - }; - } ]; ensureDatabases = [ "data_science_dev" "hass" "gitea" - "hedgedoc" "math" "n8n" "richie" -- 2.54.0 From 62890007665e60d7418e7ce9d7510c31daf843af Mon Sep 17 00:00:00 2001 From: Richie Cahill Date: Thu, 11 Jun 2026 18:08:31 -0400 Subject: [PATCH 03/28] added ebook embedding to orm --- ...10-add_ebook_search_tables_2db132cace1a.py | 200 ++++++++++++++++++ python/orm/richie/__init__.py | 16 ++ python/orm/richie/ebook.py | 130 ++++++++++++ 3 files changed, 346 insertions(+) create mode 100644 python/alembic/richie/versions/2026_06_10-add_ebook_search_tables_2db132cace1a.py create mode 100644 python/orm/richie/ebook.py diff --git a/python/alembic/richie/versions/2026_06_10-add_ebook_search_tables_2db132cace1a.py b/python/alembic/richie/versions/2026_06_10-add_ebook_search_tables_2db132cace1a.py new file mode 100644 index 0000000..f400d75 --- /dev/null +++ b/python/alembic/richie/versions/2026_06_10-add_ebook_search_tables_2db132cace1a.py @@ -0,0 +1,200 @@ +"""add ebook search tables. + +Revision ID: 2db132cace1a +Revises: b3c60cc5beb5 +Create Date: 2026-06-10 22:10:54.379159 + +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import pgvector +import sqlalchemy as sa +from alembic import op + +from python.orm import RichieBase + +if TYPE_CHECKING: + from collections.abc import Sequence + +# revision identifiers, used by Alembic. +revision: str = "2db132cace1a" +down_revision: str | None = "b3c60cc5beb5" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + +schema = RichieBase.schema_name + + +def upgrade() -> None: + """Upgrade.""" + # ### commands auto generated by Alembic - please adjust! ### + op.create_table( + "ebook_embedding_model", + sa.Column("name", sa.String(), nullable=False), + sa.Column("dimension", sa.Integer(), nullable=False), + sa.Column("is_default", sa.Boolean(), nullable=False), + sa.Column("id", sa.Integer(), nullable=False), + sa.Column("created", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False), + sa.Column("updated", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False), + sa.PrimaryKeyConstraint("id", name=op.f("pk_ebook_embedding_model")), + sa.UniqueConstraint("name", name=op.f("uq_ebook_embedding_model_name")), + schema=schema, + ) + op.create_table( + "ebook_source", + sa.Column("title", sa.String(), nullable=False), + sa.Column("author", sa.String(), nullable=True), + sa.Column("language", sa.String(), nullable=True), + sa.Column("publisher", sa.String(), nullable=True), + sa.Column("identifier", sa.String(), nullable=True), + sa.Column("file_path", sa.String(), nullable=False), + sa.Column("file_sha256", sa.String(length=64), nullable=False), + sa.Column("file_mtime", sa.DateTime(timezone=True), nullable=False), + sa.Column("file_size", sa.BigInteger(), nullable=False), + sa.Column("id", sa.Integer(), nullable=False), + sa.Column("created", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False), + sa.Column("updated", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False), + sa.PrimaryKeyConstraint("id", name=op.f("pk_ebook_source")), + sa.UniqueConstraint("file_path", name=op.f("uq_ebook_source_file_path")), + sa.UniqueConstraint("file_sha256", name=op.f("uq_ebook_source_file_sha256")), + schema=schema, + ) + op.create_table( + "ebook_chapter", + sa.Column("source_id", sa.Integer(), nullable=False), + sa.Column("spine_index", sa.Integer(), nullable=False), + sa.Column("title", sa.String(), nullable=True), + sa.Column("href", sa.String(), nullable=True), + sa.Column("id", sa.Integer(), nullable=False), + sa.Column("created", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False), + sa.Column("updated", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False), + sa.ForeignKeyConstraint( + ["source_id"], + [f"{schema}.ebook_source.id"], + name=op.f("fk_ebook_chapter_source_id_ebook_source"), + ondelete="CASCADE", + ), + sa.PrimaryKeyConstraint("id", name=op.f("pk_ebook_chapter")), + sa.UniqueConstraint("source_id", "spine_index", name=op.f("uq_ebook_chapter_source_id")), + schema=schema, + ) + op.create_table( + "ebook_chunk", + sa.Column("source_id", sa.Integer(), nullable=False), + sa.Column("chapter_id", sa.Integer(), nullable=True), + sa.Column("chunk_index", sa.Integer(), nullable=False), + sa.Column("text", sa.String(), nullable=False), + sa.Column("token_start", sa.Integer(), nullable=False), + sa.Column("token_count", sa.Integer(), nullable=False), + sa.Column("page_label", sa.String(), nullable=True), + sa.Column("content_sha256", sa.String(length=64), nullable=False), + sa.Column("search_text", sa.String(), nullable=False), + sa.Column("id", sa.BigInteger(), nullable=False), + sa.Column("created", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False), + sa.Column("updated", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False), + sa.ForeignKeyConstraint( + ["chapter_id"], + [f"{schema}.ebook_chapter.id"], + name=op.f("fk_ebook_chunk_chapter_id_ebook_chapter"), + ondelete="SET NULL", + ), + sa.ForeignKeyConstraint( + ["source_id"], + [f"{schema}.ebook_source.id"], + name=op.f("fk_ebook_chunk_source_id_ebook_source"), + ondelete="CASCADE", + ), + sa.PrimaryKeyConstraint("id", name=op.f("pk_ebook_chunk")), + sa.UniqueConstraint("source_id", "chunk_index", name="uq_ebook_chunk_source_id_chunk_index"), + sa.UniqueConstraint("source_id", "content_sha256", name="uq_ebook_chunk_source_id_content_sha256"), + schema=schema, + ) + op.create_table( + "ebook_chunk_embedding_1024", + sa.Column("chunk_id", sa.BigInteger(), nullable=False), + sa.Column("model_id", sa.Integer(), nullable=False), + sa.Column("embedding", pgvector.sqlalchemy.vector.VECTOR(dim=1024), nullable=False), + sa.Column("id", sa.BigInteger(), nullable=False), + sa.Column("created", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False), + sa.Column("updated", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False), + sa.ForeignKeyConstraint( + ["chunk_id"], + [f"{schema}.ebook_chunk.id"], + name=op.f("fk_ebook_chunk_embedding_1024_chunk_id_ebook_chunk"), + ondelete="CASCADE", + ), + sa.ForeignKeyConstraint( + ["model_id"], + [f"{schema}.ebook_embedding_model.id"], + name=op.f("fk_ebook_chunk_embedding_1024_model_id_ebook_embedding_model"), + ondelete="CASCADE", + ), + sa.PrimaryKeyConstraint("id", name=op.f("pk_ebook_chunk_embedding_1024")), + sa.UniqueConstraint("chunk_id", "model_id", name=op.f("uq_ebook_chunk_embedding_1024_chunk_id")), + schema=schema, + ) + op.create_table( + "ebook_chunk_embedding_2560", + sa.Column("chunk_id", sa.BigInteger(), nullable=False), + sa.Column("model_id", sa.Integer(), nullable=False), + sa.Column("embedding", pgvector.sqlalchemy.vector.VECTOR(dim=2560), nullable=False), + sa.Column("id", sa.BigInteger(), nullable=False), + sa.Column("created", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False), + sa.Column("updated", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False), + sa.ForeignKeyConstraint( + ["chunk_id"], + [f"{schema}.ebook_chunk.id"], + name=op.f("fk_ebook_chunk_embedding_2560_chunk_id_ebook_chunk"), + ondelete="CASCADE", + ), + sa.ForeignKeyConstraint( + ["model_id"], + [f"{schema}.ebook_embedding_model.id"], + name=op.f("fk_ebook_chunk_embedding_2560_model_id_ebook_embedding_model"), + ondelete="CASCADE", + ), + sa.PrimaryKeyConstraint("id", name=op.f("pk_ebook_chunk_embedding_2560")), + sa.UniqueConstraint("chunk_id", "model_id", name=op.f("uq_ebook_chunk_embedding_2560_chunk_id")), + schema=schema, + ) + op.create_table( + "ebook_chunk_embedding_4096", + sa.Column("chunk_id", sa.BigInteger(), nullable=False), + sa.Column("model_id", sa.Integer(), nullable=False), + sa.Column("embedding", pgvector.sqlalchemy.vector.VECTOR(dim=4096), nullable=False), + sa.Column("id", sa.BigInteger(), nullable=False), + sa.Column("created", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False), + sa.Column("updated", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False), + sa.ForeignKeyConstraint( + ["chunk_id"], + [f"{schema}.ebook_chunk.id"], + name=op.f("fk_ebook_chunk_embedding_4096_chunk_id_ebook_chunk"), + ondelete="CASCADE", + ), + sa.ForeignKeyConstraint( + ["model_id"], + [f"{schema}.ebook_embedding_model.id"], + name=op.f("fk_ebook_chunk_embedding_4096_model_id_ebook_embedding_model"), + ondelete="CASCADE", + ), + sa.PrimaryKeyConstraint("id", name=op.f("pk_ebook_chunk_embedding_4096")), + sa.UniqueConstraint("chunk_id", "model_id", name=op.f("uq_ebook_chunk_embedding_4096_chunk_id")), + schema=schema, + ) + # ### end Alembic commands ### + + +def downgrade() -> None: + """Downgrade.""" + # ### commands auto generated by Alembic - please adjust! ### + op.drop_table("ebook_chunk_embedding_4096", schema=schema) + op.drop_table("ebook_chunk_embedding_2560", schema=schema) + op.drop_table("ebook_chunk_embedding_1024", schema=schema) + op.drop_table("ebook_chunk", schema=schema) + op.drop_table("ebook_chapter", schema=schema) + op.drop_table("ebook_source", schema=schema) + op.drop_table("ebook_embedding_model", schema=schema) + # ### end Alembic commands ### diff --git a/python/orm/richie/__init__.py b/python/orm/richie/__init__.py index 47f601f..a28ce7a 100644 --- a/python/orm/richie/__init__.py +++ b/python/orm/richie/__init__.py @@ -11,6 +11,15 @@ from python.orm.richie.contact import ( Need, RelationshipType, ) +from python.orm.richie.ebook import ( + EbookChapter, + EbookChunk, + EbookChunkEmbedding1024, + EbookChunkEmbedding2560, + EbookChunkEmbedding4096, + EbookEmbeddingModel, + EbookSource, +) __all__ = [ "Audiobook", @@ -19,6 +28,13 @@ __all__ = [ "Contact", "ContactNeed", "ContactRelationship", + "EbookChapter", + "EbookChunk", + "EbookChunkEmbedding1024", + "EbookChunkEmbedding2560", + "EbookChunkEmbedding4096", + "EbookEmbeddingModel", + "EbookSource", "Need", "RelationshipType", "RichieBase", diff --git a/python/orm/richie/ebook.py b/python/orm/richie/ebook.py new file mode 100644 index 0000000..9c1e4ad --- /dev/null +++ b/python/orm/richie/ebook.py @@ -0,0 +1,130 @@ +"""EPUB search models.""" + +from __future__ import annotations + +from datetime import datetime + +from pgvector.sqlalchemy import Vector +from sqlalchemy import BigInteger, Boolean, DateTime, ForeignKey, String, UniqueConstraint +from sqlalchemy.orm import Mapped, mapped_column, relationship + +from python.orm.richie.base import TableBase, TableBaseBig + + +class EbookSource(TableBase): + """One indexed EPUB file.""" + + __tablename__ = "ebook_source" + __table_args__ = ( + UniqueConstraint("file_path"), + UniqueConstraint("file_sha256"), + ) + + title: Mapped[str] + author: Mapped[str | None] + language: Mapped[str | None] + publisher: Mapped[str | None] + identifier: Mapped[str | None] + file_path: Mapped[str] + file_sha256: Mapped[str] = mapped_column(String(64)) + file_mtime: Mapped[datetime] = mapped_column(DateTime(timezone=True)) + file_size: Mapped[int] = mapped_column(BigInteger) + + chapters: Mapped[list[EbookChapter]] = relationship( + "EbookChapter", + back_populates="source", + cascade="all, delete-orphan", + passive_deletes=True, + ) + chunks: Mapped[list[EbookChunk]] = relationship( + "EbookChunk", + back_populates="source", + cascade="all, delete-orphan", + passive_deletes=True, + ) + + +class EbookChapter(TableBase): + """A chapter or spine document inside an EPUB.""" + + __tablename__ = "ebook_chapter" + __table_args__ = (UniqueConstraint("source_id", "spine_index"),) + + source_id: Mapped[int] = mapped_column(ForeignKey("main.ebook_source.id", ondelete="CASCADE")) + spine_index: Mapped[int] + title: Mapped[str | None] + href: Mapped[str | None] + + source: Mapped[EbookSource] = relationship("EbookSource", back_populates="chapters") + chunks: Mapped[list[EbookChunk]] = relationship( + "EbookChunk", + back_populates="chapter", + cascade="all, delete-orphan", + passive_deletes=True, + ) + + +class EbookChunk(TableBaseBig): + """A searchable text chunk.""" + + __tablename__ = "ebook_chunk" + __table_args__ = ( + UniqueConstraint("source_id", "chunk_index", name="uq_ebook_chunk_source_id_chunk_index"), + UniqueConstraint("source_id", "content_sha256", name="uq_ebook_chunk_source_id_content_sha256"), + ) + + source_id: Mapped[int] = mapped_column(ForeignKey("main.ebook_source.id", ondelete="CASCADE")) + chapter_id: Mapped[int | None] = mapped_column(ForeignKey("main.ebook_chapter.id", ondelete="SET NULL")) + chunk_index: Mapped[int] + text: Mapped[str] + token_start: Mapped[int] + token_count: Mapped[int] + page_label: Mapped[str | None] + content_sha256: Mapped[str] = mapped_column(String(64)) + search_text: Mapped[str] + + source: Mapped[EbookSource] = relationship("EbookSource", back_populates="chunks") + chapter: Mapped[EbookChapter | None] = relationship("EbookChapter", back_populates="chunks") + + +class EbookEmbeddingModel(TableBase): + """A supported embedding model.""" + + __tablename__ = "ebook_embedding_model" + + name: Mapped[str] = mapped_column(String, unique=True) + dimension: Mapped[int] + is_default: Mapped[bool] = mapped_column(Boolean, default=False) + + +class EbookChunkEmbedding1024(TableBaseBig): + """1024-dimensional chunk embedding.""" + + __tablename__ = "ebook_chunk_embedding_1024" + __table_args__ = (UniqueConstraint("chunk_id", "model_id"),) + + chunk_id: Mapped[int] = mapped_column(ForeignKey("main.ebook_chunk.id", ondelete="CASCADE")) + model_id: Mapped[int] = mapped_column(ForeignKey("main.ebook_embedding_model.id", ondelete="CASCADE")) + embedding: Mapped[list[float]] = mapped_column(Vector(1024)) + + +class EbookChunkEmbedding2560(TableBaseBig): + """2560-dimensional chunk embedding.""" + + __tablename__ = "ebook_chunk_embedding_2560" + __table_args__ = (UniqueConstraint("chunk_id", "model_id"),) + + chunk_id: Mapped[int] = mapped_column(ForeignKey("main.ebook_chunk.id", ondelete="CASCADE")) + model_id: Mapped[int] = mapped_column(ForeignKey("main.ebook_embedding_model.id", ondelete="CASCADE")) + embedding: Mapped[list[float]] = mapped_column(Vector(2560)) + + +class EbookChunkEmbedding4096(TableBaseBig): + """4096-dimensional chunk embedding.""" + + __tablename__ = "ebook_chunk_embedding_4096" + __table_args__ = (UniqueConstraint("chunk_id", "model_id"),) + + chunk_id: Mapped[int] = mapped_column(ForeignKey("main.ebook_chunk.id", ondelete="CASCADE")) + model_id: Mapped[int] = mapped_column(ForeignKey("main.ebook_embedding_model.id", ondelete="CASCADE")) + embedding: Mapped[list[float]] = mapped_column(Vector(4096)) -- 2.54.0 From aa135a3af2c00e918552f4549ba81d6549ec2d7e Mon Sep 17 00:00:00 2001 From: Richie Cahill Date: Thu, 11 Jun 2026 22:08:48 -0400 Subject: [PATCH 04/28] clean up --- tests/test_audible_convert.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/test_audible_convert.py b/tests/test_audible_convert.py index 28e9c5f..d3b9fae 100644 --- a/tests/test_audible_convert.py +++ b/tests/test_audible_convert.py @@ -1021,8 +1021,6 @@ def test_existing_destination_skips_rename_and_removes_temp(tmp_path, monkeypatc def test_richie_exports_audiobook_models() -> None: - from python.orm.richie import Audiobook # noqa: PLC0415 - assert Audiobook.__tablename__ == "audiobook" -- 2.54.0 From 5087fbb4c0da79f62c130ab9b7faa8d5c8efe504 Mon Sep 17 00:00:00 2001 From: Richie Cahill Date: Fri, 12 Jun 2026 01:30:20 -0400 Subject: [PATCH 05/28] built BM25 search foundation --- python/ebook_search/bm25_corpus.py | 249 +++++++++++++++++++++++++++++ 1 file changed, 249 insertions(+) create mode 100644 python/ebook_search/bm25_corpus.py diff --git a/python/ebook_search/bm25_corpus.py b/python/ebook_search/bm25_corpus.py new file mode 100644 index 0000000..79172a4 --- /dev/null +++ b/python/ebook_search/bm25_corpus.py @@ -0,0 +1,249 @@ +"""Persisted BM25 corpus management.""" + +from __future__ import annotations + +import json +import logging +import shutil +import tempfile +from dataclasses import dataclass +from datetime import UTC, datetime +from functools import cache +from pathlib import Path +from typing import TYPE_CHECKING + +import bm25s +from sqlalchemy import func, select, union_all + +from python.orm.richie import EbookChapter, EbookChunk, EbookSource + +if TYPE_CHECKING: + from sqlalchemy.orm import Session + + from python.ebook_search.config import EbookSearchConfig + +logger = logging.getLogger(__name__) +MANIFEST_NAME = "manifest.json" +REQUIRED_INDEX_FILES = frozenset( + { + "data.csc.index.npy", + "indices.csc.index.npy", + "indptr.csc.index.npy", + "params.index.json", + "vocab.index.json", + "corpus.jsonl", + } +) + + +@dataclass(frozen=True) +class BM25Manifest: + """Metadata describing a persisted BM25 corpus.""" + + created_at: datetime + db_updated_at: datetime | None + chunk_count: int + + +@dataclass(frozen=True) +class BM25Corpus: + """Loaded persisted BM25 corpus and retriever.""" + + retriever: object | None + records: tuple[dict[str, object], ...] + manifest: BM25Manifest + + +class BM25CorpusUnavailableError(RuntimeError): + """Raised when the persisted BM25 corpus cannot be loaded.""" + + +def bm25_index_path(config: EbookSearchConfig) -> Path: + """Return the configured BM25 index path relative to the current working directory.""" + path = Path(config.bm25_index_dir).expanduser() + if path.is_absolute(): + return path + return Path.cwd() / path + + +def ensure_bm25_corpus(session: Session, config: EbookSearchConfig) -> None: + """Create or refresh the persisted BM25 corpus when it is missing or stale.""" + index_path = bm25_index_path(config) + manifest = read_bm25_manifest(index_path) + db_updated_at = corpus_last_updated_at(session) + if not bm25_index_exists(index_path, manifest): + logger.info("ebook_bm25_index_missing path=%s", index_path) + refresh_bm25_corpus(session, config, db_updated_at=db_updated_at) + return + if db_updated_at is not None and manifest is not None and manifest.created_at < db_updated_at: + logger.info( + "ebook_bm25_index_stale path=%s created_at=%s db_updated_at=%s", + index_path, + manifest.created_at.isoformat(), + db_updated_at.isoformat(), + ) + refresh_bm25_corpus(session, config, db_updated_at=db_updated_at) + return + logger.info( + "ebook_bm25_index_current path=%s chunks=%s created_at=%s", + index_path, + manifest.chunk_count if manifest else 0, + manifest.created_at.isoformat() if manifest else None, + ) + + +def refresh_bm25_corpus( + session: Session, + config: EbookSearchConfig, + *, + db_updated_at: datetime | None = None, +) -> BM25Manifest: + """Rebuild and persist the BM25 corpus from the current database chunks.""" + index_path = bm25_index_path(config) + records = fetch_bm25_corpus_records(session) + manifest = BM25Manifest( + created_at=datetime.now(tz=UTC), + db_updated_at=db_updated_at if db_updated_at is not None else corpus_last_updated_at(session), + chunk_count=len(records), + ) + write_bm25_corpus(index_path, records, manifest) + logger.info( + "ebook_bm25_index_refreshed path=%s chunks=%s created_at=%s note=%s", + index_path, + manifest.chunk_count, + manifest.created_at.isoformat(), + "restart_service_to_use_refreshed_bm25_cache", + ) + return manifest + + +@cache +def load_bm25_corpus(config: EbookSearchConfig) -> BM25Corpus: + """Load the BM25 corpus into memory once per process. + + This cache intentionally does not notice later on-disk corpus refreshes. Restart the service after rebuilding the + BM25 corpus for searches to use the new index. + """ + index_path = bm25_index_path(config) + logger.info( + "ebook_bm25_corpus_cache_load path=%s note=%s", + index_path, + "restart_service_after_bm25_refresh", + ) + manifest = read_bm25_manifest(index_path) + if manifest is None or not bm25_index_exists(index_path, manifest): + msg = f"BM25 corpus is not available: {index_path}" + raise BM25CorpusUnavailableError(msg) + if manifest.chunk_count == 0: + return BM25Corpus(retriever=None, records=(), manifest=manifest) + + retriever = bm25s.BM25.load(index_path, load_corpus=True, mmap=True) + records = tuple(dict(record) for record in retriever.corpus) + return BM25Corpus(retriever=retriever, records=records, manifest=manifest) + + +def score_bm25_corpus(query: str, corpus: BM25Corpus, *, limit: int) -> list[tuple[dict[str, object], float]]: + """Score a query against a loaded BM25 corpus.""" + if corpus.retriever is None or not corpus.records: + return [] + k = min(limit, len(corpus.records)) + documents, scores = corpus.retriever.retrieve( + bm25s.tokenize(query, show_progress=False), + corpus=list(corpus.records), + k=k, + show_progress=False, + ) + results: list[tuple[dict[str, object], float]] = [] + for document, score in zip(documents[0], scores[0], strict=True): + score_value = float(score) + if score_value <= 0: + continue + results.append((dict(document), score_value)) + return results + + +def fetch_bm25_corpus_records(session: Session) -> list[dict[str, object]]: + """Fetch BM25 corpus records from the database.""" + statement = ( + select( + EbookChunk.id.label("chunk_id"), + EbookChunk.text.label("text"), + EbookSource.title.label("source_title"), + EbookSource.author.label("source_author"), + EbookChapter.title.label("chapter_title"), + EbookChunk.page_label.label("page_label"), + func.concat_ws( + " ", + EbookSource.title, + EbookSource.author, + EbookChapter.title, + EbookChunk.search_text, + ).label("bm25_text"), + ) + .select_from(EbookChunk) + .join(EbookSource, EbookSource.id == EbookChunk.source_id) + .outerjoin(EbookChapter, EbookChapter.id == EbookChunk.chapter_id) + .order_by(EbookChunk.id) + ) + return [dict(row) for row in session.execute(statement).mappings()] + + +def corpus_last_updated_at(session: Session) -> datetime | None: + """Return the latest source/chapter/chunk update timestamp relevant to BM25 text.""" + update_times = union_all( + select(func.max(EbookSource.updated).label("updated")), + select(func.max(EbookChapter.updated).label("updated")), + select(func.max(EbookChunk.updated).label("updated")), + ).subquery() + return session.scalar(select(func.max(update_times.c.updated))) + + +def write_bm25_corpus(index_path: Path, records: list[dict[str, object]], manifest: BM25Manifest) -> None: + """Write a BM25 corpus and manifest atomically.""" + index_path.parent.mkdir(parents=True, exist_ok=True) + temp_path = Path(tempfile.mkdtemp(prefix=f"{index_path.name}.", dir=index_path.parent)) + try: + if records: + retriever = bm25s.BM25() + texts = [str(record["bm25_text"]) for record in records] + retriever.index(bm25s.tokenize(texts, show_progress=False), show_progress=False) + retriever.save(temp_path, corpus=records, show_progress=False) + write_bm25_manifest(temp_path, manifest) + if index_path.exists(): + shutil.rmtree(index_path) + temp_path.rename(index_path) + except Exception: + shutil.rmtree(temp_path, ignore_errors=True) + raise + + +def read_bm25_manifest(index_path: Path) -> BM25Manifest | None: + """Read the BM25 manifest if it exists and is valid.""" + manifest_path = index_path / MANIFEST_NAME + if not manifest_path.exists(): + return None + body = json.loads(manifest_path.read_text(encoding="utf-8")) + return BM25Manifest( + created_at=datetime.fromisoformat(str(body["created_at"])), + db_updated_at=datetime.fromisoformat(str(body["db_updated_at"])) if body.get("db_updated_at") else None, + chunk_count=int(body["chunk_count"]), + ) + + +def write_bm25_manifest(index_path: Path, manifest: BM25Manifest) -> None: + """Write the BM25 manifest to an index directory.""" + body = { + "created_at": manifest.created_at.isoformat(), + "db_updated_at": manifest.db_updated_at.isoformat() if manifest.db_updated_at else None, + "chunk_count": manifest.chunk_count, + } + (index_path / MANIFEST_NAME).write_text(json.dumps(body, indent=2, sort_keys=True), encoding="utf-8") + + +def bm25_index_exists(index_path: Path, manifest: BM25Manifest | None) -> bool: + """Return whether a usable persisted BM25 index exists.""" + if manifest is None or not index_path.is_dir(): + return False + if manifest.chunk_count == 0: + return True + return all((index_path / file_name).exists() for file_name in REQUIRED_INDEX_FILES) -- 2.54.0 From 5308ff8be6333c41bec21c410a8ee48497f0e7b5 Mon Sep 17 00:00:00 2001 From: Richie Cahill Date: Fri, 12 Jun 2026 01:57:36 -0400 Subject: [PATCH 06/28] set up embedding system --- python/ebook_search/embeddings.py | 170 ++++++++++++++++++++++++++++++ 1 file changed, 170 insertions(+) create mode 100644 python/ebook_search/embeddings.py diff --git a/python/ebook_search/embeddings.py b/python/ebook_search/embeddings.py new file mode 100644 index 0000000..be428f3 --- /dev/null +++ b/python/ebook_search/embeddings.py @@ -0,0 +1,170 @@ +"""Embedding model helpers.""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass +from typing import TYPE_CHECKING + +from sqlalchemy import func, select +from sqlalchemy.dialects.postgresql import insert + +from python.ebook_search.api.embedding_client import request_embeddings +from python.orm.richie import ( + EbookChunk, + EbookChunkEmbedding1024, + EbookChunkEmbedding2560, + EbookChunkEmbedding4096, + EbookEmbeddingModel, +) + +logger = logging.getLogger(__name__) + +if TYPE_CHECKING: + from collections.abc import Sequence + + from sqlalchemy.orm import Session + + from python.ebook_search.config import EbookSearchConfig + +MODEL_DIMENSIONS = { + "qwen3-embedding-0.6b": 1024, + "qwen3-embedding-4b": 2560, + "qwen3-embedding-8b": 4096, +} + + +def get_embedding_table( + dimension: int, +) -> type[EbookChunkEmbedding1024 | EbookChunkEmbedding2560 | EbookChunkEmbedding4096]: + """Return the embedding table mapped to an embedding dimension.""" + embedding_tables = { + 1024: EbookChunkEmbedding1024, + 2560: EbookChunkEmbedding2560, + 4096: EbookChunkEmbedding4096, + } + table = embedding_tables.get(dimension) + if not table: + msg = f"Embedding dimension {dimension} is not supported" + raise ValueError(msg) + return table + + +@dataclass(frozen=True) +class EmbeddingModelStats: + """Embedding coverage for one model.""" + + model_name: str + dimension: int + embedded_chunks: int + total_chunks: int + + @property + def missing_chunks(self) -> int: + """Return chunks missing this embedding model.""" + return max(self.total_chunks - self.embedded_chunks, 0) + + +def embed_texts(texts: Sequence[str], config: EbookSearchConfig) -> list[list[float]]: + """Embed text with the configured vLLM embedding model.""" + logger.info( + "ebook_embed_request_start base_url=%s model=%s count=%s", + config.embedding_base_url, + config.embedding_model, + len(texts), + ) + vectors = request_embeddings(texts, config) + expected_dimension = MODEL_DIMENSIONS[config.embedding_model] + for vector in vectors: + if len(vector) != expected_dimension: + msg = f"Expected {expected_dimension} dimensions, got {len(vector)}" + raise ValueError(msg) + logger.info( + "ebook_embed_request_complete model=%s count=%s dimension=%s", + config.embedding_model, + len(vectors), + expected_dimension, + ) + return vectors + + +def embed_query(query: str, config: EbookSearchConfig) -> list[float]: + """Embed a search query with the Qwen retrieval instruction.""" + instructed_query = f"Instruct: Retrieve relevant passages for the query.\nQuery: {query}" + return embed_texts([instructed_query], config)[0] + + +def ensure_embedding_models(session: Session) -> None: + """Ensure supported embedding model rows exist.""" + for name, dimension in MODEL_DIMENSIONS.items(): + existing = session.scalar(select(EbookEmbeddingModel).where(EbookEmbeddingModel.name == name)) + if existing is None: + session.add(EbookEmbeddingModel(name=name, dimension=dimension, is_default=name == "qwen3-embedding-0.6b")) + logger.info("ebook_embedding_model_created model=%s dimension=%s", name, dimension) + session.flush() + + +def embedding_model_stats(session: Session) -> list[EmbeddingModelStats]: + """Return embedding coverage counts for every supported model.""" + total_chunks = session.scalar(select(func.count(EbookChunk.id))) or 0 + models = { + model.name: model + for model in session.scalars( + select(EbookEmbeddingModel) + .where(EbookEmbeddingModel.name.in_(MODEL_DIMENSIONS)) + .order_by(EbookEmbeddingModel.name) + ) + } + + stats: list[EmbeddingModelStats] = [] + for model_name, dimension in MODEL_DIMENSIONS.items(): + model = models.get(model_name) + embedded_chunks = 0 + if model is not None: + table = get_embedding_table(dimension) + embedded_chunks = session.scalar(select(func.count(table.id)).where(table.model_id == model.id)) or 0 + stats.append( + EmbeddingModelStats( + model_name=model_name, + dimension=dimension, + embedded_chunks=embedded_chunks, + total_chunks=total_chunks, + ) + ) + return stats + + +def embed_missing_chunks(session: Session, config: EbookSearchConfig) -> int: + """Embed chunks missing embeddings for the configured model.""" + ensure_embedding_models(session) + model = session.scalar(select(EbookEmbeddingModel).where(EbookEmbeddingModel.name == config.embedding_model)) + if model is None: + supported_models = ", ".join(MODEL_DIMENSIONS) + msg = f"Unknown embedding model: {config.embedding_model}. Supported models: {supported_models}" + raise ValueError(msg) + + table = get_embedding_table(model.dimension) + chunks = list( + session.scalars( + select(EbookChunk) + .outerjoin(table, (table.chunk_id == EbookChunk.id) & (table.model_id == model.id)) + .where(table.id.is_(None)) + .order_by(EbookChunk.id) + .limit(config.embedding_batch_size) + ) + ) + if not chunks: + logger.info("ebook_embed_missing_none model=%s", config.embedding_model) + return 0 + + logger.info("ebook_embed_missing_batch_start model=%s count=%s", config.embedding_model, len(chunks)) + vectors = embed_texts([chunk.text for chunk in chunks], config) + rows = [ + {"chunk_id": chunk.id, "model_id": model.id, "embedding": vector} + for chunk, vector in zip(chunks, vectors, strict=True) + ] + statement = insert(table).values(rows).on_conflict_do_nothing(index_elements=["chunk_id", "model_id"]) + session.execute(statement) + session.flush() + logger.info("ebook_embed_missing_batch_complete model=%s count=%s", config.embedding_model, len(rows)) + return len(rows) -- 2.54.0 From ff7b2ab2fa531ba64eab59e99aa1b6b82a63306c Mon Sep 17 00:00:00 2001 From: Richie Cahill Date: Fri, 12 Jun 2026 02:45:44 -0400 Subject: [PATCH 07/28] built rag search setup --- python/ebook_search/search.py | 371 ++++++++++++++++++++++++++++++++++ python/ebook_search/timing.py | 36 ++++ 2 files changed, 407 insertions(+) create mode 100644 python/ebook_search/search.py create mode 100644 python/ebook_search/timing.py diff --git a/python/ebook_search/search.py b/python/ebook_search/search.py new file mode 100644 index 0000000..e25c17d --- /dev/null +++ b/python/ebook_search/search.py @@ -0,0 +1,371 @@ +"""Hybrid search orchestration.""" + +from __future__ import annotations + +import logging +import re +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass, replace +from typing import TYPE_CHECKING + +from pgvector.sqlalchemy import Vector +from sqlalchemy import literal, select +from sqlalchemy.orm import Session + +from python.ebook_search.bm25_corpus import ( + load_bm25_corpus, + score_bm25_corpus, +) +from python.ebook_search.embeddings import MODEL_DIMENSIONS, embed_query, get_embedding_table +from python.ebook_search.rerank import rerank_chunks +from python.ebook_search.timing import RuntimeStep, timed_result +from python.orm.richie import ( + EbookChapter, + EbookChunk, + EbookEmbeddingModel, + EbookSource, +) + +if TYPE_CHECKING: + from collections.abc import Mapping + + from sqlalchemy.engine import Engine + + from python.ebook_search.config import EbookSearchConfig + +logger = logging.getLogger(__name__) +BM25_CANDIDATE_LIMIT = 120 + + +@dataclass(frozen=True) +class SearchResult: + """One source chunk returned by search.""" + + chunk_id: int + text: str + source_title: str + score: float = 0.0 + vector_score: float | None = None + bm25_score: float | None = None + fused_score: float | None = None + rerank_score: float | None = None + source_author: str | None = None + chapter_title: str | None = None + page_label: str | None = None + rank_source: str = "Hybrid" + + +@dataclass(frozen=True) +class SearchResponse: + """Search output for the UI.""" + + query: str + results: list[SearchResult] + rank_label: str + timings: tuple[RuntimeStep, ...] = () + + @property + def total_runtime_ms(self) -> float: + """Return total measured runtime for the response.""" + return sum(step.duration_ms for step in self.timings if step.counts_toward_total) + + +@dataclass(frozen=True) +class RetrievalResponse: + """Parallel retrieval output for vector and BM25 candidates.""" + + vector_results: list[SearchResult] + lexical_results: list[SearchResult] + timings: tuple[RuntimeStep, ...] + + +def search_ebooks( + engine: Engine, + query: str, + config: EbookSearchConfig, + *, + rerank: bool = False, +) -> SearchResponse: + """Run hybrid vector/BM25 search and optional reranking.""" + if not query.strip(): + logger.info("ebook_search_empty_query") + return SearchResponse(query=query, results=[], rank_label="Hybrid") + + logger.info("ebook_search_start query_length=%s rerank=%s", len(query), rerank) + timings: list[RuntimeStep] = [] + retrieval_query, timing = timed_result("Query preparation", retrieval_query_from_text, query) + timings.append(timing) + retrieval, timing = timed_result( + "Hybrid retrieval", + parallel_retrieval, + engine, + retrieval_query, + config, + ) + timings.extend(retrieval.timings) + timings.append(timing) + fused, timing = timed_result( + "Reciprocal rank fusion", + reciprocal_rank_fusion, + retrieval.vector_results, + retrieval.lexical_results, + ) + timings.append(timing) + if config.rerank.enabled and rerank: + response, timing = timed_result("Rerank", apply_rerank, query, fused, config) + else: + response, timing = timed_result("Rerank skipped", skip_rerank, query, fused, config) + timings.append(timing) + response = replace(response, timings=tuple(timings)) + logger.info( + "ebook_search_complete vector_candidates=%s lexical_candidates=%s " + "fused_candidates=%s returned=%s rank_label=%s runtime_ms=%.1f", + len(retrieval.vector_results), + len(retrieval.lexical_results), + len(fused), + len(response.results), + response.rank_label, + response.total_runtime_ms, + ) + return response + + +def parallel_retrieval(engine: Engine, query: str, config: EbookSearchConfig) -> RetrievalResponse: + """Run vector and BM25 candidate retrieval concurrently with separate database sessions.""" + with ThreadPoolExecutor(max_workers=2, thread_name_prefix="ebook-search") as executor: + vector_future = executor.submit( + timed_result, + "Embedding + vector search", + vector_candidates, + engine, + query, + config, + ) + bm25_future = executor.submit( + timed_result, + "BM25 search", + bm25_candidates, + query, + config, + ) + vector_results, vector_timing = vector_future.result() + lexical_results, lexical_timing = bm25_future.result() + + logger.info( + "ebook_parallel_retrieval_complete vector_candidates=%s lexical_candidates=%s", + len(vector_results), + len(lexical_results), + ) + return RetrievalResponse( + vector_results=vector_results, + lexical_results=lexical_results, + timings=( + replace(vector_timing, counts_toward_total=False), + replace(lexical_timing, counts_toward_total=False), + ), + ) + + +def skip_rerank( + query: str, + candidates: list[SearchResult], + config: EbookSearchConfig, +) -> SearchResponse: + """Return fused hybrid results without reranking.""" + logger.info("ebook_rerank_skipped candidates=%s", len(candidates)) + return SearchResponse(query=query, results=candidates[: config.top_k], rank_label="Hybrid") + + +def apply_rerank( + query: str, + candidates: list[SearchResult], + config: EbookSearchConfig, +) -> SearchResponse: + """Rerank already-fused hybrid candidates.""" + reranked = rerank_chunks(query, candidates[: config.rerank.candidates], config.rerank) + logger.info( + "ebook_rerank_complete input_candidates=%s returned=%s", + min(len(candidates), config.rerank.candidates), + len(reranked), + ) + return SearchResponse( + query=query, + results=[replace(result, rank_source="Hybrid + rerank") for result in reranked[: config.top_k]], + rank_label="Hybrid + rerank", + ) + + +def vector_candidates(engine: Engine, query: str, config: EbookSearchConfig) -> list[SearchResult]: + """Return pgvector cosine candidates for a normalized query.""" + with Session(engine) as session: + model = session.scalar(select(EbookEmbeddingModel).where(EbookEmbeddingModel.name == config.embedding_model)) + if model is None: + msg = f"Embedding model is not registered: {config.embedding_model}" + raise ValueError(msg) + + expected_dimension = MODEL_DIMENSIONS[config.embedding_model] + if model.dimension != expected_dimension: + msg = f"Model row dimension {model.dimension} does not match configured dimension {expected_dimension}" + raise ValueError(msg) + + embedding = embed_query(query, config) + limit = max(config.rerank.candidates, config.top_k) * 4 + embedding_table = get_embedding_table(model.dimension) + + embedding_param = literal(embedding, type_=Vector(model.dimension)) + distance = embedding_table.embedding.op("<=>")(embedding_param) + score = (literal(1.0) - distance).label("score") + statement = ( + select( + EbookChunk.id.label("chunk_id"), + EbookChunk.text.label("text"), + EbookSource.title.label("source_title"), + EbookSource.author.label("source_author"), + EbookChapter.title.label("chapter_title"), + EbookChunk.page_label.label("page_label"), + score, + ) + .select_from(embedding_table) + .join(EbookChunk, EbookChunk.id == embedding_table.chunk_id) + .join(EbookSource, EbookSource.id == EbookChunk.source_id) + .outerjoin(EbookChapter, EbookChapter.id == EbookChunk.chapter_id) + .where(embedding_table.model_id == model.id) + .order_by(distance) + .limit(limit) + ) + rows = session.execute(statement).mappings() + results = [search_result_from_row(row) for row in rows] + logger.info( + "ebook_vector_search_complete model=%s dimension=%s candidates=%s", + config.embedding_model, + model.dimension, + len(results), + ) + return results + + +def bm25_candidates(query: str, config: EbookSearchConfig) -> list[SearchResult]: + """Return BM25-ranked lexical candidates using the persisted corpus.""" + corpus = load_bm25_corpus(config) + if not corpus.records: + logger.info("ebook_bm25_search_complete corpus=0 candidates=0") + return [] + + scored_records = score_bm25_corpus(query, corpus, limit=BM25_CANDIDATE_LIMIT) + results = [ + replace(search_result_from_row(record), score=score, vector_score=None, bm25_score=score) + for record, score in scored_records + ] + + max_score = results[0].bm25_score if results else 0.0 + logger.info( + "ebook_bm25_search_complete corpus=%s candidates=%s max_score=%.6f", + len(corpus.records), + len(results), + max_score, + ) + return results + + +def reciprocal_rank_fusion( + vector_results: list[SearchResult], + lexical_results: list[SearchResult], + *, + rank_constant: int = 60, +) -> list[SearchResult]: + """Fuse vector and lexical rankings with Reciprocal Rank Fusion.""" + by_chunk: dict[int, SearchResult] = {} + scores: dict[int, float] = {} + vector_scores: dict[int, float] = {} + bm25_scores: dict[int, float] = {} + + for rank, result in enumerate(vector_results, start=1): + by_chunk.setdefault(result.chunk_id, result) + vector_scores[result.chunk_id] = result.vector_score if result.vector_score is not None else result.score + scores[result.chunk_id] = scores.get(result.chunk_id, 0.0) + (1 / (rank_constant + rank)) + + for rank, result in enumerate(lexical_results, start=1): + by_chunk.setdefault(result.chunk_id, result) + bm25_scores[result.chunk_id] = result.bm25_score if result.bm25_score is not None else result.score + scores[result.chunk_id] = scores.get(result.chunk_id, 0.0) + (1 / (rank_constant + rank)) + + return sorted( + ( + replace( + result, + score=scores[result.chunk_id], + vector_score=vector_scores.get(result.chunk_id), + bm25_score=bm25_scores.get(result.chunk_id), + fused_score=scores[result.chunk_id], + rank_source="Hybrid", + ) + for result in by_chunk.values() + ), + key=lambda result: result.score, + reverse=True, + ) + + +def search_result_from_row(row: Mapping[str, object]) -> SearchResult: + """Convert a database row mapping into a search result.""" + return SearchResult( + chunk_id=int(row["chunk_id"]), + text=str(row["text"]), + source_title=str(row["source_title"]), + source_author=optional_str(row["source_author"]), + chapter_title=optional_str(row["chapter_title"]), + page_label=optional_str(row["page_label"]), + score=float(row["score"]) if "score" in row else 0.0, + vector_score=float(row["score"]) if "score" in row else None, + ) + + +def optional_str(value: object) -> str | None: + """Convert nullable database values to optional strings.""" + if value is None: + return None + return str(value) + + +TOKEN_RE = re.compile(r"[A-Za-z0-9_]+") + + +def tokens(text_value: str) -> list[str]: + """Extract tokens from a text value. + + This is a simple approximation of the tokenization used by PostgreSQL's full-text search, + which is sufficient for BM25 candidate retrieval. It lowercases tokens and includes alphanumeric characters and + underscores. + """ + return [match.group(0).lower() for match in TOKEN_RE.finditer(text_value)] + + +QUERY_STOP_WORDS = { + "a", + "an", + "and", + "are", + "as", + "at", + "does", + "for", + "in", + "is", + "of", + "the", + "to", + "what", + "when", + "where", + "which", + "who", + "why", +} + + +def retrieval_query_from_text(query: str) -> str: + """Remove generic question words while preserving entity and series terms.""" + keywords = [token for token in tokens(query) if token not in QUERY_STOP_WORDS] + if not keywords: + return query + return " ".join(keywords) diff --git a/python/ebook_search/timing.py b/python/ebook_search/timing.py new file mode 100644 index 0000000..eb8e474 --- /dev/null +++ b/python/ebook_search/timing.py @@ -0,0 +1,36 @@ +"""Runtime timing helpers for EPUB search.""" + +from __future__ import annotations + +from dataclasses import dataclass +from time import perf_counter +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from collections.abc import Callable + + +@dataclass(frozen=True) +class RuntimeStep: + """Elapsed runtime for one named search step.""" + + name: str + duration_ms: float + counts_toward_total: bool = True + + +def runtime_step_from_start(name: str, start_seconds: float) -> RuntimeStep: + """Create a runtime step from a prior perf_counter timestamp.""" + return RuntimeStep(name=name, duration_ms=(perf_counter() - start_seconds) * 1000) + + +def timed_result[T, **P]( + name: str, + operation: Callable[P, T], + *args: P.args, + **kwargs: P.kwargs, +) -> tuple[T, RuntimeStep]: + """Run an operation and return its result plus elapsed runtime.""" + start_seconds = perf_counter() + result = operation(*args, **kwargs) + return result, runtime_step_from_start(name, start_seconds) -- 2.54.0 From d1b59955d0e92379fcb8421c348c9bb43288132c Mon Sep 17 00:00:00 2001 From: Richie Cahill Date: Fri, 12 Jun 2026 03:06:16 -0400 Subject: [PATCH 08/28] built ingest --- python/ebook_search/epub_parse.py | 95 +++++++++++++++ python/ebook_search/ingest.py | 190 ++++++++++++++++++++++++++++++ 2 files changed, 285 insertions(+) create mode 100644 python/ebook_search/epub_parse.py create mode 100644 python/ebook_search/ingest.py diff --git a/python/ebook_search/epub_parse.py b/python/ebook_search/epub_parse.py new file mode 100644 index 0000000..919a096 --- /dev/null +++ b/python/ebook_search/epub_parse.py @@ -0,0 +1,95 @@ +"""EPUB parsing helpers.""" + +from __future__ import annotations + +import re +from dataclasses import dataclass +from typing import TYPE_CHECKING + +from bs4 import BeautifulSoup +from ebooklib import ITEM_DOCUMENT, epub + +if TYPE_CHECKING: + from pathlib import Path + +WHITESPACE_RE = re.compile(r"\s+") + + +@dataclass(frozen=True) +class ParsedChapter: + """Text extracted from one EPUB spine document.""" + + title: str | None + href: str | None + text: str + page_labels: tuple[str, ...] + + +@dataclass(frozen=True) +class ParsedEpub: + """Parsed EPUB metadata and text.""" + + title: str + author: str | None + language: str | None + publisher: str | None + identifier: str | None + chapters: tuple[ParsedChapter, ...] + + +def parse_epub(path: Path) -> ParsedEpub: + """Parse EPUB metadata and spine text.""" + book = epub.read_epub(path) + chapters = [] + for item in book.get_items_of_type(ITEM_DOCUMENT): + soup = BeautifulSoup(item.get_content(), "html.parser") + title = chapter_title(soup) + page_labels = tuple(extract_page_labels(soup)) + text = clean_text(soup.get_text(" ")) + if text: + chapters.append(ParsedChapter(title=title, href=item.get_name(), text=text, page_labels=page_labels)) + + return ParsedEpub( + title=metadata_value(book, "title") or path.stem, + author=metadata_value(book, "creator"), + language=metadata_value(book, "language"), + publisher=metadata_value(book, "publisher"), + identifier=metadata_value(book, "identifier"), + chapters=tuple(chapters), + ) + + +def metadata_value(book: epub.EpubBook, name: str) -> str | None: + """Return the first non-empty Dublin Core metadata value for a name.""" + values = book.get_metadata("DC", name) + if not values: + return None + value = values[0][0] + return str(value).strip() or None + + +def chapter_title(soup: BeautifulSoup) -> str | None: + """Extract the best available title from an EPUB document soup.""" + heading = soup.find(["h1", "h2", "h3"]) + if heading is None: + title = soup.find("title") + if title is None: + return None + return clean_text(title.get_text(" ")) or None + return clean_text(heading.get_text(" ")) or None + + +def extract_page_labels(soup: BeautifulSoup) -> list[str]: + """Extract EPUB page-break labels from a document soup.""" + labels: list[str] = [] + for tag in soup.find_all(attrs={"epub:type": "pagebreak"}): + label = tag.get("title") or tag.get("aria-label") or tag.get_text(" ") + clean = clean_text(str(label)) + if clean: + labels.append(clean) + return labels + + +def clean_text(text: str) -> str: + """Normalize whitespace in extracted EPUB text.""" + return WHITESPACE_RE.sub(" ", text).strip() diff --git a/python/ebook_search/ingest.py b/python/ebook_search/ingest.py new file mode 100644 index 0000000..2b8e44a --- /dev/null +++ b/python/ebook_search/ingest.py @@ -0,0 +1,190 @@ +"""EPUB ingestion into Richie DB.""" + +from __future__ import annotations + +import hashlib +import logging +from dataclasses import dataclass +from datetime import UTC, datetime +from pathlib import Path +from typing import TYPE_CHECKING + +import tiktoken +from sqlalchemy import or_, select + +from python.ebook_search.epub_parse import parse_epub +from python.orm.richie import EbookChapter, EbookChunk, EbookSource + +logger = logging.getLogger(__name__) +DEFAULT_CHUNK_TOKENS = 700 +DEFAULT_CHUNK_OVERLAP = 100 + +if TYPE_CHECKING: + from sqlalchemy.orm import Session + + from python.ebook_search.config import EbookSearchConfig + from python.ebook_search.epub_parse import ParsedChapter + + +@dataclass(frozen=True) +class TextChunk: + """A token-bounded chunk of text.""" + + text: str + token_start: int + token_count: int + + +def chunk_text( + text: str, + *, + chunk_tokens: int = DEFAULT_CHUNK_TOKENS, + overlap_tokens: int = DEFAULT_CHUNK_OVERLAP, +) -> list[TextChunk]: + """Split text into overlapping token chunks.""" + if chunk_tokens <= 0: + msg = "chunk_tokens must be positive" + raise ValueError(msg) + if overlap_tokens < 0 or overlap_tokens >= chunk_tokens: + msg = "overlap_tokens must be non-negative and smaller than chunk_tokens" + raise ValueError(msg) + + encoding = tiktoken.get_encoding("cl100k_base") + tokens = encoding.encode(text) + if not tokens: + return [] + + chunks: list[TextChunk] = [] + step = chunk_tokens - overlap_tokens + for start in range(0, len(tokens), step): + chunk = tokens[start : start + chunk_tokens] + if not chunk: + continue + chunks.append( + TextChunk( + text=encoding.decode(chunk).strip(), + token_start=start, + token_count=len(chunk), + ) + ) + if start + chunk_tokens >= len(tokens): + break + return [chunk for chunk in chunks if chunk.text] + + +def ingest_configured_paths(session: Session, config: EbookSearchConfig) -> int: + """Ingest every EPUB found under configured library paths.""" + count = 0 + for library_path in config.library_paths: + path = Path(library_path).expanduser() + logger.info("ebook_ingest_path_start path=%s", path) + if path.is_file() and path.suffix.lower() == ".epub": + count += int(ingest_file(session, path)) + elif path.is_dir(): + for epub_path in sorted(path.rglob("*.epub")): + count += int(ingest_file(session, epub_path)) + else: + logger.warning("ebook_ingest_path_missing path=%s", path) + logger.info("ebook_ingest_paths_complete changed_files=%s configured_paths=%s", count, len(config.library_paths)) + return count + + +def ingest_file(session: Session, path: Path) -> bool: + """Ingest one EPUB file. Return True when the database changed.""" + resolved_path = path.expanduser().resolve() + logger.info("ebook_ingest_file_start path=%s", resolved_path) + file_hash = sha256_file(resolved_path) + existing = find_existing_source(session, resolved_path, file_hash) + if existing is not None and existing.file_sha256 == file_hash: + stat = resolved_path.stat() + existing.file_path = str(resolved_path) + existing.file_mtime = datetime.fromtimestamp(stat.st_mtime, tz=UTC) + existing.file_size = stat.st_size + session.flush() + logger.info("ebook_ingest_file_unchanged source_id=%s path=%s", existing.id, resolved_path) + return False + if existing is not None: + logger.info("ebook_ingest_file_replacing source_id=%s path=%s", existing.id, resolved_path) + session.delete(existing) + session.flush() + + stat = resolved_path.stat() + parsed = parse_epub(resolved_path) + source = EbookSource( + title=parsed.title, + author=parsed.author, + language=parsed.language, + publisher=parsed.publisher, + identifier=parsed.identifier, + file_path=str(resolved_path), + file_sha256=file_hash, + file_mtime=datetime.fromtimestamp(stat.st_mtime, tz=UTC), + file_size=stat.st_size, + ) + session.add(source) + session.flush() + + chunk_index = 0 + for spine_index, parsed_chapter in enumerate(parsed.chapters): + chapter = EbookChapter( + source_id=source.id, + spine_index=spine_index, + title=parsed_chapter.title, + href=parsed_chapter.href, + ) + session.add(chapter) + session.flush() + chunk_index = add_chapter_chunks(session, source, chapter, parsed_chapter, chunk_index) + + session.flush() + logger.info( + "ebook_ingest_file_complete source_id=%s path=%s chapters=%s chunks=%s", + source.id, + resolved_path, + len(parsed.chapters), + chunk_index, + ) + return True + + +def find_existing_source(session: Session, path: Path, file_hash: str) -> EbookSource | None: + """Find an existing source by canonical path or file hash.""" + return session.scalar( + select(EbookSource).where(or_(EbookSource.file_path == str(path), EbookSource.file_sha256 == file_hash)) + ) + + +def add_chapter_chunks( + session: Session, + source: EbookSource, + chapter: EbookChapter, + parsed_chapter: ParsedChapter, + chunk_index: int, +) -> int: + """Add chunk rows for one parsed chapter and return the next chunk index.""" + page_label = parsed_chapter.page_labels[0] if parsed_chapter.page_labels else None + for text_chunk in chunk_text(parsed_chapter.text): + session.add( + EbookChunk( + source_id=source.id, + chapter_id=chapter.id, + chunk_index=chunk_index, + text=text_chunk.text, + token_start=text_chunk.token_start, + token_count=text_chunk.token_count, + page_label=page_label, + content_sha256=hashlib.sha256(text_chunk.text.encode()).hexdigest(), + search_text=f"{source.title} {source.author or ''} {chapter.title or ''} {text_chunk.text}", + ) + ) + chunk_index += 1 + return chunk_index + + +def sha256_file(path: Path) -> str: + """Calculate the SHA-256 digest for a file.""" + digest = hashlib.sha256() + with path.open("rb") as file: + for block in iter(lambda: file.read(1024 * 1024), b""): + digest.update(block) + return digest.hexdigest() -- 2.54.0 From f7b72c4053fe950c72795c4f4434cb7a662195e1 Mon Sep 17 00:00:00 2001 From: Richie Cahill Date: Fri, 12 Jun 2026 03:06:27 -0400 Subject: [PATCH 09/28] added rerank --- python/ebook_search/rerank.py | 127 ++++++++++++++++++++++++++++++++++ 1 file changed, 127 insertions(+) create mode 100644 python/ebook_search/rerank.py diff --git a/python/ebook_search/rerank.py b/python/ebook_search/rerank.py new file mode 100644 index 0000000..abb2ee0 --- /dev/null +++ b/python/ebook_search/rerank.py @@ -0,0 +1,127 @@ +"""vLLM-backed optional reranking.""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass, replace +from typing import TYPE_CHECKING + +from python.ebook_search.llm_interface import request_rerank + +if TYPE_CHECKING: + from python.ebook_search.config import RerankConfig + from python.ebook_search.search import SearchResult + +logger = logging.getLogger(__name__) + + +@dataclass(frozen=True) +class RerankResult: + """A relevance score for one candidate chunk.""" + + chunk_id: int + score: float + + +def rerank_chunks(query: str, candidates: list[SearchResult], config: RerankConfig) -> list[SearchResult]: + """Rerank candidates with a vLLM rerank endpoint.""" + if not candidates: + return [] + + logger.info( + "ebook_rerank_request_start base_url=%s model=%s candidates=%s", + config.base_url, + config.model, + len(candidates), + ) + scores = score_candidates(query, candidates, config) + results = sorted( + ( + replace( + result, + score=final_rerank_score(result, scores[result.chunk_id].score, candidates), + rerank_score=scores[result.chunk_id].score, + ) + for result in candidates + ), + key=lambda result: result.score, + reverse=True, + ) + logger.info( + "ebook_rerank_request_complete base_url=%s model=%s candidates=%s", + config.base_url, + config.model, + len(results), + ) + return results + + +def score_candidates( + query: str, + candidates: list[SearchResult], + config: RerankConfig, +) -> dict[int, RerankResult]: + """Score candidate chunks with the configured rerank API.""" + body = request_rerank(query, [candidate.text for candidate in candidates], config) + if body is None: + return zero_rerank_scores(candidates) + + scores = parse_vllm_scores(body, candidates) + for result in scores.values(): + logger.debug("ebook_rerank_candidate_scored chunk_id=%s score=%s", result.chunk_id, result.score) + return scores + + +def parse_vllm_scores(body: object, candidates: list[SearchResult]) -> dict[int, RerankResult]: + """Parse vLLM rerank scores into chunk-id keyed results.""" + if not isinstance(body, dict): + logger.debug("ebook_rerank_response_not_object", extra={"response": body}) + return zero_rerank_scores(candidates) + + results = body.get("results") or body.get("data") + if not isinstance(results, list): + logger.debug("ebook_rerank_response_missing_results", extra={"response": body}) + return zero_rerank_scores(candidates) + + scores = zero_rerank_scores(candidates) + for item in results: + if not isinstance(item, dict): + continue + index = item.get("index") + score = item.get("relevance_score", item.get("score")) + if not isinstance(index, int) or index < 0 or index >= len(candidates): + continue + if not isinstance(score, int | float): + continue + chunk_id = candidates[index].chunk_id + scores[chunk_id] = RerankResult(chunk_id=chunk_id, score=clamp_score(float(score))) + return scores + + +def zero_rerank_scores(candidates: list[SearchResult]) -> dict[int, RerankResult]: + """Return zero relevance scores for all candidate chunks.""" + return {candidate.chunk_id: RerankResult(chunk_id=candidate.chunk_id, score=0.0) for candidate in candidates} + + +def clamp_score(score: float) -> float: + """Clamp a rerank score into the supported 0.0 to 1.0 range.""" + return min(max(score, 0.0), 1.0) + + +def final_rerank_score(result: SearchResult, rerank_score: float, candidates: list[SearchResult]) -> float: + """Combine rerank relevance with normalized hybrid retrieval evidence.""" + return rerank_score * normalized_hybrid_score(result, candidates) + + +def normalized_hybrid_score(result: SearchResult, candidates: list[SearchResult]) -> float: + """Normalize a candidate hybrid score against the rerank candidate set.""" + hybrid_scores = [ + candidate.fused_score if candidate.fused_score is not None else candidate.score for candidate in candidates + ] + low = min(hybrid_scores) + high = max(hybrid_scores) + if high == low: + return 1.0 + + score = result.fused_score if result.fused_score is not None else result.score + return (score - low) / (high - low) -- 2.54.0 From 4bd61bc170538f239ede554362c2d6d53cb4d643 Mon Sep 17 00:00:00 2001 From: Richie Cahill Date: Fri, 12 Jun 2026 03:08:21 -0400 Subject: [PATCH 10/28] made llm_interface.py --- python/ebook_search/embeddings.py | 2 +- python/ebook_search/llm_interface.py | 143 +++++++++++++++++++++++++++ 2 files changed, 144 insertions(+), 1 deletion(-) create mode 100644 python/ebook_search/llm_interface.py diff --git a/python/ebook_search/embeddings.py b/python/ebook_search/embeddings.py index be428f3..f542e2b 100644 --- a/python/ebook_search/embeddings.py +++ b/python/ebook_search/embeddings.py @@ -9,7 +9,7 @@ from typing import TYPE_CHECKING from sqlalchemy import func, select from sqlalchemy.dialects.postgresql import insert -from python.ebook_search.api.embedding_client import request_embeddings +from python.ebook_search.llm_interface import request_embeddings from python.orm.richie import ( EbookChunk, EbookChunkEmbedding1024, diff --git a/python/ebook_search/llm_interface.py b/python/ebook_search/llm_interface.py new file mode 100644 index 0000000..8cfa121 --- /dev/null +++ b/python/ebook_search/llm_interface.py @@ -0,0 +1,143 @@ +"""LLM provider HTTP adapters.""" + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING + +import httpx + +if TYPE_CHECKING: + from collections.abc import Sequence + + from python.ebook_search.config import EbookSearchConfig, RerankConfig + +logger = logging.getLogger(__name__) + + +def auth_headers(api_key: str) -> dict[str, str]: + """Build authorization headers when an API key is configured.""" + if api_key == "not-needed": + return {} + return {"Authorization": f"Bearer {api_key}"} + + +def request_embeddings(texts: Sequence[str], config: EbookSearchConfig) -> list[list[float]]: + """Request embeddings from the configured OpenAI-compatible endpoint.""" + try: + response = httpx.post( + f"{config.embedding_base_url.rstrip('/')}/embeddings", + headers=auth_headers(config.embedding_api_key), + json={"model": config.embedding_model, "input": list(texts)}, + timeout=60, + ) + response.raise_for_status() + return embedding_vectors_from_response(response.json()) + except (httpx.HTTPError, ValueError, KeyError, TypeError) as error: + logger.exception( + "ebook_embed_request_failed base_url=%s model=%s count=%s", + config.embedding_base_url, + config.embedding_model, + len(texts), + ) + msg = f"Embedding request failed. base_url={config.embedding_base_url} model={config.embedding_model}" + raise RuntimeError(msg) from error + + +def embedding_vectors_from_response(body: object) -> list[list[float]]: + """Extract embedding vectors from an OpenAI-compatible embedding response.""" + if not isinstance(body, dict): + msg = "Embedding response is not an object" + raise TypeError(msg) + + data = body["data"] + if not isinstance(data, list): + msg = "Embedding response data is not a list" + raise TypeError(msg) + + vectors: list[list[float]] = [] + for item in data: + if not isinstance(item, dict): + msg = "Embedding item is not an object" + raise TypeError(msg) + embedding = item["embedding"] + if not isinstance(embedding, list): + msg = "Embedding value is not a list" + raise TypeError(msg) + vectors.append([float(value) for value in embedding]) + return vectors + + +def request_rerank( + query: str, + documents: Sequence[str], + config: RerankConfig, +) -> object | None: + """Request rerank scores from the configured vLLM endpoint.""" + payload = { + "model": config.model, + "query": query, + "documents": list(documents), + } + response = httpx.post( + f"{config.base_url.rstrip('/')}/rerank", + json=payload, + timeout=config.timeout_seconds, + ) + response.raise_for_status() + try: + return response.json() + except ValueError: + logger.debug("ebook_rerank_response_invalid_json", extra={"response": response.text}) + return None + + +def request_chat_completion( + config: EbookSearchConfig, + messages: Sequence[dict[str, str]], +) -> str: + """Request a chat completion from the configured OpenAI-compatible endpoint.""" + try: + response = httpx.post( + f"{config.vllm_base_url.rstrip('/')}/chat/completions", + headers=auth_headers(config.vllm_api_key), + json={ + "model": config.chat_model, + "messages": list(messages), + "temperature": 0, + }, + timeout=60, + ) + response.raise_for_status() + return chat_content_from_response(response.json()) + except (httpx.HTTPError, ValueError, KeyError, TypeError) as error: + msg = f"Chat request failed. base_url={config.vllm_base_url} model={config.chat_model}" + raise RuntimeError(msg) from error + + +def chat_content_from_response(body: object) -> str: + """Extract text content from an OpenAI-compatible chat response.""" + if not isinstance(body, dict): + msg = "Chat response is not an object" + raise TypeError(msg) + + choices = body["choices"] + if not isinstance(choices, list) or not choices: + msg = "Chat response has no choices" + raise ValueError(msg) + + first = choices[0] + if not isinstance(first, dict): + msg = "Chat choice is not an object" + raise TypeError(msg) + + message = first["message"] + if not isinstance(message, dict): + msg = "Chat message is not an object" + raise TypeError(msg) + + content = message.get("content") or "" + if not isinstance(content, str): + msg = "Chat content is not text" + raise TypeError(msg) + return content -- 2.54.0 From 3a5b278c158a246c21d2ce62aee6853228067f6e Mon Sep 17 00:00:00 2001 From: Richie Cahill Date: Fri, 12 Jun 2026 03:09:10 -0400 Subject: [PATCH 11/28] added __init__ --- python/ebook_search/__init__.py | 1 + 1 file changed, 1 insertion(+) create mode 100644 python/ebook_search/__init__.py diff --git a/python/ebook_search/__init__.py b/python/ebook_search/__init__.py new file mode 100644 index 0000000..b8cbc01 --- /dev/null +++ b/python/ebook_search/__init__.py @@ -0,0 +1 @@ +"""EPUB search package.""" -- 2.54.0 From 93ff2200fe73d7237e1791084c2f466dd006a31b Mon Sep 17 00:00:00 2001 From: Richie Cahill Date: Fri, 12 Jun 2026 03:09:51 -0400 Subject: [PATCH 12/28] added answer.py and config --- python/ebook_search/answer.py | 57 +++++++++++++++++ python/ebook_search/config.py | 117 ++++++++++++++++++++++++++++++++++ 2 files changed, 174 insertions(+) create mode 100644 python/ebook_search/answer.py create mode 100644 python/ebook_search/config.py diff --git a/python/ebook_search/answer.py b/python/ebook_search/answer.py new file mode 100644 index 0000000..4b85b21 --- /dev/null +++ b/python/ebook_search/answer.py @@ -0,0 +1,57 @@ +"""Grounded answer generation.""" + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING + +from python.ebook_search.llm_interface import request_chat_completion + +if TYPE_CHECKING: + from python.ebook_search.config import EbookSearchConfig + from python.ebook_search.search import SearchResult + +logger = logging.getLogger(__name__) + + +def answer_query(query: str, results: list[SearchResult], config: EbookSearchConfig) -> str: + """Answer a question using only retrieved chunks.""" + if not config.answer_enabled: + logger.info("ebook_answer_skipped_disabled") + return "Answer generation is disabled. Source chunks are shown below." + + if not results: + logger.info("ebook_answer_skipped_no_results") + return "No relevant sources were found." + + logger.info( + "ebook_answer_request_start base_url=%s model=%s sources=%s query_length=%s", + config.vllm_base_url, + config.chat_model, + len(results), + len(query), + ) + context = "\n\n".join( + f"[{index}] {result.source_title}{' - ' + result.chapter_title if result.chapter_title else ''}\n{result.text}" + for index, result in enumerate(results, start=1) + ) + content = request_chat_completion( + config, + [ + { + "role": "system", + "content": ( + "Answer only from the provided context. Cite sources with bracketed numbers like [1]. " + "If the context is insufficient, say so." + ), + }, + {"role": "user", "content": f"Question:\n{query}\n\nContext:\n{context}"}, + ], + ) + + logger.info( + "ebook_answer_request_complete model=%s answer_length=%s", + config.chat_model, + len(content), + ) + return content or "The model returned an empty answer." diff --git a/python/ebook_search/config.py b/python/ebook_search/config.py new file mode 100644 index 0000000..8d21274 --- /dev/null +++ b/python/ebook_search/config.py @@ -0,0 +1,117 @@ +"""Configuration for the EPUB search app.""" + +from __future__ import annotations + +from dataclasses import dataclass +from os import getenv + + +def getenv_bool(name: str, *, default: bool) -> bool: + """Read a boolean environment variable with a default fallback.""" + value = getenv(name) + if value is None: + return default + return value.strip().lower() in {"1", "true", "yes", "on"} + + +def getenv_int(name: str, *, default: int) -> int: + """Read an integer environment variable with a default fallback.""" + value = getenv(name) + if value is None or not value.strip(): + return default + return int(value) + + +@dataclass(frozen=True) +class RerankConfig: + """vLLM reranker settings.""" + + enabled: bool = False + base_url: str = "http://192.168.90.25:8001" + model: str = "qwen3-reranker-06b" + candidates: int = 24 + timeout_seconds: float = 30.0 + + +@dataclass(frozen=True) +class EbookSearchConfig: + """Runtime settings for EPUB search.""" + + rerank: RerankConfig + top_k: int = 12 + library_paths: tuple[str, ...] = () + vllm_base_url: str = "https://ollama.com/v1" + vllm_api_key: str = "not-needed" + chat_model: str = "deepseek-v4-flash" + answer_enabled: bool = True + embedding_base_url: str = "http://192.168.90.25:8000/v1" + embedding_api_key: str = "not-needed" + embedding_model: str = "qwen3-embedding-0.6b" + embedding_batch_size: int = 32 + bm25_index_dir: str = ".ebook_search_bm25" + bm25_refresh_delay_seconds: int = 60 + + +def load_rerank_config() -> RerankConfig: + """Load reranker config from environment variables.""" + return RerankConfig( + enabled=getenv_bool("EBOOK_SEARCH_RERANK_ENABLED", default=False), + base_url=getenv("EBOOK_SEARCH_RERANK_BASE_URL", "http://192.168.90.25:8001"), + model=getenv("EBOOK_SEARCH_RERANK_MODEL", "qwen3-reranker-06b"), + candidates=getenv_int("EBOOK_SEARCH_RERANK_CANDIDATES", default=24), + timeout_seconds=float(getenv_int("EBOOK_SEARCH_RERANK_TIMEOUT_SECONDS", default=30)), + ) + + +def load_config() -> EbookSearchConfig: + """Load EPUB search config from environment variables.""" + return EbookSearchConfig( + rerank=load_rerank_config(), + top_k=getenv_int("EBOOK_SEARCH_TOP_K", default=12), + library_paths=library_paths_from_env(), + vllm_base_url=getenv("EBOOK_SEARCH_VLLM_BASE_URL", "https://ollama.com/v1"), + vllm_api_key=getenv("EBOOK_SEARCH_VLLM_API_KEY") or getenv("OLLAMA_API_KEY") or "not-needed", + chat_model=getenv("EBOOK_SEARCH_CHAT_MODEL", "deepseek-v4-flash"), + answer_enabled=getenv_bool("EBOOK_SEARCH_ANSWER_ENABLED", default=True), + embedding_base_url=getenv("EBOOK_SEARCH_EMBEDDING_BASE_URL", "http://192.168.90.25:8000/v1"), + embedding_api_key=getenv("EBOOK_SEARCH_EMBEDDING_API_KEY", "not-needed"), + embedding_model=normalize_embedding_model(), + embedding_batch_size=getenv_int("EBOOK_SEARCH_EMBEDDING_BATCH_SIZE", default=32), + bm25_index_dir=getenv("EBOOK_SEARCH_BM25_INDEX_DIR", ".ebook_search_bm25"), + bm25_refresh_delay_seconds=getenv_int("EBOOK_SEARCH_BM25_REFRESH_DELAY_SECONDS", default=60), + ) + + +def normalize_embedding_model(default: str = "qwen3-embedding-0.6b") -> str: + """Normalize supported embedding aliases to provider model names.""" + aliases = { + "Qwen3-Embedding-0.6B": "qwen3-embedding-0.6b", + "Qwen3-Embedding-4B": "qwen3-embedding-4b", + "Qwen3-Embedding-8B": "qwen3-embedding-8b", + "Qwen/Qwen3-Embedding-0.6B": "qwen3-embedding-0.6b", + "Qwen/Qwen3-Embedding-4B": "qwen3-embedding-4b", + "Qwen/Qwen3-Embedding-8B": "qwen3-embedding-8b", + "qwen3-embedding:0.6b": "qwen3-embedding-0.6b", + "qwen3-embedding:4b": "qwen3-embedding-4b", + "qwen3-embedding:8b": "qwen3-embedding-8b", + "qwen3-embedding-0.6b": "qwen3-embedding-0.6b", + "qwen3-embedding-4b": "qwen3-embedding-4b", + "qwen3-embedding-8b": "qwen3-embedding-8b", + } + + model = getenv("EBOOK_SEARCH_EMBEDDING_MODEL", default) + standard_model = aliases.get(model) + + if standard_model is None: + error = f"Embedding model {model} is not supported. Supported models are {aliases.keys()}" + raise ValueError(error) + + return standard_model + + +def library_paths_from_env() -> tuple[str, ...]: + """Read configured EPUB library paths from the environment.""" + value = getenv("EBOOK_SEARCH_LIBRARY_PATHS") + if value is None: + return () + return tuple(path for path in value.split(":") if path) -- 2.54.0 From 07dd1922b1f8735c7d7de6f90fea0fe9d8dfd76e Mon Sep 17 00:00:00 2001 From: Richie Cahill Date: Fri, 12 Jun 2026 03:10:19 -0400 Subject: [PATCH 13/28] build api and frountend --- python/ebook_search/api/__init__.py | 1 + python/ebook_search/api/bm25_tasks.py | 58 ++++++++ python/ebook_search/api/main.py | 75 ++++++++++ python/ebook_search/api/routes/__init__.py | 16 ++ python/ebook_search/api/routes/admin.py | 116 +++++++++++++++ python/ebook_search/api/routes/page.py | 66 +++++++++ python/ebook_search/api/routes/search.py | 66 +++++++++ python/ebook_search/api/static/style.css | 140 ++++++++++++++++++ python/ebook_search/api/templates/admin.html | 57 +++++++ .../api/templates/book_detail.html | 32 ++++ python/ebook_search/api/templates/books.html | 31 ++++ .../api/templates/partials/admin_status.html | 1 + .../api/templates/partials/error.html | 1 + .../api/templates/partials/results.html | 74 +++++++++ python/ebook_search/api/templates/search.html | 30 ++++ python/ebook_search/api/web.py | 13 ++ 16 files changed, 777 insertions(+) create mode 100644 python/ebook_search/api/__init__.py create mode 100644 python/ebook_search/api/bm25_tasks.py create mode 100644 python/ebook_search/api/main.py create mode 100644 python/ebook_search/api/routes/__init__.py create mode 100644 python/ebook_search/api/routes/admin.py create mode 100644 python/ebook_search/api/routes/page.py create mode 100644 python/ebook_search/api/routes/search.py create mode 100644 python/ebook_search/api/static/style.css create mode 100644 python/ebook_search/api/templates/admin.html create mode 100644 python/ebook_search/api/templates/book_detail.html create mode 100644 python/ebook_search/api/templates/books.html create mode 100644 python/ebook_search/api/templates/partials/admin_status.html create mode 100644 python/ebook_search/api/templates/partials/error.html create mode 100644 python/ebook_search/api/templates/partials/results.html create mode 100644 python/ebook_search/api/templates/search.html create mode 100644 python/ebook_search/api/web.py diff --git a/python/ebook_search/api/__init__.py b/python/ebook_search/api/__init__.py new file mode 100644 index 0000000..297fdb0 --- /dev/null +++ b/python/ebook_search/api/__init__.py @@ -0,0 +1 @@ +"""Web and external API adapters for EPUB search.""" diff --git a/python/ebook_search/api/bm25_tasks.py b/python/ebook_search/api/bm25_tasks.py new file mode 100644 index 0000000..a211d45 --- /dev/null +++ b/python/ebook_search/api/bm25_tasks.py @@ -0,0 +1,58 @@ +"""Background BM25 refresh tasks for the web app.""" + +from __future__ import annotations + +import logging +from threading import Timer +from typing import TYPE_CHECKING + +from sqlalchemy.orm import Session + +from python.ebook_search.bm25_corpus import refresh_bm25_corpus + +if TYPE_CHECKING: + from fastapi import FastAPI + from sqlalchemy.engine import Engine + + from python.ebook_search.config import EbookSearchConfig + +logger = logging.getLogger(__name__) + + +def schedule_bm25_refresh(app: FastAPI) -> None: + """Schedule a delayed BM25 corpus refresh, replacing any pending refresh.""" + existing_timer = getattr(app.state, "bm25_refresh_timer", None) + if existing_timer is not None: + existing_timer.cancel() + + timer = Timer(app.state.config.bm25_refresh_delay_seconds, refresh_bm25_for_app, args=(app,)) + timer.daemon = True + timer.start() + app.state.bm25_refresh_timer = timer + logger.info( + "ebook_bm25_refresh_scheduled delay_seconds=%s", + app.state.config.bm25_refresh_delay_seconds, + ) + + +def cancel_bm25_refresh(app: FastAPI) -> None: + """Cancel any pending BM25 corpus refresh.""" + existing_timer = getattr(app.state, "bm25_refresh_timer", None) + if existing_timer is not None: + existing_timer.cancel() + app.state.bm25_refresh_timer = None + logger.info("ebook_bm25_refresh_cancelled") + + +def refresh_bm25_for_app(app: FastAPI) -> None: + """Refresh the BM25 corpus using the app engine and config.""" + try: + refresh_bm25_for_engine(app.state.engine, app.state.config) + except Exception: + logger.exception("ebook_bm25_refresh_failed") + + +def refresh_bm25_for_engine(engine: Engine, config: EbookSearchConfig) -> None: + """Refresh the BM25 corpus using a SQLAlchemy engine.""" + with Session(engine) as session: + refresh_bm25_corpus(session, config) diff --git a/python/ebook_search/api/main.py b/python/ebook_search/api/main.py new file mode 100644 index 0000000..9be6d99 --- /dev/null +++ b/python/ebook_search/api/main.py @@ -0,0 +1,75 @@ +"""FastAPI HTMX app for EPUB search.""" + +from __future__ import annotations + +import logging +from contextlib import asynccontextmanager +from typing import TYPE_CHECKING, Annotated + +import typer +import uvicorn +from fastapi import FastAPI +from fastapi.staticfiles import StaticFiles +from sqlalchemy.orm import Session + +from python.common import configure_logger +from python.ebook_search.api.bm25_tasks import cancel_bm25_refresh +from python.ebook_search.api.routes import register_admin_routes, register_page_routes, register_search_routes +from python.ebook_search.api.web import STATIC_DIR +from python.ebook_search.bm25_corpus import ensure_bm25_corpus +from python.ebook_search.config import load_config +from python.orm.common import get_postgres_engine + +if TYPE_CHECKING: + from collections.abc import AsyncIterator + + +logger = logging.getLogger(__name__) + + +@asynccontextmanager +async def lifespan(app: FastAPI) -> AsyncIterator[None]: + """Manage application startup and shutdown resources.""" + logger.info("ebook_search_startup") + app.state.engine = get_postgres_engine(name="RICHIE") + with Session(app.state.engine) as session: + ensure_bm25_corpus(session, app.state.config) + try: + yield + finally: + logger.info("ebook_search_shutdown") + cancel_bm25_refresh(app) + app.state.engine.dispose() + + +def create_app() -> FastAPI: + """Create the EPUB search web app.""" + app = FastAPI(title="EPUB Search", lifespan=lifespan) + app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static") + app.state.config = load_config() + logger.info( + "ebook_search_config_loaded top_k=%s embedding_model=%s rerank_enabled=%s answer_enabled=%s library_paths=%s", + app.state.config.top_k, + app.state.config.embedding_model, + app.state.config.rerank.enabled, + app.state.config.answer_enabled, + len(app.state.config.library_paths), + ) + register_page_routes(app) + register_search_routes(app) + register_admin_routes(app) + return app + + +def serve( + host: Annotated[str, typer.Option("--host", "-h", help="Host to bind to")] = "127.0.0.1", + port: Annotated[int, typer.Option("--port", "-p", help="Port to bind to")] = 8070, + log_level: Annotated[str, typer.Option("--log-level", "-l", help="Log level")] = "INFO", +) -> None: + """Start the EPUB search server.""" + configure_logger(log_level) + uvicorn.run(create_app(), host=host, port=port) + + +if __name__ == "__main__": + typer.run(serve) diff --git a/python/ebook_search/api/routes/__init__.py b/python/ebook_search/api/routes/__init__.py new file mode 100644 index 0000000..a6e49ca --- /dev/null +++ b/python/ebook_search/api/routes/__init__.py @@ -0,0 +1,16 @@ +"""EPUB search web route modules.""" + +from python.ebook_search.api.routes import admin, page, search + +register_admin_routes = admin.register_admin_routes +register_page_routes = page.register_page_routes +register_search_routes = search.register_search_routes + +__all__ = [ + "admin", + "page", + "register_admin_routes", + "register_page_routes", + "register_search_routes", + "search", +] diff --git a/python/ebook_search/api/routes/admin.py b/python/ebook_search/api/routes/admin.py new file mode 100644 index 0000000..ff83239 --- /dev/null +++ b/python/ebook_search/api/routes/admin.py @@ -0,0 +1,116 @@ +"""Admin routes for the EPUB search web UI.""" + +from __future__ import annotations + +import logging +from dataclasses import replace +from typing import TYPE_CHECKING + +from fastapi import APIRouter, Request +from fastapi.responses import HTMLResponse +from sqlalchemy.orm import Session + +from python.ebook_search.api.bm25_tasks import schedule_bm25_refresh +from python.ebook_search.api.web import templates +from python.ebook_search.embeddings import embed_missing_chunks, embedding_model_stats +from python.ebook_search.ingest import ingest_configured_paths + +if TYPE_CHECKING: + from fastapi import FastAPI + +logger = logging.getLogger(__name__) + +router = APIRouter(prefix="/admin") +EMBED_ALL_BATCH_SIZE = 32 + + +def register_admin_routes(app: FastAPI) -> None: + """Register admin routes on the app.""" + app.include_router(router) + + +@router.get("", response_class=HTMLResponse) +def admin(request: Request) -> HTMLResponse: + """Render the admin page.""" + with Session(request.app.state.engine) as session: + stats = embedding_model_stats(session) + logger.info("ebook_admin_page_loaded models=%s", len(stats)) + return templates.TemplateResponse(request, "admin.html", {"config": request.app.state.config, "stats": stats}) + + +@router.post("/scan", response_class=HTMLResponse) +def scan_library(request: Request) -> HTMLResponse: + """Scan configured library paths for EPUB changes.""" + try: + with Session(request.app.state.engine) as session: + count = ingest_configured_paths(session, request.app.state.config) + session.commit() + except Exception as error: + logger.exception("ebook_admin_scan_failed") + return templates.TemplateResponse(request, "partials/error.html", {"message": str(error)}, status_code=500) + + logger.info("ebook_admin_scan_complete changed_files=%s", count) + if count > 0: + schedule_bm25_refresh(request.app) + return templates.TemplateResponse(request, "partials/admin_status.html", {"message": f"Indexed {count} EPUBs"}) + + +@router.post("/embed-missing", response_class=HTMLResponse) +def embed_missing(request: Request) -> HTMLResponse: + """Embed chunks missing vectors for the configured model.""" + try: + with Session(request.app.state.engine) as session: + count = embed_missing_chunks(session, request.app.state.config) + session.commit() + except Exception as error: + logger.exception("ebook_admin_embed_missing_failed") + return templates.TemplateResponse(request, "partials/error.html", {"message": str(error)}, status_code=500) + + logger.info("ebook_admin_embed_missing_complete chunks=%s", count) + return templates.TemplateResponse( + request, + "partials/admin_status.html", + {"message": f"Embedded {count} chunks"}, + ) + + +@router.post("/embed-all", response_class=HTMLResponse) +def embed_all(request: Request) -> HTMLResponse: + """Embed all chunks missing vectors in fixed-size batches.""" + total = 0 + batches = 0 + config = replace(request.app.state.config, embedding_batch_size=EMBED_ALL_BATCH_SIZE) + try: + with Session(request.app.state.engine) as session: + while True: + count = embed_missing_chunks(session, config) + if count == 0: + break + session.commit() + total += count + batches += 1 + logger.info( + "ebook_admin_embed_all_batch_complete batch=%s chunks=%s total_chunks=%s", + batches, + count, + total, + ) + except Exception as error: + logger.exception( + "ebook_admin_embed_all_failed batches=%s chunks=%s", + batches, + total, + ) + return templates.TemplateResponse( + request, + "partials/error.html", + {"message": f"Embed all failed after {total} chunks in {batches} batches: {error}"}, + status_code=500, + ) + + logger.info("ebook_admin_embed_all_complete batches=%s chunks=%s", batches, total) + return templates.TemplateResponse( + request, + "partials/admin_status.html", + {"message": f"Embedded {total} chunks in {batches} batches of {EMBED_ALL_BATCH_SIZE}"}, + ) diff --git a/python/ebook_search/api/routes/page.py b/python/ebook_search/api/routes/page.py new file mode 100644 index 0000000..8e48867 --- /dev/null +++ b/python/ebook_search/api/routes/page.py @@ -0,0 +1,66 @@ +"""Page routes for the EPUB search web UI.""" + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING + +from fastapi import APIRouter, Request +from fastapi.responses import HTMLResponse +from sqlalchemy import select +from sqlalchemy.orm import Session + +from python.ebook_search.api.web import templates +from python.orm.richie import EbookSource + +if TYPE_CHECKING: + from fastapi import FastAPI + +logger = logging.getLogger(__name__) + +router = APIRouter() + + +def register_page_routes(app: FastAPI) -> None: + """Register page routes on the app.""" + app.include_router(router) + + +@router.get("/", response_class=HTMLResponse) +def index(request: Request) -> HTMLResponse: + """Render the search page.""" + return templates.TemplateResponse(request, "search.html", {"config": request.app.state.config}) + + +@router.get("/books", response_class=HTMLResponse) +def books(request: Request) -> HTMLResponse: + """Render the indexed books page.""" + with Session(request.app.state.engine) as session: + sources = list(session.scalars(select(EbookSource).order_by(EbookSource.title)).all()) + logger.info("ebook_books_page_loaded count=%s", len(sources)) + return templates.TemplateResponse(request, "books.html", {"sources": sources}) + + +@router.get("/books/{source_id}", response_class=HTMLResponse) +def book_detail(source_id: int, request: Request) -> HTMLResponse: + """Render details for one indexed book.""" + with Session(request.app.state.engine) as session: + source = session.get(EbookSource, source_id) + if source is not None: + chapter_count = len(source.chapters) + chunk_count = len(source.chunks) + else: + chapter_count = 0 + chunk_count = 0 + logger.info( + "ebook_book_detail_loaded source_id=%s found=%s chapters=%s chunks=%s", + source_id, + source is not None, + chapter_count, + chunk_count, + ) + return templates.TemplateResponse( + request, + "book_detail.html", + {"chapter_count": chapter_count, "chunk_count": chunk_count, "source": source}, + ) diff --git a/python/ebook_search/api/routes/search.py b/python/ebook_search/api/routes/search.py new file mode 100644 index 0000000..77ed022 --- /dev/null +++ b/python/ebook_search/api/routes/search.py @@ -0,0 +1,66 @@ +"""Search routes for the EPUB search web UI.""" + +from __future__ import annotations + +import logging +from dataclasses import replace +from time import perf_counter +from typing import TYPE_CHECKING, Annotated + +from fastapi import APIRouter, Form, Request +from fastapi.responses import HTMLResponse + +from python.ebook_search.answer import answer_query +from python.ebook_search.api.web import templates +from python.ebook_search.search import search_ebooks +from python.ebook_search.timing import runtime_step_from_start + +if TYPE_CHECKING: + from fastapi import FastAPI + +logger = logging.getLogger(__name__) + +router = APIRouter() + + +def register_search_routes(app: FastAPI) -> None: + """Register search routes on the app.""" + app.include_router(router) + + +@router.post("/search", response_class=HTMLResponse) +def search( + request: Request, + query: Annotated[str, Form()], + rerank: Annotated[str | None, Form()] = None, +) -> HTMLResponse: + """Run a search and render HTMX results.""" + try: + response = search_ebooks(request.app.state.engine, query, request.app.state.config, rerank=rerank == "true") + except Exception as error: + logger.exception("ebook_search_request_failed") + return templates.TemplateResponse(request, "partials/error.html", {"message": str(error)}, status_code=500) + + answer_start = perf_counter() + if request.app.state.config.answer_enabled: + try: + answer = answer_query(query, response.results, request.app.state.config) + except RuntimeError as error: + logger.warning("ebook_answer_request_failed_falling_back error=%s", error) + answer = "Answer generation failed. Source chunks are still shown below." + else: + logger.info("ebook_answer_skipped_disabled") + answer = "Answer generation is disabled. Source chunks are shown below." + answer_step_name = "Answer generation" if request.app.state.config.answer_enabled else "Answer skipped" + response = replace( + response, + timings=(*response.timings, runtime_step_from_start(answer_step_name, answer_start)), + ) + + logger.info( + "ebook_search_request_complete results=%s rank_label=%s runtime_ms=%.1f", + len(response.results), + response.rank_label, + response.total_runtime_ms, + ) + return templates.TemplateResponse(request, "partials/results.html", {"answer": answer, "response": response}) diff --git a/python/ebook_search/api/static/style.css b/python/ebook_search/api/static/style.css new file mode 100644 index 0000000..c869d55 --- /dev/null +++ b/python/ebook_search/api/static/style.css @@ -0,0 +1,140 @@ +body { + margin: 0; + background: #f7f7f4; + color: #202124; + font-family: system-ui, -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; +} + +main { + max-width: 960px; + margin: 0 auto; + padding: 24px; +} + +nav { + display: flex; + gap: 12px; + align-items: center; + margin-bottom: 20px; +} + +nav form { + margin: 0; +} + +.actions { + display: flex; + flex-wrap: wrap; + gap: 12px; + margin-bottom: 24px; +} + +textarea { + display: block; + width: 100%; + margin: 8px 0 12px; +} + +button { + padding: 8px 14px; +} + +.check { + display: inline-flex; + gap: 8px; + align-items: center; + margin-right: 12px; +} + +.rank-label { + margin-top: 24px; + font-weight: 700; +} + +.results { + padding-left: 24px; +} + +.meta, +.scores, +.status { + color: #626a73; +} + +.scores { + display: flex; + flex-wrap: wrap; + gap: 8px; + margin: 12px 0; +} + +.scores div { + display: inline-flex; + gap: 4px; + align-items: baseline; +} + +.scores dt { + font-weight: 700; +} + +.scores dd { + margin: 0; +} + +.runtime { + margin-top: 16px; +} + +.timing-chart { + display: grid; + gap: 8px; + padding: 0; + list-style: none; +} + +.timing-chart li { + display: grid; + grid-template-columns: minmax(150px, 1fr) minmax(160px, 2fr) auto auto; + gap: 8px; + align-items: center; +} + +.timing-bar { + height: 10px; + overflow: hidden; + background: #e5e5df; +} + +.timing-bar span { + display: block; + height: 100%; + background: #3767c8; +} + +.timing-value, +.timing-remaining { + color: #626a73; + font-variant-numeric: tabular-nums; +} + +table { + width: 100%; + border-collapse: collapse; +} + +th, +td { + padding: 8px; + border-bottom: 1px solid #d8d8d2; + text-align: left; +} + +th { + font-weight: 700; +} + +.error { + color: #9f1d20; + font-weight: 700; +} diff --git a/python/ebook_search/api/templates/admin.html b/python/ebook_search/api/templates/admin.html new file mode 100644 index 0000000..12e588e --- /dev/null +++ b/python/ebook_search/api/templates/admin.html @@ -0,0 +1,57 @@ + + + + + + EPUB Admin + + + + +
+ +

Admin

+
+
+
+ +
+
+ +
+
+ +
+
+
+

Embeddings

+ + + + + + + + + + + + {% for item in stats %} + + + + + + + + {% endfor %} + +
ModelDimensionsEmbeddedMissingTotal chunks
{{ item.model_name }}{{ item.dimension }}{{ item.embedded_chunks }}{{ item.missing_chunks }}{{ item.total_chunks }}
+
+
+ + diff --git a/python/ebook_search/api/templates/book_detail.html b/python/ebook_search/api/templates/book_detail.html new file mode 100644 index 0000000..735aeaf --- /dev/null +++ b/python/ebook_search/api/templates/book_detail.html @@ -0,0 +1,32 @@ + + + + + + {% if source %}{{ source.title }}{% else %}Book not found{% endif %} + + + +
+ + {% if source %} +

{{ source.title }}

+

{{ source.author or "Unknown author" }}

+
+
File
+
{{ source.file_path }}
+
Chapters
+
{{ chapter_count }}
+
Chunks
+
{{ chunk_count }}
+
+ {% else %} +

Book not found

+ {% endif %} +
+ + diff --git a/python/ebook_search/api/templates/books.html b/python/ebook_search/api/templates/books.html new file mode 100644 index 0000000..c7bc487 --- /dev/null +++ b/python/ebook_search/api/templates/books.html @@ -0,0 +1,31 @@ + + + + + + EPUB Books + + + +
+ +

Books

+ {% if sources %} +
    + {% for source in sources %} +
  1. +

    {{ source.title }}

    +

    {{ source.author or "Unknown author" }}

    +
    2. + {% endfor %} +
+ {% else %} +

No EPUBs indexed.

+ {% endif %} +
+ + diff --git a/python/ebook_search/api/templates/partials/admin_status.html b/python/ebook_search/api/templates/partials/admin_status.html new file mode 100644 index 0000000..f8fa12f --- /dev/null +++ b/python/ebook_search/api/templates/partials/admin_status.html @@ -0,0 +1 @@ +

{{ message }}

diff --git a/python/ebook_search/api/templates/partials/error.html b/python/ebook_search/api/templates/partials/error.html new file mode 100644 index 0000000..9657121 --- /dev/null +++ b/python/ebook_search/api/templates/partials/error.html @@ -0,0 +1 @@ +

{{ message }}

diff --git a/python/ebook_search/api/templates/partials/results.html b/python/ebook_search/api/templates/partials/results.html new file mode 100644 index 0000000..bc29eec --- /dev/null +++ b/python/ebook_search/api/templates/partials/results.html @@ -0,0 +1,74 @@ +
{{ response.rank_label }}
+{% if response.timings %} +
+

Runtime

+

Total {{ "%.1f"|format(response.total_runtime_ms) }} ms

+
    + {% set total = response.total_runtime_ms %} + {% set ns = namespace(remaining=total) %} + {% for step in response.timings %} + {% set width = (step.duration_ms / total * 100) if total else 0 %} + {% if step.counts_toward_total %} + {% set ns.remaining = ns.remaining - step.duration_ms %} + {% endif %} +
  1. + {{ step.name }} + + {{ "%.1f"|format(step.duration_ms) }} ms + {{ "%.1f"|format([ns.remaining, 0]|max) }} ms left +
    2. + {% endfor %} +
+
+{% endif %} +
+

Answer

+

{{ answer }}

+
+{% if response.results %} +
    + {% for result in response.results %} +
  1. +

    {{ result.source_title }}

    +

    + {% if result.source_author %}{{ result.source_author }}{% endif %} + {% if result.chapter_title %} · {{ result.chapter_title }}{% endif %} + {% if result.page_label %} · page {{ result.page_label }}{% endif %} +

    +

    {{ result.text }}

    +
    +
    +
    final
    +
    {{ "%.3f"|format(result.score) }}
    +
    + {% if result.rerank_score is not none %} +
    +
    rerank
    +
    {{ "%.3f"|format(result.rerank_score) }}
    +
    + {% endif %} + {% if result.vector_score is not none %} +
    +
    vector cosine
    +
    {{ "%.3f"|format(result.vector_score) }}
    +
    + {% endif %} + {% if result.bm25_score is not none %} +
    +
    BM25
    +
    {{ "%.6f"|format(result.bm25_score) }}
    +
    + {% endif %} + {% if result.fused_score is not none %} +
    +
    RRF
    +
    {{ "%.3f"|format(result.fused_score) }}
    +
    + {% endif %} +
    +
    2. + {% endfor %} +
+{% else %} +

No results.

+{% endif %} diff --git a/python/ebook_search/api/templates/search.html b/python/ebook_search/api/templates/search.html new file mode 100644 index 0000000..df566c5 --- /dev/null +++ b/python/ebook_search/api/templates/search.html @@ -0,0 +1,30 @@ + + + + + + EPUB Search + + + + +
+ +

EPUB Search

+
+ + + + +
+
+
+ + diff --git a/python/ebook_search/api/web.py b/python/ebook_search/api/web.py new file mode 100644 index 0000000..85f6128 --- /dev/null +++ b/python/ebook_search/api/web.py @@ -0,0 +1,13 @@ +"""Shared web UI resources for EPUB search.""" + +from __future__ import annotations + +from pathlib import Path + +from fastapi.templating import Jinja2Templates + +PACKAGE_DIR = Path(__file__).resolve().parent +TEMPLATE_DIR = PACKAGE_DIR / "templates" +STATIC_DIR = PACKAGE_DIR / "static" + +templates = Jinja2Templates(directory=TEMPLATE_DIR) -- 2.54.0 From cad3f6f79e247d75a08dddb5623ac3516486c34c Mon Sep 17 00:00:00 2001 From: Richie Cahill Date: Fri, 12 Jun 2026 03:10:53 -0400 Subject: [PATCH 14/28] setup tests --- tests/test_ebook_search_core.py | 380 ++++++++++++++++++++++++++++++ tests/test_ebook_search_http.py | 84 +++++++ tests/test_ebook_search_rerank.py | 150 ++++++++++++ tests/test_ebook_search_ui.py | 265 +++++++++++++++++++++ 4 files changed, 879 insertions(+) create mode 100644 tests/test_ebook_search_core.py create mode 100644 tests/test_ebook_search_http.py create mode 100644 tests/test_ebook_search_rerank.py create mode 100644 tests/test_ebook_search_ui.py diff --git a/tests/test_ebook_search_core.py b/tests/test_ebook_search_core.py new file mode 100644 index 0000000..75ba762 --- /dev/null +++ b/tests/test_ebook_search_core.py @@ -0,0 +1,380 @@ +"""Tests for EPUB search core helpers.""" + +from __future__ import annotations + +from dataclasses import replace +from datetime import UTC, datetime +from os import environ +from pathlib import Path +from threading import Event +from types import ModuleType + +import pytest +from sqlalchemy import create_engine, select +from sqlalchemy.orm import sessionmaker + +from python.ebook_search.answer import answer_query +from python.ebook_search.bm25_corpus import ( + BM25Corpus, + BM25CorpusUnavailableError, + BM25Manifest, + ensure_bm25_corpus, + load_bm25_corpus, +) +from python.ebook_search.config import EbookSearchConfig, RerankConfig, load_config, normalize_embedding_model +from python.ebook_search.embeddings import MODEL_DIMENSIONS, ensure_embedding_models +from python.ebook_search.ingest import chunk_text, find_existing_source +from python.ebook_search.search import ( + SearchResponse, + SearchResult, + bm25_candidates, + reciprocal_rank_fusion, + retrieval_query_from_text, + search_ebooks, +) +from python.ebook_search.timing import RuntimeStep +from python.orm.richie import EbookEmbeddingModel, EbookSource, RichieBase + + +def test_chunk_text_uses_overlap() -> None: + chunks = chunk_text(" ".join(str(index) for index in range(100)), chunk_tokens=20, overlap_tokens=5) + + assert len(chunks) > 1 + assert chunks[0].token_start == 0 + assert chunks[1].token_start == 15 + assert all(chunk.token_count <= 20 for chunk in chunks) + + +def test_reciprocal_rank_fusion_combines_vector_and_bm25_rankings() -> None: + vector_results = [ + SearchResult(chunk_id=1, text="a", source_title="A", score=0.9, vector_score=0.9), + SearchResult(chunk_id=2, text="b", source_title="B", score=0.8, vector_score=0.8), + ] + lexical_results = [ + SearchResult(chunk_id=2, text="b", source_title="B", score=4.2, bm25_score=4.2), + SearchResult(chunk_id=3, text="c", source_title="C", score=2.1, bm25_score=2.1), + ] + + fused = reciprocal_rank_fusion(vector_results, lexical_results) + + assert [result.chunk_id for result in fused] == [2, 1, 3] + assert fused[0].rank_source == "Hybrid" + assert fused[0].vector_score == 0.8 + assert fused[0].bm25_score == 4.2 + assert fused[0].fused_score == fused[0].score + + +def test_find_existing_source_matches_path_or_hash() -> None: + engine = create_engine("sqlite+pysqlite:///:memory:", future=True) + RichieBase.metadata.create_all(engine) + with sessionmaker(bind=engine, expire_on_commit=False, future=True)() as session: + source = EbookSource( + title="Book", + author=None, + language=None, + publisher=None, + identifier=None, + file_path="/old/book.epub", + file_sha256="a" * 64, + file_mtime=datetime.now(tz=UTC), + file_size=10, + ) + session.add(source) + session.commit() + + assert find_existing_source(session, Path("/old/book.epub"), "b" * 64) == source + assert find_existing_source(session, Path("/new/book.epub"), "a" * 64) == source + + +def test_reciprocal_rank_fusion_marks_hybrid_source() -> None: + vector_results = [SearchResult(chunk_id=1, text="a", source_title="A")] + lexical_results = [SearchResult(chunk_id=2, text="b", source_title="B")] + + fused = reciprocal_rank_fusion(vector_results, lexical_results) + + assert {result.rank_source for result in fused} == {"Hybrid"} + + +def test_search_response_sums_runtime_steps() -> None: + response = SearchResponse( + query="query", + results=[], + rank_label="Hybrid", + timings=( + RuntimeStep(name="A", duration_ms=1.25), + RuntimeStep(name="B", duration_ms=2.75), + RuntimeStep(name="Parallel detail", duration_ms=10.0, counts_toward_total=False), + ), + ) + + assert response.total_runtime_ms == 4.0 + + +def test_search_ebooks_runs_vector_and_bm25_in_parallel(monkeypatch) -> None: + engine = create_engine("sqlite+pysqlite:///:memory:", future=True) + vector_started = Event() + bm25_started = Event() + received_engines: list[object] = [] + + def fake_vector_candidates(received_engine, query, _config): + """Return vector candidates after confirming BM25 has started.""" + received_engines.append(received_engine) + assert query == "parallel" + vector_started.set() + assert bm25_started.wait(timeout=2) + return [SearchResult(chunk_id=1, text="vector", source_title="Vector", vector_score=0.9)] + + def fake_bm25_candidates(query, _config): + """Return BM25 candidates after confirming vector search has started.""" + assert query == "parallel" + bm25_started.set() + assert vector_started.wait(timeout=2) + return [SearchResult(chunk_id=2, text="bm25", source_title="BM25", bm25_score=2.0)] + + monkeypatch.setattr("python.ebook_search.search.vector_candidates", fake_vector_candidates) + monkeypatch.setattr("python.ebook_search.search.bm25_candidates", fake_bm25_candidates) + config = EbookSearchConfig(rerank=RerankConfig(enabled=False)) + + response = search_ebooks(engine, "parallel", config) + + timings = {step.name: step for step in response.timings} + assert [result.chunk_id for result in response.results] == [1, 2] + assert timings["Embedding + vector search"].counts_toward_total is False + assert timings["BM25 search"].counts_toward_total is False + assert timings["Hybrid retrieval"].counts_toward_total is True + assert received_engines == [engine] + + +def test_retrieval_query_keeps_entity_and_series_terms() -> None: + assert retrieval_query_from_text("what does Damien Montgomery stand for in starship mage") == ( + "damien montgomery stand starship mage" + ) + + +def test_bm25_candidates_scores_whole_corpus(monkeypatch) -> None: + record = { + "chunk_id": 2, + "text": "high", + "source_title": "B", + "source_author": None, + "chapter_title": None, + "page_label": None, + "bm25_text": "high", + } + manifest = BM25Manifest(created_at=datetime.now(tz=UTC), db_updated_at=None, chunk_count=1) + corpus = BM25Corpus(retriever=object(), records=(record,), manifest=manifest) + captured: dict[str, object] = {} + + def fake_score_bm25_corpus(query, saved_corpus, *, limit): + captured["query"] = query + captured["corpus"] = saved_corpus + captured["limit"] = limit + return [(record, 1.5)] + + monkeypatch.setattr("python.ebook_search.search.load_bm25_corpus", lambda _config: corpus) + monkeypatch.setattr("python.ebook_search.search.score_bm25_corpus", fake_score_bm25_corpus) + config = EbookSearchConfig(rerank=RerankConfig(enabled=False)) + + results = bm25_candidates("high", config) + + assert captured["query"] == "high" + assert captured["corpus"] == corpus + assert captured["limit"] == 120 + assert [result.chunk_id for result in results] == [2] + assert [result.bm25_score for result in results] == [1.5] + + +def test_bm25_candidates_raises_when_corpus_is_unavailable(monkeypatch) -> None: + def fake_load_bm25_corpus(_config): + raise BM25CorpusUnavailableError + + monkeypatch.setattr("python.ebook_search.search.load_bm25_corpus", fake_load_bm25_corpus) + config = EbookSearchConfig(rerank=RerankConfig(enabled=False)) + + with pytest.raises(BM25CorpusUnavailableError): + bm25_candidates("high", config) + + +def test_load_bm25_corpus_caches_disk_load(monkeypatch, tmp_path) -> None: + load_bm25_corpus.cache_clear() + manifest = BM25Manifest(created_at=datetime.now(tz=UTC), db_updated_at=None, chunk_count=1) + record = { + "chunk_id": 2, + "text": "cached", + "source_title": "B", + "source_author": None, + "chapter_title": None, + "page_label": None, + "bm25_text": "cached", + } + load_count = 0 + + class FakeRetriever: + """Fake persisted BM25 retriever.""" + + corpus = (record,) + + class FakeBM25: + """Fake BM25 class with observable load count.""" + + @staticmethod + def load(index_path, *, load_corpus, mmap): + nonlocal load_count + load_count += 1 + assert index_path == tmp_path + assert load_corpus is True + assert mmap is True + return FakeRetriever() + + fake_bm25s = ModuleType("bm25s") + fake_bm25s.BM25 = FakeBM25 + monkeypatch.setattr("python.ebook_search.bm25_corpus.read_bm25_manifest", lambda _path: manifest) + monkeypatch.setattr("python.ebook_search.bm25_corpus.bm25_index_exists", lambda _path, _manifest: True) + monkeypatch.setattr("python.ebook_search.bm25_corpus.bm25s", fake_bm25s) + config = EbookSearchConfig(rerank=RerankConfig(enabled=False), bm25_index_dir=str(tmp_path)) + + try: + first = load_bm25_corpus(config) + second = load_bm25_corpus(config) + finally: + load_bm25_corpus.cache_clear() + + assert first is second + assert first is not None + assert first.records == (record,) + assert load_count == 1 + + +def test_load_bm25_corpus_raises_when_index_is_missing(monkeypatch, tmp_path) -> None: + load_bm25_corpus.cache_clear() + monkeypatch.setattr("python.ebook_search.bm25_corpus.read_bm25_manifest", lambda _path: None) + monkeypatch.setattr("python.ebook_search.bm25_corpus.bm25_index_exists", lambda _path, _manifest: False) + config = EbookSearchConfig(rerank=RerankConfig(enabled=False), bm25_index_dir=str(tmp_path)) + + try: + with pytest.raises(BM25CorpusUnavailableError, match="BM25 corpus is not available"): + load_bm25_corpus(config) + finally: + load_bm25_corpus.cache_clear() + + +def test_ensure_bm25_corpus_refreshes_missing_index(monkeypatch) -> None: + refreshed: list[object] = [] + db_updated_at = datetime.now(tz=UTC) + + monkeypatch.setattr("python.ebook_search.bm25_corpus.read_bm25_manifest", lambda _path: None) + monkeypatch.setattr("python.ebook_search.bm25_corpus.bm25_index_exists", lambda _path, _manifest: False) + monkeypatch.setattr("python.ebook_search.bm25_corpus.corpus_last_updated_at", lambda _session: db_updated_at) + monkeypatch.setattr( + "python.ebook_search.bm25_corpus.refresh_bm25_corpus", + lambda session, config, *, db_updated_at: refreshed.append((session, config, db_updated_at)), + ) + + config = EbookSearchConfig(rerank=RerankConfig(enabled=False)) + session = object() + + ensure_bm25_corpus(session, config) + + assert refreshed == [(session, config, db_updated_at)] + + +def test_ensure_bm25_corpus_refreshes_stale_index(monkeypatch) -> None: + refreshed: list[object] = [] + created_at = datetime(2026, 1, 1, tzinfo=UTC) + db_updated_at = datetime(2026, 1, 2, tzinfo=UTC) + manifest = BM25Manifest(created_at=created_at, db_updated_at=created_at, chunk_count=10) + + monkeypatch.setattr("python.ebook_search.bm25_corpus.read_bm25_manifest", lambda _path: manifest) + monkeypatch.setattr("python.ebook_search.bm25_corpus.bm25_index_exists", lambda _path, _manifest: True) + monkeypatch.setattr("python.ebook_search.bm25_corpus.corpus_last_updated_at", lambda _session: db_updated_at) + monkeypatch.setattr( + "python.ebook_search.bm25_corpus.refresh_bm25_corpus", + lambda session, config, *, db_updated_at: refreshed.append((session, config, db_updated_at)), + ) + + config = EbookSearchConfig(rerank=RerankConfig(enabled=False)) + session = object() + + ensure_bm25_corpus(session, config) + + assert refreshed == [(session, config, db_updated_at)] + + +def test_supported_embedding_models_match_service_names() -> None: + assert MODEL_DIMENSIONS == { + "qwen3-embedding-0.6b": 1024, + "qwen3-embedding-4b": 2560, + "qwen3-embedding-8b": 4096, + } + + +def test_ensure_embedding_models_registers_service_names() -> None: + engine = create_engine("sqlite+pysqlite:///:memory:", future=True) + RichieBase.metadata.create_all(engine) + with sessionmaker(bind=engine, expire_on_commit=False, future=True)() as session: + ensure_embedding_models(session) + session.commit() + + models = list(session.scalars(select(EbookEmbeddingModel).order_by(EbookEmbeddingModel.name))) + + assert [(model.name, model.dimension) for model in models] == [ + ("qwen3-embedding-0.6b", 1024), + ("qwen3-embedding-4b", 2560), + ("qwen3-embedding-8b", 4096), + ] + + +def test_embedding_model_aliases_normalize_to_provider_names() -> None: + assert normalize_embedding_model() == "qwen3-embedding-0.6b" + + environ["EBOOK_SEARCH_EMBEDDING_MODEL"] = "qwen3-embedding-0.6b" + assert normalize_embedding_model() == "qwen3-embedding-0.6b" + + environ["EBOOK_SEARCH_EMBEDDING_MODEL"] = "Qwen3-Embedding-0.6B" + assert normalize_embedding_model() == "qwen3-embedding-0.6b" + + environ["EBOOK_SEARCH_EMBEDDING_MODEL"] = "Qwen/Qwen3-Embedding-4B" + + assert normalize_embedding_model() == "qwen3-embedding-4b" + + environ["EBOOK_SEARCH_EMBEDDING_MODEL"] = "qwen3-embedding:8b" + assert normalize_embedding_model() == "qwen3-embedding-8b" + + environ["EBOOK_SEARCH_EMBEDDING_MODEL"] = "qwen3-embedding-8b" + assert normalize_embedding_model() == "qwen3-embedding-8b" + + +def test_answer_generation_is_enabled_by_default(monkeypatch) -> None: + monkeypatch.delenv("EBOOK_SEARCH_ANSWER_ENABLED", raising=False) + + config = load_config() + + assert config.answer_enabled is True + + +def test_chat_defaults_use_ollama_cloud(monkeypatch) -> None: + monkeypatch.delenv("EBOOK_SEARCH_VLLM_BASE_URL", raising=False) + monkeypatch.delenv("EBOOK_SEARCH_CHAT_MODEL", raising=False) + + config = load_config() + + assert config.vllm_base_url == "https://ollama.com/v1" + assert config.chat_model == "deepseek-v4-flash" + + +def test_chat_api_key_falls_back_to_ollama_api_key(monkeypatch) -> None: + monkeypatch.delenv("EBOOK_SEARCH_VLLM_API_KEY", raising=False) + monkeypatch.setenv("OLLAMA_API_KEY", "ollama-key") + + config = load_config() + + assert config.vllm_api_key == "ollama-key" + + +def test_answer_query_does_not_call_model_when_disabled() -> None: + config = replace(load_config(), answer_enabled=False) + result = SearchResult(chunk_id=1, text="source text", source_title="Book") + + answer = answer_query("question", [result], config) + + assert "Answer generation is disabled" in answer diff --git a/tests/test_ebook_search_http.py b/tests/test_ebook_search_http.py new file mode 100644 index 0000000..993f7bd --- /dev/null +++ b/tests/test_ebook_search_http.py @@ -0,0 +1,84 @@ +"""Tests for EPUB search HTTP model adapters.""" + +from __future__ import annotations + +import httpx +import pytest + +from python.ebook_search.answer import answer_query +from python.ebook_search.config import EbookSearchConfig, RerankConfig +from python.ebook_search.embeddings import embed_texts +from python.ebook_search.search import SearchResult + + +def test_answer_query_uses_httpx_chat_completions(monkeypatch) -> None: + captured: dict[str, object] = {} + + def fake_post(url: str, **kwargs: object) -> httpx.Response: + captured["url"] = url + captured["kwargs"] = kwargs + return httpx.Response( + 200, + json={"choices": [{"message": {"content": "grounded answer"}}]}, + request=httpx.Request("POST", url), + ) + + monkeypatch.setattr(httpx, "post", fake_post) + config = EbookSearchConfig( + rerank=RerankConfig(enabled=False), + vllm_base_url="https://ollama.com/v1", + vllm_api_key="secret", + chat_model="deepseek-v4-flash", + ) + + answer = answer_query("question", [SearchResult(chunk_id=1, text="source", source_title="Book")], config) + + assert answer == "grounded answer" + assert captured["url"] == "https://ollama.com/v1/chat/completions" + kwargs = captured["kwargs"] + assert isinstance(kwargs, dict) + assert kwargs["headers"] == {"Authorization": "Bearer secret"} + payload = kwargs["json"] + assert isinstance(payload, dict) + assert payload["model"] == "deepseek-v4-flash" + + +def test_embed_texts_uses_httpx_embeddings(monkeypatch) -> None: + captured: dict[str, object] = {} + vector = [0.0] * 1024 + + def fake_post(url: str, **kwargs: object) -> httpx.Response: + captured["url"] = url + captured["kwargs"] = kwargs + return httpx.Response( + 200, + json={"data": [{"embedding": vector}]}, + request=httpx.Request("POST", url), + ) + + monkeypatch.setattr(httpx, "post", fake_post) + config = EbookSearchConfig( + rerank=RerankConfig(enabled=False), + embedding_base_url="http://bob:8000/v1", + embedding_model="qwen3-embedding-0.6b", + ) + + embeddings = embed_texts(["hello"], config) + + assert embeddings == [vector] + assert captured["url"] == "http://bob:8000/v1/embeddings" + kwargs = captured["kwargs"] + assert isinstance(kwargs, dict) + assert kwargs["headers"] == {} + assert kwargs["json"] == {"model": "qwen3-embedding-0.6b", "input": ["hello"]} + + +def test_embed_texts_rejects_bad_response_shape(monkeypatch) -> None: + def fake_post(url: str, **_kwargs: object) -> httpx.Response: + return httpx.Response(200, json={"data": [{}]}, request=httpx.Request("POST", url)) + + monkeypatch.setattr(httpx, "post", fake_post) + config = EbookSearchConfig(rerank=RerankConfig(enabled=False)) + + with pytest.raises(RuntimeError, match="Embedding request failed"): + embed_texts(["hello"], config) diff --git a/tests/test_ebook_search_rerank.py b/tests/test_ebook_search_rerank.py new file mode 100644 index 0000000..7ccae46 --- /dev/null +++ b/tests/test_ebook_search_rerank.py @@ -0,0 +1,150 @@ +"""Tests for EPUB search reranking.""" + +from __future__ import annotations + +import httpx +import pytest + +from python.ebook_search.config import EbookSearchConfig, RerankConfig, load_rerank_config +from python.ebook_search.rerank import rerank_chunks +from python.ebook_search.search import SearchResult, apply_rerank, skip_rerank + + +def candidates() -> list[SearchResult]: + return [ + SearchResult(chunk_id=1, text="alpha", source_title="A", score=0.9), + SearchResult(chunk_id=2, text="beta", source_title="B", score=0.8), + SearchResult(chunk_id=3, text="gamma", source_title="C", score=0.7), + ] + + +def rerank_response(payload: dict[str, object] | None = None, *, content: bytes | None = None) -> httpx.Response: + return httpx.Response( + 200, + content=content, + json=payload, + request=httpx.Request("POST", "http://rerank.test/rerank"), + ) + + +def test_config_defaults_keep_reranking_optional(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("EBOOK_SEARCH_RERANK_ENABLED", raising=False) + monkeypatch.delenv("EBOOK_SEARCH_RERANK_BASE_URL", raising=False) + monkeypatch.delenv("EBOOK_SEARCH_RERANK_MODEL", raising=False) + monkeypatch.delenv("EBOOK_SEARCH_RERANK_CANDIDATES", raising=False) + monkeypatch.delenv("EBOOK_SEARCH_RERANK_TIMEOUT_SECONDS", raising=False) + + config = load_rerank_config() + + assert config.enabled is False + assert config.base_url == "http://192.168.90.25:8001" + assert config.model == "qwen3-reranker-06b" + assert config.candidates == 24 + assert config.timeout_seconds == 30 + + +def test_reranking_disabled_returns_original_fused_order() -> None: + config = EbookSearchConfig(rerank=RerankConfig(enabled=False), top_k=2) + + response = skip_rerank("query", candidates(), config) + + assert response.rank_label == "Hybrid" + assert [result.chunk_id for result in response.results] == [1, 2] + + +def test_reranking_enabled_reorders_candidates(monkeypatch: pytest.MonkeyPatch) -> None: + def fake_post(_url: str, *, json: dict[str, object], timeout: float) -> httpx.Response: + assert timeout == 30 + assert json == { + "model": "qwen3-reranker-06b", + "query": "query", + "documents": ["alpha", "beta", "gamma"], + } + return rerank_response( + { + "results": [ + {"index": 0, "relevance_score": 0.1}, + {"index": 1, "relevance_score": 0.9}, + {"index": 2, "relevance_score": 0.4}, + ] + } + ) + + monkeypatch.setattr(httpx, "post", fake_post) + + results = rerank_chunks("query", candidates(), RerankConfig()) + + assert [result.chunk_id for result in results] == [2, 1, 3] + assert [round(result.score, 3) for result in results] == [0.45, 0.1, 0.0] + assert [result.rerank_score for result in results] == [0.9, 0.1, 0.4] + + +def test_reranking_cannot_ignore_hybrid_score(monkeypatch: pytest.MonkeyPatch) -> None: + candidates = [ + SearchResult(chunk_id=1, text="strong hybrid", source_title="A", score=1.0), + SearchResult(chunk_id=2, text="weak hybrid", source_title="B", score=0.1), + ] + + def fake_post(_url: str, **_kwargs: object) -> httpx.Response: + return rerank_response( + { + "results": [ + {"index": 0, "relevance_score": 0.7}, + {"index": 1, "relevance_score": 1.0}, + ] + } + ) + + monkeypatch.setattr(httpx, "post", fake_post) + + results = rerank_chunks("query", candidates, RerankConfig()) + + assert [result.chunk_id for result in results] == [1, 2] + assert results[0].score == 0.7 + assert results[1].score == 0.0 + assert results[1].rerank_score == 1.0 + + +def test_vllm_rerank_timeout_raises(monkeypatch: pytest.MonkeyPatch) -> None: + def fake_rerank_chunks( + _query: str, + _candidates: list[SearchResult], + _config: RerankConfig, + ) -> list[SearchResult]: + message = "timeout" + raise httpx.TimeoutException(message) + + monkeypatch.setattr("python.ebook_search.search.rerank_chunks", fake_rerank_chunks) + config = EbookSearchConfig(rerank=RerankConfig(enabled=True), top_k=2) + + with pytest.raises(httpx.TimeoutException, match="timeout"): + apply_rerank("query", candidates(), config) + + +def test_malformed_vllm_rerank_json_does_not_crash_search(monkeypatch: pytest.MonkeyPatch) -> None: + def fake_post(_url: str, **_kwargs: object) -> httpx.Response: + return rerank_response(content=b"not-json") + + monkeypatch.setattr(httpx, "post", fake_post) + + results = rerank_chunks("query", candidates()[:1], RerankConfig()) + + assert results[0].score == 0.0 + + +def test_vllm_rerank_scores_are_clamped(monkeypatch: pytest.MonkeyPatch) -> None: + def fake_post(_url: str, **_kwargs: object) -> httpx.Response: + return rerank_response( + { + "results": [ + {"index": 0, "relevance_score": -1}, + {"index": 1, "relevance_score": 2}, + ] + } + ) + + monkeypatch.setattr(httpx, "post", fake_post) + + results = rerank_chunks("query", candidates()[:2], RerankConfig()) + + assert [result.rerank_score for result in results] == [0.0, 1.0] diff --git a/tests/test_ebook_search_ui.py b/tests/test_ebook_search_ui.py new file mode 100644 index 0000000..03b31b2 --- /dev/null +++ b/tests/test_ebook_search_ui.py @@ -0,0 +1,265 @@ +"""Tests for EPUB search HTMX routes.""" + +from __future__ import annotations + +from fastapi.testclient import TestClient +from sqlalchemy import create_engine + +from python.ebook_search.api.main import create_app +from python.ebook_search.config import EbookSearchConfig, RerankConfig +from python.ebook_search.embeddings import EmbeddingModelStats +from python.ebook_search.search import SearchResponse, SearchResult +from python.ebook_search.timing import RuntimeStep + + +def patch_app_runtime(monkeypatch): + """Patch app startup dependencies used by UI route tests.""" + monkeypatch.setattr("python.ebook_search.api.main.get_postgres_engine", fake_get_postgres_engine) + monkeypatch.setattr("python.ebook_search.api.main.ensure_bm25_corpus", lambda _session, _config: None) + + +def fake_get_postgres_engine(**_kwargs): + """Return an in-memory engine for route tests.""" + return create_engine("sqlite+pysqlite:///:memory:", future=True) + + +def test_ui_form_passes_rerank_flag_to_search_handler(monkeypatch) -> None: + captured: dict[str, object] = {} + + def fake_search_ebooks(_engine, query, config, *, rerank=False): + captured["query"] = query + captured["rerank"] = rerank + captured["config"] = config + return SearchResponse(query=query, results=[], rank_label="Hybrid + rerank") + + monkeypatch.setattr("python.ebook_search.api.routes.search.search_ebooks", fake_search_ebooks) + monkeypatch.setattr( + "python.ebook_search.api.routes.search.answer_query", + lambda _query, _results, _config: "answer", + ) + patch_app_runtime(monkeypatch) + app = create_app() + app.state.config = EbookSearchConfig(rerank=RerankConfig(enabled=False), top_k=12, answer_enabled=True) + + with TestClient(app) as client: + response = client.post("/search", data={"query": "where is the quote?", "rerank": "true"}) + + assert response.status_code == 200 + assert "Hybrid + rerank" in response.text + assert captured["query"] == "where is the quote?" + assert captured["rerank"] is True + + +def test_ui_search_failure_returns_visible_error(monkeypatch) -> None: + def fake_search_ebooks(_engine, _query, _config, *, rerank=False): + del rerank + msg = "search exploded" + raise RuntimeError(msg) + + monkeypatch.setattr("python.ebook_search.api.routes.search.search_ebooks", fake_search_ebooks) + patch_app_runtime(monkeypatch) + app = create_app() + app.state.config = EbookSearchConfig(rerank=RerankConfig(enabled=False), top_k=12) + + with TestClient(app) as client: + response = client.post("/search", data={"query": "where is the quote?"}) + + assert response.status_code == 500 + assert "search exploded" in response.text + + +def test_ui_answer_failure_still_returns_sources(monkeypatch) -> None: + def fake_search_ebooks(_engine, query, _config, *, rerank=False): + del rerank + return SearchResponse(query=query, results=[], rank_label="Hybrid") + + def fake_answer_query(_query, _results, _config): + msg = "answer exploded" + raise RuntimeError(msg) + + monkeypatch.setattr("python.ebook_search.api.routes.search.search_ebooks", fake_search_ebooks) + monkeypatch.setattr("python.ebook_search.api.routes.search.answer_query", fake_answer_query) + patch_app_runtime(monkeypatch) + app = create_app() + app.state.config = EbookSearchConfig(rerank=RerankConfig(enabled=False), top_k=12, answer_enabled=True) + + with TestClient(app) as client: + response = client.post("/search", data={"query": "where is the quote?"}) + + assert response.status_code == 200 + assert "Answer generation failed" in response.text + + +def test_ui_skips_answer_when_disabled(monkeypatch) -> None: + called = False + + def fake_search_ebooks(_engine, query, _config, *, rerank=False): + del rerank + return SearchResponse(query=query, results=[], rank_label="Hybrid") + + def fake_answer_query(_query, _results, _config): + nonlocal called + called = True + return "answer" + + monkeypatch.setattr("python.ebook_search.api.routes.search.search_ebooks", fake_search_ebooks) + monkeypatch.setattr("python.ebook_search.api.routes.search.answer_query", fake_answer_query) + patch_app_runtime(monkeypatch) + app = create_app() + app.state.config = EbookSearchConfig(rerank=RerankConfig(enabled=False), answer_enabled=False) + + with TestClient(app) as client: + response = client.post("/search", data={"query": "where is the quote?"}) + + assert response.status_code == 200 + assert called is False + assert "Answer generation is disabled" in response.text + + +def test_ui_shows_component_scores(monkeypatch) -> None: + def fake_search_ebooks(_engine, query, _config, *, rerank=False): + del rerank + return SearchResponse( + query=query, + rank_label="Hybrid + rerank", + results=[ + SearchResult( + chunk_id=1, + text="source text", + source_title="Book", + score=0.9, + rerank_score=0.9, + vector_score=0.8, + bm25_score=2.5, + fused_score=0.03, + ) + ], + ) + + monkeypatch.setattr("python.ebook_search.api.routes.search.search_ebooks", fake_search_ebooks) + monkeypatch.setattr( + "python.ebook_search.api.routes.search.answer_query", + lambda _query, _results, _config: "answer", + ) + patch_app_runtime(monkeypatch) + app = create_app() + app.state.config = EbookSearchConfig(rerank=RerankConfig(enabled=False), answer_enabled=True) + + with TestClient(app) as client: + response = client.post("/search", data={"query": "where is the quote?"}) + + assert response.status_code == 200 + assert "rerank" in response.text + assert "vector cosine" in response.text + assert "BM25" in response.text + assert "RRF" in response.text + + +def test_ui_shows_search_runtime_chart(monkeypatch) -> None: + def fake_search_ebooks(_engine, query, _config, *, rerank=False): + del rerank + return SearchResponse( + query=query, + rank_label="Hybrid", + results=[], + timings=( + RuntimeStep(name="Embedding + vector search", duration_ms=12.5), + RuntimeStep(name="BM25 search", duration_ms=4.0), + ), + ) + + monkeypatch.setattr("python.ebook_search.api.routes.search.search_ebooks", fake_search_ebooks) + monkeypatch.setattr( + "python.ebook_search.api.routes.search.answer_query", + lambda _query, _results, _config: "answer", + ) + patch_app_runtime(monkeypatch) + app = create_app() + app.state.config = EbookSearchConfig(rerank=RerankConfig(enabled=False), answer_enabled=True) + + with TestClient(app) as client: + response = client.post("/search", data={"query": "where is the quote?"}) + + assert response.status_code == 200 + assert "Runtime" in response.text + assert "Total" in response.text + assert "Embedding + vector search" in response.text + assert "BM25 search" in response.text + assert "Answer generation" in response.text + assert "ms left" in response.text + + +def test_ui_embed_all_batches_until_complete(monkeypatch) -> None: + counts = iter([32, 32, 5, 0]) + batch_sizes: list[int] = [] + + def fake_embed_missing_chunks(_session, config): + batch_sizes.append(config.embedding_batch_size) + return next(counts) + + monkeypatch.setattr("python.ebook_search.api.routes.admin.embed_missing_chunks", fake_embed_missing_chunks) + patch_app_runtime(monkeypatch) + app = create_app() + + with TestClient(app) as client: + response = client.post("/admin/embed-all") + + assert response.status_code == 200 + assert "Embedded 69 chunks in 3 batches of 32" in response.text + assert batch_sizes == [32, 32, 32, 32] + + +def test_ui_scan_schedules_bm25_refresh_after_database_change(monkeypatch) -> None: + scheduled = False + + def fake_ingest_configured_paths(_session, _config): + return 1 + + def fake_schedule_bm25_refresh(_app): + nonlocal scheduled + scheduled = True + + monkeypatch.setattr("python.ebook_search.api.routes.admin.ingest_configured_paths", fake_ingest_configured_paths) + monkeypatch.setattr("python.ebook_search.api.routes.admin.schedule_bm25_refresh", fake_schedule_bm25_refresh) + patch_app_runtime(monkeypatch) + app = create_app() + + with TestClient(app) as client: + response = client.post("/admin/scan") + + assert response.status_code == 200 + assert "Indexed 1 EPUBs" in response.text + assert scheduled is True + + +def test_admin_page_shows_embedding_counts_by_model(monkeypatch) -> None: + def fake_embedding_model_stats(_session): + return [ + EmbeddingModelStats( + model_name="qwen3-embedding-0.6b", + dimension=1024, + embedded_chunks=40, + total_chunks=64, + ), + EmbeddingModelStats( + model_name="qwen3-embedding-4b", + dimension=2560, + embedded_chunks=8, + total_chunks=64, + ), + ] + + monkeypatch.setattr("python.ebook_search.api.routes.admin.embedding_model_stats", fake_embedding_model_stats) + patch_app_runtime(monkeypatch) + app = create_app() + + with TestClient(app) as client: + response = client.get("/admin") + + assert response.status_code == 200 + assert "qwen3-embedding-0.6b" in response.text + assert "1024" in response.text + assert "40" in response.text + assert "24" in response.text + assert "qwen3-embedding-4b" in response.text + assert "2560" in response.text -- 2.54.0 From 1efa7b047a6c6605321dbd16fc029083438855d8 Mon Sep 17 00:00:00 2001 From: Richie Cahill Date: Fri, 12 Jun 2026 03:12:26 -0400 Subject: [PATCH 15/28] updated python --- overlays/default.nix | 29 ++++++++++++++++++++++++++++- 1 file changed, 28 insertions(+), 1 deletion(-) diff --git a/overlays/default.nix b/overlays/default.nix index ee5a252..a7f635c 100644 --- a/overlays/default.nix +++ b/overlays/default.nix @@ -17,15 +17,41 @@ python-env = final: _prev: { my_python = final.python314.withPackages ( - ps: with ps; [ + ps: + let + bm25s = ps.buildPythonPackage rec { + pname = "bm25s"; + version = "0.3.9"; + pyproject = true; + + src = final.fetchPypi { + inherit pname version; + hash = "sha256-iVxnnZUrfeg1XttfPhpiCh4vKU0dQrkZvwghzOLi9Zc="; + }; + + build-system = [ ps.setuptools ]; + dependencies = with ps; [ + numpy + scipy + ]; + + pythonImportsCheck = [ "bm25s" ]; + }; + in + with ps; + [ alembic apprise apscheduler + beautifulsoup4 + ebooklib fastapi fastapi-cli httpx mypy + numpy orjson + pgvector polars psycopg pydantic @@ -39,6 +65,7 @@ scalene sqlalchemy sqlalchemy + bm25s tenacity textual tiktoken -- 2.54.0 From a38ce4505f95279982f19d1b4f2b78d9f214b4e6 Mon Sep 17 00:00:00 2001 From: Richie Cahill Date: Fri, 12 Jun 2026 13:11:17 -0400 Subject: [PATCH 16/28] add .ebook_search_bm25 to gitignore --- .gitignore | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/.gitignore b/.gitignore index 77b3d98..cf7b0f4 100644 --- a/.gitignore +++ b/.gitignore @@ -171,4 +171,5 @@ frontend/dist/ frontend/node_modules/ # data from testing llms -data/* \ No newline at end of file +data/* +.ebook_search_bm25 -- 2.54.0 From 2efc9e30a8e8b2dea2ce113734adf0723f9638e1 Mon Sep 17 00:00:00 2001 From: Richie Cahill Date: Fri, 12 Jun 2026 13:34:59 -0400 Subject: [PATCH 17/28] improved queary for vector search --- python/ebook_search/search.py | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/python/ebook_search/search.py b/python/ebook_search/search.py index e25c17d..5c179db 100644 --- a/python/ebook_search/search.py +++ b/python/ebook_search/search.py @@ -93,13 +93,14 @@ def search_ebooks( logger.info("ebook_search_start query_length=%s rerank=%s", len(query), rerank) timings: list[RuntimeStep] = [] - retrieval_query, timing = timed_result("Query preparation", retrieval_query_from_text, query) + bm25_query, timing = timed_result("BM25 query preparation", retrieval_query_from_text, query) timings.append(timing) retrieval, timing = timed_result( "Hybrid retrieval", parallel_retrieval, engine, - retrieval_query, + query, + bm25_query, config, ) timings.extend(retrieval.timings) @@ -130,7 +131,12 @@ def search_ebooks( return response -def parallel_retrieval(engine: Engine, query: str, config: EbookSearchConfig) -> RetrievalResponse: +def parallel_retrieval( + engine: Engine, + vector_query: str, + bm25_query: str, + config: EbookSearchConfig, +) -> RetrievalResponse: """Run vector and BM25 candidate retrieval concurrently with separate database sessions.""" with ThreadPoolExecutor(max_workers=2, thread_name_prefix="ebook-search") as executor: vector_future = executor.submit( @@ -138,14 +144,14 @@ def parallel_retrieval(engine: Engine, query: str, config: EbookSearchConfig) -> "Embedding + vector search", vector_candidates, engine, - query, + vector_query, config, ) bm25_future = executor.submit( timed_result, "BM25 search", bm25_candidates, - query, + bm25_query, config, ) vector_results, vector_timing = vector_future.result() @@ -196,7 +202,7 @@ def apply_rerank( def vector_candidates(engine: Engine, query: str, config: EbookSearchConfig) -> list[SearchResult]: - """Return pgvector cosine candidates for a normalized query.""" + """Return pgvector cosine candidates for a natural-language query.""" with Session(engine) as session: model = session.scalar(select(EbookEmbeddingModel).where(EbookEmbeddingModel.name == config.embedding_model)) if model is None: -- 2.54.0 From 3d582243fccf33ae1bdd3657986c81469c247534 Mon Sep 17 00:00:00 2001 From: Richie Cahill Date: Fri, 12 Jun 2026 13:35:20 -0400 Subject: [PATCH 18/28] fixed duplicat enrichment --- python/ebook_search/bm25_corpus.py | 8 +------- 1 file changed, 1 insertion(+), 7 deletions(-) diff --git a/python/ebook_search/bm25_corpus.py b/python/ebook_search/bm25_corpus.py index 79172a4..6c8a15d 100644 --- a/python/ebook_search/bm25_corpus.py +++ b/python/ebook_search/bm25_corpus.py @@ -172,13 +172,7 @@ def fetch_bm25_corpus_records(session: Session) -> list[dict[str, object]]: EbookSource.author.label("source_author"), EbookChapter.title.label("chapter_title"), EbookChunk.page_label.label("page_label"), - func.concat_ws( - " ", - EbookSource.title, - EbookSource.author, - EbookChapter.title, - EbookChunk.search_text, - ).label("bm25_text"), + EbookChunk.search_text.label("bm25_text"), ) .select_from(EbookChunk) .join(EbookSource, EbookSource.id == EbookChunk.source_id) -- 2.54.0 From 2f1affa2e54e990d836eaeeb5182cae26321607d Mon Sep 17 00:00:00 2001 From: Richie Cahill Date: Fri, 12 Jun 2026 13:36:34 -0400 Subject: [PATCH 19/28] improved reranking weights --- python/ebook_search/rerank.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/python/ebook_search/rerank.py b/python/ebook_search/rerank.py index abb2ee0..5075601 100644 --- a/python/ebook_search/rerank.py +++ b/python/ebook_search/rerank.py @@ -13,6 +13,8 @@ if TYPE_CHECKING: from python.ebook_search.search import SearchResult logger = logging.getLogger(__name__) +RERANK_SCORE_WEIGHT = 0.7 +HYBRID_SCORE_WEIGHT = 0.3 @dataclass(frozen=True) @@ -110,7 +112,7 @@ def clamp_score(score: float) -> float: def final_rerank_score(result: SearchResult, rerank_score: float, candidates: list[SearchResult]) -> float: """Combine rerank relevance with normalized hybrid retrieval evidence.""" - return rerank_score * normalized_hybrid_score(result, candidates) + return (RERANK_SCORE_WEIGHT * rerank_score) + (HYBRID_SCORE_WEIGHT * normalized_hybrid_score(result, candidates)) def normalized_hybrid_score(result: SearchResult, candidates: list[SearchResult]) -> float: -- 2.54.0 From 70f24cdbc6e043c4fd6819c79dae7fbeba8b275d Mon Sep 17 00:00:00 2001 From: Richie Cahill Date: Fri, 12 Jun 2026 13:36:45 -0400 Subject: [PATCH 20/28] updated tests --- tests/test_ebook_search_core.py | 49 +++++++++++++++++++++++++++++-- tests/test_ebook_search_rerank.py | 10 +++---- 2 files changed, 51 insertions(+), 8 deletions(-) diff --git a/tests/test_ebook_search_core.py b/tests/test_ebook_search_core.py index 75ba762..c4b05cc 100644 --- a/tests/test_ebook_search_core.py +++ b/tests/test_ebook_search_core.py @@ -19,6 +19,7 @@ from python.ebook_search.bm25_corpus import ( BM25CorpusUnavailableError, BM25Manifest, ensure_bm25_corpus, + fetch_bm25_corpus_records, load_bm25_corpus, ) from python.ebook_search.config import EbookSearchConfig, RerankConfig, load_config, normalize_embedding_model @@ -33,7 +34,7 @@ from python.ebook_search.search import ( search_ebooks, ) from python.ebook_search.timing import RuntimeStep -from python.orm.richie import EbookEmbeddingModel, EbookSource, RichieBase +from python.orm.richie import EbookChapter, EbookChunk, EbookEmbeddingModel, EbookSource, RichieBase def test_chunk_text_uses_overlap() -> None: @@ -86,6 +87,47 @@ def test_find_existing_source_matches_path_or_hash() -> None: assert find_existing_source(session, Path("/new/book.epub"), "a" * 64) == source +def test_bm25_corpus_uses_existing_search_text_without_duplicate_metadata() -> None: + engine = create_engine("sqlite+pysqlite:///:memory:", future=True) + RichieBase.metadata.create_all(engine) + with sessionmaker(bind=engine, expire_on_commit=False, future=True)() as session: + source = EbookSource( + title="Book", + author="Author", + language=None, + publisher=None, + identifier=None, + file_path="/book.epub", + file_sha256="a" * 64, + file_mtime=datetime.now(tz=UTC), + file_size=10, + ) + session.add(source) + session.flush() + chapter = EbookChapter(source_id=source.id, spine_index=0, title="Chapter", href=None) + session.add(chapter) + session.flush() + session.add( + EbookChunk( + id=1, + source_id=source.id, + chapter_id=chapter.id, + chunk_index=0, + text="content", + token_start=0, + token_count=1, + page_label=None, + content_sha256="b" * 64, + search_text="Book Author Chapter content", + ) + ) + session.commit() + + records = fetch_bm25_corpus_records(session) + + assert records[0]["bm25_text"] == "Book Author Chapter content" + + def test_reciprocal_rank_fusion_marks_hybrid_source() -> None: vector_results = [SearchResult(chunk_id=1, text="a", source_title="A")] lexical_results = [SearchResult(chunk_id=2, text="b", source_title="B")] @@ -119,7 +161,7 @@ def test_search_ebooks_runs_vector_and_bm25_in_parallel(monkeypatch) -> None: def fake_vector_candidates(received_engine, query, _config): """Return vector candidates after confirming BM25 has started.""" received_engines.append(received_engine) - assert query == "parallel" + assert query == "what is parallel" vector_started.set() assert bm25_started.wait(timeout=2) return [SearchResult(chunk_id=1, text="vector", source_title="Vector", vector_score=0.9)] @@ -135,13 +177,14 @@ def test_search_ebooks_runs_vector_and_bm25_in_parallel(monkeypatch) -> None: monkeypatch.setattr("python.ebook_search.search.bm25_candidates", fake_bm25_candidates) config = EbookSearchConfig(rerank=RerankConfig(enabled=False)) - response = search_ebooks(engine, "parallel", config) + response = search_ebooks(engine, "what is parallel", config) timings = {step.name: step for step in response.timings} assert [result.chunk_id for result in response.results] == [1, 2] assert timings["Embedding + vector search"].counts_toward_total is False assert timings["BM25 search"].counts_toward_total is False assert timings["Hybrid retrieval"].counts_toward_total is True + assert timings["BM25 query preparation"].counts_toward_total is True assert received_engines == [engine] diff --git a/tests/test_ebook_search_rerank.py b/tests/test_ebook_search_rerank.py index 7ccae46..db53333 100644 --- a/tests/test_ebook_search_rerank.py +++ b/tests/test_ebook_search_rerank.py @@ -75,7 +75,7 @@ def test_reranking_enabled_reorders_candidates(monkeypatch: pytest.MonkeyPatch) results = rerank_chunks("query", candidates(), RerankConfig()) assert [result.chunk_id for result in results] == [2, 1, 3] - assert [round(result.score, 3) for result in results] == [0.45, 0.1, 0.0] + assert [round(result.score, 3) for result in results] == [0.78, 0.37, 0.28] assert [result.rerank_score for result in results] == [0.9, 0.1, 0.4] @@ -100,8 +100,8 @@ def test_reranking_cannot_ignore_hybrid_score(monkeypatch: pytest.MonkeyPatch) - results = rerank_chunks("query", candidates, RerankConfig()) assert [result.chunk_id for result in results] == [1, 2] - assert results[0].score == 0.7 - assert results[1].score == 0.0 + assert results[0].score == pytest.approx(0.79) + assert results[1].score == 0.7 assert results[1].rerank_score == 1.0 @@ -129,7 +129,7 @@ def test_malformed_vllm_rerank_json_does_not_crash_search(monkeypatch: pytest.Mo results = rerank_chunks("query", candidates()[:1], RerankConfig()) - assert results[0].score == 0.0 + assert results[0].score == 0.3 def test_vllm_rerank_scores_are_clamped(monkeypatch: pytest.MonkeyPatch) -> None: @@ -147,4 +147,4 @@ def test_vllm_rerank_scores_are_clamped(monkeypatch: pytest.MonkeyPatch) -> None results = rerank_chunks("query", candidates()[:2], RerankConfig()) - assert [result.rerank_score for result in results] == [0.0, 1.0] + assert {result.chunk_id: result.rerank_score for result in results} == {1: 0.0, 2: 1.0} -- 2.54.0 From c5418b50fdd8b048757168532671fb68f71a3c70 Mon Sep 17 00:00:00 2001 From: Richie Cahill Date: Fri, 12 Jun 2026 13:47:43 -0400 Subject: [PATCH 21/28] added proper cache invalidation to load_bm25_corpus --- python/ebook_search/api/bm25_tasks.py | 4 +++- python/ebook_search/bm25_corpus.py | 12 +++--------- tests/test_ebook_search_ui.py | 24 ++++++++++++++++++++++++ 3 files changed, 30 insertions(+), 10 deletions(-) diff --git a/python/ebook_search/api/bm25_tasks.py b/python/ebook_search/api/bm25_tasks.py index a211d45..ff24b85 100644 --- a/python/ebook_search/api/bm25_tasks.py +++ b/python/ebook_search/api/bm25_tasks.py @@ -8,7 +8,7 @@ from typing import TYPE_CHECKING from sqlalchemy.orm import Session -from python.ebook_search.bm25_corpus import refresh_bm25_corpus +from python.ebook_search.bm25_corpus import load_bm25_corpus, refresh_bm25_corpus if TYPE_CHECKING: from fastapi import FastAPI @@ -56,3 +56,5 @@ def refresh_bm25_for_engine(engine: Engine, config: EbookSearchConfig) -> None: """Refresh the BM25 corpus using a SQLAlchemy engine.""" with Session(engine) as session: refresh_bm25_corpus(session, config) + load_bm25_corpus.cache_clear() + logger.info("ebook_bm25_corpus_cache_cleared_after_refresh") diff --git a/python/ebook_search/bm25_corpus.py b/python/ebook_search/bm25_corpus.py index 6c8a15d..2d3904f 100644 --- a/python/ebook_search/bm25_corpus.py +++ b/python/ebook_search/bm25_corpus.py @@ -108,11 +108,10 @@ def refresh_bm25_corpus( ) write_bm25_corpus(index_path, records, manifest) logger.info( - "ebook_bm25_index_refreshed path=%s chunks=%s created_at=%s note=%s", + "ebook_bm25_index_refreshed path=%s chunks=%s created_at=%s", index_path, manifest.chunk_count, manifest.created_at.isoformat(), - "restart_service_to_use_refreshed_bm25_cache", ) return manifest @@ -121,15 +120,10 @@ def refresh_bm25_corpus( def load_bm25_corpus(config: EbookSearchConfig) -> BM25Corpus: """Load the BM25 corpus into memory once per process. - This cache intentionally does not notice later on-disk corpus refreshes. Restart the service after rebuilding the - BM25 corpus for searches to use the new index. + Background refresh tasks clear this cache after rebuilding the on-disk corpus. """ index_path = bm25_index_path(config) - logger.info( - "ebook_bm25_corpus_cache_load path=%s note=%s", - index_path, - "restart_service_after_bm25_refresh", - ) + logger.info("ebook_bm25_corpus_cache_load path=%s", index_path) manifest = read_bm25_manifest(index_path) if manifest is None or not bm25_index_exists(index_path, manifest): msg = f"BM25 corpus is not available: {index_path}" diff --git a/tests/test_ebook_search_ui.py b/tests/test_ebook_search_ui.py index 03b31b2..3de8eac 100644 --- a/tests/test_ebook_search_ui.py +++ b/tests/test_ebook_search_ui.py @@ -5,6 +5,7 @@ from __future__ import annotations from fastapi.testclient import TestClient from sqlalchemy import create_engine +from python.ebook_search.api.bm25_tasks import refresh_bm25_for_engine from python.ebook_search.api.main import create_app from python.ebook_search.config import EbookSearchConfig, RerankConfig from python.ebook_search.embeddings import EmbeddingModelStats @@ -232,6 +233,29 @@ def test_ui_scan_schedules_bm25_refresh_after_database_change(monkeypatch) -> No assert scheduled is True +def test_bm25_refresh_clears_loaded_corpus_cache(monkeypatch) -> None: + refreshed: list[object] = [] + cache_cleared = False + + def fake_refresh_bm25_corpus(session, config): + refreshed.append((session, config)) + + def fake_cache_clear(): + nonlocal cache_cleared + cache_cleared = True + + monkeypatch.setattr("python.ebook_search.api.bm25_tasks.refresh_bm25_corpus", fake_refresh_bm25_corpus) + monkeypatch.setattr("python.ebook_search.api.bm25_tasks.load_bm25_corpus.cache_clear", fake_cache_clear) + engine = create_engine("sqlite+pysqlite:///:memory:", future=True) + config = EbookSearchConfig(rerank=RerankConfig(enabled=False)) + + refresh_bm25_for_engine(engine, config) + + assert len(refreshed) == 1 + assert refreshed[0][1] == config + assert cache_cleared is True + + def test_admin_page_shows_embedding_counts_by_model(monkeypatch) -> None: def fake_embedding_model_stats(_session): return [ -- 2.54.0 From 479191050e040dd76380f2aa4e3de2fe88645a84 Mon Sep 17 00:00:00 2001 From: Richie Cahill Date: Fri, 12 Jun 2026 14:45:10 -0400 Subject: [PATCH 22/28] made fastapi tools --- python/api/main.py | 2 +- python/api/routers/contact.py | 2 +- python/api/routers/views.py | 2 +- python/fastapi_tools/__init__.py | 6 ++++++ python/{api/dependencies.py => fastapi_tools/db.py} | 0 .../{api/middleware.py => fastapi_tools/zstd_middleware.py} | 2 +- 6 files changed, 10 insertions(+), 4 deletions(-) create mode 100644 python/fastapi_tools/__init__.py rename python/{api/dependencies.py => fastapi_tools/db.py} (100%) rename python/{api/middleware.py => fastapi_tools/zstd_middleware.py} (97%) diff --git a/python/api/main.py b/python/api/main.py index 3ac65ba..ce84c5b 100644 --- a/python/api/main.py +++ b/python/api/main.py @@ -9,9 +9,9 @@ import typer import uvicorn from fastapi import FastAPI -from python.api.middleware import ZstdMiddleware from python.api.routers import contact_router, views_router from python.common import configure_logger +from python.fastapi_tools import ZstdMiddleware from python.orm.common import get_postgres_engine logger = logging.getLogger(__name__) diff --git a/python/api/routers/contact.py b/python/api/routers/contact.py index 9aa398d..1cef937 100644 --- a/python/api/routers/contact.py +++ b/python/api/routers/contact.py @@ -9,7 +9,7 @@ from pydantic import BaseModel from sqlalchemy import select from sqlalchemy.orm import selectinload -from python.api.dependencies import DbSession +from python.fastapi_tools.db import DbSession from python.orm.richie.contact import Contact, ContactRelationship, Need, RelationshipType TEMPLATES_DIR = Path(__file__).parent.parent / "templates" diff --git a/python/api/routers/views.py b/python/api/routers/views.py index dc37f83..fdf451e 100644 --- a/python/api/routers/views.py +++ b/python/api/routers/views.py @@ -9,7 +9,7 @@ from fastapi.templating import Jinja2Templates from sqlalchemy import select from sqlalchemy.orm import Session, selectinload -from python.api.dependencies import DbSession +from python.fastapi_tools.db import DbSession from python.orm.richie.contact import Contact, ContactRelationship, Need, RelationshipType TEMPLATES_DIR = Path(__file__).parent.parent / "templates" diff --git a/python/fastapi_tools/__init__.py b/python/fastapi_tools/__init__.py new file mode 100644 index 0000000..d55eb3d --- /dev/null +++ b/python/fastapi_tools/__init__.py @@ -0,0 +1,6 @@ +"""Reusable FastAPI tools.""" + +from python.fastapi_tools.db import DbSession, get_db +from python.fastapi_tools.zstd_middleware import ZstdMiddleware + +__all__ = ["DbSession", "ZstdMiddleware", "get_db"] diff --git a/python/api/dependencies.py b/python/fastapi_tools/db.py similarity index 100% rename from python/api/dependencies.py rename to python/fastapi_tools/db.py diff --git a/python/api/middleware.py b/python/fastapi_tools/zstd_middleware.py similarity index 97% rename from python/api/middleware.py rename to python/fastapi_tools/zstd_middleware.py index f710a66..f273abf 100644 --- a/python/api/middleware.py +++ b/python/fastapi_tools/zstd_middleware.py @@ -1,4 +1,4 @@ -"""Middleware for the FastAPI application.""" +"""Zstd response compression middleware.""" from compression import zstd from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint -- 2.54.0 From ed45051eb581d694db5198f0cd3e30acdee4f138 Mon Sep 17 00:00:00 2001 From: Richie Cahill Date: Fri, 12 Jun 2026 14:46:00 -0400 Subject: [PATCH 23/28] reworked ebook_search routers --- python/ebook_search/api/main.py | 10 ++++++---- python/ebook_search/api/routes/__init__.py | 17 ++++++----------- python/ebook_search/api/routes/admin.py | 9 --------- python/ebook_search/api/routes/page.py | 9 --------- python/ebook_search/api/routes/search.py | 10 +--------- 5 files changed, 13 insertions(+), 42 deletions(-) diff --git a/python/ebook_search/api/main.py b/python/ebook_search/api/main.py index 9be6d99..b6c3a62 100644 --- a/python/ebook_search/api/main.py +++ b/python/ebook_search/api/main.py @@ -14,7 +14,7 @@ from sqlalchemy.orm import Session from python.common import configure_logger from python.ebook_search.api.bm25_tasks import cancel_bm25_refresh -from python.ebook_search.api.routes import register_admin_routes, register_page_routes, register_search_routes +from python.ebook_search.api.routes import admin_router, page_router, search_router from python.ebook_search.api.web import STATIC_DIR from python.ebook_search.bm25_corpus import ensure_bm25_corpus from python.ebook_search.config import load_config @@ -55,9 +55,11 @@ def create_app() -> FastAPI: app.state.config.answer_enabled, len(app.state.config.library_paths), ) - register_page_routes(app) - register_search_routes(app) - register_admin_routes(app) + + app.include_router(admin_router) + app.include_router(page_router) + app.include_router(search_router) + return app diff --git a/python/ebook_search/api/routes/__init__.py b/python/ebook_search/api/routes/__init__.py index a6e49ca..b1fc051 100644 --- a/python/ebook_search/api/routes/__init__.py +++ b/python/ebook_search/api/routes/__init__.py @@ -1,16 +1,11 @@ """EPUB search web route modules.""" -from python.ebook_search.api.routes import admin, page, search - -register_admin_routes = admin.register_admin_routes -register_page_routes = page.register_page_routes -register_search_routes = search.register_search_routes +from python.ebook_search.api.routes.admin import router as admin_router +from python.ebook_search.api.routes.page import router as page_router +from python.ebook_search.api.routes.search import router as search_router __all__ = [ - "admin", - "page", - "register_admin_routes", - "register_page_routes", - "register_search_routes", - "search", + "admin_router", + "page_router", + "search_router", ] diff --git a/python/ebook_search/api/routes/admin.py b/python/ebook_search/api/routes/admin.py index ff83239..4a14875 100644 --- a/python/ebook_search/api/routes/admin.py +++ b/python/ebook_search/api/routes/admin.py @@ -4,7 +4,6 @@ from __future__ import annotations import logging from dataclasses import replace -from typing import TYPE_CHECKING from fastapi import APIRouter, Request from fastapi.responses import HTMLResponse @@ -15,20 +14,12 @@ from python.ebook_search.api.web import templates from python.ebook_search.embeddings import embed_missing_chunks, embedding_model_stats from python.ebook_search.ingest import ingest_configured_paths -if TYPE_CHECKING: - from fastapi import FastAPI - logger = logging.getLogger(__name__) router = APIRouter(prefix="/admin") EMBED_ALL_BATCH_SIZE = 32 -def register_admin_routes(app: FastAPI) -> None: - """Register admin routes on the app.""" - app.include_router(router) - - @router.get("", response_class=HTMLResponse) def admin(request: Request) -> HTMLResponse: """Render the admin page.""" diff --git a/python/ebook_search/api/routes/page.py b/python/ebook_search/api/routes/page.py index 8e48867..b92b881 100644 --- a/python/ebook_search/api/routes/page.py +++ b/python/ebook_search/api/routes/page.py @@ -3,7 +3,6 @@ from __future__ import annotations import logging -from typing import TYPE_CHECKING from fastapi import APIRouter, Request from fastapi.responses import HTMLResponse @@ -13,19 +12,11 @@ from sqlalchemy.orm import Session from python.ebook_search.api.web import templates from python.orm.richie import EbookSource -if TYPE_CHECKING: - from fastapi import FastAPI - logger = logging.getLogger(__name__) router = APIRouter() -def register_page_routes(app: FastAPI) -> None: - """Register page routes on the app.""" - app.include_router(router) - - @router.get("/", response_class=HTMLResponse) def index(request: Request) -> HTMLResponse: """Render the search page.""" diff --git a/python/ebook_search/api/routes/search.py b/python/ebook_search/api/routes/search.py index 77ed022..235dee2 100644 --- a/python/ebook_search/api/routes/search.py +++ b/python/ebook_search/api/routes/search.py @@ -5,7 +5,7 @@ from __future__ import annotations import logging from dataclasses import replace from time import perf_counter -from typing import TYPE_CHECKING, Annotated +from typing import Annotated from fastapi import APIRouter, Form, Request from fastapi.responses import HTMLResponse @@ -15,19 +15,11 @@ from python.ebook_search.api.web import templates from python.ebook_search.search import search_ebooks from python.ebook_search.timing import runtime_step_from_start -if TYPE_CHECKING: - from fastapi import FastAPI - logger = logging.getLogger(__name__) router = APIRouter() -def register_search_routes(app: FastAPI) -> None: - """Register search routes on the app.""" - app.include_router(router) - - @router.post("/search", response_class=HTMLResponse) def search( request: Request, -- 2.54.0 From 51855725a16e20abbb2fd56519b48df91d2db873 Mon Sep 17 00:00:00 2001 From: Richie Cahill Date: Fri, 12 Jun 2026 14:50:22 -0400 Subject: [PATCH 24/28] added vector_engine to fix name postgres name space issue --- python/ebook_search/api/main.py | 2 +- python/orm/common.py | 26 ++++++++++++++++++++++++-- 2 files changed, 25 insertions(+), 3 deletions(-) diff --git a/python/ebook_search/api/main.py b/python/ebook_search/api/main.py index b6c3a62..f9937f5 100644 --- a/python/ebook_search/api/main.py +++ b/python/ebook_search/api/main.py @@ -31,7 +31,7 @@ logger = logging.getLogger(__name__) async def lifespan(app: FastAPI) -> AsyncIterator[None]: """Manage application startup and shutdown resources.""" logger.info("ebook_search_startup") - app.state.engine = get_postgres_engine(name="RICHIE") + app.state.engine = get_postgres_engine(name="RICHIE", vector_engine=True) with Session(app.state.engine) as session: ensure_bm25_corpus(session, app.state.config) try: diff --git a/python/orm/common.py b/python/orm/common.py index 6f86462..1214346 100644 --- a/python/orm/common.py +++ b/python/orm/common.py @@ -31,8 +31,24 @@ def get_connection_info(name: str) -> tuple[str, str, str, str, str | None]: return cast("tuple[str, str, str, str, str | None]", (database, host, port, username, password)) -def get_postgres_engine(*, name: str = "POSTGRES", pool_pre_ping: bool = True) -> Engine: - """Create a SQLAlchemy engine from environment variables.""" +def get_postgres_engine( + *, + name: str = "POSTGRES", + pool_pre_ping: bool = True, + vector_engine: bool = False, +) -> Engine: + """Create a SQLAlchemy engine from environment variables. + + Args: + name (str, optional): The name of the environment variable prefix. Defaults to "POSTGRES". + pool_pre_ping (bool, optional): Whether to ping the database before each connection. Defaults to True. + This fixes the issue of trying to use a conection that has timed out on the database side. + vector_engine (bool, optional): Whether to use the vector search schema. Defaults to False. + This updates the search path the incldued the vecore types and operators. + + Returns: + Engine: The SQLAlchemy engine. + """ database, host, port, username, password = get_connection_info(name) url = URL.create( @@ -44,8 +60,14 @@ def get_postgres_engine(*, name: str = "POSTGRES", pool_pre_ping: bool = True) - database=database, ) + connect_args = {} + # There more better way to do this is with separate PG account and a dedicated vector schema for the vector types + if vector_engine: + connect_args["options"] = "-csearch_path=main,public" + return create_engine( url=url, pool_pre_ping=pool_pre_ping, pool_recycle=1800, + connect_args=connect_args, ) -- 2.54.0 From 66ea18af82fb957221183a9aebe321d01b3ab654 Mon Sep 17 00:00:00 2001 From: Richie Cahill Date: Fri, 12 Jun 2026 14:57:59 -0400 Subject: [PATCH 25/28] added ZstdMiddleware to ebook_search --- python/ebook_search/api/main.py | 2 ++ tests/test_ebook_search_ui.py | 14 ++++++++++++++ 2 files changed, 16 insertions(+) diff --git a/python/ebook_search/api/main.py b/python/ebook_search/api/main.py index f9937f5..a894ca3 100644 --- a/python/ebook_search/api/main.py +++ b/python/ebook_search/api/main.py @@ -18,6 +18,7 @@ from python.ebook_search.api.routes import admin_router, page_router, search_rou from python.ebook_search.api.web import STATIC_DIR from python.ebook_search.bm25_corpus import ensure_bm25_corpus from python.ebook_search.config import load_config +from python.fastapi_tools import ZstdMiddleware from python.orm.common import get_postgres_engine if TYPE_CHECKING: @@ -45,6 +46,7 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]: def create_app() -> FastAPI: """Create the EPUB search web app.""" app = FastAPI(title="EPUB Search", lifespan=lifespan) + app.add_middleware(ZstdMiddleware) app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static") app.state.config = load_config() logger.info( diff --git a/tests/test_ebook_search_ui.py b/tests/test_ebook_search_ui.py index 3de8eac..025d867 100644 --- a/tests/test_ebook_search_ui.py +++ b/tests/test_ebook_search_ui.py @@ -2,6 +2,7 @@ from __future__ import annotations +from compression import zstd from fastapi.testclient import TestClient from sqlalchemy import create_engine @@ -24,6 +25,19 @@ def fake_get_postgres_engine(**_kwargs): return create_engine("sqlite+pysqlite:///:memory:", future=True) +def test_search_page_uses_zstd_when_requested(monkeypatch) -> None: + patch_app_runtime(monkeypatch) + app = create_app() + app.state.config = EbookSearchConfig(rerank=RerankConfig(enabled=False)) + + with TestClient(app) as client: + response = client.get("/", headers={"accept-encoding": "zstd"}) + + assert response.status_code == 200 + assert response.headers["content-encoding"] == "zstd" + assert b"EPUB Search" in zstd.decompress(response.content) + + def test_ui_form_passes_rerank_flag_to_search_handler(monkeypatch) -> None: captured: dict[str, object] = {} -- 2.54.0 From bb3c433b9da77654249fa088b4873cf89980d1a5 Mon Sep 17 00:00:00 2001 From: Richie Cahill Date: Sat, 13 Jun 2026 17:05:32 -0400 Subject: [PATCH 26/28] improved BM25 write --- python/ebook_search/bm25_corpus.py | 90 +++++++++++++++++------- tests/test_ebook_search_core.py | 107 +++++++++++++++++++++++++++-- 2 files changed, 169 insertions(+), 28 deletions(-) diff --git a/python/ebook_search/bm25_corpus.py b/python/ebook_search/bm25_corpus.py index 2d3904f..3a1752b 100644 --- a/python/ebook_search/bm25_corpus.py +++ b/python/ebook_search/bm25_corpus.py @@ -5,7 +5,6 @@ from __future__ import annotations import json import logging import shutil -import tempfile from dataclasses import dataclass from datetime import UTC, datetime from functools import cache @@ -59,13 +58,21 @@ class BM25CorpusUnavailableError(RuntimeError): def bm25_index_path(config: EbookSearchConfig) -> Path: - """Return the configured BM25 index path relative to the current working directory.""" + """Return the configured BM25 index root path relative to the current working directory.""" path = Path(config.bm25_index_dir).expanduser() if path.is_absolute(): return path return Path.cwd() / path +def get_current_bm25_index(index_path: Path) -> Path: + """Return the live BM25 index directory.""" + current_path = index_path / "current" + if current_path.exists() or current_path.is_symlink(): + return current_path + return index_path + + def ensure_bm25_corpus(session: Session, config: EbookSearchConfig) -> None: """Create or refresh the persisted BM25 corpus when it is missing or stale.""" index_path = bm25_index_path(config) @@ -100,13 +107,13 @@ def refresh_bm25_corpus( ) -> BM25Manifest: """Rebuild and persist the BM25 corpus from the current database chunks.""" index_path = bm25_index_path(config) - records = fetch_bm25_corpus_records(session) + records, texts = fetch_bm25_corpus_records(session) manifest = BM25Manifest( created_at=datetime.now(tz=UTC), db_updated_at=db_updated_at if db_updated_at is not None else corpus_last_updated_at(session), chunk_count=len(records), ) - write_bm25_corpus(index_path, records, manifest) + write_bm25_corpus(index_path, records, texts, manifest) logger.info( "ebook_bm25_index_refreshed path=%s chunks=%s created_at=%s", index_path, @@ -123,7 +130,8 @@ def load_bm25_corpus(config: EbookSearchConfig) -> BM25Corpus: Background refresh tasks clear this cache after rebuilding the on-disk corpus. """ index_path = bm25_index_path(config) - logger.info("ebook_bm25_corpus_cache_load path=%s", index_path) + active_index_path = get_current_bm25_index(index_path) + logger.info("ebook_bm25_corpus_cache_load path=%s active_path=%s", index_path, active_index_path) manifest = read_bm25_manifest(index_path) if manifest is None or not bm25_index_exists(index_path, manifest): msg = f"BM25 corpus is not available: {index_path}" @@ -131,7 +139,7 @@ def load_bm25_corpus(config: EbookSearchConfig) -> BM25Corpus: if manifest.chunk_count == 0: return BM25Corpus(retriever=None, records=(), manifest=manifest) - retriever = bm25s.BM25.load(index_path, load_corpus=True, mmap=True) + retriever = bm25s.BM25.load(active_index_path, load_corpus=True, mmap=True) records = tuple(dict(record) for record in retriever.corpus) return BM25Corpus(retriever=retriever, records=records, manifest=manifest) @@ -156,8 +164,12 @@ def score_bm25_corpus(query: str, corpus: BM25Corpus, *, limit: int) -> list[tup return results -def fetch_bm25_corpus_records(session: Session) -> list[dict[str, object]]: - """Fetch BM25 corpus records from the database.""" +def fetch_bm25_corpus_records(session: Session) -> tuple[list[dict[str, object]], list[str]]: + """Fetch persistable BM25 corpus records and their matching index texts from the database. + + search_text is only needed to build the index, so it is returned separately instead of + being persisted into the corpus records, which would double the corpus size. + """ statement = ( select( EbookChunk.id.label("chunk_id"), @@ -173,7 +185,13 @@ def fetch_bm25_corpus_records(session: Session) -> list[dict[str, object]]: .outerjoin(EbookChapter, EbookChapter.id == EbookChunk.chapter_id) .order_by(EbookChunk.id) ) - return [dict(row) for row in session.execute(statement).mappings()] + records: list[dict[str, object]] = [] + texts: list[str] = [] + for row in session.execute(statement).mappings(): + record = dict(row) + texts.append(str(record.pop("bm25_text"))) + records.append(record) + return records, texts def corpus_last_updated_at(session: Session) -> datetime | None: @@ -186,28 +204,42 @@ def corpus_last_updated_at(session: Session) -> datetime | None: return session.scalar(select(func.max(update_times.c.updated))) -def write_bm25_corpus(index_path: Path, records: list[dict[str, object]], manifest: BM25Manifest) -> None: - """Write a BM25 corpus and manifest atomically.""" - index_path.parent.mkdir(parents=True, exist_ok=True) - temp_path = Path(tempfile.mkdtemp(prefix=f"{index_path.name}.", dir=index_path.parent)) +def write_bm25_corpus( + index_path: Path, + records: list[dict[str, object]], + texts: list[str], + manifest: BM25Manifest, +) -> None: + """Write a BM25 corpus generation and publish it through the current symlink.""" + index_path.mkdir(parents=True, exist_ok=True) + + generations_path = index_path / "generations" + generations_path.mkdir(exist_ok=True) + + generation_path = next_bm25_generation_path(generations_path, manifest.created_at) + current_path = index_path / "current" + next_current_path = index_path / f".current.{generation_path.name}.tmp" try: + generation_path.mkdir() + + # Empty corpora publish a manifest-only generation so startup succeeds before any chunks exist. if records: retriever = bm25s.BM25() - texts = [str(record["bm25_text"]) for record in records] retriever.index(bm25s.tokenize(texts, show_progress=False), show_progress=False) - retriever.save(temp_path, corpus=records, show_progress=False) - write_bm25_manifest(temp_path, manifest) - if index_path.exists(): - shutil.rmtree(index_path) - temp_path.rename(index_path) + retriever.save(generation_path, corpus=records, show_progress=False) + write_bm25_manifest(generation_path, manifest) + next_current_path.unlink(missing_ok=True) + next_current_path.symlink_to(generation_path, target_is_directory=True) + next_current_path.replace(current_path) except Exception: - shutil.rmtree(temp_path, ignore_errors=True) + next_current_path.unlink(missing_ok=True) + shutil.rmtree(generation_path, ignore_errors=True) raise def read_bm25_manifest(index_path: Path) -> BM25Manifest | None: """Read the BM25 manifest if it exists and is valid.""" - manifest_path = index_path / MANIFEST_NAME + manifest_path = get_current_bm25_index(index_path) / MANIFEST_NAME if not manifest_path.exists(): return None body = json.loads(manifest_path.read_text(encoding="utf-8")) @@ -230,8 +262,20 @@ def write_bm25_manifest(index_path: Path, manifest: BM25Manifest) -> None: def bm25_index_exists(index_path: Path, manifest: BM25Manifest | None) -> bool: """Return whether a usable persisted BM25 index exists.""" - if manifest is None or not index_path.is_dir(): + active_index_path = get_current_bm25_index(index_path) + if manifest is None or not active_index_path.is_dir(): return False if manifest.chunk_count == 0: return True - return all((index_path / file_name).exists() for file_name in REQUIRED_INDEX_FILES) + return all((active_index_path / file_name).exists() for file_name in REQUIRED_INDEX_FILES) + + +def next_bm25_generation_path(generations_path: Path, created_at: datetime) -> Path: + """Return an unused dated BM25 generation path.""" + base_name = created_at.astimezone(UTC).strftime("%Y%m%dT%H%M%S.%fZ") + generation_path = generations_path / base_name + suffix = 1 + while generation_path.exists(): + generation_path = generations_path / f"{base_name}.{suffix}" + suffix += 1 + return generation_path diff --git a/tests/test_ebook_search_core.py b/tests/test_ebook_search_core.py index c4b05cc..76345ec 100644 --- a/tests/test_ebook_search_core.py +++ b/tests/test_ebook_search_core.py @@ -2,6 +2,7 @@ from __future__ import annotations +import logging from dataclasses import replace from datetime import UTC, datetime from os import environ @@ -21,6 +22,9 @@ from python.ebook_search.bm25_corpus import ( ensure_bm25_corpus, fetch_bm25_corpus_records, load_bm25_corpus, + read_bm25_manifest, + score_bm25_corpus, + write_bm25_corpus, ) from python.ebook_search.config import EbookSearchConfig, RerankConfig, load_config, normalize_embedding_model from python.ebook_search.embeddings import MODEL_DIMENSIONS, ensure_embedding_models @@ -123,9 +127,11 @@ def test_bm25_corpus_uses_existing_search_text_without_duplicate_metadata() -> N ) session.commit() - records = fetch_bm25_corpus_records(session) + records, texts = fetch_bm25_corpus_records(session) - assert records[0]["bm25_text"] == "Book Author Chapter content" + assert texts == ["Book Author Chapter content"] + assert records[0]["chunk_id"] == 1 + assert "bm25_text" not in records[0] def test_reciprocal_rank_fusion_marks_hybrid_source() -> None: @@ -227,15 +233,106 @@ def test_bm25_candidates_scores_whole_corpus(monkeypatch) -> None: assert [result.bm25_score for result in results] == [1.5] -def test_bm25_candidates_raises_when_corpus_is_unavailable(monkeypatch) -> None: +def test_bm25_candidates_returns_empty_when_corpus_is_unavailable(monkeypatch, caplog) -> None: def fake_load_bm25_corpus(_config): raise BM25CorpusUnavailableError monkeypatch.setattr("python.ebook_search.search.load_bm25_corpus", fake_load_bm25_corpus) config = EbookSearchConfig(rerank=RerankConfig(enabled=False)) - with pytest.raises(BM25CorpusUnavailableError): - bm25_candidates("high", config) + with caplog.at_level(logging.WARNING): + results = bm25_candidates("high", config) + + assert results == [] + assert "ebook_bm25_index_unavailable_skipping" in caplog.text + + +def test_write_bm25_corpus_publishes_dated_generation(tmp_path) -> None: + index_path = tmp_path / "bm25" + index_path.mkdir() + generations_path = index_path / "generations" + generations_path.mkdir() + old_generation = generations_path / "20260101T000000.000000Z" + old_generation.mkdir() + (old_generation / "sentinel").write_text("old", encoding="utf-8") + (index_path / "current").symlink_to(Path("generations") / old_generation.name, target_is_directory=True) + manifest = BM25Manifest( + created_at=datetime(2026, 6, 12, 1, 2, 3, 456789, tzinfo=UTC), + db_updated_at=None, + chunk_count=0, + ) + + write_bm25_corpus(index_path, [], [], manifest) + + current_path = index_path / "current" + assert current_path.is_symlink() + assert current_path.readlink() == generations_path / "20260612T010203.456789Z" + assert old_generation.is_dir() + assert (old_generation / "sentinel").read_text(encoding="utf-8") == "old" + assert (generations_path / "20260612T010203.456789Z").is_dir() + assert read_bm25_manifest(index_path) == manifest + + +def test_write_bm25_corpus_keeps_current_generation_when_publish_fails(monkeypatch, tmp_path) -> None: + index_path = tmp_path / "bm25" + index_path.mkdir() + generations_path = index_path / "generations" + generations_path.mkdir() + old_generation = generations_path / "20260101T000000.000000Z" + old_generation.mkdir() + (old_generation / "sentinel").write_text("old", encoding="utf-8") + current_path = index_path / "current" + current_path.symlink_to(Path("generations") / old_generation.name, target_is_directory=True) + original_replace = Path.replace + + def fail_current_replace(self, target): + if self.parent == index_path and self.name.startswith(".current.") and target == current_path: + msg = "current publish failed" + raise OSError(msg) + return original_replace(self, target) + + monkeypatch.setattr(Path, "replace", fail_current_replace) + manifest = BM25Manifest( + created_at=datetime(2026, 6, 12, 1, 2, 3, 456789, tzinfo=UTC), + db_updated_at=None, + chunk_count=0, + ) + + with pytest.raises(OSError, match="current publish failed"): + write_bm25_corpus(index_path, [], [], manifest) + + assert current_path.readlink() == Path("generations") / old_generation.name + assert (old_generation / "sentinel").read_text(encoding="utf-8") == "old" + assert not (generations_path / "20260612T010203.456789Z").exists() + + +def test_load_bm25_corpus_uses_current_generation(tmp_path) -> None: + load_bm25_corpus.cache_clear() + index_path = tmp_path / "bm25" + manifest = BM25Manifest( + created_at=datetime(2026, 6, 12, 1, 2, 3, 456789, tzinfo=UTC), + db_updated_at=None, + chunk_count=1, + ) + record = { + "chunk_id": 2, + "text": "cached", + "source_title": "B", + "source_author": None, + "chapter_title": None, + "page_label": None, + } + write_bm25_corpus(index_path, [record], ["cached phrase"], manifest) + config = EbookSearchConfig(rerank=RerankConfig(enabled=False), bm25_index_dir=str(index_path)) + + try: + corpus = load_bm25_corpus(config) + finally: + load_bm25_corpus.cache_clear() + + assert corpus.manifest == manifest + assert corpus.records[0]["chunk_id"] == 2 + assert score_bm25_corpus("cached", corpus, limit=10) def test_load_bm25_corpus_caches_disk_load(monkeypatch, tmp_path) -> None: -- 2.54.0 From 5e2252641d136fb50cfc9974c6a43d43f5063ff7 Mon Sep 17 00:00:00 2001 From: Richie Cahill Date: Sat, 13 Jun 2026 20:14:20 -0400 Subject: [PATCH 27/28] added a index for the VEctor DB --- ...ook_embedding_cosine_index_c460105682d2.py | 54 +++++++++++++++++++ python/orm/richie/ebook.py | 12 ++++- tests/test_ebook_search_core.py | 18 ++++++- 3 files changed, 81 insertions(+), 3 deletions(-) create mode 100644 python/alembic/richie/versions/2026_06_13-add_1024_ebook_embedding_cosine_index_c460105682d2.py diff --git a/python/alembic/richie/versions/2026_06_13-add_1024_ebook_embedding_cosine_index_c460105682d2.py b/python/alembic/richie/versions/2026_06_13-add_1024_ebook_embedding_cosine_index_c460105682d2.py new file mode 100644 index 0000000..8aadfa3 --- /dev/null +++ b/python/alembic/richie/versions/2026_06_13-add_1024_ebook_embedding_cosine_index_c460105682d2.py @@ -0,0 +1,54 @@ +"""add 1024 ebook embedding cosine index. + +Revision ID: c460105682d2 +Revises: 2db132cace1a +Create Date: 2026-06-13 19:53:45.680289 + +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from alembic import op + +from python.orm import RichieBase + +if TYPE_CHECKING: + from collections.abc import Sequence + +# revision identifiers, used by Alembic. +revision: str = "c460105682d2" +down_revision: str | None = "2db132cace1a" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + +schema = RichieBase.schema_name + + +def upgrade() -> None: + """Upgrade.""" + # ### commands auto generated by Alembic - please adjust! ### + op.create_index( + "ix_ebook_chunk_embedding_1024_embedding_cosine", + "ebook_chunk_embedding_1024", + ["embedding"], + unique=False, + schema=schema, + postgresql_using="hnsw", + postgresql_ops={"embedding": "vector_cosine_ops"}, + ) + # ### end Alembic commands ### + + +def downgrade() -> None: + """Downgrade.""" + # ### commands auto generated by Alembic - please adjust! ### + op.drop_index( + "ix_ebook_chunk_embedding_1024_embedding_cosine", + table_name="ebook_chunk_embedding_1024", + schema=schema, + postgresql_using="hnsw", + postgresql_ops={"embedding": "vector_cosine_ops"}, + ) + # ### end Alembic commands ### diff --git a/python/orm/richie/ebook.py b/python/orm/richie/ebook.py index 9c1e4ad..8e32409 100644 --- a/python/orm/richie/ebook.py +++ b/python/orm/richie/ebook.py @@ -5,7 +5,7 @@ from __future__ import annotations from datetime import datetime from pgvector.sqlalchemy import Vector -from sqlalchemy import BigInteger, Boolean, DateTime, ForeignKey, String, UniqueConstraint +from sqlalchemy import BigInteger, Boolean, DateTime, ForeignKey, Index, String, UniqueConstraint from sqlalchemy.orm import Mapped, mapped_column, relationship from python.orm.richie.base import TableBase, TableBaseBig @@ -101,7 +101,15 @@ class EbookChunkEmbedding1024(TableBaseBig): """1024-dimensional chunk embedding.""" __tablename__ = "ebook_chunk_embedding_1024" - __table_args__ = (UniqueConstraint("chunk_id", "model_id"),) + __table_args__ = ( + UniqueConstraint("chunk_id", "model_id"), + Index( + "ix_ebook_chunk_embedding_1024_embedding_cosine", + "embedding", + postgresql_using="hnsw", + postgresql_ops={"embedding": "vector_cosine_ops"}, + ), + ) chunk_id: Mapped[int] = mapped_column(ForeignKey("main.ebook_chunk.id", ondelete="CASCADE")) model_id: Mapped[int] = mapped_column(ForeignKey("main.ebook_embedding_model.id", ondelete="CASCADE")) diff --git a/tests/test_ebook_search_core.py b/tests/test_ebook_search_core.py index 76345ec..08b1ed8 100644 --- a/tests/test_ebook_search_core.py +++ b/tests/test_ebook_search_core.py @@ -38,7 +38,14 @@ from python.ebook_search.search import ( search_ebooks, ) from python.ebook_search.timing import RuntimeStep -from python.orm.richie import EbookChapter, EbookChunk, EbookEmbeddingModel, EbookSource, RichieBase +from python.orm.richie import ( + EbookChapter, + EbookChunk, + EbookChunkEmbedding1024, + EbookEmbeddingModel, + EbookSource, + RichieBase, +) def test_chunk_text_uses_overlap() -> None: @@ -464,6 +471,15 @@ def test_ensure_embedding_models_registers_service_names() -> None: ] +def test_1024_embedding_table_has_cosine_hnsw_index() -> None: + indexes = {index.name: index for index in EbookChunkEmbedding1024.__table__.indexes} + index = indexes["ix_ebook_chunk_embedding_1024_embedding_cosine"] + + assert [column.name for column in index.columns] == ["embedding"] + assert index.dialect_options["postgresql"]["using"] == "hnsw" + assert index.dialect_options["postgresql"]["ops"] == {"embedding": "vector_cosine_ops"} + + def test_embedding_model_aliases_normalize_to_provider_names() -> None: assert normalize_embedding_model() == "qwen3-embedding-0.6b" -- 2.54.0 From 2c366e581db30dd06cde03fbd908fb6c2616aab5 Mon Sep 17 00:00:00 2001 From: Richie Cahill Date: Sat, 13 Jun 2026 20:46:17 -0400 Subject: [PATCH 28/28] opning ports for testing --- systems/bob/default.nix | 2 ++ systems/jeeves/networking.nix | 3 +++ 2 files changed, 5 insertions(+) diff --git a/systems/bob/default.nix b/systems/bob/default.nix index 442d8f0..d5eb01e 100644 --- a/systems/bob/default.nix +++ b/systems/bob/default.nix @@ -32,6 +32,8 @@ enable = true; allowedTCPPorts = [ 8000 + 8001 + 8002 ]; }; networkmanager.enable = true; diff --git a/systems/jeeves/networking.nix b/systems/jeeves/networking.nix index a27aa7e..bacb5fa 100644 --- a/systems/jeeves/networking.nix +++ b/systems/jeeves/networking.nix @@ -17,6 +17,9 @@ allowedTCPPorts = [ ]; allowedUDPPorts = [ ]; }; + allowedTCPPorts = [ + 8070 + ]; }; useNetworkd = true; }; -- 2.54.0