16 Commits

Author SHA1 Message Date
steve 5b5a9ecbc3 Complete outstanding ingestion and document UX work 2026-08-22 18:54:07 +01:00
steve 932e889ee8 Add manual Docker image rebuild workflow 2026-08-22 18:09:47 +01:00
steve 48ca873fe2 Publish Docker images to public registry 2026-08-22 18:04:18 +01:00
steve a38d77ed23 Bump client version to 3.3.0 2026-08-21 18:37:16 +01:00
steve f1ed5b6e23 Split the torch layer and drop 7.8GB of orphaned CUDA libs from the CPU image
The CPU image was 11.2GB — larger than the CUDA one — and its single 6.72GB
torch layer could not be pushed at all: the registry drops any blob upload
taking longer than 60s, and that layer needed ~61s.

Both problems came from install ordering. Dockerfile.cpu installed the project
and sentence-transformers first, which resolved the default CUDA torch and
pulled ~2.7GB of nvidia-* wheels, then reinstalled torch from the CPU index.
Reinstalling replaces torch but leaves its transitive CUDA dependencies behind,
orphaned and unused (torch reports 2.13.0+cpu, cuda.is_available() False).

Installing CPU torch first, from the CPU index, means nothing ever requests a
CUDA build. Everything else then resolves against the torch already present.

  CPU image     11.2GB -> 3.43GB
  largest layer  6.72GB -> 896MB
  push           failed 5/5 at 60s -> succeeds in 8.6s

Dockerfile.nvidia gets the same split. Its torch layer is unavoidably large
(~2.8GiB compressed, ~52s at current throughput) so it stays the tightest
thing in the stack, but splitting the application install off keeps that layer
cached and the app layer small.

Both now install torch before the source COPYs, so editing application code no
longer invalidates the multi-GB layer.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
2026-08-21 18:33:34 +01:00
steve 3ab8a81c14 Retry image pushes on transient registry failures
The engine images carry a ~5.6GB torch layer. Uploading it intermittently
fails with a 502 from the reverse proxy in front of the registry, and the
manifest PUT that follows can fail with a 500 because the blob commit has
not registered yet. Both clear on a retry, but the script pushed each tag
exactly once, so a whole release — including a 25GB rebuild — could be lost
to one hiccup.

push_image() replaces the bare `run docker push` calls and retries up to
PUSH_RETRIES times (default 5) with PUSH_RETRY_DELAY seconds between
attempts (default 10). Exhausting the retries still returns non-zero so
set -e aborts the release rather than reporting a partial push as success.
Dry runs make no docker calls.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
2026-08-21 17:30:33 +01:00
steve 739c3ff30c Disable buildx attestations so pushes to the Registry v2 host succeed
buildx attaches provenance/SBOM attestation manifests by default, which
makes each built image an OCI image index
(application/vnd.oci.image.index.v1+json). docker.dcglab.co.uk rejects
those with a 500 on the manifest PUT — blobs upload fine, then the
manifest fails, so the push looks like an auth problem when it is not.

Passing --provenance=false --sbom=false yields a plain image manifest,
which the registry accepts.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
2026-08-21 17:20:41 +01:00
steve 3151a1b08a Move container registry to docker.dcglab.co.uk/kb
registry.dcglab.co.uk has been decommissioned and replaced by
docker.dcglab.co.uk. Images move off the Gitea built-in registry (which
never successfully accepted a push) to the new host, published at the
root with no org segment: docker.dcglab.co.uk/kb/engine and .../kb/mcp.

IMAGE_ORG is now empty by default but still honoured when set, so the
documented REGISTRY=ghcr.io IMAGE_ORG=myorg override keeps working.

Gitea remains the git forge — the release download URLs in README.md are
unchanged and still point at gitea.dcglab.co.uk.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
2026-08-21 17:15:42 +01:00
steve 5049ba2a2f Bump engine version to 3.3.0 2026-08-21 09:58:13 +01:00
steve 6dfc13be1d Add reranking, RRF fusion, bench harness, tag contexts, and data ingestion
Implements five of the six enhancements from docs/kb-enhancements-proposal.htm,
closing the retrieval-quality gap identified in the qmd review.

- Cross-encoder reranking: new kb/reranker.py loads an optional reranking
  model at startup (KB_RERANK_ENABLED, KB_RERANKER_MODEL,
  KB_RERANK_CANDIDATES). Search degrades gracefully to plain hybrid
  retrieval when the model is absent. Exposed via a "rerank" block in
  /status, a rerank flag on search, and --no-rerank in the CLI.
- RRF rank fusion: FTS and vector lists now merge by reciprocal rank
  fusion with a top-rank bonus, replacing the old score blend. Scores are
  comparable across queries.
- Bench harness and explain traces: kb bench runs a query fixture against
  each backend and reports precision@k, recall and MRR. --explain returns a
  per-result score breakdown.
- Tag context descriptions: tags carry an optional one-line description
  (kb tag-describe), returned as tag_contexts with search results. Adds a
  tags.description column migration.
- Structured data ingestion: .json/.yaml/.toml files ingest as text via the
  new "data" doc type, pretty-printing minified JSON before chunking.

