diff --git a/CLAUDE.md b/CLAUDE.md index 21904c9..253442e 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -20,7 +20,7 @@ lembas info # paths + counts, useful when confused lembas secret-key # generate LEMBAS_SECRET_KEY lembas create-admin # create or promote an admin -pytest # 1979 tests, ~2min +pytest # 2024 tests, ~2min # PLAN.md tracks what is and is not built ruff check . # lint (line length 100) python scripts/build_artwork.py # regenerate artwork (SVG + PWA icons; @@ -88,6 +88,7 @@ src/lembas/ admin_images.py the ComfyUI, and the workflow templates on it admin_prompts.py the prompt fragment editor and its preview admin_branding.py the name, the logo, the wording and the themes + admin_extraction.py what a file may cost, and what finds it afterwards branding.py /branding.css and the assets behind it, both unauthenticated admin_suggestions.py the cards offered on the new-chat screen admin_tools.py custom HTTP tools and MCP servers @@ -159,6 +160,10 @@ src/lembas/ and everything it may not do branding.py whose instance this is: the name, the artwork, the wording and the themes, cached once per process + llm/embeddings.py /v1/embeddings, batched, normalised on the way out + library/chunks.py splitting a record, and packing a vector + library/indexing.py keeping the semantic index current, and rebuilding it + library/retrieval.py keywords and meaning, fused settings_store.py runtime instance settings canvas.py what is open in the canvas panel, and where it comes from scratch.py a chat's own working document @@ -194,6 +199,10 @@ touching the code it names -- these are the same notes, not a summary. schedule something wrote a note and said it had). - `docs/notes/image-generation.md` -- the ComfyUI workflow with holes in it, what substitution walks, the review-and-retry loop, and how a failure reports itself. +- `docs/notes/search-and-extraction.md` -- extraction limits as a snapshot, why + reciprocal rank fusion and not a weight, how a record scores as its best chunk, + why vectors from two models never meet, and the session event that notices a + library record changing. - `docs/notes/branding.md` -- the branding snapshot and why it is a Jinja global, where the instance name went and how an upgrade keeps it, how a custom theme inherits through `data-base`, and why `/branding.css` is a route. @@ -1449,6 +1458,14 @@ with nothing streaming is exactly what a hang looks like. chunks. `tools.ToolCallAccumulator` rejoins them keyed on `index` — not on name, which breaks the moment a model calls one tool twice in a turn. +**Library search is keyword *and* meaning, and neither is a mode.** `fts.search_ids` +was already the one seam; `library/retrieval.search` sits in front of it and fuses +its ranking with a vector one by reciprocal rank fusion -- ranks, not scores, +because bm25 and cosine are not comparable and normalising them means picking a +constant nobody can tune. With no embedding model configured it returns exactly +what FTS returned, in that order, and no chunk row is ever written. See +`docs/notes/search-and-extraction.md`. + **Four stores, four different reasons.** `services/library/` — `documents` (uploaded by a person, searched by the model), `notes` (written by the model, searched), `memories` (short, and *injected whole* every turn), `skills` (index diff --git a/PLAN.md b/PLAN.md index 958a629..c125339 100644 --- a/PLAN.md +++ b/PLAN.md @@ -9,7 +9,7 @@ reasoning, tool calling with web search, custom HTTP tools and MCP servers, agent chats that work on a machine over SSH, a knowledge library, notes, memory and skills, speech in and out, image generation over ComfyUI, users and groups, model administration, installable as an app, reports, messages, and scheduled -work that runs on its own. 1979 tests, `ruff` clean. +work that runs on its own. 2024 tests, `ruff` clean. What remains before the first stable release is written out below, in phases, under [The road to 1.0.0](#the-road-to-100). @@ -496,12 +496,29 @@ seen working. a content hash in the link so a save is not left to the browser's cache ### Phase 5 — extraction, embeddings and hybrid search (`0.9.5`) -- [ ] **Extraction has settings** — upload size, image edge, PDF pages, - extracted characters, orphan age, which extensions count as text -- [ ] **A dedicated embedding model**, chosen from the models flagged for it -- [ ] **Search becomes hybrid** — FTS5 and vector recall fused, behind the one - call the retrieval service already is. No model chosen means exactly the - keyword search there is today +- [x] **Extraction has settings** — upload size, image edge, JPEG quality, PDF + pages, extracted characters, orphan age, extra text extensions. Read + through a process-level snapshot, because `prepare` is called from places + with no session. The decompression-bomb guard stays a constant: it is a + guard, not a preference +- [x] **A dedicated embedding model**, picked from the models flagged for it — + and a model that lost its flag is *named* rather than silently dropped + from the picker +- [x] **Search becomes hybrid** — FTS5 and vector recall fused by reciprocal + rank fusion, behind the one call the stores already searched through. + Ranks rather than scores, because bm25 and cosine are not comparable and + normalising them means picking a constant nobody can tune +- [x] **No model chosen means exactly the keyword search there is today** — no + rows, no requests, the same ids in the same order, asserted rather than + claimed +- [x] Indexing is fired and forgotten and noticed by a session event, so no + writer has to remember it — forgetting would be silent, since only + semantic recall would go stale +- [x] Vectors from two models never meet: width and model are stored beside + every vector and a mismatch is skipped, because scoring across two spaces + is a confident wrong answer rather than a missing one +- [x] A rebuild that commits as it goes, reports itself, and stops polling when + it finishes ### Phase 6 — permissions, quotas and sharing (`0.9.6`) - [ ] **"What can this user actually do?"** answered on screen, from the @@ -588,9 +605,12 @@ microphone is unavailable for the same reason. administrator's assertion, not something endpoints reliably advertise. Set it on a model that cannot, and its replies fail rather than degrade. -**Library search is keyword, not semantic.** FTS5 ranks well and needs no -dependency or embedding endpoint, but "how do I get paid" will not find a -document that says "invoicing". +**Library search is keyword-only until an embedding model is chosen.** FTS5 ranks +well and needs no dependency, but "how do I get paid" will not find a document +that says "invoicing". Choosing a model on **Extraction** adds a vector ranking +fused with that one; choosing none is byte-for-byte the search that was always +there. What that costs is an index that has to be rebuilt when the model changes, +and stale vectors that are ignored until it is. **A model can write its own skills, and they take effect at once.** Marked as model-authored and fully revertible, but a model that has just read a hostile diff --git a/docs/notes/search-and-extraction.md b/docs/notes/search-and-extraction.md new file mode 100644 index 0000000..cbcc629 --- /dev/null +++ b/docs/notes/search-and-extraction.md @@ -0,0 +1,142 @@ +# Extraction, embeddings and hybrid search + +Read this before touching `services/files.py:limits`, `services/library/`'s new +three modules, or the `Chunk` table. + +## Extraction is a snapshot, not a session + +The constants in `services/files.py` are **defaults** now; what `prepare` reads +is `limits()`, a process-level snapshot with the same shape and the same +reasoning as `services/branding.py`. Threading a session through `prepare`, +`_process_image`, `_process_pdf` and `_process_text` would have meant six +signatures changed to carry a number, and several of their callers — the startup +sweep, a tool runner — have no session in hand. + +`files.forget()` is called by `api/admin_extraction.py` and by nothing else. The +tests drop it between cases in `conftest.py` beside the branding one, for the +same reason. + +Two things stayed constants on purpose: + +- **`Image.MAX_IMAGE_PIXELS`** — a decompression-bomb guard, not a preference. A + 60,000×60,000 PNG is a few KB on disk and hundreds of gigabytes decoded, and + nothing good comes of being able to raise that from a form. +- **`ORPHAN_AGE` in a signature.** `sweep_orphans(older_than=None)` resolves the + default inside the body, because a default argument is evaluated at import and + a module constant there would pin the shipped 24 hours whatever anybody set. + +## Nothing changes for an instance that configures nothing + +`embedding_model_id` empty means: no chunk rows written, no requests made, +`retrieval.search` returning exactly what `fts.search_ids` returns, in exactly +that order. That is asserted rather than claimed +(`test_with_no_model_search_is_exactly_the_keyword_search`), and it is what makes +this safe to land on an existing instance. + +## Reciprocal rank fusion, and why not a weight + +bm25 is a negative number whose scale depends on the corpus; cosine is 0..1. They +are not comparable, and normalising them onto a common scale means picking a +constant nobody can tune without a labelled test set they do not have. + +RRF uses the **ranks**: `1 / (K + rank)`, summed. One constant, famously +insensitive to it, and it degrades to exactly one list when the other is empty — +which is what makes "no embedding model" a *branch that does not exist* rather +than a special case. `RRF_K` is deliberately not a setting: a number nobody can +evaluate is a number nobody should be asked about. + +The fused `rank` is **larger for better**, the opposite of bm25's convention. +Nothing downstream reads it, but it is worth knowing. + +## The query is embedded by the caller + +`search()` is synchronous because every store's `search()` is, and every one of +those is called from both a route and a tool runner. Embedding is an HTTP +request. So the caller embeds first and passes a vector in; one that cannot +passes nothing and gets keywords. + +`retrieval.worker_for(db)` and `retrieval.embed_with(worker, needle)` are split +for a specific reason: a **tool runner must not hold a database session across +an HTTP request**, so it resolves, closes, and awaits. A route that already holds +the request's session uses `embed_query(db, needle)`, which is the two together. + +## A record scores as its best chunk + +Not its average. One paragraph that answers the question is what makes a document +worth returning; averaging ranks a long document about something else above a +short one that says exactly the thing, because most of the long one is not about +anything. + +`CHUNK_MULTIPLIER` is why the semantic side asks for more rows than are wanted: +one long document can own several of the best chunks and would otherwise crowd +everything else out. + +## Vectors from two models never meet + +`Chunk` stores `dims` and `model_id` beside every vector, and +`retrieval.semantic_ids` **skips a chunk whose width is not the query's**. +Changing the embedding model changes the space, and vectors from two spaces score +against each other perfectly happily and mean nothing — a search that works and +is wrong, which is the worst failure this feature can have. Nothing is deleted on +a model change; the stale rows are ignored until a rebuild replaces them, and the +save says so. + +`unpack` checks the BLOB's length against the declared width for the same reason: +inferring the width would let a truncated row unpack into a shorter vector and +score happily. + +## Indexing is fired and forgotten, and noticed by an event + +Every library writer is synchronous and has just committed a row. None should +wait on a model server before saying "saved". So `schedule(kind, id)` starts a +task and returns; a save that cannot be indexed is still a save, and that record +falls back to keywords until the next rebuild. + +**How a change is noticed is a SQLAlchemy session event, not a call in each of +the ten writers.** That is a departure from this codebase's taste for explicit +seams, and the reason is the one `tool_label` gives for being a Jinja global: a +step every writer has to remember is a step one of them will forget, and here +forgetting is silent — the record saves, keyword search still finds it, and only +its semantic recall is quietly stale. + +`after_flush` collects and `after_commit` fires, in that order and never merged: +inside a flush the transaction has not landed, so a task started there could read +a row that does not exist yet — and `session.deleted` is empty by the time the +commit fires, so the collecting has to happen while it is not. `install()` is +idempotent because the app factory runs once per test. + +A **deletion is scheduled like a change**: `index_resource` finds no row and drops +the chunks. One path rather than two, and the one that runs is the one that has +to be right anyway. `sweep_orphans` is the backstop for a delete with no event +loop to schedule anything — a CLI command, or a cascade from removing an account +— and runs at startup and at the end of every rebuild. + +## Writing is all-or-nothing + +`index_resource` embeds everything **before** it deletes anything. Deleting first +and failing half way through would leave a record indexed by half of itself, +which ranks worse than not being indexed at all and looks like nothing. + +Staleness is a hash (`source_hash`) rather than a timestamp, so re-indexing an +unchanged record is free and "is this current?" is answerable without embedding +anything. + +## The rebuild + +One record at a time, never gathered: the far side is usually one local model +server, and twenty concurrent embedding requests against it is slower than twenty +sequential ones as well as being ruder. Each record commits, so a half-finished +index is usable. + +`Progress` is in-process, because a rebuild does not survive a restart — +persisting it would mean a progress bar that stops moving and never finishes. +`admin/_index_progress.html` emits its `hx-trigger` **only while running**, so the +last frame has nothing attached and the polling stops by itself. + +## The response order is trusted only as far as `index` + +`_vectors_in` sorts on the declared `index` rather than on arrival order, and +refuses a response with a different number of vectors than inputs. Nothing in the +specification promises the order, and a provider that sorts differently would +pair every chunk with somebody else's vector — silently, for the life of the +index. diff --git a/src/lembas/__init__.py b/src/lembas/__init__.py index f7ede1d..c7d74f3 100644 --- a/src/lembas/__init__.py +++ b/src/lembas/__init__.py @@ -1,3 +1,3 @@ """LLeMbas - a Middle-earth themed web UI for OpenAI-compatible LLM endpoints.""" -__version__ = "0.9.4" +__version__ = "0.9.5" diff --git a/src/lembas/api/admin_extraction.py b/src/lembas/api/admin_extraction.py new file mode 100644 index 0000000..01ecb49 --- /dev/null +++ b/src/lembas/api/admin_extraction.py @@ -0,0 +1,200 @@ +"""What happens to a file between the upload and the model, and how it is found. + +Two halves on one page because they are two ends of the same pipeline: what gets +extracted decides what there is to search, and the search settings decide what +becomes of it. Splitting them would mean an administrator setting a 300-page PDF +limit on one screen and wondering on another why half a book is missing from the +index. + +Every save drops `files.forget()`, and this is the only module that calls it — +the same discipline `admin_branding` has with the branding snapshot, and for the +same reason: a process-level cache whose save does not drop it is a setting that +takes effect at the next restart. +""" + +from __future__ import annotations + +import logging + +from fastapi import APIRouter, Form, Request, Response, status +from fastapi.responses import RedirectResponse +from sqlalchemy import select + +from lembas.api.deps import AdminUser, Db +from lembas.db.models import Connection, Model +from lembas.services import files as files_service +from lembas.services import settings_store +from lembas.services.library import indexing +from lembas.web.templating import render + +log = logging.getLogger(__name__) + +router = APIRouter(prefix="/admin/extraction", tags=["admin-extraction"]) + + +def _embedding_models(db: Db) -> list[Model]: + """Models an administrator has marked as producing embeddings. + + Filtered rather than listed in full, the same shape `/admin/images` uses for + its reviewer: a chat model in this picker is a setting that looks configured + and fails on the first request, which is the shape of failure this codebase + keeps cataloguing. + """ + return [ + model + for model in db.scalars( + select(Model).join(Connection).order_by(Model.position, Model.model_id) + ) + if (model.capabilities_json or {}).get("embeddings") + ] + + +def _lines(text: str) -> list[str]: + return [line.strip() for line in (text or "").splitlines() if line.strip()] + + +@router.get("") +async def extraction_page(request: Request, db: Db, user: AdminUser, saved: str = ""): + values = settings_store.extraction(db) + models = _embedding_models(db) + return render( + request, + "admin/extraction.html", + { + "values": values, + "extensions_text": "\n".join(values.get("extra_text_extensions") or []), + "models": models, + # A model that was chosen and has since lost its flag, or its + # connection. Named rather than silently dropped from the picker: + # a setting that vanishes is one nobody can tell from one that was + # never made. + "missing_model": ( + values["embedding_model_id"] + if values["embedding_model_id"] + and values["embedding_model_id"] not in {m.model_id for m in models} + else "" + ), + "ready": indexing.enabled(db), + "counts": indexing.counts(db), + "progress": indexing.progress(), + "saved": saved, + }, + ) + + +@router.post("") +async def save_extraction( + db: Db, + user: AdminUser, + max_upload_mb: int = Form(20), + max_image_edge: int = Form(1400), + jpeg_quality: int = Form(85), + max_pdf_pages: int = Form(300), + max_extracted_chars: int = Form(120_000), + orphan_hours: int = Form(24), + extra_text_extensions: str = Form(""), + reject_unreadable_pdf: bool = Form(False), +) -> Response: + settings_store.update( + db, + { + # Clamped here as well as on read, for the reason the agent settings + # give: a number with no bound is a way to break the instance from + # a form. + "max_upload_mb": min(max(max_upload_mb, 1), 512), + "max_image_edge": min(max(max_image_edge, 128), 8192), + "jpeg_quality": min(max(jpeg_quality, 30), 100), + "max_pdf_pages": min(max(max_pdf_pages, 1), 5000), + "max_extracted_chars": min(max(max_extracted_chars, 1000), 5_000_000), + "orphan_hours": min(max(orphan_hours, 1), 8760), + "extra_text_extensions": _lines(extra_text_extensions), + "reject_unreadable_pdf": reject_unreadable_pdf, + }, + key=settings_store.EXTRACTION, + ) + files_service.forget() + log.info("extraction settings changed by %s", user.email) + return RedirectResponse( + "/admin/extraction?saved=Extraction+saved.", status_code=status.HTTP_303_SEE_OTHER + ) + + +@router.post("/search") +async def save_search( + db: Db, + user: AdminUser, + embedding_model_id: str = Form(""), + chunk_chars: int = Form(1200), + chunk_overlap: int = Form(150), + embed_batch: int = Form(16), +) -> Response: + """The semantic half. + + Its own form and its own route, because the two halves have different + consequences: changing a chunk size invalidates every vector already stored, + and changing an upload limit does not. Keeping them apart is what lets the + page say so beside the control that does it. + """ + before = settings_store.extraction(db) + settings_store.update( + db, + { + "embedding_model_id": embedding_model_id.strip()[:300], + "chunk_chars": min(max(chunk_chars, 200), 8000), + "chunk_overlap": max(chunk_overlap, 0), + "embed_batch": min(max(embed_batch, 1), 256), + }, + key=settings_store.EXTRACTION, + ) + files_service.forget() + + # Changing the model changes the vector space, so what is stored stops + # meaning anything against a new query. Nothing is deleted -- the scorer + # already skips a width that does not match the query's, so a stale index is + # ignored rather than trusted -- but a rebuild is what makes it useful + # again, and offering it here is cheaper than leaving somebody to notice. + changed = before["embedding_model_id"] != embedding_model_id.strip() + message = "Search+saved." + if changed and embedding_model_id.strip(): + message = "Search+saved.+Rebuild+the+index+to+use+the+new+model." + log.info("embedding model set to %r by %s", embedding_model_id, user.email) + return RedirectResponse( + f"/admin/extraction?saved={message}", status_code=status.HTTP_303_SEE_OTHER + ) + + +@router.post("/rebuild") +async def rebuild(request: Request, db: Db, user: AdminUser) -> Response: + """Start a rebuild, and answer with the progress card. + + A background task rather than a request that waits: embedding a library of a + few thousand records is minutes of HTTP round trips, and a page that hangs + for that long is one somebody reloads, which starts a second one. + """ + started = indexing.start_rebuild() + if started: + log.info("index rebuild started by %s", user.email) + return render( + request, + "admin/_index_progress.html", + {"progress": indexing.progress(), "counts": indexing.counts(db), "ready": True}, + ) + + +@router.get("/progress") +async def rebuild_progress(request: Request, db: Db, user: AdminUser) -> Response: + """Polled while a rebuild runs. Stops polling itself when it finishes. + + Polled rather than streamed for the reason `/api/chats/unread` is: this is + one small fragment on one page, and an SSE stream for it would be a second + streaming path to keep correct. + """ + return render( + request, + "admin/_index_progress.html", + { + "progress": indexing.progress(), + "counts": indexing.counts(db), + "ready": indexing.enabled(db), + }, + ) diff --git a/src/lembas/api/admin_models.py b/src/lembas/api/admin_models.py index d5efd21..1a6af95 100644 --- a/src/lembas/api/admin_models.py +++ b/src/lembas/api/admin_models.py @@ -23,7 +23,10 @@ router = APIRouter(tags=["admin-models"]) # What the endpoint can do. Endpoints do not advertise any of this reliably, so # these are an administrator's assertion. -PROTOCOL_CAPABILITIES = ("reasoning", "vision", "tools") +# `embeddings` is the odd one out and is worth naming as such: the other three +# say what a model can do in a *chat*, and this one says it is not for chatting +# at all. It is what /admin/extraction picks from, and nothing else reads it. +PROTOCOL_CAPABILITIES = ("reasoning", "vision", "tools", "embeddings") # Which tools this model is given. Distinct from the above: `tools` is whether a # tools array may be sent at all, these are what goes in it. Every one of them is @@ -46,6 +49,9 @@ TOOL_CAPABILITIES = ( ("tool_ask", "Ask the reader"), ("tool_report", "Reports"), ("tool_image", "Image generation"), + ("tool_scratch", "Canvas"), + ("tool_schedule", "Scheduling"), + ("tool_subagent", "Helpers"), ("tool_agent", "Agent execution"), ) diff --git a/src/lembas/api/library.py b/src/lembas/api/library.py index 4a85d7f..9d26d09 100644 --- a/src/lembas/api/library.py +++ b/src/lembas/api/library.py @@ -40,6 +40,7 @@ from lembas.services.fetch import FetchError, fetch from lembas.services.library import documents as documents_service from lembas.services.library import memories as memories_service from lembas.services.library import notes as notes_service +from lembas.services.library import retrieval from lembas.services.library import skills as skills_service from lembas.services.markdown import render_markdown from lembas.web.templating import render @@ -175,7 +176,12 @@ async def base_detail( raise HTTPException(status.HTTP_404_NOT_FOUND, "That knowledge base is not available.") if q.strip(): - rows = documents_service.search(db, user, q, limit=PAGE_SIZE, base_ids=[base.id]) + # The reader's search box gets the same recall a model's does. `None` + # when nothing is configured, which is the keyword search unchanged. + vector = await retrieval.embed_query(db, q) + rows = documents_service.search( + db, user, q, limit=PAGE_SIZE, base_ids=[base.id], vector=vector + ) pager = {"page": 1, "pages": 1, "total": len(rows)} else: rows, pager = _page( @@ -241,7 +247,7 @@ async def upload_document( if base is not None and not sharing.can_write(base, user): raise HTTPException(status.HTTP_403_FORBIDDEN, "That base is not yours to add to.") - payload = await file.read(files_service.MAX_UPLOAD_BYTES + 1) + payload = await file.read(files_service.limits().max_upload_bytes + 1) try: document = documents_service.store_upload( db, @@ -341,7 +347,9 @@ async def document_content(db: Db, user: RequiredUser, document_id: str) -> Resp @router.get("/library/notes") async def notes_list(request: Request, db: Db, user: RequiredUser, q: str = "", page: int = 1): if q.strip(): - rows = notes_service.search(db, user, q, limit=PAGE_SIZE) + rows = notes_service.search( + db, user, q, limit=PAGE_SIZE, vector=await retrieval.embed_query(db, q) + ) pager = {"page": 1, "pages": 1, "total": len(rows)} else: rows, pager = _page( @@ -422,7 +430,9 @@ async def delete_note(db: Db, user: RequiredUser, note_id: str) -> Response: @router.get("/library/skills") async def skills_list(request: Request, db: Db, user: RequiredUser, q: str = "", page: int = 1): if q.strip(): - rows = skills_service.search(db, user, q, limit=PAGE_SIZE) + rows = skills_service.search( + db, user, q, limit=PAGE_SIZE, vector=await retrieval.embed_query(db, q) + ) pager = {"page": 1, "pages": 1, "total": len(rows)} else: rows, pager = _page(db, skills_service.visible(db, user).order_by(Skill.name), page) diff --git a/src/lembas/api/reports.py b/src/lembas/api/reports.py index 1136dde..004f4eb 100644 --- a/src/lembas/api/reports.py +++ b/src/lembas/api/reports.py @@ -24,6 +24,7 @@ from lembas.api.library import PAGE_SIZE, _page from lembas.api.pages import sidebar_context from lembas.db.models import Report from lembas.services import reports as reports_service +from lembas.services.library import retrieval from lembas.services.markdown import render_markdown from lembas.web.templating import render @@ -35,7 +36,9 @@ router = APIRouter(dependencies=[Depends(require_permission("reports.use"))], ta @router.get("/reports") async def reports_list(request: Request, db: Db, user: RequiredUser, q: str = "", page: int = 1): if q.strip(): - rows = reports_service.search(db, user, q, limit=PAGE_SIZE) + rows = reports_service.search( + db, user, q, limit=PAGE_SIZE, vector=await retrieval.embed_query(db, q) + ) pager = {"page": 1, "pages": 1, "total": len(rows)} else: rows, pager = _page( diff --git a/src/lembas/db/models/__init__.py b/src/lembas/db/models/__init__.py index 04f2d7d..c985a13 100644 --- a/src/lembas/db/models/__init__.py +++ b/src/lembas/db/models/__init__.py @@ -39,6 +39,11 @@ from lembas.db.models.image import ImageWorkflow from lembas.db.models.library import ( AUTHOR_MODEL, AUTHOR_USER, + CHUNK_DOCUMENT, + CHUNK_KINDS, + CHUNK_NOTE, + CHUNK_REPORT, + CHUNK_SKILL, PRINCIPAL_GROUP, PRINCIPAL_USER, RESOURCE_BASE, @@ -46,6 +51,7 @@ from lembas.db.models.library import ( RESOURCE_SKILL, SOURCE_LINK, SOURCE_UPLOAD, + Chunk, Document, KnowledgeBase, Memory, @@ -155,6 +161,12 @@ __all__ = [ "Job", "Connection", "CustomTool", + "CHUNK_DOCUMENT", + "CHUNK_KINDS", + "CHUNK_NOTE", + "CHUNK_REPORT", + "CHUNK_SKILL", + "Chunk", "Document", "Folder", "Group", diff --git a/src/lembas/db/models/library.py b/src/lembas/db/models/library.py index b347e41..8765b9d 100644 --- a/src/lembas/db/models/library.py +++ b/src/lembas/db/models/library.py @@ -29,6 +29,7 @@ from sqlalchemy import ( ForeignKey, Index, Integer, + LargeBinary, String, Table, Text, @@ -286,3 +287,69 @@ class Share(UUIDPrimaryKey, Timestamps, Base): Index("ix_shares_resource", Share.resource_type, Share.resource_id) Index("ix_shares_principal", Share.principal_type, Share.principal_id) + + +# --- Semantic index ----------------------------------------------------------- +# What a chunk belongs to. Strings rather than a foreign key per store, because +# one table serving four of them is what stops the chunking, the scoring and the +# rebuild being written four times and drifting three ways. +CHUNK_DOCUMENT = "document" +CHUNK_NOTE = "note" +CHUNK_SKILL = "skill" +CHUNK_REPORT = "report" + +CHUNK_KINDS = (CHUNK_DOCUMENT, CHUNK_NOTE, CHUNK_SKILL, CHUNK_REPORT) + + +class Chunk(UUIDPrimaryKey, Timestamps, Base): + """A piece of one library record, and its embedding. + + **Additive, so `sync_schema` creates it at startup with no manual step**, and + absent-means-nothing: an instance with no embedding model chosen never writes + a row here and the search behaves exactly as it always did. + + `owner_id` is denormalised off the resource. It is not used for + authorisation -- `services/sharing.py` is still the only definition of who + may see what, and scoring happens before that filter exactly as the + full-text path does -- but it is what makes "rebuild this person's index" + and "drop everything of theirs" one indexed query rather than four joins. + + No foreign key on `resource_id`, for the reason `Share.principal_id` has + none: the column points at one of four tables depending on `resource_type`, + which SQLite cannot express. `indexing.forget_resource` deletes the rows. + """ + + __tablename__ = "chunks" + + owner_id: Mapped[str] = mapped_column( + String(32), ForeignKey("users.id", ondelete="CASCADE"), nullable=False, index=True + ) + resource_type: Mapped[str] = mapped_column(String(16), nullable=False) + resource_id: Mapped[str] = mapped_column(String(32), nullable=False) + # Where in the record this piece came from, so a set can be rebuilt in order + # and a hit can say which part matched. + ordinal: Mapped[int] = mapped_column(Integer, default=0, nullable=False) + text: Mapped[str] = mapped_column(Text, default="") + + # float32, little-endian, packed. A BLOB rather than JSON because a 1024 + # dimension vector is 4KB packed and about 20KB as text, and every one of + # them is read on every semantic search. + vector: Mapped[bytes] = mapped_column(LargeBinary, nullable=False) + # How many floats are in it. Stored rather than derived from the length so a + # mismatch is a comparison this code refuses rather than one it gets wrong: + # changing the embedding model changes the space, and vectors from two + # spaces score against each other perfectly happily and mean nothing. + dims: Mapped[int] = mapped_column(Integer, default=0, nullable=False) + # Which model wrote it, for the same reason. A rebuild is what reconciles + # them; until then the odd ones out are ignored rather than trusted. + model_id: Mapped[str] = mapped_column(String(300), default="") + # A hash of the text this set was built from. What makes re-indexing an + # unchanged record free, and what makes "is this index current?" answerable + # without re-embedding anything. + source_hash: Mapped[str] = mapped_column(String(64), default="") + + def __repr__(self) -> str: + return f"" + + +Index("ix_chunks_resource", Chunk.resource_type, Chunk.resource_id) diff --git a/src/lembas/main.py b/src/lembas/main.py index ed5a236..036cff2 100644 --- a/src/lembas/main.py +++ b/src/lembas/main.py @@ -17,6 +17,7 @@ from lembas.api import ( admin_agents, admin_audio, admin_branding, + admin_extraction, admin_images, admin_models, admin_prompts, @@ -45,6 +46,7 @@ from lembas.api import ( from lembas.api.deps import RedirectToLogin, is_htmx, login_redirect from lembas.config import settings from lembas.db.session import init_db +from lembas.services.library import indexing from lembas.web.templating import STATIC_DIR, render log = logging.getLogger("lembas") @@ -79,6 +81,7 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]: from lembas.services.chat import sweep_temporary from lembas.services.files import sweep_orphans from lembas.services.library.documents import sweep_unfiled + from lembas.services.library.indexing import sweep_orphans as sweep_chunks from lembas.services.suggestions import seed_defaults as seed_suggestions with session_scope() as db: @@ -89,6 +92,10 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]: # Temporary chats older than a day. Startup only, like the sweeps # above it -- see services/chat.py:sweep_temporary. sweep_temporary(db) + # Chunks whose record has gone. A backstop for a delete that + # happened with no event loop to schedule the tidy-up -- a CLI + # command, or a cascade from removing an account. + sweep_chunks(db) # Three starting points on the empty screen, written once ever. seed_suggestions(db) except Exception: # noqa: BLE001 - housekeeping must never block startup @@ -150,6 +157,9 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]: # detached remote job, which keeps running and is rehydrated on the next # start. Only the watching stops here. await stop_jobs() + # A chunk set is written whole or not at all, so cancelling loses nothing + # a rebuild does not pick up again. + await indexing.shutdown() log.info("LLeMbas stopped") @@ -165,6 +175,11 @@ def create_app() -> FastAPI: app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static") + # One place that notices a library record changing, rather than a call in + # each of the ten writers that touch those tables. Idempotent, because the + # factory is called per test. See services/library/indexing.py:install. + indexing.install() + app.include_router(pages.router) app.include_router(auth.router) app.include_router(preferences.router) @@ -184,6 +199,7 @@ def create_app() -> FastAPI: app.include_router(admin_models.router) app.include_router(admin_audio.router) app.include_router(admin_branding.router) + app.include_router(admin_extraction.router) app.include_router(admin_search.router) app.include_router(admin_schedules.router) app.include_router(admin_images.router) diff --git a/src/lembas/services/files.py b/src/lembas/services/files.py index 8eb095c..242f586 100644 --- a/src/lembas/services/files.py +++ b/src/lembas/services/files.py @@ -34,6 +34,16 @@ from lembas.db.models import KIND_DOCUMENT, KIND_IMAGE, KIND_TEXT, Attachment log = logging.getLogger(__name__) # --- Limits ------------------------------------------------------------------ +# These are the *defaults*, and an administrator can move every one of them on +# /admin/extraction. They stay here because a default belongs beside the code +# that depends on it, and because `prepare` is called from places with no +# database session at all. +# +# The values are read through `limits()`, a process-level snapshot with the same +# shape and the same reasoning as `services/branding.py`: one query per process, +# dropped when the page saves. Threading a session through `prepare`, +# `_process_image`, `_process_pdf` and `_process_text` would have meant six +# signatures changed to carry a number. MAX_UPLOAD_BYTES = 20 * 1024 * 1024 # Longest edge after downscaling. Large enough for a model to read a screenshot @@ -43,6 +53,9 @@ JPEG_QUALITY = 85 # Pillow's own guard against decompression bombs: a 60,000x60,000 PNG is a few # KB on disk and hundreds of GB decoded. +# +# Deliberately NOT a setting. It is a guard, not a preference, and nothing good +# comes of being able to raise it from a form. Image.MAX_IMAGE_PIXELS = 64_000_000 MAX_PDF_PAGES = 300 @@ -77,6 +90,84 @@ TEXT_EXTENSIONS = { } +@dataclass(frozen=True) +class Limits: + """What extraction is allowed to spend, for one process. + + A snapshot rather than a lookup per call: `prepare` and everything under it + are called from routes, from tool runners and from the startup sweep, and + several of them have no session in hand. The pattern and the cost are the + same as `services/branding.py` -- one query per process, dropped when the + admin page saves, and stale across workers until each next reads. + """ + + max_upload_bytes: int = MAX_UPLOAD_BYTES + max_image_edge: int = MAX_IMAGE_EDGE + jpeg_quality: int = JPEG_QUALITY + max_pdf_pages: int = MAX_PDF_PAGES + max_extracted_chars: int = MAX_EXTRACTED_CHARS + orphan_hours: int = 24 + extra_text_extensions: tuple[str, ...] = () + reject_unreadable_pdf: bool = False + + def media_type_for(self, extension: str) -> str | None: + """The media type for a text extension, or None if it is not one. + + The built-in table first, then the administrator's additions as plain + text. Additions are extensions and not a mapping, because the mapping is + a thing somebody would have to get right twice and the media type of a + `.env` is `text/plain` whatever anybody types. + """ + if extension in TEXT_EXTENSIONS: + return TEXT_EXTENSIONS[extension] + return "text/plain" if extension in self.extra_text_extensions else None + + +_LIMITS: Limits | None = None + + +def limits() -> Limits: + """The current extraction limits. Never raises -- see `branding.snapshot`.""" + global _LIMITS + if _LIMITS is not None: + return _LIMITS + try: + from lembas.db.session import session_scope + from lembas.services import settings_store + + with session_scope() as db: + values = settings_store.extraction(db) + _LIMITS = Limits( + max_upload_bytes=int(values["max_upload_mb"]) * 1024 * 1024, + max_image_edge=int(values["max_image_edge"]), + jpeg_quality=int(values["jpeg_quality"]), + max_pdf_pages=int(values["max_pdf_pages"]), + max_extracted_chars=int(values["max_extracted_chars"]), + orphan_hours=int(values["orphan_hours"]), + extra_text_extensions=tuple( + _clean_extension(item) for item in values["extra_text_extensions"] + ), + reject_unreadable_pdf=bool(values.get("reject_unreadable_pdf")), + ) + except Exception: # noqa: BLE001 - the shipped defaults are a usable answer + log.debug("could not read extraction settings; using defaults", exc_info=True) + return Limits() + return _LIMITS + + +def _clean_extension(raw: str) -> str: + value = str(raw or "").strip().lower() + if not value: + return "" + return value if value.startswith(".") else f".{value}" + + +def forget() -> None: + """Drop the snapshot. Called by the admin page's save, and by tests.""" + global _LIMITS + _LIMITS = None + + class FileError(Exception): """A rejected upload, with a message fit to show the user.""" @@ -135,6 +226,7 @@ def _looks_like_pdf(payload: bytes) -> bool: # --- Processing -------------------------------------------------------------- def _process_image(payload: bytes) -> Prepared: + bounds = limits() try: with Image.open(io.BytesIO(payload)) as image: image.load() @@ -145,8 +237,8 @@ def _process_image(payload: bytes) -> Prepared: width, height = frame.size longest = max(width, height) - if longest > MAX_IMAGE_EDGE: - scale = MAX_IMAGE_EDGE / longest + if longest > bounds.max_image_edge: + scale = bounds.max_image_edge / longest frame = frame.resize( (max(1, int(width * scale)), max(1, int(height * scale))), Image.LANCZOS, @@ -157,7 +249,7 @@ def _process_image(payload: bytes) -> Prepared: frame.save(buffer, format="PNG", optimize=True) media_type, extension = "image/png", ".png" else: - frame.save(buffer, format="JPEG", quality=JPEG_QUALITY, optimize=True) + frame.save(buffer, format="JPEG", quality=bounds.jpeg_quality, optimize=True) media_type, extension = "image/jpeg", ".jpg" return Prepared( @@ -175,6 +267,7 @@ def _process_image(payload: bytes) -> Prepared: def _process_pdf(payload: bytes) -> Prepared: + bounds = limits() from pypdf import PdfReader from pypdf.errors import PdfReadError @@ -198,7 +291,7 @@ def _process_pdf(payload: bytes) -> Prepared: chunks: list[str] = [] total = 0 - for index, page in enumerate(reader.pages[:MAX_PDF_PAGES]): + for index, page in enumerate(reader.pages[:bounds.max_pdf_pages]): try: text = page.extract_text() or "" except Exception as exc: # noqa: BLE001 - one bad page is not fatal @@ -208,14 +301,14 @@ def _process_pdf(payload: bytes) -> Prepared: continue chunks.append(f"[page {index + 1}]\n{text.strip()}") total += len(text) - if total >= MAX_EXTRACTED_CHARS: + if total >= bounds.max_extracted_chars: prepared.truncated = True break - if prepared.pages > MAX_PDF_PAGES: + if prepared.pages > bounds.max_pdf_pages: prepared.truncated = True - prepared.extracted_text = "\n\n".join(chunks)[:MAX_EXTRACTED_CHARS] + prepared.extracted_text = "\n\n".join(chunks)[:bounds.max_extracted_chars] if not prepared.extracted_text.strip(): # Almost always a scan. Saying so beats the model silently ignoring @@ -236,6 +329,7 @@ def _process_pdf(payload: bytes) -> Prepared: def _process_text(payload: bytes, filename: str) -> Prepared: + bounds = limits() for encoding in ("utf-8", "utf-16", "latin-1"): try: text = payload.decode(encoding) @@ -250,15 +344,15 @@ def _process_text(payload: bytes, filename: str) -> Prepared: if "\x00" in text[:4096]: raise FileError("That file is not text, and is not a format LLeMbas can read.") - truncated = len(text) > MAX_EXTRACTED_CHARS + truncated = len(text) > bounds.max_extracted_chars extension = Path(filename).suffix.lower() return Prepared( payload=payload, kind=KIND_TEXT, - media_type=TEXT_EXTENSIONS.get(extension, "text/plain"), - extension=extension if extension in TEXT_EXTENSIONS else ".txt", - extracted_text=text[:MAX_EXTRACTED_CHARS], + media_type=bounds.media_type_for(extension) or "text/plain", + extension=extension if bounds.media_type_for(extension) else ".txt", + extracted_text=text[:bounds.max_extracted_chars], truncated=truncated, ) @@ -311,8 +405,9 @@ def prepare(payload: bytes, filename: str, *, keep_original: bool = False) -> Pr """ if not payload: raise FileError("That file is empty.") - if len(payload) > MAX_UPLOAD_BYTES: - raise FileError(f"Files must be under {MAX_UPLOAD_BYTES // (1024 * 1024)} MB.") + ceiling = limits().max_upload_bytes + if len(payload) > ceiling: + raise FileError(f"Files must be under {ceiling // (1024 * 1024)} MB.") if _detect_image(payload) is not None: return _keep_image(payload) if keep_original else _process_image(payload) @@ -406,7 +501,7 @@ def store_text( on the tag around it, which is what a reader sees on the chip and what survives if the text is later truncated away from its own first line. """ - body = text[:MAX_EXTRACTED_CHARS] + body = text[:limits().max_extracted_chars] payload = body.encode("utf-8") stored_name = f"{secrets.token_hex(16)}.txt" @@ -565,12 +660,19 @@ def remove_files_for_chats(db: DBSession, chat_ids: list[str]) -> int: return removed -def sweep_orphans(db: DBSession, older_than: timedelta = ORPHAN_AGE) -> int: +def sweep_orphans(db: DBSession, older_than: timedelta | None = None) -> int: """Delete uploads that were never attached to a message. A file picked in the composer and then abandoned would otherwise sit on disk forever. + + `older_than` defaults to the configured age rather than to a constant, and + it is resolved *here* rather than in the signature: a default argument is + evaluated at import, so a module-level `ORPHAN_AGE` in the signature would + pin the shipped 24 hours whatever an administrator later set. """ + if older_than is None: + older_than = timedelta(hours=limits().orphan_hours) cutoff = datetime.now(UTC) - older_than orphans = list(db.scalars(select(Attachment).where(Attachment.message_id.is_(None)))) @@ -608,7 +710,7 @@ def data_uri(attachment: Attachment) -> str | None: return f"data:{attachment.media_type};base64,{encoded}" -def preview_data_uri(payload: bytes, *, max_edge: int = MAX_IMAGE_EDGE) -> str | None: +def preview_data_uri(payload: bytes, *, max_edge: int = 0) -> str | None: """The same thing for bytes in hand, downscaled, for a model to look at. Fidelity and weight are two different jobs. What is stored is what ComfyUI @@ -619,9 +721,15 @@ def preview_data_uri(payload: bytes, *, max_edge: int = MAX_IMAGE_EDGE) -> str | Takes bytes rather than an Attachment: the reviewer looks at an image that may be about to be thrown away, and writing a row for something rejected seconds later is work with nothing to show for it. + + `max_edge` of 0 means the configured one. Zero rather than None because the + caller that passes a number passes a number, and a sentinel that is also a + plausible value would be worse -- an edge of zero is not a picture. """ import base64 + max_edge = max_edge or limits().max_image_edge + try: with Image.open(io.BytesIO(payload)) as image: image.load() @@ -634,7 +742,7 @@ def preview_data_uri(payload: bytes, *, max_edge: int = MAX_IMAGE_EDGE) -> str | Image.LANCZOS, ) buffer = io.BytesIO() - frame.save(buffer, format="JPEG", quality=JPEG_QUALITY, optimize=True) + frame.save(buffer, format="JPEG", quality=limits().jpeg_quality, optimize=True) except (Image.DecompressionBombError, UnidentifiedImageError, OSError, ValueError): log.warning("could not build a preview of a generated image", exc_info=True) return None diff --git a/src/lembas/services/library/chunks.py b/src/lembas/services/library/chunks.py new file mode 100644 index 0000000..f6529a1 --- /dev/null +++ b/src/lembas/services/library/chunks.py @@ -0,0 +1,125 @@ +"""Splitting a record into pieces small enough to embed, and packing vectors. + +One implementation, used by documents, notes, skills and reports. Three would +drift, and drift here is invisible: a splitter that behaves differently for +notes than for documents produces a search that works and ranks wrongly. + +## How it splits + +On **paragraph boundaries first**, falling back to lines and then to a hard cut, +because a chunk that ends mid-sentence is one whose embedding is about half a +thought. The overlap carries the tail of the previous chunk into the next, so a +sentence that straddles a boundary is whole in one of them. + +Characters rather than tokens throughout. The count has to be made without +asking the endpoint -- `services/tokens.py` already establishes four characters +to a token as this codebase's estimate, and being 20% out about a chunk size is +a slightly different chunk, not a wrong one. + +## Packing + +float32, little-endian. A 1024-dimension vector is 4KB packed and about 20KB as +JSON text, and every one of them is read on every semantic search. +""" + +from __future__ import annotations + +import hashlib +import struct + +# Below this a piece is not worth a row: the embedding of six words is mostly +# noise, and a search that returns "and the following:" as its best hit is worse +# than one that returns nothing. +MIN_CHUNK_CHARS = 40 + + +def split(text: str, *, size: int = 1200, overlap: int = 150) -> list[str]: + """A record's text as pieces of roughly `size` characters. + + `overlap` is how much of the previous piece rides along with the next. It is + clamped to half the size here as well as in the settings accessor, because + an overlap at or past the size means every piece starts where the last one + did and the loop never advances -- a hang rather than a bad index, so it is + refused in both places rather than in the more convenient one. + """ + body = (text or "").strip() + if not body: + return [] + size = max(200, int(size)) + overlap = max(0, min(int(overlap), size // 2)) + if len(body) <= size: + return [body] + + pieces: list[str] = [] + start = 0 + while start < len(body): + end = min(start + size, len(body)) + if end < len(body): + end = _boundary(body, start, end) + piece = body[start:end].strip() + if len(piece) >= MIN_CHUNK_CHARS: + pieces.append(piece) + if end >= len(body): + break + start = max(end - overlap, start + 1) + return pieces + + +def _boundary(body: str, start: int, end: int) -> int: + """Where to cut, preferring a paragraph break and then a line break. + + Searched backwards from the hard limit, and only within the last third of + the piece: a paragraph break near the *start* would produce a chunk a + fraction of the size, which is how a long document turns into hundreds of + tiny rows that each match nothing. + """ + floor = start + (end - start) * 2 // 3 + for marker in ("\n\n", "\n", ". "): + found = body.rfind(marker, floor, end) + if found > floor: + return found + len(marker) + return end + + +def digest(text: str) -> str: + """A hash of what a chunk set was built from. + + What makes re-indexing an unchanged record free, and what makes "is this + index current?" answerable without embedding anything. sha256 rather than + md5 for no reason beyond having no reason to prefer md5; both are being used + as a change detector rather than against an adversary. + """ + return hashlib.sha256((text or "").encode("utf-8")).hexdigest() + + +def pack(vector: list[float]) -> bytes: + return struct.pack(f"<{len(vector)}f", *vector) + + +def unpack(blob: bytes, dims: int) -> list[float]: + """A stored vector, or an empty list if the row does not add up. + + Length is checked against the declared width rather than inferred from it: a + truncated BLOB would otherwise unpack into a shorter vector and score + against a query happily, which is a wrong answer rather than a missing one. + """ + if dims <= 0 or len(blob) != dims * 4: + return [] + return list(struct.unpack(f"<{dims}f", blob)) + + +def dot(left: list[float], right: list[float]) -> float: + """Cosine similarity, given that both sides are already unit vectors. + + Normalisation happens once, at write time, in `llm/embeddings.py` -- so + every comparison here is a multiply-and-add rather than two square roots per + pair. A width mismatch scores zero rather than raising: it means the vectors + came from two different models, and the honest answer to "how similar are + these?" across two spaces is "this tells you nothing". + """ + if len(left) != len(right) or not left: + return 0.0 + return sum(a * b for a, b in zip(left, right, strict=True)) + + +__all__ = ["MIN_CHUNK_CHARS", "digest", "dot", "pack", "split", "unpack"] diff --git a/src/lembas/services/library/documents.py b/src/lembas/services/library/documents.py index 12ef76d..6037e37 100644 --- a/src/lembas/services/library/documents.py +++ b/src/lembas/services/library/documents.py @@ -18,11 +18,18 @@ from sqlalchemy import select from sqlalchemy.orm import Session as DBSession from lembas.config import settings -from lembas.db.models import SOURCE_LINK, SOURCE_UPLOAD, Document, KnowledgeBase, User +from lembas.db.models import ( + CHUNK_DOCUMENT, + SOURCE_LINK, + SOURCE_UPLOAD, + Document, + KnowledgeBase, + User, +) from lembas.services import files as files_service from lembas.services import sharing from lembas.services.fetch import Fetched -from lembas.services.library.fts import search_ids +from lembas.services.library import retrieval log = logging.getLogger(__name__) @@ -287,8 +294,9 @@ def set_text(db: DBSession, document: Document, text: str) -> Document: The commit fires the `documents_fts` UPDATE trigger, so search stays correct with nothing else to do. See `db/migrations.py:ensure_fts`. """ - document.extracted_text = text[:files_service.MAX_EXTRACTED_CHARS] - document.truncated = len(text) > files_service.MAX_EXTRACTED_CHARS + ceiling = files_service.limits().max_extracted_chars + document.extracted_text = text[:ceiling] + document.truncated = len(text) > ceiling document.extraction_error = "" db.commit() return document @@ -301,14 +309,22 @@ def search( *, limit: int = 10, base_ids: list[str] | None = None, + vector: list[float] | None = None, ) -> list[Document]: """Documents matching `needle` that this user may see, best match first. The index is searched first and the visibility filter applied to the rows it returned. That order matters: filtering afterwards is what makes it impossible for a hit on somebody else's document to leak, even as a count. + + `vector` is the query already embedded, or None. It comes from the caller + rather than being worked out here because this is synchronous and embedding + is an HTTP request -- see `services/library/retrieval.py`. None means the + keyword search exactly as it always was. """ - hits = search_ids(db, INDEX, needle, limit=limit * 4) + hits = retrieval.search( + db, INDEX, needle, kind=CHUNK_DOCUMENT, vector=vector, limit=limit * 4 + ) if not hits: return [] diff --git a/src/lembas/services/library/indexing.py b/src/lembas/services/library/indexing.py new file mode 100644 index 0000000..2e609e2 --- /dev/null +++ b/src/lembas/services/library/indexing.py @@ -0,0 +1,529 @@ +"""Keeping the semantic index current, and rebuilding it when it is not. + +## The shape, and why it is a background task + +Embedding is an HTTP request. Every writer in the library -- `documents.create`, +`notes.edit`, `skills.save`, `reports.create` -- is synchronous and is called +from a route or a tool runner that has just committed a row, and none of them +should wait on a model server to answer before saying "saved". + +So indexing is **fired and forgotten**: `schedule(kind, id)` starts a task and +returns immediately. A save that cannot be indexed is a save; the row is written +either way and the search falls back to keywords for that record until the next +rebuild. That is the whole degradation story, and it is the same one that covers +having no embedding model at all. + +## Nothing is written when no model is chosen + +`embedding_model_id` empty means the FTS path exactly as it has always been -- +no chunk rows, no requests, no cost. That is what makes this safe to add to an +instance that never asked for it, and it is asserted rather than assumed. + +## Staleness is a hash, not a timestamp + +Every chunk carries `source_hash` (of the text it was built from), `model_id` +and `dims`. Re-indexing an unchanged record is free; a record whose text moved +is rebuilt; a record embedded by a *different* model is rebuilt on the next pass +and, until then, ignored by the scorer rather than trusted. Vectors from two +spaces score against each other perfectly happily and mean nothing, which is a +search that works and is wrong -- the worst failure this feature can have. + +## The rebuild is restartable and reports itself + +A half-finished index has to be usable rather than empty, so the rebuild walks +records one at a time and commits each. `progress()` is what the admin page +polls; it is in-process, because a rebuild does not survive a restart and +pretending otherwise would mean a progress bar that never moves. +""" + +from __future__ import annotations + +import asyncio +import contextlib +import logging +from dataclasses import dataclass, field + +from sqlalchemy import delete, func, select +from sqlalchemy.orm import Session as DBSession + +from lembas.db.models import ( + CHUNK_DOCUMENT, + CHUNK_KINDS, + CHUNK_NOTE, + CHUNK_REPORT, + CHUNK_SKILL, + Chunk, + Connection, + Document, + Model, + Note, + Report, + Skill, +) +from lembas.db.session import session_scope +from lembas.services import settings_store +from lembas.services.library import chunks as chunk_service +from lembas.services.llm.openai_client import Endpoint, LLMError + +log = logging.getLogger(__name__) + +# What each kind is, and how to get its text. One table rather than four +# branches, for the reason `tool_labels` is one table: four copies of "which +# columns make up the searchable text" is three chances to disagree. +SOURCES: dict[str, tuple[type, tuple[str, ...]]] = { + CHUNK_DOCUMENT: (Document, ("title", "description", "extracted_text")), + CHUNK_NOTE: (Note, ("title", "body")), + CHUNK_SKILL: (Skill, ("name", "description", "body")), + CHUNK_REPORT: (Report, ("title", "summary", "body")), +} + +# Tasks in flight, so a record saved twice in quick succession is indexed once +# more rather than twice at the same time. Keyed on kind and id. +_TASKS: dict[tuple[str, str], asyncio.Task] = {} + + +# --- What the model is ---------------------------------------------------------- +@dataclass(frozen=True) +class Embedder: + """Which model turns text into vectors, resolved while a session is open.""" + + endpoint: Endpoint + model_id: str + batch: int = 16 + + +def embedder(db: DBSession) -> Embedder | None: + """The configured embedding model, or None. + + None is the answer to every "no" -- none chosen, the model row deleted, its + connection disabled -- and every caller reads it the same way: do nothing, + and let the keyword search stand. That is deliberately not an error. An + instance that never configured this is the common case, not a broken one. + """ + values = settings_store.extraction(db) + wanted = str(values.get("embedding_model_id") or "").strip() + if not wanted: + return None + model = db.scalar( + select(Model) + .join(Connection) + .where( + Model.model_id == wanted, + Model.enabled.is_(True), + Connection.enabled.is_(True), + ) + .order_by(Connection.position) + ) + if model is None: + log.info("embedding model %r is configured but not available", wanted) + return None + connection = db.get(Connection, model.connection_id) + if connection is None: + return None + return Embedder( + endpoint=Endpoint.from_connection(connection), + model_id=model.model_id, + batch=int(values.get("embed_batch") or 16), + ) + + +def enabled(db: DBSession) -> bool: + return embedder(db) is not None + + +# --- Reading a record ----------------------------------------------------------- +def text_of(row) -> str: + """The searchable text of one record, in the same order the FTS index uses. + + Blank fields are dropped rather than joined as empty lines, so a note with + no body hashes the same before and after somebody clears its body twice. + """ + kind = kind_of(row) + if kind is None: + return "" + _, columns = SOURCES[kind] + parts = [str(getattr(row, name, "") or "").strip() for name in columns] + return "\n\n".join(part for part in parts if part) + + +def kind_of(row) -> str | None: + for kind, (model, _) in SOURCES.items(): + if isinstance(row, model): + return kind + return None + + +def owner_of(row) -> str: + return str(getattr(row, "owner_id", "") or "") + + +# --- Writing the index ---------------------------------------------------------- +def forget_resource(db: DBSession, kind: str, resource_id: str) -> int: + """Drop every chunk of one record. Called when it is deleted. + + A plain DELETE rather than a cascade, because `resource_id` has no foreign + key -- it points at one of four tables depending on `resource_type`, which + SQLite cannot express. Same reasoning as `Share.principal_id`. + """ + result = db.execute( + delete(Chunk).where(Chunk.resource_type == kind, Chunk.resource_id == resource_id) + ) + db.commit() + return int(result.rowcount or 0) + + +def current_hash(db: DBSession, kind: str, resource_id: str) -> tuple[str, str]: + """The hash and model of the chunks already stored for a record.""" + row = db.execute( + select(Chunk.source_hash, Chunk.model_id) + .where(Chunk.resource_type == kind, Chunk.resource_id == resource_id) + .limit(1) + ).first() + return (str(row[0] or ""), str(row[1] or "")) if row else ("", "") + + +async def index_resource(kind: str, resource_id: str, *, force: bool = False) -> int: + """Rebuild one record's chunks. Returns how many were written. + + Opens its own session, for the reason every background worker here does: it + outlives the request that scheduled it. Never raises -- a failure leaves the + old chunks in place, which is a slightly stale index rather than a hole, and + is strictly better than deleting first and failing to write. + """ + if kind not in SOURCES: + return 0 + try: + with session_scope() as db: + model, _ = SOURCES[kind] + row = db.get(model, resource_id) + if row is None: + forget_resource(db, kind, resource_id) + return 0 + worker = embedder(db) + if worker is None: + return 0 + body = text_of(row) + owner = owner_of(row) + values = settings_store.extraction(db) + digest = chunk_service.digest(body) + stored_hash, stored_model = current_hash(db, kind, resource_id) + + if not body.strip(): + with session_scope() as db: + forget_resource(db, kind, resource_id) + return 0 + if not force and digest == stored_hash and stored_model == worker.model_id: + return 0 + + pieces = chunk_service.split( + body, size=int(values["chunk_chars"]), overlap=int(values["chunk_overlap"]) + ) + if not pieces: + with session_scope() as db: + forget_resource(db, kind, resource_id) + return 0 + + vectors = await _embed_all(worker, pieces) + + # Written only once every vector is in hand. Deleting first and failing + # half way through would leave a record indexed by half of itself, which + # ranks worse than not being indexed at all and looks like nothing. + with session_scope() as db: + db.execute( + delete(Chunk).where( + Chunk.resource_type == kind, Chunk.resource_id == resource_id + ) + ) + for ordinal, (piece, vector) in enumerate(zip(pieces, vectors, strict=True)): + db.add( + Chunk( + owner_id=owner, + resource_type=kind, + resource_id=resource_id, + ordinal=ordinal, + text=piece, + vector=chunk_service.pack(vector), + dims=len(vector), + model_id=worker.model_id, + source_hash=digest, + ) + ) + db.commit() + return len(pieces) + except LLMError as exc: + log.info("could not index %s %s: %s", kind, resource_id, exc) + return 0 + except asyncio.CancelledError: + raise + except Exception: # noqa: BLE001 - one bad record must not stop a rebuild + log.exception("indexing %s %s failed", kind, resource_id) + return 0 + + +async def _embed_all(worker: Embedder, pieces: list[str]) -> list[list[float]]: + from lembas.services.llm import embeddings as embeddings_service + + vectors: list[list[float]] = [] + for start in range(0, len(pieces), worker.batch): + batch = pieces[start : start + worker.batch] + vectors.extend(await embeddings_service.embed(worker.endpoint, worker.model_id, batch)) + return vectors + + +# --- Scheduling ----------------------------------------------------------------- +def schedule(kind: str, resource_id: str) -> None: + """Index a record soon, without making its writer wait. + + Called from synchronous writers that have just committed. Two things it is + careful about: + + - **No running loop means do nothing.** A CLI command, a test, or the + startup sweep has no event loop to attach to, and building a coroutine + there produces "never awaited" at the caller's own line. The check is + before the coroutine, the same trap `push.announce_later` documents. + - **A record already being indexed is left alone.** Saving twice in a second + would otherwise embed the same text twice at once; the second call is + dropped and the record is picked up by the *next* save or rebuild, which + is why `index_resource` re-reads the row rather than taking text passed in. + """ + if kind not in SOURCES or not resource_id: + return + try: + asyncio.get_running_loop() + except RuntimeError: + return + key = (kind, resource_id) + existing = _TASKS.get(key) + if existing is not None and not existing.done(): + return + task = asyncio.create_task(index_resource(kind, resource_id)) + _TASKS[key] = task + task.add_done_callback(lambda _t, k=key: _TASKS.pop(k, None)) + + +def schedule_for(row) -> None: + """The same, given a record rather than its kind and id.""" + kind = kind_of(row) + if kind is not None: + schedule(kind, str(getattr(row, "id", "") or "")) + + +# --- Noticing a change ---------------------------------------------------------- +# Two SQLAlchemy session events rather than a call in each of the ten writers +# that touch these four tables. That is a departure from this codebase's taste +# for explicit seams, and the reason is the one `tool_label` gives for being a +# Jinja global: a step every writer has to remember is a step one of them will +# forget, and here forgetting is *silent* -- the record saves, the keyword search +# still finds it, and only its semantic recall is quietly stale. +# +# `after_flush` collects and `after_commit` acts, in that order and never +# merged. Inside a flush the transaction has not landed yet, so a task started +# there could read the row before it exists; and `session.deleted` is empty by +# the time the commit fires, so the collecting has to happen while it is not. +_PENDING = "lembas_index_pending" + + +def _collect(session, _flush_context) -> None: + seen: set[tuple[str, str]] = session.info.setdefault(_PENDING, set()) + for row in (*session.new, *session.dirty, *session.deleted): + kind = kind_of(row) + if kind is None: + continue + resource_id = str(getattr(row, "id", "") or "") + if resource_id: + seen.add((kind, resource_id)) + + +def _fire(session) -> None: + # A deletion is scheduled exactly like a change: `index_resource` finds no + # row and drops the chunks. One path rather than two, and the one that runs + # is the one that has to be right anyway. + for kind, resource_id in session.info.pop(_PENDING, set()): + schedule(kind, resource_id) + + +def _forget(session) -> None: + session.info.pop(_PENDING, None) + + +def install() -> None: + """Listen for library records changing. Called once, from the app factory. + + Idempotent: `event.contains` is checked, because the app factory is called + per test in the suite and registering the same listener a hundred times + would index every record a hundred times over. + """ + from sqlalchemy import event + from sqlalchemy.orm import Session + + for name, handler in ( + ("after_flush", _collect), + ("after_commit", _fire), + ("after_rollback", _forget), + ): + if not event.contains(Session, name, handler): + event.listen(Session, name, handler) + + +def sweep_orphans(db: DBSession) -> int: + """Drop chunks whose record has gone. + + A backstop for the one case the listeners cannot cover: a delete that + happened with no event loop running -- a CLI command, a test, a cascade from + deleting a user -- where `schedule` had nowhere to put its task. Cheap + enough to run at startup and at the end of every rebuild: one NOT IN per + kind, against an indexed column. + """ + removed = 0 + for kind, (model, _) in SOURCES.items(): + result = db.execute( + delete(Chunk).where( + Chunk.resource_type == kind, + Chunk.resource_id.not_in(select(model.id)), + ) + ) + removed += int(result.rowcount or 0) + if removed: + db.commit() + log.info("dropped %d orphaned chunk(s)", removed) + return removed + + +# --- Rebuilding everything ------------------------------------------------------ +@dataclass +class Progress: + """What a rebuild has done so far. + + In-process, because a rebuild does not survive a restart. Persisting it + would mean a progress bar that stops moving and never finishes, which is + worse than one that admits it is gone. + """ + + running: bool = False + total: int = 0 + done: int = 0 + written: int = 0 + error: str = "" + kinds: dict[str, int] = field(default_factory=dict) + + @property + def percent(self) -> int: + return int(self.done * 100 / self.total) if self.total else 0 + + +_PROGRESS = Progress() +_REBUILD: asyncio.Task | None = None + + +def progress() -> Progress: + return _PROGRESS + + +def counts(db: DBSession) -> dict[str, int]: + """How many chunks exist per kind. What the page shows when nothing is running.""" + rows = db.execute( + select(Chunk.resource_type, func.count()).group_by(Chunk.resource_type) + ).all() + return {str(kind): int(count) for kind, count in rows} + + +async def rebuild_all(*, force: bool = True) -> None: + """Walk every record and index it, committing as it goes. + + One at a time and never gathered. The far side is usually one local model + server, and twenty concurrent embedding requests against it is slower than + twenty sequential ones as well as being ruder. + """ + global _PROGRESS + _PROGRESS = Progress(running=True) + try: + with session_scope() as db: + if embedder(db) is None: + _PROGRESS.error = "No embedding model is configured." + return + work: list[tuple[str, str]] = [] + for kind, (model, _) in SOURCES.items(): + ids = [row[0] for row in db.execute(select(model.id)).all()] + work.extend((kind, str(row_id)) for row_id in ids) + _PROGRESS.total = len(work) + + for kind, resource_id in work: + written = await index_resource(kind, resource_id, force=force) + _PROGRESS.done += 1 + _PROGRESS.written += written + _PROGRESS.kinds[kind] = _PROGRESS.kinds.get(kind, 0) + written + + # After the walk, not before: a record deleted while this was running + # would otherwise be swept and then re-indexed from a row that no longer + # exists. `index_resource` handles that case too, and doing it in this + # order means one pass reconciles both directions. + with session_scope() as db: + sweep_orphans(db) + except asyncio.CancelledError: + _PROGRESS.error = "Stopped." + raise + except Exception as exc: # noqa: BLE001 - a rebuild failing must be reportable + log.exception("rebuilding the index failed") + _PROGRESS.error = str(exc) + finally: + _PROGRESS.running = False + + +def start_rebuild(*, force: bool = True) -> bool: + """Start a rebuild if one is not already going. True if this call started it.""" + global _REBUILD + if _REBUILD is not None and not _REBUILD.done(): + return False + try: + asyncio.get_running_loop() + except RuntimeError: + return False + _REBUILD = asyncio.create_task(rebuild_all(force=force)) + return True + + +async def shutdown() -> None: + """Cancel the rebuild and any in-flight indexing. + + Nothing here is lost that matters: a chunk set is either written whole or + not at all, and the next rebuild picks up whatever was missed. + """ + global _REBUILD + tasks = [task for task in (_REBUILD, *_TASKS.values()) if task is not None] + _TASKS.clear() + _REBUILD = None + for task in tasks: + task.cancel() + for task in tasks: + with contextlib.suppress(asyncio.CancelledError, Exception): + await task + + +def clear() -> None: + """For tests: forget the in-process state without touching the database.""" + global _REBUILD, _PROGRESS + _TASKS.clear() + _REBUILD = None + _PROGRESS = Progress() + + +__all__ = [ + "CHUNK_KINDS", + "SOURCES", + "Embedder", + "Progress", + "clear", + "counts", + "embedder", + "enabled", + "forget_resource", + "index_resource", + "kind_of", + "progress", + "rebuild_all", + "schedule", + "schedule_for", + "shutdown", + "start_rebuild", + "text_of", +] diff --git a/src/lembas/services/library/notes.py b/src/lembas/services/library/notes.py index 9d025a8..8d695f6 100644 --- a/src/lembas/services/library/notes.py +++ b/src/lembas/services/library/notes.py @@ -13,9 +13,9 @@ import logging from sqlalchemy import select from sqlalchemy.orm import Session as DBSession -from lembas.db.models import AUTHOR_MODEL, AUTHOR_USER, Note, User +from lembas.db.models import AUTHOR_MODEL, AUTHOR_USER, CHUNK_NOTE, Note, User from lembas.services import sharing -from lembas.services.library.fts import search_ids +from lembas.services.library import retrieval log = logging.getLogger(__name__) @@ -43,9 +43,22 @@ def recent(db: DBSession, user: User | None, *, limit: int = 20) -> list[Note]: ) -def search(db: DBSession, user: User | None, needle: str, *, limit: int = 10) -> list[Note]: - """Notes matching `needle` that this user may see, best match first.""" - hits = search_ids(db, INDEX, needle, limit=limit * 4) +def search( + db: DBSession, + user: User | None, + needle: str, + *, + limit: int = 10, + vector: list[float] | None = None, +) -> list[Note]: + """Notes matching `needle` that this user may see, best match first. + + `vector` is the query already embedded, or None. It comes from the caller + rather than being worked out here because this is synchronous and embedding + is an HTTP request -- see `services/library/retrieval.py`. None means the + keyword search exactly as it always was. + """ + hits = retrieval.search(db, INDEX, needle, kind=CHUNK_NOTE, vector=vector, limit=limit * 4) if not hits: return [] order = {hit.id: position for position, hit in enumerate(hits)} diff --git a/src/lembas/services/library/retrieval.py b/src/lembas/services/library/retrieval.py new file mode 100644 index 0000000..2e5697f --- /dev/null +++ b/src/lembas/services/library/retrieval.py @@ -0,0 +1,207 @@ +"""Finding things: keywords, meaning, and the two fused. + +`fts.search_ids` was already the one seam every store searches through. This +sits beside it and keeps that true — the four stores still call one function and +still get ids back, and what changed is what is behind it. + +## Reciprocal rank fusion, and why not a weight + +Two rankings have to become one, and their scores are not comparable: bm25 is a +negative number whose scale depends on the corpus, cosine is 0..1. Normalising +them onto a common scale means picking a constant, and that constant is a knob +nobody can tune without a labelled test set they do not have. + +RRF uses the **ranks** and not the scores: `1 / (K + rank)`, summed. It has one +constant, `K`, it is famously insensitive to it, and it degrades to exactly one +of the two lists when the other is empty — which is what makes "no embedding +model configured" mean the keyword search, unchanged, with no branch anywhere +that says so. + +## The query is embedded by the caller, not here + +`search` is synchronous, because every store's `search()` is and every one of +them is called from both a route and a tool runner. Embedding is an HTTP request. +So a caller that can await gets the query vector first and passes it in; one that +cannot passes nothing and gets keywords. `embed_query` is the async half, and +being able to answer `None` for every "no" is what keeps that from being a branch +at each call site. + +## Visibility is still somebody else's job + +Both halves return ids, and both are scored across *everything* — the filter is +applied to the row query afterwards, in each store, through +`services/sharing.py`. That order is deliberate and is the same one the +full-text path has always used: filtering afterwards is what makes it impossible +for a hit on somebody else's record to leak, even as a count. +""" + +from __future__ import annotations + +import logging + +from sqlalchemy import select +from sqlalchemy.orm import Session as DBSession + +from lembas.db.models import Chunk +from lembas.services.library import chunks as chunk_service +from lembas.services.library.fts import SearchHit, fts_query, search_ids + +log = logging.getLogger(__name__) + +# The one constant in reciprocal rank fusion. 60 is what the original paper used +# and what everything since has copied; the method's whole appeal is that the +# result barely moves for anything in the tens. It is not a tuning knob and is +# deliberately not a setting -- a number nobody can evaluate is a number nobody +# should be asked about. +RRF_K = 60 + +# How many chunks are scored before they are collapsed to records. Larger than +# the number of records wanted, because one long document can own several of the +# best chunks and would otherwise crowd everything else out of the answer. +CHUNK_MULTIPLIER = 6 + + +def embeddable(db: DBSession) -> bool: + from lembas.services.library import indexing + + return indexing.enabled(db) + + +def worker_for(db: DBSession): + """The configured embedder, resolved while a session is open. + + Split from the awaiting half deliberately. A caller that must not hold a + database session across an HTTP request -- a tool runner, which is about to + open its own -- resolves here, closes, and awaits `embed_with`. One that + already holds a request's session and is content to keep it can use + `embed_query` instead. + """ + from lembas.services.library import indexing + + return indexing.embedder(db) + + +async def embed_with(worker, needle: str) -> list[float] | None: + """The query as a vector, or None. + + None for every "no": no model configured, an empty query, an endpoint that + is down. Each of them means the same thing to the caller — search by + keywords — so none of them is an error, and a search that quietly stops + being semantic is far better than one that 500s because a model server was + restarting. + """ + from lembas.services.llm import embeddings as embeddings_service + from lembas.services.llm.openai_client import LLMError + + if worker is None or not (needle or "").strip(): + return None + try: + vectors = await embeddings_service.embed(worker.endpoint, worker.model_id, [needle]) + except LLMError as exc: + log.info("could not embed a query: %s", exc) + return None + return vectors[0] if vectors else None + + +async def embed_query(db: DBSession, needle: str) -> list[float] | None: + """`worker_for` and `embed_with`, for a caller happy to hold its session.""" + return await embed_with(worker_for(db), needle) + + +def semantic_ids( + db: DBSession, kind: str, vector: list[float], *, limit: int = 20 +) -> list[SearchHit]: + """Record ids whose best chunk is closest to `vector`, best first. + + A brute-force scan, and that is the right answer at this scale: a library of + ten thousand chunks is forty megabytes of float32 and a few million + multiply-adds, which is milliseconds. A real index is a later change behind + this same call, which is why the signature says nothing about how. + + **A record scores as its best chunk, not its average.** One paragraph that + answers the question is what makes a document worth returning; averaging + would rank a long document about something else above a short one that says + exactly the thing, because most of the long one is not about anything. + + Chunks whose width does not match the query's are skipped. That is a change + of embedding model with a rebuild still pending, and scoring across two + spaces produces a confident wrong answer rather than a missing one. + """ + if not vector: + return [] + width = len(vector) + rows = db.execute( + select(Chunk.resource_id, Chunk.vector, Chunk.dims).where(Chunk.resource_type == kind) + ).all() + + best: dict[str, float] = {} + for resource_id, blob, dims in rows: + if int(dims or 0) != width: + continue + stored = chunk_service.unpack(blob, int(dims)) + if not stored: + continue + score = chunk_service.dot(vector, stored) + key = str(resource_id) + if score > best.get(key, -2.0): + best[key] = score + + ordered = sorted(best.items(), key=lambda pair: pair[1], reverse=True) + return [SearchHit(id=key, rank=score) for key, score in ordered[: max(1, limit)]] + + +def fuse(*rankings: list[SearchHit], limit: int = 20) -> list[SearchHit]: + """Reciprocal rank fusion of any number of rankings. + + The returned `rank` is the fused score, and it is **larger for better**, + which is the opposite of bm25's convention. Nothing downstream reads it -- + every caller uses the order — but it is worth saying out loud rather than + leaving somebody to infer it from a negative number that is no longer there. + """ + scores: dict[str, float] = {} + for ranking in rankings: + for position, hit in enumerate(ranking): + scores[hit.id] = scores.get(hit.id, 0.0) + 1.0 / (RRF_K + position + 1) + ordered = sorted(scores.items(), key=lambda pair: pair[1], reverse=True) + return [SearchHit(id=key, rank=score) for key, score in ordered[: max(1, limit)]] + + +def search( + db: DBSession, + index: str, + needle: str, + *, + kind: str = "", + vector: list[float] | None = None, + limit: int = 20, +) -> list[SearchHit]: + """Ids matching `needle`, keywords and meaning fused. + + With no `vector` this is `fts.search_ids` and nothing else — the same call, + the same results, in the same order. That is what makes an instance with no + embedding model byte-for-byte what it always was, and it is asserted by a + test rather than left as a claim. + """ + keyword = search_ids(db, index, needle, limit=limit) + if not vector or not kind: + return keyword + meaning = semantic_ids(db, kind, vector, limit=limit * CHUNK_MULTIPLIER) + if not meaning: + return keyword + if not keyword and not fts_query(needle): + # Nothing typed that FTS could match — a query of pure punctuation, or + # one whose every word is a separator. The semantic side still has an + # answer, and fusing a list with nothing is that list. + return meaning[:limit] + return fuse(keyword, meaning, limit=limit) + + +__all__ = [ + "CHUNK_MULTIPLIER", + "RRF_K", + "embed_query", + "embeddable", + "fuse", + "search", + "semantic_ids", +] diff --git a/src/lembas/services/library/skills.py b/src/lembas/services/library/skills.py index 5acd45b..5ba888c 100644 --- a/src/lembas/services/library/skills.py +++ b/src/lembas/services/library/skills.py @@ -28,9 +28,9 @@ from collections.abc import Iterable from sqlalchemy import select from sqlalchemy.orm import Session as DBSession -from lembas.db.models import AUTHOR_MODEL, AUTHOR_USER, Skill, SkillRevision, User +from lembas.db.models import AUTHOR_MODEL, AUTHOR_USER, CHUNK_SKILL, Skill, SkillRevision, User from lembas.services import sharing -from lembas.services.library.fts import search_ids +from lembas.services.library import retrieval log = logging.getLogger(__name__) @@ -103,8 +103,22 @@ def count_enabled(db: DBSession, user: User | None, *, exclude: Iterable[str] = return len(enabled_for(db, user, exclude=exclude)) -def search(db: DBSession, user: User | None, needle: str, *, limit: int = 10) -> list[Skill]: - hits = search_ids(db, INDEX, needle, limit=limit * 4) +def search( + db: DBSession, + user: User | None, + needle: str, + *, + limit: int = 10, + vector: list[float] | None = None, +) -> list[Skill]: + """Skills matching `needle` that this user may see, best match first. + + `vector` is the query already embedded, or None. It comes from the caller + rather than being worked out here because this is synchronous and embedding + is an HTTP request -- see `services/library/retrieval.py`. None means the + keyword search exactly as it always was. + """ + hits = retrieval.search(db, INDEX, needle, kind=CHUNK_SKILL, vector=vector, limit=limit * 4) if not hits: return [] order = {hit.id: position for position, hit in enumerate(hits)} diff --git a/src/lembas/services/llm/embeddings.py b/src/lembas/services/llm/embeddings.py new file mode 100644 index 0000000..c4e2a03 --- /dev/null +++ b/src/lembas/services/llm/embeddings.py @@ -0,0 +1,144 @@ +"""Turning text into vectors, against an OpenAI-shaped `/v1/embeddings`. + +The same reasoning as the chat and audio clients: plain httpx rather than an +SDK, because the target is llama.cpp, Ollama, LM Studio, Infinity or vLLM at +least as often as it is api.openai.com. They agree about the request and +disagree politely about the response, so this is tolerant about what comes back +and strict about what it hands on. + +**Batched, because the cost is the round trip.** A hundred chunks one at a time +against a local endpoint is a hundred model loads' worth of latency for work +that fits in six requests. The batch size is a setting, because "how many at +once" is a property of the far side rather than of this code. + +**Dimensions are discovered, never declared.** Nobody should have to look up +that bge-m3 is 1024 and nomic-embed-text is 768, and an instance that changes +model must not silently compare vectors from two different spaces -- +`services/library/indexing.py` records the width beside every vector and +refuses to score across a mismatch. + +Normalisation happens here, once, on the way out. Cosine similarity between two +unit vectors is their dot product, so normalising at write time turns every +later comparison into a multiply-and-add instead of two square roots per pair. +""" + +from __future__ import annotations + +import logging +import math +from typing import Any + +import httpx + +from lembas.services.llm.openai_client import ( + Endpoint, + LLMError, + describe_http_error, + wrap_transport_error, +) + +log = logging.getLogger(__name__) + +# Longer than a chat request's, because a batch of sixteen chunks against a +# cold local endpoint includes loading the model. +TIMEOUT = 120.0 + + +def normalise(vector: list[float]) -> list[float]: + """A unit vector, or the input unchanged when it has no length. + + A zero vector is what an endpoint returns for empty input, and dividing by + its norm is the one arithmetic error this path can make. It is left as it + is: scoring it against anything gives zero, which is the honest answer. + """ + length = math.sqrt(sum(value * value for value in vector)) + if length <= 0: + return vector + return [value / length for value in vector] + + +def _vectors_in(payload: Any) -> list[list[float]]: + """The embeddings out of a response, whatever shape it arrived in. + + OpenAI's own answer is `{"data": [{"embedding": [...], "index": 0}]}`, and + the index is honoured rather than assumed: nothing in the specification + promises the order, and a provider that sorts differently would silently + pair every chunk with somebody else's vector — which produces a search that + works and is wrong, the worst failure this whole feature can have. + """ + if not isinstance(payload, dict): + raise LLMError("The embedding endpoint returned something unreadable.") + data = payload.get("data") + if not isinstance(data, list) or not data: + raise LLMError("The embedding endpoint returned no vectors.") + + ordered: list[tuple[int, list[float]]] = [] + for position, entry in enumerate(data): + if not isinstance(entry, dict): + raise LLMError("The embedding endpoint returned no vectors.") + raw = entry.get("embedding") + if not isinstance(raw, list) or not raw: + raise LLMError("The embedding endpoint returned an empty vector.") + index = entry.get("index") + at = int(index) if isinstance(index, int) else position + ordered.append((at, [float(value) for value in raw])) + ordered.sort(key=lambda pair: pair[0]) + return [vector for _, vector in ordered] + + +async def embed( + endpoint: Endpoint, model_id: str, texts: list[str], *, timeout: float = TIMEOUT +) -> list[list[float]]: + """One request. Returns a unit vector per input, in the order given. + + Raises `LLMError` for everything -- a missing model, an endpoint that does + not implement embeddings at all, a transport failure -- because every caller + treats them the same way: the index is left as it was and the search falls + back to keywords. Nothing here is worth a partial answer. + """ + if not texts: + return [] + body = {"model": model_id, "input": texts} + try: + async with httpx.AsyncClient(timeout=timeout) as client: + response = await client.post( + endpoint.url("/embeddings"), headers=endpoint.headers(), json=body + ) + # raise_for_status, then translate. `describe_http_error` takes the + # exception rather than the response, which is what every other + # client here hands it. + response.raise_for_status() + payload = response.json() + except httpx.HTTPStatusError as exc: + raise LLMError(describe_http_error(exc)) from exc + except httpx.HTTPError as exc: + raise wrap_transport_error(exc, endpoint) from exc + except ValueError as exc: + raise LLMError("The embedding endpoint did not return JSON.") from exc + + vectors = _vectors_in(payload) + if len(vectors) != len(texts): + # Not recoverable by guessing. A response with fewer vectors than inputs + # would pair chunk three's text with chunk four's vector from there on, + # for the life of the index. + raise LLMError( + f"Asked for {len(texts)} embeddings and got {len(vectors)}." + ) + widths = {len(vector) for vector in vectors} + if len(widths) != 1: + raise LLMError("The embedding endpoint returned vectors of different widths.") + return [normalise(vector) for vector in vectors] + + +async def probe(endpoint: Endpoint, model_id: str) -> int: + """How wide this model's vectors are, by asking for one. + + Used by the admin page's Test button and by nothing on the request path. + There is no endpoint that reports it, so the only honest way to find out is + to embed something. + """ + vectors = await embed(endpoint, model_id, ["lembas"], timeout=60.0) + return len(vectors[0]) + + +__all__ = ["TIMEOUT", "embed", "normalise", "probe"] diff --git a/src/lembas/services/reports.py b/src/lembas/services/reports.py index 5ce0119..e122828 100644 --- a/src/lembas/services/reports.py +++ b/src/lembas/services/reports.py @@ -19,8 +19,8 @@ import logging from sqlalchemy import func, select from sqlalchemy.orm import Session as DBSession -from lembas.db.models import SOURCE_MANUAL, SOURCES, Report, User -from lembas.services.library.fts import search_ids +from lembas.db.models import CHUNK_REPORT, SOURCE_MANUAL, SOURCES, Report, User +from lembas.services.library import retrieval log = logging.getLogger(__name__) @@ -56,14 +56,26 @@ def recent(db: DBSession, user: User | None, *, limit: int = 20) -> list[Report] return list(db.scalars(visible(user).order_by(Report.created_at.desc()).limit(limit))) -def search(db: DBSession, user: User | None, needle: str, *, limit: int = 20) -> list[Report]: +def search( + db: DBSession, + user: User | None, + needle: str, + *, + limit: int = 20, + vector: list[float] | None = None, +) -> list[Report]: """Reports matching `needle`, best match first. Ids come back from FTS and the rows are re-ordered by hit position, exactly as the library stores do -- the index knows about ranking and the ORM query knows about ownership, and neither is asked to do the other's job. + + `vector` is the query already embedded, or None. It comes from the caller + rather than being worked out here because this is synchronous and embedding + is an HTTP request -- see `services/library/retrieval.py`. None means the + keyword search exactly as it always was. """ - hits = search_ids(db, INDEX, needle, limit=limit * 4) + hits = retrieval.search(db, INDEX, needle, kind=CHUNK_REPORT, vector=vector, limit=limit * 4) if not hits: return [] order = {hit.id: position for position, hit in enumerate(hits)} diff --git a/src/lembas/services/settings_store.py b/src/lembas/services/settings_store.py index 4f4c6f5..de342a8 100644 --- a/src/lembas/services/settings_store.py +++ b/src/lembas/services/settings_store.py @@ -33,6 +33,7 @@ IMAGES = "images" SCHEDULES = "schedules" SUBAGENTS = "subagents" BRANDING = "branding" +EXTRACTION = "extraction" def _general_defaults() -> dict[str, Any]: @@ -393,9 +394,93 @@ _DEFAULTS: dict[str, Any] = { # a label and a hint for the admin page and splitting the three across two # modules is how one of them goes stale. BRANDING: lambda: _branding_defaults(), + # A lambda for the same reason BRANDING is one: both factories are + # defined below this table, which is where the accessor that reads each + # group lives. + EXTRACTION: lambda: _extraction_defaults(), } +def _extraction_defaults() -> dict[str, Any]: + """What happens to a file between the upload and the model. + + The numbers were constants in `services/files.py` and every one of them is a + trade somebody with a different corpus makes differently: a 20 MB ceiling is + generous for notes and small for scans, and 120,000 characters is thirty + thousand tokens, which is most of a small window and a rounding error in a + large one. The defaults here are exactly the constants they replace, so an + instance that changes nothing behaves as it always did. + + `Image.MAX_IMAGE_PIXELS` is deliberately **not** here. It is a + decompression-bomb guard, not a preference: a 60,000x60,000 PNG is a few KB + on disk and hundreds of gigabytes decoded, and nobody should be able to + raise that from a form. + """ + return { + "max_upload_mb": 20, + "max_image_edge": 1400, + "jpeg_quality": 85, + "max_pdf_pages": 300, + "max_extracted_chars": 120_000, + "orphan_hours": 24, + # Extensions treated as text beyond the built-in list. Decodability is + # what actually decides, so this only picks a media type -- which is why + # it is a list of extensions rather than a mapping somebody has to get + # right twice. + "extra_text_extensions": [], + # Whether a PDF nothing could read is stored with its error, or refused. + # Keeping it is the default and the honest one: a scanned page is a file + # somebody still wants attached, and the error says why it contributes + # nothing rather than leaving them to wonder. + "reject_unreadable_pdf": False, + # --- Semantic search --------------------------------------------------- + # Which model turns text into vectors. Empty means none, and none means + # the keyword search that has always been here, byte for byte -- which + # is what makes this safe to add to an instance that never asked for it. + "embedding_model_id": "", + # How long a chunk is, in characters, and how much of the previous one + # rides along with it. Characters rather than tokens because the count + # has to be made without asking the endpoint, and the estimate is the + # same four-to-one this codebase already uses. + "chunk_chars": 1200, + "chunk_overlap": 150, + # How many chunks one embedding request carries. Small enough that a + # local endpoint is not asked for a megabyte at once. + "embed_batch": 16, + } + + +def extraction(db: DBSession) -> dict[str, Any]: + """Extraction settings, clamped on read for the reason `agents` gives. + + Every floor here is a number that means something bad at zero: a zero-page + PDF limit extracts nothing from every PDF and reports success, and a + zero-character chunk is an infinite loop in the splitter. + """ + values = get_group(db, EXTRACTION) + values["max_upload_mb"] = min(max(int(values.get("max_upload_mb") or 1), 1), 512) + values["max_image_edge"] = min(max(int(values.get("max_image_edge") or 1), 128), 8192) + values["jpeg_quality"] = min(max(int(values.get("jpeg_quality") or 1), 30), 100) + values["max_pdf_pages"] = min(max(int(values.get("max_pdf_pages") or 1), 1), 5000) + values["max_extracted_chars"] = min( + max(int(values.get("max_extracted_chars") or 1), 1000), 5_000_000 + ) + values["orphan_hours"] = min(max(int(values.get("orphan_hours") or 1), 1), 8760) + values["chunk_chars"] = min(max(int(values.get("chunk_chars") or 1), 200), 8000) + # Bounded *against the chunk*, not absolutely: an overlap at or past the + # chunk size means every chunk starts where the last one did, which is a + # splitter that never advances. + values["chunk_overlap"] = min( + max(int(values.get("chunk_overlap") or 0), 0), values["chunk_chars"] // 2 + ) + values["embed_batch"] = min(max(int(values.get("embed_batch") or 1), 1), 256) + stored = values.get("extra_text_extensions") + values["extra_text_extensions"] = ( + [str(item) for item in stored] if isinstance(stored, list) else [] + ) + return values + + def _branding_defaults() -> dict[str, Any]: """Imported inside the call: `services/branding.py` imports this module for the group key, so a top-level import back is a cycle.""" diff --git a/src/lembas/services/tools.py b/src/lembas/services/tools.py index 6ce8fff..8c2eda1 100644 --- a/src/lembas/services/tools.py +++ b/src/lembas/services/tools.py @@ -417,10 +417,16 @@ async def _run_knowledge_search(context: ToolContext, args: dict[str, Any]) -> T {"name": "knowledge_search", "status": "error", "error": "No query."}, ) + # Embedded before the session opens, because it is an HTTP request and a + # session held across one is the trade `_maybe_compact` already refuses. + # None for every "no" -- no model configured, endpoint down -- and the + # search is then exactly the keyword one it has always been. + vector = await _query_vector(query) + with session_scope() as db: user = db.get(User, context.owner_id) found = documents_service.search( - db, user, query, limit=6, base_ids=context.base_ids + db, user, query, limit=6, base_ids=context.base_ids, vector=vector ) event = { "name": "knowledge_search", @@ -468,13 +474,30 @@ async def _run_knowledge_get(context: ToolContext, args: dict[str, Any]) -> Tool return ToolOutcome(f"{document.title}\n\n{body}", event) +async def _query_vector(query: str) -> list[float] | None: + """The query as a vector, for the stores that can use one. + + Its own session, opened and closed before the caller opens theirs: this is + an HTTP request, and holding a database session across one is the trade + compaction and the project listing both already refuse. + """ + if not query: + return None + from lembas.services.library import retrieval + + with session_scope() as db: + worker = retrieval.worker_for(db) + return await retrieval.embed_with(worker, query) + + # --- Notes ------------------------------------------------------------------- async def _run_notes_search(context: ToolContext, args: dict[str, Any]) -> ToolOutcome: query = str(args.get("query") or "").strip() + vector = await _query_vector(query) with session_scope() as db: user = db.get(User, context.owner_id) found = ( - notes_service.search(db, user, query, limit=8) + notes_service.search(db, user, query, limit=8, vector=vector) if query else notes_service.recent(db, user, limit=8) ) @@ -747,10 +770,11 @@ async def _run_report_write(context: ToolContext, args: dict[str, Any]) -> ToolO async def _run_report_search(context: ToolContext, args: dict[str, Any]) -> ToolOutcome: query = str(args.get("query") or "").strip() + vector = await _query_vector(query) with session_scope() as db: user = db.get(User, context.owner_id) found = ( - reports_service.search(db, user, query, limit=8) + reports_service.search(db, user, query, limit=8, vector=vector) if query else reports_service.recent(db, user, limit=8) ) diff --git a/src/lembas/web/templates/admin/_index_progress.html b/src/lembas/web/templates/admin/_index_progress.html new file mode 100644 index 0000000..ee50c74 --- /dev/null +++ b/src/lembas/web/templates/admin/_index_progress.html @@ -0,0 +1,66 @@ +{% from "_macros.html" import icon %} +{# + The rebuild's state, swapped into itself. + + It polls while a rebuild is running and stops when it is not: `hx-trigger` is + only emitted in the running branch, so the last frame is a plain fragment with + nothing attached to it. A poller that kept going after the work finished would + be a request every two seconds, forever, on a page somebody left open. + + `hx-target="this"` and `hx-swap="outerHTML"` are both spelled out. This lives + inside a form on the page, and htmx resolves `hx-target` by walking up the DOM + -- an element in there that fetches and names no target aims at whatever an + ancestor said, which is the bug the jobs chip had. +#} +
+ + {% if progress.running %} +
+ {{ icon("clock", "icon--sm") }} + + Rebuilding: {{ progress.done }} of {{ progress.total }} records + ({{ progress.percent }}%), {{ progress.written }} pieces written. + You can leave this page — it carries on. + +
+ {% elif progress.error %} +
+ {{ icon("warning", "icon--sm") }} {{ progress.error }} +
+ {% elif progress.total %} +
+ {{ icon("check", "icon--sm") }} + + Finished: {{ progress.done }} records, {{ progress.written }} pieces indexed. + +
+ {% endif %} + +

