Compare commits

..
Author SHA1 Message Date
Richie 247e951a27 Remove legacy HTML templates and API service configuration for contact management system
treefmt / nix fmt (pull_request) Successful in 5s
pytest / pytest (pull_request) Successful in 32s
test ebook search / test-ebook-search (pull_request) Successful in 35s
build_systems / build-brain (pull_request) Successful in 50s
build_systems / build-bob (pull_request) Successful in 51s
build_systems / build-rhapsody-in-green (pull_request) Successful in 1m5s
build_systems / build-jeeves (pull_request) Successful in 2m26s
treefmt / nix fmt (push) Successful in 5s
pytest / pytest (push) Successful in 29s
test ebook search / test-ebook-search (push) Successful in 30s
build_systems / build-brain (push) Successful in 36s
build_systems / build-bob (push) Successful in 36s
build_systems / build-rhapsody-in-green (push) Successful in 49s
build_systems / build-jeeves (push) Successful in 2m9s
- Deleted base template (base.html) and all related contact templates (contact_detail.html, contact_form.html, contact_list.html, graph.html, need_list.html).
- Removed partial templates for managing contacts and needs (contact_table.html, manage_needs.html, manage_relationships.html, need_items.html).
- Eliminated contact API service configuration (contact_api.nix) from the NixOS setup.
2026-07-15 15:28:15 -04:00
Richie 8c3de690c9 feat(ebook-search): set up uv cache management and pruning service
treefmt / nix fmt (pull_request) Successful in 5s
pytest / pytest (pull_request) Successful in 27s
test ebook search / test-ebook-search (pull_request) Successful in 34s
build_systems / build-brain (pull_request) Successful in 48s
build_systems / build-bob (pull_request) Successful in 50s
build_systems / build-rhapsody-in-green (pull_request) Successful in 1m3s
build_systems / build-jeeves (pull_request) Successful in 2m14s
treefmt / nix fmt (push) Successful in 4s
test ebook search / test-ebook-search (push) Successful in 27s
pytest / pytest (push) Successful in 30s
build_systems / build-brain (push) Successful in 34s
build_systems / build-bob (push) Successful in 37s
build_systems / build-jeeves (push) Successful in 2m9s
build_systems / build-rhapsody-in-green (push) Successful in 48s
2026-07-12 17:26:55 -04:00
Richie 91575e3ab2 feat(ebook-search): update test workflow to use nix for dependency management 2026-07-12 17:10:47 -04:00
Richie 4799054601 ci(ebook-search): run tests via uv on the runner instead of docker
build_systems / build-brain (pull_request) Successful in 44s
build_systems / build-jeeves (pull_request) Successful in 2m54s
test ebook search / test-ebook-search (pull_request) Failing after 19s
treefmt / nix fmt (pull_request) Successful in 5s
pytest / pytest (pull_request) Successful in 27s
build_systems / build-bob (pull_request) Successful in 47s
build_systems / build-rhapsody-in-green (pull_request) Successful in 56s
The gitea runner containers have no docker access, so build the test
env with uv from the existing lockfile and run pytest directly:

- test_ebook_search workflow: uv sync --locked + uv run pytest, with
  UV_PYTHON_DOWNLOADS=never so uv uses the nix-provided python 3.14
- add uv to the runner hostPackages (needs a jeeves rebuild to apply)
- ignore nested **/.venv in .dockerignore (uv sync creates one in
  python/ebook_search/docker)
- document the uv test commands in the docker README; the docker test
  image remains for validating the image itself
2026-07-12 16:40:56 -04:00
Richie e9636d19de feat(ebook-search): add containerized test environment
pytest / pytest (pull_request) Successful in 26s
build_systems / build-brain (pull_request) Successful in 45s
build_systems / build-bob (pull_request) Successful in 45s
build_systems / build-rhapsody-in-green (pull_request) Successful in 57s
build_systems / build-jeeves (pull_request) Successful in 2m17s
treefmt / nix fmt (pull_request) Successful in 5s
test ebook search / test-ebook-search (pull_request) Failing after 4s
2026-07-12 16:12:40 -04:00
Richie 2518ec8551 feat(python-env): remove unused dependencies and clean up package list 2026-07-12 16:12:40 -04:00
Richie f5ec88f1e5 feat(ebook): add Docker packaging and lifecycle tooling
Add a self-contained docker/ package for running the ebook search app
against the existing Postgres database on jeeves:

- Dockerfile: python:3.14-slim image, non-root user, runs the FastAPI
  app on port 8070
- docker-compose.yml: service definition with library volume mount,
  BM25 index volume, .env loading, and a /health healthcheck
- containers.py: Typer CLI (ebook-search-containers) for build/start/
  stop/restart/logs/ps lifecycle management
- README.md: usage and configuration docs
2026-07-12 16:12:40 -04:00
Richie f4c4b11ff8 updated dependencies and added .dockerignore 2026-07-12 16:12:40 -04:00
Richie 37f37c41ac feat(common): add get_repo_dir function and corresponding tests
treefmt / nix fmt (pull_request) Successful in 5s
pytest / pytest (pull_request) Successful in 24s
build_systems / build-brain (pull_request) Successful in 44s
build_systems / build-bob (pull_request) Successful in 44s
build_systems / build-rhapsody-in-green (pull_request) Successful in 57s
build_systems / build-jeeves (pull_request) Successful in 2m17s
treefmt / nix fmt (push) Successful in 5s
build_systems / build-brain (push) Successful in 8s
pytest / pytest (push) Successful in 23s
build_systems / build-bob (push) Successful in 30s
build_systems / build-rhapsody-in-green (push) Successful in 42s
build_systems / build-jeeves (push) Successful in 2m2s
2026-07-12 14:24:05 -04:00
Richie 80a521f297 refactor: extract signal_alert into its own module
treefmt / nix fmt (pull_request) Successful in 5s
pytest / pytest (pull_request) Successful in 24s
build_systems / build-bob (pull_request) Successful in 45s
build_systems / build-brain (pull_request) Successful in 45s
build_systems / build-rhapsody-in-green (pull_request) Successful in 58s
build_systems / build-jeeves (pull_request) Successful in 2m15s
treefmt / nix fmt (push) Successful in 5s
build_systems / build-jeeves (push) Successful in 9s
pytest / pytest (push) Successful in 24s
build_systems / build-brain (push) Successful in 30s
build_systems / build-bob (push) Successful in 32s
build_systems / build-rhapsody-in-green (push) Successful in 44s
Move signal_alert out of python/common.py into a dedicated
python/signal_alert.py module and update its importers
(validate_system.py, snapshot_manager.py) to the new path.

Relocate the signal_alert tests from tests/test_common.py into
tests/test_signal_alert.py, repatching python.signal_alert.logger and
python.signal_alert.Apprise to match the new module.
2026-07-12 13:16:28 -04:00
Richie 4a410dbdf8 feat(postgres): add trust authentication for richie on 172.16.0.0/12
treefmt / nix fmt (pull_request) Successful in 5s
pytest / pytest (pull_request) Successful in 23s
build_systems / build-brain (pull_request) Successful in 44s
build_systems / build-bob (pull_request) Successful in 45s
build_systems / build-jeeves (pull_request) Successful in 2m20s
build_systems / build-rhapsody-in-green (pull_request) Successful in 56s
treefmt / nix fmt (push) Successful in 5s
build_systems / build-rhapsody-in-green (push) Successful in 11s
pytest / pytest (push) Successful in 24s
build_systems / build-brain (push) Successful in 30s
build_systems / build-bob (push) Successful in 33s
build_systems / build-jeeves (push) Successful in 2m7s
2026-07-12 13:06:44 -04:00
Richie 138b79bf97 feat(flake): update lock file
treefmt / nix fmt (pull_request) Successful in 5s
pytest / pytest (pull_request) Successful in 1m12s
build_systems / build-brain (pull_request) Successful in 2m46s
build_systems / build-bob (pull_request) Successful in 3m4s
build_systems / build-rhapsody-in-green (pull_request) Successful in 7m42s
build_systems / build-jeeves (pull_request) Successful in 7m44s
treefmt / nix fmt (push) Successful in 5s
pytest / pytest (push) Successful in 23s
build_systems / build-brain (push) Successful in 30s
build_systems / build-bob (push) Successful in 32s
build_systems / build-rhapsody-in-green (push) Successful in 45s
build_systems / build-jeeves (push) Successful in 2m10s
2026-07-12 12:49:57 -04:00
89 changed files with 1791 additions and 8862 deletions
+13
View File
@@ -3,6 +3,19 @@
.mypy_cache
.pytest_cache
.ruff_cache
.venv
**/.venv
.env
.cache
.claude
.coverage
.vscode
.stfolder
.literotica_data
esphome
htmlcov
data
ebooks
__pycache__
**/__pycache__
*.pyc
+26
View File
@@ -0,0 +1,26 @@
name: test ebook search
on:
workflow_dispatch:
push:
branches:
- main
pull_request:
branches:
- main
env:
UV_PYTHON_DOWNLOADS: never
UV_CACHE_DIR: /var/cache/uv
UV_LINK_MODE: copy
jobs:
test-ebook-search:
runs-on: self-hosted
steps:
- uses: actions/checkout@v4
- name: Install dependencies
run: nix develop .#devShells.x86_64-linux.ebook-search -c uv sync --locked --project python/ebook_search/docker
- name: Run ebook search tests
run: nix develop .#devShells.x86_64-linux.ebook-search -c uv run --project python/ebook_search/docker --no-sync pytest tests/ebook_search --override-ini addopts="-n auto -ra"
-6
View File
@@ -71,7 +71,6 @@
"ehci",
"emerg",
"endlessh",
"ents",
"errorlens",
"esbenp",
"esphome",
@@ -173,8 +172,6 @@
"Networkd",
"networkmanager",
"newtabpage",
"ngram",
"ngrams",
"nixfmt",
"nixos",
"nixpkgs",
@@ -301,9 +298,7 @@
"uiprotect",
"uitour",
"unifi",
"unjudged",
"unrar",
"unstorable",
"unsubmitted",
"uptimekuma",
"urlbar",
@@ -331,7 +326,6 @@
"xcursorgen",
"xdist",
"xhci",
"yake",
"yazi",
"yubikey",
"yubioath",
Generated
+18 -18
View File
@@ -8,11 +8,11 @@
},
"locked": {
"dir": "pkgs/firefox-addons",
"lastModified": 1782964936,
"narHash": "sha256-wXEBDr7/dFQYhVpDwCKc9fkrYQQE4x0bdirX1bsLBGA=",
"lastModified": 1783828963,
"narHash": "sha256-eTytzcUJCaDUZ3/9EF0+V3fvlikQMQBwiX1Sx4Gy+No=",
"owner": "rycee",
"repo": "nur-expressions",
"rev": "64feee871e0373dd6121e412c3fb12e372d1bfb5",
"rev": "8d61e9afde605cd6c22dab68b83d7a71f0a6c5b2",
"type": "gitlab"
},
"original": {
@@ -29,11 +29,11 @@
]
},
"locked": {
"lastModified": 1783005591,
"narHash": "sha256-NcLHV5uBAeggDUE2wPbKszjfyaSLsoqaYt7izOphkZw=",
"lastModified": 1783823409,
"narHash": "sha256-OI4IkRjRXa1e7hYmCGJDPDq5H/kPwhsyoS80cNUF9fI=",
"owner": "nix-community",
"repo": "home-manager",
"rev": "f469c79b955609d6a8fdd9e689be76a93b1621d7",
"rev": "7566825d4652a1b885bd4ce65bd9e8def432fec9",
"type": "github"
},
"original": {
@@ -47,11 +47,11 @@
"nixpkgs": "nixpkgs"
},
"locked": {
"lastModified": 1782562157,
"narHash": "sha256-a7+T6QSeowynwZ1ZJJbP8T8ntAytvrui8kFGJmIZt2c=",
"lastModified": 1783792734,
"narHash": "sha256-50rvY9GdFvpYDcMLcD/4cWSi0hVxArT5wsGlVsHy8eY=",
"owner": "nixos",
"repo": "nixos-hardware",
"rev": "a9cf7546a938c737b079e738de73934a13de9784",
"rev": "8efb4337e857949f4cfac86d12ef1066f417f31f",
"type": "github"
},
"original": {
@@ -76,11 +76,11 @@
},
"nixpkgs-master": {
"locked": {
"lastModified": 1783021952,
"narHash": "sha256-8PghAtSGGZ0umfVI8Qbd7ZbFrfZPiH1UwtVbgLeikDA=",
"lastModified": 1783874024,
"narHash": "sha256-Fd8rPvyBv6JjcO/nZxZiFQan6Fww/jAF4TYj0Th/Yfo=",
"owner": "nixos",
"repo": "nixpkgs",
"rev": "f136374c679c54171a3ace589d15e9e79a8bd086",
"rev": "0b4f03c64b236e4ba4252414274e92796c300124",
"type": "github"
},
"original": {
@@ -108,11 +108,11 @@
},
"nixpkgs_2": {
"locked": {
"lastModified": 1782723713,
"narHash": "sha256-oPXCU/SSUokcGaJREHibG1CBX3+s/W7orDWQOZDsEeQ=",
"lastModified": 1783776592,
"narHash": "sha256-UgCQzxeWI75XM8G+hPrPh+MKzEPjG3SpAj7dtqSbksA=",
"owner": "nixos",
"repo": "nixpkgs",
"rev": "b5aa0fbd538984f6e3d201be0005b4463d8b09f8",
"rev": "e7a3ca8092b61ff85b6a45bf863ea2b2d6a661b3",
"type": "github"
},
"original": {
@@ -141,11 +141,11 @@
]
},
"locked": {
"lastModified": 1782165805,
"narHash": "sha256-478kKQBvK6SYTOdN2h9jhKJv94nbXRbFMfuL1WshErg=",
"lastModified": 1783174389,
"narHash": "sha256-aCWC8ngycU7OdJrU2+Je3qf+1a2ykuBvpPhZT/9tXMc=",
"owner": "Mic92",
"repo": "sops-nix",
"rev": "56b24064fdcaedca53553b1a6d607fd23b613a24",
"rev": "f1406619a3884cd5c47992a70b8b35c9c0fcb4c9",
"type": "github"
},
"original": {
+1 -3
View File
@@ -17,9 +17,7 @@
python-env = final: _prev: {
my_python = final.python314.withPackages (
ps:
with ps;
[
ps: with ps; [
alembic
apprise
apscheduler
+2 -5
View File
@@ -25,7 +25,7 @@ dependencies = [
"pydantic",
"pydantic-settings",
"python-multipart",
"sqlalchemy[asyncio]",
"sqlalchemy",
"tenacity",
"tiktoken",
"tinytuya",
@@ -41,10 +41,8 @@ whisper-transcribe = "python.tools.whisper.transcribe:main"
[dependency-groups]
dev = [
"aiosqlite",
"mypy",
"pyfakefs",
"pytest-asyncio",
"pytest-cov",
"pytest-mock",
"pytest-xdist",
@@ -119,7 +117,6 @@ exclude_lines = [
]
[tool.pytest.ini_options]
addopts = "-n auto -ra"
asyncio_mode = "auto"
addopts = "-n auto -ra --ignore=tests/ebook_search"
testpaths = ["tests"]
# --cov=system_tools --cov-report=term-missing --cov-report=xml --cov-report=html --cov-branch
@@ -1,206 +0,0 @@
"""adding Phrase metadata tables.
Revision ID: dddee09eddcc
Revises: 96d72c748c24
Create Date: 2026-06-29 00:49:07.344159
"""
from __future__ import annotations
from typing import TYPE_CHECKING
import sqlalchemy as sa
from alembic import op
from sqlalchemy.dialects import postgresql
from python.orm import RichieBase
if TYPE_CHECKING:
from collections.abc import Sequence
# revision identifiers, used by Alembic.
revision: str = "dddee09eddcc"
down_revision: str | None = "96d72c748c24"
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(
"candidate_phrases",
sa.Column("book_id", sa.Integer(), nullable=False),
sa.Column("series_id", sa.Integer(), nullable=True),
sa.Column("phrase_text", sa.Text(), nullable=False),
sa.Column("phrase_norm", sa.Text(), nullable=False),
sa.Column("token_count", sa.Integer(), nullable=False),
sa.Column("source_raw_ngram", sa.Boolean(), nullable=False),
sa.Column("source_yake", sa.Boolean(), nullable=False),
sa.Column("source_spacy_ner", sa.Boolean(), nullable=False),
sa.Column("source_spacy_noun_chunk", sa.Boolean(), nullable=False),
sa.Column("source_capitalized", sa.Boolean(), nullable=False),
sa.Column("source_metadata", sa.Boolean(), nullable=False),
sa.Column("spacy_label", sa.String(), nullable=True),
sa.Column("raw_count", sa.Integer(), nullable=False),
sa.Column("chapter_count", sa.Integer(), nullable=False),
sa.Column("yake_score", sa.Float(), nullable=True),
sa.Column("candidate_score", sa.Float(), nullable=False),
sa.Column(
"sample_contexts",
sa.JSON().with_variant(postgresql.JSONB(astext_type=sa.Text()), "postgresql"),
nullable=True,
),
sa.Column("llm_judged", sa.Boolean(), nullable=False),
sa.Column("llm_keep", sa.Boolean(), nullable=True),
sa.Column("llm_confidence", sa.Float(), nullable=True),
sa.Column("llm_category", sa.String(), nullable=True),
sa.Column("llm_reason", sa.Text(), 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(
["book_id"],
[f"{schema}.ebook_source.id"],
name=op.f("fk_candidate_phrases_book_id_ebook_source"),
ondelete="CASCADE",
),
sa.PrimaryKeyConstraint("id", name=op.f("pk_candidate_phrases")),
sa.UniqueConstraint("book_id", "phrase_norm", name="uq_candidate_phrases_book_id_phrase_norm"),
schema=schema,
)
op.create_index(
"candidate_phrases_book_norm_idx", "candidate_phrases", ["book_id", "phrase_norm"], unique=False, schema=schema
)
op.create_index(
"candidate_phrases_book_score_idx",
"candidate_phrases",
["book_id", "candidate_score"],
unique=False,
schema=schema,
)
op.create_table(
"protected_phrases",
sa.Column("book_id", sa.Integer(), nullable=True),
sa.Column("series_id", sa.Integer(), nullable=True),
sa.Column("phrase_text", sa.Text(), nullable=False),
sa.Column("phrase_norm", sa.Text(), nullable=False),
sa.Column("canonical_id", sa.String(), nullable=False),
sa.Column("phrase_type", sa.String(), nullable=True),
sa.Column("token_count", sa.Integer(), nullable=False),
sa.Column("confidence", sa.Float(), nullable=False),
sa.Column("importance", sa.Float(), nullable=False),
sa.Column("allow_nested", sa.Boolean(), nullable=False),
sa.Column("suppress_children", sa.Boolean(), nullable=False),
sa.Column("source_candidate_id", sa.Integer(), 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(
["book_id"],
[f"{schema}.ebook_source.id"],
name=op.f("fk_protected_phrases_book_id_ebook_source"),
ondelete="CASCADE",
),
sa.ForeignKeyConstraint(
["source_candidate_id"],
[f"{schema}.candidate_phrases.id"],
name=op.f("fk_protected_phrases_source_candidate_id_candidate_phrases"),
ondelete="SET NULL",
),
sa.PrimaryKeyConstraint("id", name=op.f("pk_protected_phrases")),
sa.UniqueConstraint("book_id", "phrase_norm", name="uq_protected_phrases_book_id_phrase_norm"),
schema=schema,
)
op.create_index(
"protected_phrases_book_norm_idx", "protected_phrases", ["book_id", "phrase_norm"], unique=False, schema=schema
)
op.create_index("protected_phrases_norm_idx", "protected_phrases", ["phrase_norm"], unique=False, schema=schema)
op.create_index(
"protected_phrases_series_norm_idx",
"protected_phrases",
["series_id", "phrase_norm"],
unique=False,
schema=schema,
)
op.create_table(
"chunk_phrase_mentions",
sa.Column("chunk_id", sa.BigInteger(), nullable=False),
sa.Column("phrase_id", sa.Integer(), nullable=False),
sa.Column("book_id", sa.Integer(), nullable=True),
sa.Column("series_id", sa.Integer(), nullable=True),
sa.Column("start_char", sa.Integer(), nullable=False),
sa.Column("end_char", sa.Integer(), 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(
["book_id"],
[f"{schema}.ebook_source.id"],
name=op.f("fk_chunk_phrase_mentions_book_id_ebook_source"),
ondelete="CASCADE",
),
sa.ForeignKeyConstraint(
["chunk_id"],
[f"{schema}.ebook_chunk.id"],
name=op.f("fk_chunk_phrase_mentions_chunk_id_ebook_chunk"),
ondelete="CASCADE",
),
sa.ForeignKeyConstraint(
["phrase_id"],
[f"{schema}.protected_phrases.id"],
name=op.f("fk_chunk_phrase_mentions_phrase_id_protected_phrases"),
ondelete="CASCADE",
),
sa.PrimaryKeyConstraint("id", name=op.f("pk_chunk_phrase_mentions")),
sa.UniqueConstraint("chunk_id", "phrase_id", "start_char", name="uq_chunk_phrase_mentions_chunk_phrase_start"),
schema=schema,
)
op.create_index(
"chunk_phrase_mentions_chunk_idx", "chunk_phrase_mentions", ["chunk_id"], unique=False, schema=schema
)
op.create_index(
"chunk_phrase_mentions_phrase_idx", "chunk_phrase_mentions", ["phrase_id"], unique=False, schema=schema
)
op.create_table(
"phrase_aliases",
sa.Column("phrase_id", sa.Integer(), nullable=False),
sa.Column("alias_text", sa.Text(), nullable=False),
sa.Column("alias_norm", sa.Text(), nullable=False),
sa.Column("confidence", sa.Float(), 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.ForeignKeyConstraint(
["phrase_id"],
[f"{schema}.protected_phrases.id"],
name=op.f("fk_phrase_aliases_phrase_id_protected_phrases"),
ondelete="CASCADE",
),
sa.PrimaryKeyConstraint("id", name=op.f("pk_phrase_aliases")),
sa.UniqueConstraint("phrase_id", "alias_norm", name="uq_phrase_aliases_phrase_id_alias_norm"),
schema=schema,
)
op.create_index("phrase_aliases_norm_idx", "phrase_aliases", ["alias_norm"], unique=False, schema=schema)
# ### end Alembic commands ###
def downgrade() -> None:
"""Downgrade."""
# ### commands auto generated by Alembic - please adjust! ###
op.drop_index("phrase_aliases_norm_idx", table_name="phrase_aliases", schema=schema)
op.drop_table("phrase_aliases", schema=schema)
op.drop_index("chunk_phrase_mentions_phrase_idx", table_name="chunk_phrase_mentions", schema=schema)
op.drop_index("chunk_phrase_mentions_chunk_idx", table_name="chunk_phrase_mentions", schema=schema)
op.drop_table("chunk_phrase_mentions", schema=schema)
op.drop_index("protected_phrases_series_norm_idx", table_name="protected_phrases", schema=schema)
op.drop_index("protected_phrases_norm_idx", table_name="protected_phrases", schema=schema)
op.drop_index("protected_phrases_book_norm_idx", table_name="protected_phrases", schema=schema)
op.drop_table("protected_phrases", schema=schema)
op.drop_index("candidate_phrases_book_score_idx", table_name="candidate_phrases", schema=schema)
op.drop_index("candidate_phrases_book_norm_idx", table_name="candidate_phrases", schema=schema)
op.drop_table("candidate_phrases", schema=schema)
# ### end Alembic commands ###
-1
View File
@@ -1 +0,0 @@
"""FastAPI applications."""
-56
View File
@@ -1,56 +0,0 @@
"""FastAPI interface for Contact database."""
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 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
if TYPE_CHECKING:
from collections.abc import AsyncIterator
logger = logging.getLogger(__name__)
def create_app() -> FastAPI:
"""Create and configure the FastAPI application."""
@asynccontextmanager
async def lifespan(app: FastAPI) -> AsyncIterator[None]:
"""Manage application lifespan."""
app.state.engine = get_postgres_engine()
yield
app.state.engine.dispose()
app = FastAPI(title="Contact Database API", lifespan=lifespan)
app.add_middleware(ZstdMiddleware)
app.include_router(contact_router)
app.include_router(views_router)
return app
def serve(
host: Annotated[str, typer.Option("--host", "-h", help="Host to bind to")],
port: Annotated[int, typer.Option("--port", "-p", help="Port to bind to")] = 8000,
log_level: Annotated[str, typer.Option("--log-level", "-l", help="Log level")] = "INFO",
) -> None:
"""Start the Contact API server."""
configure_logger(log_level)
app = create_app()
uvicorn.run(app, host=host, port=port)
if __name__ == "__main__":
typer.run(serve)
-6
View File
@@ -1,6 +0,0 @@
"""API routers."""
from python.api.routers.contact import router as contact_router
from python.api.routers.views import router as views_router
__all__ = ["contact_router", "views_router"]
-481
View File
@@ -1,481 +0,0 @@
"""Contact API router."""
from pathlib import Path
from fastapi import APIRouter, HTTPException, Request
from fastapi.responses import HTMLResponse
from fastapi.templating import Jinja2Templates
from pydantic import BaseModel
from sqlalchemy import select
from sqlalchemy.orm import selectinload
from python.fastapi_tools.db import DbSession # noqa: TC001 this is a FastAPI needed at runtime
from python.orm.richie.contact import Contact, ContactRelationship, Need, RelationshipType
TEMPLATES_DIR = Path(__file__).parent.parent / "templates"
templates = Jinja2Templates(directory=TEMPLATES_DIR)
def _is_htmx(request: Request) -> bool:
"""Check if the request is from HTMX."""
return request.headers.get("HX-Request") == "true"
class NeedBase(BaseModel):
"""Base schema for Need."""
name: str
description: str | None = None
class NeedCreate(NeedBase):
"""Schema for creating a Need."""
class NeedResponse(NeedBase):
"""Schema for Need response."""
id: int
model_config = {"from_attributes": True}
class ContactRelationshipCreate(BaseModel):
"""Schema for creating a contact relationship."""
related_contact_id: int
relationship_type: RelationshipType
closeness_weight: int | None = None
class ContactRelationshipUpdate(BaseModel):
"""Schema for updating a contact relationship."""
relationship_type: RelationshipType | None = None
closeness_weight: int | None = None
class ContactRelationshipResponse(BaseModel):
"""Schema for contact relationship response."""
contact_id: int
related_contact_id: int
relationship_type: str
closeness_weight: int
model_config = {"from_attributes": True}
class RelationshipTypeInfo(BaseModel):
"""Information about a relationship type."""
value: str
display_name: str
default_weight: int
class GraphNode(BaseModel):
"""Node in the relationship graph."""
id: int
name: str
current_job: str | None = None
class GraphEdge(BaseModel):
"""Edge in the relationship graph."""
source: int
target: int
relationship_type: str
closeness_weight: int
class GraphData(BaseModel):
"""Complete graph data for visualization."""
nodes: list[GraphNode]
edges: list[GraphEdge]
class ContactBase(BaseModel):
"""Base schema for Contact."""
name: str
age: int | None = None
bio: str | None = None
current_job: str | None = None
gender: str | None = None
goals: str | None = None
legal_name: str | None = None
profile_pic: str | None = None
safe_conversation_starters: str | None = None
self_sufficiency_score: int | None = None
social_structure_style: str | None = None
ssn: str | None = None
suffix: str | None = None
timezone: str | None = None
topics_to_avoid: str | None = None
class ContactCreate(ContactBase):
"""Schema for creating a Contact."""
need_ids: list[int] = []
class ContactUpdate(BaseModel):
"""Schema for updating a Contact."""
name: str | None = None
age: int | None = None
bio: str | None = None
current_job: str | None = None
gender: str | None = None
goals: str | None = None
legal_name: str | None = None
profile_pic: str | None = None
safe_conversation_starters: str | None = None
self_sufficiency_score: int | None = None
social_structure_style: str | None = None
ssn: str | None = None
suffix: str | None = None
timezone: str | None = None
topics_to_avoid: str | None = None
need_ids: list[int] | None = None
class ContactResponse(ContactBase):
"""Schema for Contact response with relationships."""
id: int
needs: list[NeedResponse] = []
related_to: list[ContactRelationshipResponse] = []
related_from: list[ContactRelationshipResponse] = []
model_config = {"from_attributes": True}
class ContactListResponse(ContactBase):
"""Schema for Contact list response."""
id: int
model_config = {"from_attributes": True}
router = APIRouter(prefix="/api", tags=["contacts"])
@router.post("/needs", response_model=NeedResponse)
def create_need(need: NeedCreate, db: DbSession) -> Need:
"""Create a new need."""
db_need = Need(name=need.name, description=need.description)
db.add(db_need)
db.commit()
db.refresh(db_need)
return db_need
@router.get("/needs", response_model=list[NeedResponse])
def list_needs(db: DbSession) -> list[Need]:
"""List all needs."""
return list(db.scalars(select(Need)).all())
@router.get("/needs/{need_id}", response_model=NeedResponse)
def get_need(need_id: int, db: DbSession) -> Need:
"""Get a need by ID."""
need = db.get(Need, need_id)
if not need:
raise HTTPException(status_code=404, detail="Need not found")
return need
@router.delete("/needs/{need_id}", response_model=None)
def delete_need(need_id: int, request: Request, db: DbSession) -> dict[str, bool] | HTMLResponse:
"""Delete a need by ID."""
need = db.get(Need, need_id)
if not need:
raise HTTPException(status_code=404, detail="Need not found")
db.delete(need)
db.commit()
if _is_htmx(request):
return HTMLResponse("")
return {"deleted": True}
@router.post("/contacts", response_model=ContactResponse)
def create_contact(contact: ContactCreate, db: DbSession) -> Contact:
"""Create a new contact."""
need_ids = contact.need_ids
contact_data = contact.model_dump(exclude={"need_ids"})
db_contact = Contact(**contact_data)
if need_ids:
needs = list(db.scalars(select(Need).where(Need.id.in_(need_ids))).all())
db_contact.needs = needs
db.add(db_contact)
db.commit()
db.refresh(db_contact)
return db_contact
@router.get("/contacts", response_model=list[ContactListResponse])
def list_contacts(
db: DbSession,
skip: int = 0,
limit: int = 100,
) -> list[Contact]:
"""List all contacts with pagination."""
return list(db.scalars(select(Contact).offset(skip).limit(limit)).all())
@router.get("/contacts/{contact_id}", response_model=ContactResponse)
def get_contact(contact_id: int, db: DbSession) -> Contact:
"""Get a contact by ID with all relationships."""
contact = db.scalar(
select(Contact)
.where(Contact.id == contact_id)
.options(
selectinload(Contact.needs),
selectinload(Contact.related_to),
selectinload(Contact.related_from),
)
)
if not contact:
raise HTTPException(status_code=404, detail="Contact not found")
return contact
@router.patch("/contacts/{contact_id}", response_model=ContactResponse)
def update_contact(
contact_id: int,
contact: ContactUpdate,
db: DbSession,
) -> Contact:
"""Update a contact by ID."""
db_contact = db.get(Contact, contact_id)
if not db_contact:
raise HTTPException(status_code=404, detail="Contact not found")
update_data = contact.model_dump(exclude_unset=True)
need_ids = update_data.pop("need_ids", None)
for key, value in update_data.items():
setattr(db_contact, key, value)
if need_ids is not None:
needs = list(db.scalars(select(Need).where(Need.id.in_(need_ids))).all())
db_contact.needs = needs
db.commit()
db.refresh(db_contact)
return db_contact
@router.delete("/contacts/{contact_id}", response_model=None)
def delete_contact(contact_id: int, request: Request, db: DbSession) -> dict[str, bool] | HTMLResponse:
"""Delete a contact by ID."""
contact = db.get(Contact, contact_id)
if not contact:
raise HTTPException(status_code=404, detail="Contact not found")
db.delete(contact)
db.commit()
if _is_htmx(request):
return HTMLResponse("")
return {"deleted": True}
@router.post("/contacts/{contact_id}/needs/{need_id}")
def add_need_to_contact(
contact_id: int,
need_id: int,
db: DbSession,
) -> dict[str, bool]:
"""Add a need to a contact."""
contact = db.get(Contact, contact_id)
if not contact:
raise HTTPException(status_code=404, detail="Contact not found")
need = db.get(Need, need_id)
if not need:
raise HTTPException(status_code=404, detail="Need not found")
if need not in contact.needs:
contact.needs.append(need)
db.commit()
return {"added": True}
@router.delete("/contacts/{contact_id}/needs/{need_id}", response_model=None)
def remove_need_from_contact(
contact_id: int,
need_id: int,
request: Request,
db: DbSession,
) -> dict[str, bool] | HTMLResponse:
"""Remove a need from a contact."""
contact = db.get(Contact, contact_id)
if not contact:
raise HTTPException(status_code=404, detail="Contact not found")
need = db.get(Need, need_id)
if not need:
raise HTTPException(status_code=404, detail="Need not found")
if need in contact.needs:
contact.needs.remove(need)
db.commit()
if _is_htmx(request):
return HTMLResponse("")
return {"removed": True}
@router.post(
"/contacts/{contact_id}/relationships",
response_model=ContactRelationshipResponse,
)
def add_contact_relationship(
contact_id: int,
relationship: ContactRelationshipCreate,
db: DbSession,
) -> ContactRelationship:
"""Add a relationship between two contacts."""
contact = db.get(Contact, contact_id)
if not contact:
raise HTTPException(status_code=404, detail="Contact not found")
related_contact = db.get(Contact, relationship.related_contact_id)
if not related_contact:
raise HTTPException(status_code=404, detail="Related contact not found")
if contact_id == relationship.related_contact_id:
raise HTTPException(status_code=400, detail="Cannot relate contact to itself")
# Use provided weight or default from relationship type
weight = relationship.closeness_weight
if weight is None:
weight = relationship.relationship_type.default_weight
db_relationship = ContactRelationship(
contact_id=contact_id,
related_contact_id=relationship.related_contact_id,
relationship_type=relationship.relationship_type.value,
closeness_weight=weight,
)
db.add(db_relationship)
db.commit()
db.refresh(db_relationship)
return db_relationship
@router.get(
"/contacts/{contact_id}/relationships",
response_model=list[ContactRelationshipResponse],
)
def get_contact_relationships(
contact_id: int,
db: DbSession,
) -> list[ContactRelationship]:
"""Get all relationships for a contact."""
contact = db.get(Contact, contact_id)
if not contact:
raise HTTPException(status_code=404, detail="Contact not found")
outgoing = list(db.scalars(select(ContactRelationship).where(ContactRelationship.contact_id == contact_id)).all())
incoming = list(
db.scalars(select(ContactRelationship).where(ContactRelationship.related_contact_id == contact_id)).all()
)
return outgoing + incoming
@router.patch(
"/contacts/{contact_id}/relationships/{related_contact_id}",
response_model=ContactRelationshipResponse,
)
def update_contact_relationship(
contact_id: int,
related_contact_id: int,
update: ContactRelationshipUpdate,
db: DbSession,
) -> ContactRelationship:
"""Update a relationship between two contacts."""
relationship = db.scalar(
select(ContactRelationship).where(
ContactRelationship.contact_id == contact_id,
ContactRelationship.related_contact_id == related_contact_id,
)
)
if not relationship:
raise HTTPException(status_code=404, detail="Relationship not found")
if update.relationship_type is not None:
relationship.relationship_type = update.relationship_type.value
if update.closeness_weight is not None:
relationship.closeness_weight = update.closeness_weight
db.commit()
db.refresh(relationship)
return relationship
@router.delete("/contacts/{contact_id}/relationships/{related_contact_id}", response_model=None)
def remove_contact_relationship(
contact_id: int,
related_contact_id: int,
request: Request,
db: DbSession,
) -> dict[str, bool] | HTMLResponse:
"""Remove a relationship between two contacts."""
relationship = db.scalar(
select(ContactRelationship).where(
ContactRelationship.contact_id == contact_id,
ContactRelationship.related_contact_id == related_contact_id,
)
)
if not relationship:
raise HTTPException(status_code=404, detail="Relationship not found")
db.delete(relationship)
db.commit()
if _is_htmx(request):
return HTMLResponse("")
return {"deleted": True}
@router.get("/relationship-types")
def list_relationship_types() -> list[RelationshipTypeInfo]:
"""List all available relationship types with their default weights."""
return [
RelationshipTypeInfo(
value=rt.value,
display_name=rt.display_name,
default_weight=rt.default_weight,
)
for rt in RelationshipType
]
@router.get("/graph")
def get_relationship_graph(db: DbSession) -> GraphData:
"""Get all contacts and relationships as graph data for visualization."""
contacts = list(db.scalars(select(Contact)).all())
relationships = list(db.scalars(select(ContactRelationship)).all())
nodes = [GraphNode(id=c.id, name=c.name, current_job=c.current_job) for c in contacts]
edges = [
GraphEdge(
source=rel.contact_id,
target=rel.related_contact_id,
relationship_type=rel.relationship_type,
closeness_weight=rel.closeness_weight,
)
for rel in relationships
]
return GraphData(nodes=nodes, edges=edges)
-345
View File
@@ -1,345 +0,0 @@
"""HTMX server-rendered view router."""
from pathlib import Path
from typing import Annotated, Any
from fastapi import APIRouter, Form, HTTPException, Request
from fastapi.responses import HTMLResponse, RedirectResponse
from fastapi.templating import Jinja2Templates
from sqlalchemy import select
from sqlalchemy.orm import Session, selectinload
from python.fastapi_tools.db import DbSession # noqa: TC001 this is a FastAPI needed at runtime
from python.orm.richie.contact import Contact, ContactRelationship, Need, RelationshipType
TEMPLATES_DIR = Path(__file__).parent.parent / "templates"
templates = Jinja2Templates(directory=TEMPLATES_DIR)
router = APIRouter(tags=["views"])
FAMILIAL_TYPES = {
"parent",
"child",
"sibling",
"grandparent",
"grandchild",
"aunt_uncle",
"niece_nephew",
"cousin",
"in_law",
}
FRIEND_TYPES = {"best_friend", "close_friend", "friend", "acquaintance", "neighbor"}
PARTNER_TYPES = {"spouse", "partner"}
PROFESSIONAL_TYPES = {"mentor", "mentee", "business_partner", "colleague", "manager", "direct_report", "client"}
CONTACT_STRING_FIELDS = (
"name",
"legal_name",
"suffix",
"gender",
"current_job",
"timezone",
"profile_pic",
"bio",
"goals",
"social_structure_style",
"safe_conversation_starters",
"topics_to_avoid",
"ssn",
)
CONTACT_INT_FIELDS = ("age", "self_sufficiency_score")
def _group_relationships(relationships: list[ContactRelationship]) -> dict[str, list[ContactRelationship]]:
"""Group relationships by category."""
groups: dict[str, list[ContactRelationship]] = {
"familial": [],
"partners": [],
"friends": [],
"professional": [],
"other": [],
}
for rel in relationships:
if rel.relationship_type in FAMILIAL_TYPES:
groups["familial"].append(rel)
elif rel.relationship_type in PARTNER_TYPES:
groups["partners"].append(rel)
elif rel.relationship_type in FRIEND_TYPES:
groups["friends"].append(rel)
elif rel.relationship_type in PROFESSIONAL_TYPES:
groups["professional"].append(rel)
else:
groups["other"].append(rel)
return groups
def _build_contact_name_map(database: Session, contact: Contact) -> dict[int, str]:
"""Build a mapping of contact IDs to names for relationship display."""
related_ids = {rel.related_contact_id for rel in contact.related_to}
related_ids |= {rel.contact_id for rel in contact.related_from}
related_ids.discard(contact.id)
if not related_ids:
return {}
related_contacts = list(database.scalars(select(Contact).where(Contact.id.in_(related_ids))).all())
return {related.id: related.name for related in related_contacts}
def _get_relationship_type_display() -> dict[str, str]:
"""Build a mapping of relationship type values to display names."""
return {rel_type.value: rel_type.display_name for rel_type in RelationshipType}
async def _parse_contact_form(request: Request) -> dict[str, Any]:
"""Parse contact form data from a multipart/form request."""
form_data = await request.form()
result: dict[str, Any] = {}
for field in CONTACT_STRING_FIELDS:
value = form_data.get(field, "")
result[field] = str(value) if value else None
for field in CONTACT_INT_FIELDS:
value = form_data.get(field, "")
result[field] = int(value) if value else None
result["need_ids"] = [int(value) for value in form_data.getlist("need_ids")]
return result
def _save_contact_from_form(database: Session, contact: Contact, form_result: dict[str, Any]) -> None:
"""Apply parsed form data to a Contact and save associated needs."""
need_ids = form_result.pop("need_ids")
for key, value in form_result.items():
setattr(contact, key, value)
if need_ids:
contact.needs = list(database.scalars(select(Need).where(Need.id.in_(need_ids))).all())
else:
contact.needs = []
@router.get("/", response_class=HTMLResponse)
@router.get("/contacts", response_class=HTMLResponse)
def contact_list_page(request: Request, database: DbSession) -> HTMLResponse:
"""Render the contacts list page."""
contacts = list(database.scalars(select(Contact)).all())
return templates.TemplateResponse(request, "contact_list.html", {"contacts": contacts})
@router.get("/contacts/new", response_class=HTMLResponse)
def new_contact_page(request: Request, database: DbSession) -> HTMLResponse:
"""Render the new contact form page."""
all_needs = list(database.scalars(select(Need)).all())
return templates.TemplateResponse(request, "contact_form.html", {"contact": None, "all_needs": all_needs})
@router.post("/htmx/contacts/new")
async def create_contact_form(request: Request, database: DbSession) -> RedirectResponse:
"""Handle the create contact form submission."""
form_result = await _parse_contact_form(request)
contact = Contact()
_save_contact_from_form(database, contact, form_result)
database.add(contact)
database.commit()
database.refresh(contact)
return RedirectResponse(url=f"/contacts/{contact.id}", status_code=303)
@router.get("/contacts/{contact_id}", response_class=HTMLResponse)
def contact_detail_page(contact_id: int, request: Request, database: DbSession) -> HTMLResponse:
"""Render the contact detail page."""
contact = database.scalar(
select(Contact)
.where(Contact.id == contact_id)
.options(
selectinload(Contact.needs),
selectinload(Contact.related_to),
selectinload(Contact.related_from),
)
)
if not contact:
raise HTTPException(status_code=404, detail="Contact not found")
contact_names = _build_contact_name_map(database, contact)
grouped_relationships = _group_relationships(contact.related_to)
all_contacts = list(database.scalars(select(Contact)).all())
all_needs = list(database.scalars(select(Need)).all())
available_needs = [need for need in all_needs if need not in contact.needs]
return templates.TemplateResponse(
request,
"contact_detail.html",
{
"contact": contact,
"contact_names": contact_names,
"grouped_relationships": grouped_relationships,
"all_contacts": all_contacts,
"available_needs": available_needs,
"relationship_types": list(RelationshipType),
},
)
@router.get("/contacts/{contact_id}/edit", response_class=HTMLResponse)
def edit_contact_page(contact_id: int, request: Request, database: DbSession) -> HTMLResponse:
"""Render the edit contact form page."""
contact = database.scalar(select(Contact).where(Contact.id == contact_id).options(selectinload(Contact.needs)))
if not contact:
raise HTTPException(status_code=404, detail="Contact not found")
all_needs = list(database.scalars(select(Need)).all())
return templates.TemplateResponse(request, "contact_form.html", {"contact": contact, "all_needs": all_needs})
@router.post("/htmx/contacts/{contact_id}/edit")
async def update_contact_form(contact_id: int, request: Request, database: DbSession) -> RedirectResponse:
"""Handle the edit contact form submission."""
contact = database.get(Contact, contact_id)
if not contact:
raise HTTPException(status_code=404, detail="Contact not found")
form_result = await _parse_contact_form(request)
_save_contact_from_form(database, contact, form_result)
database.commit()
return RedirectResponse(url=f"/contacts/{contact_id}", status_code=303)
@router.post("/htmx/contacts/{contact_id}/add-need", response_class=HTMLResponse)
def add_need_to_contact_htmx(
contact_id: int,
request: Request,
database: DbSession,
need_id: Annotated[int, Form()],
) -> HTMLResponse:
"""Add a need to a contact and return updated manage-needs partial."""
contact = database.scalar(select(Contact).where(Contact.id == contact_id).options(selectinload(Contact.needs)))
if not contact:
raise HTTPException(status_code=404, detail="Contact not found")
need = database.get(Need, need_id)
if not need:
raise HTTPException(status_code=404, detail="Need not found")
if need not in contact.needs:
contact.needs.append(need)
database.commit()
database.refresh(contact)
return templates.TemplateResponse(request, "partials/manage_needs.html", {"contact": contact})
@router.post("/htmx/contacts/{contact_id}/add-relationship", response_class=HTMLResponse)
def add_relationship_htmx(
contact_id: int,
request: Request,
database: DbSession,
related_contact_id: Annotated[int, Form()],
relationship_type: Annotated[str, Form()],
) -> HTMLResponse:
"""Add a relationship and return updated manage-relationships partial."""
contact = database.scalar(select(Contact).where(Contact.id == contact_id).options(selectinload(Contact.related_to)))
if not contact:
raise HTTPException(status_code=404, detail="Contact not found")
related_contact = database.get(Contact, related_contact_id)
if not related_contact:
raise HTTPException(status_code=404, detail="Related contact not found")
rel_type = RelationshipType(relationship_type)
weight = rel_type.default_weight
relationship = ContactRelationship(
contact_id=contact_id,
related_contact_id=related_contact_id,
relationship_type=relationship_type,
closeness_weight=weight,
)
database.add(relationship)
database.commit()
database.refresh(contact)
contact_names = _build_contact_name_map(database, contact)
return templates.TemplateResponse(
request,
"partials/manage_relationships.html",
{"contact": contact, "contact_names": contact_names},
)
@router.post("/htmx/contacts/{contact_id}/relationships/{related_contact_id}/weight")
def update_relationship_weight_htmx(
contact_id: int,
related_contact_id: int,
database: DbSession,
closeness_weight: Annotated[int, Form()],
) -> HTMLResponse:
"""Update a relationship's closeness weight from HTMX range input."""
relationship = database.scalar(
select(ContactRelationship).where(
ContactRelationship.contact_id == contact_id,
ContactRelationship.related_contact_id == related_contact_id,
)
)
if not relationship:
raise HTTPException(status_code=404, detail="Relationship not found")
relationship.closeness_weight = closeness_weight
database.commit()
return HTMLResponse("")
@router.post("/htmx/needs", response_class=HTMLResponse)
def create_need_htmx(
request: Request,
database: DbSession,
name: Annotated[str, Form()],
description: Annotated[str, Form()] = "",
) -> HTMLResponse:
"""Create a need via form data and return updated needs list."""
need = Need(name=name, description=description or None)
database.add(need)
database.commit()
needs = list(database.scalars(select(Need)).all())
return templates.TemplateResponse(request, "partials/need_items.html", {"needs": needs})
@router.get("/needs", response_class=HTMLResponse)
def needs_page(request: Request, database: DbSession) -> HTMLResponse:
"""Render the needs list page."""
needs = list(database.scalars(select(Need)).all())
return templates.TemplateResponse(request, "need_list.html", {"needs": needs})
@router.get("/graph", response_class=HTMLResponse)
def graph_page(request: Request, database: DbSession) -> HTMLResponse:
"""Render the relationship graph page."""
contacts = list(database.scalars(select(Contact)).all())
relationships = list(database.scalars(select(ContactRelationship)).all())
graph_data = {
"nodes": [{"id": contact.id, "name": contact.name, "current_job": contact.current_job} for contact in contacts],
"edges": [
{
"source": rel.contact_id,
"target": rel.related_contact_id,
"relationship_type": rel.relationship_type,
"closeness_weight": rel.closeness_weight,
}
for rel in relationships
],
}
return templates.TemplateResponse(
request,
"graph.html",
{
"graph_data": graph_data,
"relationship_type_display": _get_relationship_type_display(),
},
)
-198
View File
@@ -1,198 +0,0 @@
<!DOCTYPE html>
<html lang="en" data-theme="light">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>{% block title %}Contact Database{% endblock %}</title>
<script src="https://unpkg.com/htmx.org@2.0.4"></script>
<style>
:root {
--color-bg: #f5f5f5;
--color-bg-card: #ffffff;
--color-bg-hover: #f0f0f0;
--color-bg-muted: #f9f9f9;
--color-bg-error: #ffe0e0;
--color-text: #333333;
--color-text-muted: #666666;
--color-text-error: #cc0000;
--color-border: #dddddd;
--color-border-light: #eeeeee;
--color-border-lighter: #f0f0f0;
--color-primary: #0066cc;
--color-primary-hover: #0055aa;
--color-danger: #cc3333;
--color-danger-hover: #aa2222;
--color-tag-bg: #e0e0e0;
--shadow: 0 1px 3px rgba(0, 0, 0, 0.1);
font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, sans-serif;
line-height: 1.5;
color: var(--color-text);
background-color: var(--color-bg);
}
[data-theme="dark"] {
--color-bg: #1a1a1a;
--color-bg-card: #2d2d2d;
--color-bg-hover: #3d3d3d;
--color-bg-muted: #252525;
--color-bg-error: #4a2020;
--color-text: #e0e0e0;
--color-text-muted: #a0a0a0;
--color-text-error: #ff6b6b;
--color-border: #404040;
--color-border-light: #353535;
--color-border-lighter: #303030;
--color-primary: #4da6ff;
--color-primary-hover: #7dbfff;
--color-danger: #ff6b6b;
--color-danger-hover: #ff8a8a;
--color-tag-bg: #404040;
--shadow: 0 1px 3px rgba(0, 0, 0, 0.3);
}
* { box-sizing: border-box; }
body { margin: 0; background: var(--color-bg); color: var(--color-text); }
.app { max-width: 1000px; margin: 0 auto; padding: 20px; }
nav { display: flex; align-items: center; gap: 20px; padding: 15px 0; border-bottom: 1px solid var(--color-border); margin-bottom: 20px; }
nav a { color: var(--color-primary); text-decoration: none; font-weight: 500; }
nav a:hover { text-decoration: underline; }
.theme-toggle { margin-left: auto; }
main { background: var(--color-bg-card); padding: 20px; border-radius: 8px; box-shadow: var(--shadow); }
.header { display: flex; justify-content: space-between; align-items: center; margin-bottom: 20px; }
.header h1 { margin: 0; }
a { color: var(--color-primary); }
a:hover { text-decoration: underline; }
.btn { display: inline-block; padding: 8px 16px; border: 1px solid var(--color-border); border-radius: 4px; background: var(--color-bg-card); color: var(--color-text); text-decoration: none; cursor: pointer; font-size: 14px; margin-left: 8px; }
.btn:hover { background: var(--color-bg-hover); }
.btn-primary { background: var(--color-primary); border-color: var(--color-primary); color: white; }
.btn-primary:hover { background: var(--color-primary-hover); }
.btn-danger { background: var(--color-danger); border-color: var(--color-danger); color: white; }
.btn-danger:hover { background: var(--color-danger-hover); }
.btn-small { padding: 4px 8px; font-size: 12px; }
.btn:disabled { opacity: 0.6; cursor: not-allowed; }
table { width: 100%; border-collapse: collapse; }
th, td { padding: 12px; text-align: left; border-bottom: 1px solid var(--color-border-light); }
th { font-weight: 600; background: var(--color-bg-muted); }
tr:hover { background: var(--color-bg-muted); }
.error { background: var(--color-bg-error); color: var(--color-text-error); padding: 10px; border-radius: 4px; margin-bottom: 20px; }
.tag { display: inline-block; background: var(--color-tag-bg); padding: 2px 8px; border-radius: 12px; font-size: 12px; color: var(--color-text-muted); }
.add-form { display: flex; gap: 10px; margin-top: 15px; flex-wrap: wrap; }
.add-form select, .add-form input { padding: 8px; border: 1px solid var(--color-border); border-radius: 4px; min-width: 200px; background: var(--color-bg-card); color: var(--color-text); }
.form-group { margin-bottom: 20px; }
.form-group label { display: block; font-weight: 500; margin-bottom: 5px; }
.form-group input, .form-group textarea, .form-group select { width: 100%; padding: 10px; border: 1px solid var(--color-border); border-radius: 4px; font-size: 14px; background: var(--color-bg-card); color: var(--color-text); }
.form-group textarea { resize: vertical; }
.form-row { display: grid; grid-template-columns: 1fr 1fr; gap: 20px; }
.checkbox-group { display: flex; flex-wrap: wrap; gap: 15px; }
.checkbox-label { display: flex; align-items: center; gap: 5px; cursor: pointer; }
.form-actions { display: flex; gap: 10px; margin-top: 30px; padding-top: 20px; border-top: 1px solid var(--color-border-light); }
.need-form { background: var(--color-bg-muted); padding: 20px; border-radius: 4px; margin-bottom: 20px; }
.need-items { list-style: none; padding: 0; }
.need-items li { display: flex; justify-content: space-between; align-items: flex-start; padding: 15px; border: 1px solid var(--color-border-light); border-radius: 4px; margin-bottom: 10px; }
.need-info p { margin: 5px 0 0; color: var(--color-text-muted); font-size: 14px; }
.graph-container { width: 100%; }
.graph-hint { color: var(--color-text-muted); font-size: 14px; margin-bottom: 15px; }
.selected-info { margin-top: 15px; padding: 15px; background: var(--color-bg-muted); border-radius: 8px; }
.selected-info h3 { margin: 0 0 10px; }
.selected-info p { margin: 5px 0; color: var(--color-text-muted); }
.legend { margin-top: 20px; padding: 15px; background: var(--color-bg-muted); border-radius: 8px; }
.legend h4 { margin: 0 0 10px; font-size: 14px; }
.legend-items { display: flex; flex-wrap: wrap; gap: 15px; }
.legend-item { display: flex; align-items: center; gap: 8px; font-size: 12px; color: var(--color-text-muted); }
.legend-line { width: 30px; border-radius: 2px; }
.id-card { width: 100%; }
.id-card-inner { background: linear-gradient(135deg, #0a0a0f 0%, #1a1a2e 50%, #0a0a0f 100%); background-image: radial-gradient(white 1px, transparent 1px), linear-gradient(135deg, #0a0a0f 0%, #1a1a2e 50%, #0a0a0f 100%); background-size: 50px 50px, 100% 100%; color: #fff; border-radius: 12px; padding: 25px; min-height: 500px; position: relative; overflow: hidden; }
.id-card-header { display: flex; justify-content: space-between; align-items: flex-start; margin-bottom: 15px; }
.id-card-header-left { flex: 1; }
.id-card-header-right { display: flex; flex-direction: column; align-items: flex-end; gap: 10px; }
.id-card-title { font-size: 2.5rem; font-weight: 700; margin: 0; color: #fff; text-shadow: 2px 2px 4px rgba(0,0,0,0.5); }
.id-profile-pic { width: 80px; height: 80px; border-radius: 8px; object-fit: cover; border: 2px solid rgba(255,255,255,0.3); }
.id-profile-placeholder { width: 80px; height: 80px; border-radius: 8px; background: linear-gradient(135deg, #4ecdc4 0%, #44a8a0 100%); display: flex; align-items: center; justify-content: center; border: 2px solid rgba(255,255,255,0.3); }
.id-profile-placeholder span { font-size: 2rem; font-weight: 700; color: #fff; text-shadow: 1px 1px 2px rgba(0,0,0,0.3); }
.id-card-actions { display: flex; gap: 8px; }
.id-card-actions .btn { background: rgba(255,255,255,0.1); border-color: rgba(255,255,255,0.3); color: #fff; }
.id-card-actions .btn:hover { background: rgba(255,255,255,0.2); }
.id-card-body { display: grid; grid-template-columns: 1fr 1.5fr; gap: 30px; }
.id-card-left { display: flex; flex-direction: column; gap: 8px; }
.id-field { font-size: 1rem; line-height: 1.4; }
.id-field-block { margin-top: 15px; font-size: 0.95rem; line-height: 1.5; }
.id-label { color: #4ecdc4; font-weight: 500; }
.id-card-right { display: flex; flex-direction: column; gap: 20px; }
.id-bio { font-size: 0.9rem; line-height: 1.6; color: #e0e0e0; }
.id-relationships { margin-top: 10px; }
.id-section-title { font-size: 1.5rem; margin: 0 0 15px; color: #fff; border-bottom: 1px solid rgba(255,255,255,0.2); padding-bottom: 8px; }
.id-rel-group { margin-bottom: 12px; font-size: 0.9rem; line-height: 1.6; }
.id-rel-label { color: #a0a0a0; }
.id-rel-group a { color: #4ecdc4; text-decoration: none; }
.id-rel-group a:hover { text-decoration: underline; }
.id-rel-type { color: #888; font-size: 0.85em; }
.id-card-warnings { margin-top: 30px; padding-top: 20px; border-top: 1px solid rgba(255,255,255,0.2); display: flex; flex-wrap: wrap; gap: 20px; }
.id-warning { display: flex; align-items: center; gap: 8px; font-size: 0.9rem; color: #ff6b6b; }
.warning-dot { width: 8px; height: 8px; background: #ff6b6b; border-radius: 50%; flex-shrink: 0; }
.warning-desc { color: #ccc; }
.id-card-manage { margin-top: 20px; background: var(--color-bg-muted); border-radius: 8px; padding: 15px; }
.id-card-manage summary { cursor: pointer; font-weight: 600; font-size: 1.1rem; padding: 5px 0; }
.id-card-manage[open] summary { margin-bottom: 15px; border-bottom: 1px solid var(--color-border-light); padding-bottom: 10px; }
.manage-section { margin-bottom: 25px; }
.manage-section h3 { margin: 0 0 15px; font-size: 1rem; }
.manage-relationships { display: flex; flex-direction: column; gap: 10px; margin-bottom: 15px; }
.manage-rel-item { display: flex; align-items: center; gap: 12px; padding: 10px; background: var(--color-bg-card); border-radius: 6px; flex-wrap: wrap; }
.manage-rel-item a { font-weight: 500; min-width: 120px; }
.weight-control { display: flex; align-items: center; gap: 8px; font-size: 12px; color: var(--color-text-muted); }
.weight-control input[type="range"] { width: 80px; cursor: pointer; }
.weight-value { min-width: 20px; text-align: center; font-weight: 600; }
.manage-needs-list { list-style: none; padding: 0; margin: 0 0 15px; }
.manage-needs-list li { display: flex; align-items: center; gap: 12px; padding: 10px; background: var(--color-bg-card); border-radius: 6px; margin-bottom: 8px; }
.manage-needs-list li .btn { margin-left: auto; }
.htmx-indicator { display: none; }
.htmx-request .htmx-indicator { display: inline; }
.htmx-request.htmx-indicator { display: inline; }
@media (max-width: 768px) {
.id-card-body { grid-template-columns: 1fr; }
.id-card-title { font-size: 1.8rem; }
.id-card-header { flex-direction: column; gap: 15px; }
}
</style>
</head>
<body>
<div class="app">
<nav>
<a href="/contacts">Contacts</a>
<a href="/graph">Graph</a>
<a href="/needs">Needs</a>
<button class="btn btn-small theme-toggle" onclick="toggleTheme()">
<span id="theme-label">Dark</span>
</button>
</nav>
<main id="main-content">
{% block content %}{% endblock %}
</main>
</div>
<script>
function toggleTheme() {
const html = document.documentElement;
const current = html.getAttribute('data-theme');
const next = current === 'light' ? 'dark' : 'light';
html.setAttribute('data-theme', next);
localStorage.setItem('theme', next);
document.getElementById('theme-label').textContent = next === 'light' ? 'Dark' : 'Light';
}
(function() {
const saved = localStorage.getItem('theme') || 'light';
document.documentElement.setAttribute('data-theme', saved);
document.getElementById('theme-label').textContent = saved === 'light' ? 'Dark' : 'Light';
})();
</script>
</body>
</html>
-204
View File
@@ -1,204 +0,0 @@
{% extends "base.html" %}
{% block title %}{{ contact.name }}{% endblock %}
{% block content %}
<div class="id-card">
<div class="id-card-inner">
<div class="id-card-header">
<div class="id-card-header-left">
<h1 class="id-card-title">I.D.: {{ contact.name }}</h1>
</div>
<div class="id-card-header-right">
{% if contact.profile_pic %}
<img src="{{ contact.profile_pic }}" alt="{{ contact.name }}'s profile" class="id-profile-pic">
{% else %}
<div class="id-profile-placeholder">
<span>{{ contact.name[0]|upper }}</span>
</div>
{% endif %}
<div class="id-card-actions">
<a href="/contacts/{{ contact.id }}/edit" class="btn btn-small">Edit</a>
<a href="/contacts" class="btn btn-small">Back</a>
</div>
</div>
</div>
<div class="id-card-body">
<div class="id-card-left">
{% if contact.legal_name %}
<div class="id-field">Legal name: {{ contact.legal_name }}</div>
{% endif %}
{% if contact.suffix %}
<div class="id-field">Suffix: {{ contact.suffix }}</div>
{% endif %}
{% if contact.gender %}
<div class="id-field">Gender: {{ contact.gender }}</div>
{% endif %}
{% if contact.age %}
<div class="id-field">Age: {{ contact.age }}</div>
{% endif %}
{% if contact.current_job %}
<div class="id-field">Job: {{ contact.current_job }}</div>
{% endif %}
{% if contact.social_structure_style %}
<div class="id-field">Social style: {{ contact.social_structure_style }}</div>
{% endif %}
{% if contact.self_sufficiency_score is not none %}
<div class="id-field">Self-Sufficiency: {{ contact.self_sufficiency_score }}</div>
{% endif %}
{% if contact.timezone %}
<div class="id-field">Timezone: {{ contact.timezone }}</div>
{% endif %}
{% if contact.safe_conversation_starters %}
<div class="id-field-block">
<span class="id-label">Safe con starters:</span> {{ contact.safe_conversation_starters }}
</div>
{% endif %}
{% if contact.topics_to_avoid %}
<div class="id-field-block">
<span class="id-label">Topics to avoid:</span> {{ contact.topics_to_avoid }}
</div>
{% endif %}
{% if contact.goals %}
<div class="id-field-block">
<span class="id-label">Goals:</span> {{ contact.goals }}
</div>
{% endif %}
</div>
<div class="id-card-right">
{% if contact.bio %}
<div class="id-bio">
<span class="id-label">Bio:</span> {{ contact.bio }}
</div>
{% endif %}
<div class="id-relationships">
<h2 class="id-section-title">Relationships</h2>
{% if grouped_relationships.familial %}
<div class="id-rel-group">
<span class="id-rel-label">Familial:</span>
{% for rel in grouped_relationships.familial %}
<a href="/contacts/{{ rel.related_contact_id }}">{{ contact_names[rel.related_contact_id] }}</a><span class="id-rel-type">({{ rel.relationship_type|replace("_", " ")|title }})</span>{% if not loop.last %}, {% endif %}
{% endfor %}
</div>
{% endif %}
{% if grouped_relationships.partners %}
<div class="id-rel-group">
<span class="id-rel-label">Partners:</span>
{% for rel in grouped_relationships.partners %}
<a href="/contacts/{{ rel.related_contact_id }}">{{ contact_names[rel.related_contact_id] }}</a>{% if not loop.last %}, {% endif %}
{% endfor %}
</div>
{% endif %}
{% if grouped_relationships.friends %}
<div class="id-rel-group">
<span class="id-rel-label">Friends:</span>
{% for rel in grouped_relationships.friends %}
<a href="/contacts/{{ rel.related_contact_id }}">{{ contact_names[rel.related_contact_id] }}</a>{% if not loop.last %}, {% endif %}
{% endfor %}
</div>
{% endif %}
{% if grouped_relationships.professional %}
<div class="id-rel-group">
<span class="id-rel-label">Professional:</span>
{% for rel in grouped_relationships.professional %}
<a href="/contacts/{{ rel.related_contact_id }}">{{ contact_names[rel.related_contact_id] }}</a><span class="id-rel-type">({{ rel.relationship_type|replace("_", " ")|title }})</span>{% if not loop.last %}, {% endif %}
{% endfor %}
</div>
{% endif %}
{% if grouped_relationships.other %}
<div class="id-rel-group">
<span class="id-rel-label">Other:</span>
{% for rel in grouped_relationships.other %}
<a href="/contacts/{{ rel.related_contact_id }}">{{ contact_names[rel.related_contact_id] }}</a><span class="id-rel-type">({{ rel.relationship_type|replace("_", " ")|title }})</span>{% if not loop.last %}, {% endif %}
{% endfor %}
</div>
{% endif %}
{% if contact.related_from %}
<div class="id-rel-group">
<span class="id-rel-label">Known by:</span>
{% for rel in contact.related_from %}
<a href="/contacts/{{ rel.contact_id }}">{{ contact_names[rel.contact_id] }}</a>{% if not loop.last %}, {% endif %}
{% endfor %}
</div>
{% endif %}
</div>
</div>
</div>
{% if contact.needs %}
<div class="id-card-warnings">
{% for need in contact.needs %}
<div class="id-warning">
<span class="warning-dot"></span>
Warning: {{ need.name }}
{% if need.description %}<span class="warning-desc"> - {{ need.description }}</span>{% endif %}
</div>
{% endfor %}
</div>
{% endif %}
</div>
<details class="id-card-manage">
<summary>Manage Contact</summary>
<div class="manage-section">
<h3>Manage Relationships</h3>
<div id="manage-relationships" class="manage-relationships">
{% include "partials/manage_relationships.html" %}
</div>
{% if all_contacts %}
<form hx-post="/htmx/contacts/{{ contact.id }}/add-relationship"
hx-target="#manage-relationships"
hx-swap="innerHTML"
class="add-form">
<select name="related_contact_id" required>
<option value="">Select contact...</option>
{% for other in all_contacts %}
{% if other.id != contact.id %}
<option value="{{ other.id }}">{{ other.name }}</option>
{% endif %}
{% endfor %}
</select>
<select name="relationship_type" required>
<option value="">Select relationship type...</option>
{% for rel_type in relationship_types %}
<option value="{{ rel_type.value }}">{{ rel_type.display_name }}</option>
{% endfor %}
</select>
<button type="submit" class="btn btn-primary">Add Relationship</button>
</form>
{% endif %}
</div>
<div class="manage-section">
<h3>Manage Needs/Warnings</h3>
<div id="manage-needs">
{% include "partials/manage_needs.html" %}
</div>
{% if available_needs %}
<form hx-post="/htmx/contacts/{{ contact.id }}/add-need"
hx-target="#manage-needs"
hx-swap="innerHTML"
class="add-form">
<select name="need_id" required>
<option value="">Select a need...</option>
{% for need in available_needs %}
<option value="{{ need.id }}">{{ need.name }}</option>
{% endfor %}
</select>
<button type="submit" class="btn btn-primary">Add Need</button>
</form>
{% endif %}
</div>
</details>
</div>
{% endblock %}
-115
View File
@@ -1,115 +0,0 @@
{% extends "base.html" %}
{% block title %}{{ "Edit " + contact.name if contact else "New Contact" }}{% endblock %}
{% block content %}
<div class="contact-form">
<h1>{{ "Edit Contact" if contact else "New Contact" }}</h1>
{% if contact %}
<form method="post" action="/htmx/contacts/{{ contact.id }}/edit">
{% else %}
<form method="post" action="/htmx/contacts/new">
{% endif %}
<div class="form-group">
<label for="name">Name *</label>
<input id="name" name="name" type="text" value="{{ contact.name if contact else '' }}" required>
</div>
<div class="form-row">
<div class="form-group">
<label for="legal_name">Legal Name</label>
<input id="legal_name" name="legal_name" type="text" value="{{ contact.legal_name or '' }}">
</div>
<div class="form-group">
<label for="suffix">Suffix</label>
<input id="suffix" name="suffix" type="text" value="{{ contact.suffix or '' }}">
</div>
</div>
<div class="form-row">
<div class="form-group">
<label for="age">Age</label>
<input id="age" name="age" type="number" value="{{ contact.age if contact and contact.age is not none else '' }}">
</div>
<div class="form-group">
<label for="gender">Gender</label>
<input id="gender" name="gender" type="text" value="{{ contact.gender or '' }}">
</div>
</div>
<div class="form-group">
<label for="current_job">Current Job</label>
<input id="current_job" name="current_job" type="text" value="{{ contact.current_job or '' }}">
</div>
<div class="form-group">
<label for="timezone">Timezone</label>
<input id="timezone" name="timezone" type="text" value="{{ contact.timezone or '' }}">
</div>
<div class="form-group">
<label for="profile_pic">Profile Picture URL</label>
<input id="profile_pic" name="profile_pic" type="url" placeholder="https://example.com/photo.jpg" value="{{ contact.profile_pic or '' }}">
</div>
<div class="form-group">
<label for="bio">Bio</label>
<textarea id="bio" name="bio" rows="3">{{ contact.bio or '' }}</textarea>
</div>
<div class="form-group">
<label for="goals">Goals</label>
<textarea id="goals" name="goals" rows="3">{{ contact.goals or '' }}</textarea>
</div>
<div class="form-group">
<label for="social_structure_style">Social Structure Style</label>
<input id="social_structure_style" name="social_structure_style" type="text" value="{{ contact.social_structure_style or '' }}">
</div>
<div class="form-group">
<label for="self_sufficiency_score">Self-Sufficiency Score (1-10)</label>
<input id="self_sufficiency_score" name="self_sufficiency_score" type="number" min="1" max="10" value="{{ contact.self_sufficiency_score if contact and contact.self_sufficiency_score is not none else '' }}">
</div>
<div class="form-group">
<label for="safe_conversation_starters">Safe Conversation Starters</label>
<textarea id="safe_conversation_starters" name="safe_conversation_starters" rows="2">{{ contact.safe_conversation_starters or '' }}</textarea>
</div>
<div class="form-group">
<label for="topics_to_avoid">Topics to Avoid</label>
<textarea id="topics_to_avoid" name="topics_to_avoid" rows="2">{{ contact.topics_to_avoid or '' }}</textarea>
</div>
<div class="form-group">
<label for="ssn">SSN</label>
<input id="ssn" name="ssn" type="text" value="{{ contact.ssn or '' }}">
</div>
{% if all_needs %}
<div class="form-group">
<label>Needs/Accommodations</label>
<div class="checkbox-group">
{% for need in all_needs %}
<label class="checkbox-label">
<input type="checkbox" name="need_ids" value="{{ need.id }}"
{% if contact and need in contact.needs %}checked{% endif %}>
{{ need.name }}
</label>
{% endfor %}
</div>
</div>
{% endif %}
<div class="form-actions">
<button type="submit" class="btn btn-primary">Save</button>
{% if contact %}
<a href="/contacts/{{ contact.id }}" class="btn">Cancel</a>
{% else %}
<a href="/contacts" class="btn">Cancel</a>
{% endif %}
</div>
</form>
</div>
{% endblock %}
-14
View File
@@ -1,14 +0,0 @@
{% extends "base.html" %}
{% block title %}Contacts{% endblock %}
{% block content %}
<div class="contact-list">
<div class="header">
<h1>Contacts</h1>
<a href="/contacts/new" class="btn btn-primary">Add Contact</a>
</div>
<div id="contact-table">
{% include "partials/contact_table.html" %}
</div>
</div>
{% endblock %}
-198
View File
@@ -1,198 +0,0 @@
{% extends "base.html" %}
{% block title %}Relationship Graph{% endblock %}
{% block content %}
<div class="graph-container">
<div class="header">
<h1>Relationship Graph</h1>
</div>
<p class="graph-hint">Drag nodes to reposition. Closer relationships have shorter, darker edges.</p>
<canvas id="graph-canvas" width="900" height="600"
style="border: 1px solid var(--color-border); border-radius: 8px; background: var(--color-bg); cursor: grab;">
</canvas>
<div id="selected-info"></div>
<div class="legend">
<h4>Relationship Closeness (1-10)</h4>
<div class="legend-items">
<div class="legend-item">
<span class="legend-line" style="background: hsl(220, 70%, 40%); height: 4px; display: inline-block;"></span>
<span>10 - Very Close (Spouse, Partner)</span>
</div>
<div class="legend-item">
<span class="legend-line" style="background: hsl(220, 70%, 52%); height: 3px; display: inline-block;"></span>
<span>7 - Close (Family, Best Friend)</span>
</div>
<div class="legend-item">
<span class="legend-line" style="background: hsl(220, 70%, 64%); height: 2px; display: inline-block;"></span>
<span>4 - Moderate (Friend, Colleague)</span>
</div>
<div class="legend-item">
<span class="legend-line" style="background: hsl(220, 70%, 72%); height: 1px; display: inline-block;"></span>
<span>2 - Distant (Acquaintance)</span>
</div>
</div>
</div>
</div>
<script>
(function() {
const RELATIONSHIP_DISPLAY = {{ relationship_type_display|tojson }};
const graphData = {{ graph_data|tojson }};
const canvas = document.getElementById('graph-canvas');
const ctx = canvas.getContext('2d');
const width = canvas.width;
const height = canvas.height;
const centerX = width / 2;
const centerY = height / 2;
const nodes = graphData.nodes.map(function(node) {
return Object.assign({}, node, {
x: centerX + (Math.random() - 0.5) * 300,
y: centerY + (Math.random() - 0.5) * 300,
vx: 0,
vy: 0
});
});
const nodeMap = new Map(nodes.map(function(node) { return [node.id, node]; }));
const edges = graphData.edges.map(function(edge) {
const sourceNode = nodeMap.get(edge.source);
const targetNode = nodeMap.get(edge.target);
if (!sourceNode || !targetNode) return null;
return Object.assign({}, edge, { sourceNode: sourceNode, targetNode: targetNode });
}).filter(function(edge) { return edge !== null; });
let dragNode = null;
let selectedNode = null;
const repulsion = 5000;
const springStrength = 0.05;
const baseSpringLength = 150;
const damping = 0.9;
const centerPull = 0.01;
function simulate() {
for (const node of nodes) { node.vx = 0; node.vy = 0; }
for (let i = 0; i < nodes.length; i++) {
for (let j = i + 1; j < nodes.length; j++) {
const dx = nodes[j].x - nodes[i].x;
const dy = nodes[j].y - nodes[i].y;
const dist = Math.sqrt(dx * dx + dy * dy) || 1;
const force = repulsion / (dist * dist);
const fx = (dx / dist) * force;
const fy = (dy / dist) * force;
nodes[i].vx -= fx; nodes[i].vy -= fy;
nodes[j].vx += fx; nodes[j].vy += fy;
}
}
for (const edge of edges) {
const dx = edge.targetNode.x - edge.sourceNode.x;
const dy = edge.targetNode.y - edge.sourceNode.y;
const dist = Math.sqrt(dx * dx + dy * dy) || 1;
const normalizedWeight = edge.closeness_weight / 10;
const idealLength = baseSpringLength * (1.5 - normalizedWeight);
const displacement = dist - idealLength;
const force = springStrength * displacement;
const fx = (dx / dist) * force;
const fy = (dy / dist) * force;
edge.sourceNode.vx += fx; edge.sourceNode.vy += fy;
edge.targetNode.vx -= fx; edge.targetNode.vy -= fy;
}
for (const node of nodes) {
node.vx += (centerX - node.x) * centerPull;
node.vy += (centerY - node.y) * centerPull;
}
for (const node of nodes) {
if (node === dragNode) continue;
node.x += node.vx * damping;
node.y += node.vy * damping;
node.x = Math.max(30, Math.min(width - 30, node.x));
node.y = Math.max(30, Math.min(height - 30, node.y));
}
}
function getEdgeColor(weight) {
const normalized = weight / 10;
return 'hsl(220, 70%, ' + (80 - normalized * 40) + '%)';
}
function draw() {
ctx.clearRect(0, 0, width, height);
for (const edge of edges) {
const lineWidth = 1 + (edge.closeness_weight / 10) * 3;
ctx.strokeStyle = getEdgeColor(edge.closeness_weight);
ctx.lineWidth = lineWidth;
ctx.beginPath();
ctx.moveTo(edge.sourceNode.x, edge.sourceNode.y);
ctx.lineTo(edge.targetNode.x, edge.targetNode.y);
ctx.stroke();
const midX = (edge.sourceNode.x + edge.targetNode.x) / 2;
const midY = (edge.sourceNode.y + edge.targetNode.y) / 2;
ctx.fillStyle = '#666';
ctx.font = '10px sans-serif';
ctx.textAlign = 'center';
const label = RELATIONSHIP_DISPLAY[edge.relationship_type] || edge.relationship_type;
ctx.fillText(label, midX, midY - 5);
}
for (const node of nodes) {
const isSelected = node === selectedNode;
const radius = isSelected ? 25 : 20;
ctx.beginPath();
ctx.arc(node.x, node.y, radius, 0, Math.PI * 2);
ctx.fillStyle = isSelected ? '#0066cc' : '#fff';
ctx.fill();
ctx.strokeStyle = '#0066cc';
ctx.lineWidth = 2;
ctx.stroke();
ctx.fillStyle = isSelected ? '#fff' : '#333';
ctx.font = '12px sans-serif';
ctx.textAlign = 'center';
ctx.textBaseline = 'middle';
const name = node.name.length > 10 ? node.name.slice(0, 9) + '\u2026' : node.name;
ctx.fillText(name, node.x, node.y);
}
}
function animate() {
simulate();
draw();
requestAnimationFrame(animate);
}
animate();
function getNodeAt(x, y) {
for (const node of nodes) {
const dx = x - node.x;
const dy = y - node.y;
if (dx * dx + dy * dy < 400) return node;
}
return null;
}
canvas.addEventListener('mousedown', function(event) {
const rect = canvas.getBoundingClientRect();
const node = getNodeAt(event.clientX - rect.left, event.clientY - rect.top);
if (node) {
dragNode = node;
selectedNode = node;
const infoDiv = document.getElementById('selected-info');
let html = '<div class="selected-info"><h3>' + node.name + '</h3>';
if (node.current_job) html += '<p>Job: ' + node.current_job + '</p>';
html += '<a href="/contacts/' + node.id + '">View details</a></div>';
infoDiv.innerHTML = html;
}
});
canvas.addEventListener('mousemove', function(event) {
if (!dragNode) return;
const rect = canvas.getBoundingClientRect();
dragNode.x = event.clientX - rect.left;
dragNode.y = event.clientY - rect.top;
});
canvas.addEventListener('mouseup', function() { dragNode = null; });
canvas.addEventListener('mouseleave', function() { dragNode = null; });
})();
</script>
{% endblock %}
-31
View File
@@ -1,31 +0,0 @@
{% extends "base.html" %}
{% block title %}Needs{% endblock %}
{% block content %}
<div class="need-list">
<div class="header">
<h1>Needs / Accommodations</h1>
<button class="btn btn-primary" onclick="document.getElementById('need-form').toggleAttribute('hidden')">Add Need</button>
</div>
<form id="need-form" hidden
hx-post="/htmx/needs"
hx-target="#need-items"
hx-swap="innerHTML"
hx-on::after-request="if(event.detail.successful) this.reset()"
class="need-form">
<div class="form-group">
<label for="name">Name *</label>
<input id="name" name="name" type="text" placeholder="e.g., Light Sensitive, ADHD" required>
</div>
<div class="form-group">
<label for="description">Description</label>
<textarea id="description" name="description" placeholder="Optional description..." rows="2"></textarea>
</div>
<button type="submit" class="btn btn-primary">Create</button>
</form>
<div id="need-items">
{% include "partials/need_items.html" %}
</div>
</div>
{% endblock %}
@@ -1,33 +0,0 @@
{% if contacts %}
<table>
<thead>
<tr>
<th>Name</th>
<th>Job</th>
<th>Timezone</th>
<th>Actions</th>
</tr>
</thead>
<tbody>
{% for contact in contacts %}
<tr id="contact-row-{{ contact.id }}">
<td><a href="/contacts/{{ contact.id }}">{{ contact.name }}</a></td>
<td>{{ contact.current_job or "-" }}</td>
<td>{{ contact.timezone or "-" }}</td>
<td>
<a href="/contacts/{{ contact.id }}/edit" class="btn">Edit</a>
<button class="btn btn-danger"
hx-delete="/api/contacts/{{ contact.id }}"
hx-target="#contact-row-{{ contact.id }}"
hx-swap="outerHTML"
hx-confirm="Delete this contact?">
Delete
</button>
</td>
</tr>
{% endfor %}
</tbody>
</table>
{% else %}
<p>No contacts yet.</p>
{% endif %}
@@ -1,14 +0,0 @@
<ul class="manage-needs-list">
{% for need in contact.needs %}
<li id="contact-need-{{ need.id }}">
<strong>{{ need.name }}</strong>
{% if need.description %}<span> - {{ need.description }}</span>{% endif %}
<button class="btn btn-small btn-danger"
hx-delete="/api/contacts/{{ contact.id }}/needs/{{ need.id }}"
hx-target="#contact-need-{{ need.id }}"
hx-swap="outerHTML">
Remove
</button>
</li>
{% endfor %}
</ul>
@@ -1,23 +0,0 @@
{% for rel in contact.related_to %}
<div class="manage-rel-item" id="rel-{{ contact.id }}-{{ rel.related_contact_id }}">
<a href="/contacts/{{ rel.related_contact_id }}">{{ contact_names[rel.related_contact_id] }}</a>
<span class="tag">{{ rel.relationship_type|replace("_", " ")|title }}</span>
<label class="weight-control">
<span>Closeness:</span>
<input type="range" min="1" max="10" value="{{ rel.closeness_weight }}"
hx-post="/htmx/contacts/{{ contact.id }}/relationships/{{ rel.related_contact_id }}/weight"
hx-trigger="change"
hx-include="this"
name="closeness_weight"
hx-swap="none"
oninput="this.nextElementSibling.textContent = this.value">
<span class="weight-value">{{ rel.closeness_weight }}</span>
</label>
<button class="btn btn-small btn-danger"
hx-delete="/api/contacts/{{ contact.id }}/relationships/{{ rel.related_contact_id }}"
hx-target="#rel-{{ contact.id }}-{{ rel.related_contact_id }}"
hx-swap="outerHTML">
Remove
</button>
</div>
{% endfor %}
@@ -1,21 +0,0 @@
{% if needs %}
<ul class="need-items">
{% for need in needs %}
<li id="need-item-{{ need.id }}">
<div class="need-info">
<strong>{{ need.name }}</strong>
{% if need.description %}<p>{{ need.description }}</p>{% endif %}
</div>
<button class="btn btn-danger"
hx-delete="/api/needs/{{ need.id }}"
hx-target="#need-item-{{ need.id }}"
hx-swap="outerHTML"
hx-confirm="Delete this need?">
Delete
</button>
</li>
{% endfor %}
</ul>
{% else %}
<p>No needs defined yet.</p>
{% endif %}
+2 -10
View File
@@ -8,20 +8,13 @@ from typing import TYPE_CHECKING
from python.ebook_search.llm_interface import request_chat_completion
if TYPE_CHECKING:
import httpx
from python.ebook_search.config import EbookSearchConfig
from python.ebook_search.search import SearchResult
logger = logging.getLogger(__name__)
async def answer_query(
client: httpx.AsyncClient,
query: str,
results: list[SearchResult],
config: EbookSearchConfig,
) -> str:
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")
@@ -42,8 +35,7 @@ async def answer_query(
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 = await request_chat_completion(
client,
content = request_chat_completion(
config,
[
{
+19 -32
View File
@@ -1,22 +1,18 @@
"""Background BM25 refresh tasks for the web app.
The refresh is scheduled on the event loop instead of a thread because the async psycopg
driver only works from the loop; a bare thread cannot open a session on the async engine.
"""
"""Background BM25 refresh tasks for the web app."""
from __future__ import annotations
import asyncio
import logging
from threading import Timer
from typing import TYPE_CHECKING
from sqlalchemy.ext.asyncio import AsyncSession
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.ext.asyncio import AsyncEngine
from sqlalchemy.engine import Engine
from python.ebook_search.config import EbookSearchConfig
@@ -24,18 +20,15 @@ logger = logging.getLogger(__name__)
def schedule_bm25_refresh(app: FastAPI) -> None:
"""Schedule a delayed BM25 corpus refresh, replacing any pending refresh.
"""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()
Only called from route handlers, so a running event loop is guaranteed.
"""
cancel_bm25_refresh(app)
loop = asyncio.get_running_loop()
def start_refresh() -> None:
app.state.bm25_refresh_task = loop.create_task(refresh_bm25_for_app(app))
app.state.bm25_refresh_timer = loop.call_later(app.state.config.bm25_refresh_delay_seconds, start_refresh)
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,
@@ -43,31 +36,25 @@ def schedule_bm25_refresh(app: FastAPI) -> None:
def cancel_bm25_refresh(app: FastAPI) -> None:
"""Cancel any pending BM25 corpus refresh timer and in-flight refresh task."""
"""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")
existing_task = getattr(app.state, "bm25_refresh_task", None)
if existing_task is not None:
if not existing_task.done():
existing_task.cancel()
app.state.bm25_refresh_task = None
async def refresh_bm25_for_app(app: FastAPI) -> None:
def refresh_bm25_for_app(app: FastAPI) -> None:
"""Refresh the BM25 corpus using the app engine and config."""
try:
await refresh_bm25_for_engine(app.state.engine, app.state.config)
refresh_bm25_for_engine(app.state.engine, app.state.config)
except Exception:
logger.exception("ebook_bm25_refresh_failed")
async def refresh_bm25_for_engine(engine: AsyncEngine, config: EbookSearchConfig) -> None:
"""Refresh the BM25 corpus using an async SQLAlchemy engine."""
async with AsyncSession(engine) as session:
await refresh_bm25_corpus(session, config)
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")
+3 -10
View File
@@ -4,9 +4,8 @@ from __future__ import annotations
from typing import Annotated
import httpx
from fastapi import Depends, Request
from sqlalchemy.ext.asyncio import AsyncEngine
from sqlalchemy.engine import Engine
from python.ebook_search.config import EbookSearchConfig
@@ -16,16 +15,10 @@ def get_config(request: Request) -> EbookSearchConfig:
return request.app.state.config
def get_engine(request: Request) -> AsyncEngine:
def get_engine(request: Request) -> Engine:
"""Get the database engine from app state."""
return request.app.state.engine
def get_http_client(request: Request) -> httpx.AsyncClient:
"""Get the shared LLM HTTP client from app state."""
return request.app.state.http_client
AppConfig = Annotated[EbookSearchConfig, Depends(get_config)]
AppEngine = Annotated[AsyncEngine, Depends(get_engine)]
AppHttpClient = Annotated[httpx.AsyncClient, Depends(get_http_client)]
AppEngine = Annotated[Engine, Depends(get_engine)]
-131
View File
@@ -1,131 +0,0 @@
"""Background phrase-judging tasks for the web app.
Judging a book sends one LLM request per candidate phrase, which can take minutes, so it must
not run inside the request where it would block the UI. Judgments run as async FastAPI
background tasks, awaited on the event loop after the response is sent, and are tracked per
book in app state so a second judge request for a book that is already being judged is
rejected instead of doubling the work.
State is loop-confined: every read and mutation happens on the event loop (async route
handlers and async background tasks) and no critical section contains an ``await``, so each
mutation is atomic per loop iteration and no locking is needed.
"""
from __future__ import annotations
import logging
from dataclasses import dataclass, field
from typing import TYPE_CHECKING
from python.ebook_search.protected_phrases.judge_ngrams import judge_candidate_phrases_for_books
if TYPE_CHECKING:
from fastapi import BackgroundTasks, FastAPI
logger = logging.getLogger(__name__)
@dataclass
class JudgeTaskState:
"""Running book judgments and last outcome messages, keyed by book id."""
running_book_ids: set[int] = field(default_factory=set)
outcome_messages: dict[int, str] = field(default_factory=dict)
def get_judge_task_state(app: FastAPI) -> JudgeTaskState:
"""Return the app's judge task state, creating it on first use.
Args:
app (FastAPI): App whose state holds the judge task registry.
Returns:
JudgeTaskState: The shared judge task state for this app.
"""
state = getattr(app.state, "judge_tasks", None)
if state is None:
state = JudgeTaskState()
app.state.judge_tasks = state
return state
def start_book_phrase_judgment(app: FastAPI, background_tasks: BackgroundTasks, source_id: int) -> bool:
"""Queue judging of one book's candidate phrases as a FastAPI background task.
The book is claimed before the response returns, so a repeated judge request cannot queue
a second run while one is pending or running.
Args:
app (FastAPI): App supplying the engine, config, and judge task state.
background_tasks (BackgroundTasks): Request's background tasks to queue the judgment on.
source_id (int): Book to judge candidates for.
Returns:
bool: True when a judgment was queued, False when one is already running for this book.
"""
state = get_judge_task_state(app)
if source_id in state.running_book_ids:
logger.info("ebook_book_phrase_judgment_already_running source_id=%s", source_id)
return False
state.running_book_ids.add(source_id)
state.outcome_messages.pop(source_id, None)
background_tasks.add_task(judge_book_phrases_for_app, app, source_id)
logger.info("ebook_book_phrase_judgment_queued source_id=%s", source_id)
return True
async def judge_book_phrases_for_app(app: FastAPI, source_id: int) -> None:
"""Judge one book using the app engine and config, recording the outcome message.
Args:
app (FastAPI): App supplying the engine, config, and judge task state.
source_id (int): Book to judge candidates for.
"""
state = get_judge_task_state(app)
try:
result = await judge_candidate_phrases_for_books(app.state.engine, app.state.config, source_ids=[source_id])
logger.info(
"ebook_book_phrase_judgment_complete source_id=%s judged=%s protected=%s mentions=%s failed=%s",
source_id,
result.candidates_judged,
result.protected_phrases,
result.phrase_mentions,
result.books_failed,
)
if result.books_failed:
message = "Judging failed; see server logs for details"
else:
message = (
f"Judged {result.candidates_judged} candidates; {result.protected_phrases} protected phrases promoted"
)
except Exception:
logger.exception("ebook_book_phrase_judgment_task_failed source_id=%s", source_id)
message = "Judging failed; see server logs for details"
state.running_book_ids.discard(source_id)
state.outcome_messages[source_id] = message
def is_judging_book(app: FastAPI, source_id: int) -> bool:
"""Report whether a judgment is currently queued or running for one book.
Args:
app (FastAPI): App supplying the judge task state.
source_id (int): Book to check.
Returns:
bool: True while the book's judgment is pending or running.
"""
return source_id in get_judge_task_state(app).running_book_ids
def pop_book_judgment_outcome(app: FastAPI, source_id: int) -> str | None:
"""Return and clear the outcome message from one book's last finished judgment.
Args:
app (FastAPI): App supplying the judge task state.
source_id (int): Book to fetch the outcome for.
Returns:
str | None: The outcome message, or None when there is nothing new to report.
"""
return get_judge_task_state(app).outcome_messages.pop(source_id, None)
+7 -19
View File
@@ -6,12 +6,11 @@ import logging
from contextlib import asynccontextmanager
from typing import TYPE_CHECKING, Annotated
import httpx
import typer
import uvicorn
from fastapi import FastAPI
from fastapi.staticfiles import StaticFiles
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import Session
from python.common import configure_logger
from python.ebook_search.api.bm25_tasks import cancel_bm25_refresh
@@ -19,9 +18,8 @@ from python.ebook_search.api.routes import admin_router, health_router, page_rou
from python.ebook_search.api.web import STATIC_DIR
from python.ebook_search.bm25_corpus import ensure_bm25_corpus
from python.ebook_search.config import load_config
from python.ebook_search.protected_phrases.pool import shutdown_extraction_pool
from python.fastapi_tools import ZstdMiddleware
from python.orm.common import get_async_postgres_engine
from python.orm.common import get_postgres_engine
if TYPE_CHECKING:
from collections.abc import AsyncIterator
@@ -38,36 +36,26 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]:
app.state.config = config
logger.info(
"ebook_search_config_loaded top_k=%s embedding_model=%s embedding_base_url=%s vllm_base_url=%s "
"rerank_enabled=%s phrase_matching_enabled=%s answer_enabled=%s library_paths=%s",
"rerank_enabled=%s answer_enabled=%s library_paths=%s",
config.top_k,
config.embedding_model,
config.embedding_base_url,
config.vllm_base_url,
config.rerank.enabled,
config.phrase_matching_enabled,
config.answer_enabled,
len(config.library_paths),
)
if not config.library_paths:
logger.warning("ebook_search_no_library_paths_configured")
# Concurrent phrase judging opens one session per book worker on this engine, so size the pool
# to cover those plus headroom for ordinary web requests.
app.state.engine = get_async_postgres_engine(
name="RICHIE",
vector_engine=True,
pool_size=config.phrase_judge_book_workers + 10,
)
app.state.http_client = httpx.AsyncClient()
async with AsyncSession(app.state.engine, expire_on_commit=False) as session:
await ensure_bm25_corpus(session, config)
app.state.engine = get_postgres_engine(name="RICHIE", vector_engine=True)
with Session(app.state.engine) as session:
ensure_bm25_corpus(session, config)
try:
yield
finally:
logger.info("ebook_search_shutdown")
cancel_bm25_refresh(app)
shutdown_extraction_pool()
await app.state.http_client.aclose()
await app.state.engine.dispose()
app.state.engine.dispose()
def create_app() -> FastAPI:
+16 -176
View File
@@ -8,18 +8,13 @@ from fastapi import APIRouter, Request
from fastapi.responses import HTMLResponse
from python.ebook_search.api.bm25_tasks import schedule_bm25_refresh
from python.ebook_search.api.dependencies import ( # noqa: TC001 FastAPI resolves these annotated dependencies at runtime
AppConfig,
AppEngine,
AppHttpClient,
from python.ebook_search.api.dependencies import (
AppConfig, # noqa: TC001 FastAPI resolves this annotated dependency at runtime
)
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
from python.ebook_search.protected_phrases.generate_ngrams import generate_candidate_phrases_for_books
from python.ebook_search.protected_phrases.judge_ngrams import judge_candidate_phrases_for_books
from python.ebook_search.protected_phrases.store import book_ids_pending_first_judgment, corpus_phrase_stats
from python.fastapi_tools import AsyncDbSession # noqa: TC001 FastAPI resolves this annotated dependency at runtime
from python.fastapi_tools import DbSession # noqa: TC001 FastAPI resolves this annotated dependency at runtime
logger = logging.getLogger(__name__)
@@ -27,29 +22,19 @@ router = APIRouter(prefix="/admin")
@router.get("", response_class=HTMLResponse)
async def admin(request: Request, config: AppConfig, session: AsyncDbSession) -> HTMLResponse:
def admin(request: Request, config: AppConfig, session: DbSession) -> HTMLResponse:
"""Render the admin page."""
stats = await embedding_model_stats(session)
phrase_stats = await corpus_phrase_stats(session)
logger.info(
"ebook_admin_page_loaded models=%s candidate_phrases=%s protected_phrases=%s",
len(stats),
phrase_stats.candidate_phrases,
phrase_stats.protected_phrases,
)
return templates.TemplateResponse(
request,
"admin.html",
{"config": config, "stats": stats, "phrase_stats": phrase_stats},
)
stats = embedding_model_stats(session)
logger.info("ebook_admin_page_loaded models=%s", len(stats))
return templates.TemplateResponse(request, "admin.html", {"config": config, "stats": stats})
@router.post("/scan", response_class=HTMLResponse)
async def scan_library(request: Request, config: AppConfig, session: AsyncDbSession) -> HTMLResponse:
def scan_library(request: Request, config: AppConfig, session: DbSession) -> HTMLResponse:
"""Scan configured library paths for EPUB changes."""
try:
count = await ingest_configured_paths(session, config)
await session.commit()
count = ingest_configured_paths(session, 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)
@@ -60,152 +45,12 @@ async def scan_library(request: Request, config: AppConfig, session: AsyncDbSess
return templates.TemplateResponse(request, "partials/admin_status.html", {"message": f"Indexed {count} EPUBs"})
@router.post("/phrases/generate-all", response_class=HTMLResponse)
async def generate_all_phrases(request: Request, config: AppConfig, session: AsyncDbSession) -> HTMLResponse:
"""Regenerate candidate phrases for every indexed book without LLM judging."""
return await run_phrase_generation(request, config, session, only_missing=False)
@router.post("/phrases/generate-missing", response_class=HTMLResponse)
async def generate_missing_phrases(request: Request, config: AppConfig, session: AsyncDbSession) -> HTMLResponse:
"""Generate candidate phrases only for books that have none yet."""
return await run_phrase_generation(request, config, session, only_missing=True)
async def run_phrase_generation(
request: Request,
config: AppConfig,
session: AsyncDbSession,
*,
only_missing: bool,
) -> HTMLResponse:
"""Run candidate phrase generation and render the outcome as an admin status partial.
Args:
request (Request): Current request, for template rendering.
config (AppConfig): Runtime phrase-tuning settings.
session (AsyncDbSession): Active database session.
only_missing (bool): Only generate for books without candidates instead of every book.
Returns:
HTMLResponse: Status partial describing the generation outcome.
"""
try:
result = await generate_candidate_phrases_for_books(session, config, only_missing=only_missing)
await session.commit()
except Exception as error:
await session.rollback()
logger.exception("ebook_admin_generate_phrases_failed only_missing=%s", only_missing)
return templates.TemplateResponse(request, "partials/error.html", {"message": str(error)}, status_code=500)
logger.info(
"ebook_admin_generate_phrases_complete only_missing=%s books_seen=%s books_built=%s candidates=%s",
only_missing,
result.books_seen,
result.books_built,
result.candidate_phrases,
)
if only_missing and result.books_seen == 0:
return templates.TemplateResponse(
request,
"partials/admin_status.html",
{"message": "All books already have candidate phrases"},
)
return templates.TemplateResponse(
request,
"partials/admin_status.html",
{
"message": (
f"Generated phrases for {result.books_built} of {result.books_seen} books; "
f"{result.candidate_phrases} candidates stored"
)
},
)
@router.post("/phrases/judge-all", response_class=HTMLResponse)
async def judge_all_phrases(request: Request, engine: AppEngine, config: AppConfig) -> HTMLResponse:
"""Judge unjudged candidate phrases across every indexed book."""
return await run_phrase_judgment(request, engine, config, source_ids=None)
@router.post("/phrases/judge-missing", response_class=HTMLResponse)
async def judge_missing_phrases(
request: Request,
engine: AppEngine,
config: AppConfig,
session: AsyncDbSession,
) -> HTMLResponse:
"""Judge candidate phrases only for books where judging has never run."""
source_ids = await book_ids_pending_first_judgment(session)
if not source_ids:
return templates.TemplateResponse(
request,
"partials/admin_status.html",
{"message": "All books with candidate phrases have been judged"},
)
return await run_phrase_judgment(request, engine, config, source_ids=source_ids)
async def run_phrase_judgment(
request: Request,
engine: AppEngine,
config: AppConfig,
*,
source_ids: list[int] | None,
) -> HTMLResponse:
"""Run LLM judging for candidate phrases and render the outcome as an admin status partial.
Args:
request (Request): Current request, for template rendering.
engine (AppEngine): Engine used to open per-book judging sessions.
config (AppConfig): Runtime phrase-tuning settings.
source_ids (list[int] | None): Books to judge; ``None`` judges every indexed book.
Returns:
HTMLResponse: Status partial describing the judging outcome.
"""
try:
result = await judge_candidate_phrases_for_books(engine, config, source_ids=source_ids)
except Exception as error:
logger.exception("ebook_admin_judge_phrases_failed")
return templates.TemplateResponse(request, "partials/error.html", {"message": str(error)}, status_code=500)
logger.info(
"ebook_admin_judge_phrases_complete books_seen=%s books_judged=%s books_failed=%s candidates_judged=%s "
"protected=%s mentions=%s",
result.books_seen,
result.books_judged,
result.books_failed,
result.candidates_judged,
result.protected_phrases,
result.phrase_mentions,
)
return templates.TemplateResponse(
request,
"partials/admin_status.html",
{
"message": (
f"Judged {result.candidates_judged} candidates across {result.books_judged} of "
f"{result.books_seen} books; {result.protected_phrases} protected phrases, "
f"{result.phrase_mentions} mentions"
+ (f"; {result.books_failed} books failed" if result.books_failed else "")
)
},
)
@router.post("/embed-missing", response_class=HTMLResponse)
async def embed_missing(
request: Request,
config: AppConfig,
session: AsyncDbSession,
client: AppHttpClient,
) -> HTMLResponse:
def embed_missing(request: Request, config: AppConfig, session: DbSession) -> HTMLResponse:
"""Embed chunks missing vectors for the configured model."""
try:
count = await embed_missing_chunks(session, client, config)
await session.commit()
count = embed_missing_chunks(session, 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)
@@ -219,21 +64,16 @@ async def embed_missing(
@router.post("/embed-all", response_class=HTMLResponse)
async def embed_all(
request: Request,
config: AppConfig,
session: AsyncDbSession,
client: AppHttpClient,
) -> HTMLResponse:
def embed_all(request: Request, config: AppConfig, session: DbSession) -> HTMLResponse:
"""Embed all chunks missing vectors in fixed-size batches."""
total = 0
batches = 0
try:
while True:
count = await embed_missing_chunks(session, client, config)
count = embed_missing_chunks(session, config)
if count == 0:
break
await session.commit()
session.commit()
total += count
batches += 1
logger.info(
+13 -15
View File
@@ -11,17 +11,15 @@ from fastapi.responses import JSONResponse
from sqlalchemy import literal, select
from sqlalchemy.exc import SQLAlchemyError
from python.ebook_search.api.dependencies import ( # noqa: TC001 FastAPI resolves these annotated dependencies at runtime
AppConfig,
AppHttpClient,
from python.ebook_search.api.dependencies import (
AppConfig, # noqa: TC001 FastAPI resolves this annotated dependency at runtime
)
from python.ebook_search.bm25_corpus import bm25_index_exists, bm25_index_path, read_bm25_manifest
from python.ebook_search.llm_interface import check_chat_endpoint, check_embedding_endpoint
from python.fastapi_tools import AsyncDbSession # noqa: TC001 FastAPI resolves this annotated dependency at runtime
from python.fastapi_tools import DbSession # noqa: TC001 FastAPI resolves this annotated dependency at runtime
if TYPE_CHECKING:
import httpx
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import Session
from python.ebook_search.config import EbookSearchConfig
@@ -31,17 +29,17 @@ router = APIRouter()
@router.get("/health")
async def health() -> dict[str, str]:
def health() -> dict[str, str]:
"""Liveness probe that returns ok without touching dependencies."""
return {"status": "ok"}
@router.get("/ready")
async def ready(config: AppConfig, session: AsyncDbSession, client: AppHttpClient) -> JSONResponse:
def ready(config: AppConfig, session: DbSession) -> JSONResponse:
"""Readiness probe reporting database, embedding endpoint, and BM25 index status."""
database_ok = await check_database(session)
embedding_ok = await check_embedding_endpoint(client, config)
chat_status = await chat_endpoint_status(client, config)
database_ok = check_database(session)
embedding_ok = check_embedding_endpoint(config)
chat_status = chat_endpoint_status(config)
bm25_status = check_bm25_status(config)
checks = {
@@ -71,17 +69,17 @@ async def ready(config: AppConfig, session: AsyncDbSession, client: AppHttpClien
return JSONResponse(content={"status": status, "checks": checks}, status_code=status_code)
async def chat_endpoint_status(client: httpx.AsyncClient, config: EbookSearchConfig) -> str:
def chat_endpoint_status(config: EbookSearchConfig) -> str:
"""Return the answering chat endpoint status, or disabled when answers are off."""
if not config.answer_enabled:
return "disabled"
return "ok" if await check_chat_endpoint(client, config) else "fail"
return "ok" if check_chat_endpoint(config) else "fail"
async def check_database(session: AsyncSession) -> bool:
def check_database(session: Session) -> bool:
"""Return whether the database answers a trivial query."""
try:
await session.execute(select(literal(1)))
session.execute(select(literal(1)))
except SQLAlchemyError as error:
logger.warning("ebook_ready_database_unavailable error=%s", error)
return False
+14 -158
View File
@@ -3,23 +3,17 @@
from __future__ import annotations
import logging
from typing import TYPE_CHECKING
from fastapi import APIRouter, BackgroundTasks, HTTPException, Request
from fastapi.responses import HTMLResponse, RedirectResponse
from sqlalchemy import func, select
from fastapi import APIRouter, Request
from fastapi.responses import HTMLResponse
from sqlalchemy import select
from python.ebook_search.api.dependencies import (
AppConfig, # noqa: TC001 FastAPI resolves this annotated dependency at runtime
)
from python.ebook_search.api.judge_tasks import is_judging_book, pop_book_judgment_outcome, start_book_phrase_judgment
from python.ebook_search.api.web import templates
from python.ebook_search.protected_phrases.generate_ngrams import recalculate_candidate_phrases_for_book
from python.fastapi_tools import AsyncDbSession # noqa: TC001 FastAPI resolves this annotated dependency at runtime
from python.orm.richie import EbookCandidatePhrase, EbookChapter, EbookChunk, EbookProtectedPhrase, EbookSource
if TYPE_CHECKING:
from sqlalchemy.ext.asyncio import AsyncSession
from python.fastapi_tools import DbSession # noqa: TC001 FastAPI resolves this annotated dependency at runtime
from python.orm.richie import EbookSource
logger = logging.getLogger(__name__)
@@ -27,176 +21,38 @@ router = APIRouter()
@router.get("/", response_class=HTMLResponse)
async def index(request: Request, config: AppConfig) -> HTMLResponse:
def index(request: Request, config: AppConfig) -> HTMLResponse:
"""Render the search page."""
return templates.TemplateResponse(request, "search.html", {"config": config})
@router.get("/books", response_class=HTMLResponse)
async def books(request: Request, session: AsyncDbSession) -> HTMLResponse:
def books(request: Request, session: DbSession) -> HTMLResponse:
"""Render the indexed books page."""
sources = list((await session.scalars(select(EbookSource).order_by(EbookSource.title))).all())
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})
async def get_chapter_count(session: AsyncSession, book_id: int) -> int:
"""Return the number of indexed chapters for one book."""
return await session.scalar(select(func.count(EbookChapter.id)).where(EbookChapter.source_id == book_id)) or 0
async def get_chunk_count(session: AsyncSession, book_id: int) -> int:
"""Return the number of indexed chunks for one book."""
return await session.scalar(select(func.count(EbookChunk.id)).where(EbookChunk.source_id == book_id)) or 0
async def get_candidate_count(session: AsyncSession, book_id: int) -> int:
"""Return the number of indexed candidates for one book."""
return (
await session.scalar(select(func.count(EbookCandidatePhrase.id)).where(EbookCandidatePhrase.book_id == book_id))
or 0
)
async def get_judged_candidate_count(session: AsyncSession, book_id: int) -> int:
"""Return the number of judged candidates for one book."""
return (
await session.scalar(
select(func.count(EbookCandidatePhrase.id)).where(
EbookCandidatePhrase.book_id == book_id,
EbookCandidatePhrase.llm_judged.is_(True),
)
)
or 0
)
async def get_protected_count(session: AsyncSession, book_id: int) -> int:
"""Return the number of protected phrases for one book."""
return (
await session.scalar(select(func.count(EbookProtectedPhrase.id)).where(EbookProtectedPhrase.book_id == book_id))
or 0
)
async def get_candidates(session: AsyncSession, book_id: int) -> list[EbookCandidatePhrase]:
"""Return the indexed candidates for one book."""
return list(
await session.scalars(
select(EbookCandidatePhrase)
.where(EbookCandidatePhrase.book_id == book_id)
.order_by(EbookCandidatePhrase.candidate_score.desc())
.limit(100)
)
)
async def get_protected_phrases(session: AsyncSession, book_id: int) -> list[EbookProtectedPhrase]:
"""Return the protected phrases for one book."""
return list(
await session.scalars(
select(EbookProtectedPhrase)
.where(EbookProtectedPhrase.book_id == book_id)
.order_by(EbookProtectedPhrase.importance.desc())
.limit(100)
)
)
@router.get("/books/{source_id}", response_class=HTMLResponse)
async def book_detail(source_id: int, request: Request, session: AsyncDbSession) -> HTMLResponse:
def book_detail(source_id: int, request: Request, session: DbSession) -> HTMLResponse:
"""Render details for one indexed book."""
source = await session.get(EbookSource, source_id)
phrase_status_message = None
recalculated = request.query_params.get("phrases_recalculated")
if recalculated is not None:
phrase_status_message = f"Recalculated phrases; {recalculated} candidates generated"
judgment_outcome = pop_book_judgment_outcome(request.app, source_id)
if judgment_outcome is not None:
phrase_status_message = judgment_outcome
judging_in_progress = is_judging_book(request.app, source_id)
if judging_in_progress:
phrase_status_message = "Judging candidate phrases in the background; refresh to see progress"
source = session.get(EbookSource, source_id)
if source is not None:
chapter_count = await get_chapter_count(session, source.id)
chunk_count = await get_chunk_count(session, source.id)
candidate_count = await get_candidate_count(session, source.id)
judged_candidate_count = await get_judged_candidate_count(session, source.id)
protected_count = await get_protected_count(session, source.id)
candidates = await get_candidates(session, source.id)
protected_phrases = await get_protected_phrases(session, source.id)
chapter_count = len(source.chapters)
chunk_count = len(source.chunks)
else:
chapter_count = 0
chunk_count = 0
candidate_count = 0
judged_candidate_count = 0
protected_count = 0
candidates = []
protected_phrases = []
logger.info(
"ebook_book_detail_loaded source_id=%s found=%s chapters=%s chunks=%s candidates=%s judged=%s protected=%s",
"ebook_book_detail_loaded source_id=%s found=%s chapters=%s chunks=%s",
source_id,
source is not None,
chapter_count,
chunk_count,
candidate_count,
judged_candidate_count,
protected_count,
)
return templates.TemplateResponse(
request,
"book_detail.html",
{
"candidate_count": candidate_count,
"candidates": candidates,
"chapter_count": chapter_count,
"chunk_count": chunk_count,
"judged_candidate_count": judged_candidate_count,
"judging_in_progress": judging_in_progress,
"protected_count": protected_count,
"protected_phrases": protected_phrases,
"phrase_status_message": phrase_status_message,
"source": source,
},
{"chapter_count": chapter_count, "chunk_count": chunk_count, "source": source},
)
@router.post("/books/{source_id}/recalculate-phrases")
async def recalculate_book_phrases(source_id: int, config: AppConfig, session: AsyncDbSession) -> RedirectResponse:
"""Clear and regenerate candidate phrases for one indexed book."""
source = await session.get(EbookSource, source_id)
if source is None:
raise HTTPException(status_code=404, detail="Book not found")
result = await recalculate_candidate_phrases_for_book(session, source, config, use_process_pool=True)
logger.info(
"ebook_book_phrase_recalculation_complete source_id=%s candidates=%s deleted_candidates=%s "
"deleted_protected=%s deleted_aliases=%s deleted_mentions=%s",
source_id,
result.candidate_phrases,
result.deleted_candidates,
result.deleted_protected_phrases,
result.deleted_aliases,
result.deleted_mentions,
)
return RedirectResponse(
url=f"/books/{source_id}?phrases_recalculated={result.candidate_phrases}",
status_code=303,
)
@router.post("/books/{source_id}/judge-phrases")
async def judge_book_phrases(
source_id: int,
request: Request,
background_tasks: BackgroundTasks,
session: AsyncDbSession,
) -> RedirectResponse:
"""Queue background judging of one book's candidate phrases and return immediately."""
source = await session.get(EbookSource, source_id)
if source is None:
raise HTTPException(status_code=404, detail="Book not found")
started = start_book_phrase_judgment(request.app, background_tasks, source.id)
logger.info("ebook_book_phrase_judgment_requested source_id=%s started=%s", source_id, started)
return RedirectResponse(url=f"/books/{source_id}", status_code=303)
+5 -18
View File
@@ -14,7 +14,6 @@ from python.ebook_search.answer import answer_query
from python.ebook_search.api.dependencies import ( # noqa: TC001 FastAPI resolves these annotated dependencies at runtime
AppConfig,
AppEngine,
AppHttpClient,
)
from python.ebook_search.api.web import templates
from python.ebook_search.guardrails import (
@@ -27,8 +26,6 @@ from python.ebook_search.search import SearchResponse, search_ebooks
from python.ebook_search.timing import runtime_step_from_start
if TYPE_CHECKING:
import httpx
from python.ebook_search.config import EbookSearchConfig
logger = logging.getLogger(__name__)
@@ -36,8 +33,7 @@ logger = logging.getLogger(__name__)
router = APIRouter()
async def build_answer(
client: httpx.AsyncClient,
def build_answer(
query: str,
response: SearchResponse,
config: EbookSearchConfig,
@@ -60,7 +56,7 @@ async def build_answer(
return answer, True, None
try:
answer = await answer_query(client, query, response.results, config)
answer = answer_query(query, response.results, config)
except RuntimeError as error:
logger.warning("ebook_answer_request_failed_falling_back error=%s", error)
return "Answer generation failed. Source chunks are still shown below.", False, None
@@ -78,31 +74,22 @@ async def build_answer(
@router.post("/search", response_class=HTMLResponse)
async def search(
def search(
request: Request,
config: AppConfig,
engine: AppEngine,
client: AppHttpClient,
query: Annotated[str, Form()],
rerank: Annotated[str | None, Form()] = None,
phrase_matching: Annotated[str | None, Form()] = None,
) -> HTMLResponse:
"""Run a search and render HTMX results."""
try:
response = await search_ebooks(
engine,
client,
query,
config,
rerank=rerank == "true",
phrase_matching=phrase_matching == "true",
)
response = search_ebooks(engine, query, 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()
answer, low_confidence, citation_report = await build_answer(client, query, response, config)
answer, low_confidence, citation_report = build_answer(query, response, config)
answer_step_name = "Answer generation" if config.answer_enabled else "Answer skipped"
response = replace(
response,
-33
View File
@@ -181,12 +181,6 @@ textarea:focus {
flex-wrap: wrap;
}
.search-toggles {
display: flex;
flex-wrap: wrap;
gap: 14px;
}
button {
padding: 10px 20px;
font: inherit;
@@ -218,11 +212,6 @@ button:hover {
margin-bottom: 24px;
}
.actions-grid {
display: grid;
grid-template-columns: repeat(2, max-content);
}
/* Answer + results */
#results {
display: block;
@@ -314,28 +303,6 @@ button:hover {
font-variant-numeric: tabular-nums;
}
.phrase-matches {
display: flex;
flex-wrap: wrap;
gap: 8px;
align-items: baseline;
margin: 10px 0 0;
font-size: 0.78rem;
}
.phrase-matches-label {
color: var(--muted);
font-weight: 600;
}
.phrase-match {
padding: 3px 10px;
background: var(--bg);
border: 1px solid var(--border);
border-radius: 999px;
color: var(--accent);
}
/* Runtime — developer diagnostics, hidden unless dev mode is on */
.runtime {
display: none;
+44 -109
View File
@@ -1,110 +1,45 @@
{% extends "base.html" %} {% block title %}EPUB Admin{% endblock %} {% block
head %}
<script src="https://unpkg.com/htmx.org@2.0.4"></script>
{% endblock %} {% block content %}
<h1>Admin</h1>
<section id="admin-status"></section>
<section class="actions">
<form hx-post="/admin/scan" hx-target="#admin-status" hx-swap="innerHTML">
<button type="submit">Scan</button>
</form>
</section>
<section>
<h2>Embeddings</h2>
<section class="actions">
<form
hx-post="/admin/embed-missing"
hx-target="#admin-status"
hx-swap="innerHTML"
>
<button type="submit">Embed</button>
</form>
<form
hx-post="/admin/embed-all"
hx-target="#admin-status"
hx-swap="innerHTML"
>
<button type="submit">Embed all</button>
</form>
</section>
<table>
<thead>
<tr>
<th>Model</th>
<th>Dimensions</th>
<th>Embedded</th>
<th>Missing</th>
<th>Total chunks</th>
</tr>
</thead>
<tbody>
{% for item in stats %}
<tr>
<td>{{ item.model_name }}</td>
<td>{{ item.dimension }}</td>
<td>{{ item.embedded_chunks }}</td>
<td>{{ item.missing_chunks }}</td>
<td>{{ item.total_chunks }}</td>
</tr>
{% endfor %}
</tbody>
</table>
</section>
<section>
<h2>Protected phrases</h2>
<section class="actions actions-grid">
<form
hx-post="/admin/phrases/generate-all"
hx-target="#admin-status"
hx-swap="innerHTML"
>
<button type="submit">Regenerate all phrases</button>
</form>
<form
hx-post="/admin/phrases/generate-missing"
hx-target="#admin-status"
hx-swap="innerHTML"
>
<button type="submit">Add missing phrases</button>
</form>
<form
hx-post="/admin/phrases/judge-all"
hx-target="#admin-status"
hx-swap="innerHTML"
>
<button type="submit">Judge all phrases</button>
</form>
<form
hx-post="/admin/phrases/judge-missing"
hx-target="#admin-status"
hx-swap="innerHTML"
>
<button type="submit">Judge missing phrases</button>
</form>
</section>
<table>
<thead>
<tr>
<th>Candidates</th>
<th>Judged</th>
<th>Unjudged</th>
<th>Protected</th>
<th>Books indexed</th>
<th>Books generated</th>
<th>Books fully judged</th>
</tr>
</thead>
<tbody>
<tr>
<td>{{ phrase_stats.candidate_phrases }}</td>
<td>{{ phrase_stats.judged_candidates }}</td>
<td>{{ phrase_stats.unjudged_candidates }}</td>
<td>{{ phrase_stats.protected_phrases }}</td>
<td>{{ phrase_stats.total_books }}</td>
<td>{{ phrase_stats.books_with_candidates }}</td>
<td>{{ phrase_stats.books_fully_judged }}</td>
</tr>
</tbody>
</table>
</section>
{% extends "base.html" %}
{% block title %}EPUB Admin{% endblock %}
{% block head %}<script src="https://unpkg.com/htmx.org@2.0.4"></script>{% endblock %}
{% block content %}
<h1>Admin</h1>
<section id="admin-status"></section>
<section class="actions">
<form hx-post="/admin/scan" hx-target="#admin-status" hx-swap="innerHTML">
<button type="submit">Scan</button>
</form>
<form hx-post="/admin/embed-missing" hx-target="#admin-status" hx-swap="innerHTML">
<button type="submit">Embed</button>
</form>
<form hx-post="/admin/embed-all" hx-target="#admin-status" hx-swap="innerHTML">
<button type="submit">Embed all</button>
</form>
</section>
<section>
<h2>Embeddings</h2>
<table>
<thead>
<tr>
<th>Model</th>
<th>Dimensions</th>
<th>Embedded</th>
<th>Missing</th>
<th>Total chunks</th>
</tr>
</thead>
<tbody>
{% for item in stats %}
<tr>
<td>{{ item.model_name }}</td>
<td>{{ item.dimension }}</td>
<td>{{ item.embedded_chunks }}</td>
<td>{{ item.missing_chunks }}</td>
<td>{{ item.total_chunks }}</td>
</tr>
{% endfor %}
</tbody>
</table>
</section>
{% endblock %}
@@ -6,9 +6,6 @@
{% if source %}
<h1>{{ source.title }}</h1>
<p class="meta">{{ source.author or "Unknown author" }}</p>
{% if phrase_status_message %}
<p class="status">{{ phrase_status_message }}</p>
{% endif %}
<dl class="card">
<dt>File</dt>
<dd>{{ source.file_path }}</dd>
@@ -16,93 +13,7 @@
<dd>{{ chapter_count }}</dd>
<dt>Chunks</dt>
<dd>{{ chunk_count }}</dd>
<dt>Candidates</dt>
<dd>{{ candidate_count }}</dd>
<dt>Judged</dt>
<dd>{{ judged_candidate_count }}</dd>
<dt>Protected</dt>
<dd>{{ protected_count }}</dd>
</dl>
<form
method="post"
action="/books/{{ source.id }}/recalculate-phrases"
onsubmit="return confirm('Remove old phrases for this book and generate new candidates?');"
>
<button type="submit">Recalculate phrases</button>
</form>
<form
method="post"
action="/books/{{ source.id }}/judge-phrases"
onsubmit="return confirm('Judge candidate phrases for this book with the LLM?');"
>
<button type="submit"{% if judging_in_progress %} disabled{% endif %}>
{% if judging_in_progress %}Judging&hellip;{% else %}Judge phrases{% endif %}
</button>
</form>
<section>
<h2>Candidate n-grams</h2>
{% if candidates %}
<table>
<thead>
<tr>
<th>Phrase</th>
<th>Status</th>
<th>Score</th>
<th>Count</th>
<th>Chapters</th>
</tr>
</thead>
<tbody>
{% for candidate in candidates %}
<tr>
<td>{{ candidate.phrase_text }}</td>
<td>
{% if candidate.llm_judged %}
{% if candidate.llm_keep %}Kept{% else %}Rejected{% endif %}
{% else %}
Candidate
{% endif %}
</td>
<td>{{ "%.2f"|format(candidate.candidate_score) }}</td>
<td>{{ candidate.raw_count }}</td>
<td>{{ candidate.chapter_count }}</td>
</tr>
{% endfor %}
</tbody>
</table>
{% else %}
<p>No candidate n-grams.</p>
{% endif %}
</section>
<section>
<h2>Protected phrases</h2>
{% if protected_phrases %}
<table>
<thead>
<tr>
<th>Phrase</th>
<th>Type</th>
<th>Confidence</th>
<th>Importance</th>
</tr>
</thead>
<tbody>
{% for phrase in protected_phrases %}
<tr>
<td>{{ phrase.phrase_text }}</td>
<td>{{ phrase.phrase_type or "phrase" }}</td>
<td>{{ "%.2f"|format(phrase.confidence) }}</td>
<td>{{ "%.2f"|format(phrase.importance) }}</td>
</tr>
{% endfor %}
</tbody>
</table>
{% else %}
<p>No protected phrases.</p>
{% endif %}
</section>
{% else %}
<h1>Book not found</h1>
{% endif %}
@@ -82,14 +82,6 @@
</div>
{% endif %}
</dl>
{% if result.matched_phrases %}
<p class="phrase-matches">
<span class="phrase-matches-label">boosted by</span>
{% for phrase in result.matched_phrases %}
<span class="phrase-match">{{ phrase }}</span>
{% endfor %}
</p>
{% endif %}
</li>
{% endfor %}
</ol>
+5 -17
View File
@@ -7,24 +7,12 @@
<h1>Search</h1>
<form class="card" hx-post="/search" hx-target="#results" hx-swap="innerHTML">
<label for="query">What are you looking for?</label>
<textarea id="query" name="query" rows="4" placeholder="Ask a question or paste a passage…" required
onkeydown="if (event.key === 'Enter' && !event.shiftKey) { event.preventDefault(); this.form.requestSubmit(); }"></textarea>
<textarea id="query" name="query" rows="4" placeholder="Ask a question or paste a passage…" required></textarea>
<div class="form-row">
<div class="search-toggles">
<label class="check">
<input type="checkbox" name="rerank" value="true" {% if config.rerank.enabled %}checked{% endif %}>
Rerank
</label>
<label class="check">
<input
type="checkbox"
name="phrase_matching"
value="true"
{% if config.phrase_matching_enabled %}checked{% endif %}
>
Phrase matching
</label>
</div>
<label class="check">
<input type="checkbox" name="rerank" value="true" {% if config.rerank.enabled %}checked{% endif %}>
Rerank
</label>
<button type="submit">Search</button>
</div>
</form>
+15 -20
View File
@@ -2,7 +2,6 @@
from __future__ import annotations
import asyncio
import json
import logging
import shutil
@@ -18,7 +17,7 @@ from sqlalchemy import func, select, union_all
from python.orm.richie import EbookChapter, EbookChunk, EbookSource
if TYPE_CHECKING:
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import Session
from python.ebook_search.config import EbookSearchConfig
@@ -74,14 +73,14 @@ def get_current_bm25_index(index_path: Path) -> Path:
return index_path
async def ensure_bm25_corpus(session: AsyncSession, config: EbookSearchConfig) -> None:
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 = await corpus_last_updated_at(session)
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)
await refresh_bm25_corpus(session, config, db_updated_at=db_updated_at)
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(
@@ -90,7 +89,7 @@ async def ensure_bm25_corpus(session: AsyncSession, config: EbookSearchConfig) -
manifest.created_at.isoformat(),
db_updated_at.isoformat(),
)
await refresh_bm25_corpus(session, config, db_updated_at=db_updated_at)
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",
@@ -100,24 +99,21 @@ async def ensure_bm25_corpus(session: AsyncSession, config: EbookSearchConfig) -
)
async def refresh_bm25_corpus(
session: AsyncSession,
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.
The index build is CPU and disk work, so it runs in a worker thread.
"""
"""Rebuild and persist the BM25 corpus from the current database chunks."""
index_path = bm25_index_path(config)
records, texts = await fetch_bm25_corpus_records(session)
records, texts = fetch_bm25_corpus_records(session)
manifest = BM25Manifest(
created_at=datetime.now(tz=UTC),
db_updated_at=db_updated_at if db_updated_at is not None else await corpus_last_updated_at(session),
db_updated_at=db_updated_at if db_updated_at is not None else corpus_last_updated_at(session),
chunk_count=len(records),
)
await asyncio.to_thread(write_bm25_corpus, index_path, records, texts, manifest)
write_bm25_corpus(index_path, records, texts, manifest)
logger.info(
"ebook_bm25_index_refreshed path=%s chunks=%s created_at=%s",
index_path,
@@ -168,7 +164,7 @@ def score_bm25_corpus(query: str, corpus: BM25Corpus, *, limit: int) -> list[tup
return results
async def fetch_bm25_corpus_records(session: AsyncSession) -> tuple[list[dict[str, object]], list[str]]:
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
@@ -178,7 +174,6 @@ async def fetch_bm25_corpus_records(session: AsyncSession) -> tuple[list[dict[st
select(
EbookChunk.id.label("chunk_id"),
EbookChunk.text.label("text"),
EbookSource.id.label("source_id"),
EbookSource.title.label("source_title"),
EbookSource.author.label("source_author"),
EbookChapter.title.label("chapter_title"),
@@ -192,21 +187,21 @@ async def fetch_bm25_corpus_records(session: AsyncSession) -> tuple[list[dict[st
)
records: list[dict[str, object]] = []
texts: list[str] = []
for row in (await session.execute(statement)).mappings():
for row in session.execute(statement).mappings():
record = dict(row)
texts.append(str(record.pop("bm25_text")))
records.append(record)
return records, texts
async def corpus_last_updated_at(session: AsyncSession) -> datetime | None:
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 await session.scalar(select(func.max(update_times.c.updated)))
return session.scalar(select(func.max(update_times.c.updated)))
def write_bm25_corpus(
-19
View File
@@ -86,25 +86,6 @@ class EbookSearchConfig(BaseSettings):
validate_citations_enabled: bool = True
bm25_index_dir: str = ".ebook_search_bm25"
bm25_refresh_delay_seconds: int = 60
protected_phrase_max_candidates_per_book: int = 5000
protected_phrase_llm_candidates_per_book: int = 500
protected_phrase_extraction_workers: int = 16
phrase_judge_book_workers: int = 20
phrase_judge_phrase_workers: int = 100
protected_phrase_confidence_threshold: float = 0.80
phrase_matching_enabled: bool = True
phrase_hit_boost: float = 0.25
phrase_min_tokens: int = 2
phrase_max_tokens: int = 5
phrase_max_entity_tokens: int = 8
phrase_raw_ngram_min_count: int = 2
phrase_raw_count_score_threshold: int = 3
phrase_raw_count_high_score_threshold: int = 10
phrase_chapter_count_score_threshold: int = 2
phrase_chapter_count_high_score_threshold: int = 5
phrase_target_protected_per_book: int = 100
phrase_default_allow_nested: bool = False
phrase_default_suppress_children: bool = True
@field_validator("library_paths", mode="before")
@classmethod
+32 -27
View File
@@ -1,12 +1,17 @@
FROM python:3.14-slim
FROM python:3.14-slim AS base
COPY --from=ghcr.io/astral-sh/uv:0.11.26 /uv /uvx /bin/
ENV PYTHONDONTWRITEBYTECODE=1 \
PYTHONUNBUFFERED=1 \
PIP_NO_CACHE_DIR=1 \
APP_DIR=/home/richie/dotfiles \
EBOOK_SEARCH_HOST=0.0.0.0 \
EBOOK_SEARCH_PORT=8070 \
EBOOK_SEARCH_BM25_INDEX_DIR=/data/bm25
UV_PROJECT_ENVIRONMENT=/opt/venv \
UV_PYTHON_DOWNLOADS=never \
UV_NO_CACHE=1
# Separate ENV instruction so ${APP_DIR} and ${PATH} from above resolve.
ENV PYTHONPATH=${APP_DIR} \
PATH=/opt/venv/bin:${PATH}
WORKDIR ${APP_DIR}
@@ -14,29 +19,29 @@ RUN apt-get update \
&& apt-get install -y --no-install-recommends build-essential curl \
&& rm -rf /var/lib/apt/lists/*
COPY pyproject.toml README.md LICENSE ./
COPY python ./python
COPY python/ebook_search/docker/pyproject.toml python/ebook_search/docker/uv.lock ./
RUN python -m pip install --upgrade pip \
&& python -m pip install \
"alembic" \
"beautifulsoup4" \
"bm25s" \
"ebooklib" \
"fastapi" \
"httpx" \
"jinja2" \
"pgvector" \
"psycopg[binary]" \
"pydantic" \
"pydantic-settings" \
"python-multipart" \
"sqlalchemy[asyncio]" \
"tiktoken" \
"typer" \
"uvicorn[standard]" \
"yake" \
&& python -m pip install --no-deps --editable "${APP_DIR}"
RUN uv sync --locked --no-dev
FROM base AS test
RUN uv sync --locked
COPY python ./python
COPY tests/__init__.py ./tests/__init__.py
COPY tests/ebook_search ./tests/ebook_search
CMD ["pytest"]
FROM base AS runtime
ENV EBOOK_SEARCH_HOST=0.0.0.0 \
EBOOK_SEARCH_PORT=8070 \
EBOOK_SEARCH_BM25_INDEX_DIR=/data/bm25
COPY python ./python
RUN useradd --create-home --uid 10001 app \
&& mkdir -p /data \
+49 -12
View File
@@ -3,26 +3,27 @@
Run the EPUB search app against the existing Postgres database on `jeeves`:
```sh
ebook-search-containers start --library-path /path/to/epubs --build
python -m python.ebook_search.docker.containers start --library-path /path/to/epubs --build
```
All ebook-search Docker files live in this directory:
- `Dockerfile`
- `Dockerfile` — multi-stage: `test` (runs pytest) and `runtime` (default target, the app image)
- `docker-compose.yml`
- `containers.py`
- `container.py`
- `containers.py` — Typer lifecycle CLI
- `pyproject.toml` / `uv.lock` — the container's uv-locked dependencies
The app listens on `http://localhost:8070`.
Useful lifecycle commands:
```sh
ebook-search-containers build
ebook-search-containers start --library-path /path/to/epubs
ebook-search-containers logs
ebook-search-containers ps
ebook-search-containers stop
python -m python.ebook_search.docker.containers build
python -m python.ebook_search.docker.containers start --library-path /path/to/epubs
python -m python.ebook_search.docker.containers test
python -m python.ebook_search.docker.containers logs
python -m python.ebook_search.docker.containers ps
python -m python.ebook_search.docker.containers stop
```
Direct compose usage from the repo root:
@@ -31,10 +32,46 @@ Direct compose usage from the repo root:
docker compose -f python/ebook_search/docker/docker-compose.yml ps
```
The compose service also loads the repo root `.env` into the container via `env_file`.
## Dependencies
The image builds its environment with uv from `pyproject.toml` + `uv.lock` in this
directory — this is the source of truth for the container's dependencies. To add or
update a dependency, edit `pyproject.toml` here and regenerate the lock (uv is
available in the `ebook-search` dev shell):
```sh
nix develop .#ebook-search -c uv lock --project python/ebook_search/docker
```
## Tests
The main pytest suite excludes `tests/ebook_search` (its dependencies are no longer
in the nix dev shell). The `test ebook search` CI workflow runs them in a uv env
built from the lockfile in this directory — same commands work locally from the
repo root (the `--override-ini` drops the main suite's ignore):
```sh
uv sync --locked --project python/ebook_search/docker
uv run --project python/ebook_search/docker --no-sync pytest tests/ebook_search --override-ini addopts="-n auto -ra"
```
They can also run inside the Docker `test` image, which validates the image itself:
```sh
python -m python.ebook_search.docker.containers test
```
or the raw docker equivalent:
```sh
docker build --file python/ebook_search/docker/Dockerfile --target test --tag ebook-search:test .
docker run --rm ebook-search:test
```
## Configuration
The compose service loads the repo root `.env` into the container via `env_file`.
Mount your EPUB directory by setting `EBOOK_LIBRARY_HOST_PATH` in an env file or on the command line. The container sees it as `/library`, and `EBOOK_SEARCH_LIBRARY_PATHS` is set to `/library` inside the container.
Database connection settings are controlled by `RICHIE_DB`, `RICHIE_HOST`, `RICHIE_PORT`, `RICHIE_USER`, and `RICHIE_PASSWORD`. The default host is `jeeves`.
Startup runs the Richie Alembic migrations automatically after creating the `main` schema and `vector` extension.
+30
View File
@@ -73,6 +73,23 @@ def build_image() -> None:
raise RuntimeError(msg)
def build_test_image() -> None:
"""Build the ebook search test Docker image."""
dockerfile = Path(__file__).resolve().with_name("Dockerfile")
result = docker_run(["build", "--file", str(dockerfile), "--target", "test", "--tag", "ebook-search:test", "."])
if result.returncode != 0:
msg = "Failed to build ebook search test image"
raise RuntimeError(msg)
def run_test_image() -> None:
"""Run the ebook search test suite inside Docker."""
result = docker_run(["run", "--rm", "ebook-search:test"])
if result.returncode != 0:
msg = f"Ebook search tests failed with code {result.returncode}"
raise RuntimeError(msg)
def start_stack(
*,
library_path: Path | None = None,
@@ -210,6 +227,19 @@ def logs(
typer.echo(output)
@app.command("test")
def run_tests(
*,
build: Annotated[bool, typer.Option("--build/--no-build", help="Build the test image before running.")] = True,
log_level: Annotated[str, typer.Option(help="Log level.")] = "INFO",
) -> None:
"""Run ebook search tests inside the Docker test image."""
configure_logger(log_level)
if build:
build_test_image()
run_test_image()
@app.command("ps")
def ps() -> None:
"""Show ebook search container status."""
+38
View File
@@ -0,0 +1,38 @@
[project]
name = "ebook-search"
version = "0.1.0"
description = "Locked runtime environment for the ebook search container."
requires-python = "~=3.14.0"
dependencies = [
"alembic",
"beautifulsoup4",
"bm25s",
"ebooklib",
"fastapi",
"httpx",
"jinja2",
"pgvector",
"psycopg[binary]",
"pydantic",
"pydantic-settings",
"python-multipart",
"sqlalchemy",
"tiktoken",
"typer",
"uvicorn[standard]",
"yake",
]
[dependency-groups]
dev = [
"pytest",
"pytest-mock",
"pytest-xdist",
]
[tool.uv]
package = false
[tool.pytest.ini_options]
addopts = "-n auto -ra"
testpaths = ["tests/ebook_search"]
+1109
View File
File diff suppressed because it is too large Load Diff
+19 -24
View File
@@ -23,8 +23,7 @@ logger = logging.getLogger(__name__)
if TYPE_CHECKING:
from collections.abc import Sequence
import httpx
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import Session
from python.ebook_search.config import EbookSearchConfig
@@ -66,11 +65,7 @@ class EmbeddingModelStats:
return max(self.total_chunks - self.embedded_chunks, 0)
async def embed_texts(
client: httpx.AsyncClient,
texts: Sequence[str],
config: EbookSearchConfig,
) -> list[list[float]]:
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",
@@ -78,7 +73,7 @@ async def embed_texts(
config.embedding_model,
len(texts),
)
vectors = await request_embeddings(client, texts, config)
vectors = request_embeddings(texts, config)
expected_dimension = MODEL_DIMENSIONS[config.embedding_model]
for vector in vectors:
if len(vector) != expected_dimension:
@@ -93,28 +88,28 @@ async def embed_texts(
return vectors
async def embed_query(client: httpx.AsyncClient, query: str, config: EbookSearchConfig) -> list[float]:
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 (await embed_texts(client, [instructed_query], config))[0]
return embed_texts([instructed_query], config)[0]
async def ensure_embedding_models(session: AsyncSession) -> None:
def ensure_embedding_models(session: Session) -> None:
"""Ensure supported embedding model rows exist."""
for name, dimension in MODEL_DIMENSIONS.items():
existing = await session.scalar(select(EbookEmbeddingModel).where(EbookEmbeddingModel.name == name))
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)
await session.flush()
session.flush()
async def embedding_model_stats(session: AsyncSession) -> list[EmbeddingModelStats]:
def embedding_model_stats(session: Session) -> list[EmbeddingModelStats]:
"""Return embedding coverage counts for every supported model."""
total_chunks = await session.scalar(select(func.count(EbookChunk.id))) or 0
total_chunks = session.scalar(select(func.count(EbookChunk.id))) or 0
models = {
model.name: model
for model in await session.scalars(
for model in session.scalars(
select(EbookEmbeddingModel)
.where(EbookEmbeddingModel.name.in_(MODEL_DIMENSIONS))
.order_by(EbookEmbeddingModel.name)
@@ -127,7 +122,7 @@ async def embedding_model_stats(session: AsyncSession) -> list[EmbeddingModelSta
embedded_chunks = 0
if model is not None:
table = get_embedding_table(dimension)
embedded_chunks = await session.scalar(select(func.count(table.id)).where(table.model_id == model.id)) or 0
embedded_chunks = session.scalar(select(func.count(table.id)).where(table.model_id == model.id)) or 0
stats.append(
EmbeddingModelStats(
model_name=model_name,
@@ -139,10 +134,10 @@ async def embedding_model_stats(session: AsyncSession) -> list[EmbeddingModelSta
return stats
async def embed_missing_chunks(session: AsyncSession, client: httpx.AsyncClient, config: EbookSearchConfig) -> int:
def embed_missing_chunks(session: Session, config: EbookSearchConfig) -> int:
"""Embed chunks missing embeddings for the configured model."""
await ensure_embedding_models(session)
model = await session.scalar(select(EbookEmbeddingModel).where(EbookEmbeddingModel.name == config.embedding_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}"
@@ -150,7 +145,7 @@ async def embed_missing_chunks(session: AsyncSession, client: httpx.AsyncClient,
table = get_embedding_table(model.dimension)
chunks = list(
await session.scalars(
session.scalars(
select(EbookChunk)
.outerjoin(table, (table.chunk_id == EbookChunk.id) & (table.model_id == model.id))
.where(table.id.is_(None))
@@ -163,13 +158,13 @@ async def embed_missing_chunks(session: AsyncSession, client: httpx.AsyncClient,
return 0
logger.info("ebook_embed_missing_batch_start model=%s count=%s", config.embedding_model, len(chunks))
vectors = await embed_texts(client, [chunk.text for chunk in chunks], config)
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"])
await session.execute(statement)
await session.flush()
session.execute(statement)
session.flush()
logger.info("ebook_embed_missing_batch_complete model=%s count=%s", config.embedding_model, len(rows))
return len(rows)
+65 -92
View File
@@ -2,7 +2,6 @@
from __future__ import annotations
import asyncio
import hashlib
import logging
from dataclasses import dataclass
@@ -14,7 +13,6 @@ import tiktoken
from sqlalchemy import or_, select
from python.ebook_search.epub_parse import parse_epub
from python.ebook_search.protected_phrases.matching import index_chunk_phrase_mentions_for_book
from python.orm.richie import EbookChapter, EbookChunk, EbookSource
logger = logging.getLogger(__name__)
@@ -22,7 +20,7 @@ DEFAULT_CHUNK_TOKENS = 700
DEFAULT_CHUNK_OVERLAP = 100
if TYPE_CHECKING:
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import Session
from python.ebook_search.config import EbookSearchConfig
from python.ebook_search.epub_parse import ParsedChapter
@@ -74,115 +72,90 @@ def chunk_text(
return [chunk for chunk in chunks if chunk.text]
def find_library_epubs(library_path: str) -> tuple[Path, list[Path] | None]:
"""Resolve one configured library path and collect its EPUB files (blocking filesystem walk).
Returns:
tuple[Path, list[Path] | None]: The expanded path and its EPUB files, or ``None`` when
the path is neither an EPUB file nor a directory.
"""
path = Path(library_path).expanduser()
if path.is_file() and path.suffix.lower() == ".epub":
return path, [path]
if path.is_dir():
return path, sorted(path.rglob("*.epub"))
return path, None
async def ingest_configured_paths(session: AsyncSession, config: EbookSearchConfig) -> int:
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, epub_paths = await asyncio.to_thread(find_library_epubs, library_path)
path = Path(library_path).expanduser()
logger.info("ebook_ingest_path_start path=%s", path)
if epub_paths is None:
if path.is_file() and path.suffix.lower() == ".epub":
count += int(ingest_file(session, path, config))
elif path.is_dir():
for epub_path in sorted(path.rglob("*.epub")):
count += int(ingest_file(session, epub_path, config))
else:
logger.warning("ebook_ingest_path_missing path=%s", path)
continue
for epub_path in epub_paths:
count += int(await ingest_file(session, epub_path, config))
logger.info("ebook_ingest_paths_complete changed_files=%s configured_paths=%s", count, len(config.library_paths))
return count
def resolve_ingest_path(path: Path) -> Path:
"""Expand and resolve an ingest path (blocking filesystem call)."""
return path.expanduser().resolve()
async def ingest_file(session: AsyncSession, path: Path, config: EbookSearchConfig) -> bool:
def ingest_file(session: Session, path: Path, config: EbookSearchConfig) -> bool:
"""Ingest one EPUB file. Return True when the database changed."""
try:
resolved_path = await asyncio.to_thread(resolve_ingest_path, path)
logger.info("ebook_ingest_file_start path=%s", resolved_path)
file_hash = await asyncio.to_thread(sha256_file, resolved_path)
existing = await 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
await 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)
await session.delete(existing)
await session.flush()
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()
parsed = await asyncio.to_thread(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)
await 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)
await session.flush()
chunk_index = add_chapter_chunks(session, source, chapter, parsed_chapter, chunk_index, config)
await session.commit()
mention_count = await index_chunk_phrase_mentions_for_book(session, source.id, config)
logger.info(
"ebook_ingest_file_complete source_id=%s path=%s chapters=%s chunks=%s phrase_mentions=%s",
source.id,
resolved_path,
len(parsed.chapters),
chunk_index,
mention_count,
)
except Exception:
logger.exception(f"ebook_ingest_file_error path={path}")
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
else:
return True
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, config)
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
async def find_existing_source(session: AsyncSession, path: Path, file_hash: str) -> EbookSource | None:
def find_existing_source(session: Session, path: Path, file_hash: str) -> EbookSource | None:
"""Find an existing source by canonical path or file hash."""
return await session.scalar(
return session.scalar(
select(EbookSource).where(or_(EbookSource.file_path == str(path), EbookSource.file_sha256 == file_hash))
)
def add_chapter_chunks(
session: AsyncSession,
session: Session,
source: EbookSource,
chapter: EbookChapter,
parsed_chapter: ParsedChapter,
+13 -63
View File
@@ -22,26 +22,10 @@ def auth_headers(api_key: str) -> dict[str, str]:
return {"Authorization": f"Bearer {api_key}"}
async def request_embeddings(
client: httpx.AsyncClient,
texts: Sequence[str],
config: EbookSearchConfig,
) -> list[list[float]]:
"""Request embeddings from the configured OpenAI-compatible endpoint.
Args:
client (httpx.AsyncClient): Shared async client for LLM calls.
texts (Sequence[str]): Texts to embed.
config (EbookSearchConfig): Runtime settings supplying the endpoint, model, and auth.
Returns:
list[list[float]]: One embedding vector per input text.
Raises:
RuntimeError: If the request fails or the response cannot be parsed.
"""
def request_embeddings(texts: Sequence[str], config: EbookSearchConfig) -> list[list[float]]:
"""Request embeddings from the configured OpenAI-compatible endpoint."""
try:
response = await client.post(
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)},
@@ -60,15 +44,10 @@ async def request_embeddings(
raise RuntimeError(msg) from error
async def check_embedding_endpoint(
client: httpx.AsyncClient,
config: EbookSearchConfig,
*,
timeout_seconds: float = 5.0,
) -> bool:
def check_embedding_endpoint(config: EbookSearchConfig, *, timeout_seconds: float = 5.0) -> bool:
"""Return whether the configured embedding endpoint answers a model listing."""
try:
response = await client.get(
response = httpx.get(
f"{config.embedding_base_url.rstrip('/')}/models",
headers=auth_headers(config.embedding_api_key),
timeout=timeout_seconds,
@@ -80,15 +59,10 @@ async def check_embedding_endpoint(
return True
async def check_chat_endpoint(
client: httpx.AsyncClient,
config: EbookSearchConfig,
*,
timeout_seconds: float = 5.0,
) -> bool:
def check_chat_endpoint(config: EbookSearchConfig, *, timeout_seconds: float = 5.0) -> bool:
"""Return whether the configured chat (answering) endpoint answers a model listing."""
try:
response = await client.get(
response = httpx.get(
f"{config.vllm_base_url.rstrip('/')}/models",
headers=auth_headers(config.vllm_api_key),
timeout=timeout_seconds,
@@ -124,29 +98,18 @@ def embedding_vectors_from_response(body: object) -> list[list[float]]:
return vectors
async def request_rerank(
client: httpx.AsyncClient,
def request_rerank(
query: str,
documents: Sequence[str],
config: RerankConfig,
) -> object | None:
"""Request rerank scores from the configured vLLM endpoint.
Args:
client (httpx.AsyncClient): Shared async client for LLM calls.
query (str): Query the documents are scored against.
documents (Sequence[str]): Candidate documents to score.
config (RerankConfig): Rerank endpoint settings.
Returns:
object | None: The decoded response body, or ``None`` when it is not valid JSON.
"""
"""Request rerank scores from the configured vLLM endpoint."""
payload = {
"model": config.model,
"query": query,
"documents": list(documents),
}
response = await client.post(
response = httpx.post(
f"{config.base_url.rstrip('/')}/rerank",
json=payload,
timeout=config.timeout_seconds,
@@ -159,26 +122,13 @@ async def request_rerank(
return None
async def request_chat_completion(
client: httpx.AsyncClient,
def request_chat_completion(
config: EbookSearchConfig,
messages: Sequence[dict[str, str]],
) -> str:
"""Request a chat completion over a shared async client.
Args:
client (httpx.AsyncClient): Shared async client whose connection pool bounds concurrency.
config (EbookSearchConfig): Runtime settings supplying the endpoint, model, and auth.
messages (Sequence[dict[str, str]]): OpenAI-style chat messages.
Returns:
str: The assistant message text.
Raises:
RuntimeError: If the request fails or the response cannot be parsed.
"""
"""Request a chat completion from the configured OpenAI-compatible endpoint."""
try:
response = await client.post(
response = httpx.post(
f"{config.vllm_base_url.rstrip('/')}/chat/completions",
headers=auth_headers(config.vllm_api_key),
json={
@@ -1 +0,0 @@
"""Init."""
@@ -1,17 +0,0 @@
"""Protected phrase extraction, storage, and runtime matching."""
from python.ebook_search.protected_phrases.config.lib import (
get_bad_ends,
get_bad_starts,
get_ignored_phrases,
get_junk_tokens,
get_most_common_words,
)
__all__ = [
"get_bad_ends",
"get_bad_starts",
"get_ignored_phrases",
"get_junk_tokens",
"get_most_common_words",
]
@@ -1,31 +0,0 @@
tokens = [
"a",
"an",
"and",
"any",
"as",
"at",
"be",
"because",
"but",
"by",
"can",
"could",
"do",
"for",
"from",
"have",
"if",
"of",
"or",
"some",
"than",
"the",
"these",
"this",
"to",
"will",
"with",
"would",
"did",
]
@@ -1,27 +0,0 @@
tokens = [
"a",
"an",
"did",
"didn't",
"he",
"here",
"how",
"i",
"it",
"she",
"that",
"the",
"there",
"they",
"this",
"we",
"what",
"when",
"where",
"which",
"who",
"whom",
"whose",
"why",
"you",
]
@@ -1,212 +0,0 @@
phrases = [
"a little",
"across the",
"and she",
"anyone in",
"are you",
"around him",
"around the",
"as much",
"as soon",
"at all",
"at least",
"before the",
"behind him",
"between the",
"but she",
"could not",
"did he",
"did i",
"did it",
"did not believe",
"did not care",
"did not even",
"did not know what",
"did not know",
"did not like",
"did not look",
"did not mean",
"did not move",
"did not need",
"did not see",
"did not seem",
"did not think",
"did not understand",
"did not want",
"did not",
"did she",
"did so",
"did that",
"did the",
"did they",
"did what",
"did you",
"didn't answer",
"didn't care",
"didn't even",
"didn't expect",
"didn't feel",
"didn't get",
"didn't i",
"didn't know",
"didn't like",
"didn't look",
"didn't make",
"didn't mean",
"didn't need",
"didn't really",
"didn't say",
"didn't see",
"didn't seem",
"didn't think",
"didn't want",
"didn't you",
"end up",
"ended up",
"had a",
"had been",
"have been",
"he asked",
"he concluded",
"he continued",
"he couldn't",
"he did",
"he didn't",
"he felt",
"he had",
"he hadn't",
"he knew",
"he noted",
"he pointed",
"he realized",
"he replied",
"he said",
"he saw",
"he tapped",
"he told",
"he was",
"he wasn't",
"his body",
"his chair",
"his feet",
"his hands",
"his head",
"his office",
"his own",
"his pc",
"his power",
"his shield",
"his sight",
"his voice",
"his wrist",
"how many",
"i am",
"i don't",
"i said",
"i was",
"i wouldn't",
"i'm not",
"if he",
"if they",
"is in",
"is not",
"is that",
"is the",
"it had",
"it had",
"it is",
"it was",
"it wasn't",
"it wasn't",
"no one",
"of course",
"of force",
"of it",
"of magic",
"of marines",
"of power",
"of those",
"old man",
"older man",
"one of",
"out of",
"set up",
"she admitted",
"she asked",
"she had",
"she replied",
"she said",
"she snapped",
"she told",
"she was",
"she'd been",
"shook his",
"sure he",
"tell you",
"that had",
"that is",
"that she",
"that was",
"the dark",
"the door",
"the first",
"the last",
"the man",
"the one",
"the only",
"the other",
"the rest",
"the room",
"the same",
"the two",
"the way",
"the world",
"there are",
"there was",
"there were",
"they are",
"they had",
"they were",
"they weren't",
"this is",
"this place",
"though he",
"through his",
"through the",
"to find",
"to get",
"to keep",
"to stay",
"to stop",
"to tell",
"to try",
"told her",
"told him",
"under his",
"was a",
"was enough",
"was going",
"was in",
"was no",
"was not",
"was now",
"was on",
"was one",
"was only",
"was still",
"was that",
"was the",
"was there",
"were in",
"what had",
"what happened",
"what was",
"where the",
"while i",
"you are",
"you can't",
"you don't",
"you know",
"you need",
"you were",
]
@@ -1,71 +0,0 @@
tokens = [
"said",
"asked",
"replied",
"answered",
"looked",
"nodded",
"turned",
"shook",
"smiled",
"shrugged",
"pointed",
"continued",
"repeated",
"stared",
"agreed",
"glanced",
"walked",
"told",
"thought",
"knew",
"wanted",
"muttered",
"whispered",
"laughed",
"sighed",
"paused",
"gestured",
"waved",
"frowned",
"grinned",
"admitted",
"found",
"noted",
"murmured",
"ordered",
"i'm",
"i've",
"i'd",
"i'll",
"it's",
"that's",
"don't",
"didn't",
"doesn't",
"can't",
"won't",
"wouldn't",
"couldn't",
"shouldn't",
"isn't",
"wasn't",
"aren't",
"weren't",
"you're",
"you've",
"you'll",
"we're",
"we've",
"we'll",
"they're",
"they've",
"he's",
"she's",
"there's",
"what's",
"let's",
"who's",
"he'd",
"she'd",
]
@@ -1,60 +0,0 @@
"""Protected phrase extraction, storage, and runtime matching."""
from __future__ import annotations
import logging
import tomllib
from functools import cache
from pathlib import Path
from python.ebook_search.protected_phrases.text_normalization import normalize_text
logger = logging.getLogger(__name__)
def _load_toml_string_set(path: Path, key: str) -> frozenset[str]:
"""Load and validate a TOML string list as a normalized immutable set."""
with path.open("rb") as file:
body = tomllib.load(file)
values = body.get(key)
if not isinstance(values, list) or not all(isinstance(item, str) for item in values):
msg = f"{path} must contain a {key!r} string list"
raise ValueError(msg)
return frozenset(normalize_text(value) for value in values if normalize_text(value))
@cache
def _get_phrase_config_dir() -> Path:
"""Return the directory containing phrase configuration files."""
return Path(__file__).resolve().parent
@cache
def get_ignored_phrases() -> frozenset[str]:
"""Return ignored phrase strings loaded from TOML."""
return _load_toml_string_set(_get_phrase_config_dir() / "ignored_phrases.toml", "phrases")
@cache
def get_bad_ends() -> frozenset[str]:
"""Return bad phrase-ending tokens loaded from TOML."""
return _load_toml_string_set(_get_phrase_config_dir() / "bad_ends.toml", "tokens")
@cache
def get_bad_starts() -> frozenset[str]:
"""Return bad phrase-starting tokens loaded from TOML."""
return _load_toml_string_set(_get_phrase_config_dir() / "bad_starts.toml", "tokens")
@cache
def get_most_common_words() -> frozenset[str]:
"""Return the most common English words loaded from TOML."""
return _load_toml_string_set(_get_phrase_config_dir() / "most_common_words.toml", "words")
@cache
def get_junk_tokens() -> frozenset[str]:
"""Return junk tokens (dialogue verbs and pronoun contractions) loaded from TOML."""
return _load_toml_string_set(_get_phrase_config_dir() / "junk_tokens.toml", "tokens")
@@ -1,102 +0,0 @@
words = [
"the",
"be",
"to",
"of",
"and",
"a",
"in",
"that",
"have",
"I",
"it",
"for",
"not",
"on",
"with",
"he",
"as",
"you",
"do",
"at",
"this",
"but",
"his",
"by",
"from",
"they",
"we",
"say",
"her",
"she",
"or",
"an",
"will",
"my",
"one",
"all",
"would",
"there",
"their",
"what",
"so",
"up",
"out",
"if",
"about",
"who",
"get",
"which",
"go",
"me",
"when",
"make",
"can",
"like",
"time",
"no",
"just",
"him",
"know",
"take",
"people",
"into",
"year",
"your",
"good",
"some",
"could",
"them",
"see",
"other",
"than",
"then",
"now",
"look",
"only",
"come",
"its",
"over",
"think",
"also",
"back",
"after",
"use",
"two",
"how",
"our",
"work",
"first",
"well",
"way",
"even",
"new",
"want",
"because",
"any",
"these",
"give",
"day",
"most",
"us",
]
@@ -1,853 +0,0 @@
"""Candidate phrase extraction and scoring for protected phrases."""
from __future__ import annotations
import logging
import re
from collections import Counter, defaultdict
from functools import lru_cache
from time import perf_counter
from typing import TYPE_CHECKING, Protocol
from yake import KeywordExtractor
from python.ebook_search.protected_phrases.config import (
get_bad_ends,
get_bad_starts,
get_ignored_phrases,
get_junk_tokens,
get_most_common_words,
)
from python.ebook_search.protected_phrases.models import PhraseCandidate
from python.ebook_search.protected_phrases.text_normalization import tokenize, tokenize_with_offsets
if TYPE_CHECKING:
from collections.abc import Iterable, Mapping, Sequence
from python.ebook_search.config import EbookSearchConfig
logger = logging.getLogger(__name__)
BAD_START_SCORE_PENALTY = 10.0
BAD_END_SCORE_PENALTY = 10.0
MULTI_SOURCE_SCORE_BONUS = 2.0
MULTI_SOURCE_MIN_SOURCES = 2
CAPITALIZED_PHRASE_RE = re.compile(r"\b(?:[A-Z][a-zA-Z']+)(?:\s+(?:of|the|and|in|on|for|[A-Z][a-zA-Z']+)){0,6}")
class SpacySpan(Protocol):
"""Small protocol for the spaCy span attributes used by this module."""
text: str
class SpacyEntity(SpacySpan, Protocol):
"""Small protocol for the spaCy entity attributes used by this module."""
label_: str
class SpacyDoc(Protocol):
"""Small protocol for the spaCy doc attributes used by this module."""
ents: Iterable[SpacyEntity]
noun_chunks: Iterable[SpacySpan]
class SpacyLanguage(Protocol):
"""Small protocol for a callable spaCy language pipeline."""
def __call__(self, text: str) -> SpacyDoc:
"""Parse text into a spaCy-like doc."""
class YakeExtractor(Protocol):
"""Small protocol for the YAKE extractor used by this module."""
def extract_keywords(self, text: str) -> Iterable[tuple[str, float]]:
"""Return YAKE keyword tuples."""
class YakeExtractorFactory(Protocol):
"""Callable constructor protocol for YAKE keyword extractors."""
def __call__(self, *, lan: str, n: int, dedupLim: float, top: int) -> YakeExtractor: # noqa: N803
"""Create a YAKE keyword extractor.
Args:
lan (str): Language code passed to YAKE.
n (int): Maximum n-gram size to extract.
dedupLim (float): Deduplication similarity threshold.
top (int): Maximum number of keyphrases to return.
Returns:
YakeExtractor: The constructed keyword extractor.
"""
def normalize_candidate_phrase(
phrase_text: str,
config: EbookSearchConfig,
*,
min_tokens: int | None = None,
max_tokens: int | None = None,
strip_leading_article: bool = False,
) -> tuple[str, str, int] | None:
"""Normalize a candidate phrase and validate token bounds.
Args:
phrase_text (str): Raw phrase text to normalize.
config (EbookSearchConfig): Runtime phrase-tuning settings.
min_tokens (int | None): Minimum token count override; defaults to ``config.phrase_min_tokens``.
max_tokens (int | None): Maximum token count override; defaults to ``config.phrase_max_tokens``.
strip_leading_article (bool): Whether to drop a single leading English article.
Returns:
tuple[str, str, int] | None: Display text, normalized phrase, and token count, or ``None``
when the phrase falls outside the token bounds or is ignored.
"""
normalized_tokens = tokenize_with_offsets(phrase_text)
start = 0
if strip_leading_article and normalized_tokens and normalized_tokens[0].text in {"the", "a", "an"}:
start = 1
selected_tokens = normalized_tokens[start:]
min_count = config.phrase_min_tokens if min_tokens is None else min_tokens
max_count = config.phrase_max_tokens if max_tokens is None else max_tokens
if len(selected_tokens) < min_count or len(selected_tokens) > max_count:
return None
phrase_norm = " ".join(token.text for token in selected_tokens)
if phrase_norm in get_ignored_phrases():
return None
display_text = phrase_text[selected_tokens[0].start_char : selected_tokens[-1].end_char].strip()
return display_text or phrase_norm, phrase_norm, len(selected_tokens)
def count_raw_ngrams(tokens: Sequence[str], config: EbookSearchConfig) -> Counter[str]:
"""Count every n-gram window in one normalized token block.
``tokens`` are already normalized (see :func:`tokenize`), so each window's normalized form
is the joined tokens directly. Counting into a plain :class:`Counter` rather than
:class:`PhraseCandidate` objects keeps this hot loop cheap; callers filter ignored phrases
and materialize candidates per unique phrase afterwards, which is far fewer operations than
doing either per window.
Args:
tokens (Sequence[str]): Normalized tokens for one text block.
config (EbookSearchConfig): Runtime phrase-tuning settings.
Returns:
Counter[str]: Raw occurrence counts keyed by normalized phrase.
"""
return Counter(
" ".join(tokens[start : start + ngram_size])
for ngram_size in range(config.phrase_min_tokens, config.phrase_max_tokens + 1)
for start in range(len(tokens) - ngram_size + 1)
)
def extract_raw_ngrams_by_chapter(
chapters: Sequence[str],
config: EbookSearchConfig,
) -> dict[str, PhraseCandidate]:
"""Extract raw n-grams across chapters, tracking both raw counts and chapter spread.
Counting each chapter separately makes chapter spread fall out of dict membership: a phrase's
``chapter_count`` is simply how many per-chapter count maps contain it, so no per-window seen
tracking is needed. This also lets the enrichment step skip re-sliding the same n-gram sizes.
Phrases below the minimum raw count are dropped here rather than materialized: most unique
n-grams occur once, and :func:`filter_storable_candidates` would discard them as too rare
anyway, so building ``PhraseCandidate`` objects for them is wasted work.
Args:
chapters (Sequence[str]): Chapter-like text blocks to slide n-gram windows over.
config (EbookSearchConfig): Runtime phrase-tuning settings.
Returns:
dict[str, PhraseCandidate]: Candidates meeting the minimum raw count, keyed by normalized
phrase, with raw and chapter counts.
"""
chapter_count_maps = [count_raw_ngrams(tokenize(chapter), config) for chapter in chapters]
total_counts: Counter[str] = Counter()
chapter_spread: Counter[str] = Counter()
for chapter_counts in chapter_count_maps:
total_counts.update(chapter_counts)
chapter_spread.update(chapter_counts.keys())
min_raw_count = minimum_candidate_raw_count(config)
ignored = get_ignored_phrases()
return {
phrase_norm: PhraseCandidate(
phrase_text=phrase_norm,
phrase_norm=phrase_norm,
token_count=phrase_norm.count(" ") + 1,
source_raw_ngram=True,
raw_count=raw_count,
chapter_count=chapter_spread[phrase_norm],
)
for phrase_norm, raw_count in total_counts.items()
if raw_count >= min_raw_count and phrase_norm not in ignored
}
@lru_cache(maxsize=2)
def get_yake_extractor(max_ngram: int, top_k: int) -> KeywordExtractor:
"""Return a cached YAKE extractor for the given settings.
Constructing a ``KeywordExtractor`` loads the language's stopword list from disk, so it is
cached and reused across books rather than rebuilt on every call.
Args:
max_ngram (int): Maximum n-gram size to extract.
top_k (int): Maximum number of keyphrases to request.
Returns:
KeywordExtractor: A shared extractor instance for the given settings.
"""
return KeywordExtractor(lan="en", n=max_ngram, dedupLim=0.85, top=top_k)
def extract_yake_candidates(
book_text: str,
config: EbookSearchConfig,
top_k: int = 1000,
) -> dict[str, PhraseCandidate]:
"""Extract YAKE keyphrases when the optional YAKE package is installed.
Args:
book_text (str): Full book text to extract keyphrases from.
config (EbookSearchConfig): Runtime phrase-tuning settings.
top_k (int): Maximum number of YAKE keyphrases to request.
Returns:
dict[str, PhraseCandidate]: Candidates keyed by normalized phrase, with YAKE scores.
"""
extractor = get_yake_extractor(config.phrase_max_tokens, top_k)
out: dict[str, PhraseCandidate] = {}
for phrase_text, yake_score in extractor.extract_keywords(book_text):
normalized = normalize_candidate_phrase(phrase_text, config)
if normalized is None:
continue
display_text, phrase_norm, token_count = normalized
out[phrase_norm] = PhraseCandidate(
phrase_text=display_text,
phrase_norm=phrase_norm,
token_count=token_count,
source_yake=True,
yake_score=float(yake_score),
)
return out
def extract_spacy_candidates(
book_text: str,
nlp: SpacyLanguage,
config: EbookSearchConfig,
) -> dict[str, PhraseCandidate]:
"""Extract spaCy named entities and noun chunks from one text block.
Args:
book_text (str): Text block to parse with spaCy.
nlp (SpacyLanguage): Callable spaCy language pipeline.
config (EbookSearchConfig): Runtime phrase-tuning settings.
Returns:
dict[str, PhraseCandidate]: Candidates keyed by normalized phrase from entities and noun chunks.
"""
out: dict[str, PhraseCandidate] = {}
doc = nlp(book_text)
for ent in doc.ents:
normalized = normalize_candidate_phrase(
ent.text,
config,
max_tokens=config.phrase_max_entity_tokens,
)
if normalized is None:
continue
phrase_text, phrase_norm, token_count = normalized
out[phrase_norm] = PhraseCandidate(
phrase_text=phrase_text,
phrase_norm=phrase_norm,
token_count=token_count,
source_spacy_ner=True,
spacy_label=ent.label_,
)
for chunk in doc.noun_chunks:
normalized = normalize_candidate_phrase(chunk.text, config, strip_leading_article=True)
if normalized is None:
continue
phrase_text, phrase_norm, token_count = normalized
out[phrase_norm] = PhraseCandidate(
phrase_text=phrase_text,
phrase_norm=phrase_norm,
token_count=token_count,
source_spacy_noun_chunk=True,
)
return out
def extract_capitalized_phrases(original_text: str, config: EbookSearchConfig) -> dict[str, PhraseCandidate]:
"""Extract capitalized phrase runs that often carry fictional terms.
Args:
original_text (str): Original-case book text to scan for capitalized runs.
config (EbookSearchConfig): Runtime phrase-tuning settings.
Returns:
dict[str, PhraseCandidate]: Candidates keyed by normalized phrase from capitalized runs.
"""
out: dict[str, PhraseCandidate] = {}
for match in CAPITALIZED_PHRASE_RE.finditer(original_text):
phrase_text = match.group(0).strip()
normalized = normalize_candidate_phrase(
phrase_text,
config,
max_tokens=config.phrase_max_entity_tokens,
)
if normalized is None:
continue
display_text, phrase_norm, token_count = normalized
out[phrase_norm] = PhraseCandidate(
phrase_text=display_text,
phrase_norm=phrase_norm,
token_count=token_count,
source_capitalized=True,
)
return out
def extract_metadata_candidates(
metadata: Mapping[str, object] | None,
config: EbookSearchConfig,
) -> dict[str, PhraseCandidate]:
"""Extract phrases from book metadata values such as title, author, and series.
Args:
metadata (Mapping[str, object] | None): Book metadata values, or ``None`` when unavailable.
config (EbookSearchConfig): Runtime phrase-tuning settings.
Returns:
dict[str, PhraseCandidate]: Candidates keyed by normalized phrase from metadata values.
"""
if metadata is None:
return {}
out: dict[str, PhraseCandidate] = {}
for value in metadata.values():
if value is None:
continue
phrase_text = str(value).strip()
normalized = normalize_candidate_phrase(
phrase_text,
config,
max_tokens=config.phrase_max_entity_tokens,
)
if normalized is None:
continue
display_text, phrase_norm, token_count = normalized
out[phrase_norm] = PhraseCandidate(
phrase_text=display_text,
phrase_norm=phrase_norm,
token_count=token_count,
source_metadata=True,
)
return out
def merge_candidate_sources(*sources: Mapping[str, PhraseCandidate]) -> dict[str, PhraseCandidate]:
"""Merge candidate dictionaries by normalized phrase.
Args:
*sources (Mapping[str, PhraseCandidate]): Candidate maps to combine, keyed by normalized phrase.
Returns:
dict[str, PhraseCandidate]: One merged candidate per normalized phrase.
"""
merged: dict[str, PhraseCandidate] = {}
for source in sources:
for phrase_norm, item in source.items():
existing = merged.setdefault(
phrase_norm,
PhraseCandidate(
phrase_text=item.phrase_text,
phrase_norm=phrase_norm,
token_count=item.token_count,
),
)
merge_candidate(existing, item)
return merged
def merge_candidate(existing: PhraseCandidate, item: PhraseCandidate) -> None:
"""Merge one candidate into an existing candidate object.
Args:
existing (PhraseCandidate): Candidate mutated in place to absorb ``item``.
item (PhraseCandidate): Candidate whose sources, counts, and scores are merged in.
"""
existing.source_raw_ngram = existing.source_raw_ngram or item.source_raw_ngram
existing.source_yake = existing.source_yake or item.source_yake
existing.source_spacy_ner = existing.source_spacy_ner or item.source_spacy_ner
existing.source_spacy_noun_chunk = existing.source_spacy_noun_chunk or item.source_spacy_noun_chunk
existing.source_capitalized = existing.source_capitalized or item.source_capitalized
existing.source_metadata = existing.source_metadata or item.source_metadata
existing.raw_count += item.raw_count
existing.chapter_count = max(existing.chapter_count, item.chapter_count)
if item.yake_score is not None:
existing.yake_score = item.yake_score
if item.spacy_label:
existing.spacy_label = item.spacy_label
def enrich_with_frequency_and_chapter_counts(
candidates: Mapping[str, PhraseCandidate],
chapters: Sequence[str],
*,
counted_sizes: Iterable[int] = (),
) -> dict[str, PhraseCandidate]:
"""Add raw occurrence and chapter-spread counts to candidates.
Candidates whose ``token_count`` is in ``counted_sizes`` are left untouched: those counts
were already computed while sliding the chapters in :func:`extract_raw_ngrams_by_chapter`,
so re-sliding those n-gram sizes here would just duplicate that work.
Args:
candidates (Mapping[str, PhraseCandidate]): Candidates to enrich, keyed by normalized phrase.
chapters (Sequence[str]): Chapter-like text blocks used to count occurrences and spread.
counted_sizes (Iterable[int]): Token counts whose counts are already populated and should be skipped.
Returns:
dict[str, PhraseCandidate]: Candidates with updated ``raw_count`` and ``chapter_count`` values.
"""
if not candidates:
return {}
already_counted = set(counted_sizes)
candidate_sets_by_size: dict[int, set[str]] = defaultdict(set)
for phrase_norm, candidate in candidates.items():
if candidate.token_count in already_counted:
continue
candidate_sets_by_size[candidate.token_count].add(phrase_norm)
enriched = dict(candidates)
if not candidate_sets_by_size:
return enriched
total_counts, chapter_counts = count_candidate_occurrences(candidate_sets_by_size, chapters)
for phrase_norm, candidate in enriched.items():
if candidate.token_count in already_counted:
continue
candidate.raw_count = max(candidate.raw_count, total_counts[phrase_norm])
candidate.chapter_count = chapter_counts[phrase_norm]
return enriched
def count_candidate_occurrences(
candidate_sets_by_size: Mapping[int, set[str]],
chapters: Sequence[str],
) -> tuple[dict[str, int], dict[str, int]]:
"""Count total occurrences and chapter spread for candidate phrases across chapters.
Args:
candidate_sets_by_size (Mapping[int, set[str]]): Candidate normalized phrases grouped by token count.
chapters (Sequence[str]): Chapter-like text blocks to slide n-gram windows over.
Returns:
tuple[dict[str, int], dict[str, int]]: Total occurrence counts and chapter-spread counts,
each keyed by normalized phrase.
"""
total_counts: defaultdict[str, int] = defaultdict(int)
chapter_counts: defaultdict[str, int] = defaultdict(int)
for chapter in chapters:
seen_in_chapter: set[str] = set()
chapter_tokens = tokenize(chapter)
for ngram_size, candidate_norms in candidate_sets_by_size.items():
for start in range(len(chapter_tokens) - ngram_size + 1):
phrase_norm = " ".join(chapter_tokens[start : start + ngram_size])
if phrase_norm not in candidate_norms:
continue
total_counts[phrase_norm] += 1
seen_in_chapter.add(phrase_norm)
for phrase_norm in seen_in_chapter:
chapter_counts[phrase_norm] += 1
return total_counts, chapter_counts
def filter_storable_candidates(
candidates: Mapping[str, PhraseCandidate],
config: EbookSearchConfig,
) -> tuple[dict[str, PhraseCandidate], int, int, int, int]:
"""Remove candidates that should not be persisted.
Args:
candidates (Mapping[str, PhraseCandidate]): Candidates to filter, keyed by normalized phrase.
config (EbookSearchConfig): Runtime phrase-tuning settings.
Returns:
tuple[dict[str, PhraseCandidate], int, int, int, int]: The storable candidates followed by the
counts dropped for being too short, too rare, too common, and junk.
"""
min_raw_count = minimum_candidate_raw_count(config)
filtered: dict[str, PhraseCandidate] = {}
too_short = 0
too_rare = 0
too_common = 0
junk = 0
for phrase_norm, candidate in candidates.items():
if candidate.token_count < config.phrase_min_tokens:
too_short += 1
continue
if candidate.raw_count < min_raw_count:
too_rare += 1
continue
phrase_tokens = phrase_norm.split()
if is_most_common_word_phrase(phrase_tokens):
too_common += 1
continue
if is_junk_phrase(phrase_tokens):
junk += 1
continue
filtered[phrase_norm] = candidate
return filtered, too_short, too_rare, too_common, junk
def minimum_candidate_raw_count(config: EbookSearchConfig) -> int:
"""Return the minimum occurrence count required before storing a candidate.
Args:
config (EbookSearchConfig): Runtime phrase-tuning settings.
Returns:
int: The minimum raw occurrence count, never less than 1.
"""
return max(config.phrase_raw_ngram_min_count, 1)
def is_most_common_word_phrase(phrase_tokens: list[str]) -> bool:
"""Return whether every token in a normalized phrase is a common word.
Args:
phrase_tokens (list[str]): Normalized phrase tokens to inspect.
Returns:
bool: True when the phrase is non-empty and every token is a common word.
"""
common_words = get_most_common_words()
return bool(phrase_tokens) and all(token in common_words for token in phrase_tokens)
def is_junk_phrase(phrase_tokens: list[str]) -> bool:
"""Return whether a normalized phrase is lexical junk not worth LLM judging.
Judged data shows phrases containing a dialogue/action verb or a pronoun contraction are
never kept, and phrases whose tokens are mostly common words almost never are. Possessives
of proper nouns (``chapman's death``) pass because matching is by exact token, and
exactly-half-common bigrams (``data feed``) pass because the common-word rule is strict.
Args:
phrase_tokens (list[str]): Normalized phrase tokens to inspect.
Returns:
bool: True when the phrase contains a junk token or is majority common words.
"""
if not phrase_tokens:
return False
junk_tokens = get_junk_tokens()
if any(token in junk_tokens for token in phrase_tokens):
return True
common_words = get_most_common_words()
half_phrase_len = len(phrase_tokens) // 2
return sum(token in common_words for token in phrase_tokens) > half_phrase_len
def score_candidate(candidate: PhraseCandidate, config: EbookSearchConfig) -> float:
"""Score a phrase candidate before LLM judging.
Args:
candidate (PhraseCandidate): Candidate to score.
config (EbookSearchConfig): Runtime phrase-tuning settings.
Returns:
float: Combined score from sources, frequency, and length, less any penalties.
"""
score = source_score(candidate) + frequency_score(candidate, config) + token_count_score(candidate, config)
if non_raw_source_count(candidate) >= MULTI_SOURCE_MIN_SOURCES:
score += MULTI_SOURCE_SCORE_BONUS
if candidate.phrase_norm in get_ignored_phrases():
score -= 100.0
if has_bad_start(candidate.phrase_norm):
score -= BAD_START_SCORE_PENALTY
if has_bad_end(candidate.phrase_norm):
score -= BAD_END_SCORE_PENALTY
return score
def non_raw_source_count(candidate: PhraseCandidate) -> int:
"""Count the non-raw-ngram extraction sources that produced a candidate.
Args:
candidate (PhraseCandidate): Candidate whose enabled sources are counted.
Returns:
int: Number of enabled sources other than the raw n-gram slide.
"""
return sum(
(
candidate.source_yake,
candidate.source_spacy_ner,
candidate.source_spacy_noun_chunk,
candidate.source_capitalized,
candidate.source_metadata,
)
)
def has_bad_start(phrase_norm: str) -> bool:
"""Return whether a normalized phrase starts with a bad starting token.
Args:
phrase_norm (str): Normalized phrase text to inspect.
Returns:
bool: True when the first token is a known bad starting token.
"""
phrase_tokens = phrase_norm.split()
return bool(phrase_tokens and phrase_tokens[0] in get_bad_starts())
def has_bad_end(phrase_norm: str) -> bool:
"""Return whether a normalized phrase ends with a bad ending token.
Args:
phrase_norm (str): Normalized phrase text to inspect.
Returns:
bool: True when the last token is a known bad ending token.
"""
phrase_tokens = phrase_norm.split()
return bool(phrase_tokens and phrase_tokens[-1] in get_bad_ends())
def source_score(candidate: PhraseCandidate) -> float:
"""Return the score contribution from extraction sources.
Args:
candidate (PhraseCandidate): Candidate whose enabled sources are weighted.
Returns:
float: Summed weight of the candidate's enabled extraction sources.
"""
return sum(
weight
for enabled, weight in (
(candidate.source_yake, 2.0),
(candidate.source_spacy_ner, 2.5),
(candidate.source_spacy_noun_chunk, 1.5),
(candidate.source_capitalized, 2.0),
(candidate.source_metadata, 2.0),
(candidate.source_raw_ngram, 0.5),
)
if enabled
)
def frequency_score(candidate: PhraseCandidate, config: EbookSearchConfig) -> float:
"""Return the score contribution from frequency and chapter spread.
Args:
candidate (PhraseCandidate): Candidate whose counts are scored.
config (EbookSearchConfig): Runtime phrase-tuning settings holding score thresholds.
Returns:
float: Summed weight for each frequency and chapter-spread threshold the candidate meets.
"""
return sum(
weight
for count, threshold, weight in (
(candidate.raw_count, config.phrase_raw_count_score_threshold, 0.5),
(candidate.raw_count, config.phrase_raw_count_high_score_threshold, 0.5),
(candidate.chapter_count, config.phrase_chapter_count_score_threshold, 0.5),
(candidate.chapter_count, config.phrase_chapter_count_high_score_threshold, 0.5),
)
if count >= threshold
)
def token_count_score(candidate: PhraseCandidate, config: EbookSearchConfig) -> float:
"""Return the score contribution from phrase length.
Args:
candidate (PhraseCandidate): Candidate whose token count is scored.
config (EbookSearchConfig): Runtime phrase-tuning settings holding the max token bound.
Returns:
float: Length-based score contribution, which may be negative for over- or under-length phrases.
"""
if candidate.token_count == 1:
return -0.5
if candidate.token_count in {2, 3, 4}:
return 0.5
if candidate.token_count > config.phrase_max_tokens:
return -1.0
return 0.0
def get_sample_contexts(normalized_book_text: str, phrase_norm: str, max_contexts: int = 5) -> list[str]:
"""Return normalized context snippets containing a candidate phrase.
``normalized_book_text`` is expected to already be ``normalize_text``-ed by the caller
so the whole book is not re-normalized for every phrase.
Args:
normalized_book_text (str): Whole book text, already normalized, to search.
phrase_norm (str): Normalized phrase to find contexts around.
max_contexts (int): Maximum number of context snippets to return.
Returns:
list[str]: Up to ``max_contexts`` normalized snippets surrounding the phrase.
"""
contexts: list[str] = []
start = 0
while len(contexts) < max_contexts:
index = normalized_book_text.find(phrase_norm, start)
if index == -1:
break
left = max(0, index - 300)
right = min(len(normalized_book_text), index + len(phrase_norm) + 300)
contexts.append(normalized_book_text[left:right])
start = index + len(phrase_norm)
return contexts
def candidate_source_names(candidate: PhraseCandidate) -> list[str]:
"""Return enabled source names for an extracted candidate.
Args:
candidate (PhraseCandidate): Candidate whose enabled sources are listed.
Returns:
list[str]: Names of the extraction sources that produced the candidate.
"""
names: list[str] = []
if candidate.source_raw_ngram:
names.append("raw_ngram")
if candidate.source_yake:
names.append("yake")
if candidate.source_spacy_ner:
names.append("spacy_ner")
if candidate.source_spacy_noun_chunk:
names.append("spacy_noun_chunk")
if candidate.source_capitalized:
names.append("capitalized")
if candidate.source_metadata:
names.append("metadata")
return names
def extract_phrase_candidates_for_book(
book_text: str,
chapters: Sequence[str],
config: EbookSearchConfig,
*,
nlp: SpacyLanguage | None = None,
metadata: Mapping[str, object] | None = None,
) -> list[PhraseCandidate]:
"""Extract, score, and limit phrase candidates for one book.
Args:
book_text (str): Full book text used for most extraction sources.
chapters (Sequence[str]): Chapter-like text blocks used for spaCy and frequency counts.
config (EbookSearchConfig): Runtime phrase-tuning settings.
nlp (SpacyLanguage | None): Optional spaCy pipeline for entity and noun-chunk sources.
metadata (Mapping[str, object] | None): Optional book metadata used as a candidate source.
Returns:
list[PhraseCandidate]: Scored candidates sorted best-first and capped per book.
"""
started_at = perf_counter()
logger.info(
"ebook_phrase_candidate_extract_start chapters=%s chars=%s min_tokens=%s max_tokens=%s max_candidates=%s",
len(chapters),
len(book_text),
config.phrase_min_tokens,
config.phrase_max_tokens,
config.protected_phrase_max_candidates_per_book,
)
raw_started_at = perf_counter()
raw = extract_raw_ngrams_by_chapter(chapters, config)
logger.info(
"ebook_phrase_candidate_extract_raw_complete candidates=%s duration_ms=%.1f",
len(raw),
(perf_counter() - raw_started_at) * 1000,
)
yake_started_at = perf_counter()
yake_candidates = extract_yake_candidates(book_text, config)
logger.info(
"ebook_phrase_candidate_extract_yake_complete candidates=%s duration_ms=%.1f",
len(yake_candidates),
(perf_counter() - yake_started_at) * 1000,
)
spacy_candidates: dict[str, PhraseCandidate] = {}
if nlp is not None:
spacy_started_at = perf_counter()
for chapter in chapters:
spacy_candidates = merge_candidate_sources(spacy_candidates, extract_spacy_candidates(chapter, nlp, config))
logger.info(
"ebook_phrase_candidate_extract_spacy_complete candidates=%s duration_ms=%.1f",
len(spacy_candidates),
(perf_counter() - spacy_started_at) * 1000,
)
capitalized_started_at = perf_counter()
capitalized = extract_capitalized_phrases(book_text, config)
logger.info(
"ebook_phrase_candidate_extract_capitalized_complete candidates=%s duration_ms=%.1f",
len(capitalized),
(perf_counter() - capitalized_started_at) * 1000,
)
metadata_candidates = extract_metadata_candidates(metadata, config)
candidates = merge_candidate_sources(raw, yake_candidates, spacy_candidates, capitalized, metadata_candidates)
enriched_started_at = perf_counter()
# Raw n-gram sizes were already counted per chapter above, so only enrich the remaining
# (entity-length) sizes here instead of re-sliding every size over the whole book.
candidates = enrich_with_frequency_and_chapter_counts(
candidates,
chapters,
counted_sizes=range(config.phrase_min_tokens, config.phrase_max_tokens + 1),
)
pre_filter_count = len(candidates)
candidates, filtered_too_short, filtered_too_rare, filtered_too_common, filtered_junk = filter_storable_candidates(
candidates, config
)
for candidate in candidates.values():
candidate.candidate_score = score_candidate(candidate, config)
limited = sorted(candidates.values(), key=lambda item: item.candidate_score, reverse=True)[
: config.protected_phrase_max_candidates_per_book
]
logger.info(
"ebook_phrase_candidate_extract_complete raw=%s yake=%s spacy=%s capitalized=%s metadata=%s "
"merged=%s filtered_too_short=%s filtered_too_rare=%s filtered_too_common=%s filtered_junk=%s "
"min_uses=%s storable=%s limited=%s enrich_score_ms=%.1f duration_ms=%.1f",
len(raw),
len(yake_candidates),
len(spacy_candidates),
len(capitalized),
len(metadata_candidates),
pre_filter_count,
filtered_too_short,
filtered_too_rare,
filtered_too_common,
filtered_junk,
minimum_candidate_raw_count(config),
len(candidates),
len(limited),
(perf_counter() - enriched_started_at) * 1000,
(perf_counter() - started_at) * 1000,
)
return limited
@@ -1,370 +0,0 @@
"""Book-level orchestration for candidate n-gram generation and recalculation."""
from __future__ import annotations
import asyncio
import logging
from collections import deque
from time import perf_counter
from typing import TYPE_CHECKING
from sqlalchemy import select
from python.ebook_search.protected_phrases.extraction import extract_phrase_candidates_for_book
from python.ebook_search.protected_phrases.models import (
BookCandidateResult,
PhraseCandidateGenerationResult,
PhraseRecalculationResult,
)
from python.ebook_search.protected_phrases.pool import extract_phrase_candidates_in_pool, get_extraction_pool
from python.ebook_search.protected_phrases.store import (
bulk_upsert_unjudged_candidates,
delete_phrase_data_for_book,
load_book_chapter_texts,
metadata_for_source,
new_candidate_row,
prune_unstorable_unjudged_candidate_phrases,
)
from python.orm.richie import EbookCandidatePhrase, EbookSource
if TYPE_CHECKING:
from collections.abc import Mapping, Sequence
from concurrent.futures import Future
from sqlalchemy.ext.asyncio import AsyncSession
from python.ebook_search.config import EbookSearchConfig
from python.ebook_search.protected_phrases.extraction import SpacyLanguage
from python.ebook_search.protected_phrases.models import PhraseCandidate
logger = logging.getLogger(__name__)
async def generate_candidate_phrases_for_books(
session: AsyncSession,
config: EbookSearchConfig,
*,
only_missing: bool = False,
) -> PhraseCandidateGenerationResult:
"""Create or refresh candidate phrases for indexed books without calling the LLM judge.
Extraction always runs concurrently in the shared process pool so a full backfill uses
multiple cores.
Args:
session (Session): Active database session.
config (EbookSearchConfig): Runtime phrase-tuning settings.
only_missing (bool): When True, only generate for books that have no candidate phrases
yet instead of refreshing every book.
Returns:
PhraseCandidateGenerationResult: Per-corpus counts of books seen, built, and candidates stored.
"""
source_query = select(EbookSource).order_by(EbookSource.id)
if only_missing:
has_candidates = select(EbookCandidatePhrase.id).where(EbookCandidatePhrase.book_id == EbookSource.id)
source_query = source_query.where(~has_candidates.exists())
sources = (await session.scalars(source_query)).all()
books_seen = len(sources)
logger.info(
"ebook_candidate_phrase_generation_start books_seen=%s min_tokens=%s max_tokens=%s max_candidates_per_book=%s",
books_seen,
config.phrase_min_tokens,
config.phrase_max_tokens,
config.protected_phrase_max_candidates_per_book,
)
outcomes = await generate_candidates_for_sources_pooled(session, sources, config)
result = PhraseCandidateGenerationResult(
books_seen=books_seen,
books_built=sum(1 for outcome in outcomes if outcome.built),
candidate_phrases=sum(outcome.candidates for outcome in outcomes),
)
logger.info(
"ebook_candidate_phrase_generation_complete books_seen=%s books_built=%s candidate_total=%s",
result.books_seen,
result.books_built,
result.candidate_phrases,
)
return result
async def generate_candidates_for_sources_pooled(
session: AsyncSession,
sources: Sequence[EbookSource],
config: EbookSearchConfig,
) -> list[BookCandidateResult]:
"""Generate candidate phrases for many books, extracting them concurrently in worker processes.
Chapter loading and row persistence stay on the caller's session (serial), while the CPU-bound
extraction runs in the shared process pool. A bounded window of in-flight books overlaps
extraction across cores without loading every book's candidates into memory at once.
Args:
session (Session): Active database session.
sources (Sequence[EbookSource]): Indexed books to generate candidates for.
config (EbookSearchConfig): Runtime phrase-tuning settings.
Returns:
list[BookCandidateResult]: One result per book.
"""
pool = get_extraction_pool(config.protected_phrase_extraction_workers)
max_in_flight = max(1, config.protected_phrase_extraction_workers) * 2
pending: deque[tuple[EbookSource, Future[list[PhraseCandidate]]]] = deque()
outcomes: list[BookCandidateResult] = []
async def drain_one() -> None:
source, future = pending.popleft()
extracted = await asyncio.wrap_future(future)
outcomes.append(await store_source_candidates(session, source, extracted, config))
try:
for source in sources:
chapters = await load_book_chapter_texts(session, source.id)
if not chapters:
logger.warning("ebook_candidate_phrase_generation_book_empty source_id=%s", source.id)
outcomes.append(BookCandidateResult())
continue
future = pool.submit(
extract_phrase_candidates_for_book,
"\n\n".join(chapters),
chapters,
config,
metadata=metadata_for_source(source),
)
pending.append((source, future))
if len(pending) >= max_in_flight:
await drain_one()
while pending:
await drain_one()
except Exception:
for _, future in pending:
future.cancel()
await session.rollback()
logger.exception("ebook_candidate_phrase_generation_pooled_failed")
raise
return outcomes
async def store_source_candidates(
session: AsyncSession,
source: EbookSource,
limited_candidates: list[PhraseCandidate],
config: EbookSearchConfig,
) -> BookCandidateResult:
"""Persist and commit one book's already-extracted candidates.
Args:
session (AsyncSession): Active database session.
source (EbookSource): Book the candidates belong to.
limited_candidates (list[PhraseCandidate]): Scored candidates to persist.
config (EbookSearchConfig): Runtime phrase-tuning settings.
Returns:
BookCandidateResult: Candidate count and that the book was committed.
"""
book_started_at = perf_counter()
saved_count = await store_candidate_phrases_for_book(session, source.id, None, limited_candidates, config)
await session.commit()
logger.info(
"ebook_candidate_phrase_generation_book_committed source_id=%s candidates=%s duration_ms=%.1f",
source.id,
saved_count,
(perf_counter() - book_started_at) * 1000,
)
return BookCandidateResult(candidates=saved_count, built=True)
async def recalculate_candidate_phrases_for_book(
session: AsyncSession,
source: EbookSource,
config: EbookSearchConfig,
*,
nlp: SpacyLanguage | None = None,
use_process_pool: bool = False,
) -> PhraseRecalculationResult:
"""Remove all book phrase data, regenerate candidates, and commit the completed book.
Args:
session (Session): Active database session.
source (EbookSource): Indexed book to recalculate.
config (EbookSearchConfig): Runtime phrase-tuning settings.
nlp (SpacyLanguage | None): Optional spaCy pipeline for entity and noun-chunk sources.
use_process_pool (bool): Run the CPU-bound extraction in a worker process so concurrent
recalculations do not serialize behind the GIL. Defaults to in-process for callers
(tests, backfills) that do not need it.
Returns:
PhraseRecalculationResult: Deleted-row counts and the number of candidates regenerated.
"""
started_at = perf_counter()
logger.info(
"ebook_candidate_phrase_recalculation_start source_id=%s title=%r",
source.id,
source.title,
)
try:
deleted = await delete_phrase_data_for_book(session, source.id)
chapters = await load_book_chapter_texts(session, source.id)
if not chapters:
logger.warning("ebook_candidate_phrase_recalculation_book_empty source_id=%s", source.id)
await session.commit()
return PhraseRecalculationResult(
book_id=source.id,
deleted_candidates=deleted.deleted_candidates,
deleted_protected_phrases=deleted.deleted_protected_phrases,
deleted_aliases=deleted.deleted_aliases,
deleted_mentions=deleted.deleted_mentions,
candidate_phrases=0,
)
candidate_count = await generate_candidate_phrases_for_book(
session,
source.id,
series_id=None,
chapters=chapters,
config=config,
nlp=nlp,
metadata=metadata_for_source(source),
replace_all=True,
use_process_pool=use_process_pool,
)
await session.commit()
except Exception:
await session.rollback()
logger.exception("ebook_candidate_phrase_recalculation_failed source_id=%s", source.id)
raise
result = PhraseRecalculationResult(
book_id=source.id,
deleted_candidates=deleted.deleted_candidates,
deleted_protected_phrases=deleted.deleted_protected_phrases,
deleted_aliases=deleted.deleted_aliases,
deleted_mentions=deleted.deleted_mentions,
candidate_phrases=candidate_count,
)
logger.info(
"ebook_candidate_phrase_recalculation_complete source_id=%s deleted_candidates=%s "
"deleted_protected=%s deleted_aliases=%s deleted_mentions=%s candidates=%s duration_ms=%.1f",
source.id,
result.deleted_candidates,
result.deleted_protected_phrases,
result.deleted_aliases,
result.deleted_mentions,
result.candidate_phrases,
(perf_counter() - started_at) * 1000,
)
return result
async def generate_candidate_phrases_for_book(
session: AsyncSession,
book_id: int,
series_id: int | None,
chapters: Sequence[str],
config: EbookSearchConfig,
*,
nlp: SpacyLanguage | None = None,
metadata: Mapping[str, object] | None = None,
replace_all: bool = False,
use_process_pool: bool = False,
) -> int:
"""Extract and store candidate phrases for one book without LLM judging.
Args:
session (Session): Active database session.
book_id (int): Book the candidates belong to.
series_id (int | None): Series scope for the stored candidates.
chapters (Sequence[str]): Chapter-like text blocks used for extraction and frequency counts.
config (EbookSearchConfig): Runtime phrase-tuning settings.
nlp (SpacyLanguage | None): Optional spaCy pipeline for entity and noun-chunk sources.
metadata (Mapping[str, object] | None): Optional book metadata used as a candidate source.
replace_all (bool): When the caller has already cleared this book's candidates (e.g. a
recalculation), skip the per-candidate existence lookup and bulk-insert new rows.
use_process_pool (bool): Run the CPU-bound extraction in a worker process to avoid
serializing concurrent requests behind the GIL. Ignored when ``nlp`` is set, since
the spaCy pipeline cannot be sent to a worker process.
Returns:
int: Number of candidate phrase rows stored.
"""
started_at = perf_counter()
book_text = "\n\n".join(chapters)
if use_process_pool and nlp is None:
limited_candidates = await extract_phrase_candidates_in_pool(book_text, chapters, config, metadata=metadata)
else:
limited_candidates = extract_phrase_candidates_for_book(
book_text,
chapters,
config,
nlp=nlp,
metadata=metadata,
)
saved_count = await store_candidate_phrases_for_book(
session,
book_id,
series_id,
limited_candidates,
config,
replace_all=replace_all,
)
logger.info(
"ebook_candidate_phrase_generation_book_duration book_id=%s candidates=%s duration_ms=%.1f",
book_id,
saved_count,
(perf_counter() - started_at) * 1000,
)
return saved_count
async def store_candidate_phrases_for_book(
session: AsyncSession,
book_id: int,
series_id: int | None,
limited_candidates: list[PhraseCandidate],
config: EbookSearchConfig,
*,
replace_all: bool = False,
) -> int:
"""Persist already-extracted candidate phrase rows for one book without committing.
Args:
session (Session): Active database session.
book_id (int): Book the candidates belong to.
series_id (int | None): Series scope for the stored candidates.
limited_candidates (list[PhraseCandidate]): Scored candidates to persist.
config (EbookSearchConfig): Runtime phrase-tuning settings.
replace_all (bool): When the caller has already cleared this book's candidates, skip the
per-candidate existence lookup and bulk-insert new rows.
Returns:
int: Number of candidate phrase rows stored.
"""
save_started_at = perf_counter()
if replace_all:
rows = [new_candidate_row(book_id, series_id, candidate) for candidate in limited_candidates]
session.add_all(rows)
await session.flush()
saved_count = len(rows)
logger.info(
"ebook_candidate_phrase_save_start book_id=%s candidates=%s mode=bulk_insert",
book_id,
len(limited_candidates),
)
else:
pruned_count = await prune_unstorable_unjudged_candidate_phrases(session, book_id, config)
logger.info(
"ebook_candidate_phrase_save_start book_id=%s candidates=%s pruned_unstorable=%s",
book_id,
len(limited_candidates),
pruned_count,
)
saved_count = await bulk_upsert_unjudged_candidates(session, book_id, series_id, limited_candidates)
logger.info(
"ebook_candidate_phrase_save_complete book_id=%s candidates=%s save_ms=%.1f",
book_id,
saved_count,
(perf_counter() - save_started_at) * 1000,
)
return saved_count
@@ -1,511 +0,0 @@
"""Book-level orchestration for LLM judging and promotion of candidate phrases."""
from __future__ import annotations
import asyncio
import json
import logging
import re
from time import perf_counter
from typing import TYPE_CHECKING
import httpx
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from python.ebook_search.llm_interface import request_chat_completion
from python.ebook_search.protected_phrases.extraction import (
candidate_source_names,
get_sample_contexts,
is_junk_phrase,
is_most_common_word_phrase,
score_candidate,
)
from python.ebook_search.protected_phrases.matching import index_chunk_phrase_mentions_for_book
from python.ebook_search.protected_phrases.models import BookJudgmentResult, LLMJudgment, PhraseJudgmentBackfillResult
from python.ebook_search.protected_phrases.store import (
count_protected_phrases,
count_unjudged_candidates,
load_book_text,
load_candidates_for_judgment,
phrase_candidate_from_row,
save_candidate_to_db,
upsert_protected_phrase,
)
from python.ebook_search.protected_phrases.text_normalization import normalize_text
from python.orm.richie import EbookSource
if TYPE_CHECKING:
from collections.abc import Sequence
from sqlalchemy.ext.asyncio import AsyncEngine
from python.ebook_search.config import EbookSearchConfig
from python.ebook_search.protected_phrases.models import PhraseCandidate
from python.orm.richie import EbookProtectedPhrase
JSON_OBJECT_RE = re.compile(r"\{.*\}", re.DOTALL)
logger = logging.getLogger(__name__)
async def judge_candidate_phrases_for_books(
engine: AsyncEngine,
config: EbookSearchConfig,
*,
source_ids: Sequence[int] | None = None,
) -> PhraseJudgmentBackfillResult:
"""Judge candidate phrases for books, fanning LLM calls out across books and phrases.
Up to ``phrase_judge_book_workers`` books are judged at once, and within each book candidates
are judged in concurrent chunks of ``phrase_judge_phrase_workers``. Each book uses its own
short-lived sessions for reads and writes; no database connection is held while LLM calls are
in flight. For a pseudo-single-threaded run (solo testing, debugging), set both worker
settings to 1.
Args:
engine (AsyncEngine): Engine used to open one session per book.
config (EbookSearchConfig): Runtime phrase-tuning settings and chat configuration.
source_ids (Sequence[int] | None): Books to judge; ``None`` judges every indexed book.
Returns:
PhraseJudgmentBackfillResult: Per-corpus counts of books judged, failures, candidates,
protected phrases, and mentions.
"""
if source_ids is None:
async with AsyncSession(engine) as session:
source_ids = list((await session.scalars(select(EbookSource.id).order_by(EbookSource.id))).all())
books_seen = len(source_ids)
book_workers = max(1, config.phrase_judge_book_workers)
phrase_workers = max(1, config.phrase_judge_phrase_workers)
logger.info(
"ebook_candidate_phrase_judgment_start books_seen=%s book_workers=%s phrase_workers=%s "
"confidence_threshold=%.2f",
books_seen,
book_workers,
phrase_workers,
config.protected_phrase_confidence_threshold,
)
book_semaphore = asyncio.Semaphore(book_workers)
max_connections = book_workers * phrase_workers
limits = httpx.Limits(max_connections=max_connections, max_keepalive_connections=max_connections)
async with httpx.AsyncClient(limits=limits) as client:
outcomes = await asyncio.gather(
*(judge_one_book_async(engine, source_id, config, client, book_semaphore) for source_id in source_ids)
)
result = PhraseJudgmentBackfillResult(
books_seen=books_seen,
books_judged=sum(1 for outcome in outcomes if outcome.committed),
books_failed=sum(1 for outcome in outcomes if outcome.failed),
candidates_judged=sum(outcome.judged for outcome in outcomes),
protected_phrases=sum(outcome.protected for outcome in outcomes),
phrase_mentions=sum(outcome.mentions for outcome in outcomes),
)
logger.info(
"ebook_candidate_phrase_judgment_complete books_seen=%s books_judged=%s books_failed=%s "
"candidates_judged=%s protected=%s mentions=%s",
result.books_seen,
result.books_judged,
result.books_failed,
result.candidates_judged,
result.protected_phrases,
result.phrase_mentions,
)
return result
async def judge_one_book_async(
engine: AsyncEngine,
source_id: int,
config: EbookSearchConfig,
client: httpx.AsyncClient,
book_semaphore: asyncio.Semaphore,
) -> BookJudgmentResult:
"""Judge one book concurrently and persist the outcome, honoring the book-level limit.
Args:
engine (AsyncEngine): Engine used to open the book's read and write sessions.
source_id (int): Book to judge candidates for.
config (EbookSearchConfig): Runtime phrase-tuning settings.
client (httpx.AsyncClient): Shared async client for LLM calls.
book_semaphore (asyncio.Semaphore): Caps how many books judge at once.
Returns:
BookJudgmentResult: The book's judgment outcome.
"""
async with book_semaphore:
try:
prepared = await prepare_book_judgment(engine, source_id, config)
if prepared is None:
return BookJudgmentResult()
work_items, target_remaining = prepared
judged = await judge_book_candidates_async(client, config, source_id, work_items, target_remaining)
if not judged:
return BookJudgmentResult()
return await persist_book_judgments(engine, source_id, config, judged)
except Exception:
logger.exception("ebook_candidate_phrase_judgment_book_failed source_id=%s", source_id)
return BookJudgmentResult(failed=True)
async def prepare_book_judgment(
engine: AsyncEngine,
source_id: int,
config: EbookSearchConfig,
) -> tuple[list[tuple[int, PhraseCandidate]], int | None] | None:
"""Load one book's candidates to judge, with sample contexts, on a short-lived read session.
Args:
engine (AsyncEngine): Engine used to open the read session.
source_id (int): Book to load candidates for.
config (EbookSearchConfig): Runtime phrase-tuning settings.
Returns:
tuple[list[tuple[int, PhraseCandidate]], int | None] | None: Candidate rows paired with
in-memory candidates and the remaining protected-phrase target, or ``None`` when the book
has nothing to judge.
"""
judgment_limit = config.protected_phrase_llm_candidates_per_book
if judgment_limit <= 0:
return None
async with AsyncSession(engine) as session:
if not await count_unjudged_candidates(session, source_id, config):
logger.info("ebook_candidate_phrase_judgment_book_skip_no_unjudged source_id=%s", source_id)
return None
existing_protected = await count_protected_phrases(session, source_id)
target_remaining: int | None = None
if config.phrase_target_protected_per_book > 0:
target_remaining = max(config.phrase_target_protected_per_book - existing_protected, 0)
if target_remaining == 0:
logger.info(
"ebook_candidate_phrase_judgment_skipped_target_met source_id=%s existing_protected=%s target=%s",
source_id,
existing_protected,
config.phrase_target_protected_per_book,
)
return None
book_text = await load_book_text(session, source_id)
if not book_text:
logger.warning("ebook_candidate_phrase_judgment_book_empty source_id=%s", source_id)
return None
normalized_book_text = normalize_text(book_text)
# Stored rows may predate the current junk filters and score weights, so re-filter and
# rescore every unjudged row here instead of trusting the persisted candidate_score.
rows = await load_candidates_for_judgment(session, source_id, config)
scored_items: list[tuple[int, PhraseCandidate]] = []
skipped_junk = 0
for row in rows:
candidate = phrase_candidate_from_row(row)
if is_junk_phrase(candidate.phrase_norm.split()):
skipped_junk += 1
continue
candidate.candidate_score = score_candidate(candidate, config)
scored_items.append((row.id, candidate))
scored_items.sort(key=lambda item: item[1].candidate_score, reverse=True)
work_items = scored_items[:judgment_limit]
for _, candidate in work_items:
candidate.sample_contexts = candidate.sample_contexts or get_sample_contexts(
normalized_book_text, candidate.phrase_norm
)
logger.info(
"ebook_candidate_phrase_judgment_candidates_loaded source_id=%s candidates=%s skipped_junk=%s "
"unjudged_rows=%s existing_protected=%s target_remaining=%s judgment_limit=%s",
source_id,
len(work_items),
skipped_junk,
len(rows),
existing_protected,
target_remaining,
judgment_limit,
)
return work_items, target_remaining
async def judge_book_candidates_async(
client: httpx.AsyncClient,
config: EbookSearchConfig,
source_id: int,
work_items: list[tuple[int, PhraseCandidate]],
target_remaining: int | None,
) -> list[tuple[int, PhraseCandidate, LLMJudgment, bool]]:
"""Judge a book's candidates in concurrent chunks, stopping once the target is reached.
Promotion decisions are made in memory so judging can stop early without any database writes.
Args:
client (httpx.AsyncClient): Shared async client for LLM calls.
config (EbookSearchConfig): Runtime phrase-tuning settings.
source_id (int): Book being judged, for logging.
work_items (list[tuple[int, PhraseCandidate]]): Candidate row ids paired with candidates,
in best-first score order.
target_remaining (int | None): Remaining protected-phrase target, or ``None`` for no cap.
Returns:
list[tuple[int, PhraseCandidate, LLMJudgment, bool]]: Judged rows with their judgment and
whether each should be promoted.
"""
chunk_size = max(1, config.phrase_judge_phrase_workers)
judged: list[tuple[int, PhraseCandidate, LLMJudgment, bool]] = []
promoted = 0
for start in range(0, len(work_items), chunk_size):
chunk = work_items[start : start + chunk_size]
judgments = await asyncio.gather(*(judge_candidate_async(client, config, candidate) for _, candidate in chunk))
for (candidate_id, candidate), judgment in zip(chunk, judgments, strict=True):
promote = (target_remaining is None or promoted < target_remaining) and should_protect_judged_candidate(
candidate, judgment, source_id, config, candidate_id=candidate_id
)
if promote:
promoted += 1
judged.append((candidate_id, candidate, judgment, promote))
if target_remaining is not None and promoted >= target_remaining:
break
return judged
async def judge_candidate_async(
client: httpx.AsyncClient,
config: EbookSearchConfig,
candidate: PhraseCandidate,
) -> LLMJudgment:
"""Judge one candidate with the LLM over the shared async client.
Args:
client (httpx.AsyncClient): Shared async client for LLM calls.
config (EbookSearchConfig): Runtime phrase-tuning settings.
candidate (PhraseCandidate): Candidate to judge.
Returns:
LLMJudgment: The parsed judgment.
"""
content = await request_chat_completion(client, config, build_judge_messages(candidate))
return parse_llm_judgment(content, config)
async def persist_book_judgments(
engine: AsyncEngine,
source_id: int,
config: EbookSearchConfig,
judged: list[tuple[int, PhraseCandidate, LLMJudgment, bool]],
) -> BookJudgmentResult:
"""Persist one book's judgments and promotions in a single committed transaction.
Args:
engine (AsyncEngine): Engine used to open the write session.
source_id (int): Book being persisted.
config (EbookSearchConfig): Runtime phrase-tuning settings.
judged (list[tuple[int, PhraseCandidate, LLMJudgment, bool]]): Judged candidates with their
judgment and promotion flag.
Returns:
BookJudgmentResult: The book's committed counts, or a failed result on error.
"""
book_started_at = perf_counter()
async with AsyncSession(engine, expire_on_commit=False) as session:
try:
protected: list[EbookProtectedPhrase] = []
for candidate_id, candidate, judgment, promote in judged:
candidate_row = await save_candidate_to_db(session, source_id, None, candidate, judgment=judgment)
if promote:
protected.append(
await upsert_protected_phrase(session, source_id, None, candidate, judgment, candidate_row)
)
logger.info(
"ebook_candidate_phrase_judgment_candidate_complete source_id=%s candidate_id=%s phrase=%r "
"keep=%s confidence=%.3f category=%r promoted=%s",
source_id,
candidate_id,
candidate.phrase_norm,
judgment.keep,
judgment.confidence,
judgment.category,
promote,
)
await session.flush()
mentions = await index_chunk_phrase_mentions_for_book(session, source_id, config) if protected else 0
await session.commit()
except Exception:
await session.rollback()
logger.exception("ebook_candidate_phrase_judgment_book_persist_failed source_id=%s", source_id)
return BookJudgmentResult(failed=True)
logger.info(
"ebook_candidate_phrase_judgment_book_committed source_id=%s judged=%s protected=%s mentions=%s "
"duration_ms=%.1f",
source_id,
len(judged),
len(protected),
mentions,
(perf_counter() - book_started_at) * 1000,
)
return BookJudgmentResult(judged=len(judged), protected=len(protected), mentions=mentions, committed=True)
def should_protect_judged_candidate(
candidate: PhraseCandidate,
judgment: LLMJudgment,
book_id: int,
config: EbookSearchConfig,
*,
candidate_id: int,
) -> bool:
"""Report whether a judged candidate qualifies to become a protected phrase.
Args:
candidate (PhraseCandidate): In-memory candidate that was judged.
judgment (LLMJudgment): Judge decision for the candidate.
book_id (int): Book the candidate belongs to, for logging.
config (EbookSearchConfig): Runtime phrase-tuning settings.
candidate_id (int): Stored candidate row id the judgment came from, for logging.
Returns:
bool: True when the judged candidate should be promoted to a protected phrase.
"""
if not judgment.keep or judgment.confidence < config.protected_phrase_confidence_threshold:
return False
accepted_norm = normalize_text(judgment.canonical or candidate.phrase_text)
accepted_tokens = accepted_norm.split()
accepted_token_count = len(accepted_tokens)
if accepted_token_count < config.phrase_min_tokens:
logger.info(
"ebook_candidate_phrase_judgment_candidate_skip_short_canonical book_id=%s candidate_id=%s "
"phrase=%r canonical=%r token_count=%s min_tokens=%s",
book_id,
candidate_id,
candidate.phrase_norm,
accepted_norm,
accepted_token_count,
config.phrase_min_tokens,
)
return False
if is_most_common_word_phrase(accepted_tokens):
logger.info(
"ebook_candidate_phrase_judgment_candidate_skip_common_canonical book_id=%s candidate_id=%s "
"phrase=%r canonical=%r",
book_id,
candidate_id,
candidate.phrase_norm,
accepted_norm,
)
return False
return True
def build_judge_messages(candidate: PhraseCandidate) -> list[dict[str, str]]:
"""Build the chat messages used to judge one candidate phrase.
Args:
candidate (PhraseCandidate): Candidate to describe for the judge.
Returns:
list[dict[str, str]]: OpenAI-style system and user messages.
"""
payload = {
"phrase": candidate.phrase_norm,
"token_count": candidate.token_count,
"sources": candidate_source_names(candidate),
"raw_count": candidate.raw_count,
"chapter_count": candidate.chapter_count,
"contexts": candidate.sample_contexts,
}
return [
{
"role": "system",
"content": (
"Judge whether a candidate phrase from a book should be protected for RAG retrieval. "
"Do not extract new phrases. Reject common grammar fragments, ordinary nonspecific phrases, "
"unstable fragments, and phrases kept only because they are frequent. Keep people, places, "
"organizations, factions, events, technologies, fictional conditions, magic systems, formal titles, "
"named concepts, and recurring world-specific terms. Return only a JSON object with keys: keep, "
"canonical, category, aliases, confidence, importance, allow_nested, suppress_children, reason."
),
},
{"role": "user", "content": json.dumps(payload, ensure_ascii=True)},
]
def parse_llm_judgment(content: str, config: EbookSearchConfig) -> LLMJudgment:
"""Parse and validate an LLM phrase-judge response.
Args:
content (str): Raw model response text.
config (EbookSearchConfig): Runtime phrase-tuning settings supplying nesting defaults.
Returns:
LLMJudgment: The parsed and validated judgment.
Raises:
TypeError: If the decoded JSON body is not an object.
"""
body = json.loads(extract_json_object(content))
if not isinstance(body, dict):
msg = "LLM phrase judge response is not a JSON object"
raise TypeError(msg)
aliases = body.get("aliases", ())
if not isinstance(aliases, list | tuple):
aliases = ()
return LLMJudgment(
keep=bool(body.get("keep", False)),
canonical=optional_text(body.get("canonical")),
category=optional_text(body.get("category")),
aliases=tuple(str(alias) for alias in aliases if isinstance(alias, str) and alias.strip()),
confidence=clamped_float(body.get("confidence"), default=0.0),
importance=clamped_float(body.get("importance"), default=0.5),
allow_nested=bool(body.get("allow_nested", config.phrase_default_allow_nested)),
suppress_children=bool(body.get("suppress_children", config.phrase_default_suppress_children)),
reason=optional_text(body.get("reason")),
)
def extract_json_object(content: str) -> str:
"""Extract a JSON object from plain or fenced model output.
Args:
content (str): Raw model response text.
Returns:
str: The substring spanning the first JSON object.
Raises:
ValueError: If no JSON object is found in the response.
"""
stripped = content.strip()
if stripped.startswith("{") and stripped.endswith("}"):
return stripped
match = JSON_OBJECT_RE.search(stripped)
if match is None:
msg = "LLM phrase judge response did not contain a JSON object"
raise ValueError(msg)
return match.group(0)
def optional_text(value: object) -> str | None:
"""Return stripped text for a nullable JSON value.
Args:
value (object): Decoded JSON value that may or may not be a string.
Returns:
str | None: The stripped string, or ``None`` when it is not a non-empty string.
"""
if not isinstance(value, str):
return None
stripped = value.strip()
return stripped or None
def clamped_float(value: object, *, default: float) -> float:
"""Coerce a JSON number into the 0.0 to 1.0 range.
Args:
value (object): Decoded JSON value that may or may not be a number.
default (float): Fallback returned when ``value`` is not numeric.
Returns:
float: The value clamped to ``[0.0, 1.0]``, or ``default`` when non-numeric.
"""
if not isinstance(value, int | float):
return default
return min(max(float(value), 0.0), 1.0)
@@ -1,491 +0,0 @@
"""Runtime protected-phrase matching and chunk mention indexing."""
from __future__ import annotations
import logging
from collections import defaultdict
from typing import TYPE_CHECKING
from sqlalchemy import and_, delete, func, or_, select
from python.ebook_search.protected_phrases.config import get_ignored_phrases
from python.ebook_search.protected_phrases.models import (
ChunkPhraseHit,
HydratedPhraseMatch,
PhraseLookup,
PhraseMatch,
)
from python.ebook_search.protected_phrases.text_normalization import tokenize_with_offsets
from python.orm.richie import (
EbookChunk,
EbookChunkPhraseMention,
EbookPhraseAlias,
EbookProtectedPhrase,
)
if TYPE_CHECKING:
from collections.abc import Iterator, Sequence
from sqlalchemy.ext.asyncio import AsyncSession
from python.ebook_search.config import EbookSearchConfig
from python.ebook_search.protected_phrases.text_normalization import NormalizedToken
logger = logging.getLogger(__name__)
async def load_phrase_lookup(
session: AsyncSession,
config: EbookSearchConfig,
*,
book_id: int | None = None,
series_id: int | None = None,
) -> PhraseLookup:
"""Load protected phrases and aliases into RAM lookup maps.
Args:
session (AsyncSession): Active database session.
config (EbookSearchConfig): Runtime phrase-tuning settings.
book_id (int | None): Optional book scope to restrict loaded phrases.
series_id (int | None): Optional series scope to restrict loaded phrases.
Returns:
PhraseLookup: Normalized phrase and alias maps with the token-window bounds to test.
"""
norm_to_ids: defaultdict[str, list[int]] = defaultdict(list)
alias_to_ids: defaultdict[str, list[int]] = defaultdict(list)
max_tokens = config.phrase_max_tokens
phrase_statement = select(
EbookProtectedPhrase.id,
EbookProtectedPhrase.phrase_norm,
EbookProtectedPhrase.token_count,
)
scope_filter = protected_phrase_scope_filter(book_id=book_id, series_id=series_id)
if scope_filter is not None:
phrase_statement = phrase_statement.where(scope_filter)
for row in await session.execute(phrase_statement):
phrase_id = int(row.id)
phrase_norm = str(row.phrase_norm)
norm_to_ids[phrase_norm].append(phrase_id)
max_tokens = max(max_tokens, int(row.token_count))
alias_statement = select(
EbookPhraseAlias.alias_norm,
EbookPhraseAlias.phrase_id,
).join(EbookProtectedPhrase, EbookProtectedPhrase.id == EbookPhraseAlias.phrase_id)
if scope_filter is not None:
alias_statement = alias_statement.where(scope_filter)
for row in await session.execute(alias_statement):
alias_norm = str(row.alias_norm)
alias_to_ids[alias_norm].append(int(row.phrase_id))
max_tokens = max(max_tokens, len(alias_norm.split()))
return PhraseLookup(
norm_to_phrase_ids={key: tuple(values) for key, values in norm_to_ids.items()},
alias_to_phrase_ids={key: tuple(values) for key, values in alias_to_ids.items()},
min_tokens=config.phrase_min_tokens,
max_tokens=max_tokens,
)
def protected_phrase_scope_filter(*, book_id: int | None, series_id: int | None) -> object | None:
"""Build a SQLAlchemy filter for optional phrase book and series scope.
Args:
book_id (int | None): Optional book scope to include alongside global phrases.
series_id (int | None): Optional series scope to include alongside global phrases.
Returns:
object | None: A combined SQLAlchemy filter clause, or ``None`` when no scope is given.
"""
conditions = []
if book_id is not None:
conditions.append(or_(EbookProtectedPhrase.book_id.is_(None), EbookProtectedPhrase.book_id == book_id))
if series_id is not None:
conditions.append(or_(EbookProtectedPhrase.series_id.is_(None), EbookProtectedPhrase.series_id == series_id))
if not conditions:
return None
return and_(*conditions)
def generate_query_ngrams(
tokens_: Sequence[str],
min_n: int,
max_n: int,
) -> Iterator[tuple[str, int, int]]:
"""Generate normalized query windows from longest to shortest.
Args:
tokens_ (Sequence[str]): Normalized query tokens.
min_n (int): Smallest window size to yield.
max_n (int): Largest window size to yield, capped at the token count.
Yields:
tuple[str, int, int]: Normalized window text with its start and end token indices.
"""
capped_max_n = min(max_n, len(tokens_))
for ngram_size in range(capped_max_n, min_n - 1, -1):
for start in range(len(tokens_) - ngram_size + 1):
end = start + ngram_size
phrase_norm = " ".join(tokens_[start:end])
if phrase_norm in get_ignored_phrases():
continue
yield phrase_norm, start, end
def detect_phrase_candidates(query_text: str, lookup: PhraseLookup) -> list[PhraseMatch]:
"""Detect protected phrase windows in a user query using RAM hash lookups.
Args:
query_text (str): User query text to scan.
lookup (PhraseLookup): In-memory phrase and alias lookup maps.
Returns:
list[PhraseMatch]: Unhydrated phrase matches found in the query.
"""
return detect_phrase_candidates_from_tokens(tokenize_with_offsets(query_text), lookup)
def detect_phrase_candidates_in_text(text: str, lookup: PhraseLookup) -> list[PhraseMatch]:
"""Detect protected phrase windows in arbitrary text with character offsets.
Args:
text (str): Arbitrary text, such as a chunk, to scan.
lookup (PhraseLookup): In-memory phrase and alias lookup maps.
Returns:
list[PhraseMatch]: Unhydrated phrase matches found in the text.
"""
return detect_phrase_candidates_from_tokens(tokenize_with_offsets(text), lookup)
def detect_phrase_candidates_from_tokens(tokens_: Sequence[NormalizedToken], lookup: PhraseLookup) -> list[PhraseMatch]:
"""Detect protected phrase windows from already-normalized tokens.
Args:
tokens_ (Sequence[NormalizedToken]): Normalized tokens with character offsets.
lookup (PhraseLookup): In-memory phrase and alias lookup maps.
Returns:
list[PhraseMatch]: Deduplicated unhydrated phrase matches with token and character spans.
"""
matches: list[PhraseMatch] = []
seen: set[tuple[int | None, str, int, int]] = set()
token_texts = [token.text for token in tokens_]
for phrase_norm, start, end in generate_query_ngrams(token_texts, min_n=lookup.min_tokens, max_n=lookup.max_tokens):
phrase_ids = lookup.norm_to_phrase_ids.get(phrase_norm, ())
alias_ids = lookup.alias_to_phrase_ids.get(phrase_norm, ())
for phrase_id in (*phrase_ids, *alias_ids):
key = (phrase_id, phrase_norm, start, end)
if key in seen:
continue
seen.add(key)
matches.append(
PhraseMatch(
phrase_norm=phrase_norm,
phrase_id=phrase_id,
start_token=start,
end_token=end,
token_count=end - start,
start_char=tokens_[start].start_char,
end_char=tokens_[end - 1].end_char,
)
)
return matches
async def hydrate_matches(session: AsyncSession, matches: Sequence[PhraseMatch]) -> list[HydratedPhraseMatch]:
"""Fetch protected phrase metadata for raw phrase matches.
Args:
session (AsyncSession): Active database session.
matches (Sequence[PhraseMatch]): Unhydrated matches to enrich.
Returns:
list[HydratedPhraseMatch]: Matches with protected-phrase metadata attached.
"""
if not matches:
return []
phrase_ids = sorted({match.phrase_id for match in matches if match.phrase_id is not None})
if not phrase_ids:
return []
rows = {
row.id: row
for row in await session.scalars(select(EbookProtectedPhrase).where(EbookProtectedPhrase.id.in_(phrase_ids)))
}
hydrated: list[HydratedPhraseMatch] = []
for match in matches:
if match.phrase_id is None:
continue
phrase = rows.get(match.phrase_id)
if phrase is None:
continue
hydrated.append(
HydratedPhraseMatch(
phrase_id=phrase.id,
matched_norm=match.phrase_norm,
phrase_text=phrase.phrase_text,
phrase_norm=phrase.phrase_norm,
canonical_id=phrase.canonical_id,
phrase_type=phrase.phrase_type,
token_count=match.token_count,
confidence=phrase.confidence,
importance=phrase.importance,
allow_nested=phrase.allow_nested,
suppress_children=phrase.suppress_children,
start_token=match.start_token,
end_token=match.end_token,
start_char=match.start_char,
end_char=match.end_char,
book_id=phrase.book_id,
series_id=phrase.series_id,
)
)
return hydrated
def overlaps(first: HydratedPhraseMatch, second: HydratedPhraseMatch) -> bool:
"""Return whether two token spans overlap.
Args:
first (HydratedPhraseMatch): First match to compare.
second (HydratedPhraseMatch): Second match to compare.
Returns:
bool: True when the two token spans share at least one token position.
"""
return not (first.end_token <= second.start_token or first.start_token >= second.end_token)
def is_inside(child: HydratedPhraseMatch, parent: HydratedPhraseMatch) -> bool:
"""Return whether one token span is strictly inside another.
Args:
child (HydratedPhraseMatch): Candidate nested match.
parent (HydratedPhraseMatch): Candidate enclosing match.
Returns:
bool: True when ``child`` lies within ``parent`` and is not the same span.
"""
return (
child.start_token >= parent.start_token
and child.end_token <= parent.end_token
and (child.start_token, child.end_token, child.phrase_id)
!= (parent.start_token, parent.end_token, parent.phrase_id)
)
def rank_match(match: HydratedPhraseMatch) -> tuple[float, float, int]:
"""Rank phrase matches by importance, confidence, then token count.
Args:
match (HydratedPhraseMatch): Match to build a sort key for.
Returns:
tuple[float, float, int]: A comparable key of importance, confidence, and token count.
"""
return (match.importance, match.confidence, match.token_count)
def should_suppress(candidate: HydratedPhraseMatch, kept: HydratedPhraseMatch) -> bool:
"""Return whether an already-kept match should suppress a candidate.
Args:
candidate (HydratedPhraseMatch): Match being considered for keeping.
kept (HydratedPhraseMatch): Match already kept that may suppress the candidate.
Returns:
bool: True when the candidate should be dropped in favor of the kept match.
"""
if not overlaps(candidate, kept):
return False
if candidate.canonical_id == kept.canonical_id:
return rank_match(kept) >= rank_match(candidate)
if is_inside(candidate, kept) and kept.suppress_children and not candidate.allow_nested:
return True
return not candidate.allow_nested and rank_match(kept) > rank_match(candidate)
def resolve_overlaps(matches: Sequence[HydratedPhraseMatch]) -> list[HydratedPhraseMatch]:
"""Resolve overlapping phrase matches without relying only on longest match.
Args:
matches (Sequence[HydratedPhraseMatch]): Hydrated matches that may overlap.
Returns:
list[HydratedPhraseMatch]: The kept, non-suppressed matches.
"""
sorted_matches = sorted(
matches,
key=lambda match: (match.start_token, -match.token_count, -match.importance, -match.confidence),
)
kept: list[HydratedPhraseMatch] = []
for candidate in sorted_matches:
if any(should_suppress(candidate, existing) for existing in kept):
continue
kept.append(candidate)
return kept
async def detect_protected_phrases_for_query(
session: AsyncSession,
query_text: str,
config: EbookSearchConfig,
*,
lookup: PhraseLookup | None = None,
book_id: int | None = None,
series_id: int | None = None,
) -> list[HydratedPhraseMatch]:
"""Run the full online protected-phrase query-detection pipeline.
Args:
session (AsyncSession): Active database session.
query_text (str): User query text to detect phrases in.
config (EbookSearchConfig): Runtime phrase-tuning settings.
lookup (PhraseLookup | None): Optional preloaded lookup; loaded on demand when ``None``.
book_id (int | None): Optional book scope for lookup loading.
series_id (int | None): Optional series scope for lookup loading.
Returns:
list[HydratedPhraseMatch]: Hydrated, overlap-resolved phrase matches for the query.
"""
active_lookup = (
lookup
if lookup is not None
else await load_phrase_lookup(session, config, book_id=book_id, series_id=series_id)
)
return resolve_overlaps(await hydrate_matches(session, detect_phrase_candidates(query_text, active_lookup)))
async def index_chunk_phrase_mentions_for_book(
session: AsyncSession,
book_id: int,
config: EbookSearchConfig,
*,
series_id: int | None = None,
lookup: PhraseLookup | None = None,
) -> int:
"""Rebuild chunk phrase mentions for all chunks in one book.
Args:
session (AsyncSession): Active database session.
book_id (int): Book whose chunk mentions are rebuilt.
config (EbookSearchConfig): Runtime phrase-tuning settings.
series_id (int | None): Optional series scope for lookup loading.
lookup (PhraseLookup | None): Optional preloaded lookup; loaded on demand when ``None``.
Returns:
int: Total number of chunk phrase mentions indexed for the book.
"""
active_lookup = (
lookup
if lookup is not None
else await load_phrase_lookup(session, config, book_id=book_id, series_id=series_id)
)
await session.execute(delete(EbookChunkPhraseMention).where(EbookChunkPhraseMention.book_id == book_id))
chunks = await session.scalars(select(EbookChunk).where(EbookChunk.source_id == book_id).order_by(EbookChunk.id))
count = 0
for chunk in chunks:
count += await index_chunk_phrase_mentions(session, chunk, lookup=active_lookup)
await session.flush()
logger.info("ebook_chunk_phrase_mentions_indexed book_id=%s mentions=%s", book_id, count)
return count
async def index_chunk_phrase_mentions(session: AsyncSession, chunk: EbookChunk, *, lookup: PhraseLookup) -> int:
"""Store protected phrase mentions for one chunk.
Args:
session (AsyncSession): Active database session.
chunk (EbookChunk): Chunk whose text is scanned for phrase mentions.
lookup (PhraseLookup): In-memory phrase and alias lookup maps.
Returns:
int: Number of phrase mentions stored for the chunk.
"""
raw_matches = detect_phrase_candidates_in_text(chunk.text, lookup)
hydrated = resolve_overlaps(await hydrate_matches(session, raw_matches))
for match in hydrated:
session.add(
EbookChunkPhraseMention(
chunk_id=chunk.id,
phrase_id=match.phrase_id,
book_id=match.book_id if match.book_id is not None else chunk.source_id,
series_id=match.series_id,
start_char=match.start_char if match.start_char is not None else 0,
end_char=match.end_char,
)
)
return len(hydrated)
async def phrase_hits_for_chunks(
session: AsyncSession,
*,
chunk_ids: Sequence[int],
phrase_ids: Sequence[int],
) -> dict[int, tuple[ChunkPhraseHit, ...]]:
"""Return matched protected phrases with mention counts by chunk id using indexed chunk mentions.
Args:
session (AsyncSession): Active database session.
chunk_ids (Sequence[int]): Chunk ids to look up mentions for.
phrase_ids (Sequence[int]): Protected phrase ids to restrict the results to.
Returns:
dict[int, tuple[ChunkPhraseHit, ...]]: Phrase hits per chunk id, ordered by mention count.
"""
if not chunk_ids or not phrase_ids:
return {}
mention_count = func.count(EbookChunkPhraseMention.phrase_id).label("mention_count")
statement = (
select(
EbookChunkPhraseMention.chunk_id,
EbookProtectedPhrase.id.label("phrase_id"),
EbookProtectedPhrase.phrase_text,
mention_count,
)
.join(EbookProtectedPhrase, EbookProtectedPhrase.id == EbookChunkPhraseMention.phrase_id)
.where(
EbookChunkPhraseMention.chunk_id.in_(chunk_ids),
EbookChunkPhraseMention.phrase_id.in_(phrase_ids),
)
.group_by(EbookChunkPhraseMention.chunk_id, EbookProtectedPhrase.id, EbookProtectedPhrase.phrase_text)
.order_by(EbookChunkPhraseMention.chunk_id, mention_count.desc(), EbookProtectedPhrase.phrase_text)
)
hits: defaultdict[int, list[ChunkPhraseHit]] = defaultdict(list)
for row in await session.execute(statement):
hits[int(row.chunk_id)].append(
ChunkPhraseHit(
phrase_id=int(row.phrase_id),
phrase_text=str(row.phrase_text),
mention_count=int(row.mention_count),
)
)
return {chunk_id: tuple(chunk_hits) for chunk_id, chunk_hits in hits.items()}
async def phrase_hit_counts_for_chunks(
session: AsyncSession,
*,
chunk_ids: Sequence[int],
phrase_ids: Sequence[int],
) -> dict[int, int]:
"""Return phrase-hit counts by chunk id using indexed chunk mentions.
Args:
session (AsyncSession): Active database session.
chunk_ids (Sequence[int]): Chunk ids to count mentions for.
phrase_ids (Sequence[int]): Protected phrase ids to restrict the counts to.
Returns:
dict[int, int]: Total mention count per chunk id.
"""
hits = await phrase_hits_for_chunks(session, chunk_ids=chunk_ids, phrase_ids=phrase_ids)
return {chunk_id: sum(hit.mention_count for hit in chunk_hits) for chunk_id, chunk_hits in hits.items()}
@@ -1,285 +0,0 @@
"""Dataclasses shared by protected phrase extraction, judging, matching, and backfills."""
from __future__ import annotations
from dataclasses import dataclass, field
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from collections.abc import Mapping
@dataclass(slots=True)
class PhraseCandidate:
"""A phrase candidate with merged extraction-source metadata.
Attributes:
phrase_text (str): Display text for the phrase.
phrase_norm (str): Normalized phrase used as the merge key.
token_count (int): Number of normalized tokens in the phrase.
source_raw_ngram (bool): Whether the raw n-gram extractor produced the phrase.
source_yake (bool): Whether YAKE keyword extraction produced the phrase.
source_spacy_ner (bool): Whether spaCy named-entity recognition produced the phrase.
source_spacy_noun_chunk (bool): Whether spaCy noun chunking produced the phrase.
source_capitalized (bool): Whether the capitalized-run extractor produced the phrase.
source_metadata (bool): Whether book metadata produced the phrase.
spacy_label (str | None): spaCy entity label when NER produced the phrase.
raw_count (int): Occurrences counted across the book text.
chapter_count (int): Number of chapters containing the phrase.
yake_score (float | None): Raw YAKE score when available; lower is better.
candidate_score (float): Combined pre-judging score.
sample_contexts (list[str]): Normalized context snippets around occurrences.
"""
phrase_text: str
phrase_norm: str
token_count: int
source_raw_ngram: bool = False
source_yake: bool = False
source_spacy_ner: bool = False
source_spacy_noun_chunk: bool = False
source_capitalized: bool = False
source_metadata: bool = False
spacy_label: str | None = None
raw_count: int = 0
chapter_count: int = 0
yake_score: float | None = None
candidate_score: float = 0.0
sample_contexts: list[str] = field(default_factory=list)
@dataclass(frozen=True, slots=True)
class LLMJudgment:
"""A structured phrase judgment returned by the LLM judge.
Attributes:
keep (bool): Whether the judge accepted the phrase for protection.
canonical (str | None): Canonical phrase text chosen by the judge.
category (str | None): Phrase category such as person, place, or event.
aliases (tuple[str, ...]): Alternate surface forms for the phrase.
confidence (float): Judge confidence between 0.0 and 1.0.
importance (float): Judge importance between 0.0 and 1.0.
allow_nested (bool): Whether the phrase may match inside a larger kept match.
suppress_children (bool): Whether the phrase suppresses matches nested inside it.
reason (str | None): Free-text explanation from the judge.
"""
keep: bool
canonical: str | None
category: str | None
aliases: tuple[str, ...]
confidence: float
importance: float = 0.5
allow_nested: bool = False
suppress_children: bool = True
reason: str | None = None
@dataclass(frozen=True, slots=True)
class PhraseLookup:
"""In-memory lookup maps used for constant-time phrase-window checks.
Attributes:
norm_to_phrase_ids (Mapping[str, tuple[int, ...]]): Normalized phrase to protected phrase ids.
alias_to_phrase_ids (Mapping[str, tuple[int, ...]]): Normalized alias to protected phrase ids.
min_tokens (int): Smallest token-window size to test.
max_tokens (int): Largest token-window size to test.
"""
norm_to_phrase_ids: Mapping[str, tuple[int, ...]]
alias_to_phrase_ids: Mapping[str, tuple[int, ...]]
min_tokens: int
max_tokens: int
@dataclass(frozen=True, slots=True)
class PhraseMatch:
"""An unhydrated query or chunk phrase match.
Attributes:
phrase_norm (str): Normalized text of the matched window.
start_token (int): Index of the first matched token.
end_token (int): Index one past the last matched token.
token_count (int): Number of tokens in the match.
phrase_id (int | None): Matched protected phrase id when known.
start_char (int | None): Start character offset in the source text.
end_char (int | None): End character offset in the source text.
"""
phrase_norm: str
start_token: int
end_token: int
token_count: int
phrase_id: int | None = None
start_char: int | None = None
end_char: int | None = None
@dataclass(frozen=True, slots=True)
class HydratedPhraseMatch:
"""A phrase match with protected-phrase metadata attached.
Attributes:
phrase_id (int): Protected phrase id.
matched_norm (str): Normalized window text that matched.
phrase_text (str): Display text of the protected phrase.
phrase_norm (str): Normalized text of the protected phrase.
canonical_id (str): Deterministic ``category:slug`` identifier.
phrase_type (str | None): Phrase category.
token_count (int): Number of tokens in the match.
confidence (float): Stored judge confidence.
importance (float): Stored judge importance.
allow_nested (bool): Whether the phrase may match inside a larger kept match.
suppress_children (bool): Whether the phrase suppresses matches nested inside it.
start_token (int): Index of the first matched token.
end_token (int): Index one past the last matched token.
start_char (int | None): Start character offset in the source text.
end_char (int | None): End character offset in the source text.
book_id (int | None): Book scope of the phrase.
series_id (int | None): Series scope of the phrase.
"""
phrase_id: int
matched_norm: str
phrase_text: str
phrase_norm: str
canonical_id: str
phrase_type: str | None
token_count: int
confidence: float
importance: float
allow_nested: bool
suppress_children: bool
start_token: int
end_token: int
start_char: int | None = None
end_char: int | None = None
book_id: int | None = None
series_id: int | None = None
@dataclass(frozen=True, slots=True)
class ChunkPhraseHit:
"""One protected phrase with its mention count inside one retrieved chunk.
Attributes:
phrase_id (int): Protected phrase id.
phrase_text (str): Display text of the protected phrase.
mention_count (int): Indexed mentions of the phrase in the chunk.
"""
phrase_id: int
phrase_text: str
mention_count: int
@dataclass(frozen=True, slots=True)
class PhraseCandidateGenerationResult:
"""Summary of candidate phrase extraction for indexed books.
Attributes:
books_seen (int): Indexed books examined.
books_built (int): Books that had candidates generated and committed.
candidate_phrases (int): Candidate phrases stored across all books.
"""
books_seen: int
books_built: int
candidate_phrases: int
@dataclass(frozen=True, slots=True)
class CorpusPhraseStats:
"""Corpus-wide candidate and protected phrase counts for the admin page.
Attributes:
total_books (int): Indexed books in the corpus.
books_with_candidates (int): Books that have candidate phrases generated.
books_fully_judged (int): Books with candidates where every candidate has been judged.
candidate_phrases (int): Candidate phrases stored across all books.
judged_candidates (int): Candidate phrases that have been LLM judged.
unjudged_candidates (int): Candidate phrases still waiting for judgment.
protected_phrases (int): Protected phrases promoted across all books.
"""
total_books: int
books_with_candidates: int
books_fully_judged: int
candidate_phrases: int
judged_candidates: int
unjudged_candidates: int
protected_phrases: int
@dataclass(frozen=True, slots=True)
class PhraseJudgmentBackfillResult:
"""Summary of LLM judging for stored candidate phrases.
Attributes:
books_seen (int): Indexed books examined.
books_judged (int): Books with judgments committed.
books_failed (int): Books rolled back after an error.
candidates_judged (int): Candidate phrases sent to the LLM judge.
protected_phrases (int): Protected phrases promoted from candidates.
phrase_mentions (int): Chunk phrase mentions indexed across all books.
"""
books_seen: int
books_judged: int
books_failed: int
candidates_judged: int
protected_phrases: int
phrase_mentions: int
@dataclass(frozen=True, slots=True)
class BookJudgmentResult:
"""Outcome of judging one book's candidate phrases.
Attributes:
judged (int): Candidate phrases sent to the LLM judge.
protected (int): Protected phrases promoted from candidates.
mentions (int): Chunk phrase mentions indexed for the book.
committed (bool): Whether the book's judgments were committed.
failed (bool): Whether the book was rolled back after an error.
"""
judged: int = 0
protected: int = 0
mentions: int = 0
committed: bool = False
failed: bool = False
@dataclass(frozen=True, slots=True)
class BookCandidateResult:
"""Outcome of generating one book's candidate phrases.
Attributes:
candidates (int): Candidate phrases stored for the book.
built (bool): Whether candidate generation was committed.
"""
candidates: int = 0
built: bool = False
@dataclass(frozen=True, slots=True)
class PhraseRecalculationResult:
"""Summary of phrase cleanup and candidate regeneration for one book.
Attributes:
book_id (int): Book the recalculation ran against.
deleted_candidates (int): Candidate phrase rows deleted.
deleted_protected_phrases (int): Protected phrase rows deleted.
deleted_aliases (int): Phrase alias rows deleted.
deleted_mentions (int): Chunk phrase mention rows deleted.
candidate_phrases (int): Candidate phrases regenerated after cleanup.
"""
book_id: int
deleted_candidates: int
deleted_protected_phrases: int
deleted_aliases: int
deleted_mentions: int
candidate_phrases: int
@@ -1,101 +0,0 @@
"""Process pool for offloading CPU-bound phrase extraction off the request thread.
Phrase extraction is pure-Python CPU work (n-gram sliding, YAKE), so running it inline in a
sync request handler serializes concurrent recalculations behind the GIL. Submitting it to a
``ProcessPoolExecutor`` lets concurrent extractions run in parallel across cores instead. A
``spawn`` context is used so workers do not inherit the parent's database engine, connections,
or server threads.
"""
from __future__ import annotations
import asyncio
import logging
import multiprocessing
import os
from concurrent.futures import ProcessPoolExecutor
from threading import Lock
from typing import TYPE_CHECKING
from python.ebook_search.protected_phrases.extraction import extract_phrase_candidates_for_book
if TYPE_CHECKING:
from collections.abc import Mapping, Sequence
from python.ebook_search.config import EbookSearchConfig
from python.ebook_search.protected_phrases.models import PhraseCandidate
logger = logging.getLogger(__name__)
class _ExtractionPool:
"""Lazily created process-wide extraction pool and the lock guarding it."""
def __init__(self) -> None:
self.lock = Lock()
self.pool: ProcessPoolExecutor | None = None
_extraction_pool = _ExtractionPool()
def get_extraction_pool(max_workers: int) -> ProcessPoolExecutor:
"""Return the shared extraction process pool, creating it on first use.
Args:
max_workers (int): Desired worker count; values below 1 fall back to the CPU count.
Returns:
ProcessPoolExecutor: The shared pool for phrase extraction.
"""
with _extraction_pool.lock:
if _extraction_pool.pool is None:
workers = max_workers if max_workers > 0 else (os.cpu_count() or 1)
_extraction_pool.pool = ProcessPoolExecutor(
max_workers=workers,
mp_context=multiprocessing.get_context("spawn"),
)
logger.info("ebook_phrase_extraction_pool_started workers=%s", workers)
return _extraction_pool.pool
def shutdown_extraction_pool() -> None:
"""Shut down the shared extraction pool if it was started."""
with _extraction_pool.lock:
if _extraction_pool.pool is not None:
_extraction_pool.pool.shutdown(wait=False, cancel_futures=True)
_extraction_pool.pool = None
logger.info("ebook_phrase_extraction_pool_shutdown")
async def extract_phrase_candidates_in_pool(
book_text: str,
chapters: Sequence[str],
config: EbookSearchConfig,
*,
metadata: Mapping[str, object] | None,
) -> list[PhraseCandidate]:
"""Run book phrase extraction in a worker process and await the result.
Only the CPU-bound extraction runs in the worker; the caller keeps all database work in the
request process. The spaCy pipeline is not supported here because it is not picklable, so
this always runs the non-spaCy extraction path.
Args:
book_text (str): Full book text used for extraction.
chapters (Sequence[str]): Chapter-like text blocks used for frequency counts.
config (EbookSearchConfig): Runtime phrase-tuning settings.
metadata (Mapping[str, object] | None): Optional book metadata used as a candidate source.
Returns:
list[PhraseCandidate]: Scored candidates sorted best-first and capped per book.
"""
pool = get_extraction_pool(config.protected_phrase_extraction_workers)
future = pool.submit(
extract_phrase_candidates_for_book,
book_text,
list(chapters),
config,
metadata=dict(metadata) if metadata is not None else None,
)
return await asyncio.wrap_future(future)
@@ -1,700 +0,0 @@
"""Database persistence for candidate and protected phrase rows."""
from __future__ import annotations
import logging
import re
from typing import TYPE_CHECKING
from sqlalchemy import delete, func, or_, select
from sqlalchemy.dialects.postgresql import insert as pg_insert
from sqlalchemy.dialects.sqlite import insert as sqlite_insert
from python.ebook_search.protected_phrases.extraction import minimum_candidate_raw_count
from python.ebook_search.protected_phrases.models import (
CorpusPhraseStats,
PhraseCandidate,
PhraseRecalculationResult,
)
from python.ebook_search.protected_phrases.text_normalization import normalize_text
from python.orm.richie import (
EbookCandidatePhrase,
EbookChunk,
EbookChunkPhraseMention,
EbookPhraseAlias,
EbookProtectedPhrase,
EbookSource,
)
if TYPE_CHECKING:
from collections.abc import Sequence
from sqlalchemy.dialects.postgresql.dml import Insert as PostgresInsert
from sqlalchemy.dialects.sqlite.dml import Insert as SqliteInsert
from sqlalchemy.ext.asyncio import AsyncSession
from python.ebook_search.config import EbookSearchConfig
from python.ebook_search.protected_phrases.models import LLMJudgment
from python.orm.richie.base import TableBase
logger = logging.getLogger(__name__)
def dialect_insert(session: AsyncSession, table: type[TableBase]) -> PostgresInsert | SqliteInsert:
"""Return a dialect-specific INSERT construct that supports ``ON CONFLICT DO UPDATE``.
Production runs on PostgreSQL while tests run on SQLite; both support upserts with
compatible SQLAlchemy constructs, so the correct one is chosen from the bound dialect.
Args:
session (AsyncSession): Active database session whose bind selects the dialect.
table (type[TableBase]): Mapped table to insert into.
Returns:
PostgresInsert | SqliteInsert: A dialect insert exposing ``on_conflict_do_update``.
"""
if session.get_bind().dialect.name == "sqlite":
return sqlite_insert(table)
return pg_insert(table)
async def load_book_text(session: AsyncSession, book_id: int) -> str:
"""Load a book's indexed chunk text as one string for phrase extraction.
Args:
session (AsyncSession): Active database session.
book_id (int): Book whose chunk text is loaded.
Returns:
str: The book's chunk text joined into a single string.
"""
texts = await session.scalars(
select(EbookChunk.text).where(EbookChunk.source_id == book_id).order_by(EbookChunk.chunk_index)
)
return "\n\n".join(stripped for text in texts if (stripped := text.strip()))
async def load_book_chapter_texts(session: AsyncSession, book_id: int) -> list[str]:
"""Reconstruct chapter-like text blocks from indexed chunks for phrase extraction.
Args:
session (AsyncSession): Active database session.
book_id (int): Book whose chunks are grouped into chapters.
Returns:
list[str]: Non-empty chapter-like text blocks in chunk order.
"""
rows = await session.execute(
select(EbookChunk.chapter_id, EbookChunk.text)
.where(EbookChunk.source_id == book_id)
.order_by(EbookChunk.chunk_index)
)
chapters: list[str] = []
current_chapter_id: int | None = None
current_parts: list[str] = []
have_current = False
for chapter_id, text in rows:
if have_current and chapter_id != current_chapter_id:
chapter_text = "\n\n".join(current_parts).strip()
if chapter_text:
chapters.append(chapter_text)
current_parts = []
current_chapter_id = chapter_id
current_parts.append(str(text))
have_current = True
if current_parts:
chapter_text = "\n\n".join(current_parts).strip()
if chapter_text:
chapters.append(chapter_text)
return chapters
def metadata_for_source(source: EbookSource) -> dict[str, object | None]:
"""Return phrase extraction metadata for one indexed source.
Args:
source (EbookSource): Indexed source to read metadata from.
Returns:
dict[str, object | None]: Title, author, language, publisher, and identifier values.
"""
return {
"title": source.title,
"author": source.author,
"language": source.language,
"publisher": source.publisher,
"identifier": source.identifier,
}
async def metadata_for_source_id(session: AsyncSession, source_id: int) -> dict[str, object | None]:
"""Return phrase extraction metadata for one indexed source by id.
Args:
session (AsyncSession): Active database session.
source_id (int): Id of the indexed source to read metadata from.
Returns:
dict[str, object | None]: Title, author, language, publisher, and identifier values.
Raises:
ValueError: If no source exists with the given id.
"""
source = await session.get(EbookSource, source_id)
if source is None:
msg = f"No indexed source with id {source_id}"
raise ValueError(msg)
return metadata_for_source(source)
async def count_protected_phrases(session: AsyncSession, book_id: int) -> int:
"""Count stored protected phrases for one book.
Args:
session (AsyncSession): Active database session.
book_id (int): Book whose protected phrases are counted.
Returns:
int: Number of protected phrases stored for the book.
"""
return (
await session.scalars(
select(func.count(EbookProtectedPhrase.id)).where(EbookProtectedPhrase.book_id == book_id)
)
).one()
async def count_unjudged_candidates(session: AsyncSession, book_id: int, config: EbookSearchConfig) -> int:
"""Count storable candidate rows for a book that have not yet been judged.
Args:
session (AsyncSession): Active database session.
book_id (int): Book whose unjudged candidates are counted.
config (EbookSearchConfig): Runtime phrase-tuning settings supplying storage thresholds.
Returns:
int: Number of storable, unjudged candidate rows for the book.
"""
return (
await session.scalars(
select(func.count(EbookCandidatePhrase.id)).where(
EbookCandidatePhrase.book_id == book_id,
EbookCandidatePhrase.llm_judged.is_(False),
EbookCandidatePhrase.token_count >= config.phrase_min_tokens,
EbookCandidatePhrase.raw_count >= minimum_candidate_raw_count(config),
)
)
).one()
async def corpus_phrase_stats(session: AsyncSession) -> CorpusPhraseStats:
"""Summarize candidate and protected phrase coverage across the whole corpus.
Args:
session (AsyncSession): Active database session.
Returns:
CorpusPhraseStats: Corpus-wide phrase counts and per-book coverage counts.
"""
total_books = (await session.scalars(select(func.count(EbookSource.id)))).one()
candidate_phrases, judged_candidates, books_with_candidates, books_with_unjudged = (
await session.execute(
select(
func.count(EbookCandidatePhrase.id),
func.count(EbookCandidatePhrase.id).filter(EbookCandidatePhrase.llm_judged.is_(True)),
func.count(func.distinct(EbookCandidatePhrase.book_id)),
func.count(func.distinct(EbookCandidatePhrase.book_id)).filter(
EbookCandidatePhrase.llm_judged.is_(False)
),
)
)
).one()
protected_phrases = (await session.scalars(select(func.count(EbookProtectedPhrase.id)))).one()
return CorpusPhraseStats(
total_books=total_books,
books_with_candidates=books_with_candidates,
books_fully_judged=books_with_candidates - books_with_unjudged,
candidate_phrases=candidate_phrases,
judged_candidates=judged_candidates,
unjudged_candidates=candidate_phrases - judged_candidates,
protected_phrases=protected_phrases,
)
async def book_ids_pending_first_judgment(session: AsyncSession) -> list[int]:
"""Return books that have candidate phrases but no judged candidates yet.
Args:
session (AsyncSession): Active database session.
Returns:
list[int]: Book ids with candidates where judging has never run, ordered by id.
"""
judged_books = select(EbookCandidatePhrase.book_id).where(EbookCandidatePhrase.llm_judged.is_(True)).distinct()
return list(
(
await session.scalars(
select(EbookCandidatePhrase.book_id)
.where(EbookCandidatePhrase.book_id.not_in(judged_books))
.distinct()
.order_by(EbookCandidatePhrase.book_id)
)
).all()
)
async def load_candidates_for_judgment(
session: AsyncSession,
book_id: int,
config: EbookSearchConfig,
) -> Sequence[EbookCandidatePhrase]:
"""Load every storable unjudged candidate row for a book.
Rows may have been stored before the current junk filters and score weights existed, so
callers re-check :func:`is_junk_phrase` and rescore before selecting what to judge.
Args:
session (AsyncSession): Active database session.
book_id (int): Book whose candidates are loaded.
config (EbookSearchConfig): Runtime phrase-tuning settings supplying storage thresholds.
Returns:
Sequence[EbookCandidatePhrase]: Storable, unjudged candidate rows ordered by stored score.
"""
query = (
select(EbookCandidatePhrase)
.where(
EbookCandidatePhrase.book_id == book_id,
EbookCandidatePhrase.llm_judged.is_(False),
EbookCandidatePhrase.token_count >= config.phrase_min_tokens,
EbookCandidatePhrase.raw_count >= minimum_candidate_raw_count(config),
)
.order_by(
EbookCandidatePhrase.candidate_score.desc(),
EbookCandidatePhrase.raw_count.desc(),
EbookCandidatePhrase.id,
)
)
return (await session.scalars(query)).all()
def phrase_candidate_from_row(row: EbookCandidatePhrase) -> PhraseCandidate:
"""Recreate an in-memory candidate from a persisted candidate row.
Args:
row (EbookCandidatePhrase): Stored candidate row to convert.
Returns:
PhraseCandidate: An in-memory candidate mirroring the row's fields.
"""
return PhraseCandidate(
phrase_text=row.phrase_text,
phrase_norm=row.phrase_norm,
token_count=row.token_count,
source_raw_ngram=row.source_raw_ngram,
source_yake=row.source_yake,
source_spacy_ner=row.source_spacy_ner,
source_spacy_noun_chunk=row.source_spacy_noun_chunk,
source_capitalized=row.source_capitalized,
source_metadata=row.source_metadata,
spacy_label=row.spacy_label,
raw_count=row.raw_count,
chapter_count=row.chapter_count,
yake_score=row.yake_score,
candidate_score=row.candidate_score,
sample_contexts=row.sample_contexts or [],
)
def candidate_row_values(
book_id: int,
series_id: int | None,
candidate: PhraseCandidate,
*,
judgment: LLMJudgment | None,
) -> dict[str, object]:
"""Build the column values for one candidate phrase upsert.
Args:
book_id (int): Book the candidate belongs to.
series_id (int | None): Series scope stored on the row.
candidate (PhraseCandidate): Candidate whose fields are written to the row.
judgment (LLMJudgment | None): Judgment to record, or ``None`` to leave the row unjudged.
Returns:
dict[str, object]: Column values keyed by column name.
"""
values: dict[str, object] = {
"book_id": book_id,
"phrase_norm": candidate.phrase_norm,
"series_id": series_id,
"phrase_text": candidate.phrase_text,
"token_count": candidate.token_count,
"source_raw_ngram": candidate.source_raw_ngram,
"source_yake": candidate.source_yake,
"source_spacy_ner": candidate.source_spacy_ner,
"source_spacy_noun_chunk": candidate.source_spacy_noun_chunk,
"source_capitalized": candidate.source_capitalized,
"source_metadata": candidate.source_metadata,
"spacy_label": candidate.spacy_label,
"raw_count": candidate.raw_count,
"chapter_count": candidate.chapter_count,
"yake_score": candidate.yake_score,
"candidate_score": candidate.candidate_score,
"llm_judged": judgment is not None,
}
if candidate.sample_contexts:
values["sample_contexts"] = list(candidate.sample_contexts)
if judgment is not None:
values.update(
llm_keep=judgment.keep,
llm_confidence=judgment.confidence,
llm_category=judgment.category,
llm_reason=judgment.reason,
)
return values
async def save_candidate_to_db(
session: AsyncSession,
book_id: int,
series_id: int | None,
candidate: PhraseCandidate,
*,
judgment: LLMJudgment | None,
) -> EbookCandidatePhrase:
"""Insert or update one candidate phrase row.
Args:
session (AsyncSession): Active database session.
book_id (int): Book the candidate belongs to.
series_id (int | None): Series scope stored on the row.
candidate (PhraseCandidate): Candidate whose fields are written to the row.
judgment (LLMJudgment | None): Judgment to record, or ``None`` to leave the row unjudged.
Returns:
EbookCandidatePhrase: The inserted or updated candidate row.
"""
values = candidate_row_values(book_id, series_id, candidate, judgment=judgment)
# Preserve an existing judgment when this call is only refreshing candidate fields.
skip_update = {"book_id", "phrase_norm"}
if judgment is None:
skip_update.add("llm_judged")
insert_statement = dialect_insert(session, EbookCandidatePhrase).values(**values)
statement = insert_statement.on_conflict_do_update(
index_elements=["book_id", "phrase_norm"],
set_={column: insert_statement.excluded[column] for column in values if column not in skip_update},
).returning(EbookCandidatePhrase)
return (await session.scalars(statement, execution_options={"populate_existing": True})).one()
BULK_CANDIDATE_UPSERT_CHUNK = 1000
async def bulk_upsert_unjudged_candidates(
session: AsyncSession,
book_id: int,
series_id: int | None,
candidates: Sequence[PhraseCandidate],
) -> int:
"""Insert or update many freshly extracted candidate rows in chunked multi-row upserts.
Saving one row per statement costs one database round trip per candidate, which dominated
generation time for full books, so candidates are written ``BULK_CANDIDATE_UPSERT_CHUNK``
rows per statement instead. Existing judgments and sample contexts are never overwritten:
fresh extractions carry no contexts, and ``llm_judged`` plus the ``llm_*`` columns are left
out of the conflict update. Candidates must have unique ``phrase_norm`` values, as produced
by extraction, since one multi-row upsert cannot touch the same row twice.
Args:
session (Session): Active database session.
book_id (int): Book the candidates belong to.
series_id (int | None): Series scope stored on the rows.
candidates (Sequence[PhraseCandidate]): Freshly extracted candidates to persist.
Returns:
int: Number of candidate rows written.
"""
values = [
candidate_row_values(book_id, series_id, candidate, judgment=None)
for candidate in candidates
if not candidate.sample_contexts
]
if len(values) != len(candidates):
msg = "bulk_upsert_unjudged_candidates only accepts freshly extracted candidates without sample contexts"
raise ValueError(msg)
skip_update = {"book_id", "phrase_norm", "llm_judged"}
for chunk_start in range(0, len(values), BULK_CANDIDATE_UPSERT_CHUNK):
chunk = values[chunk_start : chunk_start + BULK_CANDIDATE_UPSERT_CHUNK]
insert_statement = dialect_insert(session, EbookCandidatePhrase).values(chunk)
statement = insert_statement.on_conflict_do_update(
index_elements=["book_id", "phrase_norm"],
set_={column: insert_statement.excluded[column] for column in chunk[0] if column not in skip_update},
)
await session.execute(statement)
return len(values)
def new_candidate_row(book_id: int, series_id: int | None, candidate: PhraseCandidate) -> EbookCandidatePhrase:
"""Build a fresh unjudged candidate row without checking for an existing one.
Unlike :func:`save_candidate_to_db`, this does no lookup, so it is only safe when the caller
guarantees there is no existing row for ``(book_id, candidate.phrase_norm)`` — for example
right after :func:`delete_phrase_data_for_book` has cleared the book.
Args:
book_id (int): Book the candidate belongs to.
series_id (int | None): Series scope stored on the row.
candidate (PhraseCandidate): Candidate whose fields are written to the row.
Returns:
EbookCandidatePhrase: A new, unattached candidate row.
"""
row = EbookCandidatePhrase(book_id=book_id, phrase_norm=candidate.phrase_norm)
row.llm_judged = False
row.series_id = series_id
row.phrase_text = candidate.phrase_text
row.token_count = candidate.token_count
row.source_raw_ngram = candidate.source_raw_ngram
row.source_yake = candidate.source_yake
row.source_spacy_ner = candidate.source_spacy_ner
row.source_spacy_noun_chunk = candidate.source_spacy_noun_chunk
row.source_capitalized = candidate.source_capitalized
row.source_metadata = candidate.source_metadata
row.spacy_label = candidate.spacy_label
row.raw_count = candidate.raw_count
row.chapter_count = candidate.chapter_count
row.yake_score = candidate.yake_score
row.candidate_score = candidate.candidate_score
if candidate.sample_contexts:
row.sample_contexts = list(candidate.sample_contexts)
return row
async def upsert_protected_phrase(
session: AsyncSession,
book_id: int,
series_id: int | None,
candidate: PhraseCandidate,
judgment: LLMJudgment,
source_candidate: EbookCandidatePhrase,
) -> EbookProtectedPhrase:
"""Insert or update one accepted protected phrase and its aliases.
Args:
session (AsyncSession): Active database session.
book_id (int): Book the protected phrase belongs to.
series_id (int | None): Series scope stored on the phrase.
candidate (PhraseCandidate): Candidate the phrase was promoted from.
judgment (LLMJudgment): Accepted judgment supplying canonical text, category, and aliases.
source_candidate (EbookCandidatePhrase): Candidate row the phrase was promoted from.
Returns:
EbookProtectedPhrase: The inserted or updated protected phrase row.
Raises:
ValueError: If the chosen phrase text normalizes to empty.
"""
phrase_text = judgment.canonical or candidate.phrase_text
phrase_norm = normalize_text(phrase_text)
if not phrase_norm:
msg = f"Protected phrase normalized to empty text: {phrase_text!r}"
raise ValueError(msg)
values = {
"book_id": book_id,
"phrase_norm": phrase_norm,
"series_id": series_id,
"phrase_text": phrase_text,
"canonical_id": make_canonical_id(judgment, phrase_norm),
"phrase_type": judgment.category,
"token_count": len(phrase_norm.split()),
"confidence": judgment.confidence,
"importance": judgment.importance,
"allow_nested": judgment.allow_nested,
"suppress_children": judgment.suppress_children,
"source_candidate_id": source_candidate.id,
}
insert_statement = dialect_insert(session, EbookProtectedPhrase).values(**values)
statement = insert_statement.on_conflict_do_update(
index_elements=["book_id", "phrase_norm"],
set_={
column: insert_statement.excluded[column] for column in values if column not in {"book_id", "phrase_norm"}
},
).returning(EbookProtectedPhrase)
row = (await session.scalars(statement, execution_options={"populate_existing": True})).one()
for alias_text in judgment.aliases:
await upsert_phrase_alias(session, row, alias_text)
return row
async def upsert_phrase_alias(
session: AsyncSession,
phrase: EbookProtectedPhrase,
alias_text: str,
) -> EbookPhraseAlias | None:
"""Insert or update one protected phrase alias.
Args:
session (AsyncSession): Active database session.
phrase (EbookProtectedPhrase): Protected phrase the alias points to.
alias_text (str): Alias surface form to store.
Returns:
EbookPhraseAlias | None: The alias row, or ``None`` when the alias is empty or equals the phrase.
"""
alias_norm = normalize_text(alias_text)
if not alias_norm or alias_norm == phrase.phrase_norm:
return None
insert_statement = dialect_insert(session, EbookPhraseAlias).values(
phrase_id=phrase.id,
alias_norm=alias_norm,
alias_text=alias_text,
confidence=1.0,
)
statement = insert_statement.on_conflict_do_update(
index_elements=["phrase_id", "alias_norm"],
set_={
"alias_text": insert_statement.excluded.alias_text,
"confidence": insert_statement.excluded.confidence,
},
).returning(EbookPhraseAlias)
return (await session.scalars(statement, execution_options={"populate_existing": True})).one()
def make_canonical_id(judgment: LLMJudgment, phrase_norm: str) -> str:
"""Create a deterministic canonical id from a judgment category and phrase.
Args:
judgment (LLMJudgment): Judgment supplying the phrase category.
phrase_norm (str): Normalized phrase text to slugify.
Returns:
str: A ``category:slug`` canonical identifier.
"""
category = slugify_identifier(judgment.category or "phrase")
phrase_slug = slugify_identifier(phrase_norm)
return f"{category}:{phrase_slug}"
def slugify_identifier(value: str) -> str:
"""Normalize text for use inside a canonical id.
Args:
value (str): Text to slugify.
Returns:
str: A lowercase underscore slug, or ``"unknown"`` when empty.
"""
slug = re.sub(r"[^a-z0-9]+", "_", normalize_text(value).replace("'", ""))
return slug.strip("_") or "unknown"
async def prune_unstorable_unjudged_candidate_phrases(
session: AsyncSession,
book_id: int,
config: EbookSearchConfig,
) -> int:
"""Delete old unjudged candidate rows that no longer satisfy storage filters.
Args:
session (AsyncSession): Active database session.
book_id (int): Book whose stale candidates are pruned.
config (EbookSearchConfig): Runtime phrase-tuning settings supplying storage thresholds.
Returns:
int: Number of candidate rows deleted.
"""
deleted = rowcount(
await session.execute(
delete(EbookCandidatePhrase).where(
EbookCandidatePhrase.book_id == book_id,
EbookCandidatePhrase.llm_judged.is_(False),
or_(
EbookCandidatePhrase.token_count < config.phrase_min_tokens,
EbookCandidatePhrase.raw_count < minimum_candidate_raw_count(config),
),
)
)
)
if deleted:
logger.info(
"ebook_candidate_phrase_unstorable_pruned book_id=%s deleted=%s min_tokens=%s min_uses=%s",
book_id,
deleted,
config.phrase_min_tokens,
minimum_candidate_raw_count(config),
)
return deleted
async def delete_phrase_data_for_book(session: AsyncSession, book_id: int) -> PhraseRecalculationResult:
"""Delete all candidate, protected, alias, and mention phrase data for one book.
Args:
session (AsyncSession): Active database session.
book_id (int): Book whose phrase data is deleted.
Returns:
PhraseRecalculationResult: Deleted-row counts with ``candidate_phrases`` set to 0.
"""
protected_ids = (
await session.scalars(select(EbookProtectedPhrase.id).where(EbookProtectedPhrase.book_id == book_id))
).all()
deleted_aliases = 0
if protected_ids:
deleted_aliases = rowcount(
await session.execute(delete(EbookPhraseAlias).where(EbookPhraseAlias.phrase_id.in_(protected_ids)))
)
deleted_mentions = rowcount(
await session.execute(delete(EbookChunkPhraseMention).where(EbookChunkPhraseMention.book_id == book_id))
)
if protected_ids:
deleted_mentions += rowcount(
await session.execute(
delete(EbookChunkPhraseMention).where(EbookChunkPhraseMention.phrase_id.in_(protected_ids))
)
)
deleted_protected = rowcount(
await session.execute(delete(EbookProtectedPhrase).where(EbookProtectedPhrase.book_id == book_id))
)
deleted_candidates = rowcount(
await session.execute(delete(EbookCandidatePhrase).where(EbookCandidatePhrase.book_id == book_id))
)
await session.flush()
logger.info(
"ebook_candidate_phrase_data_deleted book_id=%s candidates=%s protected=%s aliases=%s mentions=%s",
book_id,
deleted_candidates,
deleted_protected,
deleted_aliases,
deleted_mentions,
)
return PhraseRecalculationResult(
book_id=book_id,
deleted_candidates=deleted_candidates,
deleted_protected_phrases=deleted_protected,
deleted_aliases=deleted_aliases,
deleted_mentions=deleted_mentions,
candidate_phrases=0,
)
def rowcount(result: object) -> int:
"""Return a safe integer rowcount from a SQLAlchemy execution result.
Args:
result (object): SQLAlchemy execution result that may expose ``rowcount``.
Returns:
int: The result's rowcount, or 0 when it is missing or negative.
"""
count = getattr(result, "rowcount", 0)
return int(count if count is not None and count >= 0 else 0)
@@ -1,91 +0,0 @@
"""Protected phrase extraction, storage, and runtime matching."""
from __future__ import annotations
import re
from dataclasses import dataclass
JSON_OBJECT_RE = re.compile(r"\{.*\}", re.DOTALL)
@dataclass(frozen=True, slots=True)
class NormalizedToken:
"""A normalized token plus its source character span."""
text: str
start_char: int
end_char: int
def normalize_text(text: str) -> str:
"""Normalize text for phrase storage and lookup.
Args:
text (str): Raw text to normalize.
Returns:
str: Normalized tokens joined by single spaces.
"""
return " ".join(token.text for token in tokenize_with_offsets(text))
def tokenize(text: str) -> list[str]:
"""Normalize and split text into phrase-detection tokens.
Args:
text (str): Raw text to tokenize.
Returns:
list[str]: Normalized token strings.
"""
return [token.text for token in tokenize_with_offsets(text)]
def tokenize_with_offsets(text: str) -> list[NormalizedToken]:
"""Normalize text into tokens while preserving original character offsets.
Args:
text (str): Raw text to tokenize.
Returns:
list[NormalizedToken]: Normalized tokens with their source character spans.
"""
tokens: list[NormalizedToken] = []
current: list[str] = []
start_char: int | None = None
for index, char in enumerate(text):
normalized = normalize_char(char)
if normalized == " ":
if current and start_char is not None:
tokens.append(NormalizedToken(text="".join(current), start_char=start_char, end_char=index))
current = []
start_char = None
continue
if start_char is None:
start_char = index
current.append(normalized)
if current and start_char is not None:
tokens.append(NormalizedToken(text="".join(current), start_char=start_char, end_char=len(text)))
return tokens
def normalize_char(char: str) -> str:
"""Normalize one character into a token character or a separator.
Args:
char (str): Single source character to normalize.
Returns:
str: The normalized token character, or a space acting as a separator.
"""
if char in {"\u2019", "\u2018"}:
return "'"
if char in {"-", "\u2013", "\u2014"}:
return " "
lowered = char.lower()
if lowered in "abcdefghijklmnopqrstuvwxyz0123456789'":
return lowered
return " "
+4 -12
View File
@@ -9,8 +9,6 @@ from typing import TYPE_CHECKING
from python.ebook_search.llm_interface import request_rerank
if TYPE_CHECKING:
import httpx
from python.ebook_search.config import RerankConfig
from python.ebook_search.search import SearchResult
@@ -25,12 +23,7 @@ class RerankResult:
score: float
async def rerank_chunks(
client: httpx.AsyncClient,
query: str,
candidates: list[SearchResult],
config: RerankConfig,
) -> list[SearchResult]:
def rerank_chunks(query: str, candidates: list[SearchResult], config: RerankConfig) -> list[SearchResult]:
"""Rerank candidates with a vLLM rerank endpoint."""
if not candidates:
return []
@@ -41,7 +34,7 @@ async def rerank_chunks(
config.model,
len(candidates),
)
scores = await score_candidates(client, query, candidates, config)
scores = score_candidates(query, candidates, config)
results = sorted(
(
replace(
@@ -63,14 +56,13 @@ async def rerank_chunks(
return results
async def score_candidates(
client: httpx.AsyncClient,
def score_candidates(
query: str,
candidates: list[SearchResult],
config: RerankConfig,
) -> dict[int, RerankResult]:
"""Score candidate chunks with the configured rerank API."""
body = await request_rerank(client, query, [candidate.text for candidate in candidates], config)
body = request_rerank(query, [candidate.text for candidate in candidates], config)
if body is None:
return zero_rerank_scores(candidates)
+44 -158
View File
@@ -2,17 +2,16 @@
from __future__ import annotations
import asyncio
import logging
import re
from collections import defaultdict
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.exc import SQLAlchemyError
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import Session
from python.ebook_search.bm25_corpus import (
BM25CorpusUnavailableError,
@@ -20,12 +19,8 @@ from python.ebook_search.bm25_corpus import (
score_bm25_corpus,
)
from python.ebook_search.embeddings import MODEL_DIMENSIONS, embed_query, get_embedding_table
from python.ebook_search.protected_phrases.matching import (
detect_protected_phrases_for_query,
phrase_hits_for_chunks,
)
from python.ebook_search.rerank import rerank_chunks
from python.ebook_search.timing import RuntimeStep, async_timed_result, timed_result
from python.ebook_search.timing import RuntimeStep, timed_result
from python.orm.richie import (
EbookChapter,
EbookChunk,
@@ -34,13 +29,11 @@ from python.orm.richie import (
)
if TYPE_CHECKING:
from collections.abc import Mapping, Sequence
from collections.abc import Mapping
import httpx
from sqlalchemy.ext.asyncio import AsyncEngine
from sqlalchemy.engine import Engine
from python.ebook_search.config import EbookSearchConfig
from python.ebook_search.protected_phrases.models import HydratedPhraseMatch
logger = logging.getLogger(__name__)
@@ -52,14 +45,11 @@ class SearchResult:
chunk_id: int
text: str
source_title: str
source_id: int | None = None
score: float = 0.0
vector_score: float | None = None
bm25_score: float | None = None
fused_score: float | None = None
rerank_score: float | None = None
phrase_hit_count: int = 0
matched_phrases: tuple[str, ...] = ()
source_author: str | None = None
chapter_title: str | None = None
page_label: str | None = None
@@ -74,7 +64,6 @@ class SearchResponse:
results: list[SearchResult]
rank_label: str
timings: tuple[RuntimeStep, ...] = ()
phrase_matches: tuple[HydratedPhraseMatch, ...] = ()
@property
def total_runtime_ms(self) -> float:
@@ -91,38 +80,26 @@ class RetrievalResponse:
timings: tuple[RuntimeStep, ...]
async def search_ebooks(
engine: AsyncEngine,
client: httpx.AsyncClient,
def search_ebooks(
engine: Engine,
query: str,
config: EbookSearchConfig,
*,
rerank: bool = False,
phrase_matching: bool | None = None,
) -> 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")
phrase_matching_enabled = config.phrase_matching_enabled if phrase_matching is None else phrase_matching
logger.info(
"ebook_search_start query_length=%s rerank=%s phrase_matching=%s",
len(query),
rerank,
phrase_matching_enabled,
)
logger.info("ebook_search_start query_length=%s rerank=%s", len(query), rerank)
timings: list[RuntimeStep] = []
if phrase_matching_enabled:
phrase_matches, timing = await async_timed_result(
"Protected phrase detection", query_phrase_matches(engine, query, config)
)
else:
phrase_matches, timing = timed_result("Protected phrase detection skipped", skip_phrase_matches)
timings.append(timing)
retrieval, timing = await async_timed_result(
retrieval, timing = timed_result(
"Hybrid retrieval",
parallel_retrieval(engine, client, query, config),
parallel_retrieval,
engine,
query,
config,
)
timings.extend(retrieval.timings)
timings.append(timing)
@@ -134,28 +111,18 @@ async def search_ebooks(
rank_constant=config.rrf_rank_constant,
)
timings.append(timing)
if phrase_matching_enabled:
fused, timing = await async_timed_result(
"Phrase mention boost",
apply_phrase_mention_boosts(engine, fused, phrase_matches, config.phrase_hit_boost),
)
else:
fused, timing = timed_result("Phrase mention boost skipped", skip_phrase_mention_boosts, fused)
timings.append(timing)
if config.rerank.enabled and rerank:
response, timing = await async_timed_result("Rerank", apply_rerank(client, query, fused, config))
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), phrase_matches=tuple(phrase_matches))
response = replace(response, timings=tuple(timings))
logger.info(
"ebook_search_complete vector_candidates=%s lexical_candidates=%s "
"fused_candidates=%s phrase_matching=%s phrase_matches=%s returned=%s rank_label=%s runtime_ms=%.1f",
"fused_candidates=%s returned=%s rank_label=%s runtime_ms=%.1f",
len(retrieval.vector_results),
len(retrieval.lexical_results),
len(fused),
phrase_matching_enabled,
len(phrase_matches),
len(response.results),
response.rank_label,
response.total_runtime_ms,
@@ -163,100 +130,30 @@ async def search_ebooks(
return response
def skip_phrase_matches() -> list[HydratedPhraseMatch]:
"""Return no protected phrase matches when phrase matching is disabled."""
logger.info("ebook_protected_phrase_detection_skipped")
return []
async def query_phrase_matches(
engine: AsyncEngine,
query: str,
config: EbookSearchConfig,
) -> list[HydratedPhraseMatch]:
"""Detect protected phrases in a query without making search fail when phrase tables are unavailable."""
try:
async with AsyncSession(engine) as session:
return await detect_protected_phrases_for_query(session, query, config)
except SQLAlchemyError as error:
logger.warning("ebook_protected_phrase_detection_unavailable error=%s", error)
return []
def skip_phrase_mention_boosts(candidates: list[SearchResult]) -> list[SearchResult]:
"""Return candidates unchanged when phrase matching is disabled."""
logger.info("ebook_phrase_boost_skipped candidates=%s", len(candidates))
return candidates
async def apply_phrase_mention_boosts(
engine: AsyncEngine,
candidates: list[SearchResult],
phrase_matches: Sequence[HydratedPhraseMatch],
phrase_hit_boost: float,
) -> list[SearchResult]:
"""Boost retrieved chunks that have indexed mentions for detected protected phrases."""
phrase_ids = sorted({match.phrase_id for match in phrase_matches})
if not candidates or not phrase_ids or phrase_hit_boost <= 0:
return candidates
chunk_ids = [candidate.chunk_id for candidate in candidates]
try:
async with AsyncSession(engine) as session:
phrase_hits = await phrase_hits_for_chunks(session, chunk_ids=chunk_ids, phrase_ids=phrase_ids)
except SQLAlchemyError as error:
logger.warning("ebook_phrase_boost_unavailable error=%s", error)
return candidates
if not phrase_hits:
return candidates
hit_counts = {
chunk_id: sum(hit.mention_count for hit in chunk_hits) for chunk_id, chunk_hits in phrase_hits.items()
}
boosted = [
replace(
candidate,
score=candidate.score + (hit_counts.get(candidate.chunk_id, 0) * phrase_hit_boost),
fused_score=boosted_fused_score(candidate, hit_counts.get(candidate.chunk_id, 0), phrase_hit_boost),
phrase_hit_count=hit_counts.get(candidate.chunk_id, 0),
matched_phrases=tuple(hit.phrase_text for hit in phrase_hits.get(candidate.chunk_id, ())),
rank_source=phrase_rank_source(candidate.rank_source, hit_counts.get(candidate.chunk_id, 0)),
)
for candidate in candidates
]
return sorted(boosted, key=lambda candidate: candidate.score, reverse=True)
def boosted_fused_score(candidate: SearchResult, phrase_hit_count: int, phrase_hit_boost: float) -> float | None:
"""Return a fused score adjusted by phrase hits when a fused score exists."""
if candidate.fused_score is None:
return None
return candidate.fused_score + (phrase_hit_count * phrase_hit_boost)
def phrase_rank_source(rank_source: str, phrase_hit_count: int) -> str:
"""Append phrase evidence to a rank-source label when a chunk was boosted."""
if phrase_hit_count <= 0 or "phrases" in rank_source:
return rank_source
return f"{rank_source} + phrases"
async def parallel_retrieval(
engine: AsyncEngine,
client: httpx.AsyncClient,
def parallel_retrieval(
engine: Engine,
query: str,
config: EbookSearchConfig,
) -> RetrievalResponse:
"""Run vector and BM25 candidate retrieval concurrently with separate database sessions.
BM25 scoring is pure CPU work over the cached corpus, so it runs in a worker thread
instead of on the event loop.
"""
(vector_results, vector_timing), (lexical_results, lexical_timing) = await asyncio.gather(
async_timed_result("Embedding + vector search", vector_candidates(engine, client, query, config)),
async_timed_result("BM25 search", asyncio.to_thread(bm25_candidates, query, config)),
)
"""Run vector and BM25 candidate retrieval concurrently with separate database sessions."""
with ThreadPoolExecutor(max_workers=2, thread_name_prefix="ebook-search") as executor:
vector_future = executor.submit(
timed_result,
"Embedding + vector search",
vector_candidates,
engine,
query,
config,
)
bm25_future = executor.submit(
timed_result,
"BM25 search",
bm25_candidates,
query,
config,
)
vector_results, vector_timing = vector_future.result()
lexical_results, lexical_timing = bm25_future.result()
logger.info(
"ebook_parallel_retrieval_complete vector_candidates=%s lexical_candidates=%s",
@@ -283,14 +180,13 @@ def skip_rerank(
return SearchResponse(query=query, results=candidates[: config.top_k], rank_label="Hybrid")
async def apply_rerank(
client: httpx.AsyncClient,
def apply_rerank(
query: str,
candidates: list[SearchResult],
config: EbookSearchConfig,
) -> SearchResponse:
"""Rerank already-fused hybrid candidates."""
reranked = await rerank_chunks(client, query, candidates[: config.rerank.candidates], config.rerank)
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),
@@ -303,17 +199,10 @@ async def apply_rerank(
)
async def vector_candidates(
engine: AsyncEngine,
client: httpx.AsyncClient,
query: str,
config: EbookSearchConfig,
) -> list[SearchResult]:
def vector_candidates(engine: Engine, query: str, config: EbookSearchConfig) -> list[SearchResult]:
"""Return pgvector cosine candidates for a natural-language query."""
async with AsyncSession(engine) as session:
model = await session.scalar(
select(EbookEmbeddingModel).where(EbookEmbeddingModel.name == config.embedding_model)
)
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)
@@ -323,7 +212,7 @@ async def vector_candidates(
msg = f"Model row dimension {model.dimension} does not match configured dimension {expected_dimension}"
raise ValueError(msg)
embedding = await embed_query(client, query, config)
embedding = embed_query(query, config)
limit = max(config.rerank.candidates, config.top_k) * config.vector_candidate_multiplier
embedding_table = get_embedding_table(model.dimension)
@@ -334,7 +223,6 @@ async def vector_candidates(
select(
EbookChunk.id.label("chunk_id"),
EbookChunk.text.label("text"),
EbookSource.id.label("source_id"),
EbookSource.title.label("source_title"),
EbookSource.author.label("source_author"),
EbookChapter.title.label("chapter_title"),
@@ -349,7 +237,7 @@ async def vector_candidates(
.order_by(distance)
.limit(limit)
)
rows = (await session.execute(statement)).mappings()
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",
@@ -429,11 +317,9 @@ def reciprocal_rank_fusion(
def search_result_from_row(row: Mapping[str, object]) -> SearchResult:
"""Convert a database row mapping into a search result."""
source_id = row.get("source_id")
return SearchResult(
chunk_id=int(row["chunk_id"]),
text=str(row["text"]),
source_id=int(source_id) if source_id is not None else None,
source_title=str(row["source_title"]),
source_author=optional_str(row["source_author"]),
chapter_title=optional_str(row["chapter_title"]),
+1 -8
View File
@@ -7,7 +7,7 @@ from time import perf_counter
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from collections.abc import Awaitable, Callable
from collections.abc import Callable
@dataclass(frozen=True)
@@ -34,10 +34,3 @@ def timed_result[T, **P](
start_seconds = perf_counter()
result = operation(*args, **kwargs)
return result, runtime_step_from_start(name, start_seconds)
async def async_timed_result[T](name: str, awaitable: Awaitable[T]) -> tuple[T, RuntimeStep]:
"""Await an operation and return its result plus elapsed runtime."""
start_seconds = perf_counter()
result = await awaitable
return result, runtime_step_from_start(name, start_seconds)
+2 -2
View File
@@ -1,6 +1,6 @@
"""Reusable FastAPI tools."""
from python.fastapi_tools.db import AsyncDbSession, DbSession, get_async_db, get_db
from python.fastapi_tools.db import DbSession, get_db
from python.fastapi_tools.zstd_middleware import ZstdMiddleware
__all__ = ["AsyncDbSession", "DbSession", "ZstdMiddleware", "get_async_db", "get_db"]
__all__ = ["DbSession", "ZstdMiddleware", "get_db"]
+1 -13
View File
@@ -5,11 +5,10 @@ from __future__ import annotations
from typing import TYPE_CHECKING, Annotated
from fastapi import Depends, Request
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import Session
if TYPE_CHECKING:
from collections.abc import AsyncIterator, Iterator
from collections.abc import Iterator
def get_db(request: Request) -> Iterator[Session]:
@@ -18,15 +17,4 @@ def get_db(request: Request) -> Iterator[Session]:
yield session
async def get_async_db(request: Request) -> AsyncIterator[AsyncSession]:
"""Get an async database session from app state.
expire_on_commit=False keeps ORM attributes readable after commit without
triggering implicit IO, which would raise under asyncio.
"""
async with AsyncSession(request.app.state.engine, expire_on_commit=False) as session:
yield session
DbSession = Annotated[Session, Depends(get_db)]
AsyncDbSession = Annotated[AsyncSession, Depends(get_async_db)]
+12 -63
View File
@@ -7,7 +7,6 @@ from typing import cast
from sqlalchemy import create_engine
from sqlalchemy.engine import URL, Engine
from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
NAMING_CONVENTION = {
"ix": "ix_%(table_name)s_%(column_0_name)s",
@@ -32,16 +31,23 @@ 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 build_postgres_url(name: str, *, vector_engine: bool = False) -> tuple[URL, dict[str, str]]:
"""Build the Postgres connection URL and connect_args 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): The name of the environment variable prefix.
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 to include the vector types and operators.
This updates the search path the incldued the vecore types and operators.
Returns:
tuple[URL, dict[str, str]]: The SQLAlchemy URL and connect_args for create_engine.
Engine: The SQLAlchemy engine.
"""
database, host, port, username, password = get_connection_info(name)
@@ -59,66 +65,9 @@ def build_postgres_url(name: str, *, vector_engine: bool = False) -> tuple[URL,
if vector_engine:
connect_args["options"] = "-csearch_path=main,public"
return url, connect_args
def get_postgres_engine(
*,
name: str = "POSTGRES",
pool_pre_ping: bool = True,
vector_engine: bool = False,
pool_size: int = 8,
) -> 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.
pool_size (int, optional): Number of connections to keep in the pool. Defaults to 8.
Returns:
Engine: The SQLAlchemy engine.
"""
url, connect_args = build_postgres_url(name, vector_engine=vector_engine)
return create_engine(
url=url,
pool_pre_ping=pool_pre_ping,
pool_recycle=1800,
connect_args=connect_args,
pool_size=pool_size,
)
def get_async_postgres_engine(
*,
name: str = "POSTGRES",
pool_pre_ping: bool = True,
vector_engine: bool = False,
pool_size: int = 8,
) -> AsyncEngine:
"""Create an async 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.
pool_size (int, optional): Number of connections to keep in the pool. Defaults to 8.
Returns:
AsyncEngine: The async SQLAlchemy engine.
"""
url, connect_args = build_postgres_url(name, vector_engine=vector_engine)
return create_async_engine(
url=url,
pool_pre_ping=pool_pre_ping,
pool_recycle=1800,
connect_args=connect_args,
pool_size=pool_size,
)
-8
View File
@@ -12,16 +12,12 @@ from python.orm.richie.contact import (
RelationshipType,
)
from python.orm.richie.ebook import (
EbookCandidatePhrase,
EbookChapter,
EbookChunk,
EbookChunkEmbedding1024,
EbookChunkEmbedding2560,
EbookChunkEmbedding4096,
EbookChunkPhraseMention,
EbookEmbeddingModel,
EbookPhraseAlias,
EbookProtectedPhrase,
EbookSource,
)
@@ -32,16 +28,12 @@ __all__ = [
"Contact",
"ContactNeed",
"ContactRelationship",
"EbookCandidatePhrase",
"EbookChapter",
"EbookChunk",
"EbookChunkEmbedding1024",
"EbookChunkEmbedding2560",
"EbookChunkEmbedding4096",
"EbookChunkPhraseMention",
"EbookEmbeddingModel",
"EbookPhraseAlias",
"EbookProtectedPhrase",
"EbookSource",
"Need",
"RelationshipType",
+2 -108
View File
@@ -5,23 +5,11 @@ from __future__ import annotations
from datetime import datetime
from pgvector.sqlalchemy import Vector
from sqlalchemy import (
JSON,
BigInteger,
DateTime,
ForeignKey,
Index,
String,
Text,
UniqueConstraint,
)
from sqlalchemy.dialects.postgresql import JSONB
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
JSON_DOCUMENT = JSON().with_variant(JSONB, "postgresql")
class EbookSource(TableBase):
"""One indexed EPUB file."""
@@ -106,7 +94,7 @@ class EbookEmbeddingModel(TableBase):
name: Mapped[str] = mapped_column(String, unique=True)
dimension: Mapped[int]
is_default: Mapped[bool] = mapped_column(default=False)
is_default: Mapped[bool] = mapped_column(Boolean, default=False)
class EbookChunkEmbedding1024(TableBaseBig):
@@ -148,97 +136,3 @@ class EbookChunkEmbedding4096(TableBaseBig):
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))
class EbookCandidatePhrase(TableBase):
"""A high-recall phrase candidate extracted from one book."""
__tablename__ = "candidate_phrases"
__table_args__ = (
UniqueConstraint("book_id", "phrase_norm", name="uq_candidate_phrases_book_id_phrase_norm"),
Index("candidate_phrases_book_score_idx", "book_id", "candidate_score"),
Index("candidate_phrases_book_norm_idx", "book_id", "phrase_norm"),
)
book_id: Mapped[int] = mapped_column(ForeignKey("main.ebook_source.id", ondelete="CASCADE"))
series_id: Mapped[int | None]
phrase_text: Mapped[str] = mapped_column(Text)
phrase_norm: Mapped[str] = mapped_column(Text)
token_count: Mapped[int]
source_raw_ngram: Mapped[bool] = mapped_column(default=False)
source_yake: Mapped[bool] = mapped_column(default=False)
source_spacy_ner: Mapped[bool] = mapped_column(default=False)
source_spacy_noun_chunk: Mapped[bool] = mapped_column(default=False)
source_capitalized: Mapped[bool] = mapped_column(default=False)
source_metadata: Mapped[bool] = mapped_column(default=False)
spacy_label: Mapped[str | None]
raw_count: Mapped[int] = mapped_column(default=0)
chapter_count: Mapped[int] = mapped_column(default=0)
yake_score: Mapped[float | None]
candidate_score: Mapped[float] = mapped_column(default=0.0)
sample_contexts: Mapped[list[str] | None] = mapped_column(JSON_DOCUMENT)
llm_judged: Mapped[bool] = mapped_column(default=False)
llm_keep: Mapped[bool | None]
llm_confidence: Mapped[float | None]
llm_category: Mapped[str | None]
llm_reason: Mapped[str | None] = mapped_column(Text)
class EbookProtectedPhrase(TableBase):
"""A phrase accepted by the LLM judge for protected query matching."""
__tablename__ = "protected_phrases"
__table_args__ = (
UniqueConstraint("book_id", "phrase_norm", name="uq_protected_phrases_book_id_phrase_norm"),
Index("protected_phrases_norm_idx", "phrase_norm"),
Index("protected_phrases_book_norm_idx", "book_id", "phrase_norm"),
Index("protected_phrases_series_norm_idx", "series_id", "phrase_norm"),
)
book_id: Mapped[int | None] = mapped_column(ForeignKey("main.ebook_source.id", ondelete="CASCADE"))
series_id: Mapped[int | None]
phrase_text: Mapped[str] = mapped_column(Text)
phrase_norm: Mapped[str] = mapped_column(Text)
canonical_id: Mapped[str]
phrase_type: Mapped[str | None]
token_count: Mapped[int]
confidence: Mapped[float]
importance: Mapped[float] = mapped_column(default=0.5)
allow_nested: Mapped[bool] = mapped_column(default=False)
suppress_children: Mapped[bool] = mapped_column(default=True)
source_candidate_id: Mapped[int | None] = mapped_column(
ForeignKey("main.candidate_phrases.id", ondelete="SET NULL")
)
class EbookPhraseAlias(TableBase):
"""A normalized alias that maps to a protected phrase."""
__tablename__ = "phrase_aliases"
__table_args__ = (
UniqueConstraint("phrase_id", "alias_norm", name="uq_phrase_aliases_phrase_id_alias_norm"),
Index("phrase_aliases_norm_idx", "alias_norm"),
)
phrase_id: Mapped[int] = mapped_column(ForeignKey("main.protected_phrases.id", ondelete="CASCADE"))
alias_text: Mapped[str] = mapped_column(Text)
alias_norm: Mapped[str] = mapped_column(Text)
confidence: Mapped[float] = mapped_column(default=1.0)
class EbookChunkPhraseMention(TableBase):
"""A precomputed occurrence of a protected phrase inside one chunk."""
__tablename__ = "chunk_phrase_mentions"
__table_args__ = (
UniqueConstraint("chunk_id", "phrase_id", "start_char", name="uq_chunk_phrase_mentions_chunk_phrase_start"),
Index("chunk_phrase_mentions_phrase_idx", "phrase_id"),
Index("chunk_phrase_mentions_chunk_idx", "chunk_id"),
)
chunk_id: Mapped[int] = mapped_column(ForeignKey("main.ebook_chunk.id", ondelete="CASCADE"))
phrase_id: Mapped[int] = mapped_column(ForeignKey("main.protected_phrases.id", ondelete="CASCADE"))
book_id: Mapped[int | None] = mapped_column(ForeignKey("main.ebook_source.id", ondelete="CASCADE"))
series_id: Mapped[int | None]
start_char: Mapped[int]
end_char: Mapped[int | None]
+14
View File
@@ -3,6 +3,20 @@
...
}:
{
# For running the uv-managed ebook-search venv outside the container:
# PyPI manylinux wheels (numpy via bm25s) expect libstdc++.so.6 on the
# loader path, which NixOS does not provide globally.
ebook-search = pkgs.mkShell {
nativeBuildInputs = with pkgs; [
my_python
uv
];
LD_LIBRARY_PATH = pkgs.lib.makeLibraryPath [
pkgs.stdenv.cc.cc.lib
pkgs.zlib
];
};
default = pkgs.mkShell {
NIX_CONFIG = "extra-experimental-features = nix-command flakes ca-derivations";
nativeBuildInputs = with pkgs; [
+4 -1
View File
@@ -1,6 +1,9 @@
{ ... }:
{
imports = [ ./nix_builder.nix ];
imports = [
./nix_builder.nix
./uv_cache_prune.nix
];
services.nix_builder.containers = {
nix-builder-00.enable = true;
+20 -9
View File
@@ -62,6 +62,11 @@ in
mountPoint = "/run/secrets/gitea-runners";
isReadOnly = true;
};
uv-cache = {
hostPath = vars.uv_cache;
mountPoint = "/var/cache/uv";
isReadOnly = false;
};
};
config =
{
@@ -153,14 +158,20 @@ in
}
) cfg.containers;
systemd.services = builtins.listToAttrs (
map (name: {
name = "container@${name}";
value = {
requires = [ "gitea.service" ];
after = [ "gitea.service" ];
};
}) (builtins.attrNames (filterAttrs (_: c: c.enable) cfg.containers))
);
systemd = {
services = builtins.listToAttrs (
map (name: {
name = "container@${name}";
value = {
requires = [ "gitea.service" ];
after = [ "gitea.service" ];
};
}) (builtins.attrNames (filterAttrs (_: c: c.enable) cfg.containers))
);
tmpfiles.rules = [
"d ${vars.uv_cache} 0755 ${runnerUsername} ${runnerUsername} - -"
];
};
};
}
+28
View File
@@ -0,0 +1,28 @@
{ pkgs, ... }:
let
vars = import ../vars.nix;
runnerUsername = "gitea-runner";
in
{
systemd = {
services.uv-cache-prune = {
description = "Prune the shared gitea runner uv cache";
environment.UV_CACHE_DIR = vars.uv_cache;
serviceConfig = {
Type = "oneshot";
User = runnerUsername;
Group = runnerUsername;
ExecStart = "${pkgs.uv}/bin/uv cache prune";
};
};
timers.uv-cache-prune = {
description = "Monthly prune of the shared gitea runner uv cache";
wantedBy = [ "timers.target" ];
timerConfig = {
OnCalendar = "*-*-01 04:00:00";
Persistent = true;
};
};
};
}
+1
View File
@@ -30,6 +30,7 @@ sudo zfs create media/secure/share -o mountpoint=/zfs/media/share -o exec=off
# scratch datasets
sudo zfs create scratch/kafka -o mountpoint=/zfs/scratch/kafka -o recordsize=1M
sudo zfs create scratch/transmission -o mountpoint=/zfs/scratch/transmission -o recordsize=16k -o sync=disabled -o redundant_metadata=none
sudo zfs create scratch/uv_cache -o mountpoint=/zfs/scratch/uv_cache
# storage datasets
sudo zfs create storage/ollama -o recordsize=1M -o compression=zstd-19 -o sync=disabled
-43
View File
@@ -1,43 +0,0 @@
{
pkgs,
inputs,
...
}:
{
networking.firewall.allowedTCPPorts = [
8069
];
systemd.services.contact-api = {
description = "Contact Database API";
after = [
"postgresql.service"
"network.target"
];
requires = [ "postgresql.service" ];
wantedBy = [ "multi-user.target" ];
environment = {
PYTHONPATH = "${inputs.self}";
POSTGRES_DB = "richie";
POSTGRES_HOST = "/run/postgresql";
POSTGRES_USER = "richie";
POSTGRES_PORT = "5432";
};
serviceConfig = {
Type = "simple";
ExecStart = "${pkgs.my_python}/bin/python -m python.api.main --host 192.168.90.40 --port 8069";
Restart = "on-failure";
RestartSec = "5s";
StandardOutput = "journal";
StandardError = "journal";
NoNewPrivileges = true;
ProtectSystem = "strict";
ProtectHome = "read-only";
PrivateTmp = true;
ReadOnlyPaths = [
"${inputs.self}"
];
};
};
}
+6
View File
@@ -121,3 +121,9 @@ monthly = 2
hourly = 0
daily = 0
monthly = 0
["scratch/uv_cache"]
15_min = 2
hourly = 0
daily = 0
monthly = 0
+1
View File
@@ -17,5 +17,6 @@ in
transmission = "${zfs_storage}/transmission";
ollama = "${zfs_storage}/ollama";
transmission_scratch = "${zfs_scratch}/transmission";
uv_cache = "${zfs_scratch}/uv_cache";
kafka = "${zfs_scratch}/kafka";
}
+30 -35
View File
@@ -10,8 +10,8 @@ from types import ModuleType
from typing import TYPE_CHECKING
import pytest
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, create_async_engine
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 (
@@ -77,17 +77,10 @@ def test_reciprocal_rank_fusion_combines_vector_and_bm25_rankings() -> None:
assert fused[0].fused_score == fused[0].score
async def build_async_engine() -> AsyncEngine:
"""Create an in-memory async engine with the Richie schema."""
engine = create_async_engine("sqlite+aiosqlite:///:memory:")
async with engine.begin() as connection:
await connection.run_sync(RichieBase.metadata.create_all)
return engine
async def test_find_existing_source_matches_path_or_hash() -> None:
engine = await build_async_engine()
async with AsyncSession(engine, expire_on_commit=False) as session:
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,
@@ -100,15 +93,16 @@ async def test_find_existing_source_matches_path_or_hash() -> None:
file_size=10,
)
session.add(source)
await session.commit()
session.commit()
assert await find_existing_source(session, Path("/old/book.epub"), "b" * 64) == source
assert await find_existing_source(session, Path("/new/book.epub"), "a" * 64) == source
assert find_existing_source(session, Path("/old/book.epub"), "b" * 64) == source
assert find_existing_source(session, Path("/new/book.epub"), "a" * 64) == source
async def test_bm25_corpus_uses_existing_search_text_without_duplicate_metadata() -> None:
engine = await build_async_engine()
async with AsyncSession(engine, expire_on_commit=False) as session:
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",
@@ -121,10 +115,10 @@ async def test_bm25_corpus_uses_existing_search_text_without_duplicate_metadata(
file_size=10,
)
session.add(source)
await session.flush()
session.flush()
chapter = EbookChapter(source_id=source.id, spine_index=0, title="Chapter", href=None)
session.add(chapter)
await session.flush()
session.flush()
session.add(
EbookChunk(
id=1,
@@ -139,9 +133,9 @@ async def test_bm25_corpus_uses_existing_search_text_without_duplicate_metadata(
search_text="Book Author Chapter content",
)
)
await session.commit()
session.commit()
records, texts = await fetch_bm25_corpus_records(session)
records, texts = fetch_bm25_corpus_records(session)
assert texts == ["Book Author Chapter content"]
assert records[0]["chunk_id"] == 1
@@ -376,7 +370,7 @@ def test_load_bm25_corpus_raises_when_index_is_missing(mocker: MockerFixture, tm
load_bm25_corpus.cache_clear()
async def test_ensure_bm25_corpus_refreshes_missing_index(mocker: MockerFixture) -> None:
def test_ensure_bm25_corpus_refreshes_missing_index(mocker: MockerFixture) -> None:
refreshed: list[object] = []
db_updated_at = datetime.now(tz=UTC)
@@ -391,12 +385,12 @@ async def test_ensure_bm25_corpus_refreshes_missing_index(mocker: MockerFixture)
config = EbookSearchConfig(rerank=RerankConfig(enabled=False))
session = object()
await ensure_bm25_corpus(session, config)
ensure_bm25_corpus(session, config)
assert refreshed == [(session, config, db_updated_at)]
async def test_ensure_bm25_corpus_refreshes_stale_index(mocker: MockerFixture) -> None:
def test_ensure_bm25_corpus_refreshes_stale_index(mocker: MockerFixture) -> None:
refreshed: list[object] = []
created_at = datetime(2026, 1, 1, tzinfo=UTC)
db_updated_at = datetime(2026, 1, 2, tzinfo=UTC)
@@ -413,7 +407,7 @@ async def test_ensure_bm25_corpus_refreshes_stale_index(mocker: MockerFixture) -
config = EbookSearchConfig(rerank=RerankConfig(enabled=False))
session = object()
await ensure_bm25_corpus(session, config)
ensure_bm25_corpus(session, config)
assert refreshed == [(session, config, db_updated_at)]
@@ -426,13 +420,14 @@ def test_supported_embedding_models_match_service_names() -> None:
}
async def test_ensure_embedding_models_registers_service_names() -> None:
engine = await build_async_engine()
async with AsyncSession(engine, expire_on_commit=False) as session:
await ensure_embedding_models(session)
await session.commit()
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(await session.scalars(select(EbookEmbeddingModel).order_by(EbookEmbeddingModel.name)))
models = list(session.scalars(select(EbookEmbeddingModel).order_by(EbookEmbeddingModel.name)))
assert [(model.name, model.dimension) for model in models] == [
("qwen3-embedding-0.6b", 1024),
@@ -501,10 +496,10 @@ def test_chat_api_key_falls_back_to_ollama_api_key(mocker: MockerFixture) -> Non
assert config.vllm_api_key == "ollama-key"
async def test_answer_query_does_not_call_model_when_disabled(mocker: MockerFixture) -> None:
def test_answer_query_does_not_call_model_when_disabled() -> None:
config = load_config().model_copy(update={"answer_enabled": False})
result = SearchResult(chunk_id=1, text="source text", source_title="Book")
answer = await answer_query(mocker.Mock(), "question", [result], config)
answer = answer_query("question", [result], config)
assert "Answer generation is disabled" in answer
+9 -12
View File
@@ -5,7 +5,7 @@ from __future__ import annotations
from typing import TYPE_CHECKING
from fastapi.testclient import TestClient
from sqlalchemy.ext.asyncio import create_async_engine
from sqlalchemy import create_engine
from python.ebook_search.api.main import create_app
from python.ebook_search.config import EbookSearchConfig, RerankConfig
@@ -66,8 +66,8 @@ def test_is_confident_against_threshold() -> None:
def patch_app_runtime(mocker: MockerFixture):
mocker.patch(
"python.ebook_search.api.main.get_async_postgres_engine",
side_effect=lambda **_kwargs: create_async_engine("sqlite+aiosqlite:///:memory:"),
"python.ebook_search.api.main.get_postgres_engine",
side_effect=lambda **_kwargs: create_engine("sqlite+pysqlite:///:memory:", future=True),
)
mocker.patch("python.ebook_search.api.main.ensure_bm25_corpus", side_effect=lambda _session, _config: None)
@@ -75,12 +75,11 @@ def patch_app_runtime(mocker: MockerFixture):
def test_low_confidence_skips_answer_generation(mocker: MockerFixture) -> None:
called = False
def fake_search_ebooks(_engine, _client, query, _config, *, rerank=False, phrase_matching=False):
def fake_search_ebooks(_engine, query, _config, *, rerank=False):
del rerank
del phrase_matching
return SearchResponse(query=query, rank_label="Hybrid", results=make_results(1, vector_score=0.05))
def fake_answer_query(_client, _query, _results, _config):
def fake_answer_query(_query, _results, _config):
nonlocal called
called = True
return "answer"
@@ -105,15 +104,14 @@ def test_low_confidence_skips_answer_generation(mocker: MockerFixture) -> None:
def test_invalid_citation_is_flagged(mocker: MockerFixture) -> None:
def fake_search_ebooks(_engine, _client, query, _config, *, rerank=False, phrase_matching=False):
def fake_search_ebooks(_engine, query, _config, *, rerank=False):
del rerank
del phrase_matching
return SearchResponse(query=query, rank_label="Hybrid", results=make_results(2, vector_score=0.9))
mocker.patch("python.ebook_search.api.routes.search.search_ebooks", side_effect=fake_search_ebooks)
mocker.patch(
"python.ebook_search.api.routes.search.answer_query",
side_effect=lambda _client, _query, _results, _config: "Per the text [9].",
side_effect=lambda _query, _results, _config: "Per the text [9].",
)
patch_app_runtime(mocker)
app = create_app()
@@ -128,15 +126,14 @@ def test_invalid_citation_is_flagged(mocker: MockerFixture) -> None:
def test_grounded_answer_has_no_warning_badge(mocker: MockerFixture) -> None:
def fake_search_ebooks(_engine, _client, query, _config, *, rerank=False, phrase_matching=False):
def fake_search_ebooks(_engine, query, _config, *, rerank=False):
del rerank
del phrase_matching
return SearchResponse(query=query, rank_label="Hybrid", results=make_results(2, vector_score=0.9))
mocker.patch("python.ebook_search.api.routes.search.search_ebooks", side_effect=fake_search_ebooks)
mocker.patch(
"python.ebook_search.api.routes.search.answer_query",
side_effect=lambda _client, _query, _results, _config: "Grounded in [1] and [2].",
side_effect=lambda _query, _results, _config: "Grounded in [1] and [2].",
)
patch_app_runtime(mocker)
app = create_app()
+5 -5
View File
@@ -5,7 +5,7 @@ from __future__ import annotations
from typing import TYPE_CHECKING
from fastapi.testclient import TestClient
from sqlalchemy.ext.asyncio import create_async_engine
from sqlalchemy import create_engine
from python.ebook_search.api.main import create_app
from python.ebook_search.config import EbookSearchConfig, RerankConfig
@@ -18,18 +18,18 @@ if TYPE_CHECKING:
def fake_get_postgres_engine(**_kwargs):
"""Return an in-memory engine for route tests."""
return create_async_engine("sqlite+aiosqlite:///:memory:")
return create_engine("sqlite+pysqlite:///:memory:", future=True)
def patch_app_runtime(mocker: MockerFixture):
mocker.patch("python.ebook_search.api.main.get_async_postgres_engine", side_effect=fake_get_postgres_engine)
mocker.patch("python.ebook_search.api.main.get_postgres_engine", side_effect=fake_get_postgres_engine)
mocker.patch("python.ebook_search.api.main.ensure_bm25_corpus", side_effect=lambda _session, _config: None)
def patch_dependencies(mocker: MockerFixture, *, database=True, embedding=True, chat=True, bm25="ok"):
mocker.patch(f"{HEALTH_MODULE}.check_database", side_effect=lambda _session: database)
mocker.patch(f"{HEALTH_MODULE}.check_embedding_endpoint", side_effect=lambda _client, _config: embedding)
mocker.patch(f"{HEALTH_MODULE}.check_chat_endpoint", side_effect=lambda _client, _config: chat)
mocker.patch(f"{HEALTH_MODULE}.check_embedding_endpoint", side_effect=lambda _config: embedding)
mocker.patch(f"{HEALTH_MODULE}.check_chat_endpoint", side_effect=lambda _config: chat)
mocker.patch(f"{HEALTH_MODULE}.check_bm25_status", side_effect=lambda _config: bm25)
+9 -17
View File
@@ -16,14 +16,7 @@ if TYPE_CHECKING:
from pytest_mock import MockerFixture
def make_async_client(mocker: MockerFixture, fake_post) -> httpx.AsyncClient:
"""Build a mock async client whose post call is served by fake_post."""
client = mocker.MagicMock(spec=httpx.AsyncClient)
client.post = mocker.AsyncMock(side_effect=fake_post)
return client
async def test_answer_query_uses_httpx_chat_completions(mocker: MockerFixture) -> None:
def test_answer_query_uses_httpx_chat_completions(mocker: MockerFixture) -> None:
captured: dict[str, object] = {}
def fake_post(url: str, **kwargs: object) -> httpx.Response:
@@ -35,7 +28,7 @@ async def test_answer_query_uses_httpx_chat_completions(mocker: MockerFixture) -
request=httpx.Request("POST", url),
)
client = make_async_client(mocker, fake_post)
mocker.patch.object(httpx, "post", side_effect=fake_post)
config = EbookSearchConfig(
rerank=RerankConfig(enabled=False),
vllm_base_url="https://ollama.com/v1",
@@ -43,8 +36,7 @@ async def test_answer_query_uses_httpx_chat_completions(mocker: MockerFixture) -
chat_model="deepseek-v4-flash",
)
results = [SearchResult(chunk_id=1, text="source", source_title="Book")]
answer = await answer_query(client, "question", results, config)
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"
@@ -56,7 +48,7 @@ async def test_answer_query_uses_httpx_chat_completions(mocker: MockerFixture) -
assert payload["model"] == "deepseek-v4-flash"
async def test_embed_texts_uses_httpx_embeddings(mocker: MockerFixture) -> None:
def test_embed_texts_uses_httpx_embeddings(mocker: MockerFixture) -> None:
captured: dict[str, object] = {}
vector = [0.0] * 1024
@@ -69,14 +61,14 @@ async def test_embed_texts_uses_httpx_embeddings(mocker: MockerFixture) -> None:
request=httpx.Request("POST", url),
)
client = make_async_client(mocker, fake_post)
mocker.patch.object(httpx, "post", side_effect=fake_post)
config = EbookSearchConfig(
rerank=RerankConfig(enabled=False),
embedding_base_url="http://bob:8000/v1",
embedding_model="qwen3-embedding-0.6b",
)
embeddings = await embed_texts(client, ["hello"], config)
embeddings = embed_texts(["hello"], config)
assert embeddings == [vector]
assert captured["url"] == "http://bob:8000/v1/embeddings"
@@ -86,12 +78,12 @@ async def test_embed_texts_uses_httpx_embeddings(mocker: MockerFixture) -> None:
assert kwargs["json"] == {"model": "qwen3-embedding-0.6b", "input": ["hello"]}
async def test_embed_texts_rejects_bad_response_shape(mocker: MockerFixture) -> None:
def test_embed_texts_rejects_bad_response_shape(mocker: MockerFixture) -> None:
def fake_post(url: str, **_kwargs: object) -> httpx.Response:
return httpx.Response(200, json={"data": [{}]}, request=httpx.Request("POST", url))
client = make_async_client(mocker, fake_post)
mocker.patch.object(httpx, "post", side_effect=fake_post)
config = EbookSearchConfig(rerank=RerankConfig(enabled=False))
with pytest.raises(RuntimeError, match="Embedding request failed"):
await embed_texts(client, ["hello"], config)
embed_texts(["hello"], config)
@@ -1,957 +0,0 @@
"""Tests for protected phrase extraction and matching."""
from __future__ import annotations
from datetime import UTC, datetime
from typing import TYPE_CHECKING
import pytest
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, create_async_engine
from sqlalchemy.pool import StaticPool
from python.ebook_search.config import EbookSearchConfig
from python.ebook_search.protected_phrases.config import (
get_bad_ends,
get_most_common_words,
)
from python.ebook_search.protected_phrases.extraction import (
extract_capitalized_phrases,
is_junk_phrase,
score_candidate,
)
from python.ebook_search.protected_phrases.generate_ngrams import (
generate_candidate_phrases_for_books,
recalculate_candidate_phrases_for_book,
)
from python.ebook_search.protected_phrases.judge_ngrams import (
judge_candidate_phrases_for_books,
prepare_book_judgment,
)
from python.ebook_search.protected_phrases.matching import (
detect_protected_phrases_for_query,
index_chunk_phrase_mentions,
load_phrase_lookup,
phrase_hit_counts_for_chunks,
phrase_hits_for_chunks,
resolve_overlaps,
)
from python.ebook_search.protected_phrases.models import (
ChunkPhraseHit,
HydratedPhraseMatch,
LLMJudgment,
PhraseCandidate,
)
from python.ebook_search.protected_phrases.store import book_ids_pending_first_judgment, corpus_phrase_stats
from python.ebook_search.protected_phrases.text_normalization import normalize_text, tokenize
from python.orm.richie import (
EbookCandidatePhrase,
EbookChunk,
EbookChunkPhraseMention,
EbookPhraseAlias,
EbookProtectedPhrase,
EbookSource,
RichieBase,
)
if TYPE_CHECKING:
from collections.abc import AsyncGenerator
from pytest_mock import MockerFixture
@pytest.fixture
async def engine() -> AsyncGenerator[AsyncEngine]:
"""Create a shared in-memory async database engine for phrase tests."""
test_engine = create_async_engine(
"sqlite+aiosqlite:///:memory:",
connect_args={"check_same_thread": False},
poolclass=StaticPool,
)
async with test_engine.begin() as connection:
await connection.run_sync(RichieBase.metadata.create_all)
yield test_engine
await test_engine.dispose()
@pytest.fixture
async def session(engine: AsyncEngine) -> AsyncGenerator[AsyncSession]:
"""Provide a session on the shared in-memory database."""
async with AsyncSession(engine, expire_on_commit=False) as test_session:
yield test_session
@pytest.fixture
def config() -> EbookSearchConfig:
"""Provide default phrase-tuning settings for tests."""
return EbookSearchConfig()
def test_normalize_text_preserves_phrase_stopwords_and_word_order() -> None:
"""Normalization should not collapse protected phrases by removing stopwords."""
assert normalize_text("Haden\u2019s syndrome -- lock-in!") == "haden's syndrome lock in"
assert tokenize("House of the Dragon") == ["house", "of", "the", "dragon"]
def test_capitalized_phrase_keeps_bad_end_tokens_for_downranking(config: EbookSearchConfig) -> None:
"""Capitalized extraction should keep dangling endings so scoring can downrank them."""
candidates = extract_capitalized_phrases("Damion Montgomery and the left.", config)
assert "damion montgomery and the" in candidates
assert candidates["damion montgomery and the"].phrase_text == "Damion Montgomery and the"
assert candidates["damion montgomery and the"].phrase_norm.split()[-1] in get_bad_ends()
def test_is_junk_phrase_rejects_dialogue_verbs_contractions_and_majority_common() -> None:
"""Phrases the judge never keeps should be caught lexically before any LLM call."""
assert is_junk_phrase(["vann", "said"])
assert is_junk_phrase(["shook", "her", "head"])
assert is_junk_phrase(["i'm", "going"])
assert is_junk_phrase(["don't", "know"])
assert is_junk_phrase(["of", "the", "vault"])
def test_is_junk_phrase_keeps_possessives_and_half_common_world_terms() -> None:
"""Proper-noun possessives and phrasal world terms must survive the junk filter."""
assert not is_junk_phrase(["chapman's", "death"])
assert not is_junk_phrase(["boston", "bays'"])
assert not is_junk_phrase(["lock", "in"])
assert not is_junk_phrase(["data", "feed"])
assert not is_junk_phrase(["haden's", "syndrome"])
def test_score_candidate_rewards_multiple_non_raw_sources(config: EbookSearchConfig) -> None:
"""A second non-raw source should add its weight plus the multi-source bonus."""
single_source = PhraseCandidate(
phrase_text="lock in",
phrase_norm="lock in",
token_count=2,
source_capitalized=True,
raw_count=3,
chapter_count=2,
)
multi_source = PhraseCandidate(
phrase_text="lock in",
phrase_norm="lock in",
token_count=2,
source_capitalized=True,
source_yake=True,
raw_count=3,
chapter_count=2,
)
assert score_candidate(multi_source, config) == score_candidate(single_source, config) + 2.0 + 2.0
def test_score_candidate_caps_frequency_contribution(config: EbookSearchConfig) -> None:
"""Very frequent raw-only phrases should no longer out-score sourced entities."""
frequent_raw_only = PhraseCandidate(
phrase_text="mage king",
phrase_norm="mage king",
token_count=2,
source_raw_ngram=True,
raw_count=1000,
chapter_count=100,
)
assert score_candidate(frequent_raw_only, config) == 0.5 + 2.0 + 0.5
def test_score_candidate_weights_metadata_source(config: EbookSearchConfig) -> None:
"""Title and series phrases should get credit for the metadata source."""
metadata_only = PhraseCandidate(
phrase_text="lock in",
phrase_norm="lock in",
token_count=2,
source_metadata=True,
)
assert score_candidate(metadata_only, config) == 2.0 + 0.5
async def test_detect_protected_phrases_hydrates_alias_matches(
session: AsyncSession,
config: EbookSearchConfig,
) -> None:
"""Query detection should use RAM aliases and hydrate phrase metadata from the DB."""
source = await add_source(session)
phrase = await add_phrase(session, source.id, phrase_text="lock in", phrase_norm="lock in")
session.add(EbookPhraseAlias(phrase_id=phrase.id, alias_text="locked in", alias_norm="locked in"))
await session.commit()
matches = await detect_protected_phrases_for_query(session, "what is locked-in", config)
assert [(match.phrase_text, match.canonical_id, match.phrase_type) for match in matches] == [
("lock in", "condition:lock_in", "fictional_condition")
]
def test_resolve_overlaps_keeps_independent_nested_phrases() -> None:
"""Overlap resolution should keep useful nested concepts when metadata permits it."""
child = hydrated_match(
phrase_id=1,
phrase_text="mage king",
canonical_id="title:mage_king",
start_token=3,
end_token=5,
allow_nested=True,
)
parent = hydrated_match(
phrase_id=2,
phrase_text="mage king of mars",
canonical_id="entity:mage_king_of_mars",
start_token=3,
end_token=7,
suppress_children=False,
)
assert [match.phrase_text for match in resolve_overlaps([child, parent])] == [
"mage king of mars",
"mage king",
]
def test_resolve_overlaps_suppresses_weaker_same_canonical_match() -> None:
"""Same-canonical overlaps should keep the stronger evidence."""
weak = hydrated_match(
phrase_id=1,
phrase_text="lock",
canonical_id="condition:lock_in",
start_token=2,
end_token=3,
importance=0.2,
)
strong = hydrated_match(
phrase_id=2,
phrase_text="lock in",
canonical_id="condition:lock_in",
start_token=2,
end_token=4,
importance=0.9,
)
assert [match.phrase_text for match in resolve_overlaps([weak, strong])] == ["lock in"]
async def test_index_chunk_phrase_mentions_uses_normalized_window_lookup(
session: AsyncSession,
config: EbookSearchConfig,
) -> None:
"""Chunk indexing should store mention rows without scanning all phrases at query time."""
source = await add_source(session)
phrase = await add_phrase(session, source.id, phrase_text="lock in", phrase_norm="lock in")
chunk = EbookChunk(
id=1,
source_id=source.id,
chapter_id=None,
chunk_index=0,
text="Victims experienced lock-in during the crisis.",
token_start=0,
token_count=7,
page_label=None,
content_sha256="b" * 64,
search_text="Victims experienced lock-in during the crisis.",
)
session.add(chunk)
await session.commit()
lookup = await load_phrase_lookup(session, config, book_id=source.id)
count = await index_chunk_phrase_mentions(session, chunk, lookup=lookup)
await session.commit()
mention = await session.scalar(select(EbookChunkPhraseMention))
assert count == 1
assert mention is not None
assert mention.chunk_id == chunk.id
assert mention.phrase_id == phrase.id
assert chunk.text[mention.start_char : mention.end_char] == "lock-in"
assert await phrase_hit_counts_for_chunks(session, chunk_ids=[chunk.id], phrase_ids=[phrase.id]) == {chunk.id: 1}
assert await phrase_hits_for_chunks(session, chunk_ids=[chunk.id], phrase_ids=[phrase.id]) == {
chunk.id: (ChunkPhraseHit(phrase_id=phrase.id, phrase_text="lock in", mention_count=1),)
}
async def test_generate_candidate_phrases_for_books_stores_unjudged_candidates(
session: AsyncSession,
config: EbookSearchConfig,
) -> None:
"""Candidate generation should populate the per-book phrase list without calling the LLM."""
source = await add_source(session)
session.add(
EbookChunk(
id=1,
source_id=source.id,
chapter_id=None,
chunk_index=0,
text="lock in lock in lock in",
token_start=0,
token_count=6,
page_label=None,
content_sha256="d" * 64,
search_text="lock in lock in lock in",
)
)
await session.commit()
build_config = config.model_copy(
update={
"protected_phrase_max_candidates_per_book": 1,
"phrase_min_tokens": 2,
"phrase_max_tokens": 2,
}
)
result = await generate_candidate_phrases_for_books(session, build_config)
await session.commit()
candidate = await session.scalar(select(EbookCandidatePhrase))
assert result.books_seen == 1
assert result.books_built == 1
assert result.candidate_phrases == 1
assert candidate is not None
assert candidate.phrase_norm == "lock in"
assert candidate.llm_judged is False
assert await session.scalar(select(EbookProtectedPhrase)) is None
async def test_generate_candidate_phrases_for_books_filters_one_token_and_one_use_candidates(
session: AsyncSession,
config: EbookSearchConfig,
) -> None:
"""Candidate generation should not persist one-token or single-use phrases."""
source = await add_source(session)
session.add(
EbookChunk(
id=1,
source_id=source.id,
chapter_id=None,
chunk_index=0,
text=(
"Damion walked away. Damion woke up. rare phrase appeared once. "
"and then and then. lock in lock in lock in."
),
token_start=0,
token_count=16,
page_label=None,
content_sha256="i" * 64,
search_text=(
"Damion walked away. Damion woke up. rare phrase appeared once. "
"and then and then. lock in lock in lock in."
),
)
)
session.add_all(
[
EbookCandidatePhrase(
book_id=source.id,
series_id=None,
phrase_text="Damion",
phrase_norm="damion",
token_count=1,
source_capitalized=True,
raw_count=2,
chapter_count=1,
candidate_score=10.0,
llm_judged=False,
),
EbookCandidatePhrase(
book_id=source.id,
series_id=None,
phrase_text="rare phrase",
phrase_norm="rare phrase",
token_count=2,
source_raw_ngram=True,
raw_count=1,
chapter_count=1,
candidate_score=9.0,
llm_judged=False,
),
]
)
await session.commit()
build_config = config.model_copy(
update={
"protected_phrase_max_candidates_per_book": 50,
"phrase_min_tokens": 2,
"phrase_max_tokens": 2,
"phrase_raw_ngram_min_count": 2,
}
)
result = await generate_candidate_phrases_for_books(session, build_config)
await session.commit()
candidates = list(await session.scalars(select(EbookCandidatePhrase)))
phrase_norms = {candidate.phrase_norm for candidate in candidates}
common_words = get_most_common_words()
assert result.candidate_phrases == len(candidates)
assert "lock in" in phrase_norms
assert "damion" not in phrase_norms
assert "rare phrase" not in phrase_norms
assert "and then" not in phrase_norms
assert {"and", "then"}.issubset(common_words)
assert all(candidate.token_count >= 2 for candidate in candidates)
assert all(candidate.raw_count >= 2 for candidate in candidates)
assert all(not all(token in common_words for token in candidate.phrase_norm.split()) for candidate in candidates)
async def test_generate_candidate_phrases_for_books_commits_after_each_book(
session: AsyncSession,
config: EbookSearchConfig,
mocker: MockerFixture,
) -> None:
"""Candidate generation should persist each completed book independently."""
first = await add_source(session)
second = await add_source(session, file_path="/library/book-2.epub", file_sha256="z" * 64)
session.add_all(
[
EbookChunk(
id=1,
source_id=first.id,
chapter_id=None,
chunk_index=0,
text="lock in lock in lock in",
token_start=0,
token_count=6,
page_label=None,
content_sha256="f" * 64,
search_text="lock in lock in lock in",
),
EbookChunk(
id=2,
source_id=second.id,
chapter_id=None,
chunk_index=0,
text="mage king mage king mage king",
token_start=0,
token_count=6,
page_label=None,
content_sha256="g" * 64,
search_text="mage king mage king mage king",
),
]
)
await session.commit()
commit_spy = mocker.spy(session, "commit")
build_config = config.model_copy(
update={
"protected_phrase_max_candidates_per_book": 1,
"phrase_min_tokens": 2,
"phrase_max_tokens": 2,
}
)
result = await generate_candidate_phrases_for_books(session, build_config)
assert result.books_built == 2
assert result.candidate_phrases == 2
assert commit_spy.call_count == 2
async def test_judge_candidate_phrases_for_books_promotes_stored_candidates(
engine: AsyncEngine,
session: AsyncSession,
config: EbookSearchConfig,
mocker: MockerFixture,
) -> None:
"""Judging should work from existing candidate rows and retain the candidate history."""
source = await add_source(session)
session.add(
EbookChunk(
id=1,
source_id=source.id,
chapter_id=None,
chunk_index=0,
text="lock in lock in lock in",
token_start=0,
token_count=6,
page_label=None,
content_sha256="e" * 64,
search_text="lock in lock in lock in",
)
)
await session.commit()
build_config = config.model_copy(
update={
"protected_phrase_max_candidates_per_book": 1,
"protected_phrase_llm_candidates_per_book": 1,
"phrase_min_tokens": 2,
"phrase_max_tokens": 2,
"phrase_judge_book_workers": 1,
"phrase_judge_phrase_workers": 1,
}
)
await generate_candidate_phrases_for_books(session, build_config)
await session.commit()
mocker.patch(
"python.ebook_search.protected_phrases.judge_ngrams.judge_candidate_async",
return_value=LLMJudgment(
keep=True,
canonical="lock in",
category="fictional_condition",
aliases=(),
confidence=0.95,
importance=0.9,
),
)
result = await judge_candidate_phrases_for_books(engine, build_config)
session.expire_all()
candidate = await session.scalar(select(EbookCandidatePhrase))
phrase = await session.scalar(select(EbookProtectedPhrase))
mention = await session.scalar(select(EbookChunkPhraseMention))
assert result.books_seen == 1
assert result.books_judged == 1
assert result.candidates_judged == 1
assert result.protected_phrases == 1
assert candidate is not None
assert candidate.llm_judged is True
assert candidate.llm_keep is True
assert phrase is not None
assert phrase.source_candidate_id == candidate.id
assert mention is not None
assert mention.phrase_id == phrase.id
async def test_judge_candidate_phrases_for_books_logs_and_continues_after_book_failure(
engine: AsyncEngine,
session: AsyncSession,
config: EbookSearchConfig,
mocker: MockerFixture,
) -> None:
"""A failing book should be rolled back and logged while later books are still judged and committed."""
first = await add_source(session)
second = await add_source(session, file_path="/library/book-2.epub", file_sha256="z" * 64)
session.add_all(
[
EbookChunk(
id=1,
source_id=first.id,
chapter_id=None,
chunk_index=0,
text="lock in lock in lock in",
token_start=0,
token_count=6,
page_label=None,
content_sha256="f" * 64,
search_text="lock in lock in lock in",
),
EbookChunk(
id=2,
source_id=second.id,
chapter_id=None,
chunk_index=0,
text="mage king mage king mage king",
token_start=0,
token_count=6,
page_label=None,
content_sha256="g" * 64,
search_text="mage king mage king mage king",
),
]
)
await session.commit()
build_config = config.model_copy(
update={
"protected_phrase_max_candidates_per_book": 1,
"protected_phrase_llm_candidates_per_book": 1,
"phrase_min_tokens": 2,
"phrase_max_tokens": 2,
"phrase_judge_book_workers": 1,
"phrase_judge_phrase_workers": 1,
}
)
await generate_candidate_phrases_for_books(session, build_config)
await session.commit()
def judge_or_fail(_client: object, _config: EbookSearchConfig, candidate: PhraseCandidate) -> LLMJudgment:
if candidate.phrase_norm == "lock in":
message = "llm judge unavailable"
raise RuntimeError(message)
return LLMJudgment(
keep=True,
canonical="mage king",
category="title",
aliases=(),
confidence=0.95,
importance=0.9,
)
mocker.patch(
"python.ebook_search.protected_phrases.judge_ngrams.judge_candidate_async",
side_effect=judge_or_fail,
)
second_book_id = second.id
result = await judge_candidate_phrases_for_books(engine, build_config)
session.expire_all()
assert result.books_seen == 2
assert result.books_judged == 1
assert result.books_failed == 1
assert result.protected_phrases == 1
phrase = await session.scalar(select(EbookProtectedPhrase))
assert phrase is not None
assert phrase.book_id == second_book_id
async def test_judge_candidate_phrases_for_books_skips_unstorable_existing_candidates(
engine: AsyncEngine,
session: AsyncSession,
config: EbookSearchConfig,
mocker: MockerFixture,
) -> None:
"""Old one-token or single-use candidate rows should not be judged or promoted."""
source = await add_source(session)
session.add(
EbookChunk(
id=1,
source_id=source.id,
chapter_id=None,
chunk_index=0,
text="Damion Damion rare phrase",
token_start=0,
token_count=4,
page_label=None,
content_sha256="j" * 64,
search_text="Damion Damion rare phrase",
)
)
session.add_all(
[
EbookCandidatePhrase(
book_id=source.id,
series_id=None,
phrase_text="Damion",
phrase_norm="damion",
token_count=1,
source_capitalized=True,
raw_count=2,
chapter_count=1,
candidate_score=10.0,
llm_judged=False,
),
EbookCandidatePhrase(
book_id=source.id,
series_id=None,
phrase_text="rare phrase",
phrase_norm="rare phrase",
token_count=2,
source_raw_ngram=True,
raw_count=1,
chapter_count=1,
candidate_score=9.0,
llm_judged=False,
),
]
)
await session.commit()
judge_mock = mocker.patch("python.ebook_search.protected_phrases.judge_ngrams.judge_candidate_async")
build_config = config.model_copy(
update={
"phrase_judge_book_workers": 1,
"phrase_judge_phrase_workers": 1,
}
)
result = await judge_candidate_phrases_for_books(engine, build_config)
assert result.books_judged == 0
assert result.candidates_judged == 0
assert result.protected_phrases == 0
judge_mock.assert_not_called()
assert await session.scalar(select(EbookProtectedPhrase)) is None
async def test_prepare_book_judgment_skips_stored_junk_and_rescores_stale_rows(
engine: AsyncEngine,
session: AsyncSession,
config: EbookSearchConfig,
) -> None:
"""Judgment selection should drop junk rows and rank by fresh scores, not stored ones."""
source = await add_source(session)
session.add(
EbookChunk(
id=1,
source_id=source.id,
chapter_id=None,
chunk_index=0,
text="lock in lock in lock in",
token_start=0,
token_count=6,
page_label=None,
content_sha256="k" * 64,
search_text="lock in lock in lock in",
)
)
session.add_all(
[
EbookCandidatePhrase(
book_id=source.id,
series_id=None,
phrase_text="vann said",
phrase_norm="vann said",
token_count=2,
source_raw_ngram=True,
raw_count=50,
chapter_count=10,
candidate_score=10.0,
llm_judged=False,
),
EbookCandidatePhrase(
book_id=source.id,
series_id=None,
phrase_text="boring pair",
phrase_norm="boring pair",
token_count=2,
source_raw_ngram=True,
raw_count=50,
chapter_count=10,
candidate_score=9.0,
llm_judged=False,
),
EbookCandidatePhrase(
book_id=source.id,
series_id=None,
phrase_text="lock in",
phrase_norm="lock in",
token_count=2,
source_capitalized=True,
source_yake=True,
raw_count=3,
chapter_count=2,
candidate_score=1.0,
llm_judged=False,
),
]
)
await session.commit()
build_config = config.model_copy(
update={
"protected_phrase_llm_candidates_per_book": 2,
"phrase_min_tokens": 2,
"phrase_max_tokens": 2,
}
)
prepared = await prepare_book_judgment(engine, source.id, build_config)
assert prepared is not None
work_items, target_remaining = prepared
judged_norms = [candidate.phrase_norm for _, candidate in work_items]
assert judged_norms == ["lock in", "boring pair"]
assert target_remaining == build_config.phrase_target_protected_per_book
fresh_scores = [candidate.candidate_score for _, candidate in work_items]
assert fresh_scores == sorted(fresh_scores, reverse=True)
assert fresh_scores[0] > 1.0
async def test_recalculate_candidate_phrases_for_book_removes_old_phrase_data(
session: AsyncSession,
config: EbookSearchConfig,
) -> None:
"""Book-level recalculation should clear stale candidates, protected phrases, aliases, and mentions."""
source = await add_source(session)
session.add(
EbookChunk(
id=1,
source_id=source.id,
chapter_id=None,
chunk_index=0,
text="new phrase new phrase new phrase",
token_start=0,
token_count=6,
page_label=None,
content_sha256="h" * 64,
search_text="new phrase new phrase new phrase",
)
)
old_candidate = EbookCandidatePhrase(
book_id=source.id,
series_id=None,
phrase_text="old phrase",
phrase_norm="old phrase",
token_count=2,
source_raw_ngram=True,
raw_count=1,
chapter_count=1,
candidate_score=1.0,
llm_judged=False,
)
session.add(old_candidate)
await session.flush()
old_phrase = await add_phrase(session, source.id, phrase_text="old phrase", phrase_norm="old phrase")
old_phrase.source_candidate_id = old_candidate.id
session.add(EbookPhraseAlias(phrase_id=old_phrase.id, alias_text="old alias", alias_norm="old alias"))
session.add(
EbookChunkPhraseMention(
chunk_id=1,
phrase_id=old_phrase.id,
book_id=source.id,
series_id=None,
start_char=0,
end_char=10,
)
)
await session.commit()
build_config = config.model_copy(
update={
"protected_phrase_max_candidates_per_book": 1,
"phrase_min_tokens": 2,
"phrase_max_tokens": 2,
}
)
result = await recalculate_candidate_phrases_for_book(session, source, build_config)
candidates = list(await session.scalars(select(EbookCandidatePhrase)))
assert result.deleted_candidates == 1
assert result.deleted_protected_phrases == 1
assert result.deleted_aliases == 1
assert result.deleted_mentions == 1
assert result.candidate_phrases == 1
assert [candidate.phrase_norm for candidate in candidates] == ["new phrase"]
assert await session.scalar(select(EbookProtectedPhrase)) is None
assert await session.scalar(select(EbookPhraseAlias)) is None
assert await session.scalar(select(EbookChunkPhraseMention)) is None
async def test_corpus_phrase_stats_counts_phrases_and_book_coverage(session: AsyncSession) -> None:
"""Corpus stats should count phrases plus how many books are generated and fully judged."""
unjudged_book = await add_source(session)
judged_book = await add_source(session, file_path="/library/judged.epub", file_sha256="b" * 64)
mixed_book = await add_source(session, file_path="/library/mixed.epub", file_sha256="c" * 64)
await add_source(session, file_path="/library/empty.epub", file_sha256="d" * 64)
await add_candidate(session, unjudged_book.id, phrase_norm="rare phrase", llm_judged=False)
await add_candidate(session, judged_book.id, phrase_norm="lock in", llm_judged=True)
await add_candidate(session, mixed_book.id, phrase_norm="haden's syndrome", llm_judged=True)
await add_candidate(session, mixed_book.id, phrase_norm="boston bays", llm_judged=False)
await add_phrase(session, judged_book.id, phrase_text="Lock In", phrase_norm="lock in")
stats = await corpus_phrase_stats(session)
assert stats.total_books == 4
assert stats.books_with_candidates == 3
assert stats.books_fully_judged == 1
assert stats.candidate_phrases == 4
assert stats.judged_candidates == 2
assert stats.unjudged_candidates == 2
assert stats.protected_phrases == 1
async def test_book_ids_pending_first_judgment_returns_only_never_judged_books(session: AsyncSession) -> None:
"""Only books whose candidates are all unjudged should be pending a first judgment."""
unjudged_book = await add_source(session)
judged_book = await add_source(session, file_path="/library/judged.epub", file_sha256="b" * 64)
mixed_book = await add_source(session, file_path="/library/mixed.epub", file_sha256="c" * 64)
await add_candidate(session, unjudged_book.id, phrase_norm="rare phrase", llm_judged=False)
await add_candidate(session, judged_book.id, phrase_norm="lock in", llm_judged=True)
await add_candidate(session, mixed_book.id, phrase_norm="haden's syndrome", llm_judged=True)
await add_candidate(session, mixed_book.id, phrase_norm="boston bays", llm_judged=False)
assert await book_ids_pending_first_judgment(session) == [unjudged_book.id]
async def add_candidate(
session: AsyncSession,
book_id: int,
*,
phrase_norm: str,
llm_judged: bool,
) -> EbookCandidatePhrase:
"""Add a minimal candidate phrase row."""
candidate = EbookCandidatePhrase(
book_id=book_id,
series_id=None,
phrase_text=phrase_norm,
phrase_norm=phrase_norm,
token_count=len(phrase_norm.split()),
source_raw_ngram=True,
raw_count=3,
chapter_count=2,
candidate_score=5.0,
llm_judged=llm_judged,
)
session.add(candidate)
await session.flush()
return candidate
async def add_source(
session: AsyncSession,
*,
file_path: str = "/library/book.epub",
file_sha256: str = "a" * 64,
) -> EbookSource:
"""Add a minimal ebook source."""
source = EbookSource(
title="Book",
author="Author",
language=None,
publisher=None,
identifier=None,
file_path=file_path,
file_sha256=file_sha256,
file_mtime=datetime.now(tz=UTC),
file_size=10,
)
session.add(source)
await session.flush()
return source
async def add_phrase(
session: AsyncSession,
book_id: int,
*,
phrase_text: str,
phrase_norm: str,
) -> EbookProtectedPhrase:
"""Add a protected phrase row."""
phrase = EbookProtectedPhrase(
book_id=book_id,
series_id=None,
phrase_text=phrase_text,
phrase_norm=phrase_norm,
canonical_id="condition:lock_in",
phrase_type="fictional_condition",
token_count=len(phrase_norm.split()),
confidence=0.91,
importance=0.85,
allow_nested=False,
suppress_children=True,
source_candidate_id=None,
)
session.add(phrase)
await session.flush()
return phrase
def hydrated_match(
*,
phrase_id: int,
phrase_text: str,
canonical_id: str,
start_token: int,
end_token: int,
importance: float = 0.8,
allow_nested: bool = False,
suppress_children: bool = True,
) -> HydratedPhraseMatch:
"""Build a hydrated match for overlap tests."""
return HydratedPhraseMatch(
phrase_id=phrase_id,
matched_norm=phrase_text,
phrase_text=phrase_text,
phrase_norm=phrase_text,
canonical_id=canonical_id,
phrase_type=None,
token_count=end_token - start_token,
confidence=0.9,
importance=importance,
allow_nested=allow_nested,
suppress_children=suppress_children,
start_token=start_token,
end_token=end_token,
)
+6 -30
View File
@@ -2,11 +2,10 @@
from __future__ import annotations
import asyncio
from threading import Event
from typing import TYPE_CHECKING
from sqlalchemy.ext.asyncio import create_async_engine
from sqlalchemy import create_engine
from python.ebook_search.config import EbookSearchConfig, RerankConfig
from python.ebook_search.search import SearchResult, search_ebooks
@@ -15,18 +14,18 @@ if TYPE_CHECKING:
from pytest_mock import MockerFixture
async def test_search_ebooks_runs_vector_and_bm25_in_parallel(mocker: MockerFixture) -> None:
engine = create_async_engine("sqlite+aiosqlite:///:memory:")
def test_search_ebooks_runs_vector_and_bm25_in_parallel(mocker: MockerFixture) -> None:
engine = create_engine("sqlite+pysqlite:///:memory:", future=True)
vector_started = Event()
bm25_started = Event()
received_engines: list[object] = []
async def fake_vector_candidates(received_engine, _client, query, _config):
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 await asyncio.to_thread(bm25_started.wait, 2)
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):
@@ -40,7 +39,7 @@ async def test_search_ebooks_runs_vector_and_bm25_in_parallel(mocker: MockerFixt
mocker.patch("python.ebook_search.search.bm25_candidates", side_effect=fake_bm25_candidates)
config = EbookSearchConfig(rerank=RerankConfig(enabled=False))
response = await search_ebooks(engine, mocker.Mock(), "what is parallel", config)
response = search_ebooks(engine, "what is parallel", config)
timings = {step.name: step for step in response.timings}
assert [result.chunk_id for result in response.results] == [1, 2]
@@ -48,26 +47,3 @@ async def test_search_ebooks_runs_vector_and_bm25_in_parallel(mocker: MockerFixt
assert timings["BM25 search"].counts_toward_total is False
assert timings["Hybrid retrieval"].counts_toward_total is True
assert received_engines == [engine]
async def test_search_ebooks_skips_phrase_matching_when_disabled(mocker: MockerFixture) -> None:
"""Phrase matching can be disabled for one search request."""
engine = create_async_engine("sqlite+aiosqlite:///:memory:")
mocker.patch(
"python.ebook_search.search.vector_candidates",
return_value=[SearchResult(chunk_id=1, text="vector", source_title="Vector", vector_score=0.9)],
)
mocker.patch("python.ebook_search.search.bm25_candidates", return_value=[])
detect_mock = mocker.patch("python.ebook_search.search.query_phrase_matches")
boost_mock = mocker.patch("python.ebook_search.search.apply_phrase_mention_boosts")
config = EbookSearchConfig(rerank=RerankConfig(enabled=False))
response = await search_ebooks(engine, mocker.Mock(), "what is parallel", config, phrase_matching=False)
timing_names = {step.name for step in response.timings}
assert [result.chunk_id for result in response.results] == [1]
assert response.phrase_matches == ()
assert "Protected phrase detection skipped" in timing_names
assert "Phrase mention boost skipped" in timing_names
detect_mock.assert_not_called()
boost_mock.assert_not_called()
+14 -22
View File
@@ -24,13 +24,6 @@ def candidates() -> list[SearchResult]:
]
def make_async_client(mocker: MockerFixture, fake_post) -> httpx.AsyncClient:
"""Build a mock async client whose post call is served by fake_post."""
client = mocker.MagicMock(spec=httpx.AsyncClient)
client.post = mocker.AsyncMock(side_effect=fake_post)
return client
def rerank_response(payload: dict[str, object] | None = None, *, content: bytes | None = None) -> httpx.Response:
return httpx.Response(
200,
@@ -66,7 +59,7 @@ def test_reranking_disabled_returns_original_fused_order() -> None:
assert [result.chunk_id for result in response.results] == [1, 2]
async def test_reranking_enabled_reorders_candidates(mocker: MockerFixture) -> None:
def test_reranking_enabled_reorders_candidates(mocker: MockerFixture) -> None:
def fake_post(_url: str, *, json: dict[str, object], timeout: float) -> httpx.Response:
assert timeout == 30
assert json == {
@@ -84,16 +77,16 @@ async def test_reranking_enabled_reorders_candidates(mocker: MockerFixture) -> N
}
)
client = make_async_client(mocker, fake_post)
mocker.patch.object(httpx, "post", side_effect=fake_post)
results = await rerank_chunks(client, "query", candidates(), RerankConfig())
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]
async def test_reranking_cannot_ignore_hybrid_score(mocker: MockerFixture) -> None:
def test_reranking_cannot_ignore_hybrid_score(mocker: MockerFixture) -> 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),
@@ -109,9 +102,9 @@ async def test_reranking_cannot_ignore_hybrid_score(mocker: MockerFixture) -> No
}
)
client = make_async_client(mocker, fake_post)
mocker.patch.object(httpx, "post", side_effect=fake_post)
results = await rerank_chunks(client, "query", candidates, RerankConfig())
results = rerank_chunks("query", candidates, RerankConfig())
assert [result.chunk_id for result in results] == [1, 2]
assert results[0].score == pytest.approx(0.79)
@@ -119,9 +112,8 @@ async def test_reranking_cannot_ignore_hybrid_score(mocker: MockerFixture) -> No
assert results[1].rerank_score == 1.0
async def test_vllm_rerank_timeout_raises(mocker: MockerFixture) -> None:
def test_vllm_rerank_timeout_raises(mocker: MockerFixture) -> None:
def fake_rerank_chunks(
_client: httpx.AsyncClient,
_query: str,
_candidates: list[SearchResult],
_config: RerankConfig,
@@ -133,21 +125,21 @@ async def test_vllm_rerank_timeout_raises(mocker: MockerFixture) -> None:
config = EbookSearchConfig(rerank=RerankConfig(enabled=True), top_k=2)
with pytest.raises(httpx.TimeoutException, match="timeout"):
await apply_rerank(mocker.Mock(), "query", candidates(), config)
apply_rerank("query", candidates(), config)
async def test_malformed_vllm_rerank_json_does_not_crash_search(mocker: MockerFixture) -> None:
def test_malformed_vllm_rerank_json_does_not_crash_search(mocker: MockerFixture) -> None:
def fake_post(_url: str, **_kwargs: object) -> httpx.Response:
return rerank_response(content=b"not-json")
client = make_async_client(mocker, fake_post)
mocker.patch.object(httpx, "post", side_effect=fake_post)
results = await rerank_chunks(client, "query", candidates()[:1], RerankConfig())
results = rerank_chunks("query", candidates()[:1], RerankConfig())
assert results[0].score == 0.3
async def test_vllm_rerank_scores_are_clamped(mocker: MockerFixture) -> None:
def test_vllm_rerank_scores_are_clamped(mocker: MockerFixture) -> None:
def fake_post(_url: str, **_kwargs: object) -> httpx.Response:
return rerank_response(
{
@@ -158,8 +150,8 @@ async def test_vllm_rerank_scores_are_clamped(mocker: MockerFixture) -> None:
}
)
client = make_async_client(mocker, fake_post)
mocker.patch.object(httpx, "post", side_effect=fake_post)
results = await rerank_chunks(client, "query", candidates()[:2], RerankConfig())
results = rerank_chunks("query", candidates()[:2], RerankConfig())
assert {result.chunk_id: result.rerank_score for result in results} == {1: 0.0, 2: 1.0}
+20 -369
View File
@@ -2,49 +2,32 @@
from __future__ import annotations
import asyncio
from compression import zstd
from datetime import UTC, datetime
from typing import TYPE_CHECKING
from fastapi import BackgroundTasks
from fastapi.testclient import TestClient
from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
from sqlalchemy.pool import StaticPool
from sqlalchemy import create_engine
from python.ebook_search.api.bm25_tasks import refresh_bm25_for_engine
from python.ebook_search.api.judge_tasks import (
is_judging_book,
judge_book_phrases_for_app,
pop_book_judgment_outcome,
start_book_phrase_judgment,
)
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.protected_phrases.models import (
CorpusPhraseStats,
PhraseCandidateGenerationResult,
PhraseJudgmentBackfillResult,
)
from python.ebook_search.search import SearchResponse, SearchResult
from python.ebook_search.timing import RuntimeStep
from python.orm.richie import EbookSource, RichieBase
if TYPE_CHECKING:
from pytest_mock import MockerFixture
from sqlalchemy.ext.asyncio import AsyncEngine
def patch_app_runtime(mocker: MockerFixture):
"""Patch app startup dependencies used by UI route tests."""
mocker.patch("python.ebook_search.api.main.get_async_postgres_engine", side_effect=fake_get_postgres_engine)
mocker.patch("python.ebook_search.api.main.get_postgres_engine", side_effect=fake_get_postgres_engine)
mocker.patch("python.ebook_search.api.main.ensure_bm25_corpus", side_effect=lambda _session, _config: None)
def fake_get_postgres_engine(**_kwargs):
"""Return an in-memory engine for route tests."""
return create_async_engine("sqlite+aiosqlite:///:memory:")
return create_engine("sqlite+pysqlite:///:memory:", future=True)
def test_search_page_uses_zstd_when_requested(mocker: MockerFixture) -> None:
@@ -60,68 +43,36 @@ def test_search_page_uses_zstd_when_requested(mocker: MockerFixture) -> None:
assert b"EPUB Search" in zstd.decompress(response.content)
def test_ui_form_passes_search_toggles_to_search_handler(mocker: MockerFixture) -> None:
def test_ui_form_passes_rerank_flag_to_search_handler(mocker: MockerFixture) -> None:
captured: dict[str, object] = {}
def fake_search_ebooks(_engine, _client, query, config, *, rerank=False, phrase_matching=False):
def fake_search_ebooks(_engine, query, config, *, rerank=False):
captured["query"] = query
captured["rerank"] = rerank
captured["phrase_matching"] = phrase_matching
captured["config"] = config
return SearchResponse(query=query, results=[], rank_label="Hybrid + rerank")
mocker.patch("python.ebook_search.api.routes.search.search_ebooks", side_effect=fake_search_ebooks)
mocker.patch(
"python.ebook_search.api.routes.search.answer_query",
side_effect=lambda _client, _query, _results, _config: "answer",
side_effect=lambda _query, _results, _config: "answer",
)
patch_app_runtime(mocker)
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", "phrase_matching": "true"},
)
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
assert captured["phrase_matching"] is True
def test_ui_form_can_disable_phrase_matching(mocker: MockerFixture) -> None:
captured: dict[str, object] = {}
def fake_search_ebooks(_engine, _client, query, _config, *, rerank=False, phrase_matching=False):
del rerank
captured["query"] = query
captured["phrase_matching"] = phrase_matching
return SearchResponse(query=query, results=[], rank_label="Hybrid")
mocker.patch("python.ebook_search.api.routes.search.search_ebooks", side_effect=fake_search_ebooks)
mocker.patch(
"python.ebook_search.api.routes.search.answer_query",
side_effect=lambda _client, _query, _results, _config: "answer",
)
patch_app_runtime(mocker)
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 captured["query"] == "where is the quote?"
assert captured["phrase_matching"] is False
def test_ui_search_failure_returns_visible_error(mocker: MockerFixture) -> None:
def fake_search_ebooks(_engine, _client, _query, _config, *, rerank=False, phrase_matching=False):
def fake_search_ebooks(_engine, _query, _config, *, rerank=False):
del rerank
del phrase_matching
msg = "search exploded"
raise RuntimeError(msg)
@@ -138,12 +89,11 @@ def test_ui_search_failure_returns_visible_error(mocker: MockerFixture) -> None:
def test_ui_answer_failure_still_returns_sources(mocker: MockerFixture) -> None:
def fake_search_ebooks(_engine, _client, query, _config, *, rerank=False, phrase_matching=False):
def fake_search_ebooks(_engine, query, _config, *, rerank=False):
del rerank
del phrase_matching
return SearchResponse(query=query, results=[], rank_label="Hybrid")
def fake_answer_query(_client, _query, _results, _config):
def fake_answer_query(_query, _results, _config):
msg = "answer exploded"
raise RuntimeError(msg)
@@ -163,12 +113,11 @@ def test_ui_answer_failure_still_returns_sources(mocker: MockerFixture) -> None:
def test_ui_skips_answer_when_disabled(mocker: MockerFixture) -> None:
called = False
def fake_search_ebooks(_engine, _client, query, _config, *, rerank=False, phrase_matching=False):
def fake_search_ebooks(_engine, query, _config, *, rerank=False):
del rerank
del phrase_matching
return SearchResponse(query=query, results=[], rank_label="Hybrid")
def fake_answer_query(_client, _query, _results, _config):
def fake_answer_query(_query, _results, _config):
nonlocal called
called = True
return "answer"
@@ -189,9 +138,8 @@ def test_ui_skips_answer_when_disabled(mocker: MockerFixture) -> None:
def test_ui_shows_component_scores(mocker: MockerFixture) -> None:
def fake_search_ebooks(_engine, _client, query, _config, *, rerank=False, phrase_matching=False):
def fake_search_ebooks(_engine, query, _config, *, rerank=False):
del rerank
del phrase_matching
return SearchResponse(
query=query,
rank_label="Hybrid + rerank",
@@ -212,7 +160,7 @@ def test_ui_shows_component_scores(mocker: MockerFixture) -> None:
mocker.patch("python.ebook_search.api.routes.search.search_ebooks", side_effect=fake_search_ebooks)
mocker.patch(
"python.ebook_search.api.routes.search.answer_query",
side_effect=lambda _client, _query, _results, _config: "answer",
side_effect=lambda _query, _results, _config: "answer",
)
patch_app_runtime(mocker)
app = create_app()
@@ -228,47 +176,9 @@ def test_ui_shows_component_scores(mocker: MockerFixture) -> None:
assert "RRF" in response.text
def test_ui_shows_matched_phrases_that_boosted_a_result(mocker: MockerFixture) -> None:
def fake_search_ebooks(_engine, _client, query, _config, *, rerank=False, phrase_matching=False):
del rerank
del phrase_matching
return SearchResponse(
query=query,
rank_label="Hybrid",
results=[
SearchResult(
chunk_id=1,
text="source text",
source_title="Book",
score=0.9,
phrase_hit_count=3,
matched_phrases=("lock in", "haden's syndrome"),
)
],
)
mocker.patch("python.ebook_search.api.routes.search.search_ebooks", side_effect=fake_search_ebooks)
mocker.patch(
"python.ebook_search.api.routes.search.answer_query",
side_effect=lambda _client, _query, _results, _config: "answer",
)
patch_app_runtime(mocker)
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": "what is lock in?"})
assert response.status_code == 200
assert "boosted by" in response.text
assert "lock in" in response.text
assert "haden&#39;s syndrome" in response.text
def test_ui_shows_search_runtime_chart(mocker: MockerFixture) -> None:
def fake_search_ebooks(_engine, _client, query, _config, *, rerank=False, phrase_matching=False):
def fake_search_ebooks(_engine, query, _config, *, rerank=False):
del rerank
del phrase_matching
return SearchResponse(
query=query,
rank_label="Hybrid",
@@ -282,7 +192,7 @@ def test_ui_shows_search_runtime_chart(mocker: MockerFixture) -> None:
mocker.patch("python.ebook_search.api.routes.search.search_ebooks", side_effect=fake_search_ebooks)
mocker.patch(
"python.ebook_search.api.routes.search.answer_query",
side_effect=lambda _client, _query, _results, _config: "answer",
side_effect=lambda _query, _results, _config: "answer",
)
patch_app_runtime(mocker)
app = create_app()
@@ -304,7 +214,7 @@ def test_ui_embed_all_batches_until_complete(mocker: MockerFixture) -> None:
counts = iter([32, 32, 5, 0])
batch_sizes: list[int] = []
def fake_embed_missing_chunks(_session, _client, config):
def fake_embed_missing_chunks(_session, config):
batch_sizes.append(config.embedding_batch_size)
return next(counts)
@@ -346,7 +256,7 @@ def test_ui_scan_schedules_bm25_refresh_after_database_change(mocker: MockerFixt
assert scheduled is True
async def test_bm25_refresh_clears_loaded_corpus_cache(mocker: MockerFixture) -> None:
def test_bm25_refresh_clears_loaded_corpus_cache(mocker: MockerFixture) -> None:
refreshed: list[object] = []
cache_cleared = False
@@ -359,146 +269,16 @@ async def test_bm25_refresh_clears_loaded_corpus_cache(mocker: MockerFixture) ->
mocker.patch("python.ebook_search.api.bm25_tasks.refresh_bm25_corpus", side_effect=fake_refresh_bm25_corpus)
mocker.patch("python.ebook_search.api.bm25_tasks.load_bm25_corpus.cache_clear", side_effect=fake_cache_clear)
engine = create_async_engine("sqlite+aiosqlite:///:memory:")
engine = create_engine("sqlite+pysqlite:///:memory:", future=True)
config = EbookSearchConfig(rerank=RerankConfig(enabled=False))
await refresh_bm25_for_engine(engine, config)
refresh_bm25_for_engine(engine, config)
assert len(refreshed) == 1
assert refreshed[0][1] == config
assert cache_cleared is True
def build_engine_with_book() -> AsyncEngine:
"""Create a shareable in-memory async engine holding one indexed book."""
engine = create_async_engine(
"sqlite+aiosqlite:///:memory:",
connect_args={"check_same_thread": False},
poolclass=StaticPool,
)
async def seed() -> None:
async with engine.begin() as connection:
await connection.run_sync(RichieBase.metadata.create_all)
async with AsyncSession(engine) as session:
session.add(
EbookSource(
title="Book",
author="Author",
language=None,
publisher=None,
identifier=None,
file_path="/library/book.epub",
file_sha256="a" * 64,
file_mtime=datetime.now(tz=UTC),
file_size=10,
)
)
await session.commit()
asyncio.run(seed())
return engine
def test_ui_judge_phrases_redirects_and_judges_in_background(mocker: MockerFixture) -> None:
mocker.patch("python.ebook_search.api.main.get_async_postgres_engine", return_value=build_engine_with_book())
mocker.patch("python.ebook_search.api.main.ensure_bm25_corpus", side_effect=lambda _session, _config: None)
judged_source_ids: list[list[int]] = []
def fake_judge(_engine: object, _config: object, *, source_ids: list[int]) -> PhraseJudgmentBackfillResult:
judged_source_ids.append(source_ids)
return PhraseJudgmentBackfillResult(
books_seen=1,
books_judged=1,
books_failed=0,
candidates_judged=3,
protected_phrases=2,
phrase_mentions=4,
)
mocker.patch(
"python.ebook_search.api.judge_tasks.judge_candidate_phrases_for_books",
side_effect=fake_judge,
)
app = create_app()
with TestClient(app) as client:
response = client.post("/books/1/judge-phrases", follow_redirects=False)
detail_after = client.get("/books/1")
detail_again = client.get("/books/1")
assert response.status_code == 303
assert response.headers["location"] == "/books/1"
assert judged_source_ids == [[1]]
assert "Judged 3 candidates; 2 protected phrases promoted" in detail_after.text
assert "Judged 3 candidates" not in detail_again.text
def test_ui_book_detail_shows_judging_in_progress(mocker: MockerFixture) -> None:
mocker.patch("python.ebook_search.api.main.get_async_postgres_engine", return_value=build_engine_with_book())
mocker.patch("python.ebook_search.api.main.ensure_bm25_corpus", side_effect=lambda _session, _config: None)
mocker.patch("python.ebook_search.api.routes.page.is_judging_book", return_value=True)
app = create_app()
with TestClient(app) as client:
response = client.get("/books/1")
assert response.status_code == 200
assert "Judging candidate phrases in the background" in response.text
assert "disabled" in response.text
def test_book_phrase_judgment_rejects_duplicate_while_queued(mocker: MockerFixture) -> None:
mocker.patch(
"python.ebook_search.api.judge_tasks.judge_candidate_phrases_for_books",
return_value=PhraseJudgmentBackfillResult(
books_seen=1,
books_judged=1,
books_failed=0,
candidates_judged=3,
protected_phrases=2,
phrase_mentions=4,
),
)
app = create_app()
app.state.engine = None
app.state.config = EbookSearchConfig(rerank=RerankConfig(enabled=False))
background_tasks = BackgroundTasks()
assert start_book_phrase_judgment(app, background_tasks, 1) is True
assert is_judging_book(app, 1) is True
assert start_book_phrase_judgment(app, background_tasks, 1) is False
assert len(background_tasks.tasks) == 1
asyncio.run(judge_book_phrases_for_app(app, 1))
assert is_judging_book(app, 1) is False
assert pop_book_judgment_outcome(app, 1) == "Judged 3 candidates; 2 protected phrases promoted"
assert pop_book_judgment_outcome(app, 1) is None
assert start_book_phrase_judgment(app, background_tasks, 1) is True
def test_book_phrase_judgment_records_failure_outcome(mocker: MockerFixture) -> None:
def fake_judge(_engine: object, _config: object, *, source_ids: list[int]) -> PhraseJudgmentBackfillResult:
del source_ids
message = "llm judge unavailable"
raise RuntimeError(message)
mocker.patch(
"python.ebook_search.api.judge_tasks.judge_candidate_phrases_for_books",
side_effect=fake_judge,
)
app = create_app()
app.state.engine = None
app.state.config = EbookSearchConfig(rerank=RerankConfig(enabled=False))
start_book_phrase_judgment(app, BackgroundTasks(), 7)
asyncio.run(judge_book_phrases_for_app(app, 7))
assert is_judging_book(app, 7) is False
assert pop_book_judgment_outcome(app, 7) == "Judging failed; see server logs for details"
def test_admin_page_shows_embedding_counts_by_model(mocker: MockerFixture) -> None:
def fake_embedding_model_stats(_session):
return [
@@ -517,10 +297,6 @@ def test_admin_page_shows_embedding_counts_by_model(mocker: MockerFixture) -> No
]
mocker.patch("python.ebook_search.api.routes.admin.embedding_model_stats", side_effect=fake_embedding_model_stats)
mocker.patch(
"python.ebook_search.api.routes.admin.corpus_phrase_stats",
return_value=fake_corpus_phrase_stats(),
)
patch_app_runtime(mocker)
app = create_app()
@@ -534,128 +310,3 @@ def test_admin_page_shows_embedding_counts_by_model(mocker: MockerFixture) -> No
assert "24" in response.text
assert "qwen3-embedding-4b" in response.text
assert "2560" in response.text
def fake_corpus_phrase_stats() -> CorpusPhraseStats:
"""Build distinctive corpus phrase stats for admin page assertions."""
return CorpusPhraseStats(
total_books=17,
books_with_candidates=13,
books_fully_judged=11,
candidate_phrases=901,
judged_candidates=703,
unjudged_candidates=198,
protected_phrases=157,
)
def test_admin_page_shows_protected_phrase_stats(mocker: MockerFixture) -> None:
mocker.patch("python.ebook_search.api.routes.admin.embedding_model_stats", return_value=[])
mocker.patch(
"python.ebook_search.api.routes.admin.corpus_phrase_stats",
return_value=fake_corpus_phrase_stats(),
)
patch_app_runtime(mocker)
app = create_app()
with TestClient(app) as client:
response = client.get("/admin")
assert response.status_code == 200
assert "Protected phrases" in response.text
for value in ("17", "13", "11", "901", "703", "198", "157"):
assert value in response.text
def test_ui_add_missing_phrases_generates_only_missing_books(mocker: MockerFixture) -> None:
captured: dict[str, object] = {}
def fake_generate(_session, _config, *, only_missing=False):
captured["only_missing"] = only_missing
return PhraseCandidateGenerationResult(books_seen=3, books_built=2, candidate_phrases=42)
mocker.patch(
"python.ebook_search.api.routes.admin.generate_candidate_phrases_for_books",
side_effect=fake_generate,
)
patch_app_runtime(mocker)
app = create_app()
with TestClient(app) as client:
response = client.post("/admin/phrases/generate-missing")
assert response.status_code == 200
assert captured["only_missing"] is True
assert "42 candidates stored" in response.text
def test_ui_regenerate_all_phrases_generates_every_book(mocker: MockerFixture) -> None:
captured: dict[str, object] = {}
def fake_generate(_session, _config, *, only_missing=False):
captured["only_missing"] = only_missing
return PhraseCandidateGenerationResult(books_seen=5, books_built=5, candidate_phrases=99)
mocker.patch(
"python.ebook_search.api.routes.admin.generate_candidate_phrases_for_books",
side_effect=fake_generate,
)
patch_app_runtime(mocker)
app = create_app()
with TestClient(app) as client:
response = client.post("/admin/phrases/generate-all")
assert response.status_code == 200
assert captured["only_missing"] is False
assert "5 of 5 books" in response.text
def test_ui_judge_missing_phrases_judges_only_pending_books(mocker: MockerFixture) -> None:
captured: dict[str, object] = {}
async def fake_judge(_engine, _config, *, source_ids=None):
captured["source_ids"] = source_ids
return PhraseJudgmentBackfillResult(
books_seen=2,
books_judged=2,
books_failed=0,
candidates_judged=10,
protected_phrases=4,
phrase_mentions=9,
)
mocker.patch(
"python.ebook_search.api.routes.admin.judge_candidate_phrases_for_books",
side_effect=fake_judge,
)
mocker.patch(
"python.ebook_search.api.routes.admin.book_ids_pending_first_judgment",
return_value=[3, 5],
)
patch_app_runtime(mocker)
app = create_app()
with TestClient(app) as client:
response = client.post("/admin/phrases/judge-missing")
assert response.status_code == 200
assert captured["source_ids"] == [3, 5]
assert "4 protected phrases" in response.text
def test_ui_judge_missing_phrases_reports_when_nothing_is_pending(mocker: MockerFixture) -> None:
judge = mocker.patch("python.ebook_search.api.routes.admin.judge_candidate_phrases_for_books")
mocker.patch(
"python.ebook_search.api.routes.admin.book_ids_pending_first_judgment",
return_value=[],
)
patch_app_runtime(mocker)
app = create_app()
with TestClient(app) as client:
response = client.post("/admin/phrases/judge-missing")
assert response.status_code == 200
assert "have been judged" in response.text
judge.assert_not_called()
-5
View File
@@ -1,5 +0,0 @@
{
home.sessionPath = [
"/home/richie/app_images/"
];
}
@@ -76,7 +76,6 @@
"cSpell.userWords": [
"Cahill",
"Corvidae",
"dedup",
"drivername",
"fastapi",
"Michal",
@@ -1,6 +1,5 @@
{
imports = [
../home/app_image_path.nix
../home/global.nix
../home/gui
];