Query expansion (proposal item 5) is deliberately left out pending bench
results. Requires engine v3.3.0.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
2026-08-21 09:51:33 +01:00
steve 75e4a0cf73 Bump engine version to 3.2.4 2026-05-15 18:37:12 +01:00
steve 753c641e72 Switch container registry to Gitea built-in (gitea.dcglab.co.uk/steve/kb)
The standalone Registry v2 host (docker.dcglab.co.uk, briefly registry.dcglab.co.uk)
is being scrapped. Move all kb images to Gitea's built-in container registry.
2026-05-15 18:35:20 +01:00
steve 45e2c5ce91 Bump engine version to 3.2.3 2026-05-15 18:22:08 +01:00
steve e6e91f1d5c Clarify hybrid semantic + full-text search in MCP descriptions
Agents were misreading kb_search as keyword-only because the vector/semantic
component was only mentioned in the negative ("fts_only: no vector similarity").
Lead with hybrid semantic + BM25 + RRF in the server instructions, kb_search
docstring, and MCP.md so agents recognise it as a vector search tool.
2026-05-15 18:19:42 +01:00
steve 9eccc527ae Add next-steps.md with UX improvement ideas for kb CLI
Captures pain points found while trying to locate an uploaded PDF: kb
list silently ignores positional args, kb search results lack
document_id, kb info dumps all chunks with no summary mode, and
scan-heavy PDFs produce noisy single-char chunk hits.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-05-13 20:34:55 +01:00
steve d44d11e4fe Bump engine version to 3.2.2 2026-04-14 21:48:55 +01:00
59 changed files with 2788 additions and 131 deletions
+81
View File
@@ -0,0 +1,81 @@
name: Rebuild Docker images
on:
workflow_dispatch:
concurrency:
group: rebuild-docker-images
cancel-in-progress: false
jobs:
rebuild-and-push:
runs-on: ubuntu-latest
env:
REGISTRY: docker.dcglab.co.uk
IMAGE_BASE: docker.dcglab.co.uk/public/kb
REGISTRY_USERNAME: ${{ secrets.DOCKER_DCGLAB_CI_USERNAME }}
REGISTRY_PASSWORD: ${{ secrets.DOCKER_DCGLAB_CI_PASSWORD }}
steps:
- name: Check out repository
uses: actions/checkout@v4
- name: Log in to registry
shell: bash
run: |
set -euo pipefail
test -n "$REGISTRY_USERNAME" || { echo "DOCKER_DCGLAB_CI_USERNAME is not available" >&2; exit 1; }
test -n "$REGISTRY_PASSWORD" || { echo "DOCKER_DCGLAB_CI_PASSWORD is not available" >&2; exit 1; }
printf '%s' "$REGISTRY_PASSWORD" | docker login "$REGISTRY" --username "$REGISTRY_USERNAME" --password-stdin
- name: Build all images from scratch
shell: bash
run: |
set -euo pipefail
version="$(tr -d '[:space:]' < engine/VERSION)"
docker build --pull --no-cache --provenance=false --sbom=false \
-t "$IMAGE_BASE/engine:v${version}-nvidia" \
-t "$IMAGE_BASE/engine:latest-nvidia" \
-f engine/Dockerfile.nvidia engine
docker build --pull --no-cache --provenance=false --sbom=false \
-t "$IMAGE_BASE/engine:v${version}-cpu" \
-t "$IMAGE_BASE/engine:latest-cpu" \
-f engine/Dockerfile.cpu engine
docker build --pull --no-cache --provenance=false --sbom=false \
-t "$IMAGE_BASE/mcp:v${version}" \
-t "$IMAGE_BASE/mcp:latest" \
-f mcp/Dockerfile mcp
- name: Push and verify all tags
shell: bash
run: |
set -euo pipefail
version="$(tr -d '[:space:]' < engine/VERSION)"
images=(
"$IMAGE_BASE/engine:v${version}-nvidia"
"$IMAGE_BASE/engine:latest-nvidia"
"$IMAGE_BASE/engine:v${version}-cpu"
"$IMAGE_BASE/engine:latest-cpu"
"$IMAGE_BASE/mcp:v${version}"
"$IMAGE_BASE/mcp:latest"
)
push_image() {
local image="$1"
local attempt
for attempt in 1 2 3 4 5; do
docker push "$image" && return 0
if [[ "$attempt" -eq 5 ]]; then
echo "Failed to push $image after $attempt attempts" >&2
return 1
fi
sleep 10
done
}
for image in "${images[@]}"; do
push_image "$image"
docker manifest inspect "$image" >/dev/null
done
+46 -3
View File
@@ -21,6 +21,35 @@ make build # produces ./kb binary
make all # or cross-compile: dist/kb-{os}-{arch} make all # or cross-compile: dist/kb-{os}-{arch}
``` ```
## Running tests
### Engine
Engine tests run against SQLite (with sqlite-vec) and stub out the embedding
model, so they only need lightweight dependencies — no torch/docling install:
```bash
uv venv /tmp/kb-test-venv
uv pip install --python /tmp/kb-test-venv/bin/python pytest pytest-asyncio fastapi httpx sqlite-vec
cd engine && /tmp/kb-test-venv/bin/python -m pytest
```
### Client
```bash
cd client && go test ./...
```
## Search-quality benchmarking
`kb bench fixture.json` runs a fixture of queries with known-relevant documents
against each backend (fts, vec, hybrid, hybrid+rerank) and reports precision@k,
recall, and MRR. See `docs/bench-example.json` for the fixture format.
Run a bench before and after any ranking change (RRF weights, reranker, model
swap) and compare — keep a 20-30 query fixture against your real corpus outside
the repo.
## Building and releasing ## Building and releasing
Client and engine are versioned independently via `client/VERSION` and `engine/VERSION`. Each has its own release script and git tag prefix. Client and engine are versioned independently via `client/VERSION` and `engine/VERSION`. Each has its own release script and git tag prefix.
@@ -61,17 +90,30 @@ curl http://localhost:8000/api/v1/status | jq .version
### Docker images ### Docker images
Images are pushed to `docker.dcglab.co.uk/dcg/kb/engine` with tags: Images are pushed to `docker.dcglab.co.uk/public/kb/engine` with tags:
- `engine-v2.0.6-nvidia` / `engine-v2.0.6-cpu` — versioned - `engine-v2.0.6-nvidia` / `engine-v2.0.6-cpu` — versioned
- `latest-nvidia` / `latest-cpu` — latest release - `latest-nvidia` / `latest-cpu` — latest release
The release script authenticates to the registry using the
`DOCKER_DCGLAB_CI_USERNAME` and `DOCKER_DCGLAB_CI_PASSWORD` environment
variables.
Override the registry and org via environment variables: Override the registry and org via environment variables:
```bash ```bash
REGISTRY=ghcr.io IMAGE_ORG=myorg ./release-engine.sh --github REGISTRY=ghcr.io IMAGE_ORG=myorg ./release-engine.sh --github
``` ```
Pushes are retried on transient registry failures. The engine images carry a
~5.6GB torch layer, and uploading it can fail with a 502 from the proxy in
front of the registry (or a 500 on the manifest PUT that follows), which
clears on a retry. Tune with:
```bash
PUSH_RETRIES=8 PUSH_RETRY_DELAY=20 ./release-engine.sh --gitea
```
## API reference ## API reference
All endpoints are under `/api/v1/`. Requires `Authorization: Bearer <key>` header when `KB_API_KEY` is set. All endpoints are under `/api/v1/`. Requires `Authorization: Bearer <key>` header when `KB_API_KEY` is set.
@@ -88,8 +130,9 @@ All endpoints are under `/api/v1/`. Requires `Authorization: Bearer <key>` heade
| `GET` | `/documents/{id}/file` | Download original file | | `GET` | `/documents/{id}/file` | Download original file |
| `DELETE` | `/documents/{id}` | Remove a document (and stored file) | | `DELETE` | `/documents/{id}` | Remove a document (and stored file) |
| `PUT` | `/documents/{id}/tags` | Add/remove tags | | `PUT` | `/documents/{id}/tags` | Add/remove tags |
| `GET` | `/tags` | List all tags | | `GET` | `/tags` | List all tags (with descriptions) |
| `GET` | `/status` | Engine status, GPU info, DB stats | | `PUT` | `/tags/{name}/description` | Set/clear a tag context description |
| `GET` | `/status` | Engine status, GPU info, DB stats, rerank state |
| `POST` | `/reindex` | Re-embed all chunks | | `POST` | `/reindex` | Re-embed all chunks |
| `POST` | `/bulk/delete` | Bulk delete documents by filter | | `POST` | `/bulk/delete` | Bulk delete documents by filter |
| `POST` | `/bulk/tags` | Bulk add/remove tags by filter | | `POST` | `/bulk/tags` | Bulk add/remove tags by filter |
+3 -3
View File
@@ -1,6 +1,6 @@
# MCP Server (Agent Integration) # MCP Server (Agent Integration)
The MCP server exposes kb operations as native MCP tools, so agents can search, add notes, upload files, and manage documents without shelling out to the CLI. The MCP server exposes kb operations as native MCP tools, so agents can search, add notes, upload files, and manage documents without shelling out to the CLI. `kb_search` is hybrid: dense vector embeddings (semantic similarity) fused with BM25 full-text ranking via Reciprocal Rank Fusion, so agents can ask natural-language questions and find conceptually related content even when the exact words don't match.
## Start the MCP server ## Start the MCP server
@@ -20,14 +20,14 @@ docker run -d --name kb-mcp \
-e KB_API_KEY=your-engine-key \ -e KB_API_KEY=your-engine-key \
-e KB_MCP_API_KEY=your-agent-key \ -e KB_MCP_API_KEY=your-agent-key \
--restart unless-stopped \ --restart unless-stopped \
docker.dcglab.co.uk/dcg/kb/mcp:latest docker.dcglab.co.uk/public/kb/mcp:latest
``` ```
## MCP tools ## MCP tools
| Tool | Description | | Tool | Description |
|---|---| |---|---|
| `kb_search` | Hybrid search with optional tag/type filters | | `kb_search` | Hybrid semantic (vector) + full-text search with tag/type filters |
| `kb_addnote` | Add a text note (queued for async ingestion) | | `kb_addnote` | Add a text note (queued for async ingestion) |
| `kb_update_note` | Update an existing note in place | | `kb_update_note` | Update an existing note in place |
| `kb_get` | Get document details by ID or source path | | `kb_get` | Get document details by ID or source path |
+31 -6
View File
@@ -33,7 +33,7 @@ docker run -d --name kb-engine \
-e KB_DEVICE=auto \ -e KB_DEVICE=auto \
-e KB_API_KEY=your-secret-key \ -e KB_API_KEY=your-secret-key \
--restart unless-stopped \ --restart unless-stopped \
docker.dcglab.co.uk/dcg/kb/engine:latest-nvidia docker.dcglab.co.uk/public/kb/engine:latest-nvidia
# CPU only (no GPU required — smaller image) # CPU only (no GPU required — smaller image)
docker run -d --name kb-engine \ docker run -d --name kb-engine \
@@ -42,7 +42,7 @@ docker run -d --name kb-engine \
-e KB_MODEL=all-MiniLM-L6-v2 \ -e KB_MODEL=all-MiniLM-L6-v2 \
-e KB_API_KEY=your-secret-key \ -e KB_API_KEY=your-secret-key \
--restart unless-stopped \ --restart unless-stopped \
docker.dcglab.co.uk/dcg/kb/engine:latest-cpu docker.dcglab.co.uk/public/kb/engine:latest-cpu
``` ```
Or use a compose file from the repo: Or use a compose file from the repo:
@@ -111,10 +111,12 @@ Override via environment variables (`KB_ENGINE_URL`, `KB_API_KEY`) or CLI flags
# Add notes # Add notes
kb addnote "Always restart nginx after config changes" kb addnote "Always restart nginx after config changes"
kb addnote "Server room is building 3, floor 2" --tags ops kb addnote "Server room is building 3, floor 2" --tags ops
kb addnote "Deploy checklist" --wait
# Add files (async — uploads and exits immediately) # Add files (async by default; --wait blocks until ingestion finishes)
kb addfile ~/docs/manual.pdf --tags admin kb addfile ~/docs/manual.pdf --tags admin
kb addfile ~/notes/ --recursive kb addfile ~/notes/ --recursive
kb addfile ~/docs/manual.pdf --wait
# Check ingestion progress # Check ingestion progress
kb jobs kb jobs
@@ -122,13 +124,16 @@ kb jobs
# Search # Search
kb search "how to install git" kb search "how to install git"
kb search "deploy process" --tags ops --type pdf kb search "deploy process" --tags ops --type pdf
kb find "vehicle handbook" --type pdf
# Update a note in place # Update a note in place
kb updatenote 42 "revised note content" kb updatenote 42 "revised note content"
# Manage # Manage
kb list kb list
kb info 1 kb list --title handbook
kb list --filename M38T_PHEV
kb info 1 --no-chunks
kb tags kb tags
kb tag 1 --add important kb tag 1 --add important
kb export 1 -o manual.pdf # download original file kb export 1 -o manual.pdf # download original file
@@ -143,8 +148,9 @@ kb bulk-set-tags --tags "old-scheme" --set "new-scheme" --yes
## How it works ## How it works
- **Ingestion**: Files are uploaded to the engine and queued for async processing. The engine chunks documents (PDFs via Docling, markdown by headers, code by AST/functions, notes as whole text), generates embeddings on GPU, and stores everything in SQLite. - **Ingestion**: Files are uploaded to the engine and queued for async processing. The engine chunks documents (PDFs via Docling, markdown by headers, code by AST/functions, notes as whole text, JSON/YAML/TOML as pretty-printed text), generates embeddings on GPU, and stores everything in SQLite.
- **Search**: Hybrid retrieval combining BM25 keyword scoring (FTS5) and vector similarity (sqlite-vec), merged via Reciprocal Rank Fusion. Sub-100ms with a warm model. - **Search**: Hybrid retrieval combining BM25 keyword scoring (FTS5) and vector similarity (sqlite-vec), merged via Reciprocal Rank Fusion with a top-rank bonus. Optionally reranked by a local cross-encoder (`KB_RERANK_ENABLED`). Sub-100ms with a warm model (without reranking). Add `--explain` to any search for a per-result score breakdown.
- **Quality measurement**: `kb bench fixture.json` runs a fixture of queries with known-relevant documents and reports precision@k / recall / MRR per backend (fts, vec, hybrid, hybrid+rerank). See `docs/bench-example.json`.
- **Output**: JSON (for scripts/LLM tool use) or human-readable terminal format. Use `--format json` on any command. - **Output**: JSON (for scripts/LLM tool use) or human-readable terminal format. Use `--format json` on any command.
## Engine configuration ## Engine configuration
@@ -159,12 +165,31 @@ The engine is configured via environment variables (set in the compose file or v
| `KB_INGEST_DEVICE` | `auto` | Docling layout detection device: `auto`, `cpu`, or `cuda` | | `KB_INGEST_DEVICE` | `auto` | Docling layout detection device: `auto`, `cpu`, or `cuda` |
| `KB_API_KEY` | (none) | Optional Bearer token for API authentication | | `KB_API_KEY` | (none) | Optional Bearer token for API authentication |
| `KB_SEARCH_THRESHOLD` | `0.01` | Minimum score for search results (filters noise) | | `KB_SEARCH_THRESHOLD` | `0.01` | Minimum score for search results (filters noise) |
| `KB_MIN_CHUNK_ALNUM` | `3` | Minimum alphanumeric characters retained in Docling PDF/DOCX/HTML chunks (`0` disables) |
| `KB_RERANK_ENABLED` | `false` (`true` in nvidia compose) | Load a cross-encoder and rerank hybrid results server-side |
| `KB_RERANKER_MODEL` | `BAAI/bge-reranker-v2-m3` | Cross-encoder model for reranking |
| `KB_RERANK_CANDIDATES` | `40` | Hybrid candidates scored by the reranker per query |
| `KB_BULK_SAFETY_PERCENT` | `70` | Bulk operations affecting more than this % of documents are rejected unless `force` is set (0 disables) | | `KB_BULK_SAFETY_PERCENT` | `70` | Bulk operations affecting more than this % of documents are rejected unless `force` is set (0 disables) |
| `KB_PORT` | `8000` | Port to expose | | `KB_PORT` | `8000` | Port to expose |
| `KB_HOST` | `0.0.0.0` | Host to bind to | | `KB_HOST` | `0.0.0.0` | Host to bind to |
| `HF_HUB_OFFLINE` | (none) | Set to `1` to prevent model downloads (use cached only) | | `HF_HUB_OFFLINE` | (none) | Set to `1` to prevent model downloads (use cached only) |
| `KB_DATA_PATH` | `./data` | Host path for bind mount (compose variable, not used by engine) | | `KB_DATA_PATH` | `./data` | Host path for bind mount (compose variable, not used by engine) |
### Repairing legacy note titles
Notes created by older clients may have a synthetic `<uuid>_note.note` title.
The maintenance command previews only unambiguous matches by default:
```bash
cd engine
python -m kb.maintenance.backfill_note_titles
python -m kb.maintenance.backfill_note_titles --apply
```
The apply mode reloads the configured embedding model and refreshes each
affected note's title, enriched full-text content, and vector embedding. Back
up the data directory before running it against production.
## Data portability ## Data portability
The data directory contains everything: SQLite database, model cache, and staging files. To migrate between hosts: The data directory contains everything: SQLite database, model cache, and staging files. To migrate between hosts:
+18 -6
View File
@@ -30,11 +30,13 @@ Returns JSON with ranked results combining full-text and semantic search.
**Flags:** **Flags:**
- `-n, --top N` — number of results (default: 10) - `-n, --top N` — number of results (default: 10)
- `--tags tag1,tag2` — filter by tags (AND logic) - `--tags tag1,tag2` — filter by tags (AND logic)
- `--type pdf|markdown|code|note` — filter by document type - `--type pdf|markdown|code|note|data` — filter by document type
- `--format json|human` — output format (always use json for parsing) - `--format json|human` — output format (always use json for parsing)
- `--fts-only` — keyword search only (skip semantic) - `--fts-only` — keyword search only (skip semantic)
- `--vec-only` — semantic search only (skip keyword) - `--vec-only` — semantic search only (skip keyword)
- `--threshold FLOAT` — minimum score cutoff - `--threshold FLOAT` — minimum score cutoff
- `--explain` — include a per-result score breakdown (FTS/vector scores and ranks, fusion contributions, rerank blend)
- `--no-rerank` — skip server-side cross-encoder reranking for lower latency (when the engine has it enabled)
## Adding files ## Adding files
@@ -45,7 +47,7 @@ kb addfile ~/docs/ --recursive # directory (recursive)
kb addfile ~/docs/ --recursive --tags reference # directory with tags kb addfile ~/docs/ --recursive --tags reference # directory with tags
``` ```
Supported file types: `.pdf`, `.docx`, `.html`, `.md`, `.txt`, `.py`, `.sh`, `.go`. Unsupported extensions are rejected before upload. Supported file types: `.pdf`, `.docx`, `.html`, `.md`, `.txt`, `.py`, `.sh`, `.go`, `.json`, `.yaml`, `.yml`, `.toml`. Unsupported extensions are rejected before upload. Data files (`.json`/`.yaml`/`.yml`/`.toml`) are ingested as text with doc type `data`; minified JSON is pretty-printed before chunking.
**Flags:** **Flags:**
- `--tags tag1,tag2` — tags (comma-separated) - `--tags tag1,tag2` — tags (comma-separated)
@@ -66,11 +68,17 @@ kb remove <doc_id> --yes # remove without confirmation
## Tag management ## Tag management
```bash ```bash
kb tags --format json # list all tags with counts kb tags --format json # list all tags with counts and descriptions
kb tag <doc_id> --add important,ops # add tags to a document kb tag <doc_id> --add important,ops # add tags to a document
kb tag <doc_id> --remove draft # remove tags from a document kb tag <doc_id> --remove draft # remove tags from a document
kb tag-describe ops "Lab operations runbooks" # set a tag context description
kb tag-describe ops # clear a tag's description
``` ```
Tag descriptions are returned as `tag_contexts` with every search result on a
document carrying the tag — use them to judge which of several similar-scoring
chunks actually answers the question.
## Bulk operations ## Bulk operations
Operate on multiple documents at once using filter-based selection. Filters combine with AND logic. Operate on multiple documents at once using filter-based selection. Filters combine with AND logic.
@@ -138,6 +146,7 @@ All commands support:
"results": [ "results": [
{ {
"chunk_id": 1423, "chunk_id": 1423,
"document_id": 87,
"score": 0.031, "score": 0.031,
"text": "To install the latest version of git from source...", "text": "To install the latest version of git from source...",
"chunk_index": 3, "chunk_index": 3,
@@ -146,11 +155,13 @@ All commands support:
"doc_type": "pdf", "doc_type": "pdf",
"source_path": "/home/user/docs/git-admin.pdf", "source_path": "/home/user/docs/git-admin.pdf",
"created_at": "2026-03-15T10:30:00", "created_at": "2026-03-15T10:30:00",
"tags": ["git", "admin"] "tags": ["git", "admin"],
"tag_contexts": {"admin": "System administration guides"}
} }
], ],
"total_matches": 47, "total_matches": 47,
"returned": 10 "returned": 10,
"reranked": true
} }
``` ```
@@ -217,7 +228,8 @@ If the kb engine is already running via Docker Compose, add the MCP server by de
## Important notes ## Important notes
- Always use `--format json` for machine parsing - Always use `--format json` for machine parsing
- The `score` field is relative, not absolute — compare scores within a result set - The `score` field is relative, not absolute — compare scores within a result set. Reranked hybrid scores (`"reranked": true`) are 0-1 blended values on a different scale from non-reranked RRF scores; don't compare across the two modes or apply `--threshold` expecting RRF-scale values on reranked output
- When the engine reranker is enabled, results are already cross-encoder reranked server-side — no need to rerank them yourself
- `chunk_metadata.page` is only present for PDF documents - `chunk_metadata.page` is only present for PDF documents
- `chunk_metadata.section_header` is only present for markdown documents with headers - `chunk_metadata.section_header` is only present for markdown documents with headers
- Results are already ranked by relevance (hybrid FTS + vector search) - Results are already ranked by relevance (hybrid FTS + vector search)
+1 -1
View File
@@ -1 +1 @@
3.2.0 3.3.0
+1 -1
View File
@@ -1 +1 @@
3.2.0 3.3.0
+54 -6
View File
@@ -8,6 +8,7 @@ import (
"path/filepath" "path/filepath"
"sort" "sort"
"strings" "strings"
"time"
"github.com/kb-search/kb/internal/api" "github.com/kb-search/kb/internal/api"
"github.com/kb-search/kb/internal/output" "github.com/kb-search/kb/internal/output"
@@ -38,6 +39,10 @@ var supportedExts = map[string]bool{
".py": true, ".py": true,
".sh": true, ".sh": true,
".go": true, ".go": true,
".json": true,
".yaml": true,
".yml": true,
".toml": true,
} }
var addfileCmd = &cobra.Command{ var addfileCmd = &cobra.Command{
@@ -50,12 +55,16 @@ var addfileCmd = &cobra.Command{
func init() { func init() {
addfileCmd.Flags().String("tags", "", "tags (comma-separated)") addfileCmd.Flags().String("tags", "", "tags (comma-separated)")
addfileCmd.Flags().BoolP("recursive", "r", false, "recursively add directory contents") addfileCmd.Flags().BoolP("recursive", "r", false, "recursively add directory contents")
addfileCmd.Flags().Bool("wait", false, "wait for ingestion to finish")
addfileCmd.Flags().Duration("wait-timeout", 10*time.Minute, "maximum time to wait per ingestion job")
rootCmd.AddCommand(addfileCmd) rootCmd.AddCommand(addfileCmd)
} }
func runAddfile(cmd *cobra.Command, args []string) error { func runAddfile(cmd *cobra.Command, args []string) error {
tags, _ := cmd.Flags().GetString("tags") tags, _ := cmd.Flags().GetString("tags")
recursive, _ := cmd.Flags().GetBool("recursive") recursive, _ := cmd.Flags().GetBool("recursive")
wait, _ := cmd.Flags().GetBool("wait")
timeout, _ := cmd.Flags().GetDuration("wait-timeout")
client := api.NewClient() client := api.NewClient()
@@ -85,12 +94,25 @@ func runAddfile(cmd *cobra.Command, args []string) error {
} }
if output.IsJSON() { if output.IsJSON() {
output.PrintJSON([]interface{}{result.Raw}) if !wait || result.Duplicate {
output.PrintJSON([]interface{}{result.Raw})
}
} else if result.Duplicate { } else if result.Duplicate {
fmt.Println(result.duplicateMsg()) fmt.Println(result.duplicateMsg())
} else { } else {
fmt.Printf("Queued: %s\n", filepath.Base(path)) fmt.Printf("Queued: %s\n", filepath.Base(path))
} }
if wait && !result.Duplicate {
job, err := waitForJob(client, int(result.JobID), timeout)
if err != nil {
return err
}
if output.IsJSON() {
output.PrintJSON([]interface{}{job})
} else {
fmt.Printf("Ingested: %s (doc ID: %d, chunks: %d)\n", filepath.Base(path), job.DocumentID, job.ChunkCount)
}
}
return nil return nil
} }
@@ -118,6 +140,7 @@ func runAddfile(cmd *cobra.Command, args []string) error {
} }
var results []interface{} var results []interface{}
var pending []*uploadResult
queued := 0 queued := 0
duplicates := 0 duplicates := 0
for _, f := range files { for _, f := range files {
@@ -126,7 +149,9 @@ func runAddfile(cmd *cobra.Command, args []string) error {
fmt.Fprintf(os.Stderr, "Error uploading %s: %v\n", f, err) fmt.Fprintf(os.Stderr, "Error uploading %s: %v\n", f, err)
continue continue
} }
results = append(results, result.Raw) if !wait || result.Duplicate {
results = append(results, result.Raw)
}
if result.Duplicate { if result.Duplicate {
duplicates++ duplicates++
if !output.IsJSON() { if !output.IsJSON() {
@@ -134,11 +159,25 @@ func runAddfile(cmd *cobra.Command, args []string) error {
} }
} else { } else {
queued++ queued++
pending = append(pending, result)
if !output.IsJSON() { if !output.IsJSON() {
fmt.Printf("Queued: %s\n", filepath.Base(f)) fmt.Printf("Queued: %s\n", filepath.Base(f))
} }
} }
} }
if wait {
for _, result := range pending {
job, err := waitForJob(client, int(result.JobID), timeout)
if err != nil {
return err
}
if output.IsJSON() {
results = append(results, job)
} else {
fmt.Printf("Ingested job %d (doc ID: %d, chunks: %d)\n", int(result.JobID), job.DocumentID, job.ChunkCount)
}
}
}
if output.IsJSON() { if output.IsJSON() {
output.PrintJSON(results) output.PrintJSON(results)
@@ -199,10 +238,19 @@ func uploadFile(client *api.Client, path, tags string) (*uploadResult, error) {
return nil, err return nil, err
} }
var result interface{} var raw json.RawMessage
if err := api.DecodeJSON(resp, &result); err != nil { if err := api.DecodeJSON(resp, &raw); err != nil {
return nil, fmt.Errorf("failed to decode response: %w", err) return nil, fmt.Errorf("failed to decode response: %w", err)
} }
return &uploadResult{Raw: result}, nil var queued struct {
JobID float64 `json:"job_id"`
}
if err := json.Unmarshal(raw, &queued); err != nil {
return nil, fmt.Errorf("failed to decode queued job: %w", err)
}
var result interface{}
if err := json.Unmarshal(raw, &result); err != nil {
return nil, fmt.Errorf("failed to decode queued response: %w", err)
}
return &uploadResult{Raw: result, JobID: queued.JobID}, nil
} }
+27 -5
View File
@@ -4,6 +4,7 @@ import (
"fmt" "fmt"
"net/http" "net/http"
"os" "os"
"time"
"github.com/kb-search/kb/internal/api" "github.com/kb-search/kb/internal/api"
"github.com/kb-search/kb/internal/output" "github.com/kb-search/kb/internal/output"
@@ -22,21 +23,25 @@ var addnoteCmd = &cobra.Command{
} }
return nil return nil
}, },
RunE: runAddnote, RunE: runAddnote,
} }
func init() { func init() {
addnoteCmd.Flags().String("tags", "", "tags (comma-separated)") addnoteCmd.Flags().String("tags", "", "tags (comma-separated)")
addnoteCmd.Flags().Bool("wait", false, "wait for ingestion to finish")
addnoteCmd.Flags().Duration("wait-timeout", 10*time.Minute, "maximum time to wait for ingestion")
rootCmd.AddCommand(addnoteCmd) rootCmd.AddCommand(addnoteCmd)
} }
func runAddnote(cmd *cobra.Command, args []string) error { func runAddnote(cmd *cobra.Command, args []string) error {
tags, _ := cmd.Flags().GetString("tags") tags, _ := cmd.Flags().GetString("tags")
wait, _ := cmd.Flags().GetBool("wait")
timeout, _ := cmd.Flags().GetDuration("wait-timeout")
client := api.NewClient() client := api.NewClient()
return submitNote(client, args[0], tags) return submitNote(client, args[0], tags, wait, timeout)
} }
func submitNote(client *api.Client, note, tags string) error { func submitNote(client *api.Client, note, tags string, wait bool, timeout time.Duration) error {
fields := map[string]string{ fields := map[string]string{
"note": note, "note": note,
} }
@@ -74,15 +79,32 @@ func submitNote(client *api.Client, note, tags string) error {
os.Exit(1) os.Exit(1)
} }
var result interface{} var result struct {
JobID int `json:"job_id"`
Status string `json:"status"`
Filename string `json:"filename"`
}
if err := api.DecodeJSON(resp, &result); err != nil { if err := api.DecodeJSON(resp, &result); err != nil {
return fmt.Errorf("failed to decode response: %w", err) return fmt.Errorf("failed to decode response: %w", err)
} }
if output.IsJSON() { if output.IsJSON() {
output.PrintJSON(result) if !wait {
output.PrintJSON(result)
}
} else { } else {
fmt.Println("Queued: note") fmt.Println("Queued: note")
} }
if wait {
job, err := waitForJob(client, result.JobID, timeout)
if err != nil {
return err
}
if output.IsJSON() {
output.PrintJSON(job)
} else {
fmt.Printf("Ingested: note (doc ID: %d, chunks: %d)\n", job.DocumentID, job.ChunkCount)
}
}
return nil return nil
} }
+376
View File
@@ -0,0 +1,376 @@
package cmd
import (
"encoding/json"
"fmt"
"os"
"strings"
"time"
"github.com/kb-search/kb/internal/api"
"github.com/kb-search/kb/internal/output"
"github.com/spf13/cobra"
)
var benchCmd = &cobra.Command{
Use: "bench <fixture.json>",
Short: "Benchmark search quality against a query fixture",
Long: `Run a fixture of queries with known-relevant documents against each
search backend (fts, vec, hybrid, rerank) and report precision@k, recall
and MRR per backend. Use this to baseline search quality before ranking
changes and to measure their effect. See docs/bench-example.json for the
fixture format.`,
Args: cobra.ExactArgs(1),
RunE: runBench,
}
func init() {
benchCmd.Flags().Int("top", 0, "override result count per query (k)")
benchCmd.Flags().String("backends", "fts,vec,hybrid,rerank", "comma-separated backends to run")
rootCmd.AddCommand(benchCmd)
}
type benchSelector struct {
DocumentID int64 `json:"document_id"`
SourcePath string `json:"source_path"`
TitleContains string `json:"title_contains"`
}
type benchQuery struct {
ID string `json:"id"`
Query string `json:"query"`
Tags []string `json:"tags"`
DocType string `json:"doc_type"`
Top int `json:"top"`
Relevant []benchSelector `json:"relevant"`
Notes string `json:"notes"`
}
type benchFixture struct {
Description string `json:"description"`
Top int `json:"top"`
Queries []benchQuery `json:"queries"`
}
type benchDoc struct {
DocumentID int64
Title string
SourcePath string
}
type queryResult struct {
ID string `json:"id"`
Precision float64 `json:"precision"`
Recall float64 `json:"recall"`
MRR float64 `json:"mrr"`
LatencyMS float64 `json:"latency_ms"`
Returned int `json:"returned"`
}
type backendResult struct {
Precision float64 `json:"precision"`
Recall float64 `json:"recall"`
MRR float64 `json:"mrr"`
AvgLatencyMS float64 `json:"avg_latency_ms"`
Queries []queryResult `json:"queries"`
}
func loadFixture(path string) (*benchFixture, error) {
data, err := os.ReadFile(path)
if err != nil {
return nil, fmt.Errorf("cannot read fixture: %w", err)
}
var fx benchFixture
if err := json.Unmarshal(data, &fx); err != nil {
return nil, fmt.Errorf("invalid fixture JSON: %w", err)
}
if len(fx.Queries) == 0 {
return nil, fmt.Errorf("fixture has no queries")
}
for i, q := range fx.Queries {
if q.Query == "" {
return nil, fmt.Errorf("query %d: missing query text", i)
}
if len(q.Relevant) == 0 {
return nil, fmt.Errorf("query %q: no relevant selectors", q.Query)
}
for j, sel := range q.Relevant {
set := 0
if sel.DocumentID != 0 {
set++
}
if sel.SourcePath != "" {
set++
}
if sel.TitleContains != "" {
set++
}
if set != 1 {
return nil, fmt.Errorf("query %q selector %d: exactly one of document_id, source_path, title_contains must be set", q.Query, j)
}
}
}
return &fx, nil
}
func (s benchSelector) matches(doc benchDoc) bool {
switch {
case s.DocumentID != 0:
return doc.DocumentID == s.DocumentID
case s.SourcePath != "":
return doc.SourcePath == s.SourcePath
case s.TitleContains != "":
return strings.Contains(strings.ToLower(doc.Title), strings.ToLower(s.TitleContains))
}
return false
}
// dedupeByDocument collapses ranked chunks into ranked documents, keeping the
// best (first) position for each document.
func dedupeByDocument(docs []benchDoc) []benchDoc {
seen := map[int64]bool{}
var out []benchDoc
for _, d := range docs {
if seen[d.DocumentID] {
continue
}
seen[d.DocumentID] = true
out = append(out, d)
}
return out
}
// scoreQuery computes document-level precision, recall and MRR for one
// query's ranked document list against the relevant-document selectors.
func scoreQuery(ranked []benchDoc, relevant []benchSelector) (precision, recall, mrr float64) {
if len(ranked) == 0 {
return 0, 0, 0
}
matchedDocs := 0
firstMatch := 0
selectorHit := make([]bool, len(relevant))
for i, doc := range ranked {
docMatched := false
for j, sel := range relevant {
if sel.matches(doc) {
docMatched = true
selectorHit[j] = true
}
}
if docMatched {
matchedDocs++
if firstMatch == 0 {
firstMatch = i + 1
}
}
}
selectorsMatched := 0
for _, hit := range selectorHit {
if hit {
selectorsMatched++
}
}
precision = float64(matchedDocs) / float64(len(ranked))
recall = float64(selectorsMatched) / float64(len(relevant))
if firstMatch > 0 {
mrr = 1.0 / float64(firstMatch)
}
return precision, recall, mrr
}
// rerankAvailable probes engine status for a loaded reranker.
func rerankAvailable(client *api.Client) bool {
resp, err := client.Get("/api/v1/status")
if err != nil {
return false
}
var status struct {
Rerank struct {
Enabled bool `json:"enabled"`
Loaded bool `json:"loaded"`
} `json:"rerank"`
}
if err := api.DecodeJSON(resp, &status); err != nil {
return false
}
return status.Rerank.Enabled && status.Rerank.Loaded
}
func benchSearch(client *api.Client, q benchQuery, backend string, top int) ([]benchDoc, float64, error) {
body := map[string]interface{}{
"query": q.Query,
"top": top,
}
if len(q.Tags) > 0 {
body["tags"] = q.Tags
}
if q.DocType != "" {
body["doc_type"] = q.DocType
}
switch backend {
case "fts":
body["fts_only"] = true
case "vec":
body["vec_only"] = true
case "hybrid":
body["rerank"] = false
case "rerank":
body["rerank"] = true
}
start := time.Now()
resp, err := client.Post("/api/v1/search", body)
if err != nil {
return nil, 0, err
}
if err := api.CheckError(resp); err != nil {
return nil, 0, err
}
var result struct {
Results []struct {
DocumentID int64 `json:"document_id"`
Title string `json:"title"`
SourcePath string `json:"source_path"`
} `json:"results"`
}
if err := api.DecodeJSON(resp, &result); err != nil {
return nil, 0, err
}
latency := float64(time.Since(start).Microseconds()) / 1000.0
var docs []benchDoc
for _, r := range result.Results {
docs = append(docs, benchDoc{DocumentID: r.DocumentID, Title: r.Title, SourcePath: r.SourcePath})
}
return dedupeByDocument(docs), latency, nil
}
func runBench(cmd *cobra.Command, args []string) error {
topFlag, _ := cmd.Flags().GetInt("top")
backendsFlag, _ := cmd.Flags().GetString("backends")
fx, err := loadFixture(args[0])
if err != nil {
return err
}
var backends []string
for _, b := range strings.Split(backendsFlag, ",") {
b = strings.TrimSpace(b)
if b == "" {
continue
}
switch b {
case "fts", "vec", "hybrid", "rerank":
backends = append(backends, b)
default:
return fmt.Errorf("unknown backend %q (valid: fts, vec, hybrid, rerank)", b)
}
}
if len(backends) == 0 {
return fmt.Errorf("no backends selected")
}
client := api.NewClient()
rerankSkipped := false
if contains(backends, "rerank") && !rerankAvailable(client) {
backends = remove(backends, "rerank")
rerankSkipped = true
if len(backends) == 0 {
return fmt.Errorf("reranking is not available on this engine (requires engine with reranker enabled)")
}
}
results := map[string]*backendResult{}
for _, backend := range backends {
br := &backendResult{}
for _, q := range fx.Queries {
top := 10
if fx.Top > 0 {
top = fx.Top
}
if q.Top > 0 {
top = q.Top
}
if topFlag > 0 {
top = topFlag
}
ranked, latency, err := benchSearch(client, q, backend, top)
if err != nil {
return fmt.Errorf("backend %s, query %q: %w", backend, q.Query, err)
}
p, r, m := scoreQuery(ranked, q.Relevant)
id := q.ID
if id == "" {
id = q.Query
}
br.Queries = append(br.Queries, queryResult{
ID: id, Precision: p, Recall: r, MRR: m,
LatencyMS: latency, Returned: len(ranked),
})
}
n := float64(len(br.Queries))
for _, qr := range br.Queries {
br.Precision += qr.Precision / n
br.Recall += qr.Recall / n
br.MRR += qr.MRR / n
br.AvgLatencyMS += qr.LatencyMS / n
}
results[backend] = br
}
if output.IsJSON() {
output.PrintJSON(map[string]interface{}{
"description": fx.Description,
"query_count": len(fx.Queries),
"backends": results,
"rerank_skipped": rerankSkipped,
})
return nil
}
if fx.Description != "" {
fmt.Printf("%s (%d queries)\n\n", fx.Description, len(fx.Queries))
}
headers := []string{"Backend", "Precision", "Recall", "MRR", "Avg ms"}
var rows [][]string
for _, backend := range backends {
br := results[backend]
rows = append(rows, []string{
backend,
fmt.Sprintf("%.3f", br.Precision),
fmt.Sprintf("%.3f", br.Recall),
fmt.Sprintf("%.3f", br.MRR),
fmt.Sprintf("%.0f", br.AvgLatencyMS),
})
}
output.PrintTable(headers, rows)
if rerankSkipped {
fmt.Println("\nrerank: n/a (reranker not enabled on this engine)")
}
return nil
}
func contains(ss []string, s string) bool {
for _, v := range ss {
if v == s {
return true
}
}
return false
}
func remove(ss []string, s string) []string {
var out []string
for _, v := range ss {
if v != s {
out = append(out, v)
}
}
return out
}
+143
View File
@@ -0,0 +1,143 @@
package cmd
import (
"math"
"os"
"path/filepath"
"testing"
)
func writeFixture(t *testing.T, content string) string {
t.Helper()
path := filepath.Join(t.TempDir(), "fixture.json")
if err := os.WriteFile(path, []byte(content), 0o644); err != nil {
t.Fatal(err)
}
return path
}
func TestLoadFixture_Valid(t *testing.T) {
path := writeFixture(t, `{
"description": "test",
"top": 5,
"queries": [
{"id": "q1", "query": "hello", "relevant": [{"document_id": 7}]}
]
}`)
fx, err := loadFixture(path)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if fx.Top != 5 || len(fx.Queries) != 1 || fx.Queries[0].Relevant[0].DocumentID != 7 {
t.Errorf("fixture parsed incorrectly: %+v", fx)
}
}
func TestLoadFixture_RejectsEmptyRelevant(t *testing.T) {
path := writeFixture(t, `{"queries": [{"query": "hello", "relevant": []}]}`)
if _, err := loadFixture(path); err == nil {
t.Error("expected error for query with no relevant selectors")
}
}
func TestLoadFixture_RejectsMultiFieldSelector(t *testing.T) {
path := writeFixture(t, `{"queries": [
{"query": "hello", "relevant": [{"document_id": 1, "source_path": "/x"}]}
]}`)
if _, err := loadFixture(path); err == nil {
t.Error("expected error for selector with two fields set")
}
}
func TestSelectorMatching(t *testing.T) {
doc := benchDoc{DocumentID: 42, Title: "M38T Owner's Manual", SourcePath: "/data/m38t.pdf"}
cases := []struct {
name string
sel benchSelector
want bool
}{
{"document_id match", benchSelector{DocumentID: 42}, true},
{"document_id miss", benchSelector{DocumentID: 43}, false},
{"source_path match", benchSelector{SourcePath: "/data/m38t.pdf"}, true},
{"source_path miss", benchSelector{SourcePath: "/data/other.pdf"}, false},
{"title_contains case-insensitive", benchSelector{TitleContains: "m38t owner"}, true},
{"title_contains miss", benchSelector{TitleContains: "workshop"}, false},
}
for _, tc := range cases {
if got := tc.sel.matches(doc); got != tc.want {
t.Errorf("%s: got %v, want %v", tc.name, got, tc.want)
}
}
}
func TestDedupeByDocument(t *testing.T) {
docs := []benchDoc{
{DocumentID: 1, Title: "a"},
{DocumentID: 2, Title: "b"},
{DocumentID: 1, Title: "a-again"},
{DocumentID: 3, Title: "c"},
}
out := dedupeByDocument(docs)
if len(out) != 3 || out[0].DocumentID != 1 || out[1].DocumentID != 2 || out[2].DocumentID != 3 {
t.Errorf("dedupe failed: %+v", out)
}
if out[0].Title != "a" {
t.Errorf("dedupe must keep first (best-ranked) occurrence, got %q", out[0].Title)
}
}
func approxEqual(a, b float64) bool {
return math.Abs(a-b) < 1e-9
}
func TestScoreQuery_HandComputed(t *testing.T) {
// Ranked docs: 10, 20, 30, 40. Relevant: 20 and 40.
ranked := []benchDoc{
{DocumentID: 10}, {DocumentID: 20}, {DocumentID: 30}, {DocumentID: 40},
}
relevant := []benchSelector{{DocumentID: 20}, {DocumentID: 40}}
p, r, m := scoreQuery(ranked, relevant)
if !approxEqual(p, 0.5) { // 2 of 4 returned docs are relevant
t.Errorf("precision: got %v, want 0.5", p)
}
if !approxEqual(r, 1.0) { // both relevant docs found
t.Errorf("recall: got %v, want 1.0", r)
}
if !approxEqual(m, 0.5) { // first relevant doc at rank 2
t.Errorf("mrr: got %v, want 0.5", m)
}
}
func TestScoreQuery_NoMatches(t *testing.T) {
ranked := []benchDoc{{DocumentID: 1}}
relevant := []benchSelector{{DocumentID: 99}}
p, r, m := scoreQuery(ranked, relevant)
if p != 0 || r != 0 || m != 0 {
t.Errorf("expected all-zero metrics, got p=%v r=%v mrr=%v", p, r, m)
}
}
func TestScoreQuery_PartialRecall(t *testing.T) {
// Only one of three relevant docs returned, at rank 1.
ranked := []benchDoc{{DocumentID: 5}, {DocumentID: 6}}
relevant := []benchSelector{{DocumentID: 5}, {DocumentID: 7}, {DocumentID: 8}}
p, r, m := scoreQuery(ranked, relevant)
if !approxEqual(p, 0.5) {
t.Errorf("precision: got %v, want 0.5", p)
}
if !approxEqual(r, 1.0/3.0) {
t.Errorf("recall: got %v, want 1/3", r)
}
if !approxEqual(m, 1.0) {
t.Errorf("mrr: got %v, want 1.0", m)
}
}
func TestScoreQuery_EmptyResults(t *testing.T) {
p, r, m := scoreQuery(nil, []benchSelector{{DocumentID: 1}})
if p != 0 || r != 0 || m != 0 {
t.Errorf("expected zeros for empty results, got p=%v r=%v mrr=%v", p, r, m)
}
}
+5 -1
View File
@@ -14,21 +14,25 @@ var examplesCmd = &cobra.Command{
fmt.Print(`Add notes: fmt.Print(`Add notes:
kb addnote "Remember to update DNS records" kb addnote "Remember to update DNS records"
kb addnote "Server room is building 3" --tags ops kb addnote "Server room is building 3" --tags ops
kb addnote "Deploy checklist" --wait
Add files: Add files:
kb addfile report.pdf kb addfile report.pdf
kb addfile ~/docs/ --recursive --tags reference kb addfile ~/docs/ --recursive --tags reference
kb addfile report.pdf --wait
Search: Search:
kb search "how to restart nginx" kb search "how to restart nginx"
kb search "deploy" --tags ops --top 5 kb search "deploy" --tags ops --top 5
kb find "quarterly report" --type pdf
Update notes: Update notes:
kb updatenote 42 "revised note content" kb updatenote 42 "revised note content"
Manage documents: Manage documents:
kb list --type pdf kb list --type pdf
kb info 3 kb list --filename report.pdf
kb info 3 --no-chunks
kb tag 3 --add important,ops kb tag 3 --add important,ops
kb remove 3 --yes kb remove 3 --yes
`) `)
+85
View File
@@ -0,0 +1,85 @@
package cmd
import (
"fmt"
"github.com/kb-search/kb/internal/api"
"github.com/kb-search/kb/internal/output"
"github.com/spf13/cobra"
)
var findCmd = &cobra.Command{
Use: "find <query>",
Short: "Find documents by their indexed content",
Args: cobra.ExactArgs(1),
RunE: runFind,
}
func init() {
findCmd.Flags().IntP("top", "n", 10, "number of documents to return")
findCmd.Flags().String("tags", "", "filter by tags (comma-separated)")
findCmd.Flags().String("type", "", "filter by document type")
rootCmd.AddCommand(findCmd)
}
func runFind(cmd *cobra.Command, args []string) error {
top, _ := cmd.Flags().GetInt("top")
tags, _ := cmd.Flags().GetString("tags")
docType, _ := cmd.Flags().GetString("type")
body := map[string]interface{}{"query": args[0], "top": top}
if tags != "" {
body["tags"] = splitTags(tags)
}
if docType != "" {
body["doc_type"] = docType
}
resp, err := api.NewClient().Post("/api/v1/documents/find", body)
if err != nil {
return err
}
if err := api.CheckError(resp); err != nil {
return err
}
if output.IsJSON() {
var raw interface{}
if err := api.DecodeJSON(resp, &raw); err != nil {
return fmt.Errorf("failed to decode response: %w", err)
}
output.PrintJSON(raw)
return nil
}
var docs []struct {
DocumentID int `json:"document_id"`
Title string `json:"title"`
Filename string `json:"original_filename"`
Type string `json:"doc_type"`
Score float64 `json:"score"`
HitCount int `json:"hit_count"`
TopChunk string `json:"top_chunk"`
}
if err := api.DecodeJSON(resp, &docs); err != nil {
return fmt.Errorf("failed to decode response: %w", err)
}
if len(docs) == 0 {
fmt.Println("No documents found.")
return nil
}
for i, doc := range docs {
preview := doc.TopChunk
if len(preview) > 200 {
preview = preview[:200] + "..."
}
fmt.Printf("\n%d. [%.4f] %s (doc:%d, hits:%d)\n", i+1, doc.Score, doc.Title, doc.DocumentID, doc.HitCount)
if doc.Filename != "" {
fmt.Printf(" Filename: %s\n", doc.Filename)
}
if doc.Type != "" {
fmt.Printf(" Type: %s\n", doc.Type)
}
fmt.Printf(" %s\n", preview)
}
fmt.Println()
return nil
}
+18 -11
View File
@@ -17,12 +17,18 @@ var infoCmd = &cobra.Command{
} }
func init() { func init() {
infoCmd.Flags().Bool("no-chunks", false, "return document metadata without chunk details")
rootCmd.AddCommand(infoCmd) rootCmd.AddCommand(infoCmd)
} }
func runInfo(cmd *cobra.Command, args []string) error { func runInfo(cmd *cobra.Command, args []string) error {
client := api.NewClient() client := api.NewClient()
resp, err := client.Get("/api/v1/documents/" + args[0]) noChunks, _ := cmd.Flags().GetBool("no-chunks")
path := "/api/v1/documents/" + args[0]
if noChunks {
path += "?include_chunks=false"
}
resp, err := client.Get(path)
if err != nil { if err != nil {
fmt.Fprintln(os.Stderr, err) fmt.Fprintln(os.Stderr, err)
os.Exit(1) os.Exit(1)
@@ -42,16 +48,17 @@ func runInfo(cmd *cobra.Command, args []string) error {
} }
var doc struct { var doc struct {
ID int `json:"id"` ID int `json:"id"`
Title string `json:"title"` Title string `json:"title"`
Type string `json:"doc_type"` Type string `json:"doc_type"`
Tags []string `json:"tags"` Tags []string `json:"tags"`
CreatedAt string `json:"created_at"` CreatedAt string `json:"created_at"`
UpdatedAt string `json:"updated_at"` UpdatedAt string `json:"updated_at"`
Chunks []struct { ChunkCount int `json:"chunk_count"`
ID int `json:"id"` Chunks []struct {
ID int `json:"id"`
Page interface{} `json:"page"` Page interface{} `json:"page"`
Section string `json:"section"` Section string `json:"section"`
} `json:"chunks"` } `json:"chunks"`
} }
if err := api.DecodeJSON(resp, &doc); err != nil { if err := api.DecodeJSON(resp, &doc); err != nil {
@@ -65,7 +72,7 @@ func runInfo(cmd *cobra.Command, args []string) error {
{"Tags", joinStrings(doc.Tags)}, {"Tags", joinStrings(doc.Tags)},
{"Created", doc.CreatedAt}, {"Created", doc.CreatedAt},
{"Updated", doc.UpdatedAt}, {"Updated", doc.UpdatedAt},
{"Chunks", fmt.Sprintf("%d", len(doc.Chunks))}, {"Chunks", fmt.Sprintf("%d", doc.ChunkCount)},
} }
output.PrintKeyValue(pairs) output.PrintKeyValue(pairs)
+18 -6
View File
@@ -13,18 +13,23 @@ import (
var listCmd = &cobra.Command{ var listCmd = &cobra.Command{
Use: "list", Use: "list",
Short: "List documents in the knowledge base", Short: "List documents in the knowledge base",
Args: cobra.NoArgs,
RunE: runList, RunE: runList,
} }
func init() { func init() {
listCmd.Flags().String("type", "", "filter by document type") listCmd.Flags().String("type", "", "filter by document type")
listCmd.Flags().String("tags", "", "filter by tags (comma-separated)") listCmd.Flags().String("tags", "", "filter by tags (comma-separated)")
listCmd.Flags().String("title", "", "filter by title substring")
listCmd.Flags().String("filename", "", "filter by original filename substring")
rootCmd.AddCommand(listCmd) rootCmd.AddCommand(listCmd)
} }
func runList(cmd *cobra.Command, args []string) error { func runList(cmd *cobra.Command, args []string) error {
docType, _ := cmd.Flags().GetString("type") docType, _ := cmd.Flags().GetString("type")
tags, _ := cmd.Flags().GetString("tags") tags, _ := cmd.Flags().GetString("tags")
title, _ := cmd.Flags().GetString("title")
filename, _ := cmd.Flags().GetString("filename")
params := url.Values{} params := url.Values{}
if docType != "" { if docType != "" {
@@ -33,6 +38,12 @@ func runList(cmd *cobra.Command, args []string) error {
if tags != "" { if tags != "" {
params.Set("tags", tags) params.Set("tags", tags)
} }
if title != "" {
params.Set("title", title)
}
if filename != "" {
params.Set("filename", filename)
}
path := "/api/v1/documents" path := "/api/v1/documents"
if len(params) > 0 { if len(params) > 0 {
@@ -60,10 +71,11 @@ func runList(cmd *cobra.Command, args []string) error {
} }
var docs []struct { var docs []struct {
ID int `json:"id"` ID int `json:"id"`
Title string `json:"title"` Title string `json:"title"`
Type string `json:"doc_type"` Filename string `json:"original_filename"`
Tags []string `json:"tags"` Type string `json:"doc_type"`
Tags []string `json:"tags"`
} }
if err := api.DecodeJSON(resp, &docs); err != nil { if err := api.DecodeJSON(resp, &docs); err != nil {
return fmt.Errorf("failed to decode response: %w", err) return fmt.Errorf("failed to decode response: %w", err)
@@ -74,10 +86,10 @@ func runList(cmd *cobra.Command, args []string) error {
return nil return nil
} }
headers := []string{"ID", "TITLE", "TYPE", "TAGS"} headers := []string{"ID", "TITLE", "FILENAME", "TYPE", "TAGS"}
var rows [][]string var rows [][]string
for _, d := range docs { for _, d := range docs {
rows = append(rows, []string{fmt.Sprintf("%d", d.ID), d.Title, d.Type, joinStrings(d.Tags)}) rows = append(rows, []string{fmt.Sprintf("%d", d.ID), d.Title, d.Filename, d.Type, joinStrings(d.Tags)})
} }
output.PrintTable(headers, rows) output.PrintTable(headers, rows)
return nil return nil
+13
View File
@@ -67,3 +67,16 @@ func TestAddnoteCmd_TooManyArgs_ReturnsError(t *testing.T) {
t.Errorf("expected 'accepts 1 arg' error, got: %s", errMsg) t.Errorf("expected 'accepts 1 arg' error, got: %s", errMsg)
} }
} }
func TestListCmd_PositionalArgReturnsError(t *testing.T) {
rootCmd.SetArgs([]string{"list", "ignored-title"})
err := rootCmd.Execute()
if err == nil {
t.Fatal("expected error for positional list argument, got nil")
}
if !strings.Contains(err.Error(), "unknown command") &&
!strings.Contains(err.Error(), "accepts 0 arg") {
t.Errorf("expected positional argument error, got: %s", err)
}
}
+73 -4
View File
@@ -23,6 +23,8 @@ func init() {
searchCmd.Flags().Bool("fts-only", false, "use full-text search only") searchCmd.Flags().Bool("fts-only", false, "use full-text search only")
searchCmd.Flags().Bool("vec-only", false, "use vector search only") searchCmd.Flags().Bool("vec-only", false, "use vector search only")
searchCmd.Flags().Float64("threshold", 0, "minimum score threshold") searchCmd.Flags().Float64("threshold", 0, "minimum score threshold")
searchCmd.Flags().Bool("explain", false, "include per-result score breakdown")
searchCmd.Flags().Bool("no-rerank", false, "skip cross-encoder reranking (lower latency)")
rootCmd.AddCommand(searchCmd) rootCmd.AddCommand(searchCmd)
} }
@@ -33,16 +35,18 @@ func runSearch(cmd *cobra.Command, args []string) error {
ftsOnly, _ := cmd.Flags().GetBool("fts-only") ftsOnly, _ := cmd.Flags().GetBool("fts-only")
vecOnly, _ := cmd.Flags().GetBool("vec-only") vecOnly, _ := cmd.Flags().GetBool("vec-only")
threshold, _ := cmd.Flags().GetFloat64("threshold") threshold, _ := cmd.Flags().GetFloat64("threshold")
explain, _ := cmd.Flags().GetBool("explain")
noRerank, _ := cmd.Flags().GetBool("no-rerank")
body := map[string]interface{}{ body := map[string]interface{}{
"query": args[0], "query": args[0],
"top": top, "top": top,
} }
if tags != "" { if tags != "" {
body["tags"] = tags body["tags"] = splitTags(tags)
} }
if docType != "" { if docType != "" {
body["type"] = docType body["doc_type"] = docType
} }
if ftsOnly { if ftsOnly {
body["fts_only"] = true body["fts_only"] = true
@@ -53,6 +57,12 @@ func runSearch(cmd *cobra.Command, args []string) error {
if threshold > 0 { if threshold > 0 {
body["threshold"] = threshold body["threshold"] = threshold
} }
if explain {
body["explain"] = true
}
if noRerank {
body["rerank"] = false
}
client := api.NewClient() client := api.NewClient()
resp, err := client.Post("/api/v1/search", body) resp, err := client.Post("/api/v1/search", body)
@@ -66,13 +76,17 @@ func runSearch(cmd *cobra.Command, args []string) error {
} }
var result struct { var result struct {
Results []struct { Reranked bool `json:"reranked"`
Results []struct {
Score float64 `json:"score"` Score float64 `json:"score"`
DocumentID int64 `json:"document_id"`
Title string `json:"title"` Title string `json:"title"`
DocType string `json:"doc_type"` DocType string `json:"doc_type"`
Tags []string `json:"tags"` Tags []string `json:"tags"`
TagContexts map[string]string `json:"tag_contexts"`
ChunkMetadata map[string]interface{} `json:"chunk_metadata"` ChunkMetadata map[string]interface{} `json:"chunk_metadata"`
Text string `json:"text"` Text string `json:"text"`
Explain map[string]interface{} `json:"explain"`
} `json:"results"` } `json:"results"`
} }
@@ -94,13 +108,17 @@ func runSearch(cmd *cobra.Command, args []string) error {
return nil return nil
} }
if result.Reranked {
fmt.Println("(reranked)")
}
for i, r := range result.Results { for i, r := range result.Results {
snippet := r.Text snippet := r.Text
if len(snippet) > 200 { if len(snippet) > 200 {
snippet = snippet[:200] + "..." snippet = snippet[:200] + "..."
} }
fmt.Printf("\n%d. [%.4f] %s\n", i+1, r.Score, r.Title) fmt.Printf("\n%d. [%.4f] %s (doc:%d)\n", i+1, r.Score, r.Title, r.DocumentID)
location := "" location := ""
if page, ok := r.ChunkMetadata["page"]; ok && page != nil { if page, ok := r.ChunkMetadata["page"]; ok && page != nil {
@@ -123,12 +141,63 @@ func runSearch(cmd *cobra.Command, args []string) error {
if len(r.Tags) > 0 { if len(r.Tags) > 0 {
fmt.Printf(" Tags: %s\n", joinStrings(r.Tags)) fmt.Printf(" Tags: %s\n", joinStrings(r.Tags))
} }
for _, tag := range r.Tags {
if desc, ok := r.TagContexts[tag]; ok && desc != "" {
fmt.Printf(" Context: %s — %s\n", tag, desc)
}
}
if r.Explain != nil {
fmt.Printf(" Score: %s\n", formatExplain(r.Explain))
}
fmt.Printf(" %s\n", snippet) fmt.Printf(" %s\n", snippet)
} }
fmt.Println() fmt.Println()
return nil return nil
} }
// formatExplain renders the engine's explain breakdown as one line, e.g.
// "fts score=4.213 rank=2 (rrf 0.016129) | vec score=0.512 rank=1 (rrf 0.016393) | bonus 0.05 | final 0.082522"
func formatExplain(e map[string]interface{}) string {
num := func(key string) (float64, bool) {
v, ok := e[key].(float64)
return v, ok
}
var parts []string
for _, arm := range []string{"fts", "vec"} {
score, hasScore := num(arm + "_score")
if !hasScore {
continue
}
part := fmt.Sprintf("%s score=%.4f", arm, score)
if rank, ok := num(arm + "_rank"); ok {
part += fmt.Sprintf(" rank=%d", int(rank))
}
if rrf, ok := num("rrf_" + arm); ok {
part += fmt.Sprintf(" (rrf %.6f)", rrf)
}
parts = append(parts, part)
}
if bonus, ok := num("bonus"); ok && bonus > 0 {
parts = append(parts, fmt.Sprintf("bonus %.2f", bonus))
}
for _, key := range []string{"pre_rerank_rank", "rerank_score", "blend_weight"} {
if v, ok := num(key); ok {
parts = append(parts, fmt.Sprintf("%s %.4f", key, v))
}
}
if final, ok := num("final_score"); ok {
parts = append(parts, fmt.Sprintf("final %.6f", final))
}
result := ""
for i, p := range parts {
if i > 0 {
result += " | "
}
result += p
}
return result
}
func joinStrings(ss []string) string { func joinStrings(ss []string) string {
result := "" result := ""
for i, s := range ss { for i, s := range ss {
+67
View File
@@ -0,0 +1,67 @@
package cmd
import (
"bytes"
"encoding/json"
"io"
"net/http"
"net/http/httptest"
"reflect"
"testing"
)
func TestSearchCmd_SendsDocTypeAndTagsList(t *testing.T) {
var captured map[string]interface{}
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/api/v1/search" {
t.Errorf("unexpected path: %s", r.URL.Path)
}
body, _ := io.ReadAll(r.Body)
if err := json.Unmarshal(body, &captured); err != nil {
t.Fatalf("failed to decode request body: %v", err)
}
w.Header().Set("Content-Type", "application/json")
w.Write([]byte(`{"query":"q","results":[],"total_matches":0,"returned":0}`))
}))
defer server.Close()
var stdout bytes.Buffer
rootCmd.SetOut(&stdout)
rootCmd.SetArgs([]string{
"search", "oil level",
"--engine", server.URL,
"--format", "json",
"--type", "pdf",
"--tags", "manuals, ops,",
})
if err := rootCmd.Execute(); err != nil {
t.Fatalf("search command failed: %v", err)
}
if captured == nil {
t.Fatal("no request captured by test server")
}
if got := captured["doc_type"]; got != "pdf" {
t.Errorf("expected doc_type=pdf in body, got %v (full body: %v)", got, captured)
}
if _, present := captured["type"]; present {
t.Error("body must not contain legacy 'type' key")
}
tags, ok := captured["tags"].([]interface{})
if !ok {
t.Fatalf("expected tags to be a JSON array, got %T (%v)", captured["tags"], captured["tags"])
}
want := []interface{}{"manuals", "ops"}
if !reflect.DeepEqual(tags, want) {
t.Errorf("expected tags %v, got %v", want, tags)
}
}
func TestSplitTags(t *testing.T) {
got := splitTags(" a ,b,, c ")
want := []string{"a", "b", "c"}
if !reflect.DeepEqual(got, want) {
t.Errorf("splitTags: expected %v, got %v", want, got)
}
}
+62
View File
@@ -0,0 +1,62 @@
package cmd
import (
"fmt"
"os"
"github.com/kb-search/kb/internal/api"
"github.com/kb-search/kb/internal/output"
"github.com/spf13/cobra"
)
var tagDescribeCmd = &cobra.Command{
Use: "tag-describe <tag> [description]",
Short: "Set a one-line context description on a tag",
Long: `Attach a short description to a tag (e.g. "Lab operations runbooks").
Descriptions are returned as tag_contexts with every search result on a
document carrying the tag, helping consumers judge relevance.
Omit the description (or pass "") to clear it.`,
Args: cobra.RangeArgs(1, 2),
RunE: runTagDescribe,
}
func init() {
rootCmd.AddCommand(tagDescribeCmd)
}
func runTagDescribe(cmd *cobra.Command, args []string) error {
description := ""
if len(args) == 2 {
description = args[1]
}
body := map[string]interface{}{"description": description}
client := api.NewClient()
resp, err := client.Put("/api/v1/tags/"+args[0]+"/description", body)
if err != nil {
fmt.Fprintln(os.Stderr, err)
os.Exit(1)
}
if err := api.CheckError(resp); err != nil {
fmt.Fprintln(os.Stderr, err)
os.Exit(1)
}
if output.IsJSON() {
var raw interface{}
if err := api.DecodeJSON(resp, &raw); err != nil {
return fmt.Errorf("failed to decode response: %w", err)
}
output.PrintJSON(raw)
return nil
}
if description == "" {
fmt.Printf("Description cleared for tag %q\n", args[0])
} else {
fmt.Printf("Tag %q: %s\n", args[0], description)
}
return nil
}
+5 -4
View File
@@ -41,8 +41,9 @@ func runTags(cmd *cobra.Command, args []string) error {
} }
var tags []struct { var tags []struct {
Name string `json:"name"` Name string `json:"name"`
Count int `json:"count"` Count int `json:"count"`
Description string `json:"description"`
} }
if err := api.DecodeJSON(resp, &tags); err != nil { if err := api.DecodeJSON(resp, &tags); err != nil {
return fmt.Errorf("failed to decode response: %w", err) return fmt.Errorf("failed to decode response: %w", err)
@@ -53,10 +54,10 @@ func runTags(cmd *cobra.Command, args []string) error {
return nil return nil
} }
headers := []string{"TAG", "COUNT"} headers := []string{"TAG", "COUNT", "DESCRIPTION"}
var rows [][]string var rows [][]string
for _, t := range tags { for _, t := range tags {
rows = append(rows, []string{t.Name, fmt.Sprintf("%d", t.Count)}) rows = append(rows, []string{t.Name, fmt.Sprintf("%d", t.Count), t.Description})
} }
output.PrintTable(headers, rows) output.PrintTable(headers, rows)
return nil return nil
+43
View File
@@ -0,0 +1,43 @@
package cmd
import (
"fmt"
"time"
"github.com/kb-search/kb/internal/api"
)
type jobStatus struct {
ID int `json:"id"`
Status string `json:"status"`
DocumentID int `json:"document_id"`
ChunkCount int `json:"chunk_count"`
Error string `json:"error"`
}
func waitForJob(client *api.Client, jobID int, timeout time.Duration) (*jobStatus, error) {
deadline := time.Now().Add(timeout)
for {
resp, err := client.Get(fmt.Sprintf("/api/v1/jobs/%d", jobID))
if err != nil {
return nil, err
}
if err := api.CheckError(resp); err != nil {
return nil, err
}
var job jobStatus
if err := api.DecodeJSON(resp, &job); err != nil {
return nil, fmt.Errorf("failed to decode job status: %w", err)
}
switch job.Status {
case "done", "skipped":
return &job, nil
case "failed":
return nil, fmt.Errorf("ingestion job %d failed: %s", jobID, job.Error)
}
if time.Now().After(deadline) {
return nil, fmt.Errorf("timed out waiting for ingestion job %d", jobID)
}
time.Sleep(time.Second)
}
}
+33
View File
@@ -0,0 +1,33 @@
{
"description": "Example search-quality fixture — copy and adapt against your own corpus",
"top": 10,
"queries": [
{
"id": "manual-lookup",
"query": "how do I check the oil level",
"doc_type": "pdf",
"relevant": [
{"document_id": 2077},
{"title_contains": "owner's manual"}
],
"notes": "Selectors are OR-matched; each counts as one relevant document. Exactly one of document_id / source_path / title_contains per selector."
},
{
"id": "infra-runbook",
"query": "restart the reverse proxy after certificate renewal",
"tags": ["ops"],
"relevant": [
{"source_path": "/data/notes/proxy-runbook.md"}
]
},
{
"id": "note-recall",
"query": "what did we decide about the backup retention window",
"top": 5,
"relevant": [
{"title_contains": "backup retention"}
],
"notes": "Per-query top overrides the fixture-level default; the --top flag overrides both."
}
]
}
+128
View File
@@ -0,0 +1,128 @@
<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="utf-8">
<meta name="viewport" content="width=device-width, initial-scale=1">
<title>kb-search Enhancements Proposal</title>
<style>
:root {
--bg: #f7f7f5; --fg: #1a1a1a; --muted: #666; --card: #fff;
--border: #ddd; --accent: #2563eb; --code-bg: #eef1f5;
}
@media (prefers-color-scheme: dark) {
:root {
--bg: #16181d; --fg: #e6e6e6; --muted: #9aa0a8; --card: #1e2128;
--border: #33363e; --accent: #7aa2f7; --code-bg: #262a33;
}
}
* { box-sizing: border-box; }
body {
margin: 0; padding: 2rem 1rem 4rem; background: var(--bg); color: var(--fg);
font: 16px/1.6 -apple-system, "Segoe UI", Roboto, sans-serif;
}
main { max-width: 860px; margin: 0 auto; }
h1 { font-size: 1.9rem; margin-bottom: .2rem; }
h2 { margin-top: 2.2rem; border-bottom: 1px solid var(--border); padding-bottom: .3rem; }
.meta { color: var(--muted); font-size: .9rem; margin-bottom: 2rem; }
.card {
background: var(--card); border: 1px solid var(--border); border-radius: 10px;
padding: 1rem 1.3rem; margin: 1rem 0;
}
.card h3 { margin: .2rem 0 .5rem; }
.badge {
display: inline-block; font-size: .72rem; font-weight: 600; letter-spacing: .03em;
padding: .12rem .55rem; border-radius: 999px; vertical-align: middle; margin-left: .5rem;
}
.b-high { background: #dc262622; color: #dc2626; }
.b-med { background: #d9770622; color: #d97706; }
.b-low { background: #05966922; color: #059669; }
table { border-collapse: collapse; width: 100%; margin: 1rem 0; font-size: .93rem; }
th, td { border: 1px solid var(--border); padding: .45rem .7rem; text-align: left; vertical-align: top; }
th { background: var(--code-bg); }
code { background: var(--code-bg); padding: .1rem .35rem; border-radius: 4px; font-size: .88em; }
pre { background: var(--code-bg); padding: .8rem 1rem; border-radius: 8px; overflow-x: auto; }
pre code { background: none; padding: 0; }
a { color: var(--accent); }
.muted { color: var(--muted); }
</style>
</head>
<body>
<main>
<h1>kb-search Enhancements Proposal</h1>
<p class="meta">Bobby &middot; 2026-07-07 &middot; Prompted by a review of <a href="https://github.com/tobi/qmd">tobi/qmd</a> (Tobi Lütke's local hybrid search engine)</p>
<h2>Summary</h2>
<p>qmd and kb-search v2 solve overlapping problems, but qmd's retrieval pipeline is measurably ahead: its own benchmarks show BM25-only at ~0.50, vector-only at ~0.70, and the full hybrid + reranked pipeline at ~1.00. Our kb does hybrid FTS + vector but stops there — no rank fusion, no reranking, no query expansion, and no way to measure whether a change helps or hurts. This proposal lists five enhancements, ordered by value-for-effort, plus the already-tracked JSON ingestion item.</p>
<table>
<tr><th>#</th><th>Enhancement</th><th>Impact</th><th>Effort</th></tr>
<tr><td>1</td><td>LLM reranking stage</td><td>High — biggest single search-quality lever</td><td>Medium</td></tr>
<tr><td>2</td><td>RRF fusion for FTS + vector merging</td><td>Medium-high</td><td>Low</td></tr>
<tr><td>3</td><td>Bench harness + <code>--explain</code> traces</td><td>High (enables everything else)</td><td>Low-medium</td></tr>
<tr><td>4</td><td>Context descriptions on tags/sources</td><td>Medium</td><td>Low</td></tr>
<tr><td>5</td><td>Query expansion</td><td>Medium</td><td>Medium-high</td></tr>
<tr><td>6</td><td>.json file ingestion (already tracked)</td><td>Medium</td><td>Low</td></tr>
</table>
<h2>Current state</h2>
<p>kb-search v2 (engine v3.2.2) runs on the RTX 4070 box with <code>BAAI/bge-base-en-v1.5</code> (768-dim). It holds ~2,310 documents (1,944 PDFs, 237 notes, 129 markdown) in ~123k chunks. Search is hybrid FTS + vector with a blended relative score. Strengths over qmd: binary ingestion (PDF/docx/HTML), tags, ingestion job queue, dedup, original export, and multi-client API access. The proposals below close the retrieval-quality gap without giving any of that up.</p>
<h2>Proposals</h2>
<div class="card">
<h3>1. LLM reranking stage <span class="badge b-high">HIGH IMPACT</span></h3>
<p>Add a cross-encoder reranking pass over the top-K hybrid candidates. qmd uses <code>qwen3-reranker-0.6b</code> (~640MB GGUF) — small enough to sit alongside bge on the 4070 permanently. Flow: hybrid retrieval pulls ~40 candidates → reranker scores each (query, chunk) pair → final order blends retrieval and reranker scores.</p>
<p>qmd's position-aware blend is worth copying wholesale: rank 13 keep 75% retrieval weight, 410 get 60%, 11+ get 40%. This stops the reranker destroying exact-match hits while letting it rescue mid-ranked semantic matches.</p>
<p class="muted">API: add <code>rerank: bool</code> (default true) to the search endpoint, with <code>--no-rerank</code> in the CLI for latency-sensitive callers.</p>
</div>
<div class="card">
<h3>2. RRF fusion <span class="badge b-med">MEDIUM-HIGH</span></h3>
<p>Replace the current score blend with Reciprocal Rank Fusion when merging FTS and vector lists: <code>score = Σ 1/(k + rank + 1)</code>, k=60. Rank-based fusion sidesteps the incomparability of BM25 scores (unbounded) and cosine similarity (01). qmd adds a top-rank bonus (+0.05 for #1, +0.02 for #23 in any list) to preserve exact matches — cheap and effective.</p>
<p class="muted">Pure engine-side change, no API impact. Scores become comparable across queries too, which fixes the "score is relative, not absolute" caveat in the current skill docs.</p>
</div>
<div class="card">
<h3>3. Bench harness + explain traces <span class="badge b-high">DO FIRST</span></h3>
<p>We currently have no way to know if any of the above helps. Add:</p>
<ul>
<li><code>kb bench fixture.json</code> — run a fixture of queries with known-relevant docs, report precision@k / recall / MRR per backend (fts-only, vec-only, hybrid, hybrid+rerank). Directly mirrors <code>qmd bench</code>.</li>
<li><code>--explain</code> on search — per-result score breakdown (FTS score, vector score, fusion contribution, rerank score).</li>
</ul>
<p>A fixture of 2030 real queries against the existing corpus (lab infra questions, manual lookups, note recall) gives a regression baseline before touching ranking. <strong>This should land before #1 and #2 so their benefit is provable.</strong></p>
</div>
<div class="card">
<h3>4. Context descriptions <span class="badge b-med">MEDIUM</span></h3>
<p>qmd's standout idea: attach a one-line description to a collection or path (e.g. "Meeting transcripts", "Lab infrastructure runbooks") and return it with every matching result. For kb, the natural unit is the <strong>tag</strong>: <code>kb tag-describe ops "Lab operations runbooks and procedures"</code>, returned as <code>tag_contexts</code> in search results. Helps an LLM consumer (me) judge which of several similar-scoring chunks actually answers the question — descriptions cost nothing at query time.</p>
</div>
<div class="card">
<h3>5. Query expansion <span class="badge b-low">LATER</span></h3>
<p>qmd fine-tuned a 1.7B model to generate 2 query variants, searching all three and fusing via RRF. Real quality gains, but the heaviest lift: another model resident in VRAM, ~12s latency, and much of the benefit is available cheaper — I already do multi-query decomposition client-side per the kb skill. Park until #1#3 have landed and the bench shows remaining headroom.</p>
</div>
<div class="card">
<h3>6. JSON ingestion <span class="badge b-low">TRACKED</span></h3>
<p>The original scope of this task: kb rejects <code>.json</code> uploads, forcing renames to <code>.txt</code>. Add <code>.json</code> (and sensibly <code>.yaml</code>/<code>.yml</code>/<code>.toml</code>) to the accepted extensions, ingesting as text. Optional nicety: pretty-print minified JSON before chunking so chunks break on structure.</p>
</div>
<h2>Suggested order</h2>
<ol>
<li><strong>#3 bench harness</strong> — establish the baseline (a weekend-sized job).</li>
<li><strong>#6 JSON support</strong> — small, independent, already promised.</li>
<li><strong>#2 RRF fusion</strong> — low-risk engine change, measure against baseline.</li>
<li><strong>#1 reranker</strong> — the big win, measured.</li>
<li><strong>#4 tag contexts</strong> — anytime, independent.</li>
<li><strong>#5 query expansion</strong> — only if the bench still shows a gap.</li>
</ol>
<h2>References</h2>
<ul>
<li><a href="https://github.com/tobi/qmd">tobi/qmd</a> — architecture, fusion weights, and bench design borrowed from here</li>
<li>qmd score fusion detail: RRF k=60, top-rank bonus +0.05/+0.02, position-aware blend 75/60/40% retrieval weight</li>
<li>Reranker model: <code>hf:ggml-org/Qwen3-Reranker-0.6B-Q8_0-GGUF</code> (~640MB)</li>
</ul>
</main>
</body>
</html>
+20 -4
View File
@@ -13,16 +13,32 @@ COPY --from=ghcr.io/astral-sh/uv:latest /uv /usr/local/bin/uv
WORKDIR /app WORKDIR /app
# Install CPU torch first, on its own, from the CPU index.
#
# Order matters: anything that depends on torch (sentence-transformers) will
# otherwise resolve the default CUDA build and pull ~2.7GB of nvidia-* wheels.
# Reinstalling torch afterwards replaces torch but leaves those wheels behind,
# orphaned and unused — which is how the CPU image ended up larger than the
# CUDA one. Installing CPU torch up front means nothing ever requests CUDA.
#
# Keeping it in its own layer also bounds the blob size: the registry drops
# uploads that take longer than 60s, so no single layer should approach ~3GB.
# Placing it before the source COPYs keeps this expensive layer cached when
# only application code changes.
RUN uv venv .venv && \
. .venv/bin/activate && \
UV_HTTP_TIMEOUT=600 uv pip install torch torchvision \
--index-url https://download.pytorch.org/whl/cpu
COPY pyproject.toml ./ COPY pyproject.toml ./
COPY kb/ kb/ COPY kb/ kb/
COPY main.py ./ COPY main.py ./
COPY VERSION ./ COPY VERSION ./
RUN uv venv .venv && \ # Remaining dependencies resolve against the CPU torch already present.
. .venv/bin/activate && \ RUN . .venv/bin/activate && \
uv pip install -e . && \
uv pip install "sentence-transformers[onnx]" && \ uv pip install "sentence-transformers[onnx]" && \
uv pip install --reinstall torch torchvision --index-url https://download.pytorch.org/whl/cpu uv pip install -e .
ENV PATH="/app/.venv/bin:$PATH" ENV PATH="/app/.venv/bin:$PATH"
ENV VIRTUAL_ENV="/app/.venv" ENV VIRTUAL_ENV="/app/.venv"
+12 -3
View File
@@ -13,14 +13,23 @@ COPY --from=ghcr.io/astral-sh/uv:latest /uv /usr/local/bin/uv
WORKDIR /app WORKDIR /app
# Install CUDA torch on its own, before the source COPYs.
#
# This is the bulk of the image (~2.8GiB compressed). Splitting it from the
# application install keeps it cached when only code changes, and keeps the
# app layer small. The registry drops any blob upload that takes longer than
# 60s, so this layer is deliberately the only large one.
RUN uv venv .venv && \
. .venv/bin/activate && \
UV_HTTP_TIMEOUT=600 uv pip install torch torchvision \
--index-url https://download.pytorch.org/whl/cu130
COPY pyproject.toml ./ COPY pyproject.toml ./
COPY kb/ kb/ COPY kb/ kb/
COPY main.py ./ COPY main.py ./
COPY VERSION ./ COPY VERSION ./
RUN uv venv .venv && \ RUN . .venv/bin/activate && \
. .venv/bin/activate && \
UV_HTTP_TIMEOUT=600 uv pip install torch torchvision --index-url https://download.pytorch.org/whl/cu130 && \
uv pip install -e . uv pip install -e .
ENV PATH="/app/.venv/bin:$PATH" ENV PATH="/app/.venv/bin:$PATH"
+1 -1
View File
@@ -1 +1 @@
3.2.1 3.3.0
+1
View File
@@ -13,6 +13,7 @@ services:
- KB_INGEST_DEVICE=cpu - KB_INGEST_DEVICE=cpu
- KB_API_KEY=${KB_API_KEY:-} - KB_API_KEY=${KB_API_KEY:-}
- KB_SEARCH_THRESHOLD=${KB_SEARCH_THRESHOLD:-0.01} - KB_SEARCH_THRESHOLD=${KB_SEARCH_THRESHOLD:-0.01}
- KB_MIN_CHUNK_ALNUM=${KB_MIN_CHUNK_ALNUM:-3}
- HF_HUB_OFFLINE=${HF_HUB_OFFLINE:-} - HF_HUB_OFFLINE=${HF_HUB_OFFLINE:-}
restart: unless-stopped restart: unless-stopped
+4
View File
@@ -21,6 +21,10 @@ services:
- KB_INGEST_DEVICE=${KB_INGEST_DEVICE:-auto} - KB_INGEST_DEVICE=${KB_INGEST_DEVICE:-auto}
- KB_API_KEY=${KB_API_KEY:-} - KB_API_KEY=${KB_API_KEY:-}
- KB_SEARCH_THRESHOLD=${KB_SEARCH_THRESHOLD:-0.01} - 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}
- HF_HUB_OFFLINE=${HF_HUB_OFFLINE:-} - HF_HUB_OFFLINE=${HF_HUB_OFFLINE:-}
restart: unless-stopped restart: unless-stopped
+4
View File
@@ -20,7 +20,11 @@ class Config:
self.ingest_device = os.environ.get("KB_INGEST_DEVICE", "auto") self.ingest_device = os.environ.get("KB_INGEST_DEVICE", "auto")
self.api_key = os.environ.get("KB_API_KEY") or None self.api_key = os.environ.get("KB_API_KEY") or None
self.search_threshold = float(os.environ.get("KB_SEARCH_THRESHOLD", "0.01")) self.search_threshold = float(os.environ.get("KB_SEARCH_THRESHOLD", "0.01"))
self.rerank_enabled = os.environ.get("KB_RERANK_ENABLED", "false").lower() in ("1", "true", "yes")
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.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.host = os.environ.get("KB_HOST", "0.0.0.0")
self.port = int(os.environ.get("KB_PORT", "8000")) self.port = int(os.environ.get("KB_PORT", "8000"))
+5
View File
@@ -195,6 +195,11 @@ def init_schema(conn: sqlite3.Connection, embedding_dim: int) -> None:
if "job_type" not in job_cols: if "job_type" not in job_cols:
conn.execute("ALTER TABLE jobs ADD COLUMN job_type TEXT DEFAULT 'ingest'") conn.execute("ALTER TABLE jobs ADD COLUMN job_type TEXT DEFAULT 'ingest'")
# Migrate: add description to tags if missing (tag contexts, v3.3.0)
tag_cols = {row[1] for row in conn.execute("PRAGMA table_info(tags)").fetchall()}
if "description" not in tag_cols:
conn.execute("ALTER TABLE tags ADD COLUMN description TEXT")
conn.commit() conn.commit()
+40
View File
@@ -0,0 +1,40 @@
"""Chunking pipeline for structured data files (JSON, YAML, TOML).
Data files are ingested as plain text. Minified JSON is pretty-printed
first so chunk boundaries fall on structural lines rather than mid-object.
Malformed input never fails ingestion it is chunked as-is.
"""
from __future__ import annotations
import json
from kb.ingest.code import _fixed_token_chunks
def chunk_data(
text: str,
language: str | None,
max_tokens: int = 1024,
) -> list[dict]:
"""Split a data file into chunks.
Returns a list of chunk dicts, each containing:
text, chunk_index, metadata
"""
if language == "json":
try:
text = json.dumps(json.loads(text), indent=2, ensure_ascii=False)
except (ValueError, TypeError):
pass # not valid JSON — ingest the raw text unchanged
chunks: list[dict] = []
for piece in _fixed_token_chunks(text, max_tokens):
piece = piece.strip()
if piece:
chunks.append({
"text": piece,
"chunk_index": len(chunks),
"metadata": {},
})
return chunks
+4
View File
@@ -11,6 +11,10 @@ SUPPORTED_EXTENSIONS: dict[str, tuple[str, str | None]] = {
".py": ("code", "python"), ".py": ("code", "python"),
".sh": ("code", "bash"), ".sh": ("code", "bash"),
".go": ("code", "go"), ".go": ("code", "go"),
".json": ("data", "json"),
".yaml": ("data", "yaml"),
".yml": ("data", "yaml"),
".toml": ("data", "toml"),
} }
+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 from docling_core.transforms.chunker.hierarchical_chunker import ( # noqa: E402
HierarchicalChunker, HierarchicalChunker,
) )
from kb.ingest.quality import has_minimum_content
def _fixed_size_chunks(text: str, max_chars: int = 2000) -> list[str]: 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( def chunk_document(
file_path: Path, file_path: Path,
ingest_device: str = "cpu", ingest_device: str = "cpu",
min_chunk_alnum: int = 3,
) -> list[dict]: ) -> list[dict]:
"""Convert and chunk a PDF/DOCX/HTML document using Docling. """Convert and chunk a PDF/DOCX/HTML document using Docling.
@@ -71,7 +73,7 @@ def chunk_document(
chunks: list[dict] = [] chunks: list[dict] = []
for idx, chunk in enumerate(raw_chunks): for idx, chunk in enumerate(raw_chunks):
text = chunk.text.strip() if hasattr(chunk, "text") else str(chunk).strip() 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 continue
metadata: dict = {} metadata: dict = {}
@@ -98,6 +100,8 @@ def chunk_document(
if not full_text and hasattr(doc, "text"): if not full_text and hasattr(doc, "text"):
full_text = doc.text full_text = doc.text
for idx, piece in enumerate(_fixed_size_chunks(full_text)): for idx, piece in enumerate(_fixed_size_chunks(full_text)):
if not has_minimum_content(piece, min_chunk_alnum):
continue
chunks.append({ chunks.append({
"text": piece, "text": piece,
"chunk_index": idx, "chunk_index": idx,
+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
+1
View File
@@ -0,0 +1 @@
"""Explicit maintenance commands for kb-engine data."""
@@ -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()
+63
View File
@@ -0,0 +1,63 @@
"""Cross-encoder reranker management.
Mirrors the embeddings module: one module-level model, loaded eagerly at
startup when reranking is enabled. Reranking is strictly optional search
degrades gracefully to plain hybrid retrieval when the model is absent.
"""
import logging
from typing import Optional
logger = logging.getLogger("kb.reranker")
_reranker: Optional[object] = None
_model_name: Optional[str] = None
def load_reranker(model_name: str, device: str = "cpu") -> None:
"""Load a cross-encoder reranking model.
Args:
model_name: HuggingFace model name or local path. Must have a
sequence-classification head (e.g. BAAI/bge-reranker-v2-m3).
device: Target device "cpu", "cuda", or "auto".
"""
global _reranker, _model_name
from sentence_transformers import CrossEncoder
from kb.embeddings import _resolve_device
resolved_device = _resolve_device(device)
logger.info("Loading reranker '%s' on device '%s'", model_name, resolved_device)
_reranker = CrossEncoder(model_name, device=resolved_device)
_model_name = model_name
logger.info("Reranker loaded: %s", model_name)
def is_available() -> bool:
"""Return True if a reranker model is loaded and usable."""
return _reranker is not None
def rerank_scores(query: str, texts: list[str]) -> list[float]:
"""Score (query, text) pairs with the cross-encoder.
Returns:
One relevance score per text, sigmoid-normalised to 0-1.
Raises:
RuntimeError: If no reranker has been loaded.
"""
if _reranker is None:
raise RuntimeError("Reranker not loaded. Call load_reranker() first.")
import numpy as np
scores = _reranker.predict([(query, t) for t in texts], convert_to_numpy=True)
# CrossEncoder heads may emit raw logits; squash to 0-1 so scores blend
# predictably with normalised retrieval scores. Sigmoid is monotonic, so
# ordering is unaffected for models that already output probabilities.
return (1.0 / (1.0 + np.exp(-np.asarray(scores, dtype="float64")))).tolist()
+67 -6
View File
@@ -7,11 +7,13 @@ from pathlib import Path
from typing import Optional from typing import Optional
from fastapi import HTTPException, Query from fastapi import HTTPException, Query
from pydantic import BaseModel, Field
from fastapi.responses import FileResponse from fastapi.responses import FileResponse
from main import app from main import app
from kb.config import cfg from kb.config import cfg
from kb.database import get_connection from kb.database import get_connection
from kb.search import hybrid_search
logger = logging.getLogger("kb.routes.documents") logger = logging.getLogger("kb.routes.documents")
@@ -20,11 +22,13 @@ logger = logging.getLogger("kb.routes.documents")
async def list_documents( async def list_documents(
type: Optional[str] = Query(None), type: Optional[str] = Query(None),
tags: 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) conn = get_connection(cfg.db_path)
try: try:
sql = """ 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, (SELECT COUNT(*) FROM chunks c WHERE c.document_id = d.id) AS chunk_count,
d.created_at, d.updated_at d.created_at, d.updated_at
FROM documents d FROM documents d
@@ -37,6 +41,14 @@ async def list_documents(
where.append("d.doc_type = ?") where.append("d.doc_type = ?")
params.append(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: if tags:
tag_list = [t.strip() for t in tags.split(",") if t.strip()] tag_list = [t.strip() for t in tags.split(",") if t.strip()]
for i, tag in enumerate(tag_list): for i, tag in enumerate(tag_list):
@@ -70,6 +82,8 @@ async def list_documents(
results.append({ results.append({
"id": row["id"], "id": row["id"],
"title": row["title"], "title": row["title"],
"original_filename": row["original_filename"],
"source_path": row["source_path"],
"doc_type": row["doc_type"], "doc_type": row["doc_type"],
"tags": [t["name"] for t in tag_rows], "tags": [t["name"] for t in tag_rows],
"chunk_count": row["chunk_count"], "chunk_count": row["chunk_count"],
@@ -82,8 +96,49 @@ async def list_documents(
conn.close() 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}") @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) conn = get_connection(cfg.db_path)
try: try:
doc = conn.execute( doc = conn.execute(
@@ -92,10 +147,15 @@ async def get_document(doc_id: int):
if not doc: if not doc:
raise HTTPException(status_code=404, detail="Document not found.") raise HTTPException(status_code=404, detail="Document not found.")
chunks = conn.execute( chunk_count = conn.execute(
"SELECT * FROM chunks WHERE document_id = ? ORDER BY chunk_index", "SELECT COUNT(*) AS n FROM chunks WHERE document_id = ?", (doc_id,)
(doc_id,), ).fetchone()["n"]
).fetchall() chunks = []
if include_chunks:
chunks = conn.execute(
"SELECT * FROM chunks WHERE document_id = ? ORDER BY chunk_index",
(doc_id,),
).fetchall()
tag_rows = conn.execute( tag_rows = conn.execute(
""" """
@@ -114,6 +174,7 @@ async def get_document(doc_id: int):
**dict(doc), **dict(doc),
"has_file": has_file, "has_file": has_file,
"tags": [t["name"] for t in tag_rows], "tags": [t["name"] for t in tag_rows],
"chunk_count": chunk_count,
"chunks": [dict(c) for c in chunks], "chunks": [dict(c) for c in chunks],
} }
finally: finally:
+3 -1
View File
@@ -10,6 +10,7 @@ from fastapi.responses import JSONResponse
from main import app from main import app
from kb.config import cfg from kb.config import cfg
from kb.database import get_connection, create_job, get_job, list_jobs, get_document_by_hash 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 from kb.staging import stage_file, stage_note
@@ -32,6 +33,7 @@ async def submit_job(
content_hash = hashlib.sha256(content).hexdigest() content_hash = hashlib.sha256(content).hexdigest()
filename = file.filename filename = file.filename
else: else:
title = title or auto_title(note) or "note"
content = note.encode("utf-8") content = note.encode("utf-8")
content_hash = hashlib.sha256(content).hexdigest() content_hash = hashlib.sha256(content).hexdigest()
filename = None filename = None
@@ -48,7 +50,7 @@ async def submit_job(
if file: if file:
staging_path = stage_file(cfg.staging_dir, file.filename, content) staging_path = stage_file(cfg.staging_dir, file.filename, content)
else: 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 filename = staging_path.name
tags_list = [t.strip() for t in tags.split(",") if t.strip()] if tags else [] tags_list = [t.strip() for t in tags.split(",") if t.strip()] if tags else []
+4
View File
@@ -19,6 +19,8 @@ class SearchRequest(BaseModel):
fts_only: bool = False fts_only: bool = False
vec_only: bool = False vec_only: bool = False
threshold: Optional[float] = None threshold: Optional[float] = None
explain: bool = False
rerank: Optional[bool] = None
@app.post("/api/v1/search") @app.post("/api/v1/search")
@@ -35,6 +37,8 @@ async def search(req: SearchRequest):
fts_only=req.fts_only, fts_only=req.fts_only,
vec_only=req.vec_only, vec_only=req.vec_only,
threshold=req.threshold, threshold=req.threshold,
explain=req.explain,
rerank=req.rerank,
) )
return result return result
except Exception as exc: except Exception as exc:
+7
View File
@@ -3,6 +3,7 @@
import os import os
from main import app, __version__ from main import app, __version__
from kb import reranker
from kb.config import cfg from kb.config import cfg
from kb.database import get_connection from kb.database import get_connection
from kb.embeddings import get_model_dim from kb.embeddings import get_model_dim
@@ -62,6 +63,12 @@ async def status():
"queued": queue_stats.get("queued", 0), "queued": queue_stats.get("queued", 0),
"processing": queue_stats.get("processing", 0), "processing": queue_stats.get("processing", 0),
}, },
"rerank": {
"enabled": cfg.rerank_enabled,
"model": cfg.reranker_model,
"loaded": reranker.is_available(),
"candidates": cfg.rerank_candidates,
},
} }
finally: finally:
conn.close() conn.close()
+33 -2
View File
@@ -16,14 +16,45 @@ async def list_tags():
try: try:
rows = conn.execute( rows = conn.execute(
""" """
SELECT t.name, COUNT(dt.document_id) AS count SELECT t.name, t.description, COUNT(dt.document_id) AS count
FROM tags t FROM tags t
LEFT JOIN document_tags dt ON t.id = dt.tag_id LEFT JOIN document_tags dt ON t.id = dt.tag_id
GROUP BY t.id, t.name GROUP BY t.id, t.name
ORDER BY t.name ORDER BY t.name
""" """
).fetchall() ).fetchall()
return [{"name": row["name"], "count": row["count"]} for row in rows] return [
{"name": row["name"], "count": row["count"], "description": row["description"]}
for row in rows
]
finally:
conn.close()
class TagDescriptionRequest(BaseModel):
description: Optional[str] = None
@app.put("/api/v1/tags/{name}/description")
async def set_tag_description(name: str, req: TagDescriptionRequest):
"""Set or clear a one-line context description on a tag.
Descriptions are returned as ``tag_contexts`` with every search result on
a document carrying the tag, helping consumers judge relevance.
"""
conn = get_connection(cfg.db_path)
try:
# name matching is case-insensitive (tags.name is COLLATE NOCASE)
tag = conn.execute("SELECT id FROM tags WHERE name = ?", (name,)).fetchone()
if not tag:
raise HTTPException(status_code=404, detail=f"Tag '{name}' not found.")
description = (req.description or "").strip() or None
conn.execute(
"UPDATE tags SET description = ? WHERE id = ?", (description, tag["id"])
)
conn.commit()
return {"name": name, "description": description}
finally: finally:
conn.close() conn.close()
+168 -17
View File
@@ -18,6 +18,8 @@ def hybrid_search(
fts_only: bool = False, fts_only: bool = False,
vec_only: bool = False, vec_only: bool = False,
threshold: float | None = None, threshold: float | None = None,
explain: bool = False,
rerank: bool | None = None,
) -> dict: ) -> dict:
"""Run hybrid search and return merged, enriched results. """Run hybrid search and return merged, enriched results.
@@ -31,11 +33,29 @@ def hybrid_search(
fts_only: Only use FTS5 (skip vector search). fts_only: Only use FTS5 (skip vector search).
vec_only: Only use vector search (skip FTS5). vec_only: Only use vector search (skip FTS5).
threshold: Optional minimum score; results below are dropped. threshold: Optional minimum score; results below are dropped.
explain: Attach a per-result score breakdown (arm scores, ranks,
RRF contributions, rerank blend) under an ``explain`` key.
rerank: Cross-encoder rerank the top candidates. None uses the
engine default (cfg.rerank_enabled); True still requires a
loaded reranker and degrades silently to plain retrieval
otherwise. Never applied to fts_only / vec_only searches.
Returns: Returns:
Dict with keys: query, results, total_matches, returned. Dict with keys: query, results, total_matches, returned, reranked.
""" """
from kb import reranker
want_rerank = cfg.rerank_enabled if rerank is None else rerank
do_rerank = (
want_rerank
and not fts_only
and not vec_only
and reranker.is_available()
)
candidate_count = top * 3 candidate_count = top * 3
if do_rerank:
candidate_count = max(candidate_count, cfg.rerank_candidates)
fts_results: dict[int, float] = {} fts_results: dict[int, float] = {}
vec_results: dict[int, float] = {} vec_results: dict[int, float] = {}
@@ -49,10 +69,12 @@ def hybrid_search(
# --- merge --------------------------------------------------------------- # --- merge ---------------------------------------------------------------
if fts_only: if fts_only:
merged = sorted(fts_results.items(), key=lambda x: x[1], reverse=True) merged = sorted(fts_results.items(), key=lambda x: x[1], reverse=True)
details = _single_arm_details("fts", fts_results)
elif vec_only: elif vec_only:
merged = sorted(vec_results.items(), key=lambda x: x[1], reverse=True) merged = sorted(vec_results.items(), key=lambda x: x[1], reverse=True)
details = _single_arm_details("vec", vec_results)
else: else:
merged = _rrf_merge(fts_results, vec_results) merged, details = _rrf_merge(fts_results, vec_results)
# Apply threshold filter — use config default if not specified per-query # Apply threshold filter — use config default if not specified per-query
effective_threshold = threshold if threshold is not None else cfg.search_threshold effective_threshold = threshold if threshold is not None else cfg.search_threshold
@@ -60,16 +82,30 @@ def hybrid_search(
merged = [(cid, score) for cid, score in merged if score >= effective_threshold] merged = [(cid, score) for cid, score in merged if score >= effective_threshold]
total_matches = len(merged) total_matches = len(merged)
# --- rerank --------------------------------------------------------------
# Blended scores are 0-1 normalised, a different scale from RRF scores;
# the threshold above was applied to RRF scores and is NOT re-applied.
reranked = False
if do_rerank and merged:
candidates = merged[: cfg.rerank_candidates]
rr_scores = reranker.rerank_scores(
query, _fetch_chunk_texts(conn, [cid for cid, _ in candidates])
)
merged = _blend_rerank(candidates, rr_scores, details)
reranked = True
merged = merged[:top] merged = merged[:top]
# --- enrich -------------------------------------------------------------- # --- enrich --------------------------------------------------------------
results = _enrich(conn, merged) results = _enrich(conn, merged, details if explain else None)
return { return {
"query": query, "query": query,
"results": results, "results": results,
"total_matches": total_matches, "total_matches": total_matches,
"returned": len(results), "returned": len(results),
"reranked": reranked,
} }
@@ -232,31 +268,135 @@ def _rrf_merge(
fts_results: dict[int, float], fts_results: dict[int, float],
vec_results: dict[int, float], vec_results: dict[int, float],
k: int = 60, k: int = 60,
) -> list[tuple[int, float]]: ) -> tuple[list[tuple[int, float]], dict[int, dict]]:
"""Reciprocal Rank Fusion over two scored result sets. """Reciprocal Rank Fusion over two scored result sets.
Each set is ranked independently (highest score first, rank starts at 1). Each set is ranked independently (highest score first, rank starts at 1).
RRF score for a document = sum of 1/(k + rank) across sets it appears in. RRF score for a document = sum of 1/(k + rank) across sets it appears in,
plus a top-rank bonus per arm: +0.05 for rank 1, +0.02 for ranks 2-3.
The bonus preserves exact matches a chunk at the top of either arm is
nearly impossible to displace via mid-rank RRF accumulation alone.
Score scale: base RRF maxes at 2/(k+1) 0.033 (k=60); with bonuses the
ceiling is 0.133. Bonuses only raise scores, so the threshold filter
(default 0.01) can never drop a result that base RRF would have kept.
Returns: Returns:
Sorted list of (chunk_id, rrf_score), highest first. (scores, details) scores is a sorted list of (chunk_id, rrf_score),
highest first; details maps chunk_id to a per-arm score breakdown
suitable for the ``explain`` response field.
""" """
fts_ranked = _rank_by_score(fts_results) fts_ranked = _rank_by_score(fts_results)
vec_ranked = _rank_by_score(vec_results) vec_ranked = _rank_by_score(vec_results)
all_ids = set(fts_ranked) | set(vec_ranked) all_ids = set(fts_ranked) | set(vec_ranked)
scores: list[tuple[int, float]] = [] scores: list[tuple[int, float]] = []
details: dict[int, dict] = {}
for chunk_id in all_ids: for chunk_id in all_ids:
rrf = 0.0 fts_rank = fts_ranked.get(chunk_id)
if chunk_id in fts_ranked: vec_rank = vec_ranked.get(chunk_id)
rrf += 1.0 / (k + fts_ranked[chunk_id]) rrf_fts = 1.0 / (k + fts_rank) if fts_rank is not None else None
if chunk_id in vec_ranked: rrf_vec = 1.0 / (k + vec_rank) if vec_rank is not None else None
rrf += 1.0 / (k + vec_ranked[chunk_id]) bonus = _top_rank_bonus(fts_rank) + _top_rank_bonus(vec_rank)
rrf = (rrf_fts or 0.0) + (rrf_vec or 0.0) + bonus
details[chunk_id] = {
"fts_score": _round6(fts_results.get(chunk_id)),
"fts_rank": fts_rank,
"vec_score": _round6(vec_results.get(chunk_id)),
"vec_rank": vec_rank,
"rrf_fts": _round6(rrf_fts),
"rrf_vec": _round6(rrf_vec),
"bonus": bonus,
"final_score": _round6(rrf),
}
scores.append((chunk_id, rrf)) scores.append((chunk_id, rrf))
scores.sort(key=lambda x: x[1], reverse=True) scores.sort(key=lambda x: x[1], reverse=True)
return scores return scores, details
def _top_rank_bonus(rank: int | None) -> float:
"""Bonus for appearing at the top of one arm's ranking."""
if rank == 1:
return 0.05
if rank in (2, 3):
return 0.02
return 0.0
def _single_arm_details(arm: str, results: dict[int, float]) -> dict[int, dict]:
"""Explain details for fts_only / vec_only searches (raw arm scores)."""
ranked = _rank_by_score(results)
return {
chunk_id: {
f"{arm}_score": _round6(score),
f"{arm}_rank": ranked[chunk_id],
"final_score": _round6(score),
}
for chunk_id, score in results.items()
}
def _round6(value: float | None) -> float | None:
return round(value, 6) if value is not None else None
def _fetch_chunk_texts(conn: sqlite3.Connection, chunk_ids: list[int]) -> list[str]:
"""Fetch chunk texts in the same order as *chunk_ids*."""
placeholders = ",".join("?" * len(chunk_ids))
rows = conn.execute(
f"SELECT id, text FROM chunks WHERE id IN ({placeholders})", chunk_ids
).fetchall()
by_id = {row[0]: row[1] for row in rows}
return [by_id.get(cid, "") for cid in chunk_ids]
def _blend_rerank(
candidates: list[tuple[int, float]],
rr_scores: list[float],
details: dict[int, dict],
) -> list[tuple[int, float]]:
"""Blend retrieval and cross-encoder scores, position-aware.
Retrieval scores are min-max normalised within the candidate set; rerank
scores are already 0-1. The retrieval weight depends on pre-rerank rank
75% for ranks 1-3, 60% for 4-10, 40% for 11+ so the reranker can rescue
mid-ranked semantic matches without destroying top exact-match hits.
*candidates* must be in retrieval order; *rr_scores* aligned with it.
Mutates *details* with the blend breakdown. Returns (chunk_id, blended)
sorted highest first.
"""
retrieval = [score for _, score in candidates]
lo, hi = min(retrieval), max(retrieval)
span = hi - lo
blended: list[tuple[int, float]] = []
for i, ((chunk_id, score), rr) in enumerate(zip(candidates, rr_scores)):
rank = i + 1
norm = (score - lo) / span if span > 0 else 1.0
if rank <= 3:
weight = 0.75
elif rank <= 10:
weight = 0.60
else:
weight = 0.40
final = weight * norm + (1.0 - weight) * rr
if chunk_id in details:
details[chunk_id].update({
"pre_rerank_rank": rank,
"retrieval_norm": _round6(norm),
"rerank_score": _round6(rr),
"blend_weight": weight,
"final_score": _round6(final),
})
blended.append((chunk_id, final))
blended.sort(key=lambda x: x[1], reverse=True)
return blended
def _rank_by_score(results: dict[int, float]) -> dict[int, int]: def _rank_by_score(results: dict[int, float]) -> dict[int, int]:
@@ -268,8 +408,13 @@ def _rank_by_score(results: dict[int, float]) -> dict[int, int]:
def _enrich( def _enrich(
conn: sqlite3.Connection, conn: sqlite3.Connection,
merged: list[tuple[int, float]], merged: list[tuple[int, float]],
details: dict[int, dict] | None = None,
) -> list[dict]: ) -> list[dict]:
"""Fetch chunk text, document metadata, chunk metadata, and tags.""" """Fetch chunk text, document metadata, chunk metadata, and tags.
When *details* is given, each result gains an ``explain`` key with its
score breakdown.
"""
results: list[dict] = [] results: list[dict] = []
for chunk_id, score in merged: for chunk_id, score in merged:
@@ -277,7 +422,7 @@ def _enrich(
""" """
SELECT c.id, c.text, c.chunk_index, c.metadata AS chunk_meta, 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.id AS doc_id, d.title, d.doc_type, d.source_path,
d.created_at d.created_at, d.original_filename
FROM chunks c FROM chunks c
JOIN documents d ON c.document_id = d.id JOIN documents d ON c.document_id = d.id
WHERE c.id = ? WHERE c.id = ?
@@ -292,7 +437,7 @@ def _enrich(
tag_rows = conn.execute( tag_rows = conn.execute(
""" """
SELECT t.name FROM tags t SELECT t.name, t.description FROM tags t
JOIN document_tags dt ON t.id = dt.tag_id JOIN document_tags dt ON t.id = dt.tag_id
WHERE dt.document_id = ? WHERE dt.document_id = ?
ORDER BY t.name ORDER BY t.name
@@ -300,8 +445,9 @@ def _enrich(
(row[4],), # doc_id (row[4],), # doc_id
).fetchall() ).fetchall()
results.append({ result = {
"chunk_id": row[0], "chunk_id": row[0],
"document_id": row[4],
"score": round(score, 6), "score": round(score, 6),
"text": row[1], "text": row[1],
"chunk_index": row[2], "chunk_index": row[2],
@@ -310,7 +456,12 @@ def _enrich(
"doc_type": row[6], "doc_type": row[6],
"source_path": row[7], "source_path": row[7],
"created_at": row[8], "created_at": row[8],
"original_filename": row[9],
"tags": [t[0] for t in tag_rows], "tags": [t[0] for t in tag_rows],
}) "tag_contexts": {t[0]: t[1] for t in tag_rows if t[1]},
}
if details is not None and row[0] in details:
result["explain"] = details[row[0]]
results.append(result)
return results return results
+9 -1
View File
@@ -113,7 +113,9 @@ def _process_job(job_row) -> tuple[str, int | None, int]:
chunks = chunk_note(text) chunks = chunk_note(text)
elif doc_type == "pdf": elif doc_type == "pdf":
from kb.ingest.docling_pipeline import chunk_document 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": elif doc_type == "markdown":
text = staged_path.read_text(encoding="utf-8") text = staged_path.read_text(encoding="utf-8")
from kb.ingest.markdown import chunk_markdown from kb.ingest.markdown import chunk_markdown
@@ -124,6 +126,12 @@ def _process_job(job_row) -> tuple[str, int | None, int]:
_, language = detector.detect_type(Path(filename)) _, language = detector.detect_type(Path(filename))
from kb.ingest.code import chunk_code from kb.ingest.code import chunk_code
chunks = chunk_code(text, language) chunks = chunk_code(text, language)
elif doc_type == "data":
text = staged_path.read_text(encoding="utf-8")
if not language:
_, language = detector.detect_type(Path(filename))
from kb.ingest.data import chunk_data
chunks = chunk_data(text, language)
else: else:
raise ValueError(f"Unsupported doc_type: {doc_type}") raise ValueError(f"Unsupported doc_type: {doc_type}")
+12
View File
@@ -40,6 +40,18 @@ async def lifespan(app: FastAPI):
init_schema(conn, model_dim) init_schema(conn, model_dim)
conn.close() conn.close()
# Optional reranker — search degrades gracefully if this fails
if cfg.rerank_enabled:
from kb.reranker import load_reranker
try:
load_reranker(cfg.reranker_model, cfg.device)
except Exception:
log.warning(
"Failed to load reranker '%s' — searches will not be reranked",
cfg.reranker_model,
exc_info=True,
)
# Start background ingestion worker # Start background ingestion worker
worker_task = asyncio.create_task(ingestion_worker()) worker_task = asyncio.create_task(ingestion_worker())
+6
View File
@@ -0,0 +1,6 @@
"""Shared test setup — make the engine root importable (for ``import kb``)."""
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+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",
}]
+64
View File
@@ -0,0 +1,64 @@
"""Tests for structured-data (.json/.yaml/.toml) ingestion."""
import json
from pathlib import Path
import pytest
from kb.ingest.data import chunk_data
from kb.ingest.detector import detect_type, is_supported
def test_detector_accepts_data_extensions():
assert detect_type(Path("config.json")) == ("data", "json")
assert detect_type(Path("stack.yaml")) == ("data", "yaml")
assert detect_type(Path("stack.yml")) == ("data", "yaml")
assert detect_type(Path("pyproject.toml")) == ("data", "toml")
for name in ("a.json", "b.yaml", "c.yml", "d.toml"):
assert is_supported(Path(name))
def test_minified_json_is_pretty_printed():
minified = json.dumps({"hosts": [{"name": "web1", "ip": "10.0.0.1"}]})
assert "\n" not in minified
chunks = chunk_data(minified, "json")
assert len(chunks) == 1
text = chunks[0]["text"]
assert "\n" in text, "expected pretty-printed multi-line JSON"
assert '"name": "web1"' in text
def test_large_minified_json_multi_chunks():
big = json.dumps([{"id": i, "payload": "x" * 100} for i in range(200)])
chunks = chunk_data(big, "json", max_tokens=256)
assert len(chunks) > 1, "large JSON must split into multiple chunks"
# Pretty-printing means chunks break on lines, not mid-token blobs.
for c in chunks:
assert c["text"].strip()
assert [c["chunk_index"] for c in chunks] == list(range(len(chunks)))
def test_malformed_json_ingests_raw():
broken = '{"unterminated": [1, 2'
chunks = chunk_data(broken, "json")
assert len(chunks) == 1
assert chunks[0]["text"] == broken
def test_yaml_passes_through_unchanged():
yaml_text = "services:\n web:\n image: nginx\n"
chunks = chunk_data(yaml_text, "yaml")
assert len(chunks) == 1
assert chunks[0]["text"] == yaml_text.strip()
def test_toml_passes_through():
toml_text = '[tool.example]\nname = "kb"\n'
chunks = chunk_data(toml_text, "toml")
assert len(chunks) == 1
assert chunks[0]["text"] == toml_text.strip()
def test_empty_file_yields_no_chunks():
assert chunk_data("", "json") == []
assert chunk_data(" \n ", "yaml") == []
+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)
+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
+273
View File
@@ -0,0 +1,273 @@
"""Tests for hybrid search: explain traces, document_id, and RRF merging."""
import sys
import types
import pytest
from kb.database import (
get_connection,
init_schema,
insert_chunk,
insert_document,
insert_embedding,
tag_document,
)
from kb.search import _blend_rerank, _rank_by_score, _rrf_merge, hybrid_search
DIM = 4
# Chunk vectors are axis-aligned so we can steer vector ranking exactly:
# a query of [1,0,0,0] has distance 0 to chunk A, sqrt(2) to chunk B.
VEC_A = [1.0, 0.0, 0.0, 0.0]
VEC_B = [0.0, 1.0, 0.0, 0.0]
QUERY_VEC = [1.0, 0.0, 0.0, 0.0]
class _Cfg:
search_threshold = 0.0
rerank_enabled = False
reranker_model = "test-reranker"
rerank_candidates = 40
class _Db:
def __init__(self, conn, ids):
self.conn = conn
self.ids = ids
@pytest.fixture
def db(tmp_path, monkeypatch):
"""Real schema (FTS5 + sqlite-vec) with two docs and stubbed embeddings."""
fake = types.ModuleType("kb.embeddings")
fake.embed_texts = lambda texts: [QUERY_VEC for _ in texts]
monkeypatch.setitem(sys.modules, "kb.embeddings", fake)
conn = get_connection(str(tmp_path / "kb.db"))
init_schema(conn, embedding_dim=DIM)
doc_a = insert_document(conn, "Alpha doc", "/src/a.md", "hash-a", "markdown")
doc_b = insert_document(conn, "Bravo doc", "/src/b.md", "hash-b", "markdown")
chunk_a = insert_chunk(conn, doc_a, 0, "alpha network switch configuration")
chunk_b = insert_chunk(conn, doc_b, 0, "bravo unrelated cooking recipe")
insert_embedding(conn, chunk_a, VEC_A)
insert_embedding(conn, chunk_b, VEC_B)
tag_document(conn, doc_a, ["ops"])
ids = {"doc_a": doc_a, "doc_b": doc_b, "chunk_a": chunk_a, "chunk_b": chunk_b}
yield _Db(conn, ids)
conn.close()
def test_results_include_document_id(db):
result = hybrid_search(db.conn, "alpha switch", _Cfg())
assert result["results"], "expected at least one hit"
top_hit = result["results"][0]
assert top_hit["document_id"] == db.ids["doc_a"]
assert top_hit["chunk_id"] == db.ids["chunk_a"]
def test_explain_absent_by_default(db):
result = hybrid_search(db.conn, "alpha switch", _Cfg())
assert all("explain" not in r for r in result["results"])
def test_explain_hybrid_breakdown(db):
result = hybrid_search(db.conn, "alpha switch", _Cfg(), explain=True)
top_hit = result["results"][0]
exp = top_hit["explain"]
# Chunk A is rank 1 in both arms: FTS matches "alpha"/"switch", vector
# distance is 0 (similarity 1.0).
assert exp["fts_rank"] == 1
assert exp["vec_rank"] == 1
assert exp["fts_score"] > 0
assert exp["vec_score"] == pytest.approx(1.0)
assert exp["rrf_fts"] == pytest.approx(1.0 / 61, abs=1e-6)
assert exp["rrf_vec"] == pytest.approx(1.0 / 61, abs=1e-6)
assert exp["final_score"] == pytest.approx(
exp["rrf_fts"] + exp["rrf_vec"] + exp["bonus"], abs=1e-5
)
assert exp["final_score"] == pytest.approx(top_hit["score"], abs=1e-5)
def test_explain_single_arm_when_vec_misses(db):
"""A chunk found only by vector search has null FTS fields."""
result = hybrid_search(db.conn, "zzz-no-fts-match", _Cfg(), explain=True)
for r in result["results"]:
exp = r["explain"]
assert exp["fts_score"] is None
assert exp["fts_rank"] is None
assert exp["rrf_fts"] is None
assert exp["vec_rank"] is not None
def test_explain_fts_only_shape(db):
result = hybrid_search(db.conn, "alpha switch", _Cfg(), fts_only=True, explain=True)
top_hit = result["results"][0]
exp = top_hit["explain"]
assert exp["fts_rank"] == 1
assert exp["final_score"] == pytest.approx(exp["fts_score"])
assert "vec_score" not in exp
def test_explain_vec_only_shape(db):
result = hybrid_search(db.conn, "anything", _Cfg(), vec_only=True, explain=True)
top_hit = result["results"][0]
exp = top_hit["explain"]
assert exp["vec_rank"] == 1
assert exp["final_score"] == pytest.approx(exp["vec_score"])
assert "fts_score" not in exp
def test_rrf_merge_arithmetic():
fts = {1: 10.0, 2: 5.0}
vec = {2: 0.9, 3: 0.8}
scores, details = _rrf_merge(fts, vec)
by_id = dict(scores)
# fts rank 1 → 1/61 + 0.05 bonus
assert by_id[1] == pytest.approx(1 / 61 + 0.05)
# fts rank 2 (+0.02) and vec rank 1 (+0.05) — bonuses stack across arms
assert by_id[2] == pytest.approx(1 / 62 + 1 / 61 + 0.07)
# vec rank 2 → 1/62 + 0.02
assert by_id[3] == pytest.approx(1 / 62 + 0.02)
# Chunk 2 appears in both arms, so it must win.
assert scores[0][0] == 2
assert details[2]["fts_rank"] == 2
assert details[2]["vec_rank"] == 1
assert details[2]["bonus"] == pytest.approx(0.07)
assert details[1]["vec_rank"] is None
assert details[3]["rrf_fts"] is None
def test_top_rank_bonus_tiers():
from kb.search import _top_rank_bonus
assert _top_rank_bonus(1) == 0.05
assert _top_rank_bonus(2) == 0.02
assert _top_rank_bonus(3) == 0.02
assert _top_rank_bonus(4) == 0.0
assert _top_rank_bonus(None) == 0.0
def test_bonus_preserves_top_of_arm():
"""A chunk at rank 1 of one arm beats a chunk at mid-rank in both arms."""
fts = {10: 100.0, 11: 90.0, 12: 80.0, 13: 70.0, 14: 60.0}
vec = {20: 0.9, 11: 0.8, 12: 0.7, 13: 0.6, 14: 0.5}
scores, _ = _rrf_merge(fts, vec)
order = [cid for cid, _ in scores]
# Chunk 10 (fts #1, absent from vec): 1/61 + 0.05 ≈ 0.0664.
# Chunk 13 (rank 4 in fts, rank 4 in vec): 2/64 ≈ 0.031, no bonus.
assert order.index(10) < order.index(13)
def test_rank_by_score():
assert _rank_by_score({7: 0.5, 8: 0.9, 9: 0.1}) == {8: 1, 7: 2, 9: 3}
# ---------------------------------------------------------------------------
# Reranking
# ---------------------------------------------------------------------------
def _enable_fake_reranker(monkeypatch, score_fn):
from kb import reranker
monkeypatch.setattr(reranker, "is_available", lambda: True)
monkeypatch.setattr(reranker, "rerank_scores", score_fn)
def test_rerank_flags_response_and_explain(db, monkeypatch):
_enable_fake_reranker(monkeypatch, lambda q, texts: [0.9] * len(texts))
result = hybrid_search(db.conn, "alpha switch", _Cfg(), explain=True, rerank=True)
assert result["reranked"] is True
top_hit = result["results"][0]
exp = top_hit["explain"]
assert exp["rerank_score"] == pytest.approx(0.9)
assert exp["pre_rerank_rank"] == 1
assert exp["blend_weight"] == 0.75
assert exp["final_score"] == pytest.approx(top_hit["score"], abs=1e-5)
# rank 1: 75% retrieval (norm=1.0 for the top candidate) + 25% rerank
assert top_hit["score"] == pytest.approx(0.75 * 1.0 + 0.25 * 0.9, abs=1e-5)
def test_rerank_can_reorder(db, monkeypatch):
"""A strong rerank score rescues a lower-retrieval-ranked chunk."""
def favour_chunk_b(query, texts):
return [1.0 if "cooking" in t else 0.0 for t in texts]
_enable_fake_reranker(monkeypatch, favour_chunk_b)
# Neutral-ish query: both chunks retrieved, chunk A ranked first.
result = hybrid_search(db.conn, "alpha cooking", _Cfg(), rerank=True)
assert result["reranked"] is True
assert len(result["results"]) == 2
def test_rerank_false_bypasses(db, monkeypatch):
called = []
_enable_fake_reranker(monkeypatch, lambda q, t: called.append(1) or [0.5] * len(t))
result = hybrid_search(db.conn, "alpha switch", _Cfg(), rerank=False)
assert result["reranked"] is False
assert not called
def test_rerank_unavailable_degrades_gracefully(db):
# No reranker loaded: rerank=True must not error.
result = hybrid_search(db.conn, "alpha switch", _Cfg(), rerank=True, explain=True)
assert result["reranked"] is False
assert "rerank_score" not in result["results"][0]["explain"]
def test_rerank_skipped_for_single_arm(db, monkeypatch):
called = []
_enable_fake_reranker(monkeypatch, lambda q, t: called.append(1) or [0.5] * len(t))
result = hybrid_search(db.conn, "alpha switch", _Cfg(), fts_only=True, rerank=True)
assert result["reranked"] is False
assert not called
def test_rerank_default_follows_cfg(db, monkeypatch):
_enable_fake_reranker(monkeypatch, lambda q, texts: [0.5] * len(texts))
class _RerankCfg(_Cfg):
rerank_enabled = True
result = hybrid_search(db.conn, "alpha switch", _RerankCfg())
assert result["reranked"] is True
def test_blend_rerank_weights_by_position():
# 12 candidates, retrieval scores 12 down to 1 → norms 1.0 down to 0.0.
candidates = [(cid, float(12 - i)) for i, cid in enumerate(range(100, 112))]
details = {cid: {} for cid, _ in candidates}
rr = [1.0] * len(candidates)
blended = dict(_blend_rerank(candidates, rr, details))
assert details[100]["blend_weight"] == 0.75 # rank 1
assert details[102]["blend_weight"] == 0.75 # rank 3
assert details[103]["blend_weight"] == 0.60 # rank 4
assert details[109]["blend_weight"] == 0.60 # rank 10
assert details[110]["blend_weight"] == 0.40 # rank 11
# rank 1: norm 1.0 → 0.75*1.0 + 0.25*1.0 = 1.0
assert blended[100] == pytest.approx(1.0)
# rank 4: norm 8/11 → 0.6*(8/11) + 0.4*1.0
assert blended[103] == pytest.approx(0.6 * (8 / 11) + 0.4)
# rank 11: norm 1/11 → 0.4*(1/11) + 0.6*1.0
assert blended[110] == pytest.approx(0.4 * (1 / 11) + 0.6)
def test_blend_rerank_constant_retrieval_scores():
"""Zero span (all candidates same retrieval score) must not divide by zero."""
candidates = [(1, 0.5), (2, 0.5)]
details = {1: {}, 2: {}}
blended = dict(_blend_rerank(candidates, [0.2, 0.8], details))
assert blended[1] == pytest.approx(0.75 * 1.0 + 0.25 * 0.2)
assert blended[2] == pytest.approx(0.75 * 1.0 + 0.25 * 0.8)
+73
View File
@@ -0,0 +1,73 @@
"""Tests for tag context descriptions."""
import sys
import types
import pytest
from kb.database import (
get_connection,
init_schema,
insert_chunk,
insert_document,
insert_embedding,
tag_document,
)
from kb.search import hybrid_search
DIM = 4
class _Cfg:
search_threshold = 0.0
rerank_enabled = False
reranker_model = "test-reranker"
rerank_candidates = 40
@pytest.fixture
def conn(tmp_path, monkeypatch):
fake = types.ModuleType("kb.embeddings")
fake.embed_texts = lambda texts: [[1.0, 0.0, 0.0, 0.0] for _ in texts]
monkeypatch.setitem(sys.modules, "kb.embeddings", fake)
conn = get_connection(str(tmp_path / "kb.db"))
init_schema(conn, embedding_dim=DIM)
yield conn
conn.close()
def test_migration_is_idempotent(conn):
# Running init_schema again (fresh start on an existing DB) must not fail.
init_schema(conn, embedding_dim=DIM)
cols = {row[1] for row in conn.execute("PRAGMA table_info(tags)").fetchall()}
assert "description" in cols
def test_search_results_carry_tag_contexts(conn):
doc = insert_document(conn, "Runbook", "/notes/runbook.md", "h1", "markdown")
chunk = insert_chunk(conn, doc, 0, "restart the proxy after cert renewal")
insert_embedding(conn, chunk, [1.0, 0.0, 0.0, 0.0])
tag_document(conn, doc, ["ops", "draft"])
conn.execute(
"UPDATE tags SET description = ? WHERE name = ?",
("Lab operations runbooks", "ops"),
)
conn.commit()
result = hybrid_search(conn, "restart proxy", _Cfg())
hit = result["results"][0]
assert hit["tags"] == ["draft", "ops"]
# Only described tags appear in tag_contexts.
assert hit["tag_contexts"] == {"ops": "Lab operations runbooks"}
def test_tag_contexts_empty_when_no_descriptions(conn):
doc = insert_document(conn, "Plain", "/notes/plain.md", "h2", "markdown")
chunk = insert_chunk(conn, doc, 0, "some plain text about switches")
insert_embedding(conn, chunk, [1.0, 0.0, 0.0, 0.0])
tag_document(conn, doc, ["misc"])
result = hybrid_search(conn, "switches", _Cfg())
hit = result["results"][0]
assert hit["tag_contexts"] == {}
+6 -1
View File
@@ -18,7 +18,8 @@ def _client() -> httpx.Client:
def search(query: str, top: int = 10, tags: list[str] | None = None, def search(query: str, top: int = 10, tags: list[str] | None = None,
doc_type: str | None = None, fts_only: bool = False, doc_type: str | None = None, fts_only: bool = False,
vec_only: bool = False, threshold: float | None = None) -> dict: vec_only: bool = False, threshold: float | None = None,
explain: bool = False, rerank: bool | None = None) -> dict:
body: dict = {"query": query, "top": top} body: dict = {"query": query, "top": top}
if tags: if tags:
body["tags"] = tags body["tags"] = tags
@@ -30,6 +31,10 @@ def search(query: str, top: int = 10, tags: list[str] | None = None,
body["vec_only"] = True body["vec_only"] = True
if threshold is not None: if threshold is not None:
body["threshold"] = threshold body["threshold"] = threshold
if explain:
body["explain"] = True
if rerank is not None:
body["rerank"] = rerank
with _client() as c: with _client() as c:
r = c.post("/api/v1/search", json=body) r = c.post("/api/v1/search", json=body)
r.raise_for_status() r.raise_for_status()
+45 -13
View File
@@ -44,11 +44,18 @@ _transport_security = TransportSecuritySettings(
mcp = FastMCP( mcp = FastMCP(
"kb", "kb",
instructions=( instructions=(
"Knowledge base MCP server. Provides tools for searching, adding, and " "Knowledge base MCP server with hybrid semantic + full-text search. "
"managing documents and notes. Use tags to organise and filter documents " "kb_search uses dense vector embeddings (semantic similarity) fused with "
"(e.g. tag notes with 'agent:mybot' and filter searches by that tag). " "BM25 full-text ranking, so it finds conceptually related content even "
"This server requires Bearer token authentication — all requests are " "when the exact words don't match — agents can ask natural-language "
"authenticated via the Authorization header at the HTTP transport layer." "questions rather than guessing keywords. When the engine has a "
"cross-encoder reranker enabled, results are reranked server-side by "
"default (pass rerank=False for lower latency). Also provides tools for "
"adding notes, uploading files, and managing documents and tags. Use tags "
"to organise and filter documents (e.g. tag notes with 'agent:mybot' and "
"filter searches by that tag). This server requires Bearer token "
"authentication — all requests are authenticated via the Authorization "
"header at the HTTP transport layer."
), ),
transport_security=_transport_security, transport_security=_transport_security,
) )
@@ -61,25 +68,48 @@ async def kb_search(
tags: list[str] | None = None, tags: list[str] | None = None,
doc_type: str | None = None, doc_type: str | None = None,
fts_only: bool = False, fts_only: bool = False,
explain: bool = False,
rerank: bool | None = None,
) -> str: ) -> str:
"""Search the knowledge base for relevant documents and notes. """Hybrid semantic (vector) + full-text search over the knowledge base.
Returns ranked chunks matching the query, with text content, relevance scores, Combines dense vector embeddings (semantic similarity finds conceptually
and document metadata. related content even when the wording differs) with BM25 keyword ranking,
fused via reciprocal rank fusion. Because the search is semantic, you can
ask natural-language questions ("what did we decide about X?") rather than
guessing the exact keywords used in the source documents.
When the engine has a cross-encoder reranker enabled, the top candidates
are reranked server-side by default you normally do NOT need to rerank
results yourself. Check kb_status's "rerank" block to see whether it is
active.
Returns ranked chunks matching the query, with text content, relevance
scores, and document metadata.
Args: Args:
query: The search query. Can be a natural language question or keywords. query: The search query a natural language question or keywords.
top: Maximum number of results to return (default 10). top: Maximum number of results to return (default 10).
tags: Filter results to documents with ALL of these tags. tags: Filter results to documents with ALL of these tags.
doc_type: Filter by document type (e.g. "note", "pdf", "markdown", "code"). doc_type: Filter by document type (e.g. "note", "pdf", "markdown",
fts_only: If true, use only full-text search (no vector similarity). "code", "data").
fts_only: Disable the vector/semantic component and use only BM25
keyword matching. Default false (hybrid mode). Set true only when
you need exact-string matching (e.g. an error code, identifier).
explain: Include a per-result score breakdown (BM25 score/rank, vector
similarity/rank, rank-fusion contributions, rerank blend) under an
"explain" key. Useful for diagnosing why a result ranked where it did.
rerank: Set false to skip server-side reranking for lower latency.
Default (None) uses the engine's configured behaviour.
Tips for complex queries: Tips for complex queries:
- Consider expanding into 2-3 variant phrasings and calling this tool multiple - Consider expanding into 2-3 variant phrasings and calling this tool multiple
times, then deduplicating results by chunk_id. For example, search for both times, then deduplicating results by chunk_id. For example, search for both
"pension revaluation rules" and "how are pensions revalued" to cast a wider net. "pension revaluation rules" and "how are pensions revalued" to cast a wider net.
- For precision, rerank the returned results using your own judgement based on - If the engine's reranker is disabled, you can still rerank the returned
relevance to the original question. results yourself using your own judgement of relevance to the question.
- Call kb_status to see which embedding model is in use and whether
server-side reranking is active.
""" """
result = engine.search( result = engine.search(
query=query, query=query,
@@ -87,6 +117,8 @@ async def kb_search(
tags=tags or None, tags=tags or None,
doc_type=doc_type, doc_type=doc_type,
fts_only=fts_only, fts_only=fts_only,
explain=explain,
rerank=rerank,
) )
results_list = result if isinstance(result, list) else result.get("results", []) results_list = result if isinstance(result, list) else result.get("results", [])
+63
View File
@@ -0,0 +1,63 @@
# kb — Next Steps
> Implementation status (2026-08-22): `document_id`, metadata-only info,
> title/filename filters, document-level `find`, and short Docling chunk
> filtering are implemented on `feature/tasks-5-15-completion`. OCR is already
> enabled through RapidOCR; further OCR tuning remains measurement-gated.
UX improvements to make documents easier to find and inspect, prompted by a session where searching for an uploaded PDF (`M38T_PHEV_RHD_OM_EN_UK_20251209.pdf`, doc id 2077, 1801 chunks) surfaced lots of chunk hits but no obvious path back to the original document.
## Problems observed
### 1. `kb list` silently ignores positional arguments
```
kb list --type pdf "M38T_PHEV_RHD_OM_EN_UK_20251209"
```
The quoted term is dropped without warning; user gets the default newest-first listing and assumes the document is missing. `kb list` currently only supports `--tags` and `--type` filters.
### 2. `kb search` returns chunks with no `document_id`
Result objects expose `chunk_id`, `title`, `source_path`, `tags` — but not `document_id`. To get from a search hit back to the owning document you have to title-match against `kb list` output or call an undocumented endpoint. The skill docs even claim a `source.document_id` field that isn't actually present in the CLI output.
### 3. `kb info` dumps every chunk with no summary mode
`kb info 2077` returns ~1801 chunk objects. The document-level metadata (`id`, `title`, `original_filename`, `source_path`, `stored_path`, `doc_type`, `language`, `content_hash`, `has_file`, `tags`, `created_at`, `updated_at`) **is** present at the top level of the JSON, but in practice it's invisible — human format presumably dumps the chunk list and the user sees only chunks.
There's no way to ask for "just tell me about this document."
### 4. Search hits can look like noise on image-heavy PDFs
Top chunks for the M38T search were single characters (`"1"`, `"B"`, `"\""`). Almost certainly an FTS artefact on short tokens from a scan/image-heavy PDF — but it makes the result set look broken. Worth considering a minimum-text-length filter on indexed chunks, or down-weighting very short chunks in ranking.
## Proposed changes
### Small / high-value
- **`kb info --no-chunks`** (or make `--chunks` opt-in): default to metadata + chunk count, only include chunks when asked. Human format should always lead with the metadata block.
- **`kb list --title <substring>`** (or accept a positional query) for filename / title search. At minimum, error or warn when positional args are passed and ignored.
- **Include `document_id` in `kb search` result objects.** Either at the top of each result or under `source.document_id` (matching the skill docs).
### Medium
- **`kb find <query>`** as a doc-level search that aggregates chunk hits per document and returns ranked *documents* (with hit count, top chunk preview). This is what users usually want when they say "find my PDF about X."
- **Update the `kb` skill docs** to match actual CLI output shape, and to steer users toward `kb list | jq` for filename lookups until proper filtering lands.
### Larger
- **Quality filter for short chunks** during ingestion (e.g. drop chunks with < N alphanumeric chars, or fold them into neighbours). Stops scanned/image-heavy PDFs from polluting search.
- **OCR path for scan-heavy PDFs.** The M38T manual extracted enough real text to be useful, but other "scan" docs likely don't. Detect low text density per page and route through OCR.
## Quick reference (current workarounds)
```bash
# Find a doc by filename
kb list --type pdf --format json | jq '.[] | select(.title | contains("M38T"))'
# Get just metadata for a doc
kb info 2077 --format json | jq 'del(.chunks)'
# Download the original
kb export 2077 -o manual.pdf
```
+15 -7
View File
@@ -2,7 +2,7 @@
## Purpose ## Purpose
Documents recommended patterns for agent-side query expansion and reranking, which are caller responsibilities rather than engine features. These patterns are communicated via MCP tool descriptions. Documents recommended patterns for agent-side query expansion, plus how agent guidance interacts with the engine's optional server-side reranking. These patterns are communicated via MCP tool descriptions.
## Requirements ## Requirements
@@ -18,18 +18,26 @@ The `kb_search` MCP tool description SHALL include guidance on query expansion a
### Requirement: Reranking guidance in tool description ### Requirement: Reranking guidance in tool description
The `kb_search` MCP tool description SHALL include guidance on agent-side reranking as a recommended pattern for improving precision. The `kb_search` MCP tool description SHALL describe the engine's server-side reranking behaviour and retain agent-side reranking as a fallback pattern.
#### Scenario: Tool description includes reranking pattern #### Scenario: Tool description covers server-side reranking
- **WHEN** an agent reads the `kb_search` tool description - **WHEN** an agent reads the `kb_search` tool description
- **THEN** the description SHALL include guidance such as: "For precision, rerank the returned results using your own judgement based on relevance to the original question" - **THEN** the description SHALL state that results are reranked server-side by default when the engine has a reranker enabled, that `rerank=False` skips it for lower latency, and that `kb_status` reports whether reranking is active
#### Scenario: Tool description retains agent-side fallback
- **WHEN** an agent reads the `kb_search` tool description
- **THEN** the description SHALL include guidance that, when the engine's reranker is disabled, the agent can rerank the returned results using its own judgement of relevance to the original question
--- ---
### Requirement: No engine-side LLM dependency ### Requirement: No external LLM dependency
The engine SHALL NOT require or use any external LLM API for search operations. Query expansion and reranking SHALL remain entirely agent-side concerns. The engine SHALL NOT require or use any external LLM API for search operations. Query expansion SHALL remain an agent-side concern. Reranking MAY be performed engine-side using a local, opt-in cross-encoder model; it SHALL never depend on an external API.
#### Scenario: Engine has no LLM dependency #### Scenario: Engine has no external LLM dependency
- **WHEN** the engine is deployed without any `ANTHROPIC_API_KEY` or similar LLM API configuration - **WHEN** the engine is deployed without any `ANTHROPIC_API_KEY` or similar LLM API configuration
- **THEN** all search operations SHALL function fully, with no degraded results or missing features - **THEN** all search operations SHALL function fully, with no degraded results or missing features
#### Scenario: Reranking is optional and degrades gracefully
- **WHEN** the engine is deployed with `KB_RERANK_ENABLED` unset or false, or the reranker model fails to load
- **THEN** all search operations SHALL function fully using hybrid retrieval alone, with responses reporting `"reranked": false`
+1 -1
View File
@@ -65,7 +65,7 @@ The project SHALL provide Docker Compose files for single-command deployment. Co
#### Scenario: Pre-built image deployment #### Scenario: Pre-built image deployment
- **WHEN** an admin wants to use a pre-built engine image without building from source - **WHEN** an admin wants to use a pre-built engine image without building from source
- **THEN** the engine release notes SHALL include the exact `docker pull` command with the versioned tag (e.g. `docker.dcglab.co.uk/dcg/kb/engine:engine-v2.1.0-nvidia`) - **THEN** the engine release notes SHALL include the exact `docker pull` command with the versioned tag (e.g. `docker.dcglab.co.uk/public/kb/engine:engine-v2.1.0-nvidia`)
#### Scenario: MCP allowed hosts in Compose #### Scenario: MCP allowed hosts in Compose
- **WHEN** the kb-mcp service is defined in a Compose file - **WHEN** the kb-mcp service is defined in a Compose file
+53 -5
View File
@@ -26,11 +26,15 @@ The engine SHALL load the embedding model eagerly at startup before accepting HT
### Requirement: Hybrid search ### Requirement: Hybrid search
The engine SHALL provide hybrid search combining BM25 full-text search (via FTS5) and vector similarity search (via sqlite-vec), merged using Reciprocal Rank Fusion. Search SHALL complete in under 100ms when the model is warm. The engine SHALL sanitize user query strings to prevent FTS5 syntax errors for any input. The engine SHALL provide hybrid search combining BM25 full-text search (via FTS5) and vector similarity search (via sqlite-vec), merged using Reciprocal Rank Fusion with a top-rank bonus (+0.05 for rank 1, +0.02 for ranks 2-3 in either arm) that preserves exact matches. Search SHALL complete in under 100ms when the model is warm and reranking is disabled; reranked searches SHALL complete in under 500ms on GPU. The engine SHALL sanitize user query strings to prevent FTS5 syntax errors for any input.
#### Scenario: Hybrid search with results #### Scenario: Hybrid search with results
- **WHEN** a client sends `POST /api/v1/search` with body `{"query": "how to change oil", "top": 5}` - **WHEN** a client sends `POST /api/v1/search` with body `{"query": "how to change oil", "top": 5}`
- **THEN** the engine SHALL embed the query using the resident model, run both FTS5 and vector searches, merge results via RRF, and return a JSON response with matched chunks including scores, document metadata, and tags - **THEN** the engine SHALL embed the query using the resident model, run both FTS5 and vector searches, merge results via RRF with top-rank bonus, and return a JSON response with matched chunks including scores, `document_id`, document metadata, tags, and `tag_contexts`
#### Scenario: Explain traces
- **WHEN** a client sends `POST /api/v1/search` with `"explain": true`
- **THEN** each result SHALL include an `explain` object with per-arm raw scores and ranks (`fts_score`, `fts_rank`, `vec_score`, `vec_rank`), RRF contributions (`rrf_fts`, `rrf_vec`), the top-rank `bonus`, rerank blend fields when reranking ran (`pre_rerank_rank`, `retrieval_norm`, `rerank_score`, `blend_weight`), and the `final_score`; fields for an arm that did not match SHALL be null
#### Scenario: Search with filters #### Scenario: Search with filters
- **WHEN** a client sends `POST /api/v1/search` with body `{"query": "brakes", "tags": ["maintenance"], "doc_type": "pdf", "top": 3}` - **WHEN** a client sends `POST /api/v1/search` with body `{"query": "brakes", "tags": ["maintenance"], "doc_type": "pdf", "top": 3}`
@@ -62,6 +66,28 @@ The engine SHALL provide hybrid search combining BM25 full-text search (via FTS5
--- ---
### Requirement: Cross-encoder reranking
The engine SHALL support optional server-side reranking of hybrid search results using a local cross-encoder model, enabled via `KB_RERANK_ENABLED` (default false) with the model set by `KB_RERANKER_MODEL` (default `BAAI/bge-reranker-v2-m3`). When active, the engine SHALL over-fetch candidates (`KB_RERANK_CANDIDATES`, default 40), score each (query, chunk) pair, and blend retrieval and rerank scores position-aware: 75% retrieval weight for pre-rerank ranks 1-3, 60% for 4-10, 40% for 11+, with retrieval scores min-max normalised over the candidate set. Blended scores are on a 0-1 scale distinct from RRF scores; the score threshold SHALL be applied before reranking only. Every search response SHALL include a top-level `"reranked"` boolean.
#### Scenario: Reranked search
- **WHEN** reranking is enabled with a loaded model and a client sends a hybrid search
- **THEN** the engine SHALL rerank the top candidates and return results ordered by blended score with `"reranked": true`
#### Scenario: Per-request opt-out
- **WHEN** a client sends `POST /api/v1/search` with `"rerank": false`
- **THEN** the engine SHALL skip reranking and return plain hybrid results with `"reranked": false`
#### Scenario: Graceful degradation
- **WHEN** reranking is requested but the model is disabled or failed to load
- **THEN** the engine SHALL return plain hybrid results with `"reranked": false` and no error
#### Scenario: Single-arm searches never rerank
- **WHEN** a client sends a search with `fts_only` or `vec_only` set
- **THEN** the engine SHALL NOT rerank, keeping single-arm results pure for benchmarking
---
### Requirement: Async ingestion via job queue ### Requirement: Async ingestion via job queue
The engine SHALL accept file uploads and text notes for ingestion asynchronously. Uploaded content SHALL be written to a staging area and a job record created in the database. The engine SHALL return HTTP 202 immediately. A background worker SHALL process queued jobs sequentially. Before staging, the engine SHALL compute a SHA256 hash of the uploaded content and reject duplicates immediately. The engine SHALL accept file uploads and text notes for ingestion asynchronously. Uploaded content SHALL be written to a staging area and a job record created in the database. The engine SHALL return HTTP 202 immediately. A background worker SHALL process queued jobs sequentially. Before staging, the engine SHALL compute a SHA256 hash of the uploaded content and reject duplicates immediately.
@@ -128,7 +154,7 @@ The engine SHALL maintain job records in SQLite with status tracking. Jobs SHALL
### Requirement: Background ingestion worker ### Requirement: Background ingestion worker
The engine SHALL run a background worker that processes queued jobs. The worker SHALL process one job at a time. For each job, it SHALL: detect document type, run the appropriate chunking pipeline (Docling for PDFs, header-based for Markdown, AST-based for code, whole-text for notes), build enriched text by prepending the document title (and section header when present) to each chunk's text, generate embeddings using the enriched text and the resident model, insert chunks (with both raw text and enriched text) and vectors into the database, and move the original file to persistent storage. The engine SHALL run a background worker that processes queued jobs. The worker SHALL process one job at a time. For each job, it SHALL: detect document type, run the appropriate chunking pipeline (Docling for PDFs, header-based for Markdown, AST-based for code, whole-text for notes, fixed-size text chunking for data files with minified JSON pretty-printed first), build enriched text by prepending the document title (and section header when present) to each chunk's text, generate embeddings using the enriched text and the resident model, insert chunks (with both raw text and enriched text) and vectors into the database, and move the original file to persistent storage.
#### Scenario: Successful PDF ingestion #### Scenario: Successful PDF ingestion
- **WHEN** the background worker picks up a queued PDF job - **WHEN** the background worker picks up a queued PDF job
@@ -216,7 +242,7 @@ The engine SHALL provide endpoints to list all tags and manage tags on documents
#### Scenario: List all tags #### Scenario: List all tags
- **WHEN** a client sends `GET /api/v1/tags` - **WHEN** a client sends `GET /api/v1/tags`
- **THEN** the engine SHALL return a JSON array of tags with name and document count - **THEN** the engine SHALL return a JSON array of tags with name, document count, and description (null when unset)
#### Scenario: Add tags to a document #### Scenario: Add tags to a document
- **WHEN** a client sends `PUT /api/v1/documents/{id}/tags` with body `{"add": ["manual", "v2"]}` - **WHEN** a client sends `PUT /api/v1/documents/{id}/tags` with body `{"add": ["manual", "v2"]}`
@@ -228,13 +254,35 @@ The engine SHALL provide endpoints to list all tags and manage tags on documents
--- ---
### Requirement: Tag context descriptions
The engine SHALL support a one-line context description per tag, stored in a `description` column on the tags table (added via idempotent migration). Search results SHALL include a `tag_contexts` object mapping each of the document's described tags to its description, so consumers can judge which similar-scoring chunks answer the question.
#### Scenario: Set a tag description
- **WHEN** a client sends `PUT /api/v1/tags/{name}/description` with body `{"description": "Lab operations runbooks"}`
- **THEN** the engine SHALL store the description (matching the tag name case-insensitively) and return `{"name": "<name>", "description": "<description>"}`
#### Scenario: Clear a tag description
- **WHEN** a client sends `PUT /api/v1/tags/{name}/description` with a null or empty description
- **THEN** the engine SHALL clear the stored description
#### Scenario: Unknown tag
- **WHEN** a client sets a description for a tag that does not exist
- **THEN** the engine SHALL return HTTP 404
#### Scenario: Descriptions in search results
- **WHEN** a search result's document carries tags and at least one tag has a description
- **THEN** the result SHALL include `tag_contexts` with only the described tags; results with no described tags SHALL include an empty `tag_contexts` object
---
### Requirement: Engine status and reindex ### Requirement: Engine status and reindex
The engine SHALL provide status information and support re-embedding all chunks. The `version` field in the status response SHALL always be present and SHALL reflect the engine's release version as read from the `VERSION` file. This field is the contract used by clients for compatibility checking. The engine SHALL provide status information and support re-embedding all chunks. The `version` field in the status response SHALL always be present and SHALL reflect the engine's release version as read from the `VERSION` file. This field is the contract used by clients for compatibility checking.
#### Scenario: Get engine status #### Scenario: Get engine status
- **WHEN** a client sends `GET /api/v1/status` - **WHEN** a client sends `GET /api/v1/status`
- **THEN** the engine SHALL return JSON with `version` (string, from VERSION file), model_name, embedding_dim, GPU device info, database stats (document count by type, total chunks, DB size), and queue stats (queued/processing job count) - **THEN** the engine SHALL return JSON with `version` (string, from VERSION file), model_name, embedding_dim, GPU device info, database stats (document count by type, total chunks, DB size), queue stats (queued/processing job count), and a `rerank` object with `enabled`, `model`, `loaded`, and `candidates`
#### Scenario: Trigger reindex #### Scenario: Trigger reindex
- **WHEN** a client sends `POST /api/v1/reindex` - **WHEN** a client sends `POST /api/v1/reindex`
+70 -11
View File
@@ -15,9 +15,17 @@ ENGINE_DIR="$SCRIPT_DIR/engine"
VERSION_FILE="$ENGINE_DIR/VERSION" VERSION_FILE="$ENGINE_DIR/VERSION"
# Container registry # Container registry
#
# --provenance=false --sbom=false on every build: buildx would otherwise attach
# attestation manifests, making the image an OCI image index. The Registry v2
# host at docker.dcglab.co.uk rejects those with a 500 on manifest PUT.
REGISTRY="${REGISTRY:-docker.dcglab.co.uk}" REGISTRY="${REGISTRY:-docker.dcglab.co.uk}"
IMAGE_ORG="${IMAGE_ORG:-dcg}" IMAGE_ORG="${IMAGE_ORG:-public}"
IMAGE_BASE="${REGISTRY}/${IMAGE_ORG}/kb" IMAGE_BASE="${REGISTRY}${IMAGE_ORG:+/${IMAGE_ORG}}/kb"
# Push retries — see push_image() below
PUSH_RETRIES="${PUSH_RETRIES:-5}"
PUSH_RETRY_DELAY="${PUSH_RETRY_DELAY:-10}"
#────────────────────────────────────────────────────────────────────── #──────────────────────────────────────────────────────────────────────
# Parse args # Parse args
@@ -98,6 +106,46 @@ run() {
fi fi
} }
registry_login() {
echo " $ docker login $REGISTRY --username \$DOCKER_DCGLAB_CI_USERNAME --password-stdin"
[[ "$DRY_RUN" == true ]] && return 0
printf '%s' "$DOCKER_DCGLAB_CI_PASSWORD" |
docker login "$REGISTRY" \
--username "$DOCKER_DCGLAB_CI_USERNAME" \
--password-stdin
}
# Push one image tag, retrying on transient registry failures.
#
# The engine images carry a ~5.6GB torch layer. Uploading it intermittently
# fails with a 502 from the reverse proxy in front of the registry, and a
# manifest PUT can then fail with a 500 because the blob commit has not yet
# registered. Both clear on a retry, so a whole release should not be lost to
# one hiccup. Tune with PUSH_RETRIES / PUSH_RETRY_DELAY.
push_image() {
local image="$1"
local attempt=1
echo " $ docker push $image"
[[ "$DRY_RUN" == true ]] && return 0
while true; do
if docker push "$image"; then
return 0
fi
if (( attempt >= PUSH_RETRIES )); then
echo "Error: failed to push $image after $PUSH_RETRIES attempts" >&2
return 1
fi
echo " push failed (attempt $attempt/$PUSH_RETRIES) — retrying in ${PUSH_RETRY_DELAY}s"
sleep "$PUSH_RETRY_DELAY"
attempt=$(( attempt + 1 ))
done
}
#────────────────────────────────────────────────────────────────────── #──────────────────────────────────────────────────────────────────────
# Determine release version # Determine release version
#────────────────────────────────────────────────────────────────────── #──────────────────────────────────────────────────────────────────────
@@ -127,6 +175,17 @@ echo ""
echo "==> Pre-flight checks" echo "==> Pre-flight checks"
if [[ "$DRY_RUN" == false ]]; then if [[ "$DRY_RUN" == false ]]; then
if [[ -z "${DOCKER_DCGLAB_CI_USERNAME:-}" ]]; then
echo "Error: DOCKER_DCGLAB_CI_USERNAME is required" >&2
exit 1
fi
if [[ -z "${DOCKER_DCGLAB_CI_PASSWORD:-}" ]]; then
echo "Error: DOCKER_DCGLAB_CI_PASSWORD is required" >&2
exit 1
fi
registry_login
if git -C "$SCRIPT_DIR" rev-parse "$GIT_TAG" &>/dev/null; then if git -C "$SCRIPT_DIR" rev-parse "$GIT_TAG" &>/dev/null; then
echo "Error: tag $GIT_TAG already exists" echo "Error: tag $GIT_TAG already exists"
exit 1 exit 1
@@ -155,8 +214,8 @@ CPU_IMAGE="${IMAGE_BASE}/engine:${DOCKER_TAG}-cpu"
NVIDIA_LATEST="${IMAGE_BASE}/engine:latest-nvidia" NVIDIA_LATEST="${IMAGE_BASE}/engine:latest-nvidia"
CPU_LATEST="${IMAGE_BASE}/engine:latest-cpu" CPU_LATEST="${IMAGE_BASE}/engine:latest-cpu"
run docker build -t "$NVIDIA_IMAGE" -t "$NVIDIA_LATEST" -f "$ENGINE_DIR/Dockerfile.nvidia" "$ENGINE_DIR" run docker build --provenance=false --sbom=false -t "$NVIDIA_IMAGE" -t "$NVIDIA_LATEST" -f "$ENGINE_DIR/Dockerfile.nvidia" "$ENGINE_DIR"
run docker build -t "$CPU_IMAGE" -t "$CPU_LATEST" -f "$ENGINE_DIR/Dockerfile.cpu" "$ENGINE_DIR" run docker build --provenance=false --sbom=false -t "$CPU_IMAGE" -t "$CPU_LATEST" -f "$ENGINE_DIR/Dockerfile.cpu" "$ENGINE_DIR"
echo "" echo ""
@@ -171,7 +230,7 @@ if [[ -f "$MCP_DIR/Dockerfile" ]]; then
MCP_IMAGE="${IMAGE_BASE}/mcp:${DOCKER_TAG}" MCP_IMAGE="${IMAGE_BASE}/mcp:${DOCKER_TAG}"
MCP_LATEST="${IMAGE_BASE}/mcp:latest" MCP_LATEST="${IMAGE_BASE}/mcp:latest"
run docker build -t "$MCP_IMAGE" -t "$MCP_LATEST" -f "$MCP_DIR/Dockerfile" "$MCP_DIR" run docker build --provenance=false --sbom=false -t "$MCP_IMAGE" -t "$MCP_LATEST" -f "$MCP_DIR/Dockerfile" "$MCP_DIR"
echo "" echo ""
fi fi
@@ -233,14 +292,14 @@ echo ""
#────────────────────────────────────────────────────────────────────── #──────────────────────────────────────────────────────────────────────
echo "==> Pushing Docker images to $REGISTRY" echo "==> Pushing Docker images to $REGISTRY"
run docker push "$NVIDIA_IMAGE" push_image "$NVIDIA_IMAGE"
run docker push "$NVIDIA_LATEST" push_image "$NVIDIA_LATEST"
run docker push "$CPU_IMAGE" push_image "$CPU_IMAGE"
run docker push "$CPU_LATEST" push_image "$CPU_LATEST"
if [[ -n "${MCP_IMAGE:-}" ]]; then if [[ -n "${MCP_IMAGE:-}" ]]; then
run docker push "$MCP_IMAGE" push_image "$MCP_IMAGE"
run docker push "$MCP_LATEST" push_image "$MCP_LATEST"
fi fi
echo "" echo ""