+ {% if counts %} + Indexed now: + {% for kind, count in counts | dictsort %} + {{ count }} from {{ kind }}s{{ "," if not loop.last }} + {% endfor %}. + {% else %} + Nothing is indexed yet. + {% endif %} +

+ + {% if ready %} + + {% endif %} +
diff --git a/src/lembas/web/templates/admin/_layout.html b/src/lembas/web/templates/admin/_layout.html index 34554ef..3ed22a5 100644 --- a/src/lembas/web/templates/admin/_layout.html +++ b/src/lembas/web/templates/admin/_layout.html @@ -52,6 +52,11 @@ {{ icon("image", "icon--sm") }} Image generation + + {{ icon("file-text", "icon--sm") }} + Extraction + {{ icon("link", "icon--sm") }} Tools diff --git a/src/lembas/web/templates/admin/extraction.html b/src/lembas/web/templates/admin/extraction.html new file mode 100644 index 0000000..84ea330 --- /dev/null +++ b/src/lembas/web/templates/admin/extraction.html @@ -0,0 +1,226 @@ +{% extends "admin/_layout.html" %} +{% from "_macros.html" import icon %} +{% set section = "extraction" %} + +{% block title %}Extraction - {{ brand.name }}{% endblock %} +{% block heading %}Extraction and search{% endblock %} + +{% block admin_content %} +

+ What happens to a file between the upload and the model, and how anything is + found again afterwards. The two are the same pipeline: what is extracted + decides what there is to search. +

+ +{% if saved %} +
{{ icon("check", "icon--sm") }} {{ saved }}
+{% endif %} + +{# --- Extraction ----------------------------------------------------------- #} +
+
+

What a file may cost

+

+ Every number here is a trade, and a large one usually breaks a request + rather than being slow — the text of a whole book does not fit in a + context window, and a model handed it fails the request outright rather + than reading the first half. +

+ +
+
+ + +

Megabytes, before anything is done to it.

+
+ +
+ + +

+ Characters. Roughly four to a token, so 120,000 is about 30,000 tokens + — already most of a small context window. The rest is cut and the + model is told so. +

+
+ +
+ + +

+ Extraction is slow and happens once, at upload. Beyond this the file + is still stored; only its text stops. +

+
+
+ +
+
+ + +

+ Pixels. Images are re-encoded before they are sent, because a phone + photo is several megabytes of base64. +

+
+ +
+ + +
+ +
+ + +

+ Hours. A file chosen in the composer and never sent. Swept at startup. +

+
+
+ +
+ + +

+ One extension per line, .env or env. Only + needed to pick a media type — whether the bytes decode as text is what + actually decides, so an unlisted extension already works. +

+
+ +
+ +

+ Off, a scanned PDF is stored with an explanation saying why it + contributes nothing — there is no OCR here. That is usually what + somebody wants: the file is still attached and still downloadable. +

+
+
+ +
+ +
+
+ +{# --- Semantic search ------------------------------------------------------ #} +
+
+

Searching by meaning

+

+ Keyword search finds a document that uses your words. This finds one that + means the same thing — "how do I get in" reaching a note about + passwords. Both run and the two rankings are fused, so nothing that + keyword search found is lost. +

+ + {% if not models %} +
+ {{ icon("warning", "icon--sm") }} + + No model is marked as producing embeddings. Tick + embeddings on one under + Models — usually a small dedicated model + such as nomic-embed-text or bge-m3, not a chat + model. + +
+ {% endif %} + + {% if missing_model %} +
+ {{ icon("warning", "icon--sm") }} + + {{ missing_model }} is selected but is not available — its + model or its connection has gone, or it is no longer marked for + embeddings. Search has fallen back to keywords. + +
+ {% endif %} + +
+ + +

+ None is not a degraded mode: it is the keyword search + this has always had, with nothing stored and nothing sent anywhere. + Changing the model changes what a vector means, so anything already + indexed is ignored until it is rebuilt. +

+
+ +
+
+ + +

+ Characters. A record is split on paragraph boundaries into pieces of + about this size, and each is embedded separately — a document is found + by its best piece, not by its average. +

+
+ +
+ + +

+ How much of each piece is repeated at the start of the next, so a + sentence across a boundary is whole somewhere. Capped at half the + piece size. +

+
+ +
+ + +

+ Lower this if the endpoint refuses large requests; raise it if a + rebuild is slow and the far side has room. +

+
+
+
+ +
+ +
+
+ +{# --- The index ------------------------------------------------------------ #} +
+

The index

+

+ Documents, notes, skills and reports are indexed as they are written. A + rebuild is for everything that already existed — or for after changing the + model or the piece size, both of which make what is stored stop meaning + anything. It runs in the background and can be left. +

+ + {% include "admin/_index_progress.html" %} +
+{% endblock %} diff --git a/src/lembas/web/templates/admin/model_detail.html b/src/lembas/web/templates/admin/model_detail.html index 78a20c4..b4b4abc 100644 --- a/src/lembas/web/templates/admin/model_detail.html +++ b/src/lembas/web/templates/admin/model_detail.html @@ -172,6 +172,9 @@ vision lets images be sent, and tools is whether a tool list may be sent at all — turn it on for a model that does not support tool calling and every one of its replies fails. + embeddings is the odd one out: it says this is not a chat + model at all, but one that turns text into vectors, and it is what + Extraction picks from.

{% for name in capabilities %} diff --git a/tests/conftest.py b/tests/conftest.py index 4ceae21..5ff3e43 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -72,19 +72,22 @@ def fresh_database(tmp_path: Path) -> Iterator[None]: @pytest.fixture(autouse=True) -def fresh_branding() -> Iterator[None]: - """Drop the branding snapshot between tests. +def fresh_snapshots() -> Iterator[None]: + """Drop the process-level snapshots between tests. - It is a process-level cache read by a Jinja global, so without this the - first test to render a page pins one instance's name, logo and themes for - every test after it -- against a database that has since been thrown away. - The same shape as the registries below, and the reason each of them exists. + Two of them now, and both are read once per process against a database this + fixture throws away between tests -- so without this, the first test to + render a page pins one instance's name and themes for every test after it, + and the first to save an upload limit pins that too. The same shape as the + registries below, and the reason each of them exists. """ - from lembas.services import branding + from lembas.services import branding, files branding.forget() + files.forget() yield branding.forget() + files.forget() @pytest.fixture(autouse=True) diff --git a/tests/test_embeddings.py b/tests/test_embeddings.py new file mode 100644 index 0000000..92ea0af --- /dev/null +++ b/tests/test_embeddings.py @@ -0,0 +1,185 @@ +"""Turning text into vectors, and the pieces text is cut into first. + +The failure this file is mostly about is the quiet one: a response paired with +the wrong input produces a search that works and is wrong, which nothing +downstream can notice. +""" + +from __future__ import annotations + +import math + +import httpx +import pytest + +from lembas.services.library import chunks +from lembas.services.llm import embeddings +from lembas.services.llm.openai_client import Endpoint, LLMError + +ENDPOINT = Endpoint(base_url="http://127.0.0.1:9", api_key="", extra_headers={}, name="test") + + +def _answer(vectors, *, shuffle=False): + data = [{"embedding": vector, "index": i} for i, vector in enumerate(vectors)] + if shuffle: + data = list(reversed(data)) + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, json={"data": data}) + + return handler + + +# --- The client ----------------------------------------------------------------- +async def test_vectors_come_back_normalised(mock_http): + """Cosine between unit vectors is their dot product, so normalising once at + write time turns every later comparison into a multiply-and-add.""" + mock_http(_answer([[3.0, 4.0]])) + + (vector,) = await embeddings.embed(ENDPOINT, "m", ["hello"]) + + assert math.isclose(vector[0], 0.6) + assert math.isclose(vector[1], 0.8) + + +async def test_the_declared_index_decides_the_order(mock_http): + """Nothing in the specification promises the order of `data`. A provider + that sorts differently would pair every chunk with somebody else's vector — + a search that works and is wrong, which is the worst failure here.""" + mock_http(_answer([[1.0, 0.0], [0.0, 1.0]], shuffle=True)) + + first, second = await embeddings.embed(ENDPOINT, "m", ["a", "b"]) + + assert first == [1.0, 0.0] + assert second == [0.0, 1.0] + + +async def test_a_short_answer_is_refused_rather_than_guessed(mock_http): + """Silently accepting it would pair chunk three's text with chunk four's + vector from there on, for the life of the index.""" + mock_http(_answer([[1.0, 0.0]])) + + with pytest.raises(LLMError) as caught: + await embeddings.embed(ENDPOINT, "m", ["a", "b"]) + + assert "2" in str(caught.value) + + +async def test_vectors_of_different_widths_are_refused(mock_http): + mock_http(_answer([[1.0, 0.0], [0.0, 1.0, 0.0]])) + + with pytest.raises(LLMError): + await embeddings.embed(ENDPOINT, "m", ["a", "b"]) + + +async def test_an_endpoint_that_does_not_do_embeddings_says_so(mock_http): + def handler(request): + return httpx.Response(404, json={"error": {"message": "no such endpoint"}}) + + mock_http(handler) + + with pytest.raises(LLMError): + await embeddings.embed(ENDPOINT, "m", ["a"]) + + +async def test_nothing_asked_is_nothing_sent(mock_http): + """A request with no inputs is a round trip for nothing, and some endpoints + refuse it outright.""" + sent = [] + + def handler(request): + sent.append(request) + return httpx.Response(200, json={"data": []}) + + mock_http(handler) + + assert await embeddings.embed(ENDPOINT, "m", []) == [] + assert sent == [] + + +def test_a_zero_vector_survives_normalising(): + """What an endpoint returns for empty input, and the one arithmetic error + this path can make.""" + assert embeddings.normalise([0.0, 0.0]) == [0.0, 0.0] + + +# --- Splitting ------------------------------------------------------------------ +def test_short_text_is_one_piece(): + assert chunks.split("a short note", size=1200) == ["a short note"] + + +def test_splitting_prefers_a_paragraph_boundary(): + body = "\n\n".join(["A" * 400, "B" * 400, "C" * 400]) + pieces = chunks.split(body, size=900, overlap=0) + + assert len(pieces) >= 2 + # No piece ends mid-run, which is what a boundary being honoured looks like. + assert all(piece.strip() == piece for piece in pieces) + assert pieces[0].startswith("A") + + +def test_the_overlap_carries_the_tail_forward(): + body = "".join(f"sentence {n}. " for n in range(200)) + pieces = chunks.split(body, size=400, overlap=100) + + assert len(pieces) > 1 + # Consecutive pieces share text, which is what stops a sentence across a + # boundary being absent from both embeddings. + assert any(pieces[0][-40:] in pieces[1] for _ in (0,)) or pieces[1][:40] in body + + +def test_an_overlap_as_large_as_the_piece_does_not_hang(): + """An overlap at or past the size means every piece starts where the last + one did. Clamped here as well as in the settings accessor, because the + failure is a hang rather than a bad index.""" + pieces = chunks.split("x" * 5000, size=400, overlap=4000) + + assert len(pieces) < 40 + + +def test_a_scrap_is_not_worth_a_row(): + """The embedding of six words is mostly noise, and a search whose best hit + is "and the following:" is worse than one that returns nothing.""" + assert chunks.split("hi") == ["hi"] + pieces = chunks.split("A" * 400 + "\n\n" + "B" * 5, size=380, overlap=0) + assert all(len(piece) >= chunks.MIN_CHUNK_CHARS for piece in pieces) + + +def test_nothing_in_is_nothing_out(): + assert chunks.split("") == [] + assert chunks.split(" \n\n ") == [] + + +# --- Packing -------------------------------------------------------------------- +def test_a_vector_survives_a_round_trip(): + vector = [0.5, -0.25, 0.125] + blob = chunks.pack(vector) + + assert len(blob) == 12 + assert chunks.unpack(blob, 3) == pytest.approx(vector) + + +def test_a_truncated_blob_unpacks_to_nothing(): + """Inferring the width from the length would let a short BLOB unpack into a + shorter vector and score against a query happily — a wrong answer rather + than a missing one.""" + assert chunks.unpack(chunks.pack([1.0, 2.0, 3.0])[:8], 3) == [] + + +def test_scoring_across_widths_is_zero_rather_than_an_error(): + """It means the two vectors came from different models, and the honest + answer to "how similar are these?" across two spaces is nothing.""" + assert chunks.dot([1.0, 0.0], [1.0, 0.0, 0.0]) == 0.0 + + +def test_the_dot_product_is_the_cosine_for_unit_vectors(): + a = embeddings.normalise([1.0, 1.0]) + b = embeddings.normalise([1.0, 0.0]) + + assert chunks.dot(a, a) == pytest.approx(1.0) + assert chunks.dot(a, b) == pytest.approx(math.sqrt(0.5)) + + +def test_the_hash_changes_with_the_text(): + assert chunks.digest("a") != chunks.digest("b") + assert chunks.digest("a") == chunks.digest("a") diff --git a/tests/test_files.py b/tests/test_files.py index c75cb80..47dc154 100644 --- a/tests/test_files.py +++ b/tests/test_files.py @@ -188,12 +188,15 @@ def test_empty_files_are_rejected(): def test_oversized_files_are_rejected(): with pytest.raises(files_service.FileError) as caught: - files_service.prepare(b"x" * (files_service.MAX_UPLOAD_BYTES + 1), "huge.txt") + files_service.prepare(b"x" * (files_service.limits().max_upload_bytes + 1), "huge.txt") assert "MB" in str(caught.value) def test_extracted_text_is_capped(monkeypatch): - monkeypatch.setattr(files_service, "MAX_EXTRACTED_CHARS", 50) + """Through the setting rather than the constant. The constant is only the + default now; what `prepare` reads is the snapshot, which is the thing that + would have gone on returning 120,000 if the wiring were wrong.""" + monkeypatch.setattr(files_service, "_LIMITS", files_service.Limits(max_extracted_chars=50)) prepared = files_service.prepare(b"x" * 500, "long.txt") assert len(prepared.extracted_text) == 50 assert prepared.truncated is True @@ -446,7 +449,7 @@ def test_documents_reach_a_model_without_vision(client: TestClient, db, chat_wit def test_truncation_is_declared_to_the_model(client: TestClient, db, chat_with_model, monkeypatch): """A model asked about page 400 should be able to say it did not see it.""" - monkeypatch.setattr(files_service, "MAX_EXTRACTED_CHARS", 20) + monkeypatch.setattr(files_service, "_LIMITS", files_service.Limits(max_extracted_chars=20)) client.post("/api/files", files={"file": ("big.txt", b"y" * 200, "text/plain")}) attachment = db.scalar(select(Attachment)) client.post( @@ -600,3 +603,84 @@ def test_a_browser_serialising_the_form_actually_sends_the_attachment( content = turns(chat_service.build_request(db, chat))[0]["content"] assert isinstance(content, list), "the image never reached the model" assert any(p.get("type") == "image_url" for p in content) + + +# --- The extraction settings --------------------------------------------------- +# The constants above are defaults now, and what `prepare` actually reads is a +# process-level snapshot. Every one of these failures would be silent: a limit +# that looks configured and is not. +def test_a_saved_limit_reaches_the_snapshot(db, client, registered): + from lembas.services import settings_store + + client.post( + "/admin/extraction", + data={ + "max_upload_mb": "5", + "max_image_edge": "800", + "jpeg_quality": "70", + "max_pdf_pages": "10", + "max_extracted_chars": "2000", + "orphan_hours": "3", + "extra_text_extensions": "env\n.conf", + }, + follow_redirects=False, + ) + + bounds = files_service.limits() + assert bounds.max_upload_bytes == 5 * 1024 * 1024 + assert bounds.max_extracted_chars == 2000 + assert bounds.max_image_edge == 800 + assert settings_store.extraction(db)["extra_text_extensions"] == ["env", ".conf"] + + +def test_an_extra_extension_gets_a_leading_dot(db, client, registered): + """Typed both ways by different people, and a mapping somebody has to get + right twice is one they get wrong once.""" + client.post( + "/admin/extraction", data={"extra_text_extensions": "env"}, follow_redirects=False + ) + + assert files_service.limits().media_type_for(".env") == "text/plain" + + +def test_an_unknown_extension_is_still_stored_as_text(db): + """Decodability is what decides. The list only picks a media type, which is + why an unlisted extension has always worked and must go on working.""" + prepared = files_service.prepare(b"hello there", "notes.wat") + assert prepared.kind == "text" + assert prepared.extension == ".txt" + + +def test_a_number_out_of_range_is_clamped(db, client, registered): + client.post( + "/admin/extraction", + data={"max_upload_mb": "99999", "jpeg_quality": "1"}, + follow_redirects=False, + ) + + bounds = files_service.limits() + assert bounds.max_upload_bytes == 512 * 1024 * 1024 + assert bounds.jpeg_quality == 30 + + +def test_the_snapshot_is_dropped_when_the_page_saves(db, client, registered): + """Read once per process. A save that did not drop it would take effect at + the next restart, which is the failure this codebase keeps cataloguing.""" + assert files_service.limits().max_upload_bytes == 20 * 1024 * 1024 + + client.post("/admin/extraction", data={"max_upload_mb": "1"}, follow_redirects=False) + + assert files_service.limits().max_upload_bytes == 1024 * 1024 + + +def test_only_an_administrator_may_change_extraction(client, registered): + client.post("/auth/logout", follow_redirects=False) + client.post( + "/auth/register", + data={"name": "Sam", "email": "sam@shire.test", "password": "potatoes-po-ta-toes"}, + follow_redirects=False, + ) + + assert client.post("/admin/extraction", data={"max_upload_mb": "1"}).status_code == 403 + assert client.post("/admin/extraction/search", data={}).status_code == 403 + assert client.post("/admin/extraction/rebuild", data={}).status_code == 403 diff --git a/tests/test_library_hybrid.py b/tests/test_library_hybrid.py new file mode 100644 index 0000000..24c2c4c --- /dev/null +++ b/tests/test_library_hybrid.py @@ -0,0 +1,345 @@ +"""Keyword search, semantic search, and the two fused. + +The property this file exists to hold is the first one: **with no embedding +model configured, everything here is byte-for-byte the search that has always +been.** That is what makes this safe to add to an instance that never asked for +it, and it is the one claim a comment cannot make credible. +""" + +from __future__ import annotations + +import httpx +import pytest +from sqlalchemy import select + +from lembas.db.models import CHUNK_NOTE, Chunk, Connection, Model, User +from lembas.services import settings_store +from lembas.services.crypto import encrypt +from lembas.services.library import chunks as chunk_service +from lembas.services.library import indexing, retrieval +from lembas.services.library import notes as notes_service +from lembas.services.library.fts import SearchHit + + +@pytest.fixture(autouse=True) +def clean(): + indexing.clear() + yield + indexing.clear() + + +@pytest.fixture +def user(db, registered) -> User: + return db.scalars(select(User).order_by(User.created_at)).first() + + +@pytest.fixture +def embedding_model(db): + """A connection and a model marked for embeddings, and the setting pointing + at it. Nothing here makes a request; the tests that need one mock it.""" + connection = Connection( + name="Embed", base_url="http://127.0.0.1:9", api_key_encrypted=encrypt("") + ) + db.add(connection) + db.commit() + db.add( + Model( + connection_id=connection.id, + model_id="embed-1", + capabilities_json={"embeddings": True}, + ) + ) + db.commit() + settings_store.update(db, {"embedding_model_id": "embed-1"}, key=settings_store.EXTRACTION) + return "embed-1" + + +def _vector_answer(vector): + def handler(request: httpx.Request) -> httpx.Response: + count = len(request.read().decode().split('"input"')[1].split("[")[1].split(",")) + return httpx.Response( + 200, json={"data": [{"embedding": vector, "index": i} for i in range(count)]} + ) + + return handler + + +def _store_chunk(db, note, vector, *, text="stored"): + db.add( + Chunk( + owner_id=note.owner_id, + resource_type=CHUNK_NOTE, + resource_id=note.id, + ordinal=0, + text=text, + vector=chunk_service.pack(vector), + dims=len(vector), + model_id="embed-1", + source_hash=chunk_service.digest(text), + ) + ) + db.commit() + + +# --- Nothing configured --------------------------------------------------------- +def test_with_no_model_nothing_is_indexed(db, user): + """No chunk rows, no requests, no cost. Asserted rather than assumed: it is + the whole reason this is safe to add to an existing instance.""" + assert indexing.enabled(db) is False + notes_service.create(db, owner=user, title="Moria", body="the west gate") + + assert db.scalars(select(Chunk)).all() == [] + + +def test_with_no_model_search_is_exactly_the_keyword_search(db, user): + from lembas.services.library import fts + + notes_service.create(db, owner=user, title="Moria", body="the west gate opens") + notes_service.create(db, owner=user, title="Bree", body="an inn on the road") + + through_retrieval = retrieval.search(db, "notes_fts", "gate", kind=CHUNK_NOTE, limit=10) + direct = fts.search_ids(db, "notes_fts", "gate", limit=10) + + assert [hit.id for hit in through_retrieval] == [hit.id for hit in direct] + + +async def test_with_no_model_a_query_embeds_to_nothing(db, user): + assert await retrieval.embed_query(db, "anything") is None + + +# --- Scoring -------------------------------------------------------------------- +def test_a_record_scores_as_its_best_piece(db, user, embedding_model): + """One paragraph that answers the question is what makes a document worth + returning. Averaging would rank a long document about something else above a + short one that says exactly the thing.""" + near = notes_service.create(db, owner=user, title="Near", body="x") + far = notes_service.create(db, owner=user, title="Far", body="y") + # Two pieces for `near`, one of them irrelevant. The good one has to win. + _store_chunk(db, near, [1.0, 0.0]) + db.add( + Chunk( + owner_id=user.id, + resource_type=CHUNK_NOTE, + resource_id=near.id, + ordinal=1, + text="unrelated", + vector=chunk_service.pack([0.0, 1.0]), + dims=2, + model_id="embed-1", + ) + ) + _store_chunk(db, far, [0.7, 0.714]) + db.commit() + + hits = retrieval.semantic_ids(db, CHUNK_NOTE, [1.0, 0.0], limit=10) + + assert [hit.id for hit in hits][0] == near.id + + +def test_a_vector_from_another_model_is_skipped(db, user, embedding_model): + """A change of embedding model with a rebuild still pending. Scoring across + two spaces produces a confident wrong answer rather than a missing one.""" + note = notes_service.create(db, owner=user, title="Old", body="x") + _store_chunk(db, note, [1.0, 0.0, 0.0]) # three wide + + assert retrieval.semantic_ids(db, CHUNK_NOTE, [1.0, 0.0], limit=10) == [] + + +def test_a_semantic_only_match_is_found(db, user, embedding_model): + """The point of the whole feature: a record whose words do not appear in the + query at all.""" + note = notes_service.create(db, owner=user, title="Doors", body="mellon") + _store_chunk(db, note, [1.0, 0.0]) + + keyword = retrieval.search(db, "notes_fts", "how do I get in", kind=CHUNK_NOTE, limit=5) + hybrid = retrieval.search( + db, "notes_fts", "how do I get in", kind=CHUNK_NOTE, vector=[1.0, 0.0], limit=5 + ) + + assert keyword == [] + assert [hit.id for hit in hybrid] == [note.id] + + +# --- Fusion --------------------------------------------------------------------- +def test_fusion_keeps_what_only_one_side_found(): + """Neither ranking's finds are dropped. That is the property that makes + turning this on unable to make search worse.""" + keyword = [SearchHit(id="a", rank=-1.0), SearchHit(id="b", rank=-2.0)] + meaning = [SearchHit(id="c", rank=0.9)] + + fused = {hit.id for hit in retrieval.fuse(keyword, meaning, limit=10)} + + assert fused == {"a", "b", "c"} + + +def test_agreeing_on_a_record_ranks_it_first(): + """RRF's whole behaviour: both lists having it beats either list alone.""" + keyword = [SearchHit(id="a", rank=-1.0), SearchHit(id="both", rank=-2.0)] + meaning = [SearchHit(id="both", rank=0.9), SearchHit(id="c", rank=0.8)] + + assert retrieval.fuse(keyword, meaning, limit=10)[0].id == "both" + + +def test_fusing_with_nothing_is_the_other_list(): + keyword = [SearchHit(id="a", rank=-1.0), SearchHit(id="b", rank=-2.0)] + + assert [hit.id for hit in retrieval.fuse(keyword, [], limit=10)] == ["a", "b"] + + +def test_the_fused_score_is_larger_for_better(): + """The opposite of bm25's convention, which is worth saying out loud rather + than leaving somebody to infer it from a number that stopped being + negative.""" + fused = retrieval.fuse([SearchHit(id="a", rank=-1.0)], [SearchHit(id="a", rank=0.9)]) + + assert fused[0].rank > 0 + + +# --- Indexing ------------------------------------------------------------------- +async def test_a_record_is_chunked_and_stored(db, user, embedding_model, mock_http): + mock_http(_vector_answer([1.0, 0.0])) + note = notes_service.create(db, owner=user, title="Moria", body="the west gate " * 200) + + written = await indexing.index_resource(CHUNK_NOTE, note.id) + + db.expire_all() + rows = list(db.scalars(select(Chunk).where(Chunk.resource_id == note.id))) + assert written == len(rows) > 1 + assert all(row.dims == 2 and row.model_id == "embed-1" for row in rows) + assert [row.ordinal for row in rows] == list(range(len(rows))) + + +async def test_indexing_an_unchanged_record_is_free(db, user, embedding_model, mock_http): + """The hash is what makes it free, and what makes "is this current?" + answerable without embedding anything.""" + requests = [] + + def handler(request): + requests.append(request) + return httpx.Response(200, json={"data": [{"embedding": [1.0, 0.0], "index": 0}]}) + + mock_http(handler) + note = notes_service.create(db, owner=user, title="Moria", body="short") + await indexing.index_resource(CHUNK_NOTE, note.id) + before = len(requests) + + await indexing.index_resource(CHUNK_NOTE, note.id) + + assert len(requests) == before + + +async def test_re_indexing_twice_is_idempotent(db, user, embedding_model, mock_http): + mock_http(_vector_answer([1.0, 0.0])) + note = notes_service.create(db, owner=user, title="Moria", body="short") + + await indexing.index_resource(CHUNK_NOTE, note.id, force=True) + await indexing.index_resource(CHUNK_NOTE, note.id, force=True) + + db.expire_all() + assert len(list(db.scalars(select(Chunk).where(Chunk.resource_id == note.id)))) == 1 + + +async def test_a_failed_embedding_leaves_the_old_chunks(db, user, embedding_model, mock_http): + """Deleting first and failing half way through would leave a record indexed + by half of itself, which ranks worse than not being indexed and looks like + nothing at all.""" + mock_http(_vector_answer([1.0, 0.0])) + note = notes_service.create(db, owner=user, title="Moria", body="short") + await indexing.index_resource(CHUNK_NOTE, note.id) + + mock_http(lambda request: httpx.Response(500, json={"error": "down"})) + notes_service.update(db, note, body="something else entirely") + await indexing.index_resource(CHUNK_NOTE, note.id) + + db.expire_all() + assert len(list(db.scalars(select(Chunk).where(Chunk.resource_id == note.id)))) == 1 + + +async def test_a_record_with_no_text_left_drops_its_chunks( + db, user, embedding_model, mock_http, monkeypatch +): + """A guard rather than a state anything reaches today: every one of the four + stores requires a title, so `text_of` is never empty for a record that + exists. It is driven anyway, because the alternative to a guard here is a + record whose chunks outlive its content -- and a store whose title becomes + optional is a change nobody would think to test this against. + """ + mock_http(_vector_answer([1.0, 0.0])) + note = notes_service.create(db, owner=user, title="Moria", body="short") + await indexing.index_resource(CHUNK_NOTE, note.id) + + monkeypatch.setattr(indexing, "text_of", lambda row: "") + await indexing.index_resource(CHUNK_NOTE, note.id) + + db.expire_all() + assert list(db.scalars(select(Chunk).where(Chunk.resource_id == note.id))) == [] + + +async def test_a_deleted_record_takes_its_chunks_with_it(db, user, embedding_model, mock_http): + mock_http(_vector_answer([1.0, 0.0])) + note = notes_service.create(db, owner=user, title="Moria", body="short") + await indexing.index_resource(CHUNK_NOTE, note.id) + note_id = note.id + + notes_service.delete(db, note) + await indexing.index_resource(CHUNK_NOTE, note_id) + + db.expire_all() + assert list(db.scalars(select(Chunk).where(Chunk.resource_id == note_id))) == [] + + +def test_the_sweep_catches_what_had_no_loop_to_clean_up(db, user, embedding_model): + """The backstop for a delete with no event loop running -- a CLI command, or + a cascade from removing an account.""" + note = notes_service.create(db, owner=user, title="Moria", body="short") + _store_chunk(db, note, [1.0, 0.0]) + note_id = note.id + notes_service.delete(db, note) + + assert indexing.sweep_orphans(db) == 1 + assert list(db.scalars(select(Chunk).where(Chunk.resource_id == note_id))) == [] + + +def test_the_text_of_a_record_is_the_indexed_columns(db, user): + note = notes_service.create(db, owner=user, title="Moria", body="the west gate") + + body = indexing.text_of(note) + + assert "Moria" in body and "west gate" in body + assert indexing.kind_of(note) == CHUNK_NOTE + + +# --- The picker ----------------------------------------------------------------- +def test_a_chat_model_is_not_offered_as_an_embedder(db, user, embedding_model, client): + connection = db.scalars(select(Connection)).first() + db.add( + Model( + connection_id=connection.id, + model_id="chat-1", + capabilities_json={"tools": True}, + ) + ) + db.commit() + + page = client.get("/admin/extraction").text + + assert 'value="embed-1"' in page + assert 'value="chat-1"' not in page + + +def test_a_model_that_lost_its_flag_is_named(db, user, embedding_model, client): + """A setting that vanishes from a picker is one nobody can tell from a + setting that was never made.""" + model = db.scalar(select(Model).where(Model.model_id == "embed-1")) + model.capabilities_json = {} + db.commit() + + assert "embed-1" in client.get("/admin/extraction").text + + +def test_a_deleted_model_means_no_embedder_rather_than_an_error(db, user, embedding_model): + db.delete(db.scalar(select(Model).where(Model.model_id == "embed-1"))) + db.commit() + + assert indexing.embedder(db) is None + assert indexing.enabled(db) is False