117 lines
4.2 KiB
Python
117 lines
4.2 KiB
Python
"""Admin routes for the EPUB search web UI."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from dataclasses import replace
|
|
from typing import TYPE_CHECKING
|
|
|
|
from fastapi import APIRouter, Request
|
|
from fastapi.responses import HTMLResponse
|
|
from sqlalchemy.orm import Session
|
|
|
|
from python.ebook_search.api.bm25_tasks import schedule_bm25_refresh
|
|
from python.ebook_search.api.web import templates
|
|
from python.ebook_search.embeddings import embed_missing_chunks, embedding_model_stats
|
|
from python.ebook_search.ingest import ingest_configured_paths
|
|
|
|
if TYPE_CHECKING:
|
|
from fastapi import FastAPI
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
router = APIRouter(prefix="/admin")
|
|
EMBED_ALL_BATCH_SIZE = 32
|
|
|
|
|
|
def register_admin_routes(app: FastAPI) -> None:
|
|
"""Register admin routes on the app."""
|
|
app.include_router(router)
|
|
|
|
|
|
@router.get("", response_class=HTMLResponse)
|
|
def admin(request: Request) -> HTMLResponse:
|
|
"""Render the admin page."""
|
|
with Session(request.app.state.engine) as session:
|
|
stats = embedding_model_stats(session)
|
|
logger.info("ebook_admin_page_loaded models=%s", len(stats))
|
|
return templates.TemplateResponse(request, "admin.html", {"config": request.app.state.config, "stats": stats})
|
|
|
|
|
|
@router.post("/scan", response_class=HTMLResponse)
|
|
def scan_library(request: Request) -> HTMLResponse:
|
|
"""Scan configured library paths for EPUB changes."""
|
|
try:
|
|
with Session(request.app.state.engine) as session:
|
|
count = ingest_configured_paths(session, request.app.state.config)
|
|
session.commit()
|
|
except Exception as error:
|
|
logger.exception("ebook_admin_scan_failed")
|
|
return templates.TemplateResponse(request, "partials/error.html", {"message": str(error)}, status_code=500)
|
|
|
|
logger.info("ebook_admin_scan_complete changed_files=%s", count)
|
|
if count > 0:
|
|
schedule_bm25_refresh(request.app)
|
|
return templates.TemplateResponse(request, "partials/admin_status.html", {"message": f"Indexed {count} EPUBs"})
|
|
|
|
|
|
@router.post("/embed-missing", response_class=HTMLResponse)
|
|
def embed_missing(request: Request) -> HTMLResponse:
|
|
"""Embed chunks missing vectors for the configured model."""
|
|
try:
|
|
with Session(request.app.state.engine) as session:
|
|
count = embed_missing_chunks(session, request.app.state.config)
|
|
session.commit()
|
|
except Exception as error:
|
|
logger.exception("ebook_admin_embed_missing_failed")
|
|
return templates.TemplateResponse(request, "partials/error.html", {"message": str(error)}, status_code=500)
|
|
|
|
logger.info("ebook_admin_embed_missing_complete chunks=%s", count)
|
|
return templates.TemplateResponse(
|
|
request,
|
|
"partials/admin_status.html",
|
|
{"message": f"Embedded {count} chunks"},
|
|
)
|
|
|
|
|
|
@router.post("/embed-all", response_class=HTMLResponse)
|
|
def embed_all(request: Request) -> HTMLResponse:
|
|
"""Embed all chunks missing vectors in fixed-size batches."""
|
|
total = 0
|
|
batches = 0
|
|
config = replace(request.app.state.config, embedding_batch_size=EMBED_ALL_BATCH_SIZE)
|
|
try:
|
|
with Session(request.app.state.engine) as session:
|
|
while True:
|
|
count = embed_missing_chunks(session, config)
|
|
if count == 0:
|
|
break
|
|
session.commit()
|
|
total += count
|
|
batches += 1
|
|
logger.info(
|
|
"ebook_admin_embed_all_batch_complete batch=%s chunks=%s total_chunks=%s",
|
|
batches,
|
|
count,
|
|
total,
|
|
)
|
|
except Exception as error:
|
|
logger.exception(
|
|
"ebook_admin_embed_all_failed batches=%s chunks=%s",
|
|
batches,
|
|
total,
|
|
)
|
|
return templates.TemplateResponse(
|
|
request,
|
|
"partials/error.html",
|
|
{"message": f"Embed all failed after {total} chunks in {batches} batches: {error}"},
|
|
status_code=500,
|
|
)
|
|
|
|
logger.info("ebook_admin_embed_all_complete batches=%s chunks=%s", batches, total)
|
|
return templates.TemplateResponse(
|
|
request,
|
|
"partials/admin_status.html",
|
|
{"message": f"Embedded {total} chunks in {batches} batches of {EMBED_ALL_BATCH_SIZE}"},
|
|
)
|