steve/kb
1
0
Fork 0
Code Issues 1 Pull Requests Actions Projects Releases 20 Wiki Activity

Complete outstanding ingestion and document UX work

Browse Source
This commit is contained in:
Steve Cliff
2026-08-22 18:54:07 +01:00
parent 932e889ee8
commit 5b5a9ecbc3
24 changed files with 696 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files
engine/compose.cpu.yaml
+1
View File
@@ -13,6 +13,7 @@ services:
- KB_INGEST_DEVICE=cpu
- KB_API_KEY=${KB_API_KEY:-}
- KB_SEARCH_THRESHOLD=${KB_SEARCH_THRESHOLD:-0.01}
- KB_MIN_CHUNK_ALNUM=${KB_MIN_CHUNK_ALNUM:-3}
- HF_HUB_OFFLINE=${HF_HUB_OFFLINE:-}
restart: unless-stopped
engine/compose.nvidia.yaml
+1
View File
@@ -21,6 +21,7 @@ services:
- KB_INGEST_DEVICE=${KB_INGEST_DEVICE:-auto}
- KB_API_KEY=${KB_API_KEY:-}
- KB_SEARCH_THRESHOLD=${KB_SEARCH_THRESHOLD:-0.01}
- KB_MIN_CHUNK_ALNUM=${KB_MIN_CHUNK_ALNUM:-3}
- KB_RERANK_ENABLED=${KB_RERANK_ENABLED:-true}
- KB_RERANKER_MODEL=${KB_RERANKER_MODEL:-BAAI/bge-reranker-v2-m3}
- KB_RERANK_CANDIDATES=${KB_RERANK_CANDIDATES:-40}
engine/kb/config.py
+1
View File
@@ -24,6 +24,7 @@ class Config:
self.reranker_model = os.environ.get("KB_RERANKER_MODEL", "BAAI/bge-reranker-v2-m3")
self.rerank_candidates = int(os.environ.get("KB_RERANK_CANDIDATES", "40"))
self.bulk_safety_percent = int(os.environ.get("KB_BULK_SAFETY_PERCENT", "70"))
self.min_chunk_alnum = int(os.environ.get("KB_MIN_CHUNK_ALNUM", "3"))
self.host = os.environ.get("KB_HOST", "0.0.0.0")
self.port = int(os.environ.get("KB_PORT", "8000"))
engine/kb/ingest/docling_pipeline.py
+5 -1
View File
@@ -25,6 +25,7 @@ from docling.document_converter import DocumentConverter, PdfFormatOption # noq
from docling_core.transforms.chunker.hierarchical_chunker import ( # noqa: E402
HierarchicalChunker,
)
from kb.ingest.quality import has_minimum_content
def _fixed_size_chunks(text: str, max_chars: int = 2000) -> list[str]:
@@ -40,6 +41,7 @@ def _fixed_size_chunks(text: str, max_chars: int = 2000) -> list[str]:
def chunk_document(
file_path: Path,
ingest_device: str = "cpu",
min_chunk_alnum: int = 3,
) -> list[dict]:
"""Convert and chunk a PDF/DOCX/HTML document using Docling.
@@ -71,7 +73,7 @@ def chunk_document(
chunks: list[dict] = []
for idx, chunk in enumerate(raw_chunks):
text = chunk.text.strip() if hasattr(chunk, "text") else str(chunk).strip()
if not text:
if not text or not has_minimum_content(text, min_chunk_alnum):
continue
metadata: dict = {}
@@ -98,6 +100,8 @@ def chunk_document(
if not full_text and hasattr(doc, "text"):
full_text = doc.text
for idx, piece in enumerate(_fixed_size_chunks(full_text)):
if not has_minimum_content(piece, min_chunk_alnum):
continue
chunks.append({
"text": piece,
"chunk_index": idx,
engine/kb/ingest/quality.py
+14
View File
@@ -0,0 +1,14 @@
"""Small, conservative ingestion-quality checks."""
from __future__ import annotations
def has_minimum_content(text: str, min_alnum: int = 3) -> bool:
"""Return whether text contains enough letters/numbers to be searchable.
Counting alphanumeric characters avoids indexing OCR fragments consisting
only of punctuation or one-character labels while retaining short IDs.
"""
if min_alnum <= 0:
return True
return sum(character.isalnum() for character in text) >= min_alnum
engine/kb/maintenance/__init__.py
+1
View File
@@ -0,0 +1 @@
"""Explicit maintenance commands for kb-engine data."""
engine/kb/maintenance/backfill_note_titles.py
+115
View File
@@ -0,0 +1,115 @@
"""Repair notes whose title was generated from the old ``note`` fallback.
Preview changes by default::
python -m kb.maintenance.backfill_note_titles
Apply them, including refreshed FTS text and embeddings::
python -m kb.maintenance.backfill_note_titles --apply
"""
from __future__ import annotations
import argparse
import json
import re
import struct
from kb import database, embeddings
from kb.config import cfg
from kb.ingest.note import auto_title
_SYNTHETIC_NOTE = re.compile(
r"^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}_note\.note$",
re.IGNORECASE,
)
def find_repairs(conn) -> list[dict]:
"""Return unambiguous synthetic note titles and their derived replacements."""
rows = conn.execute(
"""
SELECT d.id, d.title, d.original_filename, c.text, c.metadata
FROM documents d
JOIN chunks c ON c.document_id = d.id AND c.chunk_index = 0
WHERE d.doc_type = 'note'
ORDER BY d.id
"""
).fetchall()
repairs = []
for row in rows:
if not _SYNTHETIC_NOTE.fullmatch(row["title"] or ""):
continue
title = auto_title(row["text"] or "")
if not title:
continue
repairs.append({
"document_id": row["id"],
"old_title": row["title"],
"new_title": title,
"original_filename": row["original_filename"],
})
return repairs
def apply_repairs(conn, repairs: list[dict]) -> None:
"""Update titles, enriched text, FTS, and vectors for selected notes."""
for repair in repairs:
doc_id = repair["document_id"]
title = repair["new_title"]
chunks = conn.execute(
"SELECT id, text, metadata FROM chunks WHERE document_id = ? ORDER BY chunk_index",
(doc_id,),
).fetchall()
enriched = []
for chunk in chunks:
metadata = json.loads(chunk["metadata"] or "{}")
enriched.append(database.build_enriched_text(title, chunk["text"], metadata))
vectors = embeddings.embed_texts(enriched)
conn.execute(
"UPDATE documents SET title = ?, updated_at = current_timestamp WHERE id = ?",
(title, doc_id),
)
for chunk, text, vector in zip(chunks, enriched, vectors):
conn.execute(
"UPDATE chunks SET enriched_text = ? WHERE id = ?", (text, chunk["id"])
)
conn.execute("DELETE FROM chunks_vec WHERE chunk_id = ?", (chunk["id"],))
blob = struct.pack(f"{len(vector)}f", *vector)
conn.execute(
"INSERT INTO chunks_vec(embedding, chunk_id) VALUES (?, ?)",
(blob, chunk["id"]),
)
conn.commit()
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument(
"--apply", action="store_true", help="apply repairs (the default is preview only)"
)
args = parser.parse_args()
conn = database.get_connection(cfg.db_path)
try:
repairs = find_repairs(conn)
for repair in repairs:
print(
f'{repair["document_id"]}: {repair["old_title"]!r} -> '
f'{repair["new_title"]!r}'
)
if not args.apply:
print(f"Previewed {len(repairs)} repair(s); rerun with --apply to update them.")
return
embeddings.load_model(cfg.model, cfg.device)
apply_repairs(conn, repairs)
print(f"Repaired {len(repairs)} note title(s).")
finally:
conn.close()
if __name__ == "__main__":
main()
engine/kb/routes/documents.py
+67 -6
View File
@@ -7,11 +7,13 @@ from pathlib import Path
from typing import Optional
from fastapi import HTTPException, Query
from pydantic import BaseModel, Field
from fastapi.responses import FileResponse
from main import app
from kb.config import cfg
from kb.database import get_connection
from kb.search import hybrid_search
logger = logging.getLogger("kb.routes.documents")
@@ -20,11 +22,13 @@ logger = logging.getLogger("kb.routes.documents")
async def list_documents(
type: Optional[str] = Query(None),
tags: Optional[str] = Query(None),
title: Optional[str] = Query(None),
filename: Optional[str] = Query(None),
):
conn = get_connection(cfg.db_path)
try:
sql = """
SELECT d.id, d.title, d.doc_type,
SELECT d.id, d.title, d.original_filename, d.source_path, d.doc_type,
(SELECT COUNT(*) FROM chunks c WHERE c.document_id = d.id) AS chunk_count,
d.created_at, d.updated_at
FROM documents d
@@ -37,6 +41,14 @@ async def list_documents(
where.append("d.doc_type = ?")
params.append(type)
if title:
where.append("d.title LIKE ? COLLATE NOCASE")
params.append(f"%{title}%")
if filename:
where.append("d.original_filename LIKE ? COLLATE NOCASE")
params.append(f"%{filename}%")
if tags:
tag_list = [t.strip() for t in tags.split(",") if t.strip()]
for i, tag in enumerate(tag_list):
@@ -70,6 +82,8 @@ async def list_documents(
results.append({
"id": row["id"],
"title": row["title"],
"original_filename": row["original_filename"],
"source_path": row["source_path"],
"doc_type": row["doc_type"],
"tags": [t["name"] for t in tag_rows],
"chunk_count": row["chunk_count"],
@@ -82,8 +96,49 @@ async def list_documents(
conn.close()
class DocumentFindRequest(BaseModel):
query: str
top: int = Field(default=10, ge=1, le=100)
tags: Optional[list[str]] = None
doc_type: Optional[str] = None
@app.post("/api/v1/documents/find")
async def find_documents(req: DocumentFindRequest):
"""Return document-level results aggregated from hybrid chunk search."""
conn = get_connection(cfg.db_path)
try:
search_result = hybrid_search(
conn,
req.query,
cfg,
top=max(req.top * 10, 50),
tags=req.tags,
doc_type=req.doc_type,
)
documents: dict[int, dict] = {}
for result in search_result["results"]:
doc_id = result["document_id"]
if doc_id not in documents:
documents[doc_id] = {
"document_id": doc_id,
"title": result["title"],
"doc_type": result["doc_type"],
"source_path": result["source_path"],
"original_filename": result["original_filename"],
"tags": result["tags"],
"score": result["score"],
"hit_count": 0,
"top_chunk": result["text"],
}
documents[doc_id]["hit_count"] += 1
return list(documents.values())[: req.top]
finally:
conn.close()
@app.get("/api/v1/documents/{doc_id}")
async def get_document(doc_id: int):
async def get_document(doc_id: int, include_chunks: bool = Query(True)):
conn = get_connection(cfg.db_path)
try:
doc = conn.execute(
@@ -92,10 +147,15 @@ async def get_document(doc_id: int):
if not doc:
raise HTTPException(status_code=404, detail="Document not found.")
chunks = conn.execute(
"SELECT * FROM chunks WHERE document_id = ? ORDER BY chunk_index",
(doc_id,),
).fetchall()
chunk_count = conn.execute(
"SELECT COUNT(*) AS n FROM chunks WHERE document_id = ?", (doc_id,)
).fetchone()["n"]
chunks = []
if include_chunks:
chunks = conn.execute(
"SELECT * FROM chunks WHERE document_id = ? ORDER BY chunk_index",
(doc_id,),
).fetchall()
tag_rows = conn.execute(
"""
@@ -114,6 +174,7 @@ async def get_document(doc_id: int):
**dict(doc),
"has_file": has_file,
"tags": [t["name"] for t in tag_rows],
"chunk_count": chunk_count,
"chunks": [dict(c) for c in chunks],
}
finally:
engine/kb/routes/jobs.py
+3 -1
View File
@@ -10,6 +10,7 @@ from fastapi.responses import JSONResponse
from main import app
from kb.config import cfg
from kb.database import get_connection, create_job, get_job, list_jobs, get_document_by_hash
from kb.ingest.note import auto_title
from kb.staging import stage_file, stage_note
@@ -32,6 +33,7 @@ async def submit_job(
content_hash = hashlib.sha256(content).hexdigest()
filename = file.filename
else:
title = title or auto_title(note) or "note"
content = note.encode("utf-8")
content_hash = hashlib.sha256(content).hexdigest()
filename = None
@@ -48,7 +50,7 @@ async def submit_job(
if file:
staging_path = stage_file(cfg.staging_dir, file.filename, content)
else:
staging_path = stage_note(cfg.staging_dir, title or "note", note)
staging_path = stage_note(cfg.staging_dir, title, note)
filename = staging_path.name
tags_list = [t.strip() for t in tags.split(",") if t.strip()] if tags else []
engine/kb/search.py
+2 -1
View File
@@ -422,7 +422,7 @@ def _enrich(
"""
SELECT c.id, c.text, c.chunk_index, c.metadata AS chunk_meta,
d.id AS doc_id, d.title, d.doc_type, d.source_path,
d.created_at
d.created_at, d.original_filename
FROM chunks c
JOIN documents d ON c.document_id = d.id
WHERE c.id = ?
@@ -456,6 +456,7 @@ def _enrich(
"doc_type": row[6],
"source_path": row[7],
"created_at": row[8],
"original_filename": row[9],
"tags": [t[0] for t in tag_rows],
"tag_contexts": {t[0]: t[1] for t in tag_rows if t[1]},
}
engine/kb/worker.py
+3 -1
View File
@@ -113,7 +113,9 @@ def _process_job(job_row) -> tuple[str, int | None, int]:
chunks = chunk_note(text)
elif doc_type == "pdf":
from kb.ingest.docling_pipeline import chunk_document
chunks = chunk_document(staged_path, cfg.ingest_device)
chunks = chunk_document(
staged_path, cfg.ingest_device, cfg.min_chunk_alnum
)
elif doc_type == "markdown":
text = staged_path.read_text(encoding="utf-8")
from kb.ingest.markdown import chunk_markdown
engine/tests/test_document_routes.py
+102
View File
@@ -0,0 +1,102 @@
"""Focused tests for document metadata lookup and aggregation."""
import pytest
from kb import database
from kb.routes import documents
@pytest.fixture
def document_db(tmp_path, monkeypatch):
db_path = tmp_path / "kb.db"
conn = database.get_connection(db_path)
database.init_schema(conn, 3)
conn.execute(
"""
INSERT INTO documents(title, source_path, content_hash, doc_type, original_filename)
VALUES ('Vehicle Guide', '/data/staging/random.pdf', 'hash', 'pdf', 'M38T_manual.pdf')
"""
)
doc_id = conn.execute("SELECT id FROM documents").fetchone()["id"]
conn.execute(
"INSERT INTO chunks(document_id, chunk_index, text, enriched_text) VALUES (?, 0, 'body', 'body')",
(doc_id,),
)
conn.commit()
conn.close()
monkeypatch.setattr(documents.cfg, "data_dir", tmp_path)
return doc_id
@pytest.mark.asyncio
async def test_list_filters_title_and_original_filename(document_db):
by_title = await documents.list_documents(
type=None, tags=None, title="vehicle", filename=None
)
by_filename = await documents.list_documents(
type=None, tags=None, title=None, filename="m38t"
)
assert [item["id"] for item in by_title] == [document_db]
assert by_filename[0]["original_filename"] == "M38T_manual.pdf"
@pytest.mark.asyncio
async def test_info_can_omit_chunks_without_losing_count(document_db):
result = await documents.get_document(document_db, include_chunks=False)
assert result["chunk_count"] == 1
assert result["chunks"] == []
@pytest.mark.asyncio
async def test_find_aggregates_chunk_hits_by_document(monkeypatch):
class Connection:
def close(self):
pass
monkeypatch.setattr(documents, "get_connection", lambda _path: Connection())
monkeypatch.setattr(
documents,
"hybrid_search",
lambda *_args, **_kwargs: {
"results": [
{
"document_id": 7,
"title": "Guide",
"doc_type": "pdf",
"source_path": "/staging/file",
"original_filename": "guide.pdf",
"tags": [],
"score": 0.8,
"text": "best hit",
},
{
"document_id": 7,
"title": "Guide",
"doc_type": "pdf",
"source_path": "/staging/file",
"original_filename": "guide.pdf",
"tags": [],
"score": 0.7,
"text": "second hit",
},
]
},
)
result = await documents.find_documents(
documents.DocumentFindRequest(query="guide")
)
assert result == [{
"document_id": 7,
"title": "Guide",
"doc_type": "pdf",
"source_path": "/staging/file",
"original_filename": "guide.pdf",
"tags": [],
"score": 0.8,
"hit_count": 2,
"top_chunk": "best hit",
}]
engine/tests/test_ingest_quality.py
+19
View File
@@ -0,0 +1,19 @@
"""Tests for filtering noisy OCR fragments."""
from kb.ingest.quality import has_minimum_content
def test_short_ocr_fragments_are_rejected():
assert not has_minimum_content('"')
assert not has_minimum_content("B")
assert not has_minimum_content("12")
def test_short_identifiers_and_real_text_are_retained():
assert has_minimum_content("BID")
assert has_minimum_content("00:15")
assert has_minimum_content("Useful text")
def test_filter_can_be_disabled():
assert has_minimum_content("B", min_alnum=0)
engine/tests/test_note_titles.py
+75
View File
@@ -0,0 +1,75 @@
"""Tests for automatic and backfilled note titles."""
import sqlite3
from kb.ingest.note import auto_title
from kb.maintenance import backfill_note_titles
from kb.maintenance.backfill_note_titles import find_repairs
def test_auto_title_strips_markdown_and_limits_length():
assert auto_title("## Useful heading\nBody") == "Useful heading"
assert auto_title("x" * 100) == "x" * 80
def test_find_repairs_only_selects_synthetic_note_titles():
conn = sqlite3.connect(":memory:")
conn.row_factory = sqlite3.Row
conn.executescript(
"""
CREATE TABLE documents (
id INTEGER PRIMARY KEY, title TEXT, original_filename TEXT, doc_type TEXT
);
CREATE TABLE chunks (
id INTEGER PRIMARY KEY, document_id INTEGER, chunk_index INTEGER,
text TEXT, metadata TEXT
);
INSERT INTO documents VALUES
(1, '7dec828a-1234-4567-89ab-123456789abc_note.note',
'7dec828a-1234-4567-89ab-123456789abc_note.note', 'note'),
(2, 'A deliberate title', 'named.note', 'note');
INSERT INTO chunks VALUES
(1, 1, 0, '# Derived title\nBody', '{}'),
(2, 2, 0, 'Must not replace', '{}');
"""
)
assert find_repairs(conn) == [{
"document_id": 1,
"old_title": "7dec828a-1234-4567-89ab-123456789abc_note.note",
"new_title": "Derived title",
"original_filename": "7dec828a-1234-4567-89ab-123456789abc_note.note",
}]
def test_apply_repairs_refreshes_title_search_text_and_vector(monkeypatch):
conn = sqlite3.connect(":memory:")
conn.row_factory = sqlite3.Row
conn.executescript(
"""
CREATE TABLE documents (id INTEGER PRIMARY KEY, title TEXT, updated_at TEXT);
CREATE TABLE chunks (
id INTEGER PRIMARY KEY, document_id INTEGER, chunk_index INTEGER,
text TEXT, metadata TEXT, enriched_text TEXT
);
CREATE TABLE chunks_vec (embedding BLOB, chunk_id INTEGER);
INSERT INTO documents VALUES (1, 'old', NULL);
INSERT INTO chunks VALUES (3, 1, 0, 'New title\nBody', '{}', 'old text');
INSERT INTO chunks_vec VALUES (X'00', 3);
"""
)
monkeypatch.setattr(
backfill_note_titles.embeddings,
"embed_texts",
lambda texts: [[1.0, 2.0, 3.0] for _ in texts],
)
backfill_note_titles.apply_repairs(conn, [{
"document_id": 1, "new_title": "New title"
}])
assert conn.execute("SELECT title FROM documents").fetchone()[0] == "New title"
assert conn.execute("SELECT enriched_text FROM chunks").fetchone()[0] == (
"New title\n\nNew title\nBody"
)
assert len(conn.execute("SELECT embedding FROM chunks_vec").fetchone()[0]) == 12