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 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 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/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/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/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.""" 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/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..ff24b85 --- /dev/null +++ b/python/ebook_search/api/bm25_tasks.py @@ -0,0 +1,60 @@ +"""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 load_bm25_corpus, 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) + load_bm25_corpus.cache_clear() + logger.info("ebook_bm25_corpus_cache_cleared_after_refresh") diff --git a/python/ebook_search/api/main.py b/python/ebook_search/api/main.py new file mode 100644 index 0000000..a894ca3 --- /dev/null +++ b/python/ebook_search/api/main.py @@ -0,0 +1,79 @@ +"""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 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 +from python.fastapi_tools import ZstdMiddleware +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", vector_engine=True) + 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.add_middleware(ZstdMiddleware) + 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), + ) + + app.include_router(admin_router) + app.include_router(page_router) + app.include_router(search_router) + + 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..b1fc051 --- /dev/null +++ b/python/ebook_search/api/routes/__init__.py @@ -0,0 +1,11 @@ +"""EPUB search web route modules.""" + +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_router", + "page_router", + "search_router", +] diff --git a/python/ebook_search/api/routes/admin.py b/python/ebook_search/api/routes/admin.py new file mode 100644 index 0000000..4a14875 --- /dev/null +++ b/python/ebook_search/api/routes/admin.py @@ -0,0 +1,107 @@ +"""Admin routes for the EPUB search web UI.""" + +from __future__ import annotations + +import logging +from dataclasses import replace + +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 + +logger = logging.getLogger(__name__) + +router = APIRouter(prefix="/admin") +EMBED_ALL_BATCH_SIZE = 32 + + +@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..b92b881 --- /dev/null +++ b/python/ebook_search/api/routes/page.py @@ -0,0 +1,57 @@ +"""Page routes for the EPUB search web UI.""" + +from __future__ import annotations + +import logging + +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 + +logger = logging.getLogger(__name__) + +router = APIRouter() + + +@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..235dee2 --- /dev/null +++ b/python/ebook_search/api/routes/search.py @@ -0,0 +1,58 @@ +"""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 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 + +logger = logging.getLogger(__name__) + +router = APIRouter() + + +@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) diff --git a/python/ebook_search/bm25_corpus.py b/python/ebook_search/bm25_corpus.py new file mode 100644 index 0000000..3a1752b --- /dev/null +++ b/python/ebook_search/bm25_corpus.py @@ -0,0 +1,281 @@ +"""Persisted BM25 corpus management.""" + +from __future__ import annotations + +import json +import logging +import shutil +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 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) + 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, 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, texts, manifest) + logger.info( + "ebook_bm25_index_refreshed path=%s chunks=%s created_at=%s", + index_path, + manifest.chunk_count, + manifest.created_at.isoformat(), + ) + return manifest + + +@cache +def load_bm25_corpus(config: EbookSearchConfig) -> BM25Corpus: + """Load the BM25 corpus into memory once per process. + + Background refresh tasks clear this cache after rebuilding the on-disk corpus. + """ + index_path = bm25_index_path(config) + 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}" + raise BM25CorpusUnavailableError(msg) + if manifest.chunk_count == 0: + return BM25Corpus(retriever=None, records=(), manifest=manifest) + + 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) + + +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) -> 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"), + 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"), + 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) + ) + 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: + """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]], + 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() + retriever.index(bm25s.tokenize(texts, show_progress=False), show_progress=False) + 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: + 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 = 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")) + 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.""" + 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((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/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) diff --git a/python/ebook_search/embeddings.py b/python/ebook_search/embeddings.py new file mode 100644 index 0000000..f542e2b --- /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.llm_interface 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) 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() 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 diff --git a/python/ebook_search/rerank.py b/python/ebook_search/rerank.py new file mode 100644 index 0000000..5075601 --- /dev/null +++ b/python/ebook_search/rerank.py @@ -0,0 +1,129 @@ +"""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__) +RERANK_SCORE_WEIGHT = 0.7 +HYBRID_SCORE_WEIGHT = 0.3 + + +@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_WEIGHT * rerank_score) + (HYBRID_SCORE_WEIGHT * 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) diff --git a/python/ebook_search/search.py b/python/ebook_search/search.py new file mode 100644 index 0000000..5c179db --- /dev/null +++ b/python/ebook_search/search.py @@ -0,0 +1,377 @@ +"""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] = [] + 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, + query, + bm25_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, + 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( + timed_result, + "Embedding + vector search", + vector_candidates, + engine, + vector_query, + config, + ) + bm25_future = executor.submit( + timed_result, + "BM25 search", + bm25_candidates, + bm25_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 natural-language 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) 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 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, ) 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..8e32409 --- /dev/null +++ b/python/orm/richie/ebook.py @@ -0,0 +1,138 @@ +"""EPUB search models.""" + +from __future__ import annotations + +from datetime import datetime + +from pgvector.sqlalchemy import Vector +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 + + +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"), + 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")) + 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)) 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/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/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; }; 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" 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" 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" diff --git a/tests/test_ebook_search_core.py b/tests/test_ebook_search_core.py new file mode 100644 index 0000000..08b1ed8 --- /dev/null +++ b/tests/test_ebook_search_core.py @@ -0,0 +1,536 @@ +"""Tests for EPUB search core helpers.""" + +from __future__ import annotations + +import logging +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, + 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 +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 ( + EbookChapter, + EbookChunk, + EbookChunkEmbedding1024, + 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_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, texts = fetch_bm25_corpus_records(session) + + 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: + 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 == "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)] + + 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, "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] + + +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_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 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: + 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_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" + + 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..db53333 --- /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.78, 0.37, 0.28] + 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 == pytest.approx(0.79) + assert results[1].score == 0.7 + 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.3 + + +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.chunk_id: result.rerank_score for result in results} == {1: 0.0, 2: 1.0} diff --git a/tests/test_ebook_search_ui.py b/tests/test_ebook_search_ui.py new file mode 100644 index 0000000..025d867 --- /dev/null +++ b/tests/test_ebook_search_ui.py @@ -0,0 +1,303 @@ +"""Tests for EPUB search HTMX routes.""" + +from __future__ import annotations + +from compression import zstd +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 +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_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] = {} + + 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_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 [ + 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