Compare commits

..

54 Commits

Author SHA1 Message Date
Richie a5d7c3be4f fixed fomat issue
treefmt / nix fmt (pull_request) Successful in 5s
pytest / pytest (pull_request) Successful in 28s
build_systems / build-brain (pull_request) Successful in 46s
build_systems / build-bob (pull_request) Successful in 47s
build_systems / build-leviathan (pull_request) Successful in 53s
build_systems / build-rhapsody-in-green (pull_request) Successful in 59s
build_systems / build-jeeves (pull_request) Successful in 2m32s
2026-06-14 15:42:05 -04:00
Richie 2995a75748 fixed test
treefmt / nix fmt (pull_request) Failing after 5s
pytest / pytest (pull_request) Successful in 28s
build_systems / build-bob (pull_request) Successful in 48s
build_systems / build-brain (pull_request) Successful in 53s
build_systems / build-leviathan (pull_request) Successful in 54s
build_systems / build-rhapsody-in-green (pull_request) Successful in 1m1s
build_systems / build-jeeves (pull_request) Successful in 2m43s
2026-06-14 15:41:09 -04:00
Richie dce1838163 opning ports for testing 2026-06-14 15:41:09 -04:00
Richie 121eb979a4 added a index for the VEctor DB 2026-06-14 15:41:09 -04:00
Richie c88315e9b6 improved BM25 write 2026-06-14 15:41:09 -04:00
Richie 50795ab7fc added ZstdMiddleware to ebook_search 2026-06-14 15:41:09 -04:00
Richie e9a80a0308 added vector_engine to fix name postgres name space issue 2026-06-14 15:41:09 -04:00
Richie eb76edb740 reworked ebook_search routers 2026-06-14 15:41:09 -04:00
Richie c53afb3c70 made fastapi tools 2026-06-14 15:41:09 -04:00
Richie 8301db39e5 added proper cache invalidation to load_bm25_corpus 2026-06-14 15:40:31 -04:00
Richie 45a9e90524 updated tests 2026-06-14 15:40:31 -04:00
Richie dd67de3993 improved reranking weights 2026-06-14 15:40:31 -04:00
Richie 6e3635ca01 fixed duplicat enrichment 2026-06-14 15:40:31 -04:00
Richie 11cbe31152 improved queary for vector search 2026-06-14 15:40:31 -04:00
Richie 5544dfa61c add .ebook_search_bm25 to gitignore 2026-06-14 15:40:31 -04:00
Richie 911df63513 updated python 2026-06-14 15:40:31 -04:00
Richie 07eb170b34 setup tests 2026-06-14 15:40:04 -04:00
Richie 6715bbf0a5 build api and frountend 2026-06-14 15:40:04 -04:00
Richie 26ff1f0fd3 added answer.py and config 2026-06-14 15:40:04 -04:00
Richie 666ea97754 added __init__ 2026-06-14 15:40:04 -04:00
Richie e01c687625 made llm_interface.py 2026-06-14 15:40:04 -04:00
Richie 82a367a2b6 added rerank 2026-06-14 15:40:04 -04:00
Richie ad1834c537 built ingest 2026-06-14 15:40:04 -04:00
Richie aed1e14d95 built rag search setup 2026-06-14 15:40:04 -04:00
Richie 0a2d4c08cb set up embedding system 2026-06-14 15:40:04 -04:00
Richie db98bd3559 built BM25 search foundation 2026-06-14 15:40:04 -04:00
Richie cdded5da12 added ebook embedding to orm 2026-06-14 15:40:04 -04:00
Richie d022251a58 removed hedgedoc 2026-06-14 15:40:04 -04:00
Richie e1ef4de6a3 adding embedding Models to jeeves 2026-06-14 15:40:04 -04:00
Richie 5c230a267c ran treefmt
build_systems / build-brain (pull_request) Successful in 46s
build_systems / build-bob (pull_request) Successful in 46s
build_systems / build-leviathan (pull_request) Successful in 54s
build_systems / build-rhapsody-in-green (pull_request) Successful in 58s
build_systems / build-jeeves (pull_request) Successful in 2m33s
treefmt / nix fmt (push) Successful in 4s
build_systems / build-brain (push) Successful in 30s
build_systems / build-bob (push) Successful in 9s
pytest / pytest (push) Successful in 27s
build_systems / build-leviathan (push) Successful in 40s
build_systems / build-rhapsody-in-green (push) Successful in 44s
build_systems / build-jeeves (push) Successful in 2m25s
treefmt / nix fmt (pull_request) Successful in 4s
pytest / pytest (pull_request) Successful in 25s
2026-06-14 13:53:02 -04:00
Richie 5215d66d40 adding noqa to DbSession
build_systems / build-jeeves (pull_request) Successful in 2m38s
treefmt / nix fmt (pull_request) Failing after 5s
pytest / pytest (pull_request) Successful in 26s
build_systems / build-bob (pull_request) Successful in 46s
build_systems / build-leviathan (pull_request) Successful in 53s
build_systems / build-brain (pull_request) Successful in 55s
build_systems / build-rhapsody-in-green (pull_request) Successful in 57s
2026-06-14 13:50:49 -04:00
Richie 7ad198416b added TYPE_CHECKING to contact main.py 2026-06-14 13:50:27 -04:00
Richie 1461c2552a removed van-inventory from pyproject.toml 2026-06-14 13:48:17 -04:00
Richie 736717c2f8 added TYPE_CHECKING to middleware.py 2026-06-14 13:48:02 -04:00
Richie ab2521867e added TYPE_CHECKING to dependencies.py 2026-06-14 13:47:50 -04:00
Richie 8e0ab4190b added TYPE_CHECKING to heater main.py 2026-06-14 13:46:40 -04:00
Richie 734fd7641e moved fetch_weather to masked lat lon and 2026-06-14 13:46:03 -04:00
Richie e898e08c48 added noqa to validate system 2026-06-14 13:41:50 -04:00
Richie d916ea903c removed dead code 2026-06-14 13:39:55 -04:00
Richie d8e916dbe6 fixed type bug in get_snapshots 2026-06-14 13:38:14 -04:00
Richie 48e9f0199d deleting van_inventory 2026-06-14 13:37:55 -04:00
Richie a526420c8d fixed un needed noqa's 2026-06-14 12:58:50 -04:00
Richie 41e3e265af cleaned up audiobook.py mapped_column 2026-06-14 12:53:54 -04:00
Richie 38a17f6146 cleaned uo python dependencies 2026-06-14 12:49:22 -04:00
Richie fe48d4c1ad removing splendor 2026-06-14 12:00:39 -04:00
Richie 9290cb46ee updated series_index to float and added UniqueConstraint to audiobook and audiobook_author
treefmt / nix fmt (push) Successful in 5s
build_systems / build-bob (push) Successful in 32s
build_systems / build-leviathan (push) Successful in 41s
build_systems / build-rhapsody-in-green (push) Successful in 44s
pytest / pytest (push) Successful in 27s
build_systems / build-brain (push) Successful in 31s
build_systems / build-jeeves (push) Successful in 2m28s
pytest / pytest (pull_request) Successful in 26s
build_systems / build-bob (pull_request) Successful in 47s
treefmt / nix fmt (pull_request) Successful in 5s
build_systems / build-brain (pull_request) Successful in 44s
build_systems / build-leviathan (pull_request) Successful in 52s
build_systems / build-rhapsody-in-green (pull_request) Successful in 58s
build_systems / build-jeeves (pull_request) Successful in 2m27s
2026-06-13 22:29:56 -04:00
Richie acd3f2d3ac fixed omnibus for audio books 2026-06-13 22:29:56 -04:00
Richie 08e716f66a deleted frontend dir 2026-06-13 22:29:56 -04:00
Richie d197731af4 added llm_tool_calling.py 2026-06-13 22:29:56 -04:00
Richie 1ffc48bb02 built workflow 2026-06-13 22:29:56 -04:00
Richie b6395ef18f Add catalog.py for manually adding authors and series to the database. 2026-06-13 22:29:56 -04:00
Richie aff6f4e1bd adding audiobook data to DB 2026-06-13 22:29:56 -04:00
Richie a9a96db944 cleaned up old_installer.py
treefmt / nix fmt (pull_request) Successful in 5s
pytest / pytest (pull_request) Successful in 25s
pytest / pytest (push) Successful in 24s
build_systems / build-rhapsody-in-green (pull_request) Successful in 1m5s
treefmt / nix fmt (push) Successful in 4s
build_systems / build-brain (pull_request) Successful in 49s
build_systems / build-bob (pull_request) Successful in 50s
build_systems / build-brain (push) Successful in 44s
build_systems / build-bob (push) Successful in 45s
build_systems / build-leviathan (pull_request) Successful in 57s
build_systems / build-leviathan (push) Successful in 53s
build_systems / build-rhapsody-in-green (push) Successful in 55s
build_systems / build-jeeves (pull_request) Successful in 2m41s
build_systems / build-jeeves (push) Successful in 2m37s
2026-06-13 22:27:11 -04:00
Richie d34154541d moved installer.py to old_installer.py 2026-06-13 22:20:58 -04:00
93 changed files with 4099 additions and 3150 deletions
+1
View File
@@ -172,3 +172,4 @@ frontend/node_modules/
# data from testing llms # data from testing llms
data/* data/*
.ebook_search_bm25
+28 -2
View File
@@ -17,15 +17,41 @@
python-env = final: _prev: { python-env = final: _prev: {
my_python = final.python314.withPackages ( my_python = final.python314.withPackages (
ps: with ps; [ ps:
let
bm25s = ps.buildPythonPackage rec {
pname = "bm25s";
version = "0.3.9";
pyproject = true;
src = final.fetchPypi {
inherit pname version;
hash = "sha256-iVxnnZUrfeg1XttfPhpiCh4vKU0dQrkZvwghzOLi9Zc=";
};
build-system = [ ps.setuptools ];
dependencies = with ps; [
numpy
scipy
];
pythonImportsCheck = [ "bm25s" ];
};
in
with ps;
[
alembic alembic
apprise apprise
apscheduler apscheduler
beautifulsoup4
ebooklib
fastapi fastapi
fastapi-cli fastapi-cli
httpx httpx
mypy mypy
numpy
orjson orjson
pgvector
polars polars
psycopg psycopg
pydantic pydantic
@@ -38,7 +64,7 @@
ruff ruff
scalene scalene
sqlalchemy sqlalchemy
sqlalchemy bm25s
tenacity tenacity
textual textual
tiktoken tiktoken
+7 -5
View File
@@ -3,7 +3,7 @@ name = "system_tools"
version = "0.1.0" version = "0.1.0"
description = "" description = ""
authors = [{ name = "Richie Cahill", email = "richie@tmmworkshop.com" }] authors = [{ name = "Richie Cahill", email = "richie@tmmworkshop.com" }]
requires-python = "~=3.13.0" requires-python = "~=3.14.0"
readme = "README.md" readme = "README.md"
license = "MIT" license = "MIT"
# these dependencies are a best effort and aren't guaranteed to work # these dependencies are a best effort and aren't guaranteed to work
@@ -12,20 +12,22 @@ dependencies = [
"alembic", "alembic",
"apprise", "apprise",
"apscheduler", "apscheduler",
"fastapi",
"fastapi-cli",
"httpx", "httpx",
"python-multipart",
"polars", "polars",
"psycopg[binary]", "psycopg[binary]",
"pydantic", "pydantic",
"pyyaml", "python-multipart",
"sqlalchemy", "sqlalchemy",
"tenacity",
"tinytuya",
"typer", "typer",
"websockets", "websockets",
] ]
[project.scripts] [project.scripts]
database = "python.database_cli:app" database = "python.database_cli:app"
van-inventory = "python.van_inventory.main:serve"
whisper-transcribe = "python.tools.whisper.transcribe:main" whisper-transcribe = "python.tools.whisper.transcribe:main"
[dependency-groups] [dependency-groups]
@@ -41,7 +43,7 @@ dev = [
[tool.ruff] [tool.ruff]
target-version = "py313" target-version = "py314"
line-length = 120 line-length = 120
@@ -0,0 +1,200 @@
"""add ebook search tables.
Revision ID: 2db132cace1a
Revises: b3c60cc5beb5
Create Date: 2026-06-10 22:10:54.379159
"""
from __future__ import annotations
from typing import TYPE_CHECKING
import pgvector
import sqlalchemy as sa
from alembic import op
from python.orm import RichieBase
if TYPE_CHECKING:
from collections.abc import Sequence
# revision identifiers, used by Alembic.
revision: str = "2db132cace1a"
down_revision: str | None = "b3c60cc5beb5"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
schema = RichieBase.schema_name
def upgrade() -> None:
"""Upgrade."""
# ### commands auto generated by Alembic - please adjust! ###
op.create_table(
"ebook_embedding_model",
sa.Column("name", sa.String(), nullable=False),
sa.Column("dimension", sa.Integer(), nullable=False),
sa.Column("is_default", sa.Boolean(), nullable=False),
sa.Column("id", sa.Integer(), nullable=False),
sa.Column("created", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
sa.Column("updated", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
sa.PrimaryKeyConstraint("id", name=op.f("pk_ebook_embedding_model")),
sa.UniqueConstraint("name", name=op.f("uq_ebook_embedding_model_name")),
schema=schema,
)
op.create_table(
"ebook_source",
sa.Column("title", sa.String(), nullable=False),
sa.Column("author", sa.String(), nullable=True),
sa.Column("language", sa.String(), nullable=True),
sa.Column("publisher", sa.String(), nullable=True),
sa.Column("identifier", sa.String(), nullable=True),
sa.Column("file_path", sa.String(), nullable=False),
sa.Column("file_sha256", sa.String(length=64), nullable=False),
sa.Column("file_mtime", sa.DateTime(timezone=True), nullable=False),
sa.Column("file_size", sa.BigInteger(), nullable=False),
sa.Column("id", sa.Integer(), nullable=False),
sa.Column("created", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
sa.Column("updated", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
sa.PrimaryKeyConstraint("id", name=op.f("pk_ebook_source")),
sa.UniqueConstraint("file_path", name=op.f("uq_ebook_source_file_path")),
sa.UniqueConstraint("file_sha256", name=op.f("uq_ebook_source_file_sha256")),
schema=schema,
)
op.create_table(
"ebook_chapter",
sa.Column("source_id", sa.Integer(), nullable=False),
sa.Column("spine_index", sa.Integer(), nullable=False),
sa.Column("title", sa.String(), nullable=True),
sa.Column("href", sa.String(), nullable=True),
sa.Column("id", sa.Integer(), nullable=False),
sa.Column("created", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
sa.Column("updated", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
sa.ForeignKeyConstraint(
["source_id"],
[f"{schema}.ebook_source.id"],
name=op.f("fk_ebook_chapter_source_id_ebook_source"),
ondelete="CASCADE",
),
sa.PrimaryKeyConstraint("id", name=op.f("pk_ebook_chapter")),
sa.UniqueConstraint("source_id", "spine_index", name=op.f("uq_ebook_chapter_source_id")),
schema=schema,
)
op.create_table(
"ebook_chunk",
sa.Column("source_id", sa.Integer(), nullable=False),
sa.Column("chapter_id", sa.Integer(), nullable=True),
sa.Column("chunk_index", sa.Integer(), nullable=False),
sa.Column("text", sa.String(), nullable=False),
sa.Column("token_start", sa.Integer(), nullable=False),
sa.Column("token_count", sa.Integer(), nullable=False),
sa.Column("page_label", sa.String(), nullable=True),
sa.Column("content_sha256", sa.String(length=64), nullable=False),
sa.Column("search_text", sa.String(), nullable=False),
sa.Column("id", sa.BigInteger(), nullable=False),
sa.Column("created", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
sa.Column("updated", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
sa.ForeignKeyConstraint(
["chapter_id"],
[f"{schema}.ebook_chapter.id"],
name=op.f("fk_ebook_chunk_chapter_id_ebook_chapter"),
ondelete="SET NULL",
),
sa.ForeignKeyConstraint(
["source_id"],
[f"{schema}.ebook_source.id"],
name=op.f("fk_ebook_chunk_source_id_ebook_source"),
ondelete="CASCADE",
),
sa.PrimaryKeyConstraint("id", name=op.f("pk_ebook_chunk")),
sa.UniqueConstraint("source_id", "chunk_index", name="uq_ebook_chunk_source_id_chunk_index"),
sa.UniqueConstraint("source_id", "content_sha256", name="uq_ebook_chunk_source_id_content_sha256"),
schema=schema,
)
op.create_table(
"ebook_chunk_embedding_1024",
sa.Column("chunk_id", sa.BigInteger(), nullable=False),
sa.Column("model_id", sa.Integer(), nullable=False),
sa.Column("embedding", pgvector.sqlalchemy.vector.VECTOR(dim=1024), nullable=False),
sa.Column("id", sa.BigInteger(), nullable=False),
sa.Column("created", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
sa.Column("updated", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
sa.ForeignKeyConstraint(
["chunk_id"],
[f"{schema}.ebook_chunk.id"],
name=op.f("fk_ebook_chunk_embedding_1024_chunk_id_ebook_chunk"),
ondelete="CASCADE",
),
sa.ForeignKeyConstraint(
["model_id"],
[f"{schema}.ebook_embedding_model.id"],
name=op.f("fk_ebook_chunk_embedding_1024_model_id_ebook_embedding_model"),
ondelete="CASCADE",
),
sa.PrimaryKeyConstraint("id", name=op.f("pk_ebook_chunk_embedding_1024")),
sa.UniqueConstraint("chunk_id", "model_id", name=op.f("uq_ebook_chunk_embedding_1024_chunk_id")),
schema=schema,
)
op.create_table(
"ebook_chunk_embedding_2560",
sa.Column("chunk_id", sa.BigInteger(), nullable=False),
sa.Column("model_id", sa.Integer(), nullable=False),
sa.Column("embedding", pgvector.sqlalchemy.vector.VECTOR(dim=2560), nullable=False),
sa.Column("id", sa.BigInteger(), nullable=False),
sa.Column("created", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
sa.Column("updated", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
sa.ForeignKeyConstraint(
["chunk_id"],
[f"{schema}.ebook_chunk.id"],
name=op.f("fk_ebook_chunk_embedding_2560_chunk_id_ebook_chunk"),
ondelete="CASCADE",
),
sa.ForeignKeyConstraint(
["model_id"],
[f"{schema}.ebook_embedding_model.id"],
name=op.f("fk_ebook_chunk_embedding_2560_model_id_ebook_embedding_model"),
ondelete="CASCADE",
),
sa.PrimaryKeyConstraint("id", name=op.f("pk_ebook_chunk_embedding_2560")),
sa.UniqueConstraint("chunk_id", "model_id", name=op.f("uq_ebook_chunk_embedding_2560_chunk_id")),
schema=schema,
)
op.create_table(
"ebook_chunk_embedding_4096",
sa.Column("chunk_id", sa.BigInteger(), nullable=False),
sa.Column("model_id", sa.Integer(), nullable=False),
sa.Column("embedding", pgvector.sqlalchemy.vector.VECTOR(dim=4096), nullable=False),
sa.Column("id", sa.BigInteger(), nullable=False),
sa.Column("created", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
sa.Column("updated", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False),
sa.ForeignKeyConstraint(
["chunk_id"],
[f"{schema}.ebook_chunk.id"],
name=op.f("fk_ebook_chunk_embedding_4096_chunk_id_ebook_chunk"),
ondelete="CASCADE",
),
sa.ForeignKeyConstraint(
["model_id"],
[f"{schema}.ebook_embedding_model.id"],
name=op.f("fk_ebook_chunk_embedding_4096_model_id_ebook_embedding_model"),
ondelete="CASCADE",
),
sa.PrimaryKeyConstraint("id", name=op.f("pk_ebook_chunk_embedding_4096")),
sa.UniqueConstraint("chunk_id", "model_id", name=op.f("uq_ebook_chunk_embedding_4096_chunk_id")),
schema=schema,
)
# ### end Alembic commands ###
def downgrade() -> None:
"""Downgrade."""
# ### commands auto generated by Alembic - please adjust! ###
op.drop_table("ebook_chunk_embedding_4096", schema=schema)
op.drop_table("ebook_chunk_embedding_2560", schema=schema)
op.drop_table("ebook_chunk_embedding_1024", schema=schema)
op.drop_table("ebook_chunk", schema=schema)
op.drop_table("ebook_chapter", schema=schema)
op.drop_table("ebook_source", schema=schema)
op.drop_table("ebook_embedding_model", schema=schema)
# ### end Alembic commands ###
@@ -0,0 +1,54 @@
"""add 1024 ebook embedding cosine index.
Revision ID: c460105682d2
Revises: 2db132cace1a
Create Date: 2026-06-13 19:53:45.680289
"""
from __future__ import annotations
from typing import TYPE_CHECKING
from alembic import op
from python.orm import RichieBase
if TYPE_CHECKING:
from collections.abc import Sequence
# revision identifiers, used by Alembic.
revision: str = "c460105682d2"
down_revision: str | None = "2db132cace1a"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
schema = RichieBase.schema_name
def upgrade() -> None:
"""Upgrade."""
# ### commands auto generated by Alembic - please adjust! ###
op.create_index(
"ix_ebook_chunk_embedding_1024_embedding_cosine",
"ebook_chunk_embedding_1024",
["embedding"],
unique=False,
schema=schema,
postgresql_using="hnsw",
postgresql_ops={"embedding": "vector_cosine_ops"},
)
# ### end Alembic commands ###
def downgrade() -> None:
"""Downgrade."""
# ### commands auto generated by Alembic - please adjust! ###
op.drop_index(
"ix_ebook_chunk_embedding_1024_embedding_cosine",
table_name="ebook_chunk_embedding_1024",
schema=schema,
postgresql_using="hnsw",
postgresql_ops={"embedding": "vector_cosine_ops"},
)
# ### end Alembic commands ###
+7 -3
View File
@@ -1,19 +1,23 @@
"""FastAPI interface for Contact database.""" """FastAPI interface for Contact database."""
from __future__ import annotations
import logging import logging
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager from contextlib import asynccontextmanager
from typing import Annotated from typing import TYPE_CHECKING, Annotated
import typer import typer
import uvicorn import uvicorn
from fastapi import FastAPI from fastapi import FastAPI
from python.api.middleware import ZstdMiddleware
from python.api.routers import contact_router, views_router from python.api.routers import contact_router, views_router
from python.common import configure_logger from python.common import configure_logger
from python.fastapi_tools import ZstdMiddleware
from python.orm.common import get_postgres_engine from python.orm.common import get_postgres_engine
if TYPE_CHECKING:
from collections.abc import AsyncIterator
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
+1 -1
View File
@@ -9,7 +9,7 @@ from pydantic import BaseModel
from sqlalchemy import select from sqlalchemy import select
from sqlalchemy.orm import selectinload from sqlalchemy.orm import selectinload
from python.api.dependencies import DbSession 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 from python.orm.richie.contact import Contact, ContactRelationship, Need, RelationshipType
TEMPLATES_DIR = Path(__file__).parent.parent / "templates" TEMPLATES_DIR = Path(__file__).parent.parent / "templates"
+1 -1
View File
@@ -9,7 +9,7 @@ from fastapi.templating import Jinja2Templates
from sqlalchemy import select from sqlalchemy import select
from sqlalchemy.orm import Session, selectinload from sqlalchemy.orm import Session, selectinload
from python.api.dependencies import DbSession 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 from python.orm.richie.contact import Contact, ContactRelationship, Need, RelationshipType
TEMPLATES_DIR = Path(__file__).parent.parent / "templates" TEMPLATES_DIR = Path(__file__).parent.parent / "templates"
+3 -15
View File
@@ -4,12 +4,10 @@ Usage:
database <db_name> <command> [args...] database <db_name> <command> [args...]
Examples: Examples:
database van_inventory upgrade head
database van_inventory downgrade head-1
database van_inventory revision --autogenerate -m "add meals table"
database van_inventory check
database richie check database richie check
database richie upgrade head database richie upgrade head
database richie downgrade head-1
database richie revision --autogenerate -m "add meals table"
""" """
from __future__ import annotations from __future__ import annotations
@@ -48,10 +46,7 @@ class DatabaseConfig:
def alembic_config(self) -> Config: def alembic_config(self) -> Config:
"""Build an alembic Config for this database.""" """Build an alembic Config for this database."""
# Runtime import needed — Config is in TYPE_CHECKING for the return type annotation cfg = Config()
from alembic.config import Config as AlembicConfig # noqa: PLC0415
cfg = AlembicConfig()
cfg.set_main_option("script_location", self.script_location) cfg.set_main_option("script_location", self.script_location)
cfg.set_main_option("file_template", self.file_template) cfg.set_main_option("file_template", self.file_template)
cfg.set_main_option("prepend_sys_path", ".") cfg.set_main_option("prepend_sys_path", ".")
@@ -76,13 +71,6 @@ DATABASES: dict[str, DatabaseConfig] = {
base_class_name="RichieBase", base_class_name="RichieBase",
models_module="python.orm.richie", models_module="python.orm.richie",
), ),
"van_inventory": DatabaseConfig(
env_prefix="VAN_INVENTORY",
version_location="python/alembic/van_inventory/versions",
base_module="python.orm.van_inventory.base",
base_class_name="VanInventoryBase",
models_module="python.orm.van_inventory.models",
),
} }
+1
View File
@@ -0,0 +1 @@
"""EPUB search package."""
+57
View File
@@ -0,0 +1,57 @@
"""Grounded answer generation."""
from __future__ import annotations
import logging
from typing import TYPE_CHECKING
from python.ebook_search.llm_interface import request_chat_completion
if TYPE_CHECKING:
from python.ebook_search.config import EbookSearchConfig
from python.ebook_search.search import SearchResult
logger = logging.getLogger(__name__)
def answer_query(query: str, results: list[SearchResult], config: EbookSearchConfig) -> str:
"""Answer a question using only retrieved chunks."""
if not config.answer_enabled:
logger.info("ebook_answer_skipped_disabled")
return "Answer generation is disabled. Source chunks are shown below."
if not results:
logger.info("ebook_answer_skipped_no_results")
return "No relevant sources were found."
logger.info(
"ebook_answer_request_start base_url=%s model=%s sources=%s query_length=%s",
config.vllm_base_url,
config.chat_model,
len(results),
len(query),
)
context = "\n\n".join(
f"[{index}] {result.source_title}{' - ' + result.chapter_title if result.chapter_title else ''}\n{result.text}"
for index, result in enumerate(results, start=1)
)
content = request_chat_completion(
config,
[
{
"role": "system",
"content": (
"Answer only from the provided context. Cite sources with bracketed numbers like [1]. "
"If the context is insufficient, say so."
),
},
{"role": "user", "content": f"Question:\n{query}\n\nContext:\n{context}"},
],
)
logger.info(
"ebook_answer_request_complete model=%s answer_length=%s",
config.chat_model,
len(content),
)
return content or "The model returned an empty answer."
+1
View File
@@ -0,0 +1 @@
"""Web and external API adapters for EPUB search."""
+60
View File
@@ -0,0 +1,60 @@
"""Background BM25 refresh tasks for the web app."""
from __future__ import annotations
import logging
from threading import Timer
from typing import TYPE_CHECKING
from sqlalchemy.orm import Session
from python.ebook_search.bm25_corpus import load_bm25_corpus, refresh_bm25_corpus
if TYPE_CHECKING:
from fastapi import FastAPI
from sqlalchemy.engine import Engine
from python.ebook_search.config import EbookSearchConfig
logger = logging.getLogger(__name__)
def schedule_bm25_refresh(app: FastAPI) -> None:
"""Schedule a delayed BM25 corpus refresh, replacing any pending refresh."""
existing_timer = getattr(app.state, "bm25_refresh_timer", None)
if existing_timer is not None:
existing_timer.cancel()
timer = Timer(app.state.config.bm25_refresh_delay_seconds, refresh_bm25_for_app, args=(app,))
timer.daemon = True
timer.start()
app.state.bm25_refresh_timer = timer
logger.info(
"ebook_bm25_refresh_scheduled delay_seconds=%s",
app.state.config.bm25_refresh_delay_seconds,
)
def cancel_bm25_refresh(app: FastAPI) -> None:
"""Cancel any pending BM25 corpus refresh."""
existing_timer = getattr(app.state, "bm25_refresh_timer", None)
if existing_timer is not None:
existing_timer.cancel()
app.state.bm25_refresh_timer = None
logger.info("ebook_bm25_refresh_cancelled")
def refresh_bm25_for_app(app: FastAPI) -> None:
"""Refresh the BM25 corpus using the app engine and config."""
try:
refresh_bm25_for_engine(app.state.engine, app.state.config)
except Exception:
logger.exception("ebook_bm25_refresh_failed")
def refresh_bm25_for_engine(engine: Engine, config: EbookSearchConfig) -> None:
"""Refresh the BM25 corpus using a SQLAlchemy engine."""
with Session(engine) as session:
refresh_bm25_corpus(session, config)
load_bm25_corpus.cache_clear()
logger.info("ebook_bm25_corpus_cache_cleared_after_refresh")
+79
View File
@@ -0,0 +1,79 @@
"""FastAPI HTMX app for EPUB search."""
from __future__ import annotations
import logging
from contextlib import asynccontextmanager
from typing import TYPE_CHECKING, Annotated
import typer
import uvicorn
from fastapi import FastAPI
from fastapi.staticfiles import StaticFiles
from sqlalchemy.orm import Session
from python.common import configure_logger
from python.ebook_search.api.bm25_tasks import cancel_bm25_refresh
from python.ebook_search.api.routes import admin_router, page_router, search_router
from python.ebook_search.api.web import STATIC_DIR
from python.ebook_search.bm25_corpus import ensure_bm25_corpus
from python.ebook_search.config import load_config
from python.fastapi_tools import ZstdMiddleware
from python.orm.common import get_postgres_engine
if TYPE_CHECKING:
from collections.abc import AsyncIterator
logger = logging.getLogger(__name__)
@asynccontextmanager
async def lifespan(app: FastAPI) -> AsyncIterator[None]:
"""Manage application startup and shutdown resources."""
logger.info("ebook_search_startup")
app.state.engine = get_postgres_engine(name="RICHIE", vector_engine=True)
with Session(app.state.engine) as session:
ensure_bm25_corpus(session, app.state.config)
try:
yield
finally:
logger.info("ebook_search_shutdown")
cancel_bm25_refresh(app)
app.state.engine.dispose()
def create_app() -> FastAPI:
"""Create the EPUB search web app."""
app = FastAPI(title="EPUB Search", lifespan=lifespan)
app.add_middleware(ZstdMiddleware)
app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")
app.state.config = load_config()
logger.info(
"ebook_search_config_loaded top_k=%s embedding_model=%s rerank_enabled=%s answer_enabled=%s library_paths=%s",
app.state.config.top_k,
app.state.config.embedding_model,
app.state.config.rerank.enabled,
app.state.config.answer_enabled,
len(app.state.config.library_paths),
)
app.include_router(admin_router)
app.include_router(page_router)
app.include_router(search_router)
return app
def serve(
host: Annotated[str, typer.Option("--host", "-h", help="Host to bind to")] = "127.0.0.1",
port: Annotated[int, typer.Option("--port", "-p", help="Port to bind to")] = 8070,
log_level: Annotated[str, typer.Option("--log-level", "-l", help="Log level")] = "INFO",
) -> None:
"""Start the EPUB search server."""
configure_logger(log_level)
uvicorn.run(create_app(), host=host, port=port)
if __name__ == "__main__":
typer.run(serve)
@@ -0,0 +1,11 @@
"""EPUB search web route modules."""
from python.ebook_search.api.routes.admin import router as admin_router
from python.ebook_search.api.routes.page import router as page_router
from python.ebook_search.api.routes.search import router as search_router
__all__ = [
"admin_router",
"page_router",
"search_router",
]
+107
View File
@@ -0,0 +1,107 @@
"""Admin routes for the EPUB search web UI."""
from __future__ import annotations
import logging
from dataclasses import replace
from fastapi import APIRouter, Request
from fastapi.responses import HTMLResponse
from sqlalchemy.orm import Session
from python.ebook_search.api.bm25_tasks import schedule_bm25_refresh
from python.ebook_search.api.web import templates
from python.ebook_search.embeddings import embed_missing_chunks, embedding_model_stats
from python.ebook_search.ingest import ingest_configured_paths
logger = logging.getLogger(__name__)
router = APIRouter(prefix="/admin")
EMBED_ALL_BATCH_SIZE = 32
@router.get("", response_class=HTMLResponse)
def admin(request: Request) -> HTMLResponse:
"""Render the admin page."""
with Session(request.app.state.engine) as session:
stats = embedding_model_stats(session)
logger.info("ebook_admin_page_loaded models=%s", len(stats))
return templates.TemplateResponse(request, "admin.html", {"config": request.app.state.config, "stats": stats})
@router.post("/scan", response_class=HTMLResponse)
def scan_library(request: Request) -> HTMLResponse:
"""Scan configured library paths for EPUB changes."""
try:
with Session(request.app.state.engine) as session:
count = ingest_configured_paths(session, request.app.state.config)
session.commit()
except Exception as error:
logger.exception("ebook_admin_scan_failed")
return templates.TemplateResponse(request, "partials/error.html", {"message": str(error)}, status_code=500)
logger.info("ebook_admin_scan_complete changed_files=%s", count)
if count > 0:
schedule_bm25_refresh(request.app)
return templates.TemplateResponse(request, "partials/admin_status.html", {"message": f"Indexed {count} EPUBs"})
@router.post("/embed-missing", response_class=HTMLResponse)
def embed_missing(request: Request) -> HTMLResponse:
"""Embed chunks missing vectors for the configured model."""
try:
with Session(request.app.state.engine) as session:
count = embed_missing_chunks(session, request.app.state.config)
session.commit()
except Exception as error:
logger.exception("ebook_admin_embed_missing_failed")
return templates.TemplateResponse(request, "partials/error.html", {"message": str(error)}, status_code=500)
logger.info("ebook_admin_embed_missing_complete chunks=%s", count)
return templates.TemplateResponse(
request,
"partials/admin_status.html",
{"message": f"Embedded {count} chunks"},
)
@router.post("/embed-all", response_class=HTMLResponse)
def embed_all(request: Request) -> HTMLResponse:
"""Embed all chunks missing vectors in fixed-size batches."""
total = 0
batches = 0
config = replace(request.app.state.config, embedding_batch_size=EMBED_ALL_BATCH_SIZE)
try:
with Session(request.app.state.engine) as session:
while True:
count = embed_missing_chunks(session, config)
if count == 0:
break
session.commit()
total += count
batches += 1
logger.info(
"ebook_admin_embed_all_batch_complete batch=%s chunks=%s total_chunks=%s",
batches,
count,
total,
)
except Exception as error:
logger.exception(
"ebook_admin_embed_all_failed batches=%s chunks=%s",
batches,
total,
)
return templates.TemplateResponse(
request,
"partials/error.html",
{"message": f"Embed all failed after {total} chunks in {batches} batches: {error}"},
status_code=500,
)
logger.info("ebook_admin_embed_all_complete batches=%s chunks=%s", batches, total)
return templates.TemplateResponse(
request,
"partials/admin_status.html",
{"message": f"Embedded {total} chunks in {batches} batches of {EMBED_ALL_BATCH_SIZE}"},
)
+57
View File
@@ -0,0 +1,57 @@
"""Page routes for the EPUB search web UI."""
from __future__ import annotations
import logging
from fastapi import APIRouter, Request
from fastapi.responses import HTMLResponse
from sqlalchemy import select
from sqlalchemy.orm import Session
from python.ebook_search.api.web import templates
from python.orm.richie import EbookSource
logger = logging.getLogger(__name__)
router = APIRouter()
@router.get("/", response_class=HTMLResponse)
def index(request: Request) -> HTMLResponse:
"""Render the search page."""
return templates.TemplateResponse(request, "search.html", {"config": request.app.state.config})
@router.get("/books", response_class=HTMLResponse)
def books(request: Request) -> HTMLResponse:
"""Render the indexed books page."""
with Session(request.app.state.engine) as session:
sources = list(session.scalars(select(EbookSource).order_by(EbookSource.title)).all())
logger.info("ebook_books_page_loaded count=%s", len(sources))
return templates.TemplateResponse(request, "books.html", {"sources": sources})
@router.get("/books/{source_id}", response_class=HTMLResponse)
def book_detail(source_id: int, request: Request) -> HTMLResponse:
"""Render details for one indexed book."""
with Session(request.app.state.engine) as session:
source = session.get(EbookSource, source_id)
if source is not None:
chapter_count = len(source.chapters)
chunk_count = len(source.chunks)
else:
chapter_count = 0
chunk_count = 0
logger.info(
"ebook_book_detail_loaded source_id=%s found=%s chapters=%s chunks=%s",
source_id,
source is not None,
chapter_count,
chunk_count,
)
return templates.TemplateResponse(
request,
"book_detail.html",
{"chapter_count": chapter_count, "chunk_count": chunk_count, "source": source},
)
+58
View File
@@ -0,0 +1,58 @@
"""Search routes for the EPUB search web UI."""
from __future__ import annotations
import logging
from dataclasses import replace
from time import perf_counter
from typing import Annotated
from fastapi import APIRouter, Form, Request
from fastapi.responses import HTMLResponse
from python.ebook_search.answer import answer_query
from python.ebook_search.api.web import templates
from python.ebook_search.search import search_ebooks
from python.ebook_search.timing import runtime_step_from_start
logger = logging.getLogger(__name__)
router = APIRouter()
@router.post("/search", response_class=HTMLResponse)
def search(
request: Request,
query: Annotated[str, Form()],
rerank: Annotated[str | None, Form()] = None,
) -> HTMLResponse:
"""Run a search and render HTMX results."""
try:
response = search_ebooks(request.app.state.engine, query, request.app.state.config, rerank=rerank == "true")
except Exception as error:
logger.exception("ebook_search_request_failed")
return templates.TemplateResponse(request, "partials/error.html", {"message": str(error)}, status_code=500)
answer_start = perf_counter()
if request.app.state.config.answer_enabled:
try:
answer = answer_query(query, response.results, request.app.state.config)
except RuntimeError as error:
logger.warning("ebook_answer_request_failed_falling_back error=%s", error)
answer = "Answer generation failed. Source chunks are still shown below."
else:
logger.info("ebook_answer_skipped_disabled")
answer = "Answer generation is disabled. Source chunks are shown below."
answer_step_name = "Answer generation" if request.app.state.config.answer_enabled else "Answer skipped"
response = replace(
response,
timings=(*response.timings, runtime_step_from_start(answer_step_name, answer_start)),
)
logger.info(
"ebook_search_request_complete results=%s rank_label=%s runtime_ms=%.1f",
len(response.results),
response.rank_label,
response.total_runtime_ms,
)
return templates.TemplateResponse(request, "partials/results.html", {"answer": answer, "response": response})
+140
View File
@@ -0,0 +1,140 @@
body {
margin: 0;
background: #f7f7f4;
color: #202124;
font-family: system-ui, -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif;
}
main {
max-width: 960px;
margin: 0 auto;
padding: 24px;
}
nav {
display: flex;
gap: 12px;
align-items: center;
margin-bottom: 20px;
}
nav form {
margin: 0;
}
.actions {
display: flex;
flex-wrap: wrap;
gap: 12px;
margin-bottom: 24px;
}
textarea {
display: block;
width: 100%;
margin: 8px 0 12px;
}
button {
padding: 8px 14px;
}
.check {
display: inline-flex;
gap: 8px;
align-items: center;
margin-right: 12px;
}
.rank-label {
margin-top: 24px;
font-weight: 700;
}
.results {
padding-left: 24px;
}
.meta,
.scores,
.status {
color: #626a73;
}
.scores {
display: flex;
flex-wrap: wrap;
gap: 8px;
margin: 12px 0;
}
.scores div {
display: inline-flex;
gap: 4px;
align-items: baseline;
}
.scores dt {
font-weight: 700;
}
.scores dd {
margin: 0;
}
.runtime {
margin-top: 16px;
}
.timing-chart {
display: grid;
gap: 8px;
padding: 0;
list-style: none;
}
.timing-chart li {
display: grid;
grid-template-columns: minmax(150px, 1fr) minmax(160px, 2fr) auto auto;
gap: 8px;
align-items: center;
}
.timing-bar {
height: 10px;
overflow: hidden;
background: #e5e5df;
}
.timing-bar span {
display: block;
height: 100%;
background: #3767c8;
}
.timing-value,
.timing-remaining {
color: #626a73;
font-variant-numeric: tabular-nums;
}
table {
width: 100%;
border-collapse: collapse;
}
th,
td {
padding: 8px;
border-bottom: 1px solid #d8d8d2;
text-align: left;
}
th {
font-weight: 700;
}
.error {
color: #9f1d20;
font-weight: 700;
}
@@ -0,0 +1,57 @@
<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>EPUB Admin</title>
<script src="https://unpkg.com/htmx.org@2.0.4"></script>
<link rel="stylesheet" href="/static/style.css">
</head>
<body>
<main>
<nav>
<a href="/">Search</a>
<a href="/books">Books</a>
<a href="/admin">Admin</a>
</nav>
<h1>Admin</h1>
<section id="admin-status"></section>
<section class="actions">
<form hx-post="/admin/scan" hx-target="#admin-status" hx-swap="innerHTML">
<button type="submit">Scan</button>
</form>
<form hx-post="/admin/embed-missing" hx-target="#admin-status" hx-swap="innerHTML">
<button type="submit">Embed</button>
</form>
<form hx-post="/admin/embed-all" hx-target="#admin-status" hx-swap="innerHTML">
<button type="submit">Embed all</button>
</form>
</section>
<section>
<h2>Embeddings</h2>
<table>
<thead>
<tr>
<th>Model</th>
<th>Dimensions</th>
<th>Embedded</th>
<th>Missing</th>
<th>Total chunks</th>
</tr>
</thead>
<tbody>
{% for item in stats %}
<tr>
<td>{{ item.model_name }}</td>
<td>{{ item.dimension }}</td>
<td>{{ item.embedded_chunks }}</td>
<td>{{ item.missing_chunks }}</td>
<td>{{ item.total_chunks }}</td>
</tr>
{% endfor %}
</tbody>
</table>
</section>
</main>
</body>
</html>
@@ -0,0 +1,32 @@
<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>{% if source %}{{ source.title }}{% else %}Book not found{% endif %}</title>
<link rel="stylesheet" href="/static/style.css">
</head>
<body>
<main>
<nav>
<a href="/">Search</a>
<a href="/books">Books</a>
<a href="/admin">Admin</a>
</nav>
{% if source %}
<h1>{{ source.title }}</h1>
<p class="meta">{{ source.author or "Unknown author" }}</p>
<dl>
<dt>File</dt>
<dd>{{ source.file_path }}</dd>
<dt>Chapters</dt>
<dd>{{ chapter_count }}</dd>
<dt>Chunks</dt>
<dd>{{ chunk_count }}</dd>
</dl>
{% else %}
<h1>Book not found</h1>
{% endif %}
</main>
</body>
</html>
@@ -0,0 +1,31 @@
<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>EPUB Books</title>
<link rel="stylesheet" href="/static/style.css">
</head>
<body>
<main>
<nav>
<a href="/">Search</a>
<a href="/books">Books</a>
<a href="/admin">Admin</a>
</nav>
<h1>Books</h1>
{% if sources %}
<ol class="results">
{% for source in sources %}
<li>
<h2><a href="/books/{{ source.id }}">{{ source.title }}</a></h2>
<p class="meta">{{ source.author or "Unknown author" }}</p>
</li>
{% endfor %}
</ol>
{% else %}
<p>No EPUBs indexed.</p>
{% endif %}
</main>
</body>
</html>
@@ -0,0 +1 @@
<p class="status">{{ message }}</p>
@@ -0,0 +1 @@
<p class="error">{{ message }}</p>
@@ -0,0 +1,74 @@
<div class="rank-label">{{ response.rank_label }}</div>
{% if response.timings %}
<section class="runtime">
<h2>Runtime</h2>
<p class="meta">Total {{ "%.1f"|format(response.total_runtime_ms) }} ms</p>
<ol class="timing-chart">
{% set total = response.total_runtime_ms %}
{% set ns = namespace(remaining=total) %}
{% for step in response.timings %}
{% set width = (step.duration_ms / total * 100) if total else 0 %}
{% if step.counts_toward_total %}
{% set ns.remaining = ns.remaining - step.duration_ms %}
{% endif %}
<li>
<span class="timing-label">{{ step.name }}</span>
<span class="timing-bar"><span style="width: {{ "%.2f"|format(width) }}%"></span></span>
<span class="timing-value">{{ "%.1f"|format(step.duration_ms) }} ms</span>
<span class="timing-remaining">{{ "%.1f"|format([ns.remaining, 0]|max) }} ms left</span>
</li>
{% endfor %}
</ol>
</section>
{% endif %}
<section class="answer">
<h2>Answer</h2>
<p>{{ answer }}</p>
</section>
{% if response.results %}
<ol class="results">
{% for result in response.results %}
<li>
<h2>{{ result.source_title }}</h2>
<p class="meta">
{% if result.source_author %}{{ result.source_author }}{% endif %}
{% if result.chapter_title %} · {{ result.chapter_title }}{% endif %}
{% if result.page_label %} · page {{ result.page_label }}{% endif %}
</p>
<p>{{ result.text }}</p>
<dl class="scores">
<div>
<dt>final</dt>
<dd>{{ "%.3f"|format(result.score) }}</dd>
</div>
{% if result.rerank_score is not none %}
<div>
<dt>rerank</dt>
<dd>{{ "%.3f"|format(result.rerank_score) }}</dd>
</div>
{% endif %}
{% if result.vector_score is not none %}
<div>
<dt>vector cosine</dt>
<dd>{{ "%.3f"|format(result.vector_score) }}</dd>
</div>
{% endif %}
{% if result.bm25_score is not none %}
<div>
<dt>BM25</dt>
<dd>{{ "%.6f"|format(result.bm25_score) }}</dd>
</div>
{% endif %}
{% if result.fused_score is not none %}
<div>
<dt>RRF</dt>
<dd>{{ "%.3f"|format(result.fused_score) }}</dd>
</div>
{% endif %}
</dl>
</li>
{% endfor %}
</ol>
{% else %}
<p>No results.</p>
{% endif %}
@@ -0,0 +1,30 @@
<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>EPUB Search</title>
<script src="https://unpkg.com/htmx.org@2.0.4"></script>
<link rel="stylesheet" href="/static/style.css">
</head>
<body>
<main>
<nav>
<a href="/">Search</a>
<a href="/books">Books</a>
<a href="/admin">Admin</a>
</nav>
<h1>EPUB Search</h1>
<form hx-post="/search" hx-target="#results" hx-swap="innerHTML">
<label for="query">Search</label>
<textarea id="query" name="query" rows="4" required></textarea>
<label class="check">
<input type="checkbox" name="rerank" value="true" {% if config.rerank.enabled %}checked{% endif %}>
Rerank
</label>
<button type="submit">Search</button>
</form>
<section id="results"></section>
</main>
</body>
</html>
+13
View File
@@ -0,0 +1,13 @@
"""Shared web UI resources for EPUB search."""
from __future__ import annotations
from pathlib import Path
from fastapi.templating import Jinja2Templates
PACKAGE_DIR = Path(__file__).resolve().parent
TEMPLATE_DIR = PACKAGE_DIR / "templates"
STATIC_DIR = PACKAGE_DIR / "static"
templates = Jinja2Templates(directory=TEMPLATE_DIR)
+281
View File
@@ -0,0 +1,281 @@
"""Persisted BM25 corpus management."""
from __future__ import annotations
import json
import logging
import shutil
from dataclasses import dataclass
from datetime import UTC, datetime
from functools import cache
from pathlib import Path
from typing import TYPE_CHECKING
import bm25s
from sqlalchemy import func, select, union_all
from python.orm.richie import EbookChapter, EbookChunk, EbookSource
if TYPE_CHECKING:
from sqlalchemy.orm import Session
from python.ebook_search.config import EbookSearchConfig
logger = logging.getLogger(__name__)
MANIFEST_NAME = "manifest.json"
REQUIRED_INDEX_FILES = frozenset(
{
"data.csc.index.npy",
"indices.csc.index.npy",
"indptr.csc.index.npy",
"params.index.json",
"vocab.index.json",
"corpus.jsonl",
}
)
@dataclass(frozen=True)
class BM25Manifest:
"""Metadata describing a persisted BM25 corpus."""
created_at: datetime
db_updated_at: datetime | None
chunk_count: int
@dataclass(frozen=True)
class BM25Corpus:
"""Loaded persisted BM25 corpus and retriever."""
retriever: object | None
records: tuple[dict[str, object], ...]
manifest: BM25Manifest
class BM25CorpusUnavailableError(RuntimeError):
"""Raised when the persisted BM25 corpus cannot be loaded."""
def bm25_index_path(config: EbookSearchConfig) -> Path:
"""Return the configured BM25 index root path relative to the current working directory."""
path = Path(config.bm25_index_dir).expanduser()
if path.is_absolute():
return path
return Path.cwd() / path
def get_current_bm25_index(index_path: Path) -> Path:
"""Return the live BM25 index directory."""
current_path = index_path / "current"
if current_path.exists() or current_path.is_symlink():
return current_path
return index_path
def ensure_bm25_corpus(session: Session, config: EbookSearchConfig) -> None:
"""Create or refresh the persisted BM25 corpus when it is missing or stale."""
index_path = bm25_index_path(config)
manifest = read_bm25_manifest(index_path)
db_updated_at = corpus_last_updated_at(session)
if not bm25_index_exists(index_path, manifest):
logger.info("ebook_bm25_index_missing path=%s", index_path)
refresh_bm25_corpus(session, config, db_updated_at=db_updated_at)
return
if db_updated_at is not None and manifest is not None and manifest.created_at < db_updated_at:
logger.info(
"ebook_bm25_index_stale path=%s created_at=%s db_updated_at=%s",
index_path,
manifest.created_at.isoformat(),
db_updated_at.isoformat(),
)
refresh_bm25_corpus(session, config, db_updated_at=db_updated_at)
return
logger.info(
"ebook_bm25_index_current path=%s chunks=%s created_at=%s",
index_path,
manifest.chunk_count if manifest else 0,
manifest.created_at.isoformat() if manifest else None,
)
def refresh_bm25_corpus(
session: Session,
config: EbookSearchConfig,
*,
db_updated_at: datetime | None = None,
) -> BM25Manifest:
"""Rebuild and persist the BM25 corpus from the current database chunks."""
index_path = bm25_index_path(config)
records, texts = fetch_bm25_corpus_records(session)
manifest = BM25Manifest(
created_at=datetime.now(tz=UTC),
db_updated_at=db_updated_at if db_updated_at is not None else corpus_last_updated_at(session),
chunk_count=len(records),
)
write_bm25_corpus(index_path, records, texts, manifest)
logger.info(
"ebook_bm25_index_refreshed path=%s chunks=%s created_at=%s",
index_path,
manifest.chunk_count,
manifest.created_at.isoformat(),
)
return manifest
@cache
def load_bm25_corpus(config: EbookSearchConfig) -> BM25Corpus:
"""Load the BM25 corpus into memory once per process.
Background refresh tasks clear this cache after rebuilding the on-disk corpus.
"""
index_path = bm25_index_path(config)
active_index_path = get_current_bm25_index(index_path)
logger.info("ebook_bm25_corpus_cache_load path=%s active_path=%s", index_path, active_index_path)
manifest = read_bm25_manifest(index_path)
if manifest is None or not bm25_index_exists(index_path, manifest):
msg = f"BM25 corpus is not available: {index_path}"
raise BM25CorpusUnavailableError(msg)
if manifest.chunk_count == 0:
return BM25Corpus(retriever=None, records=(), manifest=manifest)
retriever = bm25s.BM25.load(active_index_path, load_corpus=True, mmap=True)
records = tuple(dict(record) for record in retriever.corpus)
return BM25Corpus(retriever=retriever, records=records, manifest=manifest)
def score_bm25_corpus(query: str, corpus: BM25Corpus, *, limit: int) -> list[tuple[dict[str, object], float]]:
"""Score a query against a loaded BM25 corpus."""
if corpus.retriever is None or not corpus.records:
return []
k = min(limit, len(corpus.records))
documents, scores = corpus.retriever.retrieve(
bm25s.tokenize(query, show_progress=False),
corpus=list(corpus.records),
k=k,
show_progress=False,
)
results: list[tuple[dict[str, object], float]] = []
for document, score in zip(documents[0], scores[0], strict=True):
score_value = float(score)
if score_value <= 0:
continue
results.append((dict(document), score_value))
return results
def fetch_bm25_corpus_records(session: Session) -> tuple[list[dict[str, object]], list[str]]:
"""Fetch persistable BM25 corpus records and their matching index texts from the database.
search_text is only needed to build the index, so it is returned separately instead of
being persisted into the corpus records, which would double the corpus size.
"""
statement = (
select(
EbookChunk.id.label("chunk_id"),
EbookChunk.text.label("text"),
EbookSource.title.label("source_title"),
EbookSource.author.label("source_author"),
EbookChapter.title.label("chapter_title"),
EbookChunk.page_label.label("page_label"),
EbookChunk.search_text.label("bm25_text"),
)
.select_from(EbookChunk)
.join(EbookSource, EbookSource.id == EbookChunk.source_id)
.outerjoin(EbookChapter, EbookChapter.id == EbookChunk.chapter_id)
.order_by(EbookChunk.id)
)
records: list[dict[str, object]] = []
texts: list[str] = []
for row in session.execute(statement).mappings():
record = dict(row)
texts.append(str(record.pop("bm25_text")))
records.append(record)
return records, texts
def corpus_last_updated_at(session: Session) -> datetime | None:
"""Return the latest source/chapter/chunk update timestamp relevant to BM25 text."""
update_times = union_all(
select(func.max(EbookSource.updated).label("updated")),
select(func.max(EbookChapter.updated).label("updated")),
select(func.max(EbookChunk.updated).label("updated")),
).subquery()
return session.scalar(select(func.max(update_times.c.updated)))
def write_bm25_corpus(
index_path: Path,
records: list[dict[str, object]],
texts: list[str],
manifest: BM25Manifest,
) -> None:
"""Write a BM25 corpus generation and publish it through the current symlink."""
index_path.mkdir(parents=True, exist_ok=True)
generations_path = index_path / "generations"
generations_path.mkdir(exist_ok=True)
generation_path = next_bm25_generation_path(generations_path, manifest.created_at)
current_path = index_path / "current"
next_current_path = index_path / f".current.{generation_path.name}.tmp"
try:
generation_path.mkdir()
# Empty corpora publish a manifest-only generation so startup succeeds before any chunks exist.
if records:
retriever = bm25s.BM25()
retriever.index(bm25s.tokenize(texts, show_progress=False), show_progress=False)
retriever.save(generation_path, corpus=records, show_progress=False)
write_bm25_manifest(generation_path, manifest)
next_current_path.unlink(missing_ok=True)
next_current_path.symlink_to(generation_path, target_is_directory=True)
next_current_path.replace(current_path)
except Exception:
next_current_path.unlink(missing_ok=True)
shutil.rmtree(generation_path, ignore_errors=True)
raise
def read_bm25_manifest(index_path: Path) -> BM25Manifest | None:
"""Read the BM25 manifest if it exists and is valid."""
manifest_path = get_current_bm25_index(index_path) / MANIFEST_NAME
if not manifest_path.exists():
return None
body = json.loads(manifest_path.read_text(encoding="utf-8"))
return BM25Manifest(
created_at=datetime.fromisoformat(str(body["created_at"])),
db_updated_at=datetime.fromisoformat(str(body["db_updated_at"])) if body.get("db_updated_at") else None,
chunk_count=int(body["chunk_count"]),
)
def write_bm25_manifest(index_path: Path, manifest: BM25Manifest) -> None:
"""Write the BM25 manifest to an index directory."""
body = {
"created_at": manifest.created_at.isoformat(),
"db_updated_at": manifest.db_updated_at.isoformat() if manifest.db_updated_at else None,
"chunk_count": manifest.chunk_count,
}
(index_path / MANIFEST_NAME).write_text(json.dumps(body, indent=2, sort_keys=True), encoding="utf-8")
def bm25_index_exists(index_path: Path, manifest: BM25Manifest | None) -> bool:
"""Return whether a usable persisted BM25 index exists."""
active_index_path = get_current_bm25_index(index_path)
if manifest is None or not active_index_path.is_dir():
return False
if manifest.chunk_count == 0:
return True
return all((active_index_path / file_name).exists() for file_name in REQUIRED_INDEX_FILES)
def next_bm25_generation_path(generations_path: Path, created_at: datetime) -> Path:
"""Return an unused dated BM25 generation path."""
base_name = created_at.astimezone(UTC).strftime("%Y%m%dT%H%M%S.%fZ")
generation_path = generations_path / base_name
suffix = 1
while generation_path.exists():
generation_path = generations_path / f"{base_name}.{suffix}"
suffix += 1
return generation_path
+117
View File
@@ -0,0 +1,117 @@
"""Configuration for the EPUB search app."""
from __future__ import annotations
from dataclasses import dataclass
from os import getenv
def getenv_bool(name: str, *, default: bool) -> bool:
"""Read a boolean environment variable with a default fallback."""
value = getenv(name)
if value is None:
return default
return value.strip().lower() in {"1", "true", "yes", "on"}
def getenv_int(name: str, *, default: int) -> int:
"""Read an integer environment variable with a default fallback."""
value = getenv(name)
if value is None or not value.strip():
return default
return int(value)
@dataclass(frozen=True)
class RerankConfig:
"""vLLM reranker settings."""
enabled: bool = False
base_url: str = "http://192.168.90.25:8001"
model: str = "qwen3-reranker-06b"
candidates: int = 24
timeout_seconds: float = 30.0
@dataclass(frozen=True)
class EbookSearchConfig:
"""Runtime settings for EPUB search."""
rerank: RerankConfig
top_k: int = 12
library_paths: tuple[str, ...] = ()
vllm_base_url: str = "https://ollama.com/v1"
vllm_api_key: str = "not-needed"
chat_model: str = "deepseek-v4-flash"
answer_enabled: bool = True
embedding_base_url: str = "http://192.168.90.25:8000/v1"
embedding_api_key: str = "not-needed"
embedding_model: str = "qwen3-embedding-0.6b"
embedding_batch_size: int = 32
bm25_index_dir: str = ".ebook_search_bm25"
bm25_refresh_delay_seconds: int = 60
def load_rerank_config() -> RerankConfig:
"""Load reranker config from environment variables."""
return RerankConfig(
enabled=getenv_bool("EBOOK_SEARCH_RERANK_ENABLED", default=False),
base_url=getenv("EBOOK_SEARCH_RERANK_BASE_URL", "http://192.168.90.25:8001"),
model=getenv("EBOOK_SEARCH_RERANK_MODEL", "qwen3-reranker-06b"),
candidates=getenv_int("EBOOK_SEARCH_RERANK_CANDIDATES", default=24),
timeout_seconds=float(getenv_int("EBOOK_SEARCH_RERANK_TIMEOUT_SECONDS", default=30)),
)
def load_config() -> EbookSearchConfig:
"""Load EPUB search config from environment variables."""
return EbookSearchConfig(
rerank=load_rerank_config(),
top_k=getenv_int("EBOOK_SEARCH_TOP_K", default=12),
library_paths=library_paths_from_env(),
vllm_base_url=getenv("EBOOK_SEARCH_VLLM_BASE_URL", "https://ollama.com/v1"),
vllm_api_key=getenv("EBOOK_SEARCH_VLLM_API_KEY") or getenv("OLLAMA_API_KEY") or "not-needed",
chat_model=getenv("EBOOK_SEARCH_CHAT_MODEL", "deepseek-v4-flash"),
answer_enabled=getenv_bool("EBOOK_SEARCH_ANSWER_ENABLED", default=True),
embedding_base_url=getenv("EBOOK_SEARCH_EMBEDDING_BASE_URL", "http://192.168.90.25:8000/v1"),
embedding_api_key=getenv("EBOOK_SEARCH_EMBEDDING_API_KEY", "not-needed"),
embedding_model=normalize_embedding_model(),
embedding_batch_size=getenv_int("EBOOK_SEARCH_EMBEDDING_BATCH_SIZE", default=32),
bm25_index_dir=getenv("EBOOK_SEARCH_BM25_INDEX_DIR", ".ebook_search_bm25"),
bm25_refresh_delay_seconds=getenv_int("EBOOK_SEARCH_BM25_REFRESH_DELAY_SECONDS", default=60),
)
def normalize_embedding_model(default: str = "qwen3-embedding-0.6b") -> str:
"""Normalize supported embedding aliases to provider model names."""
aliases = {
"Qwen3-Embedding-0.6B": "qwen3-embedding-0.6b",
"Qwen3-Embedding-4B": "qwen3-embedding-4b",
"Qwen3-Embedding-8B": "qwen3-embedding-8b",
"Qwen/Qwen3-Embedding-0.6B": "qwen3-embedding-0.6b",
"Qwen/Qwen3-Embedding-4B": "qwen3-embedding-4b",
"Qwen/Qwen3-Embedding-8B": "qwen3-embedding-8b",
"qwen3-embedding:0.6b": "qwen3-embedding-0.6b",
"qwen3-embedding:4b": "qwen3-embedding-4b",
"qwen3-embedding:8b": "qwen3-embedding-8b",
"qwen3-embedding-0.6b": "qwen3-embedding-0.6b",
"qwen3-embedding-4b": "qwen3-embedding-4b",
"qwen3-embedding-8b": "qwen3-embedding-8b",
}
model = getenv("EBOOK_SEARCH_EMBEDDING_MODEL", default)
standard_model = aliases.get(model)
if standard_model is None:
error = f"Embedding model {model} is not supported. Supported models are {aliases.keys()}"
raise ValueError(error)
return standard_model
def library_paths_from_env() -> tuple[str, ...]:
"""Read configured EPUB library paths from the environment."""
value = getenv("EBOOK_SEARCH_LIBRARY_PATHS")
if value is None:
return ()
return tuple(path for path in value.split(":") if path)
+170
View File
@@ -0,0 +1,170 @@
"""Embedding model helpers."""
from __future__ import annotations
import logging
from dataclasses import dataclass
from typing import TYPE_CHECKING
from sqlalchemy import func, select
from sqlalchemy.dialects.postgresql import insert
from python.ebook_search.llm_interface import request_embeddings
from python.orm.richie import (
EbookChunk,
EbookChunkEmbedding1024,
EbookChunkEmbedding2560,
EbookChunkEmbedding4096,
EbookEmbeddingModel,
)
logger = logging.getLogger(__name__)
if TYPE_CHECKING:
from collections.abc import Sequence
from sqlalchemy.orm import Session
from python.ebook_search.config import EbookSearchConfig
MODEL_DIMENSIONS = {
"qwen3-embedding-0.6b": 1024,
"qwen3-embedding-4b": 2560,
"qwen3-embedding-8b": 4096,
}
def get_embedding_table(
dimension: int,
) -> type[EbookChunkEmbedding1024 | EbookChunkEmbedding2560 | EbookChunkEmbedding4096]:
"""Return the embedding table mapped to an embedding dimension."""
embedding_tables = {
1024: EbookChunkEmbedding1024,
2560: EbookChunkEmbedding2560,
4096: EbookChunkEmbedding4096,
}
table = embedding_tables.get(dimension)
if not table:
msg = f"Embedding dimension {dimension} is not supported"
raise ValueError(msg)
return table
@dataclass(frozen=True)
class EmbeddingModelStats:
"""Embedding coverage for one model."""
model_name: str
dimension: int
embedded_chunks: int
total_chunks: int
@property
def missing_chunks(self) -> int:
"""Return chunks missing this embedding model."""
return max(self.total_chunks - self.embedded_chunks, 0)
def embed_texts(texts: Sequence[str], config: EbookSearchConfig) -> list[list[float]]:
"""Embed text with the configured vLLM embedding model."""
logger.info(
"ebook_embed_request_start base_url=%s model=%s count=%s",
config.embedding_base_url,
config.embedding_model,
len(texts),
)
vectors = request_embeddings(texts, config)
expected_dimension = MODEL_DIMENSIONS[config.embedding_model]
for vector in vectors:
if len(vector) != expected_dimension:
msg = f"Expected {expected_dimension} dimensions, got {len(vector)}"
raise ValueError(msg)
logger.info(
"ebook_embed_request_complete model=%s count=%s dimension=%s",
config.embedding_model,
len(vectors),
expected_dimension,
)
return vectors
def embed_query(query: str, config: EbookSearchConfig) -> list[float]:
"""Embed a search query with the Qwen retrieval instruction."""
instructed_query = f"Instruct: Retrieve relevant passages for the query.\nQuery: {query}"
return embed_texts([instructed_query], config)[0]
def ensure_embedding_models(session: Session) -> None:
"""Ensure supported embedding model rows exist."""
for name, dimension in MODEL_DIMENSIONS.items():
existing = session.scalar(select(EbookEmbeddingModel).where(EbookEmbeddingModel.name == name))
if existing is None:
session.add(EbookEmbeddingModel(name=name, dimension=dimension, is_default=name == "qwen3-embedding-0.6b"))
logger.info("ebook_embedding_model_created model=%s dimension=%s", name, dimension)
session.flush()
def embedding_model_stats(session: Session) -> list[EmbeddingModelStats]:
"""Return embedding coverage counts for every supported model."""
total_chunks = session.scalar(select(func.count(EbookChunk.id))) or 0
models = {
model.name: model
for model in session.scalars(
select(EbookEmbeddingModel)
.where(EbookEmbeddingModel.name.in_(MODEL_DIMENSIONS))
.order_by(EbookEmbeddingModel.name)
)
}
stats: list[EmbeddingModelStats] = []
for model_name, dimension in MODEL_DIMENSIONS.items():
model = models.get(model_name)
embedded_chunks = 0
if model is not None:
table = get_embedding_table(dimension)
embedded_chunks = session.scalar(select(func.count(table.id)).where(table.model_id == model.id)) or 0
stats.append(
EmbeddingModelStats(
model_name=model_name,
dimension=dimension,
embedded_chunks=embedded_chunks,
total_chunks=total_chunks,
)
)
return stats
def embed_missing_chunks(session: Session, config: EbookSearchConfig) -> int:
"""Embed chunks missing embeddings for the configured model."""
ensure_embedding_models(session)
model = session.scalar(select(EbookEmbeddingModel).where(EbookEmbeddingModel.name == config.embedding_model))
if model is None:
supported_models = ", ".join(MODEL_DIMENSIONS)
msg = f"Unknown embedding model: {config.embedding_model}. Supported models: {supported_models}"
raise ValueError(msg)
table = get_embedding_table(model.dimension)
chunks = list(
session.scalars(
select(EbookChunk)
.outerjoin(table, (table.chunk_id == EbookChunk.id) & (table.model_id == model.id))
.where(table.id.is_(None))
.order_by(EbookChunk.id)
.limit(config.embedding_batch_size)
)
)
if not chunks:
logger.info("ebook_embed_missing_none model=%s", config.embedding_model)
return 0
logger.info("ebook_embed_missing_batch_start model=%s count=%s", config.embedding_model, len(chunks))
vectors = embed_texts([chunk.text for chunk in chunks], config)
rows = [
{"chunk_id": chunk.id, "model_id": model.id, "embedding": vector}
for chunk, vector in zip(chunks, vectors, strict=True)
]
statement = insert(table).values(rows).on_conflict_do_nothing(index_elements=["chunk_id", "model_id"])
session.execute(statement)
session.flush()
logger.info("ebook_embed_missing_batch_complete model=%s count=%s", config.embedding_model, len(rows))
return len(rows)
+95
View File
@@ -0,0 +1,95 @@
"""EPUB parsing helpers."""
from __future__ import annotations
import re
from dataclasses import dataclass
from typing import TYPE_CHECKING
from bs4 import BeautifulSoup
from ebooklib import ITEM_DOCUMENT, epub
if TYPE_CHECKING:
from pathlib import Path
WHITESPACE_RE = re.compile(r"\s+")
@dataclass(frozen=True)
class ParsedChapter:
"""Text extracted from one EPUB spine document."""
title: str | None
href: str | None
text: str
page_labels: tuple[str, ...]
@dataclass(frozen=True)
class ParsedEpub:
"""Parsed EPUB metadata and text."""
title: str
author: str | None
language: str | None
publisher: str | None
identifier: str | None
chapters: tuple[ParsedChapter, ...]
def parse_epub(path: Path) -> ParsedEpub:
"""Parse EPUB metadata and spine text."""
book = epub.read_epub(path)
chapters = []
for item in book.get_items_of_type(ITEM_DOCUMENT):
soup = BeautifulSoup(item.get_content(), "html.parser")
title = chapter_title(soup)
page_labels = tuple(extract_page_labels(soup))
text = clean_text(soup.get_text(" "))
if text:
chapters.append(ParsedChapter(title=title, href=item.get_name(), text=text, page_labels=page_labels))
return ParsedEpub(
title=metadata_value(book, "title") or path.stem,
author=metadata_value(book, "creator"),
language=metadata_value(book, "language"),
publisher=metadata_value(book, "publisher"),
identifier=metadata_value(book, "identifier"),
chapters=tuple(chapters),
)
def metadata_value(book: epub.EpubBook, name: str) -> str | None:
"""Return the first non-empty Dublin Core metadata value for a name."""
values = book.get_metadata("DC", name)
if not values:
return None
value = values[0][0]
return str(value).strip() or None
def chapter_title(soup: BeautifulSoup) -> str | None:
"""Extract the best available title from an EPUB document soup."""
heading = soup.find(["h1", "h2", "h3"])
if heading is None:
title = soup.find("title")
if title is None:
return None
return clean_text(title.get_text(" ")) or None
return clean_text(heading.get_text(" ")) or None
def extract_page_labels(soup: BeautifulSoup) -> list[str]:
"""Extract EPUB page-break labels from a document soup."""
labels: list[str] = []
for tag in soup.find_all(attrs={"epub:type": "pagebreak"}):
label = tag.get("title") or tag.get("aria-label") or tag.get_text(" ")
clean = clean_text(str(label))
if clean:
labels.append(clean)
return labels
def clean_text(text: str) -> str:
"""Normalize whitespace in extracted EPUB text."""
return WHITESPACE_RE.sub(" ", text).strip()
+190
View File
@@ -0,0 +1,190 @@
"""EPUB ingestion into Richie DB."""
from __future__ import annotations
import hashlib
import logging
from dataclasses import dataclass
from datetime import UTC, datetime
from pathlib import Path
from typing import TYPE_CHECKING
import tiktoken
from sqlalchemy import or_, select
from python.ebook_search.epub_parse import parse_epub
from python.orm.richie import EbookChapter, EbookChunk, EbookSource
logger = logging.getLogger(__name__)
DEFAULT_CHUNK_TOKENS = 700
DEFAULT_CHUNK_OVERLAP = 100
if TYPE_CHECKING:
from sqlalchemy.orm import Session
from python.ebook_search.config import EbookSearchConfig
from python.ebook_search.epub_parse import ParsedChapter
@dataclass(frozen=True)
class TextChunk:
"""A token-bounded chunk of text."""
text: str
token_start: int
token_count: int
def chunk_text(
text: str,
*,
chunk_tokens: int = DEFAULT_CHUNK_TOKENS,
overlap_tokens: int = DEFAULT_CHUNK_OVERLAP,
) -> list[TextChunk]:
"""Split text into overlapping token chunks."""
if chunk_tokens <= 0:
msg = "chunk_tokens must be positive"
raise ValueError(msg)
if overlap_tokens < 0 or overlap_tokens >= chunk_tokens:
msg = "overlap_tokens must be non-negative and smaller than chunk_tokens"
raise ValueError(msg)
encoding = tiktoken.get_encoding("cl100k_base")
tokens = encoding.encode(text)
if not tokens:
return []
chunks: list[TextChunk] = []
step = chunk_tokens - overlap_tokens
for start in range(0, len(tokens), step):
chunk = tokens[start : start + chunk_tokens]
if not chunk:
continue
chunks.append(
TextChunk(
text=encoding.decode(chunk).strip(),
token_start=start,
token_count=len(chunk),
)
)
if start + chunk_tokens >= len(tokens):
break
return [chunk for chunk in chunks if chunk.text]
def ingest_configured_paths(session: Session, config: EbookSearchConfig) -> int:
"""Ingest every EPUB found under configured library paths."""
count = 0
for library_path in config.library_paths:
path = Path(library_path).expanduser()
logger.info("ebook_ingest_path_start path=%s", path)
if path.is_file() and path.suffix.lower() == ".epub":
count += int(ingest_file(session, path))
elif path.is_dir():
for epub_path in sorted(path.rglob("*.epub")):
count += int(ingest_file(session, epub_path))
else:
logger.warning("ebook_ingest_path_missing path=%s", path)
logger.info("ebook_ingest_paths_complete changed_files=%s configured_paths=%s", count, len(config.library_paths))
return count
def ingest_file(session: Session, path: Path) -> bool:
"""Ingest one EPUB file. Return True when the database changed."""
resolved_path = path.expanduser().resolve()
logger.info("ebook_ingest_file_start path=%s", resolved_path)
file_hash = sha256_file(resolved_path)
existing = find_existing_source(session, resolved_path, file_hash)
if existing is not None and existing.file_sha256 == file_hash:
stat = resolved_path.stat()
existing.file_path = str(resolved_path)
existing.file_mtime = datetime.fromtimestamp(stat.st_mtime, tz=UTC)
existing.file_size = stat.st_size
session.flush()
logger.info("ebook_ingest_file_unchanged source_id=%s path=%s", existing.id, resolved_path)
return False
if existing is not None:
logger.info("ebook_ingest_file_replacing source_id=%s path=%s", existing.id, resolved_path)
session.delete(existing)
session.flush()
stat = resolved_path.stat()
parsed = parse_epub(resolved_path)
source = EbookSource(
title=parsed.title,
author=parsed.author,
language=parsed.language,
publisher=parsed.publisher,
identifier=parsed.identifier,
file_path=str(resolved_path),
file_sha256=file_hash,
file_mtime=datetime.fromtimestamp(stat.st_mtime, tz=UTC),
file_size=stat.st_size,
)
session.add(source)
session.flush()
chunk_index = 0
for spine_index, parsed_chapter in enumerate(parsed.chapters):
chapter = EbookChapter(
source_id=source.id,
spine_index=spine_index,
title=parsed_chapter.title,
href=parsed_chapter.href,
)
session.add(chapter)
session.flush()
chunk_index = add_chapter_chunks(session, source, chapter, parsed_chapter, chunk_index)
session.flush()
logger.info(
"ebook_ingest_file_complete source_id=%s path=%s chapters=%s chunks=%s",
source.id,
resolved_path,
len(parsed.chapters),
chunk_index,
)
return True
def find_existing_source(session: Session, path: Path, file_hash: str) -> EbookSource | None:
"""Find an existing source by canonical path or file hash."""
return session.scalar(
select(EbookSource).where(or_(EbookSource.file_path == str(path), EbookSource.file_sha256 == file_hash))
)
def add_chapter_chunks(
session: Session,
source: EbookSource,
chapter: EbookChapter,
parsed_chapter: ParsedChapter,
chunk_index: int,
) -> int:
"""Add chunk rows for one parsed chapter and return the next chunk index."""
page_label = parsed_chapter.page_labels[0] if parsed_chapter.page_labels else None
for text_chunk in chunk_text(parsed_chapter.text):
session.add(
EbookChunk(
source_id=source.id,
chapter_id=chapter.id,
chunk_index=chunk_index,
text=text_chunk.text,
token_start=text_chunk.token_start,
token_count=text_chunk.token_count,
page_label=page_label,
content_sha256=hashlib.sha256(text_chunk.text.encode()).hexdigest(),
search_text=f"{source.title} {source.author or ''} {chapter.title or ''} {text_chunk.text}",
)
)
chunk_index += 1
return chunk_index
def sha256_file(path: Path) -> str:
"""Calculate the SHA-256 digest for a file."""
digest = hashlib.sha256()
with path.open("rb") as file:
for block in iter(lambda: file.read(1024 * 1024), b""):
digest.update(block)
return digest.hexdigest()
+143
View File
@@ -0,0 +1,143 @@
"""LLM provider HTTP adapters."""
from __future__ import annotations
import logging
from typing import TYPE_CHECKING
import httpx
if TYPE_CHECKING:
from collections.abc import Sequence
from python.ebook_search.config import EbookSearchConfig, RerankConfig
logger = logging.getLogger(__name__)
def auth_headers(api_key: str) -> dict[str, str]:
"""Build authorization headers when an API key is configured."""
if api_key == "not-needed":
return {}
return {"Authorization": f"Bearer {api_key}"}
def request_embeddings(texts: Sequence[str], config: EbookSearchConfig) -> list[list[float]]:
"""Request embeddings from the configured OpenAI-compatible endpoint."""
try:
response = httpx.post(
f"{config.embedding_base_url.rstrip('/')}/embeddings",
headers=auth_headers(config.embedding_api_key),
json={"model": config.embedding_model, "input": list(texts)},
timeout=60,
)
response.raise_for_status()
return embedding_vectors_from_response(response.json())
except (httpx.HTTPError, ValueError, KeyError, TypeError) as error:
logger.exception(
"ebook_embed_request_failed base_url=%s model=%s count=%s",
config.embedding_base_url,
config.embedding_model,
len(texts),
)
msg = f"Embedding request failed. base_url={config.embedding_base_url} model={config.embedding_model}"
raise RuntimeError(msg) from error
def embedding_vectors_from_response(body: object) -> list[list[float]]:
"""Extract embedding vectors from an OpenAI-compatible embedding response."""
if not isinstance(body, dict):
msg = "Embedding response is not an object"
raise TypeError(msg)
data = body["data"]
if not isinstance(data, list):
msg = "Embedding response data is not a list"
raise TypeError(msg)
vectors: list[list[float]] = []
for item in data:
if not isinstance(item, dict):
msg = "Embedding item is not an object"
raise TypeError(msg)
embedding = item["embedding"]
if not isinstance(embedding, list):
msg = "Embedding value is not a list"
raise TypeError(msg)
vectors.append([float(value) for value in embedding])
return vectors
def request_rerank(
query: str,
documents: Sequence[str],
config: RerankConfig,
) -> object | None:
"""Request rerank scores from the configured vLLM endpoint."""
payload = {
"model": config.model,
"query": query,
"documents": list(documents),
}
response = httpx.post(
f"{config.base_url.rstrip('/')}/rerank",
json=payload,
timeout=config.timeout_seconds,
)
response.raise_for_status()
try:
return response.json()
except ValueError:
logger.debug("ebook_rerank_response_invalid_json", extra={"response": response.text})
return None
def request_chat_completion(
config: EbookSearchConfig,
messages: Sequence[dict[str, str]],
) -> str:
"""Request a chat completion from the configured OpenAI-compatible endpoint."""
try:
response = httpx.post(
f"{config.vllm_base_url.rstrip('/')}/chat/completions",
headers=auth_headers(config.vllm_api_key),
json={
"model": config.chat_model,
"messages": list(messages),
"temperature": 0,
},
timeout=60,
)
response.raise_for_status()
return chat_content_from_response(response.json())
except (httpx.HTTPError, ValueError, KeyError, TypeError) as error:
msg = f"Chat request failed. base_url={config.vllm_base_url} model={config.chat_model}"
raise RuntimeError(msg) from error
def chat_content_from_response(body: object) -> str:
"""Extract text content from an OpenAI-compatible chat response."""
if not isinstance(body, dict):
msg = "Chat response is not an object"
raise TypeError(msg)
choices = body["choices"]
if not isinstance(choices, list) or not choices:
msg = "Chat response has no choices"
raise ValueError(msg)
first = choices[0]
if not isinstance(first, dict):
msg = "Chat choice is not an object"
raise TypeError(msg)
message = first["message"]
if not isinstance(message, dict):
msg = "Chat message is not an object"
raise TypeError(msg)
content = message.get("content") or ""
if not isinstance(content, str):
msg = "Chat content is not text"
raise TypeError(msg)
return content
+129
View File
@@ -0,0 +1,129 @@
"""vLLM-backed optional reranking."""
from __future__ import annotations
import logging
from dataclasses import dataclass, replace
from typing import TYPE_CHECKING
from python.ebook_search.llm_interface import request_rerank
if TYPE_CHECKING:
from python.ebook_search.config import RerankConfig
from python.ebook_search.search import SearchResult
logger = logging.getLogger(__name__)
RERANK_SCORE_WEIGHT = 0.7
HYBRID_SCORE_WEIGHT = 0.3
@dataclass(frozen=True)
class RerankResult:
"""A relevance score for one candidate chunk."""
chunk_id: int
score: float
def rerank_chunks(query: str, candidates: list[SearchResult], config: RerankConfig) -> list[SearchResult]:
"""Rerank candidates with a vLLM rerank endpoint."""
if not candidates:
return []
logger.info(
"ebook_rerank_request_start base_url=%s model=%s candidates=%s",
config.base_url,
config.model,
len(candidates),
)
scores = score_candidates(query, candidates, config)
results = sorted(
(
replace(
result,
score=final_rerank_score(result, scores[result.chunk_id].score, candidates),
rerank_score=scores[result.chunk_id].score,
)
for result in candidates
),
key=lambda result: result.score,
reverse=True,
)
logger.info(
"ebook_rerank_request_complete base_url=%s model=%s candidates=%s",
config.base_url,
config.model,
len(results),
)
return results
def score_candidates(
query: str,
candidates: list[SearchResult],
config: RerankConfig,
) -> dict[int, RerankResult]:
"""Score candidate chunks with the configured rerank API."""
body = request_rerank(query, [candidate.text for candidate in candidates], config)
if body is None:
return zero_rerank_scores(candidates)
scores = parse_vllm_scores(body, candidates)
for result in scores.values():
logger.debug("ebook_rerank_candidate_scored chunk_id=%s score=%s", result.chunk_id, result.score)
return scores
def parse_vllm_scores(body: object, candidates: list[SearchResult]) -> dict[int, RerankResult]:
"""Parse vLLM rerank scores into chunk-id keyed results."""
if not isinstance(body, dict):
logger.debug("ebook_rerank_response_not_object", extra={"response": body})
return zero_rerank_scores(candidates)
results = body.get("results") or body.get("data")
if not isinstance(results, list):
logger.debug("ebook_rerank_response_missing_results", extra={"response": body})
return zero_rerank_scores(candidates)
scores = zero_rerank_scores(candidates)
for item in results:
if not isinstance(item, dict):
continue
index = item.get("index")
score = item.get("relevance_score", item.get("score"))
if not isinstance(index, int) or index < 0 or index >= len(candidates):
continue
if not isinstance(score, int | float):
continue
chunk_id = candidates[index].chunk_id
scores[chunk_id] = RerankResult(chunk_id=chunk_id, score=clamp_score(float(score)))
return scores
def zero_rerank_scores(candidates: list[SearchResult]) -> dict[int, RerankResult]:
"""Return zero relevance scores for all candidate chunks."""
return {candidate.chunk_id: RerankResult(chunk_id=candidate.chunk_id, score=0.0) for candidate in candidates}
def clamp_score(score: float) -> float:
"""Clamp a rerank score into the supported 0.0 to 1.0 range."""
return min(max(score, 0.0), 1.0)
def final_rerank_score(result: SearchResult, rerank_score: float, candidates: list[SearchResult]) -> float:
"""Combine rerank relevance with normalized hybrid retrieval evidence."""
return (RERANK_SCORE_WEIGHT * rerank_score) + (HYBRID_SCORE_WEIGHT * normalized_hybrid_score(result, candidates))
def normalized_hybrid_score(result: SearchResult, candidates: list[SearchResult]) -> float:
"""Normalize a candidate hybrid score against the rerank candidate set."""
hybrid_scores = [
candidate.fused_score if candidate.fused_score is not None else candidate.score for candidate in candidates
]
low = min(hybrid_scores)
high = max(hybrid_scores)
if high == low:
return 1.0
score = result.fused_score if result.fused_score is not None else result.score
return (score - low) / (high - low)
+383
View File
@@ -0,0 +1,383 @@
"""Hybrid search orchestration."""
from __future__ import annotations
import logging
import re
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass, replace
from typing import TYPE_CHECKING
from pgvector.sqlalchemy import Vector
from sqlalchemy import literal, select
from sqlalchemy.orm import Session
from python.ebook_search.bm25_corpus import (
BM25CorpusUnavailableError,
load_bm25_corpus,
score_bm25_corpus,
)
from python.ebook_search.embeddings import MODEL_DIMENSIONS, embed_query, get_embedding_table
from python.ebook_search.rerank import rerank_chunks
from python.ebook_search.timing import RuntimeStep, timed_result
from python.orm.richie import (
EbookChapter,
EbookChunk,
EbookEmbeddingModel,
EbookSource,
)
if TYPE_CHECKING:
from collections.abc import Mapping
from sqlalchemy.engine import Engine
from python.ebook_search.config import EbookSearchConfig
logger = logging.getLogger(__name__)
BM25_CANDIDATE_LIMIT = 120
@dataclass(frozen=True)
class SearchResult:
"""One source chunk returned by search."""
chunk_id: int
text: str
source_title: str
score: float = 0.0
vector_score: float | None = None
bm25_score: float | None = None
fused_score: float | None = None
rerank_score: float | None = None
source_author: str | None = None
chapter_title: str | None = None
page_label: str | None = None
rank_source: str = "Hybrid"
@dataclass(frozen=True)
class SearchResponse:
"""Search output for the UI."""
query: str
results: list[SearchResult]
rank_label: str
timings: tuple[RuntimeStep, ...] = ()
@property
def total_runtime_ms(self) -> float:
"""Return total measured runtime for the response."""
return sum(step.duration_ms for step in self.timings if step.counts_toward_total)
@dataclass(frozen=True)
class RetrievalResponse:
"""Parallel retrieval output for vector and BM25 candidates."""
vector_results: list[SearchResult]
lexical_results: list[SearchResult]
timings: tuple[RuntimeStep, ...]
def search_ebooks(
engine: Engine,
query: str,
config: EbookSearchConfig,
*,
rerank: bool = False,
) -> SearchResponse:
"""Run hybrid vector/BM25 search and optional reranking."""
if not query.strip():
logger.info("ebook_search_empty_query")
return SearchResponse(query=query, results=[], rank_label="Hybrid")
logger.info("ebook_search_start query_length=%s rerank=%s", len(query), rerank)
timings: list[RuntimeStep] = []
bm25_query, timing = timed_result("BM25 query preparation", retrieval_query_from_text, query)
timings.append(timing)
retrieval, timing = timed_result(
"Hybrid retrieval",
parallel_retrieval,
engine,
query,
bm25_query,
config,
)
timings.extend(retrieval.timings)
timings.append(timing)
fused, timing = timed_result(
"Reciprocal rank fusion",
reciprocal_rank_fusion,
retrieval.vector_results,
retrieval.lexical_results,
)
timings.append(timing)
if config.rerank.enabled and rerank:
response, timing = timed_result("Rerank", apply_rerank, query, fused, config)
else:
response, timing = timed_result("Rerank skipped", skip_rerank, query, fused, config)
timings.append(timing)
response = replace(response, timings=tuple(timings))
logger.info(
"ebook_search_complete vector_candidates=%s lexical_candidates=%s "
"fused_candidates=%s returned=%s rank_label=%s runtime_ms=%.1f",
len(retrieval.vector_results),
len(retrieval.lexical_results),
len(fused),
len(response.results),
response.rank_label,
response.total_runtime_ms,
)
return response
def parallel_retrieval(
engine: Engine,
vector_query: str,
bm25_query: str,
config: EbookSearchConfig,
) -> RetrievalResponse:
"""Run vector and BM25 candidate retrieval concurrently with separate database sessions."""
with ThreadPoolExecutor(max_workers=2, thread_name_prefix="ebook-search") as executor:
vector_future = executor.submit(
timed_result,
"Embedding + vector search",
vector_candidates,
engine,
vector_query,
config,
)
bm25_future = executor.submit(
timed_result,
"BM25 search",
bm25_candidates,
bm25_query,
config,
)
vector_results, vector_timing = vector_future.result()
lexical_results, lexical_timing = bm25_future.result()
logger.info(
"ebook_parallel_retrieval_complete vector_candidates=%s lexical_candidates=%s",
len(vector_results),
len(lexical_results),
)
return RetrievalResponse(
vector_results=vector_results,
lexical_results=lexical_results,
timings=(
replace(vector_timing, counts_toward_total=False),
replace(lexical_timing, counts_toward_total=False),
),
)
def skip_rerank(
query: str,
candidates: list[SearchResult],
config: EbookSearchConfig,
) -> SearchResponse:
"""Return fused hybrid results without reranking."""
logger.info("ebook_rerank_skipped candidates=%s", len(candidates))
return SearchResponse(query=query, results=candidates[: config.top_k], rank_label="Hybrid")
def apply_rerank(
query: str,
candidates: list[SearchResult],
config: EbookSearchConfig,
) -> SearchResponse:
"""Rerank already-fused hybrid candidates."""
reranked = rerank_chunks(query, candidates[: config.rerank.candidates], config.rerank)
logger.info(
"ebook_rerank_complete input_candidates=%s returned=%s",
min(len(candidates), config.rerank.candidates),
len(reranked),
)
return SearchResponse(
query=query,
results=[replace(result, rank_source="Hybrid + rerank") for result in reranked[: config.top_k]],
rank_label="Hybrid + rerank",
)
def vector_candidates(engine: Engine, query: str, config: EbookSearchConfig) -> list[SearchResult]:
"""Return pgvector cosine candidates for a natural-language query."""
with Session(engine) as session:
model = session.scalar(select(EbookEmbeddingModel).where(EbookEmbeddingModel.name == config.embedding_model))
if model is None:
msg = f"Embedding model is not registered: {config.embedding_model}"
raise ValueError(msg)
expected_dimension = MODEL_DIMENSIONS[config.embedding_model]
if model.dimension != expected_dimension:
msg = f"Model row dimension {model.dimension} does not match configured dimension {expected_dimension}"
raise ValueError(msg)
embedding = embed_query(query, config)
limit = max(config.rerank.candidates, config.top_k) * 4
embedding_table = get_embedding_table(model.dimension)
embedding_param = literal(embedding, type_=Vector(model.dimension))
distance = embedding_table.embedding.op("<=>")(embedding_param)
score = (literal(1.0) - distance).label("score")
statement = (
select(
EbookChunk.id.label("chunk_id"),
EbookChunk.text.label("text"),
EbookSource.title.label("source_title"),
EbookSource.author.label("source_author"),
EbookChapter.title.label("chapter_title"),
EbookChunk.page_label.label("page_label"),
score,
)
.select_from(embedding_table)
.join(EbookChunk, EbookChunk.id == embedding_table.chunk_id)
.join(EbookSource, EbookSource.id == EbookChunk.source_id)
.outerjoin(EbookChapter, EbookChapter.id == EbookChunk.chapter_id)
.where(embedding_table.model_id == model.id)
.order_by(distance)
.limit(limit)
)
rows = session.execute(statement).mappings()
results = [search_result_from_row(row) for row in rows]
logger.info(
"ebook_vector_search_complete model=%s dimension=%s candidates=%s",
config.embedding_model,
model.dimension,
len(results),
)
return results
def bm25_candidates(query: str, config: EbookSearchConfig) -> list[SearchResult]:
"""Return BM25-ranked lexical candidates using the persisted corpus."""
try:
corpus = load_bm25_corpus(config)
except BM25CorpusUnavailableError as error:
logger.warning("ebook_bm25_index_unavailable_skipping error=%s", error)
return []
if not corpus.records:
logger.info("ebook_bm25_search_complete corpus=0 candidates=0")
return []
scored_records = score_bm25_corpus(query, corpus, limit=BM25_CANDIDATE_LIMIT)
results = [
replace(search_result_from_row(record), score=score, vector_score=None, bm25_score=score)
for record, score in scored_records
]
max_score = results[0].bm25_score if results else 0.0
logger.info(
"ebook_bm25_search_complete corpus=%s candidates=%s max_score=%.6f",
len(corpus.records),
len(results),
max_score,
)
return results
def reciprocal_rank_fusion(
vector_results: list[SearchResult],
lexical_results: list[SearchResult],
*,
rank_constant: int = 60,
) -> list[SearchResult]:
"""Fuse vector and lexical rankings with Reciprocal Rank Fusion."""
by_chunk: dict[int, SearchResult] = {}
scores: dict[int, float] = {}
vector_scores: dict[int, float] = {}
bm25_scores: dict[int, float] = {}
for rank, result in enumerate(vector_results, start=1):
by_chunk.setdefault(result.chunk_id, result)
vector_scores[result.chunk_id] = result.vector_score if result.vector_score is not None else result.score
scores[result.chunk_id] = scores.get(result.chunk_id, 0.0) + (1 / (rank_constant + rank))
for rank, result in enumerate(lexical_results, start=1):
by_chunk.setdefault(result.chunk_id, result)
bm25_scores[result.chunk_id] = result.bm25_score if result.bm25_score is not None else result.score
scores[result.chunk_id] = scores.get(result.chunk_id, 0.0) + (1 / (rank_constant + rank))
return sorted(
(
replace(
result,
score=scores[result.chunk_id],
vector_score=vector_scores.get(result.chunk_id),
bm25_score=bm25_scores.get(result.chunk_id),
fused_score=scores[result.chunk_id],
rank_source="Hybrid",
)
for result in by_chunk.values()
),
key=lambda result: result.score,
reverse=True,
)
def search_result_from_row(row: Mapping[str, object]) -> SearchResult:
"""Convert a database row mapping into a search result."""
return SearchResult(
chunk_id=int(row["chunk_id"]),
text=str(row["text"]),
source_title=str(row["source_title"]),
source_author=optional_str(row["source_author"]),
chapter_title=optional_str(row["chapter_title"]),
page_label=optional_str(row["page_label"]),
score=float(row["score"]) if "score" in row else 0.0,
vector_score=float(row["score"]) if "score" in row else None,
)
def optional_str(value: object) -> str | None:
"""Convert nullable database values to optional strings."""
if value is None:
return None
return str(value)
TOKEN_RE = re.compile(r"[A-Za-z0-9_]+")
def tokens(text_value: str) -> list[str]:
"""Extract tokens from a text value.
This is a simple approximation of the tokenization used by PostgreSQL's full-text search,
which is sufficient for BM25 candidate retrieval. It lowercases tokens and includes alphanumeric characters and
underscores.
"""
return [match.group(0).lower() for match in TOKEN_RE.finditer(text_value)]
QUERY_STOP_WORDS = {
"a",
"an",
"and",
"are",
"as",
"at",
"does",
"for",
"in",
"is",
"of",
"the",
"to",
"what",
"when",
"where",
"which",
"who",
"why",
}
def retrieval_query_from_text(query: str) -> str:
"""Remove generic question words while preserving entity and series terms."""
keywords = [token for token in tokens(query) if token not in QUERY_STOP_WORDS]
if not keywords:
return query
return " ".join(keywords)
+36
View File
@@ -0,0 +1,36 @@
"""Runtime timing helpers for EPUB search."""
from __future__ import annotations
from dataclasses import dataclass
from time import perf_counter
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from collections.abc import Callable
@dataclass(frozen=True)
class RuntimeStep:
"""Elapsed runtime for one named search step."""
name: str
duration_ms: float
counts_toward_total: bool = True
def runtime_step_from_start(name: str, start_seconds: float) -> RuntimeStep:
"""Create a runtime step from a prior perf_counter timestamp."""
return RuntimeStep(name=name, duration_ms=(perf_counter() - start_seconds) * 1000)
def timed_result[T, **P](
name: str,
operation: Callable[P, T],
*args: P.args,
**kwargs: P.kwargs,
) -> tuple[T, RuntimeStep]:
"""Run an operation and return its result plus elapsed runtime."""
start_seconds = perf_counter()
result = operation(*args, **kwargs)
return result, runtime_step_from_start(name, start_seconds)
+6
View File
@@ -0,0 +1,6 @@
"""Reusable FastAPI tools."""
from python.fastapi_tools.db import DbSession, get_db
from python.fastapi_tools.zstd_middleware import ZstdMiddleware
__all__ = ["DbSession", "ZstdMiddleware", "get_db"]
@@ -1,11 +1,15 @@
"""FastAPI dependencies.""" """FastAPI dependencies."""
from collections.abc import Iterator from __future__ import annotations
from typing import Annotated
from typing import TYPE_CHECKING, Annotated
from fastapi import Depends, Request from fastapi import Depends, Request
from sqlalchemy.orm import Session from sqlalchemy.orm import Session
if TYPE_CHECKING:
from collections.abc import Iterator
def get_db(request: Request) -> Iterator[Session]: def get_db(request: Request) -> Iterator[Session]:
"""Get database session from app state.""" """Get database session from app state."""
@@ -1,10 +1,14 @@
"""Middleware for the FastAPI application.""" """Zstd response compression middleware."""
from compression import zstd from compression import zstd
from typing import TYPE_CHECKING
from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
from starlette.requests import Request
from starlette.responses import Response from starlette.responses import Response
if TYPE_CHECKING:
from starlette.requests import Request
MINIMUM_RESPONSE_SIZE = 500 MINIMUM_RESPONSE_SIZE = 500
+6 -2
View File
@@ -1,9 +1,10 @@
"""FastAPI heater control service.""" """FastAPI heater control service."""
from __future__ import annotations
import logging import logging
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager from contextlib import asynccontextmanager
from typing import Annotated from typing import TYPE_CHECKING, Annotated
import typer import typer
import uvicorn import uvicorn
@@ -13,6 +14,9 @@ from python.common import configure_logger
from python.heater.controller import HeaterController from python.heater.controller import HeaterController
from python.heater.models import ActionResult, DeviceConfig, HeaterStatus from python.heater.models import ActionResult, DeviceConfig, HeaterStatus
if TYPE_CHECKING:
from collections.abc import AsyncIterator
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
+1
View File
@@ -262,6 +262,7 @@ def installer(
): ):
run(command, check=True, stdin=test.stdout) run(command, check=True, stdin=test.stdout)
# Fixed mount point for the new system; the installer runs as root on a fresh disk
mnt_dir = "/tmp/nix_install" # noqa: S108 mnt_dir = "/tmp/nix_install" # noqa: S108
Path(mnt_dir).mkdir(parents=True, exist_ok=True) Path(mnt_dir).mkdir(parents=True, exist_ok=True)
@@ -16,9 +16,13 @@ from typing import TYPE_CHECKING
if TYPE_CHECKING: if TYPE_CHECKING:
from collections.abc import Sequence from collections.abc import Sequence
logger = logging.getLogger(__name__)
ESCAPE_KEY = 27
def configure_logger(level: str = "INFO") -> None: def configure_logger(level: str = "INFO") -> None:
"""Configure the logger. """Configure the logger.
Args: Args:
level (str, optional): The logging level. Defaults to "INFO". level (str, optional): The logging level. Defaults to "INFO".
""" """
@@ -32,15 +36,17 @@ def configure_logger(level: str = "INFO") -> None:
def bash_wrapper(command: str) -> str: def bash_wrapper(command: str) -> str:
"""Execute a bash command and capture the output. """Execute a bash command and capture the output.
Args: Args:
command (str): The bash command to be executed. command (str): The bash command to be executed.
Returns: Returns:
Tuple[str, int]: A tuple containing the output of the command (stdout) as a string, Tuple[str, int]: A tuple containing the output of the command (stdout) as a string,
the error output (stderr) as a string (optional), and the return code as an integer. the error output (stderr) as a string (optional), and the return code as an integer.
""" """
logging.debug(f"running {command=}") logger.debug(f"running {command=}")
# This is a acceptable risk # This is a acceptable risk
process = Popen(command.split(), stdout=PIPE, stderr=PIPE) # noqa: S603 process = Popen(command.split(), stdout=PIPE, stderr=PIPE)
output, _ = process.communicate() output, _ = process.communicate()
if process.returncode != 0: if process.returncode != 0:
error = f"Failed to run command {command=} return code {process.returncode=}" error = f"Failed to run command {command=} return code {process.returncode=}"
@@ -51,6 +57,7 @@ def bash_wrapper(command: str) -> str:
def partition_disk(disk: str, swap_size: int, reserve: int = 0) -> None: def partition_disk(disk: str, swap_size: int, reserve: int = 0) -> None:
"""Partition a disk. """Partition a disk.
Args: Args:
disk (str): The disk to partition. disk (str): The disk to partition.
swap_size (int): The size of the swap partition in GB. swap_size (int): The size of the swap partition in GB.
@@ -58,7 +65,7 @@ def partition_disk(disk: str, swap_size: int, reserve: int = 0) -> None:
reserve (int, optional): The size of the reserve partition in GB. Defaults to 0. reserve (int, optional): The size of the reserve partition in GB. Defaults to 0.
minimum value is 0. minimum value is 0.
""" """
logging.info(f"partitioning {disk=}") logger.info(f"partitioning {disk=}")
swap_size = max(swap_size, 1) swap_size = max(swap_size, 1)
reserve = max(reserve, 0) reserve = max(reserve, 0)
@@ -66,16 +73,16 @@ def partition_disk(disk: str, swap_size: int, reserve: int = 0) -> None:
if reserve > 0: if reserve > 0:
msg = f"Creating swap partition on {disk=} with size {swap_size=}GiB and reserve {reserve=}GiB" msg = f"Creating swap partition on {disk=} with size {swap_size=}GiB and reserve {reserve=}GiB"
logging.info(msg) logger.info(msg)
swap_start = swap_size + reserve swap_start = swap_size + reserve
swap_partition = f"mkpart swap -{swap_start}GiB -{reserve}GiB " swap_partition = f"mkpart swap -{swap_start}GiB -{reserve}GiB "
else: else:
logging.info(f"Creating swap partition on {disk=} with size {swap_size=}GiB") logger.info(f"Creating swap partition on {disk=} with size {swap_size=}GiB")
swap_start = swap_size swap_start = swap_size
swap_partition = f"mkpart swap -{swap_start}GiB 100% " swap_partition = f"mkpart swap -{swap_start}GiB 100% "
logging.debug(f"{swap_partition=}") logger.debug(f"{swap_partition=}")
create_partitions = ( create_partitions = (
f"parted --script --align=optimal {disk} -- " f"parted --script --align=optimal {disk} -- "
@@ -87,13 +94,14 @@ def partition_disk(disk: str, swap_size: int, reserve: int = 0) -> None:
) )
bash_wrapper(create_partitions) bash_wrapper(create_partitions)
logging.info(f"{disk=} successfully partitioned") logger.info(f"{disk=} successfully partitioned")
def create_zfs_pool(pool_disks: Sequence[str], mnt_dir: str) -> None: def create_zfs_pool(pool_disks: Sequence[str], mnt_dir: str) -> None:
"""Create a ZFS pool. """Create a ZFS pool.
Args: Args:
disks (Sequence[str]): A tuple of disks to use for the pool. pool_disks (Sequence[str]): A tuple of disks to use for the pool.
mnt_dir (str): The mount directory. mnt_dir (str): The mount directory.
""" """
if len(pool_disks) <= 0: if len(pool_disks) <= 0:
@@ -125,13 +133,12 @@ def create_zfs_pool(pool_disks: Sequence[str], mnt_dir: str) -> None:
bash_wrapper(zpool_create) bash_wrapper(zpool_create)
zpools = bash_wrapper("zpool list -o name") zpools = bash_wrapper("zpool list -o name")
if "root_pool" not in zpools.splitlines(): if "root_pool" not in zpools.splitlines():
logging.critical("Failed to create root_pool") logger.critical("Failed to create root_pool")
sys.exit(1) sys.exit(1)
def create_zfs_datasets() -> None: def create_zfs_datasets() -> None:
"""Create ZFS datasets.""" """Create ZFS datasets."""
bash_wrapper("zfs create -o canmount=noauto -o reservation=10G root_pool/root") bash_wrapper("zfs create -o canmount=noauto -o reservation=10G root_pool/root")
bash_wrapper("zfs create root_pool/home") bash_wrapper("zfs create root_pool/home")
bash_wrapper("zfs create root_pool/var -o reservation=1G") bash_wrapper("zfs create root_pool/var -o reservation=1G")
@@ -146,7 +153,7 @@ def create_zfs_datasets() -> None:
} }
missing_datasets = expected_datasets.difference(datasets.splitlines()) missing_datasets = expected_datasets.difference(datasets.splitlines())
if missing_datasets: if missing_datasets:
logging.critical(f"Failed to create pools {missing_datasets}") logger.critical(f"Failed to create pools {missing_datasets}")
sys.exit(1) sys.exit(1)
@@ -159,6 +166,8 @@ def get_cpu_manufacturer() -> str:
for line in output.splitlines(): for line in output.splitlines():
if "vendor_id" in line: if "vendor_id" in line:
return id_vendor[line.split(": ")[1].strip()] return id_vendor[line.split(": ")[1].strip()]
error = "Failed to get CPU manufacturer"
raise RuntimeError(error)
def get_boot_drive_id(disk: str) -> str: def get_boot_drive_id(disk: str) -> str:
@@ -167,9 +176,8 @@ def get_boot_drive_id(disk: str) -> str:
return output.splitlines()[1] return output.splitlines()[1]
def create_nix_hardware_file(mnt_dir: str, disks: Sequence[str], encrypt: bool) -> None: def create_nix_hardware_file(mnt_dir: str, disks: Sequence[str], *, encrypt: bool) -> None:
"""Create a NixOS hardware file.""" """Create a NixOS hardware file."""
cpu_manufacturer = get_cpu_manufacturer() cpu_manufacturer = get_cpu_manufacturer()
devices = "" devices = ""
@@ -193,7 +201,15 @@ def create_nix_hardware_file(mnt_dir: str, disks: Sequence[str], encrypt: bool)
' imports = [ (modulesPath + "/installer/scan/not-detected.nix") ];\n\n' ' imports = [ (modulesPath + "/installer/scan/not-detected.nix") ];\n\n'
" boot = {\n" " boot = {\n"
" initrd = {\n" " initrd = {\n"
' availableKernelModules = [ \n "ahci"\n "ehci_pci"\n "nvme"\n "sd_mod"\n "usb_storage"\n "usbhid"\n "xhci_pci"\n ];\n' " availableKernelModules = [ \n"
' "ahci"\n'
' "ehci_pci"\n'
' "nvme"\n'
' "sd_mod"\n'
' "usb_storage"\n'
' "usbhid"\n'
' "xhci_pci"\n'
" ];\n"
" kernelModules = [ ];\n" " kernelModules = [ ];\n"
f" {devices}" f" {devices}"
" };\n" " };\n"
@@ -207,11 +223,18 @@ def create_nix_hardware_file(mnt_dir: str, disks: Sequence[str], encrypt: bool)
' "/nix" = {\n device = "root_pool/nix";\n fsType = "zfs";\n };\n\n' ' "/nix" = {\n device = "root_pool/nix";\n fsType = "zfs";\n };\n\n'
' "/boot" = {\n' ' "/boot" = {\n'
f' device = "/dev/disk/by-uuid/{get_boot_drive_id(disks[0])}";\n' f' device = "/dev/disk/by-uuid/{get_boot_drive_id(disks[0])}";\n'
' fsType = "vfat";\n options = [\n "fmask=0077"\n "dmask=0077"\n ];\n };\n };\n\n' ' fsType = "vfat";\n'
" options = [\n"
' "fmask=0077"\n'
' "dmask=0077"\n'
" ];\n"
" };\n"
" };\n\n"
" swapDevices = [ ];\n\n" " swapDevices = [ ];\n\n"
" networking.useDHCP = lib.mkDefault true;\n\n" " networking.useDHCP = lib.mkDefault true;\n\n"
' nixpkgs.hostPlatform = lib.mkDefault "x86_64-linux";\n' ' nixpkgs.hostPlatform = lib.mkDefault "x86_64-linux";\n'
f" hardware.cpu.{cpu_manufacturer}.updateMicrocode = lib.mkDefault config.hardware.enableRedistributableFirmware;\n" f" hardware.cpu.{cpu_manufacturer}.updateMicrocode = lib.mkDefault "
"config.hardware.enableRedistributableFirmware;\n"
f' networking.hostId = "{host_id}";\n' f' networking.hostId = "{host_id}";\n'
"}\n" "}\n"
) )
@@ -219,7 +242,7 @@ def create_nix_hardware_file(mnt_dir: str, disks: Sequence[str], encrypt: bool)
Path(f"{mnt_dir}/etc/nixos/hardware-configuration.nix").write_text(nix_hardware) Path(f"{mnt_dir}/etc/nixos/hardware-configuration.nix").write_text(nix_hardware)
def install_nixos(mnt_dir: str, disks: Sequence[str], encrypt: bool) -> None: def install_nixos(mnt_dir: str, disks: Sequence[str], *, encrypt: bool) -> None:
"""Install NixOS.""" """Install NixOS."""
bash_wrapper(f"mount -o X-mount.mkdir -t zfs root_pool/root {mnt_dir}") bash_wrapper(f"mount -o X-mount.mkdir -t zfs root_pool/root {mnt_dir}")
bash_wrapper(f"mount -o X-mount.mkdir -t zfs root_pool/home {mnt_dir}/home") bash_wrapper(f"mount -o X-mount.mkdir -t zfs root_pool/home {mnt_dir}/home")
@@ -230,14 +253,16 @@ def install_nixos(mnt_dir: str, disks: Sequence[str], encrypt: bool) -> None:
bash_wrapper(f"mkfs.vfat -n EFI {disk}-part1") bash_wrapper(f"mkfs.vfat -n EFI {disk}-part1")
# set up mirroring afterwards if more than one disk # set up mirroring afterwards if more than one disk
boot_partition = f"mount -t vfat -o fmask=0077,dmask=0077,iocharset=iso8859-1,X-mount.mkdir {disks[0]}-part1 {mnt_dir}/boot" boot_partition = (
f"mount -t vfat -o fmask=0077,dmask=0077,iocharset=iso8859-1,X-mount.mkdir {disks[0]}-part1 {mnt_dir}/boot"
)
bash_wrapper(boot_partition) bash_wrapper(boot_partition)
bash_wrapper(f"nixos-generate-config --root {mnt_dir}") bash_wrapper(f"nixos-generate-config --root {mnt_dir}")
create_nix_hardware_file(mnt_dir, disks, encrypt) create_nix_hardware_file(mnt_dir, disks, encrypt=encrypt)
run(("nixos-install", "--root", mnt_dir), check=True) # noqa: S603 run(("nixos-install", "--root", mnt_dir), check=True)
def installer( def installer(
@@ -247,27 +272,38 @@ def installer(
encrypt_key: str | None, encrypt_key: str | None,
) -> None: ) -> None:
"""Main.""" """Main."""
logging.info("Starting installation") logger.info("Starting installation")
for disk in disks: for disk in disks:
partition_disk(disk, swap_size, reserve) partition_disk(disk, swap_size, reserve)
if encrypt_key: if encrypt_key:
sleep(1) sleep(1)
for command in ( key_input = encrypt_key.encode()
f'printf "{encrypt_key}" | cryptsetup luksFormat --type luks2 {disk}-part2 -', run(
f'printf "{encrypt_key}" | cryptsetup luksOpen {disk}-part2 luks-root-pool-{disk.split("/")[-1]}-part2 -', ("cryptsetup", "luksFormat", "--type", "luks2", f"{disk}-part2", "-"),
): input=key_input,
run(command, shell=True, check=True) check=True,
)
run(
(
"cryptsetup",
"luksOpen",
f"{disk}-part2",
f"luks-root-pool-{disk.split('/')[-1]}-part2",
"-",
),
input=key_input,
check=True,
)
# Fixed mount point for the new system; the installer runs as root on a fresh disk
mnt_dir = "/tmp/nix_install" # noqa: S108 mnt_dir = "/tmp/nix_install" # noqa: S108
Path(mnt_dir).mkdir(parents=True, exist_ok=True) Path(mnt_dir).mkdir(parents=True, exist_ok=True)
if encrypt_key: if encrypt_key:
pool_disks = [ pool_disks = [f"/dev/mapper/luks-root-pool-{disk.split('/')[-1]}-part2" for disk in disks]
f"/dev/mapper/luks-root-pool-{disk.split('/')[-1]}-part2" for disk in disks
]
else: else:
pool_disks = [f"{disk}-part2" for disk in disks] pool_disks = [f"{disk}-part2" for disk in disks]
@@ -275,57 +311,73 @@ def installer(
create_zfs_datasets() create_zfs_datasets()
install_nixos(mnt_dir, disks, encrypt_key) install_nixos(mnt_dir, disks, encrypt=bool(encrypt_key))
logging.info("Installation complete") logger.info("Installation complete")
class Cursor: class Cursor:
def __init__(self): """Track cursor position and constrain movement to screen bounds."""
def __init__(self) -> None:
"""Initialize cursor position and screen dimensions."""
self.x_position = 0 self.x_position = 0
self.y_position = 0 self.y_position = 0
self.height = 0 self.height = 0
self.width = 0 self.width = 0
def set_height(self, height: int): def set_height(self, height: int) -> None:
"""Set the maximum screen height."""
self.height = height self.height = height
def set_width(self, width: int): def set_width(self, width: int) -> None:
"""Set the maximum screen width."""
self.width = width self.width = width
def x_bounce_check(self, cursor: int) -> int: def x_bounce_check(self, cursor: int) -> int:
"""Clamp an x position to the screen width."""
cursor = max(0, cursor) cursor = max(0, cursor)
return min(self.width - 1, cursor) return min(self.width - 1, cursor)
def y_bounce_check(self, cursor: int) -> int: def y_bounce_check(self, cursor: int) -> int:
"""Clamp a y position to the screen height."""
cursor = max(0, cursor) cursor = max(0, cursor)
return min(self.height - 1, cursor) return min(self.height - 1, cursor)
def set_x(self, x: int): def set_x(self, x: int) -> None:
"""Set the cursor x position."""
self.x_position = self.x_bounce_check(x) self.x_position = self.x_bounce_check(x)
def set_y(self, y: int): def set_y(self, y: int) -> None:
"""Set the cursor y position."""
self.y_position = self.y_bounce_check(y) self.y_position = self.y_bounce_check(y)
def get_x(self) -> int: def get_x(self) -> int:
"""Get the cursor x position."""
return self.x_position return self.x_position
def get_y(self) -> int: def get_y(self) -> int:
"""Get the cursor y position."""
return self.y_position return self.y_position
def move_up(self): def move_up(self) -> None:
"""Move the cursor up one row."""
self.set_y(self.y_position - 1) self.set_y(self.y_position - 1)
def move_down(self): def move_down(self) -> None:
"""Move the cursor down one row."""
self.set_y(self.y_position + 1) self.set_y(self.y_position + 1)
def move_left(self): def move_left(self) -> None:
"""Move the cursor left one column."""
self.set_x(self.x_position - 1) self.set_x(self.x_position - 1)
def move_right(self): def move_right(self) -> None:
"""Move the cursor right one column."""
self.set_x(self.x_position + 1) self.set_x(self.x_position + 1)
def navigation(self, key: int) -> None: def navigation(self, key: int) -> None:
"""Move the cursor for a curses navigation key."""
action = { action = {
curses.KEY_DOWN: self.move_down, curses.KEY_DOWN: self.move_down,
curses.KEY_UP: self.move_up, curses.KEY_UP: self.move_up,
@@ -339,7 +391,8 @@ class Cursor:
class State: class State:
"""State class to store the state of the program.""" """State class to store the state of the program."""
def __init__(self): def __init__(self) -> None:
"""Initialize installer menu state."""
self.key = 0 self.key = 0
self.cursor = Cursor() self.cursor = Cursor()
@@ -357,11 +410,9 @@ class State:
def get_device(raw_device: str) -> dict[str, str]: def get_device(raw_device: str) -> dict[str, str]:
"""Parse an lsblk key-value device row."""
raw_device_components = raw_device.split(" ") raw_device_components = raw_device.split(" ")
return { return {thing.split("=")[0].lower(): thing.split("=")[1].strip('"') for thing in raw_device_components}
thing.split("=")[0].lower(): thing.split("=")[1].strip('"')
for thing in raw_device_components
}
def get_devices() -> list[dict[str, str]]: def get_devices() -> list[dict[str, str]]:
@@ -373,6 +424,7 @@ def get_devices() -> list[dict[str, str]]:
def get_device_id_mapping() -> dict[str, set[str]]: def get_device_id_mapping() -> dict[str, set[str]]:
"""Get a list of device ids. """Get a list of device ids.
Returns: Returns:
list[str]: the list of device ids list[str]: the list of device ids
""" """
@@ -387,9 +439,8 @@ def get_device_id_mapping() -> dict[str, set[str]]:
return device_id_mapping return device_id_mapping
def calculate_device_menu_padding( def calculate_device_menu_padding(devices: list[dict[str, str]], column: str, padding: int = 0) -> int:
devices: list[dict[str, str]], column: str, padding: int = 0 """Calculate the width needed for a device menu column."""
) -> int:
return max(len(device[column]) for device in devices) + padding return max(len(device[column]) for device in devices) + padding
@@ -401,6 +452,7 @@ def draw_device_ids(
menu_width: list[int], menu_width: list[int],
device_ids: set[str], device_ids: set[str],
) -> tuple[State, int]: ) -> tuple[State, int]:
"""Draw selectable device IDs for a device row."""
for device_id in sorted(device_ids): for device_id in sorted(device_ids):
row_number = row_number + 1 row_number = row_number + 1
if row_number == state.cursor.get_y() and state.cursor.get_x() in menu_width: if row_number == state.cursor.get_y() and state.cursor.get_x() in menu_width:
@@ -429,8 +481,9 @@ def draw_device_menu(
state: State, state: State,
menu_start_y: int = 0, menu_start_y: int = 0,
menu_start_x: int = 0, menu_start_x: int = 0,
) -> State: ) -> tuple[State, int]:
"""draw the device menu and handle user input """Draw the device menu and handle user input.
Args: Args:
std_screen (curses.window): the curses window to draw on std_screen (curses.window): the curses window to draw on
devices (list[dict[str, str]]): the list of devices to draw devices (list[dict[str, str]]): the list of devices to draw
@@ -438,6 +491,7 @@ def draw_device_menu(
state (State): the state object to update state (State): the state object to update
menu_start_y (int, optional): the y position to start drawing the menu. Defaults to 0. menu_start_y (int, optional): the y position to start drawing the menu. Defaults to 0.
menu_start_x (int, optional): the x position to start drawing the menu. Defaults to 0. menu_start_x (int, optional): the x position to start drawing the menu. Defaults to 0.
Returns: Returns:
State: the updated state object State: the updated state object
""" """
@@ -448,7 +502,9 @@ def draw_device_menu(
type_padding = calculate_device_menu_padding(devices, "type", padding) type_padding = calculate_device_menu_padding(devices, "type", padding)
mountpoints_padding = calculate_device_menu_padding(devices, "mountpoints", padding) mountpoints_padding = calculate_device_menu_padding(devices, "mountpoints", padding)
device_header = f"{'Name':{name_padding}}{'Size':{size_padding}}{'Type':{type_padding}}{'Mountpoints':{mountpoints_padding}}" device_header = (
f"{'Name':{name_padding}}{'Size':{size_padding}}{'Type':{type_padding}}{'Mountpoints':{mountpoints_padding}}"
)
menu_width = range(menu_start_x, len(device_header) + menu_start_x) menu_width = range(menu_start_x, len(device_header) + menu_start_x)
@@ -481,8 +537,9 @@ def draw_device_menu(
def debug_menu(std_screen: curses.window, key: int) -> None: def debug_menu(std_screen: curses.window, key: int) -> None:
"""Draw debug information for the current curses screen."""
height, width = std_screen.getmaxyx() height, width = std_screen.getmaxyx()
width_height = "Width: {}, Height: {}".format(width, height) width_height = f"Width: {width}, Height: {height}"
std_screen.addstr(height - 4, 0, width_height, curses.color_pair(5)) std_screen.addstr(height - 4, 0, width_height, curses.color_pair(5))
key_pressed = f"Last key pressed: {key}"[: width - 1] key_pressed = f"Last key pressed: {key}"[: width - 1]
@@ -490,7 +547,7 @@ def debug_menu(std_screen: curses.window, key: int) -> None:
key_pressed = "No key press detected..."[: width - 1] key_pressed = "No key press detected..."[: width - 1]
std_screen.addstr(height - 3, 0, key_pressed) std_screen.addstr(height - 3, 0, key_pressed)
for i in range(0, 8): for i in range(8):
std_screen.addstr(height - 2, i * 3, f"{i}██", curses.color_pair(i)) std_screen.addstr(height - 2, i * 3, f"{i}██", curses.color_pair(i))
@@ -500,12 +557,11 @@ def status_bar(
width: int, width: int,
height: int, height: int,
) -> None: ) -> None:
"""Draw the footer status bar."""
std_screen.attron(curses.A_REVERSE) std_screen.attron(curses.A_REVERSE)
std_screen.attron(curses.color_pair(3)) std_screen.attron(curses.color_pair(3))
status_bar = ( status_bar = f"Press 'q' to exit | STATUS BAR | Pos: {cursor.get_x()}, {cursor.get_y()}"
f"Press 'q' to exit | STATUS BAR | Pos: {cursor.get_x()}, {cursor.get_y()}"
)
std_screen.addstr(height - 1, 0, status_bar) std_screen.addstr(height - 1, 0, status_bar)
std_screen.addstr(height - 1, len(status_bar), " " * (width - len(status_bar) - 1)) std_screen.addstr(height - 1, len(status_bar), " " * (width - len(status_bar) - 1))
@@ -514,13 +570,15 @@ def status_bar(
def set_color() -> None: def set_color() -> None:
"""Initialize curses color pairs."""
curses.start_color() curses.start_color()
curses.use_default_colors() curses.use_default_colors()
for i in range(0, curses.COLORS): for i in range(curses.COLORS):
curses.init_pair(i + 1, i, -1) curses.init_pair(i + 1, i, -1)
def get_text_input(std_screen: curses.window, prompt: str, y: int, x: int) -> str: def get_text_input(std_screen: curses.window, prompt: str, y: int, x: int) -> str:
"""Read text input from a curses screen."""
curses.echo() curses.echo()
std_screen.addstr(y, x, prompt) std_screen.addstr(y, x, prompt)
input_str = "" input_str = ""
@@ -528,10 +586,10 @@ def get_text_input(std_screen: curses.window, prompt: str, y: int, x: int) -> st
key = std_screen.getch() key = std_screen.getch()
if key == ord("\n"): if key == ord("\n"):
break break
elif key == 27: # ESC key if key == ESCAPE_KEY:
input_str = "" input_str = ""
break break
elif key in (curses.KEY_BACKSPACE, ord("\b"), 127): if key in (curses.KEY_BACKSPACE, ord("\b"), 127):
input_str = input_str[:-1] input_str = input_str[:-1]
std_screen.addstr(y, x + len(prompt), input_str + " ") std_screen.addstr(y, x + len(prompt), input_str + " ")
else: else:
@@ -546,6 +604,7 @@ def swap_size_input(
state: State, state: State,
swap_offset: int, swap_offset: int,
) -> State: ) -> State:
"""Handle swap size input."""
swap_size_text = "Swap size (GB): " swap_size_text = "Swap size (GB): "
std_screen.addstr(swap_offset, 0, f"{swap_size_text}{state.swap_size}") std_screen.addstr(swap_offset, 0, f"{swap_size_text}{state.swap_size}")
if state.key == ord("\n") and state.cursor.get_y() == swap_offset: if state.key == ord("\n") and state.cursor.get_y() == swap_offset:
@@ -557,9 +616,7 @@ def swap_size_input(
state.swap_size = int(swap_size_str) state.swap_size = int(swap_size_str)
state.show_swap_input = False state.show_swap_input = False
except ValueError: except ValueError:
std_screen.addstr( std_screen.addstr(swap_offset, 0, "Invalid input. Press any key to continue.")
swap_offset, 0, "Invalid input. Press any key to continue."
)
std_screen.getch() std_screen.getch()
state.show_swap_input = False state.show_swap_input = False
@@ -571,22 +628,19 @@ def reserve_size_input(
state: State, state: State,
reserve_offset: int, reserve_offset: int,
) -> State: ) -> State:
"""Handle reserve size input."""
reserve_size_text = "reserve size (GB): " reserve_size_text = "reserve size (GB): "
std_screen.addstr(reserve_offset, 0, f"{reserve_size_text}{state.reserve_size}") std_screen.addstr(reserve_offset, 0, f"{reserve_size_text}{state.reserve_size}")
if state.key == ord("\n") and state.cursor.get_y() == reserve_offset: if state.key == ord("\n") and state.cursor.get_y() == reserve_offset:
state.show_reserve_input = True state.show_reserve_input = True
if state.show_reserve_input: if state.show_reserve_input:
reserve_size_str = get_text_input( reserve_size_str = get_text_input(std_screen, reserve_size_text, reserve_offset, 0)
std_screen, reserve_size_text, reserve_offset, 0
)
try: try:
state.reserve_size = int(reserve_size_str) state.reserve_size = int(reserve_size_str)
state.show_reserve_input = False state.show_reserve_input = False
except ValueError: except ValueError:
std_screen.addstr( std_screen.addstr(reserve_offset, 0, "Invalid input. Press any key to continue.")
reserve_offset, 0, "Invalid input. Press any key to continue."
)
std_screen.getch() std_screen.getch()
state.show_reserve_input = False state.show_reserve_input = False
@@ -594,9 +648,11 @@ def reserve_size_input(
def draw_menu(std_screen: curses.window) -> State: def draw_menu(std_screen: curses.window) -> State:
"""draw the menu and handle user input """Draw the menu and handle user input.
Args: Args:
std_screen (curses.window): the curses window to draw on std_screen (curses.window): the curses window to draw on
Returns: Returns:
State: the state object State: the state object
""" """
@@ -656,17 +712,18 @@ def draw_menu(std_screen: curses.window) -> State:
def main() -> None: def main() -> None:
"""Run the installer menu and start installation."""
configure_logger("DEBUG") configure_logger("DEBUG")
state = curses.wrapper(draw_menu) state = curses.wrapper(draw_menu)
encrypt_key = getenv("ENCRYPT_KEY") encrypt_key = getenv("ENCRYPT_KEY")
logging.info("installing_nixos") logger.info("installing_nixos")
logging.info(f"disks: {state.selected_device_ids}") logger.info(f"disks: {state.selected_device_ids}")
logging.info(f"swap_size: {state.swap_size}") logger.info(f"swap_size: {state.swap_size}")
logging.info(f"reserve: {state.reserve_size}") logger.info(f"reserve: {state.reserve_size}")
logging.info(f"encrypted: {bool(encrypt_key)}") logger.info(f"encrypted: {bool(encrypt_key)}")
sleep(3) sleep(3)
-2
View File
@@ -1,9 +1,7 @@
"""ORM package exports.""" """ORM package exports."""
from python.orm.richie.base import RichieBase from python.orm.richie.base import RichieBase
from python.orm.van_inventory.base import VanInventoryBase
__all__ = [ __all__ = [
"RichieBase", "RichieBase",
"VanInventoryBase",
] ]
+24 -2
View File
@@ -31,8 +31,24 @@ def get_connection_info(name: str) -> tuple[str, str, str, str, str | None]:
return cast("tuple[str, str, str, str, str | None]", (database, host, port, username, password)) return cast("tuple[str, str, str, str, str | None]", (database, host, port, username, password))
def get_postgres_engine(*, name: str = "POSTGRES", pool_pre_ping: bool = True) -> Engine: def get_postgres_engine(
"""Create a SQLAlchemy engine from environment variables.""" *,
name: str = "POSTGRES",
pool_pre_ping: bool = True,
vector_engine: bool = False,
) -> Engine:
"""Create a SQLAlchemy engine from environment variables.
Args:
name (str, optional): The name of the environment variable prefix. Defaults to "POSTGRES".
pool_pre_ping (bool, optional): Whether to ping the database before each connection. Defaults to True.
This fixes the issue of trying to use a conection that has timed out on the database side.
vector_engine (bool, optional): Whether to use the vector search schema. Defaults to False.
This updates the search path the incldued the vecore types and operators.
Returns:
Engine: The SQLAlchemy engine.
"""
database, host, port, username, password = get_connection_info(name) database, host, port, username, password = get_connection_info(name)
url = URL.create( url = URL.create(
@@ -44,8 +60,14 @@ def get_postgres_engine(*, name: str = "POSTGRES", pool_pre_ping: bool = True) -
database=database, database=database,
) )
connect_args = {}
# There more better way to do this is with separate PG account and a dedicated vector schema for the vector types
if vector_engine:
connect_args["options"] = "-csearch_path=main,public"
return create_engine( 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,
) )
+16
View File
@@ -11,6 +11,15 @@ from python.orm.richie.contact import (
Need, Need,
RelationshipType, RelationshipType,
) )
from python.orm.richie.ebook import (
EbookChapter,
EbookChunk,
EbookChunkEmbedding1024,
EbookChunkEmbedding2560,
EbookChunkEmbedding4096,
EbookEmbeddingModel,
EbookSource,
)
__all__ = [ __all__ = [
"Audiobook", "Audiobook",
@@ -19,6 +28,13 @@ __all__ = [
"Contact", "Contact",
"ContactNeed", "ContactNeed",
"ContactRelationship", "ContactRelationship",
"EbookChapter",
"EbookChunk",
"EbookChunkEmbedding1024",
"EbookChunkEmbedding2560",
"EbookChunkEmbedding4096",
"EbookEmbeddingModel",
"EbookSource",
"Need", "Need",
"RelationshipType", "RelationshipType",
"RichieBase", "RichieBase",
+4 -4
View File
@@ -2,7 +2,7 @@
from __future__ import annotations from __future__ import annotations
from sqlalchemy import ForeignKey, String, UniqueConstraint from sqlalchemy import ForeignKey, UniqueConstraint
from sqlalchemy.orm import Mapped, mapped_column, relationship from sqlalchemy.orm import Mapped, mapped_column, relationship
from python.orm.richie.base import TableBase from python.orm.richie.base import TableBase
@@ -14,7 +14,7 @@ class AudiobookAuthor(TableBase):
__tablename__ = "audiobook_author" __tablename__ = "audiobook_author"
__table_args__ = (UniqueConstraint("name"),) __table_args__ = (UniqueConstraint("name"),)
name: Mapped[str] = mapped_column(String, unique=True) name: Mapped[str]
books: Mapped[list[Audiobook]] = relationship("Audiobook", back_populates="author") books: Mapped[list[Audiobook]] = relationship("Audiobook", back_populates="author")
series: Mapped[list[AudiobookSeries]] = relationship("AudiobookSeries", back_populates="author") series: Mapped[list[AudiobookSeries]] = relationship("AudiobookSeries", back_populates="author")
@@ -26,7 +26,7 @@ class AudiobookSeries(TableBase):
__tablename__ = "audiobook_series" __tablename__ = "audiobook_series"
__table_args__ = (UniqueConstraint("author_id", "name"),) __table_args__ = (UniqueConstraint("author_id", "name"),)
name: Mapped[str] = mapped_column(String) name: Mapped[str]
author_id: Mapped[int] = mapped_column(ForeignKey("main.audiobook_author.id", ondelete="CASCADE")) author_id: Mapped[int] = mapped_column(ForeignKey("main.audiobook_author.id", ondelete="CASCADE"))
author: Mapped[AudiobookAuthor] = relationship("AudiobookAuthor", back_populates="series") author: Mapped[AudiobookAuthor] = relationship("AudiobookAuthor", back_populates="series")
@@ -46,7 +46,7 @@ class Audiobook(TableBase):
), ),
) )
title: Mapped[str] = mapped_column(String) title: Mapped[str]
author_id: Mapped[int] = mapped_column(ForeignKey("main.audiobook_author.id", ondelete="CASCADE")) author_id: Mapped[int] = mapped_column(ForeignKey("main.audiobook_author.id", ondelete="CASCADE"))
series_id: Mapped[int | None] = mapped_column(ForeignKey("main.audiobook_series.id", ondelete="SET NULL")) series_id: Mapped[int | None] = mapped_column(ForeignKey("main.audiobook_series.id", ondelete="SET NULL"))
series_index: Mapped[float] = mapped_column(default=0.0) series_index: Mapped[float] = mapped_column(default=0.0)
+138
View File
@@ -0,0 +1,138 @@
"""EPUB search models."""
from __future__ import annotations
from datetime import datetime
from pgvector.sqlalchemy import Vector
from sqlalchemy import BigInteger, Boolean, DateTime, ForeignKey, Index, String, UniqueConstraint
from sqlalchemy.orm import Mapped, mapped_column, relationship
from python.orm.richie.base import TableBase, TableBaseBig
class EbookSource(TableBase):
"""One indexed EPUB file."""
__tablename__ = "ebook_source"
__table_args__ = (
UniqueConstraint("file_path"),
UniqueConstraint("file_sha256"),
)
title: Mapped[str]
author: Mapped[str | None]
language: Mapped[str | None]
publisher: Mapped[str | None]
identifier: Mapped[str | None]
file_path: Mapped[str]
file_sha256: Mapped[str] = mapped_column(String(64))
file_mtime: Mapped[datetime] = mapped_column(DateTime(timezone=True))
file_size: Mapped[int] = mapped_column(BigInteger)
chapters: Mapped[list[EbookChapter]] = relationship(
"EbookChapter",
back_populates="source",
cascade="all, delete-orphan",
passive_deletes=True,
)
chunks: Mapped[list[EbookChunk]] = relationship(
"EbookChunk",
back_populates="source",
cascade="all, delete-orphan",
passive_deletes=True,
)
class EbookChapter(TableBase):
"""A chapter or spine document inside an EPUB."""
__tablename__ = "ebook_chapter"
__table_args__ = (UniqueConstraint("source_id", "spine_index"),)
source_id: Mapped[int] = mapped_column(ForeignKey("main.ebook_source.id", ondelete="CASCADE"))
spine_index: Mapped[int]
title: Mapped[str | None]
href: Mapped[str | None]
source: Mapped[EbookSource] = relationship("EbookSource", back_populates="chapters")
chunks: Mapped[list[EbookChunk]] = relationship(
"EbookChunk",
back_populates="chapter",
cascade="all, delete-orphan",
passive_deletes=True,
)
class EbookChunk(TableBaseBig):
"""A searchable text chunk."""
__tablename__ = "ebook_chunk"
__table_args__ = (
UniqueConstraint("source_id", "chunk_index", name="uq_ebook_chunk_source_id_chunk_index"),
UniqueConstraint("source_id", "content_sha256", name="uq_ebook_chunk_source_id_content_sha256"),
)
source_id: Mapped[int] = mapped_column(ForeignKey("main.ebook_source.id", ondelete="CASCADE"))
chapter_id: Mapped[int | None] = mapped_column(ForeignKey("main.ebook_chapter.id", ondelete="SET NULL"))
chunk_index: Mapped[int]
text: Mapped[str]
token_start: Mapped[int]
token_count: Mapped[int]
page_label: Mapped[str | None]
content_sha256: Mapped[str] = mapped_column(String(64))
search_text: Mapped[str]
source: Mapped[EbookSource] = relationship("EbookSource", back_populates="chunks")
chapter: Mapped[EbookChapter | None] = relationship("EbookChapter", back_populates="chunks")
class EbookEmbeddingModel(TableBase):
"""A supported embedding model."""
__tablename__ = "ebook_embedding_model"
name: Mapped[str] = mapped_column(String, unique=True)
dimension: Mapped[int]
is_default: Mapped[bool] = mapped_column(Boolean, default=False)
class EbookChunkEmbedding1024(TableBaseBig):
"""1024-dimensional chunk embedding."""
__tablename__ = "ebook_chunk_embedding_1024"
__table_args__ = (
UniqueConstraint("chunk_id", "model_id"),
Index(
"ix_ebook_chunk_embedding_1024_embedding_cosine",
"embedding",
postgresql_using="hnsw",
postgresql_ops={"embedding": "vector_cosine_ops"},
),
)
chunk_id: Mapped[int] = mapped_column(ForeignKey("main.ebook_chunk.id", ondelete="CASCADE"))
model_id: Mapped[int] = mapped_column(ForeignKey("main.ebook_embedding_model.id", ondelete="CASCADE"))
embedding: Mapped[list[float]] = mapped_column(Vector(1024))
class EbookChunkEmbedding2560(TableBaseBig):
"""2560-dimensional chunk embedding."""
__tablename__ = "ebook_chunk_embedding_2560"
__table_args__ = (UniqueConstraint("chunk_id", "model_id"),)
chunk_id: Mapped[int] = mapped_column(ForeignKey("main.ebook_chunk.id", ondelete="CASCADE"))
model_id: Mapped[int] = mapped_column(ForeignKey("main.ebook_embedding_model.id", ondelete="CASCADE"))
embedding: Mapped[list[float]] = mapped_column(Vector(2560))
class EbookChunkEmbedding4096(TableBaseBig):
"""4096-dimensional chunk embedding."""
__tablename__ = "ebook_chunk_embedding_4096"
__table_args__ = (UniqueConstraint("chunk_id", "model_id"),)
chunk_id: Mapped[int] = mapped_column(ForeignKey("main.ebook_chunk.id", ondelete="CASCADE"))
model_id: Mapped[int] = mapped_column(ForeignKey("main.ebook_embedding_model.id", ondelete="CASCADE"))
embedding: Mapped[list[float]] = mapped_column(Vector(4096))
-1
View File
@@ -1 +0,0 @@
"""Van inventory database ORM exports."""
-39
View File
@@ -1,39 +0,0 @@
"""Van inventory database ORM base."""
from __future__ import annotations
from datetime import datetime
from sqlalchemy import DateTime, MetaData, func
from sqlalchemy.ext.declarative import AbstractConcreteBase
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
from python.orm.common import NAMING_CONVENTION
class VanInventoryBase(DeclarativeBase):
"""Base class for van_inventory database ORM models."""
schema_name = "main"
metadata = MetaData(
schema=schema_name,
naming_convention=NAMING_CONVENTION,
)
class VanTableBase(AbstractConcreteBase, VanInventoryBase):
"""Abstract concrete base for van_inventory tables with IDs and timestamps."""
__abstract__ = True
id: Mapped[int] = mapped_column(primary_key=True)
created: Mapped[datetime] = mapped_column(
DateTime(timezone=True),
server_default=func.now(),
)
updated: Mapped[datetime] = mapped_column(
DateTime(timezone=True),
server_default=func.now(),
onupdate=func.now(),
)
-46
View File
@@ -1,46 +0,0 @@
"""Van inventory ORM models."""
from __future__ import annotations
from sqlalchemy import ForeignKey, UniqueConstraint
from sqlalchemy.orm import Mapped, mapped_column, relationship
from python.orm.van_inventory.base import VanTableBase
class Item(VanTableBase):
"""A food item in the van."""
__tablename__ = "items"
name: Mapped[str] = mapped_column(unique=True)
quantity: Mapped[float] = mapped_column(default=0)
unit: Mapped[str]
category: Mapped[str | None]
meal_ingredients: Mapped[list[MealIngredient]] = relationship(back_populates="item")
class Meal(VanTableBase):
"""A meal that can be made from items in the van."""
__tablename__ = "meals"
name: Mapped[str] = mapped_column(unique=True)
instructions: Mapped[str | None]
ingredients: Mapped[list[MealIngredient]] = relationship(back_populates="meal")
class MealIngredient(VanTableBase):
"""Links a meal to the items it requires, with quantities."""
__tablename__ = "meal_ingredients"
__table_args__ = (UniqueConstraint("meal_id", "item_id"),)
meal_id: Mapped[int] = mapped_column(ForeignKey("meals.id"))
item_id: Mapped[int] = mapped_column(ForeignKey("items.id"))
quantity_needed: Mapped[float]
meal: Mapped[Meal] = relationship(back_populates="ingredients")
item: Mapped[Item] = relationship(back_populates="meal_ingredients")
-1
View File
@@ -1 +0,0 @@
game_data/
-1
View File
@@ -1 +0,0 @@
"""init."""
-675
View File
@@ -1,675 +0,0 @@
"""Base logic for the Splendor game."""
from __future__ import annotations
import itertools
import json
import random
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Literal, Protocol
if TYPE_CHECKING:
from collections.abc import Sequence
from pathlib import Path
GemColor = Literal["white", "blue", "green", "red", "black", "gold"]
GEM_COLORS: tuple[GemColor, ...] = (
"white",
"blue",
"green",
"red",
"black",
"gold",
)
BASE_COLORS: tuple[GemColor, ...] = (
"white",
"blue",
"green",
"red",
"black",
)
GEM_ORDER: list[GemColor] = list(GEM_COLORS)
GEM_INDEX: dict[GemColor, int] = {c: i for i, c in enumerate(GEM_ORDER)}
BASE_INDEX: dict[GemColor, int] = {c: i for i, c in enumerate(BASE_COLORS)}
@dataclass(frozen=True)
class Card:
"""Development card: gives points + a permanent gem discount."""
tier: int
points: int
color: GemColor
cost: dict[GemColor, int]
@dataclass(frozen=True)
class Noble:
"""Noble tile: gives points if you have enough bonuses."""
name: str
points: int
requirements: dict[GemColor, int]
@dataclass
class PlayerState:
"""State of a player in the game."""
strategy: Strategy
tokens: dict[GemColor, int] = field(default_factory=lambda: dict.fromkeys(GEM_COLORS, 0))
discounts: dict[GemColor, int] = field(default_factory=lambda: dict.fromkeys(GEM_COLORS, 0))
cards: list[Card] = field(default_factory=list)
reserved: list[Card] = field(default_factory=list)
nobles: list[Noble] = field(default_factory=list)
card_score: int = 0
noble_score: int = 0
def total_tokens(self) -> int:
"""Total tokens in player's bank."""
return sum(self.tokens.values())
def add_noble(self, noble: Noble) -> None:
"""Add a noble to the player."""
self.nobles.append(noble)
self.noble_score = sum(noble.points for noble in self.nobles)
def add_card(self, card: Card) -> None:
"""Add a card to the player."""
self.cards.append(card)
self.card_score = sum(card.points for card in self.cards)
@property
def score(self) -> int:
"""Total points in player's cards + nobles."""
return self.card_score + self.noble_score
def can_afford(self, card: Card) -> bool:
"""Check if player can afford card, using discounts + gold."""
missing = 0
gold = self.tokens["gold"]
for color, cost in card.cost.items():
missing += max(0, cost - self.discounts.get(color, 0) - self.tokens.get(color, 0))
if missing > gold:
return False
return True
def pay_for_card(self, card: Card) -> dict[GemColor, int]:
"""Pay tokens for card, move card to tableau, return payment for bank."""
if not self.can_afford(card):
msg = f"cannot afford card {card}"
raise ValueError(msg)
payment: dict[GemColor, int] = dict.fromkeys(GEM_COLORS, 0)
gold_available = self.tokens["gold"]
for color in BASE_COLORS:
cost = card.cost.get(color, 0)
effective_cost = max(0, cost - self.discounts.get(color, 0))
use = min(self.tokens[color], effective_cost)
self.tokens[color] -= use
payment[color] += use
remaining = effective_cost - use
if remaining > 0:
use_gold = min(gold_available, remaining)
gold_available -= use_gold
self.tokens["gold"] -= use_gold
payment["gold"] += use_gold
self.add_card(card)
self.discounts[card.color] += 1
return payment
def get_default_starting_tokens(player_count: int) -> dict[GemColor, int]:
"""get_default_starting_tokens."""
token_count = (player_count * player_count - 3 * player_count + 10) // 2
return {
"white": token_count,
"blue": token_count,
"green": token_count,
"red": token_count,
"black": token_count,
"gold": 5,
}
@dataclass
class GameConfig:
"""Game configuration: gems, bank, cards, nobles, etc."""
win_score: int = 15
table_cards_per_tier: int = 4
reserve_limit: int = 3
token_limit: int = 10
turn_limit: int = 1000
minimum_tokens_to_buy_2: int = 4
max_token_take: int = 3
cards: list[Card] = field(default_factory=list)
nobles: list[Noble] = field(default_factory=list)
class GameState:
"""Game state: players, bank, decks, table, available nobles, etc."""
def __init__(
self,
config: GameConfig,
players: list[PlayerState],
bank: dict[GemColor, int],
decks_by_tier: dict[int, list[Card]],
table_by_tier: dict[int, list[Card]],
available_nobles: list[Noble],
) -> None:
"""Game state."""
self.config = config
self.players = players
self.bank = bank
self.decks_by_tier = decks_by_tier
self.table_by_tier = table_by_tier
self.available_nobles = available_nobles
self.noble_min_requirements = 0
self.get_noble_min_requirements()
self.current_player_index = 0
self.finished = False
def get_noble_min_requirements(self) -> None:
"""Find the minimum requirement for all available nobles."""
test = 0
for noble in self.available_nobles:
test = max(test, min(foo for foo in noble.requirements.values()))
self.noble_min_requirements = test
def next_player(self) -> None:
"""Advance to the next player."""
self.current_player_index = (self.current_player_index + 1) % len(self.players)
@property
def current_player(self) -> PlayerState:
"""Current player."""
return self.players[self.current_player_index]
def refill_table(self) -> None:
"""Refill face-up cards from decks."""
for tier, deck in self.decks_by_tier.items():
table = self.table_by_tier[tier]
while len(table) < self.config.table_cards_per_tier and deck:
table.append(deck.pop())
def check_winner_simple(self) -> PlayerState | None:
"""Simplified: end immediately when someone hits win_score."""
eligible = [player for player in self.players if player.score >= self.config.win_score]
if not eligible:
return None
eligible.sort(
key=lambda p: (p.score, -len(p.cards)),
reverse=True,
)
self.finished = True
return eligible[0]
class Action:
"""Marker protocol for actions."""
@dataclass
class TakeDifferent(Action):
"""Take up to 3 different gem colors."""
colors: list[GemColor]
@dataclass
class TakeDouble(Action):
"""Take two of the same color."""
color: GemColor
@dataclass
class BuyCard(Action):
"""Buy a face-up card."""
tier: int
index: int
@dataclass
class BuyCardReserved(Action):
"""Buy a face-up card."""
index: int
@dataclass
class ReserveCard(Action):
"""Reserve a face-up card."""
tier: int
index: int | None = None
from_deck: bool = False
class Strategy(Protocol):
"""Implement this to make a bot or human controller."""
def __init__(self, name: str) -> None:
"""Initialize a strategy."""
self.name = name
def choose_action(self, game: GameState, player: PlayerState) -> Action | None:
"""Return an Action, or None to concede/end."""
raise NotImplementedError
def choose_discard(
self,
game: GameState, # noqa: ARG002
player: PlayerState,
excess: int,
) -> dict[GemColor, int]:
"""Called if player has more than token_limit tokens after an action.
Default: naive auto-discard.
"""
return auto_discard_tokens(player, excess)
def choose_noble(
self,
game: GameState, # noqa: ARG002
player: PlayerState, # noqa: ARG002
nobles: list[Noble],
) -> Noble:
"""Called if player qualifies for multiple nobles. Default: first."""
return nobles[0]
def auto_discard_tokens(player: PlayerState, excess: int) -> dict[GemColor, int]:
"""Very dumb discard logic: discard from colors you have the most of."""
to_discard: dict[GemColor, int] = dict.fromkeys(GEM_COLORS, 0)
remaining = excess
while remaining > 0:
color = max(player.tokens, key=lambda c: player.tokens[c])
if player.tokens[color] == 0:
break
player.tokens[color] -= 1
to_discard[color] += 1
remaining -= 1
return to_discard
def enforce_token_limit(
game: GameState,
strategy: Strategy,
player: PlayerState,
) -> None:
"""If player has more than token_limit tokens, force discards."""
limit = game.config.token_limit
total = player.total_tokens()
if total <= limit:
return
excess = total - limit
discards = strategy.choose_discard(game, player, excess)
for color, amount in discards.items():
available = player.tokens[color]
to_remove = min(amount, available)
if to_remove <= 0:
continue
player.tokens[color] -= to_remove
game.bank[color] += to_remove
remaining = player.total_tokens() - limit
if remaining > 0:
auto = auto_discard_tokens(player, remaining)
for color, amount in auto.items():
game.bank[color] += amount
def _check_nobles_for_player(player: PlayerState, noble: Noble) -> bool:
# this rule is slower
for color, cost in noble.requirements.items(): # noqa: SIM110
if player.discounts[color] < cost:
return False
return True
def check_nobles_for_player(
game: GameState,
strategy: Strategy,
player: PlayerState,
) -> None:
"""Award at most one noble to player if they qualify."""
if game.noble_min_requirements > max(player.discounts.values()):
return
candidates = [noble for noble in game.available_nobles if _check_nobles_for_player(player, noble)]
if not candidates:
return
chosen = candidates[0] if len(candidates) == 1 else strategy.choose_noble(game, player, candidates)
if chosen not in game.available_nobles:
return
game.available_nobles.remove(chosen)
game.get_noble_min_requirements()
player.add_noble(chosen)
def apply_take_different(game: GameState, strategy: Strategy, action: TakeDifferent) -> None:
"""Mutate game state according to action."""
player = game.current_player
colors = [color for color in action.colors if color in BASE_COLORS and game.bank[color] > 0]
if not (1 <= len(colors) <= game.config.max_token_take):
return
for color in colors:
game.bank[color] -= 1
player.tokens[color] += 1
enforce_token_limit(game, strategy, player)
def apply_take_double(game: GameState, strategy: Strategy, action: TakeDouble) -> None:
"""Mutate game state according to action."""
player = game.current_player
color = action.color
if color not in BASE_COLORS:
return
if game.bank[color] < game.config.minimum_tokens_to_buy_2:
return
game.bank[color] -= 2
player.tokens[color] += 2
enforce_token_limit(game, strategy, player)
def apply_buy_card(game: GameState, _strategy: Strategy, action: BuyCard) -> None:
"""Mutate game state according to action."""
player = game.current_player
row = game.table_by_tier.get(action.tier)
if row is None or not (0 <= action.index < len(row)):
return
card = row[action.index]
if not player.can_afford(card):
return
row.pop(action.index)
payment = player.pay_for_card(card)
for color, amount in payment.items():
game.bank[color] += amount
game.refill_table()
def apply_buy_card_reserved(game: GameState, _strategy: Strategy, action: BuyCardReserved) -> None:
"""Mutate game state according to action."""
player = game.current_player
if not (0 <= action.index < len(player.reserved)):
return
card = player.reserved[action.index]
if not player.can_afford(card):
return
player.reserved.pop(action.index)
payment = player.pay_for_card(card)
for color, amount in payment.items():
game.bank[color] += amount
def apply_reserve_card(game: GameState, strategy: Strategy, action: ReserveCard) -> None:
"""Mutate game state according to action."""
player = game.current_player
if len(player.reserved) >= game.config.reserve_limit:
return
card: Card | None = None
if action.from_deck:
deck = game.decks_by_tier.get(action.tier)
if deck:
card = deck.pop()
else:
row = game.table_by_tier.get(action.tier)
if row is None:
return
if action.index is None or not (0 <= action.index < len(row)):
return
card = row.pop(action.index)
game.refill_table()
if card is None:
return
player.reserved.append(card)
if game.bank["gold"] > 0:
game.bank["gold"] -= 1
player.tokens["gold"] += 1
enforce_token_limit(game, strategy, player)
def apply_action(game: GameState, strategy: Strategy, action: Action) -> None:
"""Mutate game state according to action."""
actions = {
TakeDifferent: apply_take_different,
TakeDouble: apply_take_double,
BuyCard: apply_buy_card,
ReserveCard: apply_reserve_card,
BuyCardReserved: apply_buy_card_reserved,
}
action_func = actions.get(type(action))
if action_func is None:
msg = f"Unknown action type: {type(action)}"
raise ValueError(msg)
action_func(game, strategy, action)
# not sure how to simplify this yet
def get_legal_actions( # noqa: C901
game: GameState,
player: PlayerState | None = None,
) -> list[Action]:
"""Enumerate all syntactically legal actions for the given player.
This enforces:
- token-taking rules
- reserve limits
- affordability for buys
"""
if player is None:
player = game.players[game.current_player_index]
actions: list[Action] = []
colors_available = [c for c in BASE_COLORS if game.bank[c] > 0]
for r in (1, 2, 3):
actions.extend(TakeDifferent(colors=list(combo)) for combo in itertools.combinations(colors_available, r))
actions.extend(
TakeDouble(color=color) for color in BASE_COLORS if game.bank[color] >= game.config.minimum_tokens_to_buy_2
)
for tier, row in game.table_by_tier.items():
for idx, card in enumerate(row):
if player.can_afford(card):
actions.append(BuyCard(tier=tier, index=idx))
for idx, card in enumerate(player.reserved):
if player.can_afford(card):
actions.append(BuyCardReserved(index=idx))
if len(player.reserved) < game.config.reserve_limit:
for tier, row in game.table_by_tier.items():
for idx, _ in enumerate(row):
actions.append(
ReserveCard(tier=tier, index=idx, from_deck=False),
)
for tier, deck in game.decks_by_tier.items():
if deck:
actions.append(
ReserveCard(tier=tier, index=None, from_deck=True),
)
return actions
def create_random_cards_tier(
tier: int,
card_count: int,
cost_choices: list[int],
point_choices: list[int],
) -> list[Card]:
"""Create a random set of cards for a given tier."""
cards: list[Card] = []
for color in BASE_COLORS:
for _ in range(card_count):
cost = dict.fromkeys(GEM_COLORS, 0)
for c in BASE_COLORS:
if c == color:
continue
cost[c] = random.choice(cost_choices)
points = random.choice(point_choices)
cards.append(Card(tier=tier, points=points, color=color, cost=cost))
return cards
def create_random_cards() -> list[Card]:
"""Generate a generic but Splendor-ish set of cards.
This is not the official deck, but structured similarly enough for play.
"""
cards: list[Card] = []
cards.extend(
create_random_cards_tier(
tier=1,
card_count=5,
cost_choices=[0, 1, 1, 2],
point_choices=[0, 0, 1],
)
)
cards.extend(
create_random_cards_tier(
tier=2,
card_count=4,
cost_choices=[2, 3, 4],
point_choices=[1, 2, 2, 3],
)
)
cards.extend(
create_random_cards_tier(
tier=3,
card_count=3,
cost_choices=[4, 5, 6],
point_choices=[3, 4, 5],
)
)
random.shuffle(cards)
return cards
def create_random_nobles() -> list[Noble]:
"""A small set of noble tiles, roughly Splendor-ish."""
nobles: list[Noble] = []
base_requirements: list[dict[GemColor, int]] = [
{"white": 3, "blue": 3, "green": 3},
{"blue": 3, "green": 3, "red": 3},
{"green": 3, "red": 3, "black": 3},
{"red": 3, "black": 3, "white": 3},
{"black": 3, "white": 3, "blue": 3},
{"white": 4, "blue": 4},
{"green": 4, "red": 4},
{"blue": 4, "black": 4},
]
for idx, req in enumerate(base_requirements, start=1):
nobles.append(
Noble(
name=f"Noble {idx}",
points=3,
requirements=dict(req.items()),
),
)
return nobles
def load_nobles(file: Path) -> list[Noble]:
"""Load nobles from a file."""
nobles = json.loads(file.read_text())
return [Noble(**noble) for noble in nobles]
def load_cards(file: Path) -> list[Card]:
"""Load cards from a file."""
cards = json.loads(file.read_text())
return [Card(**card) for card in cards]
def new_game(
strategies: Sequence[Strategy],
config: GameConfig,
) -> GameState:
"""Create a new game state from a config + list of players."""
num_players = len(strategies)
bank = get_default_starting_tokens(num_players)
decks_by_tier: dict[int, list[Card]] = {1: [], 2: [], 3: []}
for card in config.cards:
decks_by_tier.setdefault(card.tier, []).append(card)
for deck in decks_by_tier.values():
random.shuffle(deck)
table_by_tier: dict[int, list[Card]] = {1: [], 2: [], 3: []}
players = [PlayerState(strategy=strategy) for strategy in strategies]
nobles = list(config.nobles)
random.shuffle(nobles)
nobles = nobles[: num_players + 1]
game = GameState(
config=config,
players=players,
bank=bank,
decks_by_tier=decks_by_tier,
table_by_tier=table_by_tier,
available_nobles=nobles,
)
game.refill_table()
return game
def run_game(game: GameState) -> tuple[PlayerState, int]:
"""Run a full game loop until someone wins or a player returns None."""
turn_count = 0
while not game.finished:
turn_count += 1
player = game.current_player
strategy = player.strategy
action = strategy.choose_action(game, player)
if action is None:
game.finished = True
break
apply_action(game, strategy, action)
check_nobles_for_player(game, strategy, player)
winner = game.check_winner_simple()
if winner is not None:
return winner, turn_count
game.next_player()
if turn_count >= game.config.turn_limit:
break
fallback = max(game.players, key=lambda player: player.score)
return fallback, turn_count
-288
View File
@@ -1,288 +0,0 @@
"""Bot for Splendor game."""
from __future__ import annotations
import random
from .base import (
BASE_COLORS,
Action,
BuyCard,
BuyCardReserved,
Card,
GameState,
GemColor,
PlayerState,
ReserveCard,
Strategy,
TakeDifferent,
TakeDouble,
auto_discard_tokens,
get_legal_actions,
)
def can_bot_afford(player: PlayerState, card: Card) -> bool:
"""Check if player can afford card, using discounts + gold."""
missing = 0
gold = player.tokens["gold"]
for color, cost in card.cost.items():
missing += max(0, cost - player.discounts.get(color, 0) - player.tokens.get(color, 0))
if missing > gold:
return False
return True
class RandomBot(Strategy):
"""Dumb bot that follows rules but doesn't think."""
def __init__(self, name: str) -> None:
"""Initialize the bot."""
super().__init__(name=name)
def choose_action(self, game: GameState, player: PlayerState) -> Action | None:
"""Choose an action for the current player."""
affordable: list[tuple[int, int]] = []
for tier, row in game.table_by_tier.items():
for idx, card in enumerate(row):
if can_bot_afford(player, card):
affordable.append((tier, idx))
if affordable and random.random() < 0.5:
tier, idx = random.choice(affordable)
return BuyCard(tier=tier, index=idx)
if random.random() < 0.2:
tier = random.choice([1, 2, 3])
row = game.table_by_tier.get(tier, [])
if row:
idx = random.randrange(len(row))
return ReserveCard(tier=tier, index=idx, from_deck=False)
if random.random() < 0.5:
colors_for_double = [c for c in BASE_COLORS if game.bank[c] >= 4]
if colors_for_double:
return TakeDouble(color=random.choice(colors_for_double))
colors_for_diff = [c for c in BASE_COLORS if game.bank[c] > 0]
random.shuffle(colors_for_diff)
return TakeDifferent(colors=colors_for_diff[:3])
def choose_discard(
self,
game: GameState, # noqa: ARG002
player: PlayerState,
excess: int,
) -> dict[GemColor, int]:
"""Choose how many tokens to discard."""
return auto_discard_tokens(player, excess)
def check_cards_in_tier(row: list[Card], player: PlayerState) -> list[int]:
"""Check if player can afford card, using discounts + gold."""
return [index for index, card in enumerate(row) if can_bot_afford(player, card)]
class PersonalizedBot(Strategy):
"""PersonalizedBot."""
"""Dumb bot that follows rules but doesn't think."""
def __init__(self, name: str) -> None:
"""Initialize the bot."""
super().__init__(name=name)
def choose_action(self, game: GameState, player: PlayerState) -> Action | None:
"""Choose an action for the current player."""
for tier in (1, 2, 3):
row = game.table_by_tier[tier]
if affordable := check_cards_in_tier(row, player):
index = random.choice(affordable)
return BuyCard(tier=tier, index=index)
colors_for_diff = [c for c in BASE_COLORS if game.bank[c] > 0]
random.shuffle(colors_for_diff)
return TakeDifferent(colors=colors_for_diff[:3])
def choose_discard(
self,
game: GameState, # noqa: ARG002
player: PlayerState,
excess: int,
) -> dict[GemColor, int]:
"""Choose how many tokens to discard."""
return auto_discard_tokens(player, excess)
class PersonalizedBot2(Strategy):
"""PersonalizedBot2."""
"""Dumb bot that follows rules but doesn't think."""
def __init__(self, name: str) -> None:
"""Initialize the bot."""
super().__init__(name=name)
def choose_action(self, game: GameState, player: PlayerState) -> Action | None:
"""Choose an action for the current player."""
tiers = (1, 2, 3)
for tier in tiers:
row = game.table_by_tier[tier]
if affordable := check_cards_in_tier(row, player):
index = random.choice(affordable)
return BuyCard(tier=tier, index=index)
if affordable := check_cards_in_tier(player.reserved, player):
index = random.choice(affordable)
return BuyCardReserved(index=index)
colors_for_diff = [c for c in BASE_COLORS if game.bank[c] > 0]
if len(colors_for_diff) >= 3:
random.shuffle(colors_for_diff)
return TakeDifferent(colors=colors_for_diff[:3])
for tier in tiers:
len_deck = len(game.decks_by_tier[tier])
if len_deck:
return ReserveCard(tier=tier, index=None, from_deck=True)
return TakeDifferent(colors=colors_for_diff[:3])
def choose_discard(
self,
game: GameState, # noqa: ARG002
player: PlayerState,
excess: int,
) -> dict[GemColor, int]:
"""Choose how many tokens to discard."""
return auto_discard_tokens(player, excess)
def buy_card_reserved(player: PlayerState) -> Action | None:
"""Buy a card reserved."""
if affordable := check_cards_in_tier(player.reserved, player):
index = random.choice(affordable)
return BuyCardReserved(index=index)
return None
def buy_card(game: GameState, player: PlayerState) -> Action | None:
"""Buy a card."""
for tier in (1, 2, 3):
row = game.table_by_tier[tier]
if affordable := check_cards_in_tier(row, player):
index = random.choice(affordable)
return BuyCard(tier=tier, index=index)
return None
def take_tokens(game: GameState) -> Action | None:
"""Take tokens."""
colors_for_diff = [color for color in BASE_COLORS if game.bank[color] > 0]
if len(colors_for_diff) >= 3:
random.shuffle(colors_for_diff)
return TakeDifferent(colors=colors_for_diff[: game.config.max_token_take])
return None
class PersonalizedBot3(Strategy):
"""PersonalizedBot3."""
"""Dumb bot that follows rules but doesn't think."""
def __init__(self, name: str) -> None:
"""Initialize the bot."""
super().__init__(name=name)
def choose_action(self, game: GameState, player: PlayerState) -> Action | None:
"""Choose an action for the current player."""
print(len(get_legal_actions(game, player)))
print(get_legal_actions(game, player))
if action := buy_card_reserved(player):
return action
if action := buy_card(game, player):
return action
colors_for_diff = [color for color in BASE_COLORS if game.bank[color] > 0]
if len(colors_for_diff) >= 3:
random.shuffle(colors_for_diff)
return TakeDifferent(colors=colors_for_diff[:3])
for tier in (1, 2, 3):
len_deck = len(game.decks_by_tier[tier])
if len_deck:
return ReserveCard(tier=tier, index=None, from_deck=True)
return TakeDifferent(colors=colors_for_diff[:3])
def choose_discard(
self,
game: GameState, # noqa: ARG002
player: PlayerState,
excess: int,
) -> dict[GemColor, int]:
"""Choose how many tokens to discard."""
return auto_discard_tokens(player, excess)
def estimate_value_of_card(game: GameState, player: PlayerState, color: GemColor) -> int:
"""Estimate value of a color in the player's bank."""
return game.bank[color] - player.discounts.get(color, 0)
def estimate_value_of_token(game: GameState, player: PlayerState, color: GemColor) -> int:
"""Estimate value of a color in the player's bank."""
return game.bank[color] - player.discounts.get(color, 0)
class PersonalizedBot4(Strategy):
"""PersonalizedBot4."""
def __init__(self, name: str) -> None:
"""Initialize the bot."""
super().__init__(name=name)
def filter_actions(self, actions: list[Action]) -> list[Action]:
"""Filter actions to only take different."""
return [
action
for action in actions
if (isinstance(action, TakeDifferent) and len(action.colors) == 3) or not isinstance(action, TakeDifferent)
]
def choose_action(self, game: GameState, player: PlayerState) -> Action | None:
"""Choose an action for the current player."""
legal_actions = get_legal_actions(game, player)
print(len(legal_actions))
good_actions = self.filter_actions(legal_actions)
print(len(good_actions))
print(good_actions)
print(len(get_legal_actions(game, player)))
if action := buy_card_reserved(player):
return action
if action := buy_card(game, player):
return action
colors_for_diff = [color for color in BASE_COLORS if game.bank[color] > 0]
if len(colors_for_diff) >= 3:
random.shuffle(colors_for_diff)
return TakeDifferent(colors=colors_for_diff[:3])
for tier in (1, 2, 3):
len_deck = len(game.decks_by_tier[tier])
if len_deck:
return ReserveCard(tier=tier, index=None, from_deck=True)
return TakeDifferent(colors=colors_for_diff[:3])
def choose_discard(
self,
game: GameState, # noqa: ARG002
player: PlayerState,
excess: int,
) -> dict[GemColor, int]:
"""Choose how many tokens to discard."""
return auto_discard_tokens(player, excess)
-724
View File
@@ -1,724 +0,0 @@
"""Splendor game."""
from __future__ import annotations
import sys
from typing import TYPE_CHECKING, Any
from textual.app import App, ComposeResult
from textual.containers import Horizontal, Vertical
from textual.widget import Widget
from textual.widgets import Footer, Header, Input, Static
from .base import (
BASE_COLORS,
GEM_COLORS,
Action,
BuyCard,
BuyCardReserved,
Card,
GameState,
GemColor,
Noble,
PlayerState,
ReserveCard,
Strategy,
TakeDifferent,
TakeDouble,
)
if TYPE_CHECKING:
from collections.abc import Mapping
# Abbreviations used when rendering costs
COST_ABBR: dict[GemColor, str] = {
"white": "W",
"blue": "B",
"green": "G",
"red": "R",
"black": "K",
"gold": "O",
}
# Abbreviations players can type on the command line
COLOR_ABBR_TO_FULL: dict[str, GemColor] = {
"w": "white",
"b": "blue",
"g": "green",
"r": "red",
"k": "black",
"o": "gold",
}
def parse_color_token(raw: str) -> GemColor:
"""Convert user input into a GemColor.
Supports:
- full names: white, blue, green, red, black, gold
- abbreviations: w, b, g, r, k, o
"""
key = raw.lower()
# full color names first
if key in BASE_COLORS:
return key # type: ignore[return-value]
# abbreviations
if key in COLOR_ABBR_TO_FULL:
return COLOR_ABBR_TO_FULL[key]
error = f"Unknown color: {raw}"
raise ValueError(error)
def format_cost(cost: Mapping[GemColor, int]) -> str:
"""Format a cost/requirements dict as colored tokens like 'B:2, R:1'.
Uses `color_token` internally so colors are guaranteed to match your bank.
"""
parts: list[str] = []
for color in GEM_COLORS:
n = cost.get(color, 0)
if not n:
continue
# color_token gives us e.g. "[blue]blue: 3[/]"
token = color_token(color, n)
# Turn the leading color name into the abbreviation (blue: 3 → B:3)
# We only replace the first occurrence.
full = f"{color}:"
abbr = f"{COST_ABBR[color]}:"
token = token.replace(full, abbr, 1)
parts.append(token)
return ", ".join(parts) if parts else "-"
def format_card(card: Card) -> str:
"""Readable card line using dataclass fields instead of __str__."""
color_abbr = COST_ABBR[card.color]
header = f"T{card.tier} {color_abbr} P{card.points}"
cost_str = format_cost(card.cost)
return f"{header} ({cost_str})"
def format_noble(noble: Noble) -> str:
"""Readable noble line using dataclass fields instead of __str__."""
cost_str = format_cost(noble.requirements)
return f"{noble.name} +{noble.points} ({cost_str})"
def format_tokens(tokens: Mapping[GemColor, int]) -> str:
"""Colored 'color: n' list for a token dict."""
return " ".join(color_token(c, tokens.get(c, 0)) for c in GEM_COLORS)
def format_discounts(discounts: Mapping[GemColor, int]) -> str:
"""Colored discounts, skipping zeros."""
parts: list[str] = []
for c in GEM_COLORS:
n = discounts.get(c, 0)
if not n:
continue
abbr = COST_ABBR[c]
fg, bg = COLOR_STYLE[c]
parts.append(f"[{fg} on {bg}]{abbr}:{n}[/{fg} on {bg}]")
return ", ".join(parts) if parts else "-"
COLOR_STYLE: dict[GemColor, tuple[str, str]] = {
"white": ("black", "white"), # fg, bg
"blue": ("bright_white", "blue"),
"green": ("bright_white", "sea_green4"),
"red": ("white", "red3"),
"black": ("white", "grey0"),
"gold": ("black", "yellow3"),
}
def fmt_gem(color: GemColor) -> str:
"""Render gem name with fg/bg matching real token color."""
fg, bg = COLOR_STYLE[color]
return f"[{fg} on {bg}] {color} [/{fg} on {bg}]"
def fmt_number(value: int) -> str:
"""Return a Rich-markup colored 'value' string."""
return f"[bold cyan]{value}[/]"
def color_token(name: GemColor, amount: int) -> str:
"""Return a Rich-markup colored 'name: n' string."""
# Map Splendor colors -> terminal colors
color_map: Mapping[GemColor, str] = {
"white": "white",
"blue": "blue",
"green": "green",
"red": "red",
"black": "grey70", # 'black' is unreadable on dark backgrounds
"gold": "yellow",
}
style = color_map.get(name, "white")
return f"[{style}]{name}: {amount}[/]"
class Board(Widget):
"""Big board widget with the layout you sketched."""
def __init__(self, game: GameState, me: PlayerState, **kwargs: Any) -> None: # noqa: ANN401
"""Initialize the board widget."""
super().__init__(**kwargs)
self.game = game
self.me = me
def compose(self) -> ComposeResult:
"""Compose the board widget."""
# Structure:
# ┌ bank row
# ├ middle row (tiers | nobles)
# └ players row
with Vertical(id="board_root"):
yield Static(id="bank_box")
with Horizontal(id="middle_row"):
with Vertical(id="tiers_box"):
yield Static(id="tier1_box")
yield Static(id="tier2_box")
yield Static(id="tier3_box")
yield Static(id="nobles_box")
yield Static(id="players_box")
def on_mount(self) -> None:
"""Refresh the board content."""
self.refresh_content()
def refresh_content(self) -> None:
"""Refresh the board content."""
self._render_bank()
self._render_tiers()
self._render_nobles()
self._render_players()
# --- sections ----------------------------------------------------
def _render_bank(self) -> None:
bank = self.game.bank
parts: list[str] = ["[b]Bank:[/b]"]
# One line, all tokens colored
parts.append(format_tokens(bank))
self.query_one("#bank_box", Static).update("\n".join(parts))
def _render_tiers(self) -> None:
for tier in (1, 2, 3):
box = self.query_one(f"#tier{tier}_box", Static)
cards: list[Card] = self.game.table_by_tier.get(tier, [])
lines: list[str] = [f"[b]Tier {tier} cards:[/b]"]
if not cards:
lines.append(" (none)")
else:
for idx, card in enumerate(cards):
lines.append(f" [{idx}] {format_card(card)}")
box.update("\n".join(lines))
def _render_nobles(self) -> None:
nobles_box = self.query_one("#nobles_box", Static)
lines: list[str] = ["[b]Nobles[/b]"]
if not self.game.available_nobles:
lines.append(" (none)")
else:
lines.extend(" - " + format_noble(noble) for noble in self.game.available_nobles)
nobles_box.update("\n".join(lines))
def _render_players(self) -> None:
players_box = self.query_one("#players_box", Static)
lines: list[str] = ["[b]Players:[/b]", ""]
for player in self.game.players:
mark = "*" if player is self.me else " "
token_str = format_tokens(player.tokens)
discount_str = format_discounts(player.discounts)
lines.append(
f"{mark} {player.name:10} Score={player.score:2d} Discounts={discount_str}",
)
lines.append(f" Tokens: {token_str}")
if player.nobles:
noble_names = ", ".join(n.name for n in player.nobles)
lines.append(f" Nobles: {noble_names}")
# Optional: show counts of cards / reserved
if player.cards:
lines.append(f" Cards: {len(player.cards)}")
if player.reserved:
lines.append(f" Reserved: {len(player.reserved)}")
lines.append("")
players_box.update("\n".join(lines))
class ActionApp(App[None]):
"""Textual app that asks for a single action command and returns an Action."""
CSS = """
Screen {
/* 3 rows: command zone, board, footer */
layout: grid;
grid-size: 1 3;
grid-rows: auto 1fr auto;
}
/* Top area with input + instructions */
#command_zone {
grid-columns: 1;
grid-rows: 1;
padding: 1 1;
}
/* Board sits in the middle row and can grow */
#board {
grid-columns: 1;
grid-rows: 2;
padding: 0 1 1 1;
}
Footer {
grid-columns: 1;
grid-rows: 3;
}
Input {
border: round $accent;
}
/* === Board layout === */
#board_root {
/* outer frame around the whole board area */
border: heavy white;
padding: 0 1;
}
/* Bank row: full width */
#bank_box {
border: heavy white;
padding: 0 1;
}
/* Middle row: tiers (left) + nobles (right) */
#middle_row {
layout: horizontal;
}
#tiers_box {
border: heavy white;
padding: 0 1;
width: 70%;
}
#tier1_box,
#tier2_box,
#tier3_box {
border-bottom: heavy white;
padding: 0 0 1 0;
margin-bottom: 1;
}
#nobles_box {
border: heavy white;
padding: 0 1;
width: 30%;
}
/* Players row: full width at bottom */
#players_box {
border: heavy white;
padding: 0 1;
}
"""
def __init__(self, game: GameState, player: PlayerState) -> None:
"""Initialize the action app."""
super().__init__()
self.game = game
self.player = player
self.result: Action | None = None
self.message: str = ""
def compose(self) -> ComposeResult:
"""Compose the action app."""
# Row 1: input + Actions text
with Vertical(id="command_zone"):
yield Input(
placeholder="Enter command, e.g. '1 white blue red' or '1 w b r' or 'q'",
id="input_line",
)
yield Static("", id="prompt")
# Row 2: board
yield Board(self.game, self.player, id="board")
# Row 3: footer
yield Footer()
def on_mount(self) -> None:
"""Mount the action app."""
self._update_prompt()
self.query_one(Input).focus()
def _update_prompt(self) -> None:
lines: list[str] = []
lines.append("[bold underline]Actions:[/]")
lines.append(
" [bold green]1[/] <colors...> - Take up to 3 different gem colors "
"(e.g. [cyan]1 white blue red[/] or [cyan]1 w b r[/])",
)
lines.append(
f" [bold green]2[/] <color> - Take 2 of the same color (needs {fmt_number(4)} in bank, "
"e.g. [cyan]2 blue[/] or [cyan]2 b[/])",
)
lines.append(
" [bold green]3[/] <tier> <idx> - Buy a face-up card (e.g. [cyan]3 1 0[/] for tier 1, index 0)",
)
lines.append(" [bold green]4[/] <idx> - Buy a reserved card")
lines.append(" [bold green]5[/] <tier> <idx> - Reserve a face-up card")
lines.append(" [bold green]6[/] <tier> - Reserve top card of a deck")
lines.append(" [bold red]q[/] - Quit game")
if self.message:
lines.append("")
lines.append(f"[bold red]Message:[/] {self.message}")
self.query_one("#prompt", Static).update("\n".join(lines))
def _cmd_1(self, parts: list[str]) -> str | None:
"""Take up to 3 different gem colors: 1 white blue red OR 1 w b r."""
color_names = parts[1:]
if not color_names:
return "Need at least one color (full name or abbreviation)."
colors: list[GemColor] = []
for name in color_names:
color = parse_color_token(name)
if self.game.bank[color] <= 0:
return f"No tokens left for color: {color}"
colors.append(color)
self.result = TakeDifferent(colors=colors[:3])
self.exit()
return None
def _cmd_2(self, parts: list[str]) -> str | None:
"""Take two of the same color."""
if len(parts) < 2:
return "Usage: 2 <color>"
color = parse_color_token(parts[1])
if self.game.bank[color] < self.game.config.minimum_tokens_to_buy_2:
return "Bank must have at least 4 of that color."
self.result = TakeDouble(color=color)
self.exit()
return None
def _cmd_3(self, parts: list[str]) -> str | None:
"""Buy face-up card."""
if len(parts) < 3:
return "Usage: 3 <tier> <index>"
tier = int(parts[1])
idx = int(parts[2])
self.result = BuyCard(tier=tier, index=idx)
self.exit()
return None
def _cmd_4(self, parts: list[str]) -> str | None:
"""Buy reserved card."""
if len(parts) < 2:
return "Usage: 4 <reserved_index>"
idx = int(parts[1])
if not (0 <= idx < len(self.player.reserved)):
return "Reserved index out of range."
self.result = BuyCardReserved(tier=0, index=idx)
self.exit()
return None
def _cmd_5(self, parts: list[str]) -> str | None:
"""Reserve face-up card."""
if len(parts) < 3:
return "Usage: 5 <tier> <index>"
tier = int(parts[1])
idx = int(parts[2])
self.result = ReserveCard(tier=tier, index=idx, from_deck=False)
self.exit()
return None
def _cmd_6(self, parts: list[str]) -> str | None:
"""Reserve top of deck."""
if len(parts) < 2:
return "Usage: 6 <tier>"
tier = int(parts[1])
self.result = ReserveCard(tier=tier, index=None, from_deck=True)
self.exit()
return None
def _unknown_cmd(self, _parts: list[str]) -> str:
return "Unknown command."
def on_input_submitted(self, event: Input.Submitted) -> None:
"""Handle user input."""
text = (event.value or "").strip()
event.input.value = ""
if not text:
return
if text.lower() in {"q", "quit", "0"}:
self.result = None
self.exit()
return
parts = text.split()
cmds = {
"1": self._cmd_1,
"2": self._cmd_2,
"3": self._cmd_3,
"4": self._cmd_4,
"5": self._cmd_5,
"6": self._cmd_6,
}
cmd = parts[0]
error = cmds.get(cmd, self._unknown_cmd)(parts)
if error:
self.message = error
self._update_prompt()
return
class DiscardApp(App[None]):
"""Textual app to choose discards when over token limit."""
CSS = """
Screen {
layout: vertical;
}
#command_zone {
padding: 1 1;
}
#board {
padding: 0 1 1 1;
}
Input {
border: round $accent;
}
"""
def __init__(self, game: GameState, player: PlayerState) -> None:
"""Initialize the discard app."""
super().__init__()
self.game = game
self.player = player
self.discards: dict[GemColor, int] = dict.fromkeys(GEM_COLORS, 0)
self.message: str = ""
def compose(self) -> ComposeResult: # type: ignore[override]
"""Compose the discard app."""
yield Header(show_clock=False)
with Vertical(id="command_zone"):
yield Input(
placeholder="Enter color to discard, e.g. 'blue' or 'b'",
id="input_line",
)
yield Static("", id="prompt")
# Board directly under the command zone
yield Board(self.game, self.player, id="board")
yield Footer()
def on_mount(self) -> None: # type: ignore[override]
"""Mount the discard app."""
self._update_prompt()
self.query_one(Input).focus()
def _remaining_to_discard(self) -> int:
return self.player.total_tokens() - sum(self.discards.values()) - self.game.config.token_limit
def _update_prompt(self) -> None:
remaining = max(self._remaining_to_discard(), 0)
lines: list[str] = []
lines.append(
"You must discard "
f"{fmt_number(remaining)} token(s) "
f"to get down to {fmt_number(self.game.config.token_limit)}.",
)
disc_str = ", ".join(f"{fmt_gem(c)}={fmt_number(self.discards[c])}" for c in GEM_COLORS)
lines.append(f"Current planned discards: {{ {disc_str} }}")
lines.append(
"Type a color name or abbreviation (e.g. 'blue' or 'b') to discard one token.",
)
if self.message:
lines.append("")
lines.append(f"[bold red]Message:[/] {self.message}")
self.query_one("#prompt", Static).update("\n".join(lines))
def on_input_submitted(self, event: Input.Submitted) -> None: # type: ignore[override]
"""Handle user input."""
raw = (event.value or "").strip()
event.input.value = ""
if not raw:
return
try:
color = parse_color_token(raw)
except ValueError:
self.message = f"Unknown color: {raw}"
self._update_prompt()
return
available = self.player.tokens[color] - self.discards[color]
if available <= 0:
self.message = f"No more {color} tokens available to discard."
self._update_prompt()
return
self.discards[color] += 1
if self._remaining_to_discard() <= 0:
self.exit()
return
self.message = ""
self._update_prompt()
# ---------------------------------------------------------------------------
# Noble choice app
# ---------------------------------------------------------------------------
class NobleChoiceApp(App[None]):
"""Textual app to choose one noble."""
CSS = """
Screen {
layout: vertical;
}
#command_zone {
padding: 1 1;
}
#board {
padding: 0 1 1 1;
}
Input {
border: round $accent;
}
"""
def __init__(
self,
game: GameState,
player: PlayerState,
nobles: list[Noble],
) -> None:
"""Initialize the noble choice app."""
super().__init__()
self.game = game
self.player = player
self.nobles = nobles
self.result: Noble | None = None
self.message: str = ""
def compose(self) -> ComposeResult: # type: ignore[override]
"""Compose the noble choice app."""
yield Header(show_clock=False)
with Vertical(id="command_zone"):
yield Input(
placeholder="Enter noble index, e.g. '0'",
id="input_line",
)
yield Static("", id="prompt")
# Board directly under the command zone
yield Board(self.game, self.player, id="board")
yield Footer()
def on_mount(self) -> None: # type: ignore[override]
"""Mount the noble choice app."""
self._update_prompt()
self.query_one(Input).focus()
def _update_prompt(self) -> None:
lines: list[str] = []
lines.append("[bold underline]You qualify for nobles:[/]")
for i, noble in enumerate(self.nobles):
lines.append(f" [bright_cyan]{i})[/] {format_noble(noble)}")
lines.append("Enter the index of the noble you want.")
if self.message:
lines.append("")
lines.append(f"[bold red]Message:[/] {self.message}")
self.query_one("#prompt", Static).update("\n".join(lines))
def on_input_submitted(self, event: Input.Submitted) -> None: # type: ignore[override]
"""Handle user input."""
raw = (event.value or "").strip()
event.input.value = ""
if not raw:
return
try:
idx = int(raw)
except ValueError:
self.message = "Please enter a valid integer index."
self._update_prompt()
return
if not (0 <= idx < len(self.nobles)):
self.message = "Index out of range."
self._update_prompt()
return
self.result = self.nobles[idx]
self.exit()
class TuiHuman(Strategy):
"""Textual-based human player Strategy with colorful board."""
def choose_action(
self,
game: GameState,
player: PlayerState,
) -> Action | None:
"""Choose an action for the player."""
if not sys.stdout.isatty():
return None
app = ActionApp(game, player)
app.run()
return app.result
def choose_discard(
self,
game: GameState,
player: PlayerState,
excess: int, # noqa: ARG002
) -> dict[GemColor, int]:
"""Choose tokens to discard."""
if not sys.stdout.isatty():
return dict.fromkeys(GEM_COLORS, 0)
app = DiscardApp(game, player)
app.run()
return app.discards
def choose_noble(
self,
game: GameState,
player: PlayerState,
nobles: list[Noble],
) -> Noble:
"""Choose a noble for the player."""
if not sys.stdout.isatty():
return nobles[0]
app = NobleChoiceApp(game, player, nobles)
app.run()
return app.result
-19
View File
@@ -1,19 +0,0 @@
"""Main entry point for Splendor game."""
from __future__ import annotations
from .base import new_game, run_game
from .bot import RandomBot
from .human import TuiHuman
def main() -> None:
"""Main entry point."""
human = TuiHuman()
bot = RandomBot()
game_state = new_game(["You", "Bot A"])
run_game(game_state, [human, bot])
if __name__ == "__main__":
main()
-111
View File
@@ -1,111 +0,0 @@
"""Public state for RL/search."""
from __future__ import annotations
from dataclasses import dataclass
from .base import (
BASE_COLORS,
BASE_INDEX,
GEM_ORDER,
Card,
GameState,
Noble,
PlayerState,
)
@dataclass(frozen=True)
class ObsCard:
"""Numeric-ish card view for RL/search."""
tier: int
points: int
color_index: int
cost: list[int]
@dataclass(frozen=True)
class ObsNoble:
"""Numeric-ish noble view for RL/search."""
points: int
requirements: list[int]
@dataclass(frozen=True)
class ObsPlayer:
"""Numeric-ish player view for RL/search."""
tokens: list[int]
discounts: list[int]
score: int
cards: list[ObsCard]
reserved: list[ObsCard]
nobles: list[ObsNoble]
@dataclass(frozen=True)
class Observation:
"""Full public state for RL/search."""
current_player: int
bank: list[int]
players: list[ObsPlayer]
table_by_tier: dict[int, list[ObsCard]]
decks_remaining: dict[int, int]
available_nobles: list[ObsNoble]
def _encode_card(card: Card) -> ObsCard:
color_index = BASE_INDEX.get(card.color, -1)
cost_vec = [card.cost.get(c, 0) for c in BASE_COLORS]
return ObsCard(
tier=card.tier,
points=card.points,
color_index=color_index,
cost=cost_vec,
)
def _encode_noble(noble: Noble) -> ObsNoble:
req_vec = [noble.requirements.get(c, 0) for c in BASE_COLORS]
return ObsNoble(
points=noble.points,
requirements=req_vec,
)
def _encode_player(player: PlayerState) -> ObsPlayer:
tokens_vec = [player.tokens[c] for c in GEM_ORDER]
discounts_vec = [player.discounts[c] for c in GEM_ORDER]
cards_enc = [_encode_card(c) for c in player.cards]
reserved_enc = [_encode_card(c) for c in player.reserved]
nobles_enc = [_encode_noble(n) for n in player.nobles]
return ObsPlayer(
tokens=tokens_vec,
discounts=discounts_vec,
score=player.score,
cards=cards_enc,
reserved=reserved_enc,
nobles=nobles_enc,
)
def to_observation(game: GameState) -> Observation:
"""Create a structured observation of the full public state."""
bank_vec = [game.bank[c] for c in GEM_ORDER]
players_enc = [_encode_player(p) for p in game.players]
table_enc: dict[int, list[ObsCard]] = {
tier: [_encode_card(c) for c in row] for tier, row in game.table_by_tier.items()
}
decks_remaining = {tier: len(deck) for tier, deck in game.decks_by_tier.items()}
nobles_enc = [_encode_noble(n) for n in game.available_nobles]
return Observation(
current_player=game.current_player_index,
bank=bank_vec,
players=players_enc,
table_by_tier=table_enc,
decks_remaining=decks_remaining,
available_nobles=nobles_enc,
)
-36
View File
@@ -1,36 +0,0 @@
"""Simulate a step in the game."""
from __future__ import annotations
import copy
from .base import Action, GameState, PlayerState, apply_action, check_nobles_for_player
from .bot import RandomBot
class SimStrategy(RandomBot):
"""Strategy used in simulate_step.
We never call choose_action here (caller chooses actions),
but we reuse discard/noble-selection logic.
"""
def choose_action(self, game: GameState, player: PlayerState) -> Action | None: # noqa: ARG002
"""Choose an action for the current player."""
msg = "SimStrategy.choose_action should not be used in simulate_step"
raise RuntimeError(msg)
def simulate_step(game: GameState, action: Action) -> GameState:
"""Return a deep-copied next state after applying action for the current player.
Useful for tree search / MCTS:
next_state = simulate_step(state, action)
"""
next_state = copy.deepcopy(game)
sim_strategy = SimStrategy()
apply_action(next_state, sim_strategy, action)
check_nobles_for_player(next_state, sim_strategy, next_state.current_player)
next_state.next_player()
return next_state
-50
View File
@@ -1,50 +0,0 @@
"""Simulator for Splendor game."""
from __future__ import annotations
from collections import defaultdict
from pathlib import Path
from statistics import mean
from .base import GameConfig, load_cards, load_nobles, new_game, run_game
from .bot import PersonalizedBot4, RandomBot
def main() -> None:
"""Main entry point."""
turn_limit = 1000
good_games = 0
games = 1
winners: dict[str, list] = defaultdict(list)
game_data = Path(__file__).parent / "game_data"
cards = load_cards(game_data / "cards/default.json")
nobles = load_nobles(game_data / "nobles/default.json")
for _ in range(games):
bot_a = RandomBot("bot_a")
bot_b = RandomBot("bot_b")
bot_c = RandomBot("bot_c")
bot_d = PersonalizedBot4("my_bot")
config = GameConfig(
cards=cards,
nobles=nobles,
turn_limit=turn_limit,
)
players = (bot_a, bot_b, bot_c, bot_d)
game_state = new_game(players, config)
winner, turns = run_game(game_state)
if turns < turn_limit:
good_games += 1
winners[winner.strategy.name].append(turns)
print(
f"out of {games} {turn_limit} turn games with {len(players)}"
f"random bots there where {good_games} games where a bot won"
)
for name, turns in winners.items():
print(f"{name} won {len(turns)} games in {mean(turns):.2f} turns")
if __name__ == "__main__":
main()
+1 -1
View File
@@ -4,7 +4,7 @@ import logging
import sys import sys
import tomllib import tomllib
from os import environ from os import environ
from pathlib import Path from pathlib import Path # noqa: TC003 This is required for the typer CLI
from socket import gethostname from socket import gethostname
import typer import typer
+1 -1
View File
@@ -451,7 +451,7 @@ def convert_aax_file_with_agent(aax_file: Path, config: ConversionConfig) -> Non
destination.parent.mkdir(parents=True, exist_ok=True) destination.parent.mkdir(parents=True, exist_ok=True)
try: try:
temp_file.replace(destination) temp_file.replace(destination)
except Exception as error: # noqa: BLE001 except OSError as error:
write_review_file( write_review_file(
destination=destination, destination=destination,
ffprobe_metadata=ffprobe_metadata, ffprobe_metadata=ffprobe_metadata,
-1
View File
@@ -169,7 +169,6 @@ def csv_id(row: dict[str, str | None], csv_path: Path, row_number: int) -> int |
except ValueError as error: except ValueError as error:
msg = f"{csv_path}:{row_number}: id must be an integer: {value}" msg = f"{csv_path}:{row_number}: id must be an integer: {value}"
raise CatalogImportError(msg) from error raise CatalogImportError(msg) from error
return None
if __name__ == "__main__": if __name__ == "__main__":
-1
View File
@@ -1 +0,0 @@
"""Van inventory FastAPI application."""
-16
View File
@@ -1,16 +0,0 @@
"""FastAPI dependencies for van inventory."""
from collections.abc import Iterator
from typing import Annotated
from fastapi import Depends, Request
from sqlalchemy.orm import Session
def get_db(request: Request) -> Iterator[Session]:
"""Get database session from app state."""
with Session(request.app.state.engine) as session:
yield session
DbSession = Annotated[Session, Depends(get_db)]
-56
View File
@@ -1,56 +0,0 @@
"""FastAPI app for van inventory."""
from __future__ import annotations
import logging
from contextlib import asynccontextmanager
from pathlib import Path
from typing import TYPE_CHECKING, Annotated
import typer
import uvicorn
from fastapi import FastAPI
from fastapi.staticfiles import StaticFiles
from python.common import configure_logger
from python.orm.common import get_postgres_engine
from python.van_inventory.routers import api_router, frontend_router
STATIC_DIR = Path(__file__).resolve().parent / "static"
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]:
app.state.engine = get_postgres_engine(name="VAN_INVENTORY")
yield
app.state.engine.dispose()
app = FastAPI(title="Van Inventory", lifespan=lifespan)
app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")
app.include_router(api_router)
app.include_router(frontend_router)
return app
def serve(
# Intentionally binds all interfaces — this is a LAN-only van server
host: Annotated[str, typer.Option("--host", "-h", help="Host to bind to")] = "0.0.0.0", # noqa: S104
port: Annotated[int, typer.Option("--port", "-p", help="Port to bind to")] = 8001,
log_level: Annotated[str, typer.Option("--log-level", "-l", help="Log level")] = "INFO",
) -> None:
"""Start the Van Inventory server."""
configure_logger(log_level)
app = create_app()
uvicorn.run(app, host=host, port=port)
if __name__ == "__main__":
typer.run(serve)
-6
View File
@@ -1,6 +0,0 @@
"""Van inventory API routers."""
from python.van_inventory.routers.api import router as api_router
from python.van_inventory.routers.frontend import router as frontend_router
__all__ = ["api_router", "frontend_router"]
-314
View File
@@ -1,314 +0,0 @@
"""Van inventory API router."""
from __future__ import annotations
from typing import TYPE_CHECKING
from fastapi import APIRouter, HTTPException
from pydantic import BaseModel, Field
from sqlalchemy import select
from sqlalchemy.orm import selectinload
from python.orm.van_inventory.models import Item, Meal, MealIngredient
if TYPE_CHECKING:
from python.van_inventory.dependencies import DbSession
# --- Schemas ---
class ItemCreate(BaseModel):
"""Schema for creating an item."""
name: str
quantity: float = Field(default=0, ge=0)
unit: str
category: str | None = None
class ItemUpdate(BaseModel):
"""Schema for updating an item."""
name: str | None = None
quantity: float | None = Field(default=None, ge=0)
unit: str | None = None
category: str | None = None
class ItemResponse(BaseModel):
"""Schema for item response."""
id: int
name: str
quantity: float
unit: str
category: str | None
model_config = {"from_attributes": True}
class IngredientCreate(BaseModel):
"""Schema for adding an ingredient to a meal."""
item_id: int
quantity_needed: float = Field(gt=0)
class MealCreate(BaseModel):
"""Schema for creating a meal."""
name: str
instructions: str | None = None
ingredients: list[IngredientCreate] = []
class MealUpdate(BaseModel):
"""Schema for updating a meal."""
name: str | None = None
instructions: str | None = None
class IngredientResponse(BaseModel):
"""Schema for ingredient response."""
item_id: int
item_name: str
quantity_needed: float
unit: str
model_config = {"from_attributes": True}
class MealResponse(BaseModel):
"""Schema for meal response."""
id: int
name: str
instructions: str | None
ingredients: list[IngredientResponse] = []
model_config = {"from_attributes": True}
@classmethod
def from_meal(cls, meal: Meal) -> MealResponse:
"""Build a MealResponse from an ORM Meal with loaded ingredients."""
return cls(
id=meal.id,
name=meal.name,
instructions=meal.instructions,
ingredients=[
IngredientResponse(
item_id=mi.item_id,
item_name=mi.item.name,
quantity_needed=mi.quantity_needed,
unit=mi.item.unit,
)
for mi in meal.ingredients
],
)
class ShoppingItem(BaseModel):
"""An item needed for a meal that is short on stock."""
item_name: str
unit: str
needed: float
have: float
short: float
class MealAvailability(BaseModel):
"""Availability status for a meal."""
meal_id: int
meal_name: str
can_make: bool
missing: list[ShoppingItem] = []
# --- Routes ---
router = APIRouter(prefix="/api", tags=["van_inventory"])
# Items
@router.post("/items", response_model=ItemResponse)
def create_item(item: ItemCreate, db: DbSession) -> Item:
"""Create a new inventory item."""
db_item = Item(**item.model_dump())
db.add(db_item)
db.commit()
db.refresh(db_item)
return db_item
@router.get("/items", response_model=list[ItemResponse])
def list_items(db: DbSession) -> list[Item]:
"""List all inventory items."""
return list(db.scalars(select(Item).order_by(Item.name)).all())
@router.get("/items/{item_id}", response_model=ItemResponse)
def get_item(item_id: int, db: DbSession) -> Item:
"""Get an item by ID."""
item = db.get(Item, item_id)
if not item:
raise HTTPException(status_code=404, detail="Item not found")
return item
@router.patch("/items/{item_id}", response_model=ItemResponse)
def update_item(item_id: int, item: ItemUpdate, db: DbSession) -> Item:
"""Update an item by ID."""
db_item = db.get(Item, item_id)
if not db_item:
raise HTTPException(status_code=404, detail="Item not found")
for key, value in item.model_dump(exclude_unset=True).items():
setattr(db_item, key, value)
db.commit()
db.refresh(db_item)
return db_item
@router.delete("/items/{item_id}")
def delete_item(item_id: int, db: DbSession) -> dict[str, bool]:
"""Delete an item by ID."""
item = db.get(Item, item_id)
if not item:
raise HTTPException(status_code=404, detail="Item not found")
db.delete(item)
db.commit()
return {"deleted": True}
# Meals
@router.post("/meals", response_model=MealResponse)
def create_meal(meal: MealCreate, db: DbSession) -> MealResponse:
"""Create a new meal with optional ingredients."""
for ing in meal.ingredients:
if not db.get(Item, ing.item_id):
raise HTTPException(status_code=422, detail=f"Item {ing.item_id} not found")
db_meal = Meal(name=meal.name, instructions=meal.instructions)
db.add(db_meal)
db.flush()
for ing in meal.ingredients:
db.add(MealIngredient(meal_id=db_meal.id, item_id=ing.item_id, quantity_needed=ing.quantity_needed))
db.commit()
db_meal = db.scalar(
select(Meal)
.where(Meal.id == db_meal.id)
.options(selectinload(Meal.ingredients).selectinload(MealIngredient.item))
)
return MealResponse.from_meal(db_meal)
@router.get("/meals", response_model=list[MealResponse])
def list_meals(db: DbSession) -> list[MealResponse]:
"""List all meals with ingredients."""
meals = list(
db.scalars(
select(Meal).options(selectinload(Meal.ingredients).selectinload(MealIngredient.item)).order_by(Meal.name)
).all()
)
return [MealResponse.from_meal(m) for m in meals]
@router.get("/meals/availability", response_model=list[MealAvailability])
def check_all_meals(db: DbSession) -> list[MealAvailability]:
"""Check which meals can be made with current inventory."""
meals = list(
db.scalars(select(Meal).options(selectinload(Meal.ingredients).selectinload(MealIngredient.item))).all()
)
return [_check_meal(m) for m in meals]
@router.get("/meals/{meal_id}", response_model=MealResponse)
def get_meal(meal_id: int, db: DbSession) -> MealResponse:
"""Get a meal by ID with ingredients."""
meal = db.scalar(
select(Meal).where(Meal.id == meal_id).options(selectinload(Meal.ingredients).selectinload(MealIngredient.item))
)
if not meal:
raise HTTPException(status_code=404, detail="Meal not found")
return MealResponse.from_meal(meal)
@router.delete("/meals/{meal_id}")
def delete_meal(meal_id: int, db: DbSession) -> dict[str, bool]:
"""Delete a meal by ID."""
meal = db.get(Meal, meal_id)
if not meal:
raise HTTPException(status_code=404, detail="Meal not found")
db.delete(meal)
db.commit()
return {"deleted": True}
@router.post("/meals/{meal_id}/ingredients", response_model=MealResponse)
def add_ingredient(meal_id: int, ingredient: IngredientCreate, db: DbSession) -> MealResponse:
"""Add an ingredient to a meal."""
meal = db.get(Meal, meal_id)
if not meal:
raise HTTPException(status_code=404, detail="Meal not found")
if not db.get(Item, ingredient.item_id):
raise HTTPException(status_code=422, detail="Item not found")
existing = db.scalar(
select(MealIngredient).where(MealIngredient.meal_id == meal_id, MealIngredient.item_id == ingredient.item_id)
)
if existing:
raise HTTPException(status_code=409, detail="Ingredient already exists for this meal")
db.add(MealIngredient(meal_id=meal_id, item_id=ingredient.item_id, quantity_needed=ingredient.quantity_needed))
db.commit()
meal = db.scalar(
select(Meal).where(Meal.id == meal_id).options(selectinload(Meal.ingredients).selectinload(MealIngredient.item))
)
return MealResponse.from_meal(meal)
@router.delete("/meals/{meal_id}/ingredients/{item_id}")
def remove_ingredient(meal_id: int, item_id: int, db: DbSession) -> dict[str, bool]:
"""Remove an ingredient from a meal."""
mi = db.scalar(select(MealIngredient).where(MealIngredient.meal_id == meal_id, MealIngredient.item_id == item_id))
if not mi:
raise HTTPException(status_code=404, detail="Ingredient not found")
db.delete(mi)
db.commit()
return {"deleted": True}
@router.get("/meals/{meal_id}/availability", response_model=MealAvailability)
def check_meal(meal_id: int, db: DbSession) -> MealAvailability:
"""Check if a specific meal can be made and what's missing."""
meal = db.scalar(
select(Meal).where(Meal.id == meal_id).options(selectinload(Meal.ingredients).selectinload(MealIngredient.item))
)
if not meal:
raise HTTPException(status_code=404, detail="Meal not found")
return _check_meal(meal)
def _check_meal(meal: Meal) -> MealAvailability:
missing = [
ShoppingItem(
item_name=mi.item.name,
unit=mi.item.unit,
needed=mi.quantity_needed,
have=mi.item.quantity,
short=mi.quantity_needed - mi.item.quantity,
)
for mi in meal.ingredients
if mi.item.quantity < mi.quantity_needed
]
return MealAvailability(
meal_id=meal.id,
meal_name=meal.name,
can_make=len(missing) == 0,
missing=missing,
)
-198
View File
@@ -1,198 +0,0 @@
"""HTMX frontend routes for van inventory."""
from __future__ import annotations
from pathlib import Path
from typing import Annotated
from fastapi import APIRouter, Form, HTTPException, Request
from fastapi.responses import HTMLResponse
from fastapi.templating import Jinja2Templates
from sqlalchemy import select
from sqlalchemy.orm import selectinload
from python.orm.van_inventory.models import Item, Meal, MealIngredient
# FastAPI needs DbSession at runtime to resolve the Depends() annotation
from python.van_inventory.dependencies import DbSession # noqa: TC001
from python.van_inventory.routers.api import _check_meal
TEMPLATE_DIR = Path(__file__).resolve().parent.parent / "templates"
templates = Jinja2Templates(directory=TEMPLATE_DIR)
router = APIRouter(tags=["frontend"])
# --- Items ---
@router.get("/", response_class=HTMLResponse)
def items_page(request: Request, db: DbSession) -> HTMLResponse:
"""Render the inventory page."""
items = list(db.scalars(select(Item).order_by(Item.name)).all())
return templates.TemplateResponse(request, "items.html", {"items": items})
@router.post("/items", response_class=HTMLResponse)
def htmx_create_item(
request: Request,
db: DbSession,
name: Annotated[str, Form()],
quantity: Annotated[float, Form()] = 0,
unit: Annotated[str, Form()] = "",
category: Annotated[str | None, Form()] = None,
) -> HTMLResponse:
"""Create an item and return updated item rows."""
if quantity < 0:
raise HTTPException(status_code=422, detail="Quantity must not be negative")
db.add(Item(name=name, quantity=quantity, unit=unit, category=category or None))
db.commit()
items = list(db.scalars(select(Item).order_by(Item.name)).all())
return templates.TemplateResponse(request, "partials/item_rows.html", {"items": items})
@router.patch("/items/{item_id}", response_class=HTMLResponse)
def htmx_update_item(
request: Request,
item_id: int,
db: DbSession,
quantity: Annotated[float, Form()],
) -> HTMLResponse:
"""Update an item's quantity and return updated item rows."""
if quantity < 0:
raise HTTPException(status_code=422, detail="Quantity must not be negative")
item = db.get(Item, item_id)
if item:
item.quantity = quantity
db.commit()
items = list(db.scalars(select(Item).order_by(Item.name)).all())
return templates.TemplateResponse(request, "partials/item_rows.html", {"items": items})
@router.delete("/items/{item_id}", response_class=HTMLResponse)
def htmx_delete_item(request: Request, item_id: int, db: DbSession) -> HTMLResponse:
"""Delete an item and return updated item rows."""
item = db.get(Item, item_id)
if item:
db.delete(item)
db.commit()
items = list(db.scalars(select(Item).order_by(Item.name)).all())
return templates.TemplateResponse(request, "partials/item_rows.html", {"items": items})
# --- Meals ---
def _load_meals(db: DbSession) -> list[Meal]:
return list(
db.scalars(
select(Meal).options(selectinload(Meal.ingredients).selectinload(MealIngredient.item)).order_by(Meal.name)
).all()
)
@router.get("/meals", response_class=HTMLResponse)
def meals_page(request: Request, db: DbSession) -> HTMLResponse:
"""Render the meals page."""
meals = _load_meals(db)
return templates.TemplateResponse(request, "meals.html", {"meals": meals})
@router.post("/meals", response_class=HTMLResponse)
def htmx_create_meal(
request: Request,
db: DbSession,
name: Annotated[str, Form()],
instructions: Annotated[str | None, Form()] = None,
) -> HTMLResponse:
"""Create a meal and return updated meal rows."""
db.add(Meal(name=name, instructions=instructions or None))
db.commit()
meals = _load_meals(db)
return templates.TemplateResponse(request, "partials/meal_rows.html", {"meals": meals})
@router.delete("/meals/{meal_id}", response_class=HTMLResponse)
def htmx_delete_meal(request: Request, meal_id: int, db: DbSession) -> HTMLResponse:
"""Delete a meal and return updated meal rows."""
meal = db.get(Meal, meal_id)
if meal:
db.delete(meal)
db.commit()
meals = _load_meals(db)
return templates.TemplateResponse(request, "partials/meal_rows.html", {"meals": meals})
# --- Meal detail ---
def _load_meal(db: DbSession, meal_id: int) -> Meal | None:
return db.scalar(
select(Meal).where(Meal.id == meal_id).options(selectinload(Meal.ingredients).selectinload(MealIngredient.item))
)
@router.get("/meals/{meal_id}", response_class=HTMLResponse)
def meal_detail_page(request: Request, meal_id: int, db: DbSession) -> HTMLResponse:
"""Render the meal detail page."""
meal = _load_meal(db, meal_id)
if not meal:
raise HTTPException(status_code=404, detail="Meal not found")
items = list(db.scalars(select(Item).order_by(Item.name)).all())
return templates.TemplateResponse(request, "meal_detail.html", {"meal": meal, "items": items})
@router.post("/meals/{meal_id}/ingredients", response_class=HTMLResponse)
def htmx_add_ingredient(
request: Request,
meal_id: int,
db: DbSession,
item_id: Annotated[int, Form()],
quantity_needed: Annotated[float, Form()],
) -> HTMLResponse:
"""Add an ingredient to a meal and return updated ingredient rows."""
if quantity_needed <= 0:
raise HTTPException(status_code=422, detail="Quantity must be positive")
meal = db.get(Meal, meal_id)
if not meal:
raise HTTPException(status_code=404, detail="Meal not found")
if not db.get(Item, item_id):
raise HTTPException(status_code=422, detail="Item not found")
existing = db.scalar(
select(MealIngredient).where(MealIngredient.meal_id == meal_id, MealIngredient.item_id == item_id)
)
if existing:
raise HTTPException(status_code=409, detail="Ingredient already exists for this meal")
db.add(MealIngredient(meal_id=meal_id, item_id=item_id, quantity_needed=quantity_needed))
db.commit()
meal = _load_meal(db, meal_id)
return templates.TemplateResponse(request, "partials/ingredient_rows.html", {"meal": meal})
@router.delete("/meals/{meal_id}/ingredients/{item_id}", response_class=HTMLResponse)
def htmx_remove_ingredient(
request: Request,
meal_id: int,
item_id: int,
db: DbSession,
) -> HTMLResponse:
"""Remove an ingredient from a meal and return updated ingredient rows."""
mi = db.scalar(select(MealIngredient).where(MealIngredient.meal_id == meal_id, MealIngredient.item_id == item_id))
if mi:
db.delete(mi)
db.commit()
meal = _load_meal(db, meal_id)
return templates.TemplateResponse(request, "partials/ingredient_rows.html", {"meal": meal})
# --- Availability ---
@router.get("/availability", response_class=HTMLResponse)
def availability_page(request: Request, db: DbSession) -> HTMLResponse:
"""Render the meal availability page."""
meals = list(
db.scalars(select(Meal).options(selectinload(Meal.ingredients).selectinload(MealIngredient.item))).all()
)
availability = [_check_meal(m) for m in meals]
return templates.TemplateResponse(request, "availability.html", {"availability": availability})
-212
View File
@@ -1,212 +0,0 @@
:root {
--neon-pink: #ff2a6d;
--neon-cyan: #05d9e8;
--neon-yellow: #f9f002;
--neon-purple: #d300c5;
--bg-dark: #0a0a0f;
--bg-panel: #0d0d1a;
--bg-input: #111128;
--border: #1a1a3e;
--text: #c0c0d0;
--text-dim: #8e8ea0;
}
* { box-sizing: border-box; margin: 0; padding: 0; }
body {
font-family: 'Share Tech Mono', monospace;
max-width: 900px;
margin: 0 auto;
padding: 1rem;
background: var(--bg-dark);
color: var(--text);
position: relative;
}
/* Scanline overlay */
body::before {
content: '';
position: fixed;
top: 0; left: 0; right: 0; bottom: 0;
background: repeating-linear-gradient(
0deg,
transparent,
transparent 2px,
rgba(0, 0, 0, 0.08) 2px,
rgba(0, 0, 0, 0.08) 4px
);
pointer-events: none;
z-index: 9999;
}
h1, h2, h3 {
font-family: 'Orbitron', sans-serif;
margin-bottom: 0.5rem;
color: var(--neon-cyan);
text-shadow: 0 0 10px rgba(5, 217, 232, 0.5), 0 0 40px rgba(5, 217, 232, 0.2);
text-transform: uppercase;
letter-spacing: 2px;
}
a { color: var(--neon-pink); text-decoration: none; transition: all 0.2s; }
a:hover {
text-shadow: 0 0 8px rgba(255, 42, 109, 0.8), 0 0 20px rgba(255, 42, 109, 0.4);
}
nav {
display: flex;
gap: 1.5rem;
padding: 1rem 0;
border-bottom: 1px solid var(--border);
margin-bottom: 1.5rem;
position: relative;
}
nav::after {
content: '';
position: absolute;
bottom: -1px;
left: 0;
right: 0;
height: 1px;
background: linear-gradient(90deg, var(--neon-pink), var(--neon-cyan), var(--neon-purple));
opacity: 0.6;
}
nav a {
font-family: 'Orbitron', sans-serif;
font-weight: 700;
font-size: 0.85rem;
letter-spacing: 1px;
text-transform: uppercase;
padding: 0.3rem 0;
border-bottom: 2px solid transparent;
transition: all 0.2s;
}
nav a:hover {
border-bottom-color: var(--neon-pink);
text-shadow: 0 0 8px rgba(255, 42, 109, 0.8);
}
table {
width: 100%;
border-collapse: collapse;
margin: 1rem 0;
border: 1px solid var(--border);
}
th, td {
text-align: left;
padding: 0.6rem 0.75rem;
border-bottom: 1px solid var(--border);
}
th {
font-family: 'Orbitron', sans-serif;
color: var(--neon-cyan);
font-size: 0.7rem;
text-transform: uppercase;
letter-spacing: 2px;
background: var(--bg-panel);
border-bottom: 1px solid var(--neon-cyan);
text-shadow: 0 0 6px rgba(5, 217, 232, 0.3);
}
tr:hover td {
background: rgba(5, 217, 232, 0.03);
}
form {
display: flex;
flex-wrap: wrap;
gap: 0.5rem;
align-items: end;
margin: 1rem 0;
padding: 1rem;
border: 1px solid var(--border);
background: var(--bg-panel);
}
input, select {
padding: 0.5rem 0.6rem;
border: 1px solid var(--border);
border-radius: 2px;
background: var(--bg-input);
color: var(--neon-cyan);
font-family: 'Share Tech Mono', monospace;
transition: all 0.2s;
}
input:focus, select:focus {
outline: none;
border-color: var(--neon-cyan);
box-shadow: 0 0 8px rgba(5, 217, 232, 0.3), inset 0 0 8px rgba(5, 217, 232, 0.05);
}
button {
padding: 0.5rem 1.2rem;
border: 1px solid var(--neon-pink);
border-radius: 2px;
background: transparent;
color: var(--neon-pink);
cursor: pointer;
font-family: 'Orbitron', sans-serif;
font-weight: 700;
font-size: 0.7rem;
letter-spacing: 1px;
text-transform: uppercase;
transition: all 0.2s;
}
button:hover {
background: var(--neon-pink);
color: var(--bg-dark);
box-shadow: 0 0 15px rgba(255, 42, 109, 0.5), 0 0 30px rgba(255, 42, 109, 0.2);
}
button.danger {
border-color: var(--text-dim);
color: var(--text-dim);
}
button.danger:hover {
border-color: var(--neon-pink);
background: var(--neon-pink);
color: var(--bg-dark);
box-shadow: 0 0 15px rgba(255, 42, 109, 0.5);
}
.badge {
display: inline-block;
padding: 0.2rem 0.6rem;
border-radius: 2px;
font-family: 'Orbitron', sans-serif;
font-size: 0.65rem;
font-weight: 700;
letter-spacing: 1px;
text-transform: uppercase;
}
.badge.yes {
background: rgba(5, 217, 232, 0.1);
color: var(--neon-cyan);
border: 1px solid var(--neon-cyan);
text-shadow: 0 0 6px rgba(5, 217, 232, 0.5);
}
.badge.no {
background: rgba(255, 42, 109, 0.1);
color: var(--neon-pink);
border: 1px solid var(--neon-pink);
text-shadow: 0 0 6px rgba(255, 42, 109, 0.5);
}
.missing-list { font-size: 0.85rem; color: var(--text-dim); }
label {
font-size: 0.75rem;
color: var(--text-dim);
display: flex;
flex-direction: column;
gap: 0.2rem;
text-transform: uppercase;
letter-spacing: 1px;
}
.flash {
padding: 0.5rem 1rem;
margin: 0.5rem 0;
border-radius: 2px;
background: rgba(5, 217, 232, 0.1);
color: var(--neon-cyan);
border: 1px solid var(--neon-cyan);
}
@@ -1,30 +0,0 @@
{% extends "base.html" %}
{% block title %}What Can I Make? - Van{% endblock %}
{% block content %}
<h1>What Can I Make?</h1>
<table>
<thead>
<tr><th>Meal</th><th>Status</th><th>Missing</th></tr>
</thead>
<tbody>
{% for meal in availability %}
<tr>
<td><a href="/meals/{{ meal.meal_id }}">{{ meal.meal_name }}</a></td>
<td>
{% if meal.can_make %}
<span class="badge yes">Ready</span>
{% else %}
<span class="badge no">Missing items</span>
{% endif %}
</td>
<td class="missing-list">
{% for m in meal.missing %}
{{ m.item_name }}: need {{ m.short }} more {{ m.unit }}{% if not loop.last %}, {% endif %}
{% endfor %}
</td>
</tr>
{% endfor %}
</tbody>
</table>
{% endblock %}
-20
View File
@@ -1,20 +0,0 @@
<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>{% block title %}Van Inventory{% endblock %}</title>
<script src="https://unpkg.com/htmx.org@2.0.4"></script>
<link rel="preconnect" href="https://fonts.googleapis.com">
<link href="https://fonts.googleapis.com/css2?family=Orbitron:wght@400;700;900&family=Share+Tech+Mono&display=swap" rel="stylesheet">
<link rel="stylesheet" href="/static/style.css">
</head>
<body>
<nav>
<a href="/">Inventory</a>
<a href="/meals">Meals</a>
<a href="/availability">What Can I Make?</a>
</nav>
{% block content %}{% endblock %}
</body>
</html>
-17
View File
@@ -1,17 +0,0 @@
{% extends "base.html" %}
{% block title %}Inventory - Van{% endblock %}
{% block content %}
<h1>Van Inventory</h1>
<form hx-post="/items" hx-target="#item-list" hx-swap="innerHTML" hx-on::after-request="if(event.detail.successful) this.reset()">
<label>Name <input type="text" name="name" required></label>
<label>Qty <input type="number" name="quantity" step="any" value="0" min="0" required></label>
<label>Unit <input type="text" name="unit" required placeholder="lbs, cans, etc"></label>
<label>Category <input type="text" name="category" placeholder="optional"></label>
<button type="submit">Add Item</button>
</form>
<div id="item-list">
{% include "partials/item_rows.html" %}
</div>
{% endblock %}
@@ -1,24 +0,0 @@
{% extends "base.html" %}
{% block title %}{{ meal.name }} - Van{% endblock %}
{% block content %}
<h1>{{ meal.name }}</h1>
{% if meal.instructions %}<p>{{ meal.instructions }}</p>{% endif %}
<h2>Ingredients</h2>
<form hx-post="/meals/{{ meal.id }}/ingredients" hx-target="#ingredient-list" hx-swap="innerHTML" hx-on::after-request="if(event.detail.successful) this.reset()">
<label>Item
<select name="item_id" required>
<option value="">--</option>
{% for item in items %}
<option value="{{ item.id }}">{{ item.name }} ({{ item.unit }})</option>
{% endfor %}
</select>
</label>
<label>Qty needed <input type="number" name="quantity_needed" step="any" min="0.01" required></label>
<button type="submit">Add</button>
</form>
<div id="ingredient-list">
{% include "partials/ingredient_rows.html" %}
</div>
{% endblock %}
-15
View File
@@ -1,15 +0,0 @@
{% extends "base.html" %}
{% block title %}Meals - Van{% endblock %}
{% block content %}
<h1>Meals</h1>
<form hx-post="/meals" hx-target="#meal-list" hx-swap="innerHTML" hx-on::after-request="if(event.detail.successful) this.reset()">
<label>Name <input type="text" name="name" required></label>
<label>Instructions <input type="text" name="instructions" placeholder="optional"></label>
<button type="submit">Add Meal</button>
</form>
<div id="meal-list">
{% include "partials/meal_rows.html" %}
</div>
{% endblock %}
@@ -1,16 +0,0 @@
<table>
<thead>
<tr><th>Item</th><th>Needed</th><th>Have</th><th>Unit</th><th></th></tr>
</thead>
<tbody>
{% for mi in meal.ingredients %}
<tr>
<td>{{ mi.item.name }}</td>
<td>{{ mi.quantity_needed }}</td>
<td>{{ mi.item.quantity }}</td>
<td>{{ mi.item.unit }}</td>
<td><button class="danger" hx-delete="/meals/{{ meal.id }}/ingredients/{{ mi.item_id }}" hx-target="#ingredient-list" hx-swap="innerHTML" hx-confirm="Remove {{ mi.item.name }}?">X</button></td>
</tr>
{% endfor %}
</tbody>
</table>
@@ -1,21 +0,0 @@
<table>
<thead>
<tr><th>Name</th><th>Qty</th><th>Unit</th><th>Category</th><th></th></tr>
</thead>
<tbody>
{% for item in items %}
<tr>
<td>{{ item.name }}</td>
<td>
<form hx-patch="/items/{{ item.id }}" hx-target="#item-list" hx-swap="innerHTML" style="display:inline; margin:0;">
<input type="number" name="quantity" value="{{ item.quantity }}" step="any" min="0" style="width:5rem">
<button type="submit" style="padding:0.2rem 0.5rem; font-size:0.8rem;">Update</button>
</form>
</td>
<td>{{ item.unit }}</td>
<td>{{ item.category or "" }}</td>
<td><button class="danger" hx-delete="/items/{{ item.id }}" hx-target="#item-list" hx-swap="innerHTML" hx-confirm="Delete {{ item.name }}?">X</button></td>
</tr>
{% endfor %}
</tbody>
</table>
@@ -1,15 +0,0 @@
<table>
<thead>
<tr><th>Name</th><th>Ingredients</th><th>Instructions</th><th></th></tr>
</thead>
<tbody>
{% for meal in meals %}
<tr>
<td><a href="/meals/{{ meal.id }}">{{ meal.name }}</a></td>
<td>{{ meal.ingredients | length }}</td>
<td>{{ (meal.instructions or "")[:50] }}</td>
<td><button class="danger" hx-delete="/meals/{{ meal.id }}" hx-target="#meal-list" hx-swap="innerHTML" hx-confirm="Delete {{ meal.name }}?">X</button></td>
</tr>
{% endfor %}
</tbody>
</table>
+1 -1
View File
@@ -257,7 +257,7 @@ def update_weather(config: Config) -> None:
logger.info(f"Masked location: {masked_lat}, {masked_lon}") logger.info(f"Masked location: {masked_lat}, {masked_lon}")
weather = fetch_weather(config.pirate_weather_api_key, lat, lon) weather = fetch_weather(config.pirate_weather_api_key, masked_lat, masked_lon)
logger.info(f"Weather: {weather.temperature}°F, {weather.condition}") logger.info(f"Weather: {weather.temperature}°F, {weather.condition}")
post_to_ha(config.ha_url, config.ha_token, weather) post_to_ha(config.ha_url, config.ha_token, weather)
+3 -1
View File
@@ -1,6 +1,8 @@
"""Models for van weather service.""" """Models for van weather service."""
from datetime import datetime from __future__ import annotations
from datetime import datetime # noqa: TC003 This is required for pydantic
from pydantic import BaseModel, field_serializer from pydantic import BaseModel, field_serializer
+2 -2
View File
@@ -108,7 +108,7 @@ class Dataset:
self.written = int(properties["written"]["value"]) self.written = int(properties["written"]["value"])
self.xattr = properties["xattr"]["value"] self.xattr = properties["xattr"]["value"]
def get_snapshots(self) -> list[Snapshot] | None: def get_snapshots(self) -> list[Snapshot]:
"""Get all snapshots from zfs and process then is test dicts of sets.""" """Get all snapshots from zfs and process then is test dicts of sets."""
snapshots_data = _zfs_list(f"zfs list -t snapshot -pHj {self.name} -o all") snapshots_data = _zfs_list(f"zfs list -t snapshot -pHj {self.name} -o all")
@@ -125,7 +125,7 @@ class Dataset:
if return_code == 0: if return_code == 0:
return "snapshot created" return "snapshot created"
if snapshots := self.get_snapshots(): snapshots = self.get_snapshots()
snapshot_names = {snapshot.name for snapshot in snapshots} snapshot_names = {snapshot.name for snapshot in snapshots}
if snapshot_name in snapshot_names: if snapshot_name in snapshot_names:
return f"Snapshot {snapshot_name} already exists for {self.name}" return f"Snapshot {snapshot_name} already exists for {self.name}"
+2
View File
@@ -32,6 +32,8 @@
enable = true; enable = true;
allowedTCPPorts = [ allowedTCPPorts = [
8000 8000
8001
8002
]; ];
}; };
networkmanager.enable = true; networkmanager.enable = true;
+1 -1
View File
@@ -4,7 +4,7 @@
host = "0.0.0.0"; host = "0.0.0.0";
enable = true; enable = true;
syncModels = true; syncModels = false;
loadModels = [ loadModels = [
"codellama:7b" "codellama:7b"
"deepscaler:1.5b" "deepscaler:1.5b"
-50
View File
@@ -1,50 +0,0 @@
{
pkgs,
inputs,
...
}:
{
networking.firewall.allowedTCPPorts = [ 8001 ];
users = {
users.vaninventory = {
isSystemUser = true;
group = "vaninventory";
};
groups.vaninventory = { };
};
systemd.services.van_inventory = {
description = "Van Inventory API";
after = [
"network.target"
"postgresql.service"
];
requires = [ "postgresql.service" ];
wantedBy = [ "multi-user.target" ];
environment = {
PYTHONPATH = "${inputs.self}/";
VAN_INVENTORY_DB = "vaninventory";
VAN_INVENTORY_USER = "vaninventory";
VAN_INVENTORY_HOST = "/run/postgresql";
VAN_INVENTORY_PORT = "5432";
};
serviceConfig = {
Type = "simple";
User = "vaninventory";
Group = "vaninventory";
ExecStart = "${pkgs.my_python}/bin/python -m python.van_inventory.main --host 0.0.0.0 --port 8001";
Restart = "on-failure";
RestartSec = "5s";
StandardOutput = "journal";
StandardError = "journal";
NoNewPrivileges = true;
ProtectSystem = "strict";
ProtectHome = "read-only";
PrivateTmp = true;
ReadOnlyPaths = [ "${inputs.self}" ];
};
};
}
+3
View File
@@ -17,6 +17,9 @@
allowedTCPPorts = [ ]; allowedTCPPorts = [ ];
allowedUDPPorts = [ ]; allowedUDPPorts = [ ];
}; };
allowedTCPPorts = [
8070
];
}; };
useNetworkd = true; useNetworkd = true;
}; };
+4 -1
View File
@@ -6,7 +6,7 @@ in
user = "ollama"; user = "ollama";
enable = true; enable = true;
host = "0.0.0.0"; host = "0.0.0.0";
syncModels = true; syncModels = false;
loadModels = [ loadModels = [
"codellama:7b" "codellama:7b"
"deepscaler:1.5b" "deepscaler:1.5b"
@@ -30,6 +30,9 @@ in
"ministral-3:14b" "ministral-3:14b"
"nemotron-3-nano:30b" "nemotron-3-nano:30b"
"qwen3-coder:30b" "qwen3-coder:30b"
"qwen3-embedding:0.6b"
"qwen3-embedding:4b"
"qwen3-embedding:8b"
"qwen3-vl:32b" "qwen3-vl:32b"
"qwen3:14b" "qwen3:14b"
"qwen3.5:35b" "qwen3.5:35b"
-11
View File
@@ -38,9 +38,6 @@ in
# signalbot # signalbot
local signalbot signalbot trust local signalbot signalbot trust
# hedgedoc
local hedgedoc hedgedoc trust
# math # math
local postgres math trust local postgres math trust
host postgres math 127.0.0.1/32 trust host postgres math 127.0.0.1/32 trust
@@ -120,19 +117,11 @@ in
login = true; login = true;
}; };
} }
{
name = "hedgedoc";
ensureDBOwnership = true;
ensureClauses = {
login = true;
};
}
]; ];
ensureDatabases = [ ensureDatabases = [
"data_science_dev" "data_science_dev"
"hass" "hass"
"gitea" "gitea"
"hedgedoc"
"math" "math"
"n8n" "n8n"
"richie" "richie"
-2
View File
@@ -1021,8 +1021,6 @@ def test_existing_destination_skips_rename_and_removes_temp(tmp_path, monkeypatc
def test_richie_exports_audiobook_models() -> None: def test_richie_exports_audiobook_models() -> None:
from python.orm.richie import Audiobook # noqa: PLC0415
assert Audiobook.__tablename__ == "audiobook" assert Audiobook.__tablename__ == "audiobook"
+536
View File
@@ -0,0 +1,536 @@
"""Tests for EPUB search core helpers."""
from __future__ import annotations
import logging
from dataclasses import replace
from datetime import UTC, datetime
from os import environ
from pathlib import Path
from threading import Event
from types import ModuleType
import pytest
from sqlalchemy import create_engine, select
from sqlalchemy.orm import sessionmaker
from python.ebook_search.answer import answer_query
from python.ebook_search.bm25_corpus import (
BM25Corpus,
BM25CorpusUnavailableError,
BM25Manifest,
ensure_bm25_corpus,
fetch_bm25_corpus_records,
load_bm25_corpus,
read_bm25_manifest,
score_bm25_corpus,
write_bm25_corpus,
)
from python.ebook_search.config import EbookSearchConfig, RerankConfig, load_config, normalize_embedding_model
from python.ebook_search.embeddings import MODEL_DIMENSIONS, ensure_embedding_models
from python.ebook_search.ingest import chunk_text, find_existing_source
from python.ebook_search.search import (
SearchResponse,
SearchResult,
bm25_candidates,
reciprocal_rank_fusion,
retrieval_query_from_text,
search_ebooks,
)
from python.ebook_search.timing import RuntimeStep
from python.orm.richie import (
EbookChapter,
EbookChunk,
EbookChunkEmbedding1024,
EbookEmbeddingModel,
EbookSource,
RichieBase,
)
def test_chunk_text_uses_overlap() -> None:
chunks = chunk_text(" ".join(str(index) for index in range(100)), chunk_tokens=20, overlap_tokens=5)
assert len(chunks) > 1
assert chunks[0].token_start == 0
assert chunks[1].token_start == 15
assert all(chunk.token_count <= 20 for chunk in chunks)
def test_reciprocal_rank_fusion_combines_vector_and_bm25_rankings() -> None:
vector_results = [
SearchResult(chunk_id=1, text="a", source_title="A", score=0.9, vector_score=0.9),
SearchResult(chunk_id=2, text="b", source_title="B", score=0.8, vector_score=0.8),
]
lexical_results = [
SearchResult(chunk_id=2, text="b", source_title="B", score=4.2, bm25_score=4.2),
SearchResult(chunk_id=3, text="c", source_title="C", score=2.1, bm25_score=2.1),
]
fused = reciprocal_rank_fusion(vector_results, lexical_results)
assert [result.chunk_id for result in fused] == [2, 1, 3]
assert fused[0].rank_source == "Hybrid"
assert fused[0].vector_score == 0.8
assert fused[0].bm25_score == 4.2
assert fused[0].fused_score == fused[0].score
def test_find_existing_source_matches_path_or_hash() -> None:
engine = create_engine("sqlite+pysqlite:///:memory:", future=True)
RichieBase.metadata.create_all(engine)
with sessionmaker(bind=engine, expire_on_commit=False, future=True)() as session:
source = EbookSource(
title="Book",
author=None,
language=None,
publisher=None,
identifier=None,
file_path="/old/book.epub",
file_sha256="a" * 64,
file_mtime=datetime.now(tz=UTC),
file_size=10,
)
session.add(source)
session.commit()
assert find_existing_source(session, Path("/old/book.epub"), "b" * 64) == source
assert find_existing_source(session, Path("/new/book.epub"), "a" * 64) == source
def test_bm25_corpus_uses_existing_search_text_without_duplicate_metadata() -> None:
engine = create_engine("sqlite+pysqlite:///:memory:", future=True)
RichieBase.metadata.create_all(engine)
with sessionmaker(bind=engine, expire_on_commit=False, future=True)() as session:
source = EbookSource(
title="Book",
author="Author",
language=None,
publisher=None,
identifier=None,
file_path="/book.epub",
file_sha256="a" * 64,
file_mtime=datetime.now(tz=UTC),
file_size=10,
)
session.add(source)
session.flush()
chapter = EbookChapter(source_id=source.id, spine_index=0, title="Chapter", href=None)
session.add(chapter)
session.flush()
session.add(
EbookChunk(
id=1,
source_id=source.id,
chapter_id=chapter.id,
chunk_index=0,
text="content",
token_start=0,
token_count=1,
page_label=None,
content_sha256="b" * 64,
search_text="Book Author Chapter content",
)
)
session.commit()
records, texts = fetch_bm25_corpus_records(session)
assert texts == ["Book Author Chapter content"]
assert records[0]["chunk_id"] == 1
assert "bm25_text" not in records[0]
def test_reciprocal_rank_fusion_marks_hybrid_source() -> None:
vector_results = [SearchResult(chunk_id=1, text="a", source_title="A")]
lexical_results = [SearchResult(chunk_id=2, text="b", source_title="B")]
fused = reciprocal_rank_fusion(vector_results, lexical_results)
assert {result.rank_source for result in fused} == {"Hybrid"}
def test_search_response_sums_runtime_steps() -> None:
response = SearchResponse(
query="query",
results=[],
rank_label="Hybrid",
timings=(
RuntimeStep(name="A", duration_ms=1.25),
RuntimeStep(name="B", duration_ms=2.75),
RuntimeStep(name="Parallel detail", duration_ms=10.0, counts_toward_total=False),
),
)
assert response.total_runtime_ms == 4.0
def test_search_ebooks_runs_vector_and_bm25_in_parallel(monkeypatch) -> None:
engine = create_engine("sqlite+pysqlite:///:memory:", future=True)
vector_started = Event()
bm25_started = Event()
received_engines: list[object] = []
def fake_vector_candidates(received_engine, query, _config):
"""Return vector candidates after confirming BM25 has started."""
received_engines.append(received_engine)
assert query == "what is parallel"
vector_started.set()
assert bm25_started.wait(timeout=2)
return [SearchResult(chunk_id=1, text="vector", source_title="Vector", vector_score=0.9)]
def fake_bm25_candidates(query, _config):
"""Return BM25 candidates after confirming vector search has started."""
assert query == "parallel"
bm25_started.set()
assert vector_started.wait(timeout=2)
return [SearchResult(chunk_id=2, text="bm25", source_title="BM25", bm25_score=2.0)]
monkeypatch.setattr("python.ebook_search.search.vector_candidates", fake_vector_candidates)
monkeypatch.setattr("python.ebook_search.search.bm25_candidates", fake_bm25_candidates)
config = EbookSearchConfig(rerank=RerankConfig(enabled=False))
response = search_ebooks(engine, "what is parallel", config)
timings = {step.name: step for step in response.timings}
assert [result.chunk_id for result in response.results] == [1, 2]
assert timings["Embedding + vector search"].counts_toward_total is False
assert timings["BM25 search"].counts_toward_total is False
assert timings["Hybrid retrieval"].counts_toward_total is True
assert timings["BM25 query preparation"].counts_toward_total is True
assert received_engines == [engine]
def test_retrieval_query_keeps_entity_and_series_terms() -> None:
assert retrieval_query_from_text("what does Damien Montgomery stand for in starship mage") == (
"damien montgomery stand starship mage"
)
def test_bm25_candidates_scores_whole_corpus(monkeypatch) -> None:
record = {
"chunk_id": 2,
"text": "high",
"source_title": "B",
"source_author": None,
"chapter_title": None,
"page_label": None,
"bm25_text": "high",
}
manifest = BM25Manifest(created_at=datetime.now(tz=UTC), db_updated_at=None, chunk_count=1)
corpus = BM25Corpus(retriever=object(), records=(record,), manifest=manifest)
captured: dict[str, object] = {}
def fake_score_bm25_corpus(query, saved_corpus, *, limit):
captured["query"] = query
captured["corpus"] = saved_corpus
captured["limit"] = limit
return [(record, 1.5)]
monkeypatch.setattr("python.ebook_search.search.load_bm25_corpus", lambda _config: corpus)
monkeypatch.setattr("python.ebook_search.search.score_bm25_corpus", fake_score_bm25_corpus)
config = EbookSearchConfig(rerank=RerankConfig(enabled=False))
results = bm25_candidates("high", config)
assert captured["query"] == "high"
assert captured["corpus"] == corpus
assert captured["limit"] == 120
assert [result.chunk_id for result in results] == [2]
assert [result.bm25_score for result in results] == [1.5]
def test_bm25_candidates_returns_empty_when_corpus_is_unavailable(monkeypatch, caplog) -> None:
def fake_load_bm25_corpus(_config):
raise BM25CorpusUnavailableError
monkeypatch.setattr("python.ebook_search.search.load_bm25_corpus", fake_load_bm25_corpus)
config = EbookSearchConfig(rerank=RerankConfig(enabled=False))
with caplog.at_level(logging.WARNING):
results = bm25_candidates("high", config)
assert results == []
assert "ebook_bm25_index_unavailable_skipping" in caplog.text
def test_write_bm25_corpus_publishes_dated_generation(tmp_path) -> None:
index_path = tmp_path / "bm25"
index_path.mkdir()
generations_path = index_path / "generations"
generations_path.mkdir()
old_generation = generations_path / "20260101T000000.000000Z"
old_generation.mkdir()
(old_generation / "sentinel").write_text("old", encoding="utf-8")
(index_path / "current").symlink_to(Path("generations") / old_generation.name, target_is_directory=True)
manifest = BM25Manifest(
created_at=datetime(2026, 6, 12, 1, 2, 3, 456789, tzinfo=UTC),
db_updated_at=None,
chunk_count=0,
)
write_bm25_corpus(index_path, [], [], manifest)
current_path = index_path / "current"
assert current_path.is_symlink()
assert current_path.readlink() == generations_path / "20260612T010203.456789Z"
assert old_generation.is_dir()
assert (old_generation / "sentinel").read_text(encoding="utf-8") == "old"
assert (generations_path / "20260612T010203.456789Z").is_dir()
assert read_bm25_manifest(index_path) == manifest
def test_write_bm25_corpus_keeps_current_generation_when_publish_fails(monkeypatch, tmp_path) -> None:
index_path = tmp_path / "bm25"
index_path.mkdir()
generations_path = index_path / "generations"
generations_path.mkdir()
old_generation = generations_path / "20260101T000000.000000Z"
old_generation.mkdir()
(old_generation / "sentinel").write_text("old", encoding="utf-8")
current_path = index_path / "current"
current_path.symlink_to(Path("generations") / old_generation.name, target_is_directory=True)
original_replace = Path.replace
def fail_current_replace(self, target):
if self.parent == index_path and self.name.startswith(".current.") and target == current_path:
msg = "current publish failed"
raise OSError(msg)
return original_replace(self, target)
monkeypatch.setattr(Path, "replace", fail_current_replace)
manifest = BM25Manifest(
created_at=datetime(2026, 6, 12, 1, 2, 3, 456789, tzinfo=UTC),
db_updated_at=None,
chunk_count=0,
)
with pytest.raises(OSError, match="current publish failed"):
write_bm25_corpus(index_path, [], [], manifest)
assert current_path.readlink() == Path("generations") / old_generation.name
assert (old_generation / "sentinel").read_text(encoding="utf-8") == "old"
assert not (generations_path / "20260612T010203.456789Z").exists()
def test_load_bm25_corpus_uses_current_generation(tmp_path) -> None:
load_bm25_corpus.cache_clear()
index_path = tmp_path / "bm25"
manifest = BM25Manifest(
created_at=datetime(2026, 6, 12, 1, 2, 3, 456789, tzinfo=UTC),
db_updated_at=None,
chunk_count=1,
)
record = {
"chunk_id": 2,
"text": "cached",
"source_title": "B",
"source_author": None,
"chapter_title": None,
"page_label": None,
}
write_bm25_corpus(index_path, [record], ["cached phrase"], manifest)
config = EbookSearchConfig(rerank=RerankConfig(enabled=False), bm25_index_dir=str(index_path))
try:
corpus = load_bm25_corpus(config)
finally:
load_bm25_corpus.cache_clear()
assert corpus.manifest == manifest
assert corpus.records[0]["chunk_id"] == 2
assert score_bm25_corpus("cached", corpus, limit=10)
def test_load_bm25_corpus_caches_disk_load(monkeypatch, tmp_path) -> None:
load_bm25_corpus.cache_clear()
manifest = BM25Manifest(created_at=datetime.now(tz=UTC), db_updated_at=None, chunk_count=1)
record = {
"chunk_id": 2,
"text": "cached",
"source_title": "B",
"source_author": None,
"chapter_title": None,
"page_label": None,
"bm25_text": "cached",
}
load_count = 0
class FakeRetriever:
"""Fake persisted BM25 retriever."""
corpus = (record,)
class FakeBM25:
"""Fake BM25 class with observable load count."""
@staticmethod
def load(index_path, *, load_corpus, mmap):
nonlocal load_count
load_count += 1
assert index_path == tmp_path
assert load_corpus is True
assert mmap is True
return FakeRetriever()
fake_bm25s = ModuleType("bm25s")
fake_bm25s.BM25 = FakeBM25
monkeypatch.setattr("python.ebook_search.bm25_corpus.read_bm25_manifest", lambda _path: manifest)
monkeypatch.setattr("python.ebook_search.bm25_corpus.bm25_index_exists", lambda _path, _manifest: True)
monkeypatch.setattr("python.ebook_search.bm25_corpus.bm25s", fake_bm25s)
config = EbookSearchConfig(rerank=RerankConfig(enabled=False), bm25_index_dir=str(tmp_path))
try:
first = load_bm25_corpus(config)
second = load_bm25_corpus(config)
finally:
load_bm25_corpus.cache_clear()
assert first is second
assert first is not None
assert first.records == (record,)
assert load_count == 1
def test_load_bm25_corpus_raises_when_index_is_missing(monkeypatch, tmp_path) -> None:
load_bm25_corpus.cache_clear()
monkeypatch.setattr("python.ebook_search.bm25_corpus.read_bm25_manifest", lambda _path: None)
monkeypatch.setattr("python.ebook_search.bm25_corpus.bm25_index_exists", lambda _path, _manifest: False)
config = EbookSearchConfig(rerank=RerankConfig(enabled=False), bm25_index_dir=str(tmp_path))
try:
with pytest.raises(BM25CorpusUnavailableError, match="BM25 corpus is not available"):
load_bm25_corpus(config)
finally:
load_bm25_corpus.cache_clear()
def test_ensure_bm25_corpus_refreshes_missing_index(monkeypatch) -> None:
refreshed: list[object] = []
db_updated_at = datetime.now(tz=UTC)
monkeypatch.setattr("python.ebook_search.bm25_corpus.read_bm25_manifest", lambda _path: None)
monkeypatch.setattr("python.ebook_search.bm25_corpus.bm25_index_exists", lambda _path, _manifest: False)
monkeypatch.setattr("python.ebook_search.bm25_corpus.corpus_last_updated_at", lambda _session: db_updated_at)
monkeypatch.setattr(
"python.ebook_search.bm25_corpus.refresh_bm25_corpus",
lambda session, config, *, db_updated_at: refreshed.append((session, config, db_updated_at)),
)
config = EbookSearchConfig(rerank=RerankConfig(enabled=False))
session = object()
ensure_bm25_corpus(session, config)
assert refreshed == [(session, config, db_updated_at)]
def test_ensure_bm25_corpus_refreshes_stale_index(monkeypatch) -> None:
refreshed: list[object] = []
created_at = datetime(2026, 1, 1, tzinfo=UTC)
db_updated_at = datetime(2026, 1, 2, tzinfo=UTC)
manifest = BM25Manifest(created_at=created_at, db_updated_at=created_at, chunk_count=10)
monkeypatch.setattr("python.ebook_search.bm25_corpus.read_bm25_manifest", lambda _path: manifest)
monkeypatch.setattr("python.ebook_search.bm25_corpus.bm25_index_exists", lambda _path, _manifest: True)
monkeypatch.setattr("python.ebook_search.bm25_corpus.corpus_last_updated_at", lambda _session: db_updated_at)
monkeypatch.setattr(
"python.ebook_search.bm25_corpus.refresh_bm25_corpus",
lambda session, config, *, db_updated_at: refreshed.append((session, config, db_updated_at)),
)
config = EbookSearchConfig(rerank=RerankConfig(enabled=False))
session = object()
ensure_bm25_corpus(session, config)
assert refreshed == [(session, config, db_updated_at)]
def test_supported_embedding_models_match_service_names() -> None:
assert MODEL_DIMENSIONS == {
"qwen3-embedding-0.6b": 1024,
"qwen3-embedding-4b": 2560,
"qwen3-embedding-8b": 4096,
}
def test_ensure_embedding_models_registers_service_names() -> None:
engine = create_engine("sqlite+pysqlite:///:memory:", future=True)
RichieBase.metadata.create_all(engine)
with sessionmaker(bind=engine, expire_on_commit=False, future=True)() as session:
ensure_embedding_models(session)
session.commit()
models = list(session.scalars(select(EbookEmbeddingModel).order_by(EbookEmbeddingModel.name)))
assert [(model.name, model.dimension) for model in models] == [
("qwen3-embedding-0.6b", 1024),
("qwen3-embedding-4b", 2560),
("qwen3-embedding-8b", 4096),
]
def test_1024_embedding_table_has_cosine_hnsw_index() -> None:
indexes = {index.name: index for index in EbookChunkEmbedding1024.__table__.indexes}
index = indexes["ix_ebook_chunk_embedding_1024_embedding_cosine"]
assert [column.name for column in index.columns] == ["embedding"]
assert index.dialect_options["postgresql"]["using"] == "hnsw"
assert index.dialect_options["postgresql"]["ops"] == {"embedding": "vector_cosine_ops"}
def test_embedding_model_aliases_normalize_to_provider_names() -> None:
assert normalize_embedding_model() == "qwen3-embedding-0.6b"
environ["EBOOK_SEARCH_EMBEDDING_MODEL"] = "qwen3-embedding-0.6b"
assert normalize_embedding_model() == "qwen3-embedding-0.6b"
environ["EBOOK_SEARCH_EMBEDDING_MODEL"] = "Qwen3-Embedding-0.6B"
assert normalize_embedding_model() == "qwen3-embedding-0.6b"
environ["EBOOK_SEARCH_EMBEDDING_MODEL"] = "Qwen/Qwen3-Embedding-4B"
assert normalize_embedding_model() == "qwen3-embedding-4b"
environ["EBOOK_SEARCH_EMBEDDING_MODEL"] = "qwen3-embedding:8b"
assert normalize_embedding_model() == "qwen3-embedding-8b"
environ["EBOOK_SEARCH_EMBEDDING_MODEL"] = "qwen3-embedding-8b"
assert normalize_embedding_model() == "qwen3-embedding-8b"
def test_answer_generation_is_enabled_by_default(monkeypatch) -> None:
monkeypatch.delenv("EBOOK_SEARCH_ANSWER_ENABLED", raising=False)
config = load_config()
assert config.answer_enabled is True
def test_chat_defaults_use_ollama_cloud(monkeypatch) -> None:
monkeypatch.delenv("EBOOK_SEARCH_VLLM_BASE_URL", raising=False)
monkeypatch.delenv("EBOOK_SEARCH_CHAT_MODEL", raising=False)
config = load_config()
assert config.vllm_base_url == "https://ollama.com/v1"
assert config.chat_model == "deepseek-v4-flash"
def test_chat_api_key_falls_back_to_ollama_api_key(monkeypatch) -> None:
monkeypatch.delenv("EBOOK_SEARCH_VLLM_API_KEY", raising=False)
monkeypatch.setenv("OLLAMA_API_KEY", "ollama-key")
config = load_config()
assert config.vllm_api_key == "ollama-key"
def test_answer_query_does_not_call_model_when_disabled() -> None:
config = replace(load_config(), answer_enabled=False)
result = SearchResult(chunk_id=1, text="source text", source_title="Book")
answer = answer_query("question", [result], config)
assert "Answer generation is disabled" in answer
+84
View File
@@ -0,0 +1,84 @@
"""Tests for EPUB search HTTP model adapters."""
from __future__ import annotations
import httpx
import pytest
from python.ebook_search.answer import answer_query
from python.ebook_search.config import EbookSearchConfig, RerankConfig
from python.ebook_search.embeddings import embed_texts
from python.ebook_search.search import SearchResult
def test_answer_query_uses_httpx_chat_completions(monkeypatch) -> None:
captured: dict[str, object] = {}
def fake_post(url: str, **kwargs: object) -> httpx.Response:
captured["url"] = url
captured["kwargs"] = kwargs
return httpx.Response(
200,
json={"choices": [{"message": {"content": "grounded answer"}}]},
request=httpx.Request("POST", url),
)
monkeypatch.setattr(httpx, "post", fake_post)
config = EbookSearchConfig(
rerank=RerankConfig(enabled=False),
vllm_base_url="https://ollama.com/v1",
vllm_api_key="secret",
chat_model="deepseek-v4-flash",
)
answer = answer_query("question", [SearchResult(chunk_id=1, text="source", source_title="Book")], config)
assert answer == "grounded answer"
assert captured["url"] == "https://ollama.com/v1/chat/completions"
kwargs = captured["kwargs"]
assert isinstance(kwargs, dict)
assert kwargs["headers"] == {"Authorization": "Bearer secret"}
payload = kwargs["json"]
assert isinstance(payload, dict)
assert payload["model"] == "deepseek-v4-flash"
def test_embed_texts_uses_httpx_embeddings(monkeypatch) -> None:
captured: dict[str, object] = {}
vector = [0.0] * 1024
def fake_post(url: str, **kwargs: object) -> httpx.Response:
captured["url"] = url
captured["kwargs"] = kwargs
return httpx.Response(
200,
json={"data": [{"embedding": vector}]},
request=httpx.Request("POST", url),
)
monkeypatch.setattr(httpx, "post", fake_post)
config = EbookSearchConfig(
rerank=RerankConfig(enabled=False),
embedding_base_url="http://bob:8000/v1",
embedding_model="qwen3-embedding-0.6b",
)
embeddings = embed_texts(["hello"], config)
assert embeddings == [vector]
assert captured["url"] == "http://bob:8000/v1/embeddings"
kwargs = captured["kwargs"]
assert isinstance(kwargs, dict)
assert kwargs["headers"] == {}
assert kwargs["json"] == {"model": "qwen3-embedding-0.6b", "input": ["hello"]}
def test_embed_texts_rejects_bad_response_shape(monkeypatch) -> None:
def fake_post(url: str, **_kwargs: object) -> httpx.Response:
return httpx.Response(200, json={"data": [{}]}, request=httpx.Request("POST", url))
monkeypatch.setattr(httpx, "post", fake_post)
config = EbookSearchConfig(rerank=RerankConfig(enabled=False))
with pytest.raises(RuntimeError, match="Embedding request failed"):
embed_texts(["hello"], config)
+150
View File
@@ -0,0 +1,150 @@
"""Tests for EPUB search reranking."""
from __future__ import annotations
import httpx
import pytest
from python.ebook_search.config import EbookSearchConfig, RerankConfig, load_rerank_config
from python.ebook_search.rerank import rerank_chunks
from python.ebook_search.search import SearchResult, apply_rerank, skip_rerank
def candidates() -> list[SearchResult]:
return [
SearchResult(chunk_id=1, text="alpha", source_title="A", score=0.9),
SearchResult(chunk_id=2, text="beta", source_title="B", score=0.8),
SearchResult(chunk_id=3, text="gamma", source_title="C", score=0.7),
]
def rerank_response(payload: dict[str, object] | None = None, *, content: bytes | None = None) -> httpx.Response:
return httpx.Response(
200,
content=content,
json=payload,
request=httpx.Request("POST", "http://rerank.test/rerank"),
)
def test_config_defaults_keep_reranking_optional(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.delenv("EBOOK_SEARCH_RERANK_ENABLED", raising=False)
monkeypatch.delenv("EBOOK_SEARCH_RERANK_BASE_URL", raising=False)
monkeypatch.delenv("EBOOK_SEARCH_RERANK_MODEL", raising=False)
monkeypatch.delenv("EBOOK_SEARCH_RERANK_CANDIDATES", raising=False)
monkeypatch.delenv("EBOOK_SEARCH_RERANK_TIMEOUT_SECONDS", raising=False)
config = load_rerank_config()
assert config.enabled is False
assert config.base_url == "http://192.168.90.25:8001"
assert config.model == "qwen3-reranker-06b"
assert config.candidates == 24
assert config.timeout_seconds == 30
def test_reranking_disabled_returns_original_fused_order() -> None:
config = EbookSearchConfig(rerank=RerankConfig(enabled=False), top_k=2)
response = skip_rerank("query", candidates(), config)
assert response.rank_label == "Hybrid"
assert [result.chunk_id for result in response.results] == [1, 2]
def test_reranking_enabled_reorders_candidates(monkeypatch: pytest.MonkeyPatch) -> None:
def fake_post(_url: str, *, json: dict[str, object], timeout: float) -> httpx.Response:
assert timeout == 30
assert json == {
"model": "qwen3-reranker-06b",
"query": "query",
"documents": ["alpha", "beta", "gamma"],
}
return rerank_response(
{
"results": [
{"index": 0, "relevance_score": 0.1},
{"index": 1, "relevance_score": 0.9},
{"index": 2, "relevance_score": 0.4},
]
}
)
monkeypatch.setattr(httpx, "post", fake_post)
results = rerank_chunks("query", candidates(), RerankConfig())
assert [result.chunk_id for result in results] == [2, 1, 3]
assert [round(result.score, 3) for result in results] == [0.78, 0.37, 0.28]
assert [result.rerank_score for result in results] == [0.9, 0.1, 0.4]
def test_reranking_cannot_ignore_hybrid_score(monkeypatch: pytest.MonkeyPatch) -> None:
candidates = [
SearchResult(chunk_id=1, text="strong hybrid", source_title="A", score=1.0),
SearchResult(chunk_id=2, text="weak hybrid", source_title="B", score=0.1),
]
def fake_post(_url: str, **_kwargs: object) -> httpx.Response:
return rerank_response(
{
"results": [
{"index": 0, "relevance_score": 0.7},
{"index": 1, "relevance_score": 1.0},
]
}
)
monkeypatch.setattr(httpx, "post", fake_post)
results = rerank_chunks("query", candidates, RerankConfig())
assert [result.chunk_id for result in results] == [1, 2]
assert results[0].score == pytest.approx(0.79)
assert results[1].score == 0.7
assert results[1].rerank_score == 1.0
def test_vllm_rerank_timeout_raises(monkeypatch: pytest.MonkeyPatch) -> None:
def fake_rerank_chunks(
_query: str,
_candidates: list[SearchResult],
_config: RerankConfig,
) -> list[SearchResult]:
message = "timeout"
raise httpx.TimeoutException(message)
monkeypatch.setattr("python.ebook_search.search.rerank_chunks", fake_rerank_chunks)
config = EbookSearchConfig(rerank=RerankConfig(enabled=True), top_k=2)
with pytest.raises(httpx.TimeoutException, match="timeout"):
apply_rerank("query", candidates(), config)
def test_malformed_vllm_rerank_json_does_not_crash_search(monkeypatch: pytest.MonkeyPatch) -> None:
def fake_post(_url: str, **_kwargs: object) -> httpx.Response:
return rerank_response(content=b"not-json")
monkeypatch.setattr(httpx, "post", fake_post)
results = rerank_chunks("query", candidates()[:1], RerankConfig())
assert results[0].score == 0.3
def test_vllm_rerank_scores_are_clamped(monkeypatch: pytest.MonkeyPatch) -> None:
def fake_post(_url: str, **_kwargs: object) -> httpx.Response:
return rerank_response(
{
"results": [
{"index": 0, "relevance_score": -1},
{"index": 1, "relevance_score": 2},
]
}
)
monkeypatch.setattr(httpx, "post", fake_post)
results = rerank_chunks("query", candidates()[:2], RerankConfig())
assert {result.chunk_id: result.rerank_score for result in results} == {1: 0.0, 2: 1.0}
+303
View File
@@ -0,0 +1,303 @@
"""Tests for EPUB search HTMX routes."""
from __future__ import annotations
from compression import zstd
from fastapi.testclient import TestClient
from sqlalchemy import create_engine
from python.ebook_search.api.bm25_tasks import refresh_bm25_for_engine
from python.ebook_search.api.main import create_app
from python.ebook_search.config import EbookSearchConfig, RerankConfig
from python.ebook_search.embeddings import EmbeddingModelStats
from python.ebook_search.search import SearchResponse, SearchResult
from python.ebook_search.timing import RuntimeStep
def patch_app_runtime(monkeypatch):
"""Patch app startup dependencies used by UI route tests."""
monkeypatch.setattr("python.ebook_search.api.main.get_postgres_engine", fake_get_postgres_engine)
monkeypatch.setattr("python.ebook_search.api.main.ensure_bm25_corpus", lambda _session, _config: None)
def fake_get_postgres_engine(**_kwargs):
"""Return an in-memory engine for route tests."""
return create_engine("sqlite+pysqlite:///:memory:", future=True)
def test_search_page_uses_zstd_when_requested(monkeypatch) -> None:
patch_app_runtime(monkeypatch)
app = create_app()
app.state.config = EbookSearchConfig(rerank=RerankConfig(enabled=False))
with TestClient(app) as client:
response = client.get("/", headers={"accept-encoding": "zstd"})
assert response.status_code == 200
assert response.headers["content-encoding"] == "zstd"
assert b"EPUB Search" in zstd.decompress(response.content)
def test_ui_form_passes_rerank_flag_to_search_handler(monkeypatch) -> None:
captured: dict[str, object] = {}
def fake_search_ebooks(_engine, query, config, *, rerank=False):
captured["query"] = query
captured["rerank"] = rerank
captured["config"] = config
return SearchResponse(query=query, results=[], rank_label="Hybrid + rerank")
monkeypatch.setattr("python.ebook_search.api.routes.search.search_ebooks", fake_search_ebooks)
monkeypatch.setattr(
"python.ebook_search.api.routes.search.answer_query",
lambda _query, _results, _config: "answer",
)
patch_app_runtime(monkeypatch)
app = create_app()
app.state.config = EbookSearchConfig(rerank=RerankConfig(enabled=False), top_k=12, answer_enabled=True)
with TestClient(app) as client:
response = client.post("/search", data={"query": "where is the quote?", "rerank": "true"})
assert response.status_code == 200
assert "Hybrid + rerank" in response.text
assert captured["query"] == "where is the quote?"
assert captured["rerank"] is True
def test_ui_search_failure_returns_visible_error(monkeypatch) -> None:
def fake_search_ebooks(_engine, _query, _config, *, rerank=False):
del rerank
msg = "search exploded"
raise RuntimeError(msg)
monkeypatch.setattr("python.ebook_search.api.routes.search.search_ebooks", fake_search_ebooks)
patch_app_runtime(monkeypatch)
app = create_app()
app.state.config = EbookSearchConfig(rerank=RerankConfig(enabled=False), top_k=12)
with TestClient(app) as client:
response = client.post("/search", data={"query": "where is the quote?"})
assert response.status_code == 500
assert "search exploded" in response.text
def test_ui_answer_failure_still_returns_sources(monkeypatch) -> None:
def fake_search_ebooks(_engine, query, _config, *, rerank=False):
del rerank
return SearchResponse(query=query, results=[], rank_label="Hybrid")
def fake_answer_query(_query, _results, _config):
msg = "answer exploded"
raise RuntimeError(msg)
monkeypatch.setattr("python.ebook_search.api.routes.search.search_ebooks", fake_search_ebooks)
monkeypatch.setattr("python.ebook_search.api.routes.search.answer_query", fake_answer_query)
patch_app_runtime(monkeypatch)
app = create_app()
app.state.config = EbookSearchConfig(rerank=RerankConfig(enabled=False), top_k=12, answer_enabled=True)
with TestClient(app) as client:
response = client.post("/search", data={"query": "where is the quote?"})
assert response.status_code == 200
assert "Answer generation failed" in response.text
def test_ui_skips_answer_when_disabled(monkeypatch) -> None:
called = False
def fake_search_ebooks(_engine, query, _config, *, rerank=False):
del rerank
return SearchResponse(query=query, results=[], rank_label="Hybrid")
def fake_answer_query(_query, _results, _config):
nonlocal called
called = True
return "answer"
monkeypatch.setattr("python.ebook_search.api.routes.search.search_ebooks", fake_search_ebooks)
monkeypatch.setattr("python.ebook_search.api.routes.search.answer_query", fake_answer_query)
patch_app_runtime(monkeypatch)
app = create_app()
app.state.config = EbookSearchConfig(rerank=RerankConfig(enabled=False), answer_enabled=False)
with TestClient(app) as client:
response = client.post("/search", data={"query": "where is the quote?"})
assert response.status_code == 200
assert called is False
assert "Answer generation is disabled" in response.text
def test_ui_shows_component_scores(monkeypatch) -> None:
def fake_search_ebooks(_engine, query, _config, *, rerank=False):
del rerank
return SearchResponse(
query=query,
rank_label="Hybrid + rerank",
results=[
SearchResult(
chunk_id=1,
text="source text",
source_title="Book",
score=0.9,
rerank_score=0.9,
vector_score=0.8,
bm25_score=2.5,
fused_score=0.03,
)
],
)
monkeypatch.setattr("python.ebook_search.api.routes.search.search_ebooks", fake_search_ebooks)
monkeypatch.setattr(
"python.ebook_search.api.routes.search.answer_query",
lambda _query, _results, _config: "answer",
)
patch_app_runtime(monkeypatch)
app = create_app()
app.state.config = EbookSearchConfig(rerank=RerankConfig(enabled=False), answer_enabled=True)
with TestClient(app) as client:
response = client.post("/search", data={"query": "where is the quote?"})
assert response.status_code == 200
assert "rerank" in response.text
assert "vector cosine" in response.text
assert "BM25" in response.text
assert "RRF" in response.text
def test_ui_shows_search_runtime_chart(monkeypatch) -> None:
def fake_search_ebooks(_engine, query, _config, *, rerank=False):
del rerank
return SearchResponse(
query=query,
rank_label="Hybrid",
results=[],
timings=(
RuntimeStep(name="Embedding + vector search", duration_ms=12.5),
RuntimeStep(name="BM25 search", duration_ms=4.0),
),
)
monkeypatch.setattr("python.ebook_search.api.routes.search.search_ebooks", fake_search_ebooks)
monkeypatch.setattr(
"python.ebook_search.api.routes.search.answer_query",
lambda _query, _results, _config: "answer",
)
patch_app_runtime(monkeypatch)
app = create_app()
app.state.config = EbookSearchConfig(rerank=RerankConfig(enabled=False), answer_enabled=True)
with TestClient(app) as client:
response = client.post("/search", data={"query": "where is the quote?"})
assert response.status_code == 200
assert "Runtime" in response.text
assert "Total" in response.text
assert "Embedding + vector search" in response.text
assert "BM25 search" in response.text
assert "Answer generation" in response.text
assert "ms left" in response.text
def test_ui_embed_all_batches_until_complete(monkeypatch) -> None:
counts = iter([32, 32, 5, 0])
batch_sizes: list[int] = []
def fake_embed_missing_chunks(_session, config):
batch_sizes.append(config.embedding_batch_size)
return next(counts)
monkeypatch.setattr("python.ebook_search.api.routes.admin.embed_missing_chunks", fake_embed_missing_chunks)
patch_app_runtime(monkeypatch)
app = create_app()
with TestClient(app) as client:
response = client.post("/admin/embed-all")
assert response.status_code == 200
assert "Embedded 69 chunks in 3 batches of 32" in response.text
assert batch_sizes == [32, 32, 32, 32]
def test_ui_scan_schedules_bm25_refresh_after_database_change(monkeypatch) -> None:
scheduled = False
def fake_ingest_configured_paths(_session, _config):
return 1
def fake_schedule_bm25_refresh(_app):
nonlocal scheduled
scheduled = True
monkeypatch.setattr("python.ebook_search.api.routes.admin.ingest_configured_paths", fake_ingest_configured_paths)
monkeypatch.setattr("python.ebook_search.api.routes.admin.schedule_bm25_refresh", fake_schedule_bm25_refresh)
patch_app_runtime(monkeypatch)
app = create_app()
with TestClient(app) as client:
response = client.post("/admin/scan")
assert response.status_code == 200
assert "Indexed 1 EPUBs" in response.text
assert scheduled is True
def test_bm25_refresh_clears_loaded_corpus_cache(monkeypatch) -> None:
refreshed: list[object] = []
cache_cleared = False
def fake_refresh_bm25_corpus(session, config):
refreshed.append((session, config))
def fake_cache_clear():
nonlocal cache_cleared
cache_cleared = True
monkeypatch.setattr("python.ebook_search.api.bm25_tasks.refresh_bm25_corpus", fake_refresh_bm25_corpus)
monkeypatch.setattr("python.ebook_search.api.bm25_tasks.load_bm25_corpus.cache_clear", fake_cache_clear)
engine = create_engine("sqlite+pysqlite:///:memory:", future=True)
config = EbookSearchConfig(rerank=RerankConfig(enabled=False))
refresh_bm25_for_engine(engine, config)
assert len(refreshed) == 1
assert refreshed[0][1] == config
assert cache_cleared is True
def test_admin_page_shows_embedding_counts_by_model(monkeypatch) -> None:
def fake_embedding_model_stats(_session):
return [
EmbeddingModelStats(
model_name="qwen3-embedding-0.6b",
dimension=1024,
embedded_chunks=40,
total_chunks=64,
),
EmbeddingModelStats(
model_name="qwen3-embedding-4b",
dimension=2560,
embedded_chunks=8,
total_chunks=64,
),
]
monkeypatch.setattr("python.ebook_search.api.routes.admin.embedding_model_stats", fake_embedding_model_stats)
patch_app_runtime(monkeypatch)
app = create_app()
with TestClient(app) as client:
response = client.get("/admin")
assert response.status_code == 200
assert "qwen3-embedding-0.6b" in response.text
assert "1024" in response.text
assert "40" in response.text
assert "24" in response.text
assert "qwen3-embedding-4b" in response.text
assert "2560" in response.text
+4 -4
View File
@@ -19,7 +19,7 @@ if TYPE_CHECKING:
class MockFuture(Future): class MockFuture(Future):
"""MockFuture.""" """MockFuture."""
def __init__(self, result: Any) -> None: # noqa: ANN401 def __init__(self, result: Any) -> None:
"""Init.""" """Init."""
super().__init__() super().__init__()
self._result = result self._result = result
@@ -31,7 +31,7 @@ class MockFuture(Future):
logging.debug(f"{timeout}=") logging.debug(f"{timeout}=")
return self._exception return self._exception
def result(self, timeout: float | None = None) -> Any: # noqa: ANN401 def result(self, timeout: float | None = None) -> Any:
"""Result.""" """Result."""
logging.debug(f"{timeout}=") logging.debug(f"{timeout}=")
return self._result return self._result
@@ -40,11 +40,11 @@ class MockFuture(Future):
class MockPoolExecutor(ThreadPoolExecutor): class MockPoolExecutor(ThreadPoolExecutor):
"""MockPoolExecutor.""" """MockPoolExecutor."""
def __init__(self, *args: Any, **kwargs: Any) -> None: # noqa: ANN401 def __init__(self, *args: Any, **kwargs: Any) -> None:
"""Initializes a new ThreadPoolExecutor instance.""" """Initializes a new ThreadPoolExecutor instance."""
super().__init__(*args, **kwargs) super().__init__(*args, **kwargs)
def submit(self, fn: Callable[..., Any], /, *args: Any, **kwargs: Any) -> Future: # noqa: ANN401 def submit(self, fn: Callable[..., Any], /, *args: Any, **kwargs: Any) -> Future:
"""Submits a callable to be executed with the given arguments. """Submits a callable to be executed with the given arguments.
Args: Args:
+6 -3
View File
@@ -21,7 +21,7 @@ def test_validate_system(mocker: MockerFixture, fs: FakeFilesystem) -> None:
"""test_validate_system.""" """test_validate_system."""
fs.create_file( fs.create_file(
"/mock_snapshot_config.toml", "/mock_snapshot_config.toml",
contents='zpool = ["root_pool", "storage", "media"]\nservices = ["docker"]\n', contents='zpools = ["root_pool", "storage", "media"]\nservices = ["docker"]\n',
) )
mocker.patch(f"{VALIDATE_SYSTEM}.systemd_tests", return_value=None) mocker.patch(f"{VALIDATE_SYSTEM}.systemd_tests", return_value=None)
@@ -33,9 +33,10 @@ def test_validate_system_errors(mocker: MockerFixture, fs: FakeFilesystem) -> No
"""test_validate_system_errors.""" """test_validate_system_errors."""
fs.create_file( fs.create_file(
"/mock_snapshot_config.toml", "/mock_snapshot_config.toml",
contents='zpool = ["root_pool", "storage", "media"]\nservices = ["docker"]\n', contents='zpools = ["root_pool", "storage", "media"]\nservices = ["docker"]\n',
) )
mocker.patch(f"{VALIDATE_SYSTEM}.signal_alert")
mocker.patch(f"{VALIDATE_SYSTEM}.systemd_tests", return_value=["systemd_tests error"]) mocker.patch(f"{VALIDATE_SYSTEM}.systemd_tests", return_value=["systemd_tests error"])
mocker.patch(f"{VALIDATE_SYSTEM}.zpool_tests", return_value=["zpool_tests error"]) mocker.patch(f"{VALIDATE_SYSTEM}.zpool_tests", return_value=["zpool_tests error"])
@@ -49,9 +50,11 @@ def test_validate_system_execution(mocker: MockerFixture, fs: FakeFilesystem) ->
"""test_validate_system_execution.""" """test_validate_system_execution."""
fs.create_file( fs.create_file(
"/mock_snapshot_config.toml", "/mock_snapshot_config.toml",
contents='zpool = ["root_pool", "storage", "media"]\nservices = ["docker"]\n', contents='zpools = ["root_pool", "storage", "media"]\nservices = ["docker"]\n',
) )
mocker.patch(f"{VALIDATE_SYSTEM}.signal_alert")
mocker.patch(f"{VALIDATE_SYSTEM}.systemd_tests", return_value=None)
mocker.patch(f"{VALIDATE_SYSTEM}.zpool_tests", side_effect=RuntimeError("zpool_tests error")) mocker.patch(f"{VALIDATE_SYSTEM}.zpool_tests", side_effect=RuntimeError("zpool_tests error"))
with pytest.raises(SystemExit) as exception_info: with pytest.raises(SystemExit) as exception_info: