Finding a thing that does not use your words
Three pieces, and the first one is that they are all optional. Extraction stops being constants. Upload size, image edge, JPEG quality, PDF pages, extracted characters, orphan age and the text-extension list are settings now, read through a process-level snapshot rather than a session -- `prepare` and everything under it are called from routes, tool runners and the startup sweep, and several of those have no session in hand. Two things deliberately stayed constants: the decompression-bomb guard, which is a guard and not a preference, and ORPHAN_AGE, which would have been evaluated at import if it stayed in the signature and pinned the shipped 24 hours whatever anybody set. An embedding model is picked from the models an administrator flagged for it, and one that has since lost its flag is *named* rather than dropped from the picker: a setting that vanishes is one nobody can tell from a setting never made. Nothing here is required. Choosing none means no chunk rows, no requests, and retrieval.search returning exactly what fts.search_ids returns in exactly that order -- asserted, because it is what makes this safe to land on an instance that never asked for it. The two rankings are fused by reciprocal rank fusion: ranks and not scores, because bm25 is a corpus-dependent negative and cosine is 0..1, and normalising them onto one scale means picking a constant nobody can tune without a labelled set they do not have. RRF's one constant is famously insensitive and degrades to whichever list is non-empty -- which is what turns "no embedding model" into a branch that does not exist. A record scores as its best chunk rather than its average, or a long document about something else outranks a short one that says the thing. Width and model are stored beside every vector and a mismatch is skipped, because vectors from two spaces score against each other perfectly happily and mean nothing -- a search that works and is wrong is the worst failure this can have, and a model change now leaves stale rows ignored rather than trusted. Indexing is fired and forgotten, and how a change is noticed is a session event rather than a call in each of the ten library writers. That is a departure from this codebase's taste for explicit seams, for the reason tool_label is a Jinja global: a step every writer has to remember is one that gets forgotten, and here forgetting is silent -- the record saves, keyword search still finds it, and only its recall goes stale. Chunks are embedded before anything is deleted, so a failure leaves the old index rather than half a new one. Also: `embeddings` joins the model capabilities, and the three tool flags that had shipped with no checkbox -- canvas, scheduling and helpers -- have one. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Jaroslav Beneš
co-authored by
Claude Opus 5
Claude Opus 5
parent
78e5717f77
commit
20bb569b00
@@ -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"]
|
||||
@@ -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 []
|
||||
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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)}
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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)}
|
||||
|
||||
Reference in New Issue
Block a user