Compare commits
12
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
247e951a27 | ||
|
|
8c3de690c9 | ||
|
|
91575e3ab2 | ||
|
|
4799054601 | ||
|
|
e9636d19de | ||
|
|
2518ec8551 | ||
|
|
f5ec88f1e5 | ||
|
|
f4c4b11ff8 | ||
|
|
37f37c41ac | ||
|
|
80a521f297 | ||
|
|
4a410dbdf8 | ||
|
|
138b79bf97 |
@@ -3,6 +3,19 @@
|
|||||||
.mypy_cache
|
.mypy_cache
|
||||||
.pytest_cache
|
.pytest_cache
|
||||||
.ruff_cache
|
.ruff_cache
|
||||||
|
.venv
|
||||||
|
**/.venv
|
||||||
|
.env
|
||||||
|
.cache
|
||||||
|
.claude
|
||||||
|
.coverage
|
||||||
|
.vscode
|
||||||
|
.stfolder
|
||||||
|
.literotica_data
|
||||||
|
esphome
|
||||||
|
htmlcov
|
||||||
|
data
|
||||||
|
ebooks
|
||||||
__pycache__
|
__pycache__
|
||||||
**/__pycache__
|
**/__pycache__
|
||||||
*.pyc
|
*.pyc
|
||||||
|
|||||||
@@ -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"
|
||||||
Vendored
-6
@@ -71,7 +71,6 @@
|
|||||||
"ehci",
|
"ehci",
|
||||||
"emerg",
|
"emerg",
|
||||||
"endlessh",
|
"endlessh",
|
||||||
"ents",
|
|
||||||
"errorlens",
|
"errorlens",
|
||||||
"esbenp",
|
"esbenp",
|
||||||
"esphome",
|
"esphome",
|
||||||
@@ -173,8 +172,6 @@
|
|||||||
"Networkd",
|
"Networkd",
|
||||||
"networkmanager",
|
"networkmanager",
|
||||||
"newtabpage",
|
"newtabpage",
|
||||||
"ngram",
|
|
||||||
"ngrams",
|
|
||||||
"nixfmt",
|
"nixfmt",
|
||||||
"nixos",
|
"nixos",
|
||||||
"nixpkgs",
|
"nixpkgs",
|
||||||
@@ -301,9 +298,7 @@
|
|||||||
"uiprotect",
|
"uiprotect",
|
||||||
"uitour",
|
"uitour",
|
||||||
"unifi",
|
"unifi",
|
||||||
"unjudged",
|
|
||||||
"unrar",
|
"unrar",
|
||||||
"unstorable",
|
|
||||||
"unsubmitted",
|
"unsubmitted",
|
||||||
"uptimekuma",
|
"uptimekuma",
|
||||||
"urlbar",
|
"urlbar",
|
||||||
@@ -331,7 +326,6 @@
|
|||||||
"xcursorgen",
|
"xcursorgen",
|
||||||
"xdist",
|
"xdist",
|
||||||
"xhci",
|
"xhci",
|
||||||
"yake",
|
|
||||||
"yazi",
|
"yazi",
|
||||||
"yubikey",
|
"yubikey",
|
||||||
"yubioath",
|
"yubioath",
|
||||||
|
|||||||
Generated
+18
-18
@@ -8,11 +8,11 @@
|
|||||||
},
|
},
|
||||||
"locked": {
|
"locked": {
|
||||||
"dir": "pkgs/firefox-addons",
|
"dir": "pkgs/firefox-addons",
|
||||||
"lastModified": 1782964936,
|
"lastModified": 1783828963,
|
||||||
"narHash": "sha256-wXEBDr7/dFQYhVpDwCKc9fkrYQQE4x0bdirX1bsLBGA=",
|
"narHash": "sha256-eTytzcUJCaDUZ3/9EF0+V3fvlikQMQBwiX1Sx4Gy+No=",
|
||||||
"owner": "rycee",
|
"owner": "rycee",
|
||||||
"repo": "nur-expressions",
|
"repo": "nur-expressions",
|
||||||
"rev": "64feee871e0373dd6121e412c3fb12e372d1bfb5",
|
"rev": "8d61e9afde605cd6c22dab68b83d7a71f0a6c5b2",
|
||||||
"type": "gitlab"
|
"type": "gitlab"
|
||||||
},
|
},
|
||||||
"original": {
|
"original": {
|
||||||
@@ -29,11 +29,11 @@
|
|||||||
]
|
]
|
||||||
},
|
},
|
||||||
"locked": {
|
"locked": {
|
||||||
"lastModified": 1783005591,
|
"lastModified": 1783823409,
|
||||||
"narHash": "sha256-NcLHV5uBAeggDUE2wPbKszjfyaSLsoqaYt7izOphkZw=",
|
"narHash": "sha256-OI4IkRjRXa1e7hYmCGJDPDq5H/kPwhsyoS80cNUF9fI=",
|
||||||
"owner": "nix-community",
|
"owner": "nix-community",
|
||||||
"repo": "home-manager",
|
"repo": "home-manager",
|
||||||
"rev": "f469c79b955609d6a8fdd9e689be76a93b1621d7",
|
"rev": "7566825d4652a1b885bd4ce65bd9e8def432fec9",
|
||||||
"type": "github"
|
"type": "github"
|
||||||
},
|
},
|
||||||
"original": {
|
"original": {
|
||||||
@@ -47,11 +47,11 @@
|
|||||||
"nixpkgs": "nixpkgs"
|
"nixpkgs": "nixpkgs"
|
||||||
},
|
},
|
||||||
"locked": {
|
"locked": {
|
||||||
"lastModified": 1782562157,
|
"lastModified": 1783792734,
|
||||||
"narHash": "sha256-a7+T6QSeowynwZ1ZJJbP8T8ntAytvrui8kFGJmIZt2c=",
|
"narHash": "sha256-50rvY9GdFvpYDcMLcD/4cWSi0hVxArT5wsGlVsHy8eY=",
|
||||||
"owner": "nixos",
|
"owner": "nixos",
|
||||||
"repo": "nixos-hardware",
|
"repo": "nixos-hardware",
|
||||||
"rev": "a9cf7546a938c737b079e738de73934a13de9784",
|
"rev": "8efb4337e857949f4cfac86d12ef1066f417f31f",
|
||||||
"type": "github"
|
"type": "github"
|
||||||
},
|
},
|
||||||
"original": {
|
"original": {
|
||||||
@@ -76,11 +76,11 @@
|
|||||||
},
|
},
|
||||||
"nixpkgs-master": {
|
"nixpkgs-master": {
|
||||||
"locked": {
|
"locked": {
|
||||||
"lastModified": 1783021952,
|
"lastModified": 1783874024,
|
||||||
"narHash": "sha256-8PghAtSGGZ0umfVI8Qbd7ZbFrfZPiH1UwtVbgLeikDA=",
|
"narHash": "sha256-Fd8rPvyBv6JjcO/nZxZiFQan6Fww/jAF4TYj0Th/Yfo=",
|
||||||
"owner": "nixos",
|
"owner": "nixos",
|
||||||
"repo": "nixpkgs",
|
"repo": "nixpkgs",
|
||||||
"rev": "f136374c679c54171a3ace589d15e9e79a8bd086",
|
"rev": "0b4f03c64b236e4ba4252414274e92796c300124",
|
||||||
"type": "github"
|
"type": "github"
|
||||||
},
|
},
|
||||||
"original": {
|
"original": {
|
||||||
@@ -108,11 +108,11 @@
|
|||||||
},
|
},
|
||||||
"nixpkgs_2": {
|
"nixpkgs_2": {
|
||||||
"locked": {
|
"locked": {
|
||||||
"lastModified": 1782723713,
|
"lastModified": 1783776592,
|
||||||
"narHash": "sha256-oPXCU/SSUokcGaJREHibG1CBX3+s/W7orDWQOZDsEeQ=",
|
"narHash": "sha256-UgCQzxeWI75XM8G+hPrPh+MKzEPjG3SpAj7dtqSbksA=",
|
||||||
"owner": "nixos",
|
"owner": "nixos",
|
||||||
"repo": "nixpkgs",
|
"repo": "nixpkgs",
|
||||||
"rev": "b5aa0fbd538984f6e3d201be0005b4463d8b09f8",
|
"rev": "e7a3ca8092b61ff85b6a45bf863ea2b2d6a661b3",
|
||||||
"type": "github"
|
"type": "github"
|
||||||
},
|
},
|
||||||
"original": {
|
"original": {
|
||||||
@@ -141,11 +141,11 @@
|
|||||||
]
|
]
|
||||||
},
|
},
|
||||||
"locked": {
|
"locked": {
|
||||||
"lastModified": 1782165805,
|
"lastModified": 1783174389,
|
||||||
"narHash": "sha256-478kKQBvK6SYTOdN2h9jhKJv94nbXRbFMfuL1WshErg=",
|
"narHash": "sha256-aCWC8ngycU7OdJrU2+Je3qf+1a2ykuBvpPhZT/9tXMc=",
|
||||||
"owner": "Mic92",
|
"owner": "Mic92",
|
||||||
"repo": "sops-nix",
|
"repo": "sops-nix",
|
||||||
"rev": "56b24064fdcaedca53553b1a6d607fd23b613a24",
|
"rev": "f1406619a3884cd5c47992a70b8b35c9c0fcb4c9",
|
||||||
"type": "github"
|
"type": "github"
|
||||||
},
|
},
|
||||||
"original": {
|
"original": {
|
||||||
|
|||||||
@@ -17,9 +17,7 @@
|
|||||||
|
|
||||||
python-env = final: _prev: {
|
python-env = final: _prev: {
|
||||||
my_python = final.python314.withPackages (
|
my_python = final.python314.withPackages (
|
||||||
ps:
|
ps: with ps; [
|
||||||
with ps;
|
|
||||||
[
|
|
||||||
alembic
|
alembic
|
||||||
apprise
|
apprise
|
||||||
apscheduler
|
apscheduler
|
||||||
|
|||||||
+2
-5
@@ -25,7 +25,7 @@ dependencies = [
|
|||||||
"pydantic",
|
"pydantic",
|
||||||
"pydantic-settings",
|
"pydantic-settings",
|
||||||
"python-multipart",
|
"python-multipart",
|
||||||
"sqlalchemy[asyncio]",
|
"sqlalchemy",
|
||||||
"tenacity",
|
"tenacity",
|
||||||
"tiktoken",
|
"tiktoken",
|
||||||
"tinytuya",
|
"tinytuya",
|
||||||
@@ -41,10 +41,8 @@ whisper-transcribe = "python.tools.whisper.transcribe:main"
|
|||||||
|
|
||||||
[dependency-groups]
|
[dependency-groups]
|
||||||
dev = [
|
dev = [
|
||||||
"aiosqlite",
|
|
||||||
"mypy",
|
"mypy",
|
||||||
"pyfakefs",
|
"pyfakefs",
|
||||||
"pytest-asyncio",
|
|
||||||
"pytest-cov",
|
"pytest-cov",
|
||||||
"pytest-mock",
|
"pytest-mock",
|
||||||
"pytest-xdist",
|
"pytest-xdist",
|
||||||
@@ -119,7 +117,6 @@ exclude_lines = [
|
|||||||
]
|
]
|
||||||
|
|
||||||
[tool.pytest.ini_options]
|
[tool.pytest.ini_options]
|
||||||
addopts = "-n auto -ra"
|
addopts = "-n auto -ra --ignore=tests/ebook_search"
|
||||||
asyncio_mode = "auto"
|
|
||||||
testpaths = ["tests"]
|
testpaths = ["tests"]
|
||||||
# --cov=system_tools --cov-report=term-missing --cov-report=xml --cov-report=html --cov-branch
|
# --cov=system_tools --cov-report=term-missing --cov-report=xml --cov-report=html --cov-branch
|
||||||
|
|||||||
-206
@@ -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 +0,0 @@
|
|||||||
"""FastAPI applications."""
|
|
||||||
@@ -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)
|
|
||||||
@@ -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"]
|
|
||||||
@@ -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)
|
|
||||||
@@ -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(),
|
|
||||||
},
|
|
||||||
)
|
|
||||||
@@ -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>
|
|
||||||
@@ -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 %}
|
|
||||||
@@ -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 %}
|
|
||||||
@@ -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 %}
|
|
||||||
@@ -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 %}
|
|
||||||
@@ -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 %}
|
|
||||||
@@ -8,20 +8,13 @@ from typing import TYPE_CHECKING
|
|||||||
from python.ebook_search.llm_interface import request_chat_completion
|
from python.ebook_search.llm_interface import request_chat_completion
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
import httpx
|
|
||||||
|
|
||||||
from python.ebook_search.config import EbookSearchConfig
|
from python.ebook_search.config import EbookSearchConfig
|
||||||
from python.ebook_search.search import SearchResult
|
from python.ebook_search.search import SearchResult
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
async def answer_query(
|
def answer_query(query: str, results: list[SearchResult], config: EbookSearchConfig) -> str:
|
||||||
client: httpx.AsyncClient,
|
|
||||||
query: str,
|
|
||||||
results: list[SearchResult],
|
|
||||||
config: EbookSearchConfig,
|
|
||||||
) -> str:
|
|
||||||
"""Answer a question using only retrieved chunks."""
|
"""Answer a question using only retrieved chunks."""
|
||||||
if not config.answer_enabled:
|
if not config.answer_enabled:
|
||||||
logger.info("ebook_answer_skipped_disabled")
|
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}"
|
f"[{index}] {result.source_title}{' - ' + result.chapter_title if result.chapter_title else ''}\n{result.text}"
|
||||||
for index, result in enumerate(results, start=1)
|
for index, result in enumerate(results, start=1)
|
||||||
)
|
)
|
||||||
content = await request_chat_completion(
|
content = request_chat_completion(
|
||||||
client,
|
|
||||||
config,
|
config,
|
||||||
[
|
[
|
||||||
{
|
{
|
||||||
|
|||||||
@@ -1,22 +1,18 @@
|
|||||||
"""Background BM25 refresh tasks for the web app.
|
"""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.
|
|
||||||
"""
|
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import asyncio
|
|
||||||
import logging
|
import logging
|
||||||
|
from threading import Timer
|
||||||
from typing import TYPE_CHECKING
|
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
|
from python.ebook_search.bm25_corpus import load_bm25_corpus, refresh_bm25_corpus
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from fastapi import FastAPI
|
from fastapi import FastAPI
|
||||||
from sqlalchemy.ext.asyncio import AsyncEngine
|
from sqlalchemy.engine import Engine
|
||||||
|
|
||||||
from python.ebook_search.config import EbookSearchConfig
|
from python.ebook_search.config import EbookSearchConfig
|
||||||
|
|
||||||
@@ -24,18 +20,15 @@ logger = logging.getLogger(__name__)
|
|||||||
|
|
||||||
|
|
||||||
def schedule_bm25_refresh(app: FastAPI) -> None:
|
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.
|
timer = Timer(app.state.config.bm25_refresh_delay_seconds, refresh_bm25_for_app, args=(app,))
|
||||||
"""
|
timer.daemon = True
|
||||||
cancel_bm25_refresh(app)
|
timer.start()
|
||||||
|
app.state.bm25_refresh_timer = timer
|
||||||
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)
|
|
||||||
logger.info(
|
logger.info(
|
||||||
"ebook_bm25_refresh_scheduled delay_seconds=%s",
|
"ebook_bm25_refresh_scheduled delay_seconds=%s",
|
||||||
app.state.config.bm25_refresh_delay_seconds,
|
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:
|
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)
|
existing_timer = getattr(app.state, "bm25_refresh_timer", None)
|
||||||
if existing_timer is not None:
|
if existing_timer is not None:
|
||||||
existing_timer.cancel()
|
existing_timer.cancel()
|
||||||
app.state.bm25_refresh_timer = None
|
app.state.bm25_refresh_timer = None
|
||||||
logger.info("ebook_bm25_refresh_cancelled")
|
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
|
|
||||||
|
|
||||||
|
def refresh_bm25_for_app(app: FastAPI) -> None:
|
||||||
async def refresh_bm25_for_app(app: FastAPI) -> None:
|
|
||||||
"""Refresh the BM25 corpus using the app engine and config."""
|
"""Refresh the BM25 corpus using the app engine and config."""
|
||||||
try:
|
try:
|
||||||
await refresh_bm25_for_engine(app.state.engine, app.state.config)
|
refresh_bm25_for_engine(app.state.engine, app.state.config)
|
||||||
except Exception:
|
except Exception:
|
||||||
logger.exception("ebook_bm25_refresh_failed")
|
logger.exception("ebook_bm25_refresh_failed")
|
||||||
|
|
||||||
|
|
||||||
async def refresh_bm25_for_engine(engine: AsyncEngine, config: EbookSearchConfig) -> None:
|
def refresh_bm25_for_engine(engine: Engine, config: EbookSearchConfig) -> None:
|
||||||
"""Refresh the BM25 corpus using an async SQLAlchemy engine."""
|
"""Refresh the BM25 corpus using a SQLAlchemy engine."""
|
||||||
async with AsyncSession(engine) as session:
|
with Session(engine) as session:
|
||||||
await refresh_bm25_corpus(session, config)
|
refresh_bm25_corpus(session, config)
|
||||||
load_bm25_corpus.cache_clear()
|
load_bm25_corpus.cache_clear()
|
||||||
logger.info("ebook_bm25_corpus_cache_cleared_after_refresh")
|
logger.info("ebook_bm25_corpus_cache_cleared_after_refresh")
|
||||||
|
|||||||
@@ -4,9 +4,8 @@ from __future__ import annotations
|
|||||||
|
|
||||||
from typing import Annotated
|
from typing import Annotated
|
||||||
|
|
||||||
import httpx
|
|
||||||
from fastapi import Depends, Request
|
from fastapi import Depends, Request
|
||||||
from sqlalchemy.ext.asyncio import AsyncEngine
|
from sqlalchemy.engine import Engine
|
||||||
|
|
||||||
from python.ebook_search.config import EbookSearchConfig
|
from python.ebook_search.config import EbookSearchConfig
|
||||||
|
|
||||||
@@ -16,16 +15,10 @@ def get_config(request: Request) -> EbookSearchConfig:
|
|||||||
return request.app.state.config
|
return request.app.state.config
|
||||||
|
|
||||||
|
|
||||||
def get_engine(request: Request) -> AsyncEngine:
|
def get_engine(request: Request) -> Engine:
|
||||||
"""Get the database engine from app state."""
|
"""Get the database engine from app state."""
|
||||||
return request.app.state.engine
|
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)]
|
AppConfig = Annotated[EbookSearchConfig, Depends(get_config)]
|
||||||
AppEngine = Annotated[AsyncEngine, Depends(get_engine)]
|
AppEngine = Annotated[Engine, Depends(get_engine)]
|
||||||
AppHttpClient = Annotated[httpx.AsyncClient, Depends(get_http_client)]
|
|
||||||
|
|||||||
@@ -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)
|
|
||||||
@@ -6,12 +6,11 @@ import logging
|
|||||||
from contextlib import asynccontextmanager
|
from contextlib import asynccontextmanager
|
||||||
from typing import TYPE_CHECKING, Annotated
|
from typing import TYPE_CHECKING, Annotated
|
||||||
|
|
||||||
import httpx
|
|
||||||
import typer
|
import typer
|
||||||
import uvicorn
|
import uvicorn
|
||||||
from fastapi import FastAPI
|
from fastapi import FastAPI
|
||||||
from fastapi.staticfiles import StaticFiles
|
from fastapi.staticfiles import StaticFiles
|
||||||
from sqlalchemy.ext.asyncio import AsyncSession
|
from sqlalchemy.orm import Session
|
||||||
|
|
||||||
from python.common import configure_logger
|
from python.common import configure_logger
|
||||||
from python.ebook_search.api.bm25_tasks import cancel_bm25_refresh
|
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.api.web import STATIC_DIR
|
||||||
from python.ebook_search.bm25_corpus import ensure_bm25_corpus
|
from python.ebook_search.bm25_corpus import ensure_bm25_corpus
|
||||||
from python.ebook_search.config import load_config
|
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.fastapi_tools import ZstdMiddleware
|
||||||
from python.orm.common import get_async_postgres_engine
|
from python.orm.common import get_postgres_engine
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from collections.abc import AsyncIterator
|
from collections.abc import AsyncIterator
|
||||||
@@ -38,36 +36,26 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]:
|
|||||||
app.state.config = config
|
app.state.config = config
|
||||||
logger.info(
|
logger.info(
|
||||||
"ebook_search_config_loaded top_k=%s embedding_model=%s embedding_base_url=%s vllm_base_url=%s "
|
"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.top_k,
|
||||||
config.embedding_model,
|
config.embedding_model,
|
||||||
config.embedding_base_url,
|
config.embedding_base_url,
|
||||||
config.vllm_base_url,
|
config.vllm_base_url,
|
||||||
config.rerank.enabled,
|
config.rerank.enabled,
|
||||||
config.phrase_matching_enabled,
|
|
||||||
config.answer_enabled,
|
config.answer_enabled,
|
||||||
len(config.library_paths),
|
len(config.library_paths),
|
||||||
)
|
)
|
||||||
if not config.library_paths:
|
if not config.library_paths:
|
||||||
logger.warning("ebook_search_no_library_paths_configured")
|
logger.warning("ebook_search_no_library_paths_configured")
|
||||||
# Concurrent phrase judging opens one session per book worker on this engine, so size the pool
|
app.state.engine = get_postgres_engine(name="RICHIE", vector_engine=True)
|
||||||
# to cover those plus headroom for ordinary web requests.
|
with Session(app.state.engine) as session:
|
||||||
app.state.engine = get_async_postgres_engine(
|
ensure_bm25_corpus(session, config)
|
||||||
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)
|
|
||||||
try:
|
try:
|
||||||
yield
|
yield
|
||||||
finally:
|
finally:
|
||||||
logger.info("ebook_search_shutdown")
|
logger.info("ebook_search_shutdown")
|
||||||
cancel_bm25_refresh(app)
|
cancel_bm25_refresh(app)
|
||||||
shutdown_extraction_pool()
|
app.state.engine.dispose()
|
||||||
await app.state.http_client.aclose()
|
|
||||||
await app.state.engine.dispose()
|
|
||||||
|
|
||||||
|
|
||||||
def create_app() -> FastAPI:
|
def create_app() -> FastAPI:
|
||||||
|
|||||||
@@ -8,18 +8,13 @@ from fastapi import APIRouter, Request
|
|||||||
from fastapi.responses import HTMLResponse
|
from fastapi.responses import HTMLResponse
|
||||||
|
|
||||||
from python.ebook_search.api.bm25_tasks import schedule_bm25_refresh
|
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
|
from python.ebook_search.api.dependencies import (
|
||||||
AppConfig,
|
AppConfig, # noqa: TC001 FastAPI resolves this annotated dependency at runtime
|
||||||
AppEngine,
|
|
||||||
AppHttpClient,
|
|
||||||
)
|
)
|
||||||
from python.ebook_search.api.web import templates
|
from python.ebook_search.api.web import templates
|
||||||
from python.ebook_search.embeddings import embed_missing_chunks, embedding_model_stats
|
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.ingest import ingest_configured_paths
|
||||||
from python.ebook_search.protected_phrases.generate_ngrams import generate_candidate_phrases_for_books
|
from python.fastapi_tools import DbSession # noqa: TC001 FastAPI resolves this annotated dependency at runtime
|
||||||
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
|
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
@@ -27,29 +22,19 @@ router = APIRouter(prefix="/admin")
|
|||||||
|
|
||||||
|
|
||||||
@router.get("", response_class=HTMLResponse)
|
@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."""
|
"""Render the admin page."""
|
||||||
stats = await embedding_model_stats(session)
|
stats = embedding_model_stats(session)
|
||||||
phrase_stats = await corpus_phrase_stats(session)
|
logger.info("ebook_admin_page_loaded models=%s", len(stats))
|
||||||
logger.info(
|
return templates.TemplateResponse(request, "admin.html", {"config": config, "stats": stats})
|
||||||
"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},
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
@router.post("/scan", response_class=HTMLResponse)
|
@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."""
|
"""Scan configured library paths for EPUB changes."""
|
||||||
try:
|
try:
|
||||||
count = await ingest_configured_paths(session, config)
|
count = ingest_configured_paths(session, config)
|
||||||
await session.commit()
|
session.commit()
|
||||||
except Exception as error:
|
except Exception as error:
|
||||||
logger.exception("ebook_admin_scan_failed")
|
logger.exception("ebook_admin_scan_failed")
|
||||||
return templates.TemplateResponse(request, "partials/error.html", {"message": str(error)}, status_code=500)
|
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"})
|
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)
|
@router.post("/embed-missing", response_class=HTMLResponse)
|
||||||
async def embed_missing(
|
def embed_missing(request: Request, config: AppConfig, session: DbSession) -> HTMLResponse:
|
||||||
request: Request,
|
|
||||||
config: AppConfig,
|
|
||||||
session: AsyncDbSession,
|
|
||||||
client: AppHttpClient,
|
|
||||||
) -> HTMLResponse:
|
|
||||||
"""Embed chunks missing vectors for the configured model."""
|
"""Embed chunks missing vectors for the configured model."""
|
||||||
try:
|
try:
|
||||||
count = await embed_missing_chunks(session, client, config)
|
count = embed_missing_chunks(session, config)
|
||||||
await session.commit()
|
session.commit()
|
||||||
except Exception as error:
|
except Exception as error:
|
||||||
logger.exception("ebook_admin_embed_missing_failed")
|
logger.exception("ebook_admin_embed_missing_failed")
|
||||||
return templates.TemplateResponse(request, "partials/error.html", {"message": str(error)}, status_code=500)
|
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)
|
@router.post("/embed-all", response_class=HTMLResponse)
|
||||||
async def embed_all(
|
def embed_all(request: Request, config: AppConfig, session: DbSession) -> HTMLResponse:
|
||||||
request: Request,
|
|
||||||
config: AppConfig,
|
|
||||||
session: AsyncDbSession,
|
|
||||||
client: AppHttpClient,
|
|
||||||
) -> HTMLResponse:
|
|
||||||
"""Embed all chunks missing vectors in fixed-size batches."""
|
"""Embed all chunks missing vectors in fixed-size batches."""
|
||||||
total = 0
|
total = 0
|
||||||
batches = 0
|
batches = 0
|
||||||
try:
|
try:
|
||||||
while True:
|
while True:
|
||||||
count = await embed_missing_chunks(session, client, config)
|
count = embed_missing_chunks(session, config)
|
||||||
if count == 0:
|
if count == 0:
|
||||||
break
|
break
|
||||||
await session.commit()
|
session.commit()
|
||||||
total += count
|
total += count
|
||||||
batches += 1
|
batches += 1
|
||||||
logger.info(
|
logger.info(
|
||||||
|
|||||||
@@ -11,17 +11,15 @@ from fastapi.responses import JSONResponse
|
|||||||
from sqlalchemy import literal, select
|
from sqlalchemy import literal, select
|
||||||
from sqlalchemy.exc import SQLAlchemyError
|
from sqlalchemy.exc import SQLAlchemyError
|
||||||
|
|
||||||
from python.ebook_search.api.dependencies import ( # noqa: TC001 FastAPI resolves these annotated dependencies at runtime
|
from python.ebook_search.api.dependencies import (
|
||||||
AppConfig,
|
AppConfig, # noqa: TC001 FastAPI resolves this annotated dependency at runtime
|
||||||
AppHttpClient,
|
|
||||||
)
|
)
|
||||||
from python.ebook_search.bm25_corpus import bm25_index_exists, bm25_index_path, read_bm25_manifest
|
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.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:
|
if TYPE_CHECKING:
|
||||||
import httpx
|
from sqlalchemy.orm import Session
|
||||||
from sqlalchemy.ext.asyncio import AsyncSession
|
|
||||||
|
|
||||||
from python.ebook_search.config import EbookSearchConfig
|
from python.ebook_search.config import EbookSearchConfig
|
||||||
|
|
||||||
@@ -31,17 +29,17 @@ router = APIRouter()
|
|||||||
|
|
||||||
|
|
||||||
@router.get("/health")
|
@router.get("/health")
|
||||||
async def health() -> dict[str, str]:
|
def health() -> dict[str, str]:
|
||||||
"""Liveness probe that returns ok without touching dependencies."""
|
"""Liveness probe that returns ok without touching dependencies."""
|
||||||
return {"status": "ok"}
|
return {"status": "ok"}
|
||||||
|
|
||||||
|
|
||||||
@router.get("/ready")
|
@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."""
|
"""Readiness probe reporting database, embedding endpoint, and BM25 index status."""
|
||||||
database_ok = await check_database(session)
|
database_ok = check_database(session)
|
||||||
embedding_ok = await check_embedding_endpoint(client, config)
|
embedding_ok = check_embedding_endpoint(config)
|
||||||
chat_status = await chat_endpoint_status(client, config)
|
chat_status = chat_endpoint_status(config)
|
||||||
bm25_status = check_bm25_status(config)
|
bm25_status = check_bm25_status(config)
|
||||||
|
|
||||||
checks = {
|
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)
|
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."""
|
"""Return the answering chat endpoint status, or disabled when answers are off."""
|
||||||
if not config.answer_enabled:
|
if not config.answer_enabled:
|
||||||
return "disabled"
|
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."""
|
"""Return whether the database answers a trivial query."""
|
||||||
try:
|
try:
|
||||||
await session.execute(select(literal(1)))
|
session.execute(select(literal(1)))
|
||||||
except SQLAlchemyError as error:
|
except SQLAlchemyError as error:
|
||||||
logger.warning("ebook_ready_database_unavailable error=%s", error)
|
logger.warning("ebook_ready_database_unavailable error=%s", error)
|
||||||
return False
|
return False
|
||||||
|
|||||||
@@ -3,23 +3,17 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import logging
|
import logging
|
||||||
from typing import TYPE_CHECKING
|
|
||||||
|
|
||||||
from fastapi import APIRouter, BackgroundTasks, HTTPException, Request
|
from fastapi import APIRouter, Request
|
||||||
from fastapi.responses import HTMLResponse, RedirectResponse
|
from fastapi.responses import HTMLResponse
|
||||||
from sqlalchemy import func, select
|
from sqlalchemy import select
|
||||||
|
|
||||||
from python.ebook_search.api.dependencies import (
|
from python.ebook_search.api.dependencies import (
|
||||||
AppConfig, # noqa: TC001 FastAPI resolves this annotated dependency at runtime
|
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.api.web import templates
|
||||||
from python.ebook_search.protected_phrases.generate_ngrams import recalculate_candidate_phrases_for_book
|
from python.fastapi_tools import DbSession # noqa: TC001 FastAPI resolves this annotated dependency at runtime
|
||||||
from python.fastapi_tools import AsyncDbSession # noqa: TC001 FastAPI resolves this annotated dependency at runtime
|
from python.orm.richie import EbookSource
|
||||||
from python.orm.richie import EbookCandidatePhrase, EbookChapter, EbookChunk, EbookProtectedPhrase, EbookSource
|
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
|
||||||
from sqlalchemy.ext.asyncio import AsyncSession
|
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
@@ -27,176 +21,38 @@ router = APIRouter()
|
|||||||
|
|
||||||
|
|
||||||
@router.get("/", response_class=HTMLResponse)
|
@router.get("/", response_class=HTMLResponse)
|
||||||
async def index(request: Request, config: AppConfig) -> HTMLResponse:
|
def index(request: Request, config: AppConfig) -> HTMLResponse:
|
||||||
"""Render the search page."""
|
"""Render the search page."""
|
||||||
return templates.TemplateResponse(request, "search.html", {"config": config})
|
return templates.TemplateResponse(request, "search.html", {"config": config})
|
||||||
|
|
||||||
|
|
||||||
@router.get("/books", response_class=HTMLResponse)
|
@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."""
|
"""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))
|
logger.info("ebook_books_page_loaded count=%s", len(sources))
|
||||||
return templates.TemplateResponse(request, "books.html", {"sources": 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)
|
@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."""
|
"""Render details for one indexed book."""
|
||||||
source = await session.get(EbookSource, source_id)
|
source = 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"
|
|
||||||
if source is not None:
|
if source is not None:
|
||||||
chapter_count = await get_chapter_count(session, source.id)
|
chapter_count = len(source.chapters)
|
||||||
chunk_count = await get_chunk_count(session, source.id)
|
chunk_count = len(source.chunks)
|
||||||
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)
|
|
||||||
else:
|
else:
|
||||||
chapter_count = 0
|
chapter_count = 0
|
||||||
chunk_count = 0
|
chunk_count = 0
|
||||||
candidate_count = 0
|
|
||||||
judged_candidate_count = 0
|
|
||||||
protected_count = 0
|
|
||||||
candidates = []
|
|
||||||
protected_phrases = []
|
|
||||||
logger.info(
|
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_id,
|
||||||
source is not None,
|
source is not None,
|
||||||
chapter_count,
|
chapter_count,
|
||||||
chunk_count,
|
chunk_count,
|
||||||
candidate_count,
|
|
||||||
judged_candidate_count,
|
|
||||||
protected_count,
|
|
||||||
)
|
)
|
||||||
return templates.TemplateResponse(
|
return templates.TemplateResponse(
|
||||||
request,
|
request,
|
||||||
"book_detail.html",
|
"book_detail.html",
|
||||||
{
|
{"chapter_count": chapter_count, "chunk_count": chunk_count, "source": source},
|
||||||
"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,
|
|
||||||
},
|
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
@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)
|
|
||||||
|
|||||||
@@ -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
|
from python.ebook_search.api.dependencies import ( # noqa: TC001 FastAPI resolves these annotated dependencies at runtime
|
||||||
AppConfig,
|
AppConfig,
|
||||||
AppEngine,
|
AppEngine,
|
||||||
AppHttpClient,
|
|
||||||
)
|
)
|
||||||
from python.ebook_search.api.web import templates
|
from python.ebook_search.api.web import templates
|
||||||
from python.ebook_search.guardrails import (
|
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
|
from python.ebook_search.timing import runtime_step_from_start
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
import httpx
|
|
||||||
|
|
||||||
from python.ebook_search.config import EbookSearchConfig
|
from python.ebook_search.config import EbookSearchConfig
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
@@ -36,8 +33,7 @@ logger = logging.getLogger(__name__)
|
|||||||
router = APIRouter()
|
router = APIRouter()
|
||||||
|
|
||||||
|
|
||||||
async def build_answer(
|
def build_answer(
|
||||||
client: httpx.AsyncClient,
|
|
||||||
query: str,
|
query: str,
|
||||||
response: SearchResponse,
|
response: SearchResponse,
|
||||||
config: EbookSearchConfig,
|
config: EbookSearchConfig,
|
||||||
@@ -60,7 +56,7 @@ async def build_answer(
|
|||||||
return answer, True, None
|
return answer, True, None
|
||||||
|
|
||||||
try:
|
try:
|
||||||
answer = await answer_query(client, query, response.results, config)
|
answer = answer_query(query, response.results, config)
|
||||||
except RuntimeError as error:
|
except RuntimeError as error:
|
||||||
logger.warning("ebook_answer_request_failed_falling_back error=%s", error)
|
logger.warning("ebook_answer_request_failed_falling_back error=%s", error)
|
||||||
return "Answer generation failed. Source chunks are still shown below.", False, None
|
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)
|
@router.post("/search", response_class=HTMLResponse)
|
||||||
async def search(
|
def search(
|
||||||
request: Request,
|
request: Request,
|
||||||
config: AppConfig,
|
config: AppConfig,
|
||||||
engine: AppEngine,
|
engine: AppEngine,
|
||||||
client: AppHttpClient,
|
|
||||||
query: Annotated[str, Form()],
|
query: Annotated[str, Form()],
|
||||||
rerank: Annotated[str | None, Form()] = None,
|
rerank: Annotated[str | None, Form()] = None,
|
||||||
phrase_matching: Annotated[str | None, Form()] = None,
|
|
||||||
) -> HTMLResponse:
|
) -> HTMLResponse:
|
||||||
"""Run a search and render HTMX results."""
|
"""Run a search and render HTMX results."""
|
||||||
try:
|
try:
|
||||||
response = await search_ebooks(
|
response = search_ebooks(engine, query, config, rerank=rerank == "true")
|
||||||
engine,
|
|
||||||
client,
|
|
||||||
query,
|
|
||||||
config,
|
|
||||||
rerank=rerank == "true",
|
|
||||||
phrase_matching=phrase_matching == "true",
|
|
||||||
)
|
|
||||||
except Exception as error:
|
except Exception as error:
|
||||||
logger.exception("ebook_search_request_failed")
|
logger.exception("ebook_search_request_failed")
|
||||||
return templates.TemplateResponse(request, "partials/error.html", {"message": str(error)}, status_code=500)
|
return templates.TemplateResponse(request, "partials/error.html", {"message": str(error)}, status_code=500)
|
||||||
|
|
||||||
answer_start = perf_counter()
|
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"
|
answer_step_name = "Answer generation" if config.answer_enabled else "Answer skipped"
|
||||||
response = replace(
|
response = replace(
|
||||||
response,
|
response,
|
||||||
|
|||||||
@@ -181,12 +181,6 @@ textarea:focus {
|
|||||||
flex-wrap: wrap;
|
flex-wrap: wrap;
|
||||||
}
|
}
|
||||||
|
|
||||||
.search-toggles {
|
|
||||||
display: flex;
|
|
||||||
flex-wrap: wrap;
|
|
||||||
gap: 14px;
|
|
||||||
}
|
|
||||||
|
|
||||||
button {
|
button {
|
||||||
padding: 10px 20px;
|
padding: 10px 20px;
|
||||||
font: inherit;
|
font: inherit;
|
||||||
@@ -218,11 +212,6 @@ button:hover {
|
|||||||
margin-bottom: 24px;
|
margin-bottom: 24px;
|
||||||
}
|
}
|
||||||
|
|
||||||
.actions-grid {
|
|
||||||
display: grid;
|
|
||||||
grid-template-columns: repeat(2, max-content);
|
|
||||||
}
|
|
||||||
|
|
||||||
/* Answer + results */
|
/* Answer + results */
|
||||||
#results {
|
#results {
|
||||||
display: block;
|
display: block;
|
||||||
@@ -314,28 +303,6 @@ button:hover {
|
|||||||
font-variant-numeric: tabular-nums;
|
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 — developer diagnostics, hidden unless dev mode is on */
|
||||||
.runtime {
|
.runtime {
|
||||||
display: none;
|
display: none;
|
||||||
|
|||||||
@@ -1,110 +1,45 @@
|
|||||||
{% extends "base.html" %} {% block title %}EPUB Admin{% endblock %} {% block
|
{% extends "base.html" %}
|
||||||
head %}
|
|
||||||
<script src="https://unpkg.com/htmx.org@2.0.4"></script>
|
{% block title %}EPUB Admin{% endblock %}
|
||||||
{% endblock %} {% block content %}
|
{% block head %}<script src="https://unpkg.com/htmx.org@2.0.4"></script>{% endblock %}
|
||||||
<h1>Admin</h1>
|
|
||||||
<section id="admin-status"></section>
|
{% block content %}
|
||||||
<section class="actions">
|
<h1>Admin</h1>
|
||||||
<form hx-post="/admin/scan" hx-target="#admin-status" hx-swap="innerHTML">
|
<section id="admin-status"></section>
|
||||||
<button type="submit">Scan</button>
|
<section class="actions">
|
||||||
</form>
|
<form hx-post="/admin/scan" hx-target="#admin-status" hx-swap="innerHTML">
|
||||||
</section>
|
<button type="submit">Scan</button>
|
||||||
<section>
|
</form>
|
||||||
<h2>Embeddings</h2>
|
<form hx-post="/admin/embed-missing" hx-target="#admin-status" hx-swap="innerHTML">
|
||||||
<section class="actions">
|
<button type="submit">Embed</button>
|
||||||
<form
|
</form>
|
||||||
hx-post="/admin/embed-missing"
|
<form hx-post="/admin/embed-all" hx-target="#admin-status" hx-swap="innerHTML">
|
||||||
hx-target="#admin-status"
|
<button type="submit">Embed all</button>
|
||||||
hx-swap="innerHTML"
|
</form>
|
||||||
>
|
</section>
|
||||||
<button type="submit">Embed</button>
|
<section>
|
||||||
</form>
|
<h2>Embeddings</h2>
|
||||||
<form
|
<table>
|
||||||
hx-post="/admin/embed-all"
|
<thead>
|
||||||
hx-target="#admin-status"
|
<tr>
|
||||||
hx-swap="innerHTML"
|
<th>Model</th>
|
||||||
>
|
<th>Dimensions</th>
|
||||||
<button type="submit">Embed all</button>
|
<th>Embedded</th>
|
||||||
</form>
|
<th>Missing</th>
|
||||||
</section>
|
<th>Total chunks</th>
|
||||||
<table>
|
</tr>
|
||||||
<thead>
|
</thead>
|
||||||
<tr>
|
<tbody>
|
||||||
<th>Model</th>
|
{% for item in stats %}
|
||||||
<th>Dimensions</th>
|
<tr>
|
||||||
<th>Embedded</th>
|
<td>{{ item.model_name }}</td>
|
||||||
<th>Missing</th>
|
<td>{{ item.dimension }}</td>
|
||||||
<th>Total chunks</th>
|
<td>{{ item.embedded_chunks }}</td>
|
||||||
</tr>
|
<td>{{ item.missing_chunks }}</td>
|
||||||
</thead>
|
<td>{{ item.total_chunks }}</td>
|
||||||
<tbody>
|
</tr>
|
||||||
{% for item in stats %}
|
{% endfor %}
|
||||||
<tr>
|
</tbody>
|
||||||
<td>{{ item.model_name }}</td>
|
</table>
|
||||||
<td>{{ item.dimension }}</td>
|
</section>
|
||||||
<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>
|
|
||||||
{% endblock %}
|
{% endblock %}
|
||||||
|
|||||||
@@ -6,9 +6,6 @@
|
|||||||
{% if source %}
|
{% if source %}
|
||||||
<h1>{{ source.title }}</h1>
|
<h1>{{ source.title }}</h1>
|
||||||
<p class="meta">{{ source.author or "Unknown author" }}</p>
|
<p class="meta">{{ source.author or "Unknown author" }}</p>
|
||||||
{% if phrase_status_message %}
|
|
||||||
<p class="status">{{ phrase_status_message }}</p>
|
|
||||||
{% endif %}
|
|
||||||
<dl class="card">
|
<dl class="card">
|
||||||
<dt>File</dt>
|
<dt>File</dt>
|
||||||
<dd>{{ source.file_path }}</dd>
|
<dd>{{ source.file_path }}</dd>
|
||||||
@@ -16,93 +13,7 @@
|
|||||||
<dd>{{ chapter_count }}</dd>
|
<dd>{{ chapter_count }}</dd>
|
||||||
<dt>Chunks</dt>
|
<dt>Chunks</dt>
|
||||||
<dd>{{ chunk_count }}</dd>
|
<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>
|
</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…{% 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 %}
|
{% else %}
|
||||||
<h1>Book not found</h1>
|
<h1>Book not found</h1>
|
||||||
{% endif %}
|
{% endif %}
|
||||||
|
|||||||
@@ -82,14 +82,6 @@
|
|||||||
</div>
|
</div>
|
||||||
{% endif %}
|
{% endif %}
|
||||||
</dl>
|
</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>
|
</li>
|
||||||
{% endfor %}
|
{% endfor %}
|
||||||
</ol>
|
</ol>
|
||||||
|
|||||||
@@ -7,24 +7,12 @@
|
|||||||
<h1>Search</h1>
|
<h1>Search</h1>
|
||||||
<form class="card" hx-post="/search" hx-target="#results" hx-swap="innerHTML">
|
<form class="card" hx-post="/search" hx-target="#results" hx-swap="innerHTML">
|
||||||
<label for="query">What are you looking for?</label>
|
<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
|
<textarea id="query" name="query" rows="4" placeholder="Ask a question or paste a passage…" required></textarea>
|
||||||
onkeydown="if (event.key === 'Enter' && !event.shiftKey) { event.preventDefault(); this.form.requestSubmit(); }"></textarea>
|
|
||||||
<div class="form-row">
|
<div class="form-row">
|
||||||
<div class="search-toggles">
|
<label class="check">
|
||||||
<label class="check">
|
<input type="checkbox" name="rerank" value="true" {% if config.rerank.enabled %}checked{% endif %}>
|
||||||
<input type="checkbox" name="rerank" value="true" {% if config.rerank.enabled %}checked{% endif %}>
|
Rerank
|
||||||
Rerank
|
</label>
|
||||||
</label>
|
|
||||||
<label class="check">
|
|
||||||
<input
|
|
||||||
type="checkbox"
|
|
||||||
name="phrase_matching"
|
|
||||||
value="true"
|
|
||||||
{% if config.phrase_matching_enabled %}checked{% endif %}
|
|
||||||
>
|
|
||||||
Phrase matching
|
|
||||||
</label>
|
|
||||||
</div>
|
|
||||||
<button type="submit">Search</button>
|
<button type="submit">Search</button>
|
||||||
</div>
|
</div>
|
||||||
</form>
|
</form>
|
||||||
|
|||||||
@@ -2,7 +2,6 @@
|
|||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import asyncio
|
|
||||||
import json
|
import json
|
||||||
import logging
|
import logging
|
||||||
import shutil
|
import shutil
|
||||||
@@ -18,7 +17,7 @@ from sqlalchemy import func, select, union_all
|
|||||||
from python.orm.richie import EbookChapter, EbookChunk, EbookSource
|
from python.orm.richie import EbookChapter, EbookChunk, EbookSource
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
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.config import EbookSearchConfig
|
||||||
|
|
||||||
@@ -74,14 +73,14 @@ def get_current_bm25_index(index_path: Path) -> Path:
|
|||||||
return index_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."""
|
"""Create or refresh the persisted BM25 corpus when it is missing or stale."""
|
||||||
index_path = bm25_index_path(config)
|
index_path = bm25_index_path(config)
|
||||||
manifest = read_bm25_manifest(index_path)
|
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):
|
if not bm25_index_exists(index_path, manifest):
|
||||||
logger.info("ebook_bm25_index_missing path=%s", index_path)
|
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
|
return
|
||||||
if db_updated_at is not None and manifest is not None and manifest.created_at < db_updated_at:
|
if db_updated_at is not None and manifest is not None and manifest.created_at < db_updated_at:
|
||||||
logger.info(
|
logger.info(
|
||||||
@@ -90,7 +89,7 @@ async def ensure_bm25_corpus(session: AsyncSession, config: EbookSearchConfig) -
|
|||||||
manifest.created_at.isoformat(),
|
manifest.created_at.isoformat(),
|
||||||
db_updated_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
|
return
|
||||||
logger.info(
|
logger.info(
|
||||||
"ebook_bm25_index_current path=%s chunks=%s created_at=%s",
|
"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(
|
def refresh_bm25_corpus(
|
||||||
session: AsyncSession,
|
session: Session,
|
||||||
config: EbookSearchConfig,
|
config: EbookSearchConfig,
|
||||||
*,
|
*,
|
||||||
db_updated_at: datetime | None = None,
|
db_updated_at: datetime | None = None,
|
||||||
) -> BM25Manifest:
|
) -> BM25Manifest:
|
||||||
"""Rebuild and persist the BM25 corpus from the current database chunks.
|
"""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.
|
|
||||||
"""
|
|
||||||
index_path = bm25_index_path(config)
|
index_path = bm25_index_path(config)
|
||||||
records, texts = await fetch_bm25_corpus_records(session)
|
records, texts = fetch_bm25_corpus_records(session)
|
||||||
manifest = BM25Manifest(
|
manifest = BM25Manifest(
|
||||||
created_at=datetime.now(tz=UTC),
|
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),
|
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(
|
logger.info(
|
||||||
"ebook_bm25_index_refreshed path=%s chunks=%s created_at=%s",
|
"ebook_bm25_index_refreshed path=%s chunks=%s created_at=%s",
|
||||||
index_path,
|
index_path,
|
||||||
@@ -168,7 +164,7 @@ def score_bm25_corpus(query: str, corpus: BM25Corpus, *, limit: int) -> list[tup
|
|||||||
return results
|
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.
|
"""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
|
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(
|
select(
|
||||||
EbookChunk.id.label("chunk_id"),
|
EbookChunk.id.label("chunk_id"),
|
||||||
EbookChunk.text.label("text"),
|
EbookChunk.text.label("text"),
|
||||||
EbookSource.id.label("source_id"),
|
|
||||||
EbookSource.title.label("source_title"),
|
EbookSource.title.label("source_title"),
|
||||||
EbookSource.author.label("source_author"),
|
EbookSource.author.label("source_author"),
|
||||||
EbookChapter.title.label("chapter_title"),
|
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]] = []
|
records: list[dict[str, object]] = []
|
||||||
texts: list[str] = []
|
texts: list[str] = []
|
||||||
for row in (await session.execute(statement)).mappings():
|
for row in session.execute(statement).mappings():
|
||||||
record = dict(row)
|
record = dict(row)
|
||||||
texts.append(str(record.pop("bm25_text")))
|
texts.append(str(record.pop("bm25_text")))
|
||||||
records.append(record)
|
records.append(record)
|
||||||
return records, texts
|
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."""
|
"""Return the latest source/chapter/chunk update timestamp relevant to BM25 text."""
|
||||||
update_times = union_all(
|
update_times = union_all(
|
||||||
select(func.max(EbookSource.updated).label("updated")),
|
select(func.max(EbookSource.updated).label("updated")),
|
||||||
select(func.max(EbookChapter.updated).label("updated")),
|
select(func.max(EbookChapter.updated).label("updated")),
|
||||||
select(func.max(EbookChunk.updated).label("updated")),
|
select(func.max(EbookChunk.updated).label("updated")),
|
||||||
).subquery()
|
).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(
|
def write_bm25_corpus(
|
||||||
|
|||||||
@@ -86,25 +86,6 @@ class EbookSearchConfig(BaseSettings):
|
|||||||
validate_citations_enabled: bool = True
|
validate_citations_enabled: bool = True
|
||||||
bm25_index_dir: str = ".ebook_search_bm25"
|
bm25_index_dir: str = ".ebook_search_bm25"
|
||||||
bm25_refresh_delay_seconds: int = 60
|
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")
|
@field_validator("library_paths", mode="before")
|
||||||
@classmethod
|
@classmethod
|
||||||
|
|||||||
@@ -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 \
|
ENV PYTHONDONTWRITEBYTECODE=1 \
|
||||||
PYTHONUNBUFFERED=1 \
|
PYTHONUNBUFFERED=1 \
|
||||||
PIP_NO_CACHE_DIR=1 \
|
|
||||||
APP_DIR=/home/richie/dotfiles \
|
APP_DIR=/home/richie/dotfiles \
|
||||||
EBOOK_SEARCH_HOST=0.0.0.0 \
|
UV_PROJECT_ENVIRONMENT=/opt/venv \
|
||||||
EBOOK_SEARCH_PORT=8070 \
|
UV_PYTHON_DOWNLOADS=never \
|
||||||
EBOOK_SEARCH_BM25_INDEX_DIR=/data/bm25
|
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}
|
WORKDIR ${APP_DIR}
|
||||||
|
|
||||||
@@ -14,29 +19,29 @@ RUN apt-get update \
|
|||||||
&& apt-get install -y --no-install-recommends build-essential curl \
|
&& apt-get install -y --no-install-recommends build-essential curl \
|
||||||
&& rm -rf /var/lib/apt/lists/*
|
&& rm -rf /var/lib/apt/lists/*
|
||||||
|
|
||||||
COPY pyproject.toml README.md LICENSE ./
|
COPY python/ebook_search/docker/pyproject.toml python/ebook_search/docker/uv.lock ./
|
||||||
COPY python ./python
|
|
||||||
|
|
||||||
RUN python -m pip install --upgrade pip \
|
RUN uv sync --locked --no-dev
|
||||||
&& python -m pip install \
|
|
||||||
"alembic" \
|
|
||||||
"beautifulsoup4" \
|
FROM base AS test
|
||||||
"bm25s" \
|
|
||||||
"ebooklib" \
|
RUN uv sync --locked
|
||||||
"fastapi" \
|
|
||||||
"httpx" \
|
COPY python ./python
|
||||||
"jinja2" \
|
COPY tests/__init__.py ./tests/__init__.py
|
||||||
"pgvector" \
|
COPY tests/ebook_search ./tests/ebook_search
|
||||||
"psycopg[binary]" \
|
|
||||||
"pydantic" \
|
CMD ["pytest"]
|
||||||
"pydantic-settings" \
|
|
||||||
"python-multipart" \
|
|
||||||
"sqlalchemy[asyncio]" \
|
FROM base AS runtime
|
||||||
"tiktoken" \
|
|
||||||
"typer" \
|
ENV EBOOK_SEARCH_HOST=0.0.0.0 \
|
||||||
"uvicorn[standard]" \
|
EBOOK_SEARCH_PORT=8070 \
|
||||||
"yake" \
|
EBOOK_SEARCH_BM25_INDEX_DIR=/data/bm25
|
||||||
&& python -m pip install --no-deps --editable "${APP_DIR}"
|
|
||||||
|
COPY python ./python
|
||||||
|
|
||||||
RUN useradd --create-home --uid 10001 app \
|
RUN useradd --create-home --uid 10001 app \
|
||||||
&& mkdir -p /data \
|
&& mkdir -p /data \
|
||||||
|
|||||||
@@ -3,26 +3,27 @@
|
|||||||
Run the EPUB search app against the existing Postgres database on `jeeves`:
|
Run the EPUB search app against the existing Postgres database on `jeeves`:
|
||||||
|
|
||||||
```sh
|
```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:
|
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`
|
- `docker-compose.yml`
|
||||||
- `containers.py`
|
- `containers.py` — Typer lifecycle CLI
|
||||||
- `container.py`
|
- `pyproject.toml` / `uv.lock` — the container's uv-locked dependencies
|
||||||
|
|
||||||
The app listens on `http://localhost:8070`.
|
The app listens on `http://localhost:8070`.
|
||||||
|
|
||||||
Useful lifecycle commands:
|
Useful lifecycle commands:
|
||||||
|
|
||||||
```sh
|
```sh
|
||||||
ebook-search-containers build
|
python -m python.ebook_search.docker.containers build
|
||||||
ebook-search-containers start --library-path /path/to/epubs
|
python -m python.ebook_search.docker.containers start --library-path /path/to/epubs
|
||||||
ebook-search-containers logs
|
python -m python.ebook_search.docker.containers test
|
||||||
ebook-search-containers ps
|
python -m python.ebook_search.docker.containers logs
|
||||||
ebook-search-containers stop
|
python -m python.ebook_search.docker.containers ps
|
||||||
|
python -m python.ebook_search.docker.containers stop
|
||||||
```
|
```
|
||||||
|
|
||||||
Direct compose usage from the repo root:
|
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
|
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.
|
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`.
|
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.
|
|
||||||
|
|||||||
@@ -73,6 +73,23 @@ def build_image() -> None:
|
|||||||
raise RuntimeError(msg)
|
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(
|
def start_stack(
|
||||||
*,
|
*,
|
||||||
library_path: Path | None = None,
|
library_path: Path | None = None,
|
||||||
@@ -210,6 +227,19 @@ def logs(
|
|||||||
typer.echo(output)
|
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")
|
@app.command("ps")
|
||||||
def ps() -> None:
|
def ps() -> None:
|
||||||
"""Show ebook search container status."""
|
"""Show ebook search container status."""
|
||||||
|
|||||||
@@ -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"]
|
||||||
Generated
+1109
File diff suppressed because it is too large
Load Diff
@@ -23,8 +23,7 @@ logger = logging.getLogger(__name__)
|
|||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from collections.abc import Sequence
|
from collections.abc import Sequence
|
||||||
|
|
||||||
import httpx
|
from sqlalchemy.orm import Session
|
||||||
from sqlalchemy.ext.asyncio import AsyncSession
|
|
||||||
|
|
||||||
from python.ebook_search.config import EbookSearchConfig
|
from python.ebook_search.config import EbookSearchConfig
|
||||||
|
|
||||||
@@ -66,11 +65,7 @@ class EmbeddingModelStats:
|
|||||||
return max(self.total_chunks - self.embedded_chunks, 0)
|
return max(self.total_chunks - self.embedded_chunks, 0)
|
||||||
|
|
||||||
|
|
||||||
async def embed_texts(
|
def embed_texts(texts: Sequence[str], config: EbookSearchConfig) -> list[list[float]]:
|
||||||
client: httpx.AsyncClient,
|
|
||||||
texts: Sequence[str],
|
|
||||||
config: EbookSearchConfig,
|
|
||||||
) -> list[list[float]]:
|
|
||||||
"""Embed text with the configured vLLM embedding model."""
|
"""Embed text with the configured vLLM embedding model."""
|
||||||
logger.info(
|
logger.info(
|
||||||
"ebook_embed_request_start base_url=%s model=%s count=%s",
|
"ebook_embed_request_start base_url=%s model=%s count=%s",
|
||||||
@@ -78,7 +73,7 @@ async def embed_texts(
|
|||||||
config.embedding_model,
|
config.embedding_model,
|
||||||
len(texts),
|
len(texts),
|
||||||
)
|
)
|
||||||
vectors = await request_embeddings(client, texts, config)
|
vectors = request_embeddings(texts, config)
|
||||||
expected_dimension = MODEL_DIMENSIONS[config.embedding_model]
|
expected_dimension = MODEL_DIMENSIONS[config.embedding_model]
|
||||||
for vector in vectors:
|
for vector in vectors:
|
||||||
if len(vector) != expected_dimension:
|
if len(vector) != expected_dimension:
|
||||||
@@ -93,28 +88,28 @@ async def embed_texts(
|
|||||||
return vectors
|
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."""
|
"""Embed a search query with the Qwen retrieval instruction."""
|
||||||
instructed_query = f"Instruct: Retrieve relevant passages for the query.\nQuery: {query}"
|
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."""
|
"""Ensure supported embedding model rows exist."""
|
||||||
for name, dimension in MODEL_DIMENSIONS.items():
|
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:
|
if existing is None:
|
||||||
session.add(EbookEmbeddingModel(name=name, dimension=dimension, is_default=name == "qwen3-embedding-0.6b"))
|
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)
|
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."""
|
"""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 = {
|
models = {
|
||||||
model.name: model
|
model.name: model
|
||||||
for model in await session.scalars(
|
for model in session.scalars(
|
||||||
select(EbookEmbeddingModel)
|
select(EbookEmbeddingModel)
|
||||||
.where(EbookEmbeddingModel.name.in_(MODEL_DIMENSIONS))
|
.where(EbookEmbeddingModel.name.in_(MODEL_DIMENSIONS))
|
||||||
.order_by(EbookEmbeddingModel.name)
|
.order_by(EbookEmbeddingModel.name)
|
||||||
@@ -127,7 +122,7 @@ async def embedding_model_stats(session: AsyncSession) -> list[EmbeddingModelSta
|
|||||||
embedded_chunks = 0
|
embedded_chunks = 0
|
||||||
if model is not None:
|
if model is not None:
|
||||||
table = get_embedding_table(dimension)
|
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(
|
stats.append(
|
||||||
EmbeddingModelStats(
|
EmbeddingModelStats(
|
||||||
model_name=model_name,
|
model_name=model_name,
|
||||||
@@ -139,10 +134,10 @@ async def embedding_model_stats(session: AsyncSession) -> list[EmbeddingModelSta
|
|||||||
return stats
|
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."""
|
"""Embed chunks missing embeddings for the configured model."""
|
||||||
await ensure_embedding_models(session)
|
ensure_embedding_models(session)
|
||||||
model = await session.scalar(select(EbookEmbeddingModel).where(EbookEmbeddingModel.name == config.embedding_model))
|
model = session.scalar(select(EbookEmbeddingModel).where(EbookEmbeddingModel.name == config.embedding_model))
|
||||||
if model is None:
|
if model is None:
|
||||||
supported_models = ", ".join(MODEL_DIMENSIONS)
|
supported_models = ", ".join(MODEL_DIMENSIONS)
|
||||||
msg = f"Unknown embedding model: {config.embedding_model}. Supported models: {supported_models}"
|
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)
|
table = get_embedding_table(model.dimension)
|
||||||
chunks = list(
|
chunks = list(
|
||||||
await session.scalars(
|
session.scalars(
|
||||||
select(EbookChunk)
|
select(EbookChunk)
|
||||||
.outerjoin(table, (table.chunk_id == EbookChunk.id) & (table.model_id == model.id))
|
.outerjoin(table, (table.chunk_id == EbookChunk.id) & (table.model_id == model.id))
|
||||||
.where(table.id.is_(None))
|
.where(table.id.is_(None))
|
||||||
@@ -163,13 +158,13 @@ async def embed_missing_chunks(session: AsyncSession, client: httpx.AsyncClient,
|
|||||||
return 0
|
return 0
|
||||||
|
|
||||||
logger.info("ebook_embed_missing_batch_start model=%s count=%s", config.embedding_model, len(chunks))
|
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 = [
|
rows = [
|
||||||
{"chunk_id": chunk.id, "model_id": model.id, "embedding": vector}
|
{"chunk_id": chunk.id, "model_id": model.id, "embedding": vector}
|
||||||
for chunk, vector in zip(chunks, vectors, strict=True)
|
for chunk, vector in zip(chunks, vectors, strict=True)
|
||||||
]
|
]
|
||||||
statement = insert(table).values(rows).on_conflict_do_nothing(index_elements=["chunk_id", "model_id"])
|
statement = insert(table).values(rows).on_conflict_do_nothing(index_elements=["chunk_id", "model_id"])
|
||||||
await session.execute(statement)
|
session.execute(statement)
|
||||||
await session.flush()
|
session.flush()
|
||||||
logger.info("ebook_embed_missing_batch_complete model=%s count=%s", config.embedding_model, len(rows))
|
logger.info("ebook_embed_missing_batch_complete model=%s count=%s", config.embedding_model, len(rows))
|
||||||
return len(rows)
|
return len(rows)
|
||||||
|
|||||||
@@ -2,7 +2,6 @@
|
|||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import asyncio
|
|
||||||
import hashlib
|
import hashlib
|
||||||
import logging
|
import logging
|
||||||
from dataclasses import dataclass
|
from dataclasses import dataclass
|
||||||
@@ -14,7 +13,6 @@ import tiktoken
|
|||||||
from sqlalchemy import or_, select
|
from sqlalchemy import or_, select
|
||||||
|
|
||||||
from python.ebook_search.epub_parse import parse_epub
|
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
|
from python.orm.richie import EbookChapter, EbookChunk, EbookSource
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
@@ -22,7 +20,7 @@ DEFAULT_CHUNK_TOKENS = 700
|
|||||||
DEFAULT_CHUNK_OVERLAP = 100
|
DEFAULT_CHUNK_OVERLAP = 100
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
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.config import EbookSearchConfig
|
||||||
from python.ebook_search.epub_parse import ParsedChapter
|
from python.ebook_search.epub_parse import ParsedChapter
|
||||||
@@ -74,115 +72,90 @@ def chunk_text(
|
|||||||
return [chunk for chunk in chunks if chunk.text]
|
return [chunk for chunk in chunks if chunk.text]
|
||||||
|
|
||||||
|
|
||||||
def find_library_epubs(library_path: str) -> tuple[Path, list[Path] | None]:
|
def ingest_configured_paths(session: Session, config: EbookSearchConfig) -> int:
|
||||||
"""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:
|
|
||||||
"""Ingest every EPUB found under configured library paths."""
|
"""Ingest every EPUB found under configured library paths."""
|
||||||
count = 0
|
count = 0
|
||||||
for library_path in config.library_paths:
|
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)
|
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)
|
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))
|
logger.info("ebook_ingest_paths_complete changed_files=%s configured_paths=%s", count, len(config.library_paths))
|
||||||
return count
|
return count
|
||||||
|
|
||||||
|
|
||||||
def resolve_ingest_path(path: Path) -> Path:
|
def ingest_file(session: Session, path: Path, config: EbookSearchConfig) -> bool:
|
||||||
"""Expand and resolve an ingest path (blocking filesystem call)."""
|
|
||||||
return path.expanduser().resolve()
|
|
||||||
|
|
||||||
|
|
||||||
async def ingest_file(session: AsyncSession, path: Path, config: EbookSearchConfig) -> bool:
|
|
||||||
"""Ingest one EPUB file. Return True when the database changed."""
|
"""Ingest one EPUB file. Return True when the database changed."""
|
||||||
try:
|
resolved_path = path.expanduser().resolve()
|
||||||
resolved_path = await asyncio.to_thread(resolve_ingest_path, path)
|
logger.info("ebook_ingest_file_start path=%s", resolved_path)
|
||||||
logger.info("ebook_ingest_file_start path=%s", resolved_path)
|
file_hash = sha256_file(resolved_path)
|
||||||
file_hash = await asyncio.to_thread(sha256_file, resolved_path)
|
existing = find_existing_source(session, resolved_path, file_hash)
|
||||||
existing = await find_existing_source(session, resolved_path, file_hash)
|
if existing is not None and existing.file_sha256 == 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()
|
|
||||||
|
|
||||||
stat = resolved_path.stat()
|
stat = resolved_path.stat()
|
||||||
parsed = await asyncio.to_thread(parse_epub, resolved_path)
|
existing.file_path = str(resolved_path)
|
||||||
source = EbookSource(
|
existing.file_mtime = datetime.fromtimestamp(stat.st_mtime, tz=UTC)
|
||||||
title=parsed.title,
|
existing.file_size = stat.st_size
|
||||||
author=parsed.author,
|
session.flush()
|
||||||
language=parsed.language,
|
logger.info("ebook_ingest_file_unchanged source_id=%s path=%s", existing.id, resolved_path)
|
||||||
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}")
|
|
||||||
return False
|
return False
|
||||||
else:
|
if existing is not None:
|
||||||
return True
|
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."""
|
"""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))
|
select(EbookSource).where(or_(EbookSource.file_path == str(path), EbookSource.file_sha256 == file_hash))
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
def add_chapter_chunks(
|
def add_chapter_chunks(
|
||||||
session: AsyncSession,
|
session: Session,
|
||||||
source: EbookSource,
|
source: EbookSource,
|
||||||
chapter: EbookChapter,
|
chapter: EbookChapter,
|
||||||
parsed_chapter: ParsedChapter,
|
parsed_chapter: ParsedChapter,
|
||||||
|
|||||||
@@ -22,26 +22,10 @@ def auth_headers(api_key: str) -> dict[str, str]:
|
|||||||
return {"Authorization": f"Bearer {api_key}"}
|
return {"Authorization": f"Bearer {api_key}"}
|
||||||
|
|
||||||
|
|
||||||
async def request_embeddings(
|
def request_embeddings(texts: Sequence[str], config: EbookSearchConfig) -> list[list[float]]:
|
||||||
client: httpx.AsyncClient,
|
"""Request embeddings from the configured OpenAI-compatible endpoint."""
|
||||||
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.
|
|
||||||
"""
|
|
||||||
try:
|
try:
|
||||||
response = await client.post(
|
response = httpx.post(
|
||||||
f"{config.embedding_base_url.rstrip('/')}/embeddings",
|
f"{config.embedding_base_url.rstrip('/')}/embeddings",
|
||||||
headers=auth_headers(config.embedding_api_key),
|
headers=auth_headers(config.embedding_api_key),
|
||||||
json={"model": config.embedding_model, "input": list(texts)},
|
json={"model": config.embedding_model, "input": list(texts)},
|
||||||
@@ -60,15 +44,10 @@ async def request_embeddings(
|
|||||||
raise RuntimeError(msg) from error
|
raise RuntimeError(msg) from error
|
||||||
|
|
||||||
|
|
||||||
async def check_embedding_endpoint(
|
def check_embedding_endpoint(config: EbookSearchConfig, *, timeout_seconds: float = 5.0) -> bool:
|
||||||
client: httpx.AsyncClient,
|
|
||||||
config: EbookSearchConfig,
|
|
||||||
*,
|
|
||||||
timeout_seconds: float = 5.0,
|
|
||||||
) -> bool:
|
|
||||||
"""Return whether the configured embedding endpoint answers a model listing."""
|
"""Return whether the configured embedding endpoint answers a model listing."""
|
||||||
try:
|
try:
|
||||||
response = await client.get(
|
response = httpx.get(
|
||||||
f"{config.embedding_base_url.rstrip('/')}/models",
|
f"{config.embedding_base_url.rstrip('/')}/models",
|
||||||
headers=auth_headers(config.embedding_api_key),
|
headers=auth_headers(config.embedding_api_key),
|
||||||
timeout=timeout_seconds,
|
timeout=timeout_seconds,
|
||||||
@@ -80,15 +59,10 @@ async def check_embedding_endpoint(
|
|||||||
return True
|
return True
|
||||||
|
|
||||||
|
|
||||||
async def check_chat_endpoint(
|
def check_chat_endpoint(config: EbookSearchConfig, *, timeout_seconds: float = 5.0) -> bool:
|
||||||
client: httpx.AsyncClient,
|
|
||||||
config: EbookSearchConfig,
|
|
||||||
*,
|
|
||||||
timeout_seconds: float = 5.0,
|
|
||||||
) -> bool:
|
|
||||||
"""Return whether the configured chat (answering) endpoint answers a model listing."""
|
"""Return whether the configured chat (answering) endpoint answers a model listing."""
|
||||||
try:
|
try:
|
||||||
response = await client.get(
|
response = httpx.get(
|
||||||
f"{config.vllm_base_url.rstrip('/')}/models",
|
f"{config.vllm_base_url.rstrip('/')}/models",
|
||||||
headers=auth_headers(config.vllm_api_key),
|
headers=auth_headers(config.vllm_api_key),
|
||||||
timeout=timeout_seconds,
|
timeout=timeout_seconds,
|
||||||
@@ -124,29 +98,18 @@ def embedding_vectors_from_response(body: object) -> list[list[float]]:
|
|||||||
return vectors
|
return vectors
|
||||||
|
|
||||||
|
|
||||||
async def request_rerank(
|
def request_rerank(
|
||||||
client: httpx.AsyncClient,
|
|
||||||
query: str,
|
query: str,
|
||||||
documents: Sequence[str],
|
documents: Sequence[str],
|
||||||
config: RerankConfig,
|
config: RerankConfig,
|
||||||
) -> object | None:
|
) -> object | None:
|
||||||
"""Request rerank scores from the configured vLLM endpoint.
|
"""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.
|
|
||||||
"""
|
|
||||||
payload = {
|
payload = {
|
||||||
"model": config.model,
|
"model": config.model,
|
||||||
"query": query,
|
"query": query,
|
||||||
"documents": list(documents),
|
"documents": list(documents),
|
||||||
}
|
}
|
||||||
response = await client.post(
|
response = httpx.post(
|
||||||
f"{config.base_url.rstrip('/')}/rerank",
|
f"{config.base_url.rstrip('/')}/rerank",
|
||||||
json=payload,
|
json=payload,
|
||||||
timeout=config.timeout_seconds,
|
timeout=config.timeout_seconds,
|
||||||
@@ -159,26 +122,13 @@ async def request_rerank(
|
|||||||
return None
|
return None
|
||||||
|
|
||||||
|
|
||||||
async def request_chat_completion(
|
def request_chat_completion(
|
||||||
client: httpx.AsyncClient,
|
|
||||||
config: EbookSearchConfig,
|
config: EbookSearchConfig,
|
||||||
messages: Sequence[dict[str, str]],
|
messages: Sequence[dict[str, str]],
|
||||||
) -> str:
|
) -> str:
|
||||||
"""Request a chat completion over a shared async client.
|
"""Request a chat completion from the configured OpenAI-compatible endpoint."""
|
||||||
|
|
||||||
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.
|
|
||||||
"""
|
|
||||||
try:
|
try:
|
||||||
response = await client.post(
|
response = httpx.post(
|
||||||
f"{config.vllm_base_url.rstrip('/')}/chat/completions",
|
f"{config.vllm_base_url.rstrip('/')}/chat/completions",
|
||||||
headers=auth_headers(config.vllm_api_key),
|
headers=auth_headers(config.vllm_api_key),
|
||||||
json={
|
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 " "
|
|
||||||
@@ -9,8 +9,6 @@ from typing import TYPE_CHECKING
|
|||||||
from python.ebook_search.llm_interface import request_rerank
|
from python.ebook_search.llm_interface import request_rerank
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
import httpx
|
|
||||||
|
|
||||||
from python.ebook_search.config import RerankConfig
|
from python.ebook_search.config import RerankConfig
|
||||||
from python.ebook_search.search import SearchResult
|
from python.ebook_search.search import SearchResult
|
||||||
|
|
||||||
@@ -25,12 +23,7 @@ class RerankResult:
|
|||||||
score: float
|
score: float
|
||||||
|
|
||||||
|
|
||||||
async def rerank_chunks(
|
def rerank_chunks(query: str, candidates: list[SearchResult], config: RerankConfig) -> list[SearchResult]:
|
||||||
client: httpx.AsyncClient,
|
|
||||||
query: str,
|
|
||||||
candidates: list[SearchResult],
|
|
||||||
config: RerankConfig,
|
|
||||||
) -> list[SearchResult]:
|
|
||||||
"""Rerank candidates with a vLLM rerank endpoint."""
|
"""Rerank candidates with a vLLM rerank endpoint."""
|
||||||
if not candidates:
|
if not candidates:
|
||||||
return []
|
return []
|
||||||
@@ -41,7 +34,7 @@ async def rerank_chunks(
|
|||||||
config.model,
|
config.model,
|
||||||
len(candidates),
|
len(candidates),
|
||||||
)
|
)
|
||||||
scores = await score_candidates(client, query, candidates, config)
|
scores = score_candidates(query, candidates, config)
|
||||||
results = sorted(
|
results = sorted(
|
||||||
(
|
(
|
||||||
replace(
|
replace(
|
||||||
@@ -63,14 +56,13 @@ async def rerank_chunks(
|
|||||||
return results
|
return results
|
||||||
|
|
||||||
|
|
||||||
async def score_candidates(
|
def score_candidates(
|
||||||
client: httpx.AsyncClient,
|
|
||||||
query: str,
|
query: str,
|
||||||
candidates: list[SearchResult],
|
candidates: list[SearchResult],
|
||||||
config: RerankConfig,
|
config: RerankConfig,
|
||||||
) -> dict[int, RerankResult]:
|
) -> dict[int, RerankResult]:
|
||||||
"""Score candidate chunks with the configured rerank API."""
|
"""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:
|
if body is None:
|
||||||
return zero_rerank_scores(candidates)
|
return zero_rerank_scores(candidates)
|
||||||
|
|
||||||
|
|||||||
+44
-158
@@ -2,17 +2,16 @@
|
|||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import asyncio
|
|
||||||
import logging
|
import logging
|
||||||
import re
|
import re
|
||||||
from collections import defaultdict
|
from collections import defaultdict
|
||||||
|
from concurrent.futures import ThreadPoolExecutor
|
||||||
from dataclasses import dataclass, replace
|
from dataclasses import dataclass, replace
|
||||||
from typing import TYPE_CHECKING
|
from typing import TYPE_CHECKING
|
||||||
|
|
||||||
from pgvector.sqlalchemy import Vector
|
from pgvector.sqlalchemy import Vector
|
||||||
from sqlalchemy import literal, select
|
from sqlalchemy import literal, select
|
||||||
from sqlalchemy.exc import SQLAlchemyError
|
from sqlalchemy.orm import Session
|
||||||
from sqlalchemy.ext.asyncio import AsyncSession
|
|
||||||
|
|
||||||
from python.ebook_search.bm25_corpus import (
|
from python.ebook_search.bm25_corpus import (
|
||||||
BM25CorpusUnavailableError,
|
BM25CorpusUnavailableError,
|
||||||
@@ -20,12 +19,8 @@ from python.ebook_search.bm25_corpus import (
|
|||||||
score_bm25_corpus,
|
score_bm25_corpus,
|
||||||
)
|
)
|
||||||
from python.ebook_search.embeddings import MODEL_DIMENSIONS, embed_query, get_embedding_table
|
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.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 (
|
from python.orm.richie import (
|
||||||
EbookChapter,
|
EbookChapter,
|
||||||
EbookChunk,
|
EbookChunk,
|
||||||
@@ -34,13 +29,11 @@ from python.orm.richie import (
|
|||||||
)
|
)
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from collections.abc import Mapping, Sequence
|
from collections.abc import Mapping
|
||||||
|
|
||||||
import httpx
|
from sqlalchemy.engine import Engine
|
||||||
from sqlalchemy.ext.asyncio import AsyncEngine
|
|
||||||
|
|
||||||
from python.ebook_search.config import EbookSearchConfig
|
from python.ebook_search.config import EbookSearchConfig
|
||||||
from python.ebook_search.protected_phrases.models import HydratedPhraseMatch
|
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
@@ -52,14 +45,11 @@ class SearchResult:
|
|||||||
chunk_id: int
|
chunk_id: int
|
||||||
text: str
|
text: str
|
||||||
source_title: str
|
source_title: str
|
||||||
source_id: int | None = None
|
|
||||||
score: float = 0.0
|
score: float = 0.0
|
||||||
vector_score: float | None = None
|
vector_score: float | None = None
|
||||||
bm25_score: float | None = None
|
bm25_score: float | None = None
|
||||||
fused_score: float | None = None
|
fused_score: float | None = None
|
||||||
rerank_score: float | None = None
|
rerank_score: float | None = None
|
||||||
phrase_hit_count: int = 0
|
|
||||||
matched_phrases: tuple[str, ...] = ()
|
|
||||||
source_author: str | None = None
|
source_author: str | None = None
|
||||||
chapter_title: str | None = None
|
chapter_title: str | None = None
|
||||||
page_label: str | None = None
|
page_label: str | None = None
|
||||||
@@ -74,7 +64,6 @@ class SearchResponse:
|
|||||||
results: list[SearchResult]
|
results: list[SearchResult]
|
||||||
rank_label: str
|
rank_label: str
|
||||||
timings: tuple[RuntimeStep, ...] = ()
|
timings: tuple[RuntimeStep, ...] = ()
|
||||||
phrase_matches: tuple[HydratedPhraseMatch, ...] = ()
|
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def total_runtime_ms(self) -> float:
|
def total_runtime_ms(self) -> float:
|
||||||
@@ -91,38 +80,26 @@ class RetrievalResponse:
|
|||||||
timings: tuple[RuntimeStep, ...]
|
timings: tuple[RuntimeStep, ...]
|
||||||
|
|
||||||
|
|
||||||
async def search_ebooks(
|
def search_ebooks(
|
||||||
engine: AsyncEngine,
|
engine: Engine,
|
||||||
client: httpx.AsyncClient,
|
|
||||||
query: str,
|
query: str,
|
||||||
config: EbookSearchConfig,
|
config: EbookSearchConfig,
|
||||||
*,
|
*,
|
||||||
rerank: bool = False,
|
rerank: bool = False,
|
||||||
phrase_matching: bool | None = None,
|
|
||||||
) -> SearchResponse:
|
) -> SearchResponse:
|
||||||
"""Run hybrid vector/BM25 search and optional reranking."""
|
"""Run hybrid vector/BM25 search and optional reranking."""
|
||||||
if not query.strip():
|
if not query.strip():
|
||||||
logger.info("ebook_search_empty_query")
|
logger.info("ebook_search_empty_query")
|
||||||
return SearchResponse(query=query, results=[], rank_label="Hybrid")
|
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", len(query), rerank)
|
||||||
logger.info(
|
|
||||||
"ebook_search_start query_length=%s rerank=%s phrase_matching=%s",
|
|
||||||
len(query),
|
|
||||||
rerank,
|
|
||||||
phrase_matching_enabled,
|
|
||||||
)
|
|
||||||
timings: list[RuntimeStep] = []
|
timings: list[RuntimeStep] = []
|
||||||
if phrase_matching_enabled:
|
retrieval, timing = timed_result(
|
||||||
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(
|
|
||||||
"Hybrid retrieval",
|
"Hybrid retrieval",
|
||||||
parallel_retrieval(engine, client, query, config),
|
parallel_retrieval,
|
||||||
|
engine,
|
||||||
|
query,
|
||||||
|
config,
|
||||||
)
|
)
|
||||||
timings.extend(retrieval.timings)
|
timings.extend(retrieval.timings)
|
||||||
timings.append(timing)
|
timings.append(timing)
|
||||||
@@ -134,28 +111,18 @@ async def search_ebooks(
|
|||||||
rank_constant=config.rrf_rank_constant,
|
rank_constant=config.rrf_rank_constant,
|
||||||
)
|
)
|
||||||
timings.append(timing)
|
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:
|
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:
|
else:
|
||||||
response, timing = timed_result("Rerank skipped", skip_rerank, query, fused, config)
|
response, timing = timed_result("Rerank skipped", skip_rerank, query, fused, config)
|
||||||
timings.append(timing)
|
timings.append(timing)
|
||||||
response = replace(response, timings=tuple(timings), phrase_matches=tuple(phrase_matches))
|
response = replace(response, timings=tuple(timings))
|
||||||
logger.info(
|
logger.info(
|
||||||
"ebook_search_complete vector_candidates=%s lexical_candidates=%s "
|
"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.vector_results),
|
||||||
len(retrieval.lexical_results),
|
len(retrieval.lexical_results),
|
||||||
len(fused),
|
len(fused),
|
||||||
phrase_matching_enabled,
|
|
||||||
len(phrase_matches),
|
|
||||||
len(response.results),
|
len(response.results),
|
||||||
response.rank_label,
|
response.rank_label,
|
||||||
response.total_runtime_ms,
|
response.total_runtime_ms,
|
||||||
@@ -163,100 +130,30 @@ async def search_ebooks(
|
|||||||
return response
|
return response
|
||||||
|
|
||||||
|
|
||||||
def skip_phrase_matches() -> list[HydratedPhraseMatch]:
|
def parallel_retrieval(
|
||||||
"""Return no protected phrase matches when phrase matching is disabled."""
|
engine: Engine,
|
||||||
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,
|
|
||||||
query: str,
|
query: str,
|
||||||
config: EbookSearchConfig,
|
config: EbookSearchConfig,
|
||||||
) -> RetrievalResponse:
|
) -> RetrievalResponse:
|
||||||
"""Run vector and BM25 candidate retrieval concurrently with separate database sessions.
|
"""Run vector and BM25 candidate retrieval concurrently with separate database sessions."""
|
||||||
|
with ThreadPoolExecutor(max_workers=2, thread_name_prefix="ebook-search") as executor:
|
||||||
BM25 scoring is pure CPU work over the cached corpus, so it runs in a worker thread
|
vector_future = executor.submit(
|
||||||
instead of on the event loop.
|
timed_result,
|
||||||
"""
|
"Embedding + vector search",
|
||||||
(vector_results, vector_timing), (lexical_results, lexical_timing) = await asyncio.gather(
|
vector_candidates,
|
||||||
async_timed_result("Embedding + vector search", vector_candidates(engine, client, query, config)),
|
engine,
|
||||||
async_timed_result("BM25 search", asyncio.to_thread(bm25_candidates, query, config)),
|
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(
|
logger.info(
|
||||||
"ebook_parallel_retrieval_complete vector_candidates=%s lexical_candidates=%s",
|
"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")
|
return SearchResponse(query=query, results=candidates[: config.top_k], rank_label="Hybrid")
|
||||||
|
|
||||||
|
|
||||||
async def apply_rerank(
|
def apply_rerank(
|
||||||
client: httpx.AsyncClient,
|
|
||||||
query: str,
|
query: str,
|
||||||
candidates: list[SearchResult],
|
candidates: list[SearchResult],
|
||||||
config: EbookSearchConfig,
|
config: EbookSearchConfig,
|
||||||
) -> SearchResponse:
|
) -> SearchResponse:
|
||||||
"""Rerank already-fused hybrid candidates."""
|
"""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(
|
logger.info(
|
||||||
"ebook_rerank_complete input_candidates=%s returned=%s",
|
"ebook_rerank_complete input_candidates=%s returned=%s",
|
||||||
min(len(candidates), config.rerank.candidates),
|
min(len(candidates), config.rerank.candidates),
|
||||||
@@ -303,17 +199,10 @@ async def apply_rerank(
|
|||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
async def vector_candidates(
|
def vector_candidates(engine: Engine, query: str, config: EbookSearchConfig) -> list[SearchResult]:
|
||||||
engine: AsyncEngine,
|
|
||||||
client: httpx.AsyncClient,
|
|
||||||
query: str,
|
|
||||||
config: EbookSearchConfig,
|
|
||||||
) -> list[SearchResult]:
|
|
||||||
"""Return pgvector cosine candidates for a natural-language query."""
|
"""Return pgvector cosine candidates for a natural-language query."""
|
||||||
async with AsyncSession(engine) as session:
|
with Session(engine) as session:
|
||||||
model = await session.scalar(
|
model = session.scalar(select(EbookEmbeddingModel).where(EbookEmbeddingModel.name == config.embedding_model))
|
||||||
select(EbookEmbeddingModel).where(EbookEmbeddingModel.name == config.embedding_model)
|
|
||||||
)
|
|
||||||
if model is None:
|
if model is None:
|
||||||
msg = f"Embedding model is not registered: {config.embedding_model}"
|
msg = f"Embedding model is not registered: {config.embedding_model}"
|
||||||
raise ValueError(msg)
|
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}"
|
msg = f"Model row dimension {model.dimension} does not match configured dimension {expected_dimension}"
|
||||||
raise ValueError(msg)
|
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
|
limit = max(config.rerank.candidates, config.top_k) * config.vector_candidate_multiplier
|
||||||
embedding_table = get_embedding_table(model.dimension)
|
embedding_table = get_embedding_table(model.dimension)
|
||||||
|
|
||||||
@@ -334,7 +223,6 @@ async def vector_candidates(
|
|||||||
select(
|
select(
|
||||||
EbookChunk.id.label("chunk_id"),
|
EbookChunk.id.label("chunk_id"),
|
||||||
EbookChunk.text.label("text"),
|
EbookChunk.text.label("text"),
|
||||||
EbookSource.id.label("source_id"),
|
|
||||||
EbookSource.title.label("source_title"),
|
EbookSource.title.label("source_title"),
|
||||||
EbookSource.author.label("source_author"),
|
EbookSource.author.label("source_author"),
|
||||||
EbookChapter.title.label("chapter_title"),
|
EbookChapter.title.label("chapter_title"),
|
||||||
@@ -349,7 +237,7 @@ async def vector_candidates(
|
|||||||
.order_by(distance)
|
.order_by(distance)
|
||||||
.limit(limit)
|
.limit(limit)
|
||||||
)
|
)
|
||||||
rows = (await session.execute(statement)).mappings()
|
rows = session.execute(statement).mappings()
|
||||||
results = [search_result_from_row(row) for row in rows]
|
results = [search_result_from_row(row) for row in rows]
|
||||||
logger.info(
|
logger.info(
|
||||||
"ebook_vector_search_complete model=%s dimension=%s candidates=%s",
|
"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:
|
def search_result_from_row(row: Mapping[str, object]) -> SearchResult:
|
||||||
"""Convert a database row mapping into a search result."""
|
"""Convert a database row mapping into a search result."""
|
||||||
source_id = row.get("source_id")
|
|
||||||
return SearchResult(
|
return SearchResult(
|
||||||
chunk_id=int(row["chunk_id"]),
|
chunk_id=int(row["chunk_id"]),
|
||||||
text=str(row["text"]),
|
text=str(row["text"]),
|
||||||
source_id=int(source_id) if source_id is not None else None,
|
|
||||||
source_title=str(row["source_title"]),
|
source_title=str(row["source_title"]),
|
||||||
source_author=optional_str(row["source_author"]),
|
source_author=optional_str(row["source_author"]),
|
||||||
chapter_title=optional_str(row["chapter_title"]),
|
chapter_title=optional_str(row["chapter_title"]),
|
||||||
|
|||||||
@@ -7,7 +7,7 @@ from time import perf_counter
|
|||||||
from typing import TYPE_CHECKING
|
from typing import TYPE_CHECKING
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from collections.abc import Awaitable, Callable
|
from collections.abc import Callable
|
||||||
|
|
||||||
|
|
||||||
@dataclass(frozen=True)
|
@dataclass(frozen=True)
|
||||||
@@ -34,10 +34,3 @@ def timed_result[T, **P](
|
|||||||
start_seconds = perf_counter()
|
start_seconds = perf_counter()
|
||||||
result = operation(*args, **kwargs)
|
result = operation(*args, **kwargs)
|
||||||
return result, runtime_step_from_start(name, start_seconds)
|
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)
|
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
"""Reusable FastAPI tools."""
|
"""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
|
from python.fastapi_tools.zstd_middleware import ZstdMiddleware
|
||||||
|
|
||||||
__all__ = ["AsyncDbSession", "DbSession", "ZstdMiddleware", "get_async_db", "get_db"]
|
__all__ = ["DbSession", "ZstdMiddleware", "get_db"]
|
||||||
|
|||||||
@@ -5,11 +5,10 @@ from __future__ import annotations
|
|||||||
from typing import TYPE_CHECKING, Annotated
|
from typing import TYPE_CHECKING, Annotated
|
||||||
|
|
||||||
from fastapi import Depends, Request
|
from fastapi import Depends, Request
|
||||||
from sqlalchemy.ext.asyncio import AsyncSession
|
|
||||||
from sqlalchemy.orm import Session
|
from sqlalchemy.orm import Session
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from collections.abc import AsyncIterator, Iterator
|
from collections.abc import Iterator
|
||||||
|
|
||||||
|
|
||||||
def get_db(request: Request) -> Iterator[Session]:
|
def get_db(request: Request) -> Iterator[Session]:
|
||||||
@@ -18,15 +17,4 @@ def get_db(request: Request) -> Iterator[Session]:
|
|||||||
yield 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)]
|
DbSession = Annotated[Session, Depends(get_db)]
|
||||||
AsyncDbSession = Annotated[AsyncSession, Depends(get_async_db)]
|
|
||||||
|
|||||||
+12
-63
@@ -7,7 +7,6 @@ from typing import cast
|
|||||||
|
|
||||||
from sqlalchemy import create_engine
|
from sqlalchemy import create_engine
|
||||||
from sqlalchemy.engine import URL, Engine
|
from sqlalchemy.engine import URL, Engine
|
||||||
from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
|
|
||||||
|
|
||||||
NAMING_CONVENTION = {
|
NAMING_CONVENTION = {
|
||||||
"ix": "ix_%(table_name)s_%(column_0_name)s",
|
"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))
|
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]]:
|
def get_postgres_engine(
|
||||||
"""Build the Postgres connection URL and connect_args from environment variables.
|
*,
|
||||||
|
name: str = "POSTGRES",
|
||||||
|
pool_pre_ping: bool = True,
|
||||||
|
vector_engine: bool = False,
|
||||||
|
) -> Engine:
|
||||||
|
"""Create a SQLAlchemy engine from environment variables.
|
||||||
|
|
||||||
Args:
|
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.
|
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:
|
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)
|
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:
|
if vector_engine:
|
||||||
connect_args["options"] = "-csearch_path=main,public"
|
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(
|
return create_engine(
|
||||||
url=url,
|
url=url,
|
||||||
pool_pre_ping=pool_pre_ping,
|
pool_pre_ping=pool_pre_ping,
|
||||||
pool_recycle=1800,
|
pool_recycle=1800,
|
||||||
connect_args=connect_args,
|
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,
|
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -12,16 +12,12 @@ from python.orm.richie.contact import (
|
|||||||
RelationshipType,
|
RelationshipType,
|
||||||
)
|
)
|
||||||
from python.orm.richie.ebook import (
|
from python.orm.richie.ebook import (
|
||||||
EbookCandidatePhrase,
|
|
||||||
EbookChapter,
|
EbookChapter,
|
||||||
EbookChunk,
|
EbookChunk,
|
||||||
EbookChunkEmbedding1024,
|
EbookChunkEmbedding1024,
|
||||||
EbookChunkEmbedding2560,
|
EbookChunkEmbedding2560,
|
||||||
EbookChunkEmbedding4096,
|
EbookChunkEmbedding4096,
|
||||||
EbookChunkPhraseMention,
|
|
||||||
EbookEmbeddingModel,
|
EbookEmbeddingModel,
|
||||||
EbookPhraseAlias,
|
|
||||||
EbookProtectedPhrase,
|
|
||||||
EbookSource,
|
EbookSource,
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -32,16 +28,12 @@ __all__ = [
|
|||||||
"Contact",
|
"Contact",
|
||||||
"ContactNeed",
|
"ContactNeed",
|
||||||
"ContactRelationship",
|
"ContactRelationship",
|
||||||
"EbookCandidatePhrase",
|
|
||||||
"EbookChapter",
|
"EbookChapter",
|
||||||
"EbookChunk",
|
"EbookChunk",
|
||||||
"EbookChunkEmbedding1024",
|
"EbookChunkEmbedding1024",
|
||||||
"EbookChunkEmbedding2560",
|
"EbookChunkEmbedding2560",
|
||||||
"EbookChunkEmbedding4096",
|
"EbookChunkEmbedding4096",
|
||||||
"EbookChunkPhraseMention",
|
|
||||||
"EbookEmbeddingModel",
|
"EbookEmbeddingModel",
|
||||||
"EbookPhraseAlias",
|
|
||||||
"EbookProtectedPhrase",
|
|
||||||
"EbookSource",
|
"EbookSource",
|
||||||
"Need",
|
"Need",
|
||||||
"RelationshipType",
|
"RelationshipType",
|
||||||
|
|||||||
+2
-108
@@ -5,23 +5,11 @@ from __future__ import annotations
|
|||||||
from datetime import datetime
|
from datetime import datetime
|
||||||
|
|
||||||
from pgvector.sqlalchemy import Vector
|
from pgvector.sqlalchemy import Vector
|
||||||
from sqlalchemy import (
|
from sqlalchemy import BigInteger, Boolean, DateTime, ForeignKey, Index, String, UniqueConstraint
|
||||||
JSON,
|
|
||||||
BigInteger,
|
|
||||||
DateTime,
|
|
||||||
ForeignKey,
|
|
||||||
Index,
|
|
||||||
String,
|
|
||||||
Text,
|
|
||||||
UniqueConstraint,
|
|
||||||
)
|
|
||||||
from sqlalchemy.dialects.postgresql import JSONB
|
|
||||||
from sqlalchemy.orm import Mapped, mapped_column, relationship
|
from sqlalchemy.orm import Mapped, mapped_column, relationship
|
||||||
|
|
||||||
from python.orm.richie.base import TableBase, TableBaseBig
|
from python.orm.richie.base import TableBase, TableBaseBig
|
||||||
|
|
||||||
JSON_DOCUMENT = JSON().with_variant(JSONB, "postgresql")
|
|
||||||
|
|
||||||
|
|
||||||
class EbookSource(TableBase):
|
class EbookSource(TableBase):
|
||||||
"""One indexed EPUB file."""
|
"""One indexed EPUB file."""
|
||||||
@@ -106,7 +94,7 @@ class EbookEmbeddingModel(TableBase):
|
|||||||
|
|
||||||
name: Mapped[str] = mapped_column(String, unique=True)
|
name: Mapped[str] = mapped_column(String, unique=True)
|
||||||
dimension: Mapped[int]
|
dimension: Mapped[int]
|
||||||
is_default: Mapped[bool] = mapped_column(default=False)
|
is_default: Mapped[bool] = mapped_column(Boolean, default=False)
|
||||||
|
|
||||||
|
|
||||||
class EbookChunkEmbedding1024(TableBaseBig):
|
class EbookChunkEmbedding1024(TableBaseBig):
|
||||||
@@ -148,97 +136,3 @@ class EbookChunkEmbedding4096(TableBaseBig):
|
|||||||
chunk_id: Mapped[int] = mapped_column(ForeignKey("main.ebook_chunk.id", ondelete="CASCADE"))
|
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"))
|
model_id: Mapped[int] = mapped_column(ForeignKey("main.ebook_embedding_model.id", ondelete="CASCADE"))
|
||||||
embedding: Mapped[list[float]] = mapped_column(Vector(4096))
|
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]
|
|
||||||
|
|||||||
@@ -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 {
|
default = pkgs.mkShell {
|
||||||
NIX_CONFIG = "extra-experimental-features = nix-command flakes ca-derivations";
|
NIX_CONFIG = "extra-experimental-features = nix-command flakes ca-derivations";
|
||||||
nativeBuildInputs = with pkgs; [
|
nativeBuildInputs = with pkgs; [
|
||||||
|
|||||||
@@ -1,6 +1,9 @@
|
|||||||
{ ... }:
|
{ ... }:
|
||||||
{
|
{
|
||||||
imports = [ ./nix_builder.nix ];
|
imports = [
|
||||||
|
./nix_builder.nix
|
||||||
|
./uv_cache_prune.nix
|
||||||
|
];
|
||||||
|
|
||||||
services.nix_builder.containers = {
|
services.nix_builder.containers = {
|
||||||
nix-builder-00.enable = true;
|
nix-builder-00.enable = true;
|
||||||
|
|||||||
@@ -62,6 +62,11 @@ in
|
|||||||
mountPoint = "/run/secrets/gitea-runners";
|
mountPoint = "/run/secrets/gitea-runners";
|
||||||
isReadOnly = true;
|
isReadOnly = true;
|
||||||
};
|
};
|
||||||
|
uv-cache = {
|
||||||
|
hostPath = vars.uv_cache;
|
||||||
|
mountPoint = "/var/cache/uv";
|
||||||
|
isReadOnly = false;
|
||||||
|
};
|
||||||
};
|
};
|
||||||
config =
|
config =
|
||||||
{
|
{
|
||||||
@@ -153,14 +158,20 @@ in
|
|||||||
}
|
}
|
||||||
) cfg.containers;
|
) cfg.containers;
|
||||||
|
|
||||||
systemd.services = builtins.listToAttrs (
|
systemd = {
|
||||||
map (name: {
|
services = builtins.listToAttrs (
|
||||||
name = "container@${name}";
|
map (name: {
|
||||||
value = {
|
name = "container@${name}";
|
||||||
requires = [ "gitea.service" ];
|
value = {
|
||||||
after = [ "gitea.service" ];
|
requires = [ "gitea.service" ];
|
||||||
};
|
after = [ "gitea.service" ];
|
||||||
}) (builtins.attrNames (filterAttrs (_: c: c.enable) cfg.containers))
|
};
|
||||||
);
|
}) (builtins.attrNames (filterAttrs (_: c: c.enable) cfg.containers))
|
||||||
|
);
|
||||||
|
|
||||||
|
tmpfiles.rules = [
|
||||||
|
"d ${vars.uv_cache} 0755 ${runnerUsername} ${runnerUsername} - -"
|
||||||
|
];
|
||||||
|
};
|
||||||
};
|
};
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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;
|
||||||
|
};
|
||||||
|
};
|
||||||
|
};
|
||||||
|
}
|
||||||
@@ -30,6 +30,7 @@ sudo zfs create media/secure/share -o mountpoint=/zfs/media/share -o exec=off
|
|||||||
# scratch datasets
|
# scratch datasets
|
||||||
sudo zfs create scratch/kafka -o mountpoint=/zfs/scratch/kafka -o recordsize=1M
|
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/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
|
# storage datasets
|
||||||
sudo zfs create storage/ollama -o recordsize=1M -o compression=zstd-19 -o sync=disabled
|
sudo zfs create storage/ollama -o recordsize=1M -o compression=zstd-19 -o sync=disabled
|
||||||
|
|||||||
@@ -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}"
|
|
||||||
];
|
|
||||||
};
|
|
||||||
};
|
|
||||||
}
|
|
||||||
@@ -121,3 +121,9 @@ monthly = 2
|
|||||||
hourly = 0
|
hourly = 0
|
||||||
daily = 0
|
daily = 0
|
||||||
monthly = 0
|
monthly = 0
|
||||||
|
|
||||||
|
["scratch/uv_cache"]
|
||||||
|
15_min = 2
|
||||||
|
hourly = 0
|
||||||
|
daily = 0
|
||||||
|
monthly = 0
|
||||||
|
|||||||
@@ -17,5 +17,6 @@ in
|
|||||||
transmission = "${zfs_storage}/transmission";
|
transmission = "${zfs_storage}/transmission";
|
||||||
ollama = "${zfs_storage}/ollama";
|
ollama = "${zfs_storage}/ollama";
|
||||||
transmission_scratch = "${zfs_scratch}/transmission";
|
transmission_scratch = "${zfs_scratch}/transmission";
|
||||||
|
uv_cache = "${zfs_scratch}/uv_cache";
|
||||||
kafka = "${zfs_scratch}/kafka";
|
kafka = "${zfs_scratch}/kafka";
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -10,8 +10,8 @@ from types import ModuleType
|
|||||||
from typing import TYPE_CHECKING
|
from typing import TYPE_CHECKING
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
from sqlalchemy import select
|
from sqlalchemy import create_engine, select
|
||||||
from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, create_async_engine
|
from sqlalchemy.orm import sessionmaker
|
||||||
|
|
||||||
from python.ebook_search.answer import answer_query
|
from python.ebook_search.answer import answer_query
|
||||||
from python.ebook_search.bm25_corpus import (
|
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
|
assert fused[0].fused_score == fused[0].score
|
||||||
|
|
||||||
|
|
||||||
async def build_async_engine() -> AsyncEngine:
|
def test_find_existing_source_matches_path_or_hash() -> None:
|
||||||
"""Create an in-memory async engine with the Richie schema."""
|
engine = create_engine("sqlite+pysqlite:///:memory:", future=True)
|
||||||
engine = create_async_engine("sqlite+aiosqlite:///:memory:")
|
RichieBase.metadata.create_all(engine)
|
||||||
async with engine.begin() as connection:
|
with sessionmaker(bind=engine, expire_on_commit=False, future=True)() as session:
|
||||||
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:
|
|
||||||
source = EbookSource(
|
source = EbookSource(
|
||||||
title="Book",
|
title="Book",
|
||||||
author=None,
|
author=None,
|
||||||
@@ -100,15 +93,16 @@ async def test_find_existing_source_matches_path_or_hash() -> None:
|
|||||||
file_size=10,
|
file_size=10,
|
||||||
)
|
)
|
||||||
session.add(source)
|
session.add(source)
|
||||||
await session.commit()
|
session.commit()
|
||||||
|
|
||||||
assert await find_existing_source(session, Path("/old/book.epub"), "b" * 64) == source
|
assert 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("/new/book.epub"), "a" * 64) == source
|
||||||
|
|
||||||
|
|
||||||
async def test_bm25_corpus_uses_existing_search_text_without_duplicate_metadata() -> None:
|
def test_bm25_corpus_uses_existing_search_text_without_duplicate_metadata() -> None:
|
||||||
engine = await build_async_engine()
|
engine = create_engine("sqlite+pysqlite:///:memory:", future=True)
|
||||||
async with AsyncSession(engine, expire_on_commit=False) as session:
|
RichieBase.metadata.create_all(engine)
|
||||||
|
with sessionmaker(bind=engine, expire_on_commit=False, future=True)() as session:
|
||||||
source = EbookSource(
|
source = EbookSource(
|
||||||
title="Book",
|
title="Book",
|
||||||
author="Author",
|
author="Author",
|
||||||
@@ -121,10 +115,10 @@ async def test_bm25_corpus_uses_existing_search_text_without_duplicate_metadata(
|
|||||||
file_size=10,
|
file_size=10,
|
||||||
)
|
)
|
||||||
session.add(source)
|
session.add(source)
|
||||||
await session.flush()
|
session.flush()
|
||||||
chapter = EbookChapter(source_id=source.id, spine_index=0, title="Chapter", href=None)
|
chapter = EbookChapter(source_id=source.id, spine_index=0, title="Chapter", href=None)
|
||||||
session.add(chapter)
|
session.add(chapter)
|
||||||
await session.flush()
|
session.flush()
|
||||||
session.add(
|
session.add(
|
||||||
EbookChunk(
|
EbookChunk(
|
||||||
id=1,
|
id=1,
|
||||||
@@ -139,9 +133,9 @@ async def test_bm25_corpus_uses_existing_search_text_without_duplicate_metadata(
|
|||||||
search_text="Book Author Chapter content",
|
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 texts == ["Book Author Chapter content"]
|
||||||
assert records[0]["chunk_id"] == 1
|
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()
|
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] = []
|
refreshed: list[object] = []
|
||||||
db_updated_at = datetime.now(tz=UTC)
|
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))
|
config = EbookSearchConfig(rerank=RerankConfig(enabled=False))
|
||||||
session = object()
|
session = object()
|
||||||
|
|
||||||
await ensure_bm25_corpus(session, config)
|
ensure_bm25_corpus(session, config)
|
||||||
|
|
||||||
assert refreshed == [(session, config, db_updated_at)]
|
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] = []
|
refreshed: list[object] = []
|
||||||
created_at = datetime(2026, 1, 1, tzinfo=UTC)
|
created_at = datetime(2026, 1, 1, tzinfo=UTC)
|
||||||
db_updated_at = datetime(2026, 1, 2, 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))
|
config = EbookSearchConfig(rerank=RerankConfig(enabled=False))
|
||||||
session = object()
|
session = object()
|
||||||
|
|
||||||
await ensure_bm25_corpus(session, config)
|
ensure_bm25_corpus(session, config)
|
||||||
|
|
||||||
assert refreshed == [(session, config, db_updated_at)]
|
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:
|
def test_ensure_embedding_models_registers_service_names() -> None:
|
||||||
engine = await build_async_engine()
|
engine = create_engine("sqlite+pysqlite:///:memory:", future=True)
|
||||||
async with AsyncSession(engine, expire_on_commit=False) as session:
|
RichieBase.metadata.create_all(engine)
|
||||||
await ensure_embedding_models(session)
|
with sessionmaker(bind=engine, expire_on_commit=False, future=True)() as session:
|
||||||
await session.commit()
|
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] == [
|
assert [(model.name, model.dimension) for model in models] == [
|
||||||
("qwen3-embedding-0.6b", 1024),
|
("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"
|
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})
|
config = load_config().model_copy(update={"answer_enabled": False})
|
||||||
result = SearchResult(chunk_id=1, text="source text", source_title="Book")
|
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
|
assert "Answer generation is disabled" in answer
|
||||||
|
|||||||
@@ -5,7 +5,7 @@ from __future__ import annotations
|
|||||||
from typing import TYPE_CHECKING
|
from typing import TYPE_CHECKING
|
||||||
|
|
||||||
from fastapi.testclient import TestClient
|
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.api.main import create_app
|
||||||
from python.ebook_search.config import EbookSearchConfig, RerankConfig
|
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):
|
def patch_app_runtime(mocker: MockerFixture):
|
||||||
mocker.patch(
|
mocker.patch(
|
||||||
"python.ebook_search.api.main.get_async_postgres_engine",
|
"python.ebook_search.api.main.get_postgres_engine",
|
||||||
side_effect=lambda **_kwargs: create_async_engine("sqlite+aiosqlite:///:memory:"),
|
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)
|
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:
|
def test_low_confidence_skips_answer_generation(mocker: MockerFixture) -> None:
|
||||||
called = False
|
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 rerank
|
||||||
del phrase_matching
|
|
||||||
return SearchResponse(query=query, rank_label="Hybrid", results=make_results(1, vector_score=0.05))
|
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
|
nonlocal called
|
||||||
called = True
|
called = True
|
||||||
return "answer"
|
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 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 rerank
|
||||||
del phrase_matching
|
|
||||||
return SearchResponse(query=query, rank_label="Hybrid", results=make_results(2, vector_score=0.9))
|
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.search_ebooks", side_effect=fake_search_ebooks)
|
||||||
mocker.patch(
|
mocker.patch(
|
||||||
"python.ebook_search.api.routes.search.answer_query",
|
"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)
|
patch_app_runtime(mocker)
|
||||||
app = create_app()
|
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 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 rerank
|
||||||
del phrase_matching
|
|
||||||
return SearchResponse(query=query, rank_label="Hybrid", results=make_results(2, vector_score=0.9))
|
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.search_ebooks", side_effect=fake_search_ebooks)
|
||||||
mocker.patch(
|
mocker.patch(
|
||||||
"python.ebook_search.api.routes.search.answer_query",
|
"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)
|
patch_app_runtime(mocker)
|
||||||
app = create_app()
|
app = create_app()
|
||||||
|
|||||||
@@ -5,7 +5,7 @@ from __future__ import annotations
|
|||||||
from typing import TYPE_CHECKING
|
from typing import TYPE_CHECKING
|
||||||
|
|
||||||
from fastapi.testclient import TestClient
|
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.api.main import create_app
|
||||||
from python.ebook_search.config import EbookSearchConfig, RerankConfig
|
from python.ebook_search.config import EbookSearchConfig, RerankConfig
|
||||||
@@ -18,18 +18,18 @@ if TYPE_CHECKING:
|
|||||||
|
|
||||||
def fake_get_postgres_engine(**_kwargs):
|
def fake_get_postgres_engine(**_kwargs):
|
||||||
"""Return an in-memory engine for route tests."""
|
"""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):
|
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)
|
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"):
|
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_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_embedding_endpoint", side_effect=lambda _config: embedding)
|
||||||
mocker.patch(f"{HEALTH_MODULE}.check_chat_endpoint", side_effect=lambda _client, _config: chat)
|
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)
|
mocker.patch(f"{HEALTH_MODULE}.check_bm25_status", side_effect=lambda _config: bm25)
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -16,14 +16,7 @@ if TYPE_CHECKING:
|
|||||||
from pytest_mock import MockerFixture
|
from pytest_mock import MockerFixture
|
||||||
|
|
||||||
|
|
||||||
def make_async_client(mocker: MockerFixture, fake_post) -> httpx.AsyncClient:
|
def test_answer_query_uses_httpx_chat_completions(mocker: MockerFixture) -> None:
|
||||||
"""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:
|
|
||||||
captured: dict[str, object] = {}
|
captured: dict[str, object] = {}
|
||||||
|
|
||||||
def fake_post(url: str, **kwargs: object) -> httpx.Response:
|
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),
|
request=httpx.Request("POST", url),
|
||||||
)
|
)
|
||||||
|
|
||||||
client = make_async_client(mocker, fake_post)
|
mocker.patch.object(httpx, "post", side_effect=fake_post)
|
||||||
config = EbookSearchConfig(
|
config = EbookSearchConfig(
|
||||||
rerank=RerankConfig(enabled=False),
|
rerank=RerankConfig(enabled=False),
|
||||||
vllm_base_url="https://ollama.com/v1",
|
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",
|
chat_model="deepseek-v4-flash",
|
||||||
)
|
)
|
||||||
|
|
||||||
results = [SearchResult(chunk_id=1, text="source", source_title="Book")]
|
answer = answer_query("question", [SearchResult(chunk_id=1, text="source", source_title="Book")], config)
|
||||||
answer = await answer_query(client, "question", results, config)
|
|
||||||
|
|
||||||
assert answer == "grounded answer"
|
assert answer == "grounded answer"
|
||||||
assert captured["url"] == "https://ollama.com/v1/chat/completions"
|
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"
|
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] = {}
|
captured: dict[str, object] = {}
|
||||||
vector = [0.0] * 1024
|
vector = [0.0] * 1024
|
||||||
|
|
||||||
@@ -69,14 +61,14 @@ async def test_embed_texts_uses_httpx_embeddings(mocker: MockerFixture) -> None:
|
|||||||
request=httpx.Request("POST", url),
|
request=httpx.Request("POST", url),
|
||||||
)
|
)
|
||||||
|
|
||||||
client = make_async_client(mocker, fake_post)
|
mocker.patch.object(httpx, "post", side_effect=fake_post)
|
||||||
config = EbookSearchConfig(
|
config = EbookSearchConfig(
|
||||||
rerank=RerankConfig(enabled=False),
|
rerank=RerankConfig(enabled=False),
|
||||||
embedding_base_url="http://bob:8000/v1",
|
embedding_base_url="http://bob:8000/v1",
|
||||||
embedding_model="qwen3-embedding-0.6b",
|
embedding_model="qwen3-embedding-0.6b",
|
||||||
)
|
)
|
||||||
|
|
||||||
embeddings = await embed_texts(client, ["hello"], config)
|
embeddings = embed_texts(["hello"], config)
|
||||||
|
|
||||||
assert embeddings == [vector]
|
assert embeddings == [vector]
|
||||||
assert captured["url"] == "http://bob:8000/v1/embeddings"
|
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"]}
|
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:
|
def fake_post(url: str, **_kwargs: object) -> httpx.Response:
|
||||||
return httpx.Response(200, json={"data": [{}]}, request=httpx.Request("POST", url))
|
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))
|
config = EbookSearchConfig(rerank=RerankConfig(enabled=False))
|
||||||
|
|
||||||
with pytest.raises(RuntimeError, match="Embedding request failed"):
|
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,
|
|
||||||
)
|
|
||||||
@@ -2,11 +2,10 @@
|
|||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import asyncio
|
|
||||||
from threading import Event
|
from threading import Event
|
||||||
from typing import TYPE_CHECKING
|
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.config import EbookSearchConfig, RerankConfig
|
||||||
from python.ebook_search.search import SearchResult, search_ebooks
|
from python.ebook_search.search import SearchResult, search_ebooks
|
||||||
@@ -15,18 +14,18 @@ if TYPE_CHECKING:
|
|||||||
from pytest_mock import MockerFixture
|
from pytest_mock import MockerFixture
|
||||||
|
|
||||||
|
|
||||||
async def test_search_ebooks_runs_vector_and_bm25_in_parallel(mocker: MockerFixture) -> None:
|
def test_search_ebooks_runs_vector_and_bm25_in_parallel(mocker: MockerFixture) -> None:
|
||||||
engine = create_async_engine("sqlite+aiosqlite:///:memory:")
|
engine = create_engine("sqlite+pysqlite:///:memory:", future=True)
|
||||||
vector_started = Event()
|
vector_started = Event()
|
||||||
bm25_started = Event()
|
bm25_started = Event()
|
||||||
received_engines: list[object] = []
|
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."""
|
"""Return vector candidates after confirming BM25 has started."""
|
||||||
received_engines.append(received_engine)
|
received_engines.append(received_engine)
|
||||||
assert query == "what is parallel"
|
assert query == "what is parallel"
|
||||||
vector_started.set()
|
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)]
|
return [SearchResult(chunk_id=1, text="vector", source_title="Vector", vector_score=0.9)]
|
||||||
|
|
||||||
def fake_bm25_candidates(query, _config):
|
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)
|
mocker.patch("python.ebook_search.search.bm25_candidates", side_effect=fake_bm25_candidates)
|
||||||
config = EbookSearchConfig(rerank=RerankConfig(enabled=False))
|
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}
|
timings = {step.name: step for step in response.timings}
|
||||||
assert [result.chunk_id for result in response.results] == [1, 2]
|
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["BM25 search"].counts_toward_total is False
|
||||||
assert timings["Hybrid retrieval"].counts_toward_total is True
|
assert timings["Hybrid retrieval"].counts_toward_total is True
|
||||||
assert received_engines == [engine]
|
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()
|
|
||||||
|
|||||||
@@ -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:
|
def rerank_response(payload: dict[str, object] | None = None, *, content: bytes | None = None) -> httpx.Response:
|
||||||
return httpx.Response(
|
return httpx.Response(
|
||||||
200,
|
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]
|
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:
|
def fake_post(_url: str, *, json: dict[str, object], timeout: float) -> httpx.Response:
|
||||||
assert timeout == 30
|
assert timeout == 30
|
||||||
assert json == {
|
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 [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 [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]
|
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 = [
|
candidates = [
|
||||||
SearchResult(chunk_id=1, text="strong hybrid", source_title="A", score=1.0),
|
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),
|
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 [result.chunk_id for result in results] == [1, 2]
|
||||||
assert results[0].score == pytest.approx(0.79)
|
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
|
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(
|
def fake_rerank_chunks(
|
||||||
_client: httpx.AsyncClient,
|
|
||||||
_query: str,
|
_query: str,
|
||||||
_candidates: list[SearchResult],
|
_candidates: list[SearchResult],
|
||||||
_config: RerankConfig,
|
_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)
|
config = EbookSearchConfig(rerank=RerankConfig(enabled=True), top_k=2)
|
||||||
|
|
||||||
with pytest.raises(httpx.TimeoutException, match="timeout"):
|
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:
|
def fake_post(_url: str, **_kwargs: object) -> httpx.Response:
|
||||||
return rerank_response(content=b"not-json")
|
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
|
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:
|
def fake_post(_url: str, **_kwargs: object) -> httpx.Response:
|
||||||
return rerank_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}
|
assert {result.chunk_id: result.rerank_score for result in results} == {1: 0.0, 2: 1.0}
|
||||||
|
|||||||
+20
-369
@@ -2,49 +2,32 @@
|
|||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import asyncio
|
|
||||||
from compression import zstd
|
from compression import zstd
|
||||||
from datetime import UTC, datetime
|
|
||||||
from typing import TYPE_CHECKING
|
from typing import TYPE_CHECKING
|
||||||
|
|
||||||
from fastapi import BackgroundTasks
|
|
||||||
from fastapi.testclient import TestClient
|
from fastapi.testclient import TestClient
|
||||||
from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
|
from sqlalchemy import create_engine
|
||||||
from sqlalchemy.pool import StaticPool
|
|
||||||
|
|
||||||
from python.ebook_search.api.bm25_tasks import refresh_bm25_for_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.api.main import create_app
|
||||||
from python.ebook_search.config import EbookSearchConfig, RerankConfig
|
from python.ebook_search.config import EbookSearchConfig, RerankConfig
|
||||||
from python.ebook_search.embeddings import EmbeddingModelStats
|
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.search import SearchResponse, SearchResult
|
||||||
from python.ebook_search.timing import RuntimeStep
|
from python.ebook_search.timing import RuntimeStep
|
||||||
from python.orm.richie import EbookSource, RichieBase
|
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from pytest_mock import MockerFixture
|
from pytest_mock import MockerFixture
|
||||||
from sqlalchemy.ext.asyncio import AsyncEngine
|
|
||||||
|
|
||||||
|
|
||||||
def patch_app_runtime(mocker: MockerFixture):
|
def patch_app_runtime(mocker: MockerFixture):
|
||||||
"""Patch app startup dependencies used by UI route tests."""
|
"""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)
|
mocker.patch("python.ebook_search.api.main.ensure_bm25_corpus", side_effect=lambda _session, _config: None)
|
||||||
|
|
||||||
|
|
||||||
def fake_get_postgres_engine(**_kwargs):
|
def fake_get_postgres_engine(**_kwargs):
|
||||||
"""Return an in-memory engine for route tests."""
|
"""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:
|
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)
|
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] = {}
|
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["query"] = query
|
||||||
captured["rerank"] = rerank
|
captured["rerank"] = rerank
|
||||||
captured["phrase_matching"] = phrase_matching
|
|
||||||
captured["config"] = config
|
captured["config"] = config
|
||||||
return SearchResponse(query=query, results=[], rank_label="Hybrid + rerank")
|
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.search_ebooks", side_effect=fake_search_ebooks)
|
||||||
mocker.patch(
|
mocker.patch(
|
||||||
"python.ebook_search.api.routes.search.answer_query",
|
"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)
|
patch_app_runtime(mocker)
|
||||||
app = create_app()
|
app = create_app()
|
||||||
app.state.config = EbookSearchConfig(rerank=RerankConfig(enabled=False), top_k=12, answer_enabled=True)
|
app.state.config = EbookSearchConfig(rerank=RerankConfig(enabled=False), top_k=12, answer_enabled=True)
|
||||||
|
|
||||||
with TestClient(app) as client:
|
with TestClient(app) as client:
|
||||||
response = client.post(
|
response = client.post("/search", data={"query": "where is the quote?", "rerank": "true"})
|
||||||
"/search",
|
|
||||||
data={"query": "where is the quote?", "rerank": "true", "phrase_matching": "true"},
|
|
||||||
)
|
|
||||||
|
|
||||||
assert response.status_code == 200
|
assert response.status_code == 200
|
||||||
assert "Hybrid + rerank" in response.text
|
assert "Hybrid + rerank" in response.text
|
||||||
assert captured["query"] == "where is the quote?"
|
assert captured["query"] == "where is the quote?"
|
||||||
assert captured["rerank"] is True
|
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 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 rerank
|
||||||
del phrase_matching
|
|
||||||
msg = "search exploded"
|
msg = "search exploded"
|
||||||
raise RuntimeError(msg)
|
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 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 rerank
|
||||||
del phrase_matching
|
|
||||||
return SearchResponse(query=query, results=[], rank_label="Hybrid")
|
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"
|
msg = "answer exploded"
|
||||||
raise RuntimeError(msg)
|
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:
|
def test_ui_skips_answer_when_disabled(mocker: MockerFixture) -> None:
|
||||||
called = False
|
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 rerank
|
||||||
del phrase_matching
|
|
||||||
return SearchResponse(query=query, results=[], rank_label="Hybrid")
|
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
|
nonlocal called
|
||||||
called = True
|
called = True
|
||||||
return "answer"
|
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 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 rerank
|
||||||
del phrase_matching
|
|
||||||
return SearchResponse(
|
return SearchResponse(
|
||||||
query=query,
|
query=query,
|
||||||
rank_label="Hybrid + rerank",
|
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.search_ebooks", side_effect=fake_search_ebooks)
|
||||||
mocker.patch(
|
mocker.patch(
|
||||||
"python.ebook_search.api.routes.search.answer_query",
|
"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)
|
patch_app_runtime(mocker)
|
||||||
app = create_app()
|
app = create_app()
|
||||||
@@ -228,47 +176,9 @@ def test_ui_shows_component_scores(mocker: MockerFixture) -> None:
|
|||||||
assert "RRF" in response.text
|
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's syndrome" in response.text
|
|
||||||
|
|
||||||
|
|
||||||
def test_ui_shows_search_runtime_chart(mocker: MockerFixture) -> None:
|
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 rerank
|
||||||
del phrase_matching
|
|
||||||
return SearchResponse(
|
return SearchResponse(
|
||||||
query=query,
|
query=query,
|
||||||
rank_label="Hybrid",
|
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.search_ebooks", side_effect=fake_search_ebooks)
|
||||||
mocker.patch(
|
mocker.patch(
|
||||||
"python.ebook_search.api.routes.search.answer_query",
|
"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)
|
patch_app_runtime(mocker)
|
||||||
app = create_app()
|
app = create_app()
|
||||||
@@ -304,7 +214,7 @@ def test_ui_embed_all_batches_until_complete(mocker: MockerFixture) -> None:
|
|||||||
counts = iter([32, 32, 5, 0])
|
counts = iter([32, 32, 5, 0])
|
||||||
batch_sizes: list[int] = []
|
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)
|
batch_sizes.append(config.embedding_batch_size)
|
||||||
return next(counts)
|
return next(counts)
|
||||||
|
|
||||||
@@ -346,7 +256,7 @@ def test_ui_scan_schedules_bm25_refresh_after_database_change(mocker: MockerFixt
|
|||||||
assert scheduled is True
|
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] = []
|
refreshed: list[object] = []
|
||||||
cache_cleared = False
|
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.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)
|
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))
|
config = EbookSearchConfig(rerank=RerankConfig(enabled=False))
|
||||||
|
|
||||||
await refresh_bm25_for_engine(engine, config)
|
refresh_bm25_for_engine(engine, config)
|
||||||
|
|
||||||
assert len(refreshed) == 1
|
assert len(refreshed) == 1
|
||||||
assert refreshed[0][1] == config
|
assert refreshed[0][1] == config
|
||||||
assert cache_cleared is True
|
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 test_admin_page_shows_embedding_counts_by_model(mocker: MockerFixture) -> None:
|
||||||
def fake_embedding_model_stats(_session):
|
def fake_embedding_model_stats(_session):
|
||||||
return [
|
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.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)
|
patch_app_runtime(mocker)
|
||||||
app = create_app()
|
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 "24" in response.text
|
||||||
assert "qwen3-embedding-4b" in response.text
|
assert "qwen3-embedding-4b" in response.text
|
||||||
assert "2560" 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()
|
|
||||||
|
|||||||
@@ -1,5 +0,0 @@
|
|||||||
{
|
|
||||||
home.sessionPath = [
|
|
||||||
"/home/richie/app_images/"
|
|
||||||
];
|
|
||||||
}
|
|
||||||
@@ -76,7 +76,6 @@
|
|||||||
"cSpell.userWords": [
|
"cSpell.userWords": [
|
||||||
"Cahill",
|
"Cahill",
|
||||||
"Corvidae",
|
"Corvidae",
|
||||||
"dedup",
|
|
||||||
"drivername",
|
"drivername",
|
||||||
"fastapi",
|
"fastapi",
|
||||||
"Michal",
|
"Michal",
|
||||||
|
|||||||
@@ -1,6 +1,5 @@
|
|||||||
{
|
{
|
||||||
imports = [
|
imports = [
|
||||||
../home/app_image_path.nix
|
|
||||||
../home/global.nix
|
../home/global.nix
|
||||||
../home/gui
|
../home/gui
|
||||||
];
|
];
|
||||||
|
|||||||
Reference in New Issue
Block a user