Add knowledgebases

Knowledgebases give bots RAG over documents, scoped like extensions: to particular bots, to a workspace, or to every bot, with their own button beside Tools in the composer. Documents are chunked and embedded by the file-ingest worker with the install's embedding model, which is now an install-wide setting rather than a per-bot one. - Word search runs on an FTS5 index ranked by BM25, with a scored LIKE fallback for scripts written without spaces. - Changing the embedding model re-indexes documents automatically. - Ingest jobs wait without spending attempts while every embedding model rests, fail at once on errors retrying can't fix, and keep passages searchable by words when the model refuses. - Settings > Knowledgebases shows each knowledgebase's and document's index state: spinner, warning, grey check (words only) or green check. Co-Authored-By: Claude Opus 5.5 <[email protected]>
2026-10-26 16:41:01 +01:00
parent 2259b6895f
commit e7b23b7454
111 changed files with 4573 additions and 148 deletions
@@ -0,0 +1,297 @@
+import BunSQLite from "@moduletrace/bun-sqlite";
+import type { BUN_SQLITE_COBBLE_KNOWLEDGEBASE_CHUNKS, BunSQLiteTables } from "@/db/types/db";
+import { AppData } from "@/src/data/app-data";
+import embedText from "@/src/harness/embeddings/embed-text";
+import HarnessError from "@/src/harness/harness-error";
+import runWithFallbackScope from "@/src/harness/usage/run-with-fallback-scope";
+import chunkText from "../bot-files/chunk-text";
+import ensureKnowledgebaseFts from "./fts/ensure-knowledgebase-fts";
+import grabKnowledgebaseDocumentText from "./grab-knowledgebase-document-text";
+
+type Params = {
+    document_id: number | string;
+    signal?: AbortSignal;
+};
+
+export type KnowledgebaseIngestResult = {
+    /**
+     * `ingested`, `empty` (no text to read), `skipped` (too large) or
+     * `words` (no embedding model: searchable by words only).
+     */
+    status: "ingested" | "empty" | "skipped" | "words" | "missing";
+    chunks: number;
+};
+
+/**
+ * Reads a knowledgebase document's text, splits it into chunks and embeds
+ * them with the install's embedding model (never a bot's: every vector in
+ * a knowledgebase must be comparable). The old passages stay searchable
+ * until the new ones are in. Without an embedding model the chunks are
+ * still kept, and indexed for ranked word matching
+ * (`knowledgebase_chunks_fts`).
+ *
+ * Failures are recorded on the document (`ingest_error`) and rethrown, so
+ * the queue can retry.
+ */
+export default async function ingestKnowledgebaseDocument({
+    document_id,
+    signal,
+}: Params): Promise<KnowledgebaseIngestResult> {
+    const document = (
+        await BunSQLite.sql<{
+            id: number;
+            knowledgebase_id: number;
+            name: string;
+            source: string;
+            media_id: number | null;
+            content: string | null;
+            mime_type: string | null;
+            size: number;
+            user_id: number | null;
+        }>({
+            sql: `SELECT id, knowledgebase_id, name, source, media_id, content, mime_type, size, user_id
+                FROM knowledgebase_documents WHERE id = ? LIMIT 1`,
+            values: [document_id],
+        })
+    ).payload?.[0];
+
+    if (!document?.id) return { status: "missing", chunks: 0 };
+
+    if (Number(document.size) > AppData.KnowledgebaseDocumentMaxBytes) {
+        await finish({
+            document_id: document.id,
+            chunks: 0,
+            model: null,
+            error: `Too large to read (over ${AppData.KnowledgebaseDocumentMaxBytes / (1024 * 1024)} MB).`,
+        });
+
+        return { status: "skipped", chunks: 0 };
+    }
+
+    /**
+     * The ingest queue has no run behind it: the embedding is charged to
+     * whoever added the document, so it shows on the usage page.
+     */
+    const scope = document.user_id ? { user_id: document.user_id } : null;
+    /**
+     * Set once the document's outcome is written, so the catch below
+     * doesn't write over it.
+     */
+    let recorded = false;
+
+    try {
+        const text = await grabKnowledgebaseDocumentText({ document });
+
+        const old_ids = await chunkIds({ document_id: document.id });
+        const chunks = text.trim() ? chunkText({ text }) : [];
+
+        if (!chunks.length) {
+            await dropChunks({ ids: old_ids });
+            await finish({ document_id: document.id, chunks: 0, model: null, error: "It has no text to read." });
+
+            return { status: "empty", chunks: 0 };
+        }
+
+        const new_ids = await insertChunks({ document, chunks });
+        let model: string | null = null;
+
+        try {
+            for (let start = 0; start < chunks.length; start += AppData.EmbeddingBatchSize) {
+                const batch = chunks.slice(start, start + AppData.EmbeddingBatchSize);
+                const embedded = await runWithFallbackScope({
+                    scope,
+                    task: () => embedText({ input: batch, signal }),
+                });
+
+                model = embedded.model || model;
+
+                for (let i = 0; i < embedded.vectors.length; i++) {
+                    await insertVector({
+                        vector: embedded.vectors[i]!,
+                        chunk_id: new_ids[start + i]!,
+                        document_id: document.id,
+                        knowledgebase_id: document.knowledgebase_id,
+                    });
+                }
+            }
+        } catch (error) {
+            if (!(error instanceof HarnessError) || error.resting) {
+                /**
+                 * Not the model's answer (a bug, a database error), or no
+                 * model was tried yet: the old passages are still whole,
+                 * the new ones go, and the job tries again.
+                 */
+                await dropChunks({ ids: new_ids });
+                throw error;
+            }
+
+            /**
+             * No embedding model, or it refused: the new passages stay,
+             * found by their words, and none of their vectors are kept
+             * (half a document embedded would rank unevenly).
+             */
+            await dropVectors({ ids: new_ids });
+
+            if (error.code != "unavailable") {
+                await dropChunks({ ids: old_ids });
+                await finish({
+                    document_id: document.id,
+                    chunks: chunks.length,
+                    model: null,
+                    error: `Found by its words only: the embedding model failed. ${error.message}`.slice(0, 500),
+                });
+                recorded = true;
+
+                throw error;
+            }
+
+            model = null;
+        }
+
+        await dropChunks({ ids: old_ids });
+        await finish({ document_id: document.id, chunks: chunks.length, model });
+
+        return { status: model ? "ingested" : "words", chunks: chunks.length };
+    } catch (error) {
+        /**
+         * Every model resting: nothing was tried, the job waits for them
+         * and the document keeps what it said before.
+         */
+        if (recorded || (error instanceof HarnessError && error.resting)) throw error;
+
+        const message = error instanceof Error ? error.message : String(error);
+
+        await BunSQLite.sql({
+            sql: `UPDATE knowledgebase_documents SET ingest_error = ?, updated_at = ? WHERE id = ?`,
+            values: [message.slice(0, 500), Date.now(), document.id],
+        }).catch(() => {});
+
+        throw error;
+    }
+}
+
+async function chunkIds({ document_id }: { document_id: number }) {
+    const res = await BunSQLite.sql<{ id: number }>({
+        sql: `SELECT id FROM knowledgebase_chunks WHERE document_id = ?`,
+        values: [document_id],
+    });
+
+    return (res.payload || []).map((row) => Number(row.id));
+}
+
+async function dropVectors({ ids }: { ids: number[] }) {
+    for (let start = 0; start < ids.length; start += 500) {
+        const batch = ids.slice(start, start + 500);
+
+        await BunSQLite.sql({
+            sql: `DELETE FROM knowledgebase_vectors WHERE chunk_id IN (${batch.map(() => "?").join(", ")})`,
+            values: batch,
+        });
+    }
+}
+
+async function dropChunks({ ids }: { ids: number[] }) {
+    for (let start = 0; start < ids.length; start += 500) {
+        const batch = ids.slice(start, start + 500);
+        const placeholders = batch.map(() => "?").join(", ");
+
+        await BunSQLite.sql({ sql: `DELETE FROM knowledgebase_vectors WHERE chunk_id IN (${placeholders})`, values: batch });
+        await BunSQLite.sql({
+            sql: `DELETE FROM knowledgebase_chunks_fts WHERE chunk_id IN (${placeholders})`,
+            values: batch,
+        });
+        await BunSQLite.sql({ sql: `DELETE FROM knowledgebase_chunks WHERE id IN (${placeholders})`, values: batch });
+    }
+}
+
+async function insertChunks({
+    document,
+    chunks,
+}: {
+    document: { id: number; knowledgebase_id: number };
+    chunks: string[];
+}) {
+    const ids: number[] = [];
+
+    await ensureKnowledgebaseFts();
+
+    for (let index = 0; index < chunks.length; index++) {
+        const text = chunks[index]!;
+
+        const res = await BunSQLite.insert<BUN_SQLITE_COBBLE_KNOWLEDGEBASE_CHUNKS, (typeof BunSQLiteTables)[number]>({
+            table: "knowledgebase_chunks",
+            data: [
+                {
+                    document_id: document.id,
+                    knowledgebase_id: Number(document.knowledgebase_id),
+                    chunk_index: index,
+                    text,
+                    token_estimate: Math.ceil(text.length / 4),
+                },
+            ],
+        });
+
+        const id = res.postInsertReturn?.insertId;
+
+        if (!res.success || !id) throw new Error(`Couldn't store a passage: ${res.msg || res.error}`);
+
+        ids.push(Number(id));
+
+        /**
+         * Its words go in the full-text index, which ranks a word search.
+         */
+        const indexed = await BunSQLite.sql({
+            sql: `INSERT INTO knowledgebase_chunks_fts (text, chunk_id, document_id, knowledgebase_id) VALUES (?, ?, ?, ?)`,
+            values: [text, Number(id), document.id, Number(document.knowledgebase_id)],
+        });
+
+        if (!indexed.success) throw new Error(`Couldn't index a passage: ${indexed.error || indexed.msg}`);
+    }
+
+    return ids;
+}
+
+async function insertVector({
+    vector,
+    chunk_id,
+    document_id,
+    knowledgebase_id,
+}: {
+    vector: Float32Array;
+    chunk_id: number;
+    document_id: number;
+    knowledgebase_id: number;
+}) {
+    /**
+     * vec0's metadata columns (its id and timestamps) are NOT NULL, so a
+     * raw insert fills them; the vector needs `vec_f32`.
+     */
+    const now = Date.now();
+
+    const res = await BunSQLite.sql({
+        sql: `INSERT INTO knowledgebase_vectors (created_at, updated_at, embedding, chunk_id, document_id, knowledgebase_id) VALUES (?, ?, vec_f32(?), ?, ?, ?)`,
+        values: [now, now, vector as any, chunk_id, document_id, knowledgebase_id],
+    });
+
+    if (!res.success) throw new Error(`Couldn't store an embedding: ${res.error || res.msg}`);
+}
+
+async function finish({
+    document_id,
+    chunks,
+    model,
+    error,
+}: {
+    document_id: number;
+    chunks: number;
+    model: string | null;
+    error?: string;
+}) {
+    const now = Date.now();
+
+    await BunSQLite.sql({
+        sql: `UPDATE knowledgebase_documents SET ingested_at = ?, chunk_count = ?, embedding_model = ?, ingest_error = ?,
+                updated_at = ? WHERE id = ?`,
+        values: [now, chunks, model, error || null, now, document_id],
+    });
+}