Compare commits
22 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 5b5a9ecbc3 | |||
| 932e889ee8 | |||
| 48ca873fe2 | |||
| a38d77ed23 | |||
| f1ed5b6e23 | |||
| 3ab8a81c14 | |||
| 739c3ff30c | |||
| 3151a1b08a | |||
| 5049ba2a2f | |||
| 6dfc13be1d | |||
| 75e4a0cf73 | |||
| 753c641e72 | |||
| 45e2c5ce91 | |||
| e6e91f1d5c | |||
| 9eccc527ae | |||
| d44d11e4fe | |||
| 574370e8d1 | |||
| 17b19999de | |||
| bb78f4ea80 | |||
| 223ff2cf5d | |||
| e9a282ddb1 | |||
| b5a203d2aa |
@@ -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
|
||||||
+52
-13
@@ -11,9 +11,6 @@ cd engine
|
|||||||
|
|
||||||
# NVIDIA GPU
|
# NVIDIA GPU
|
||||||
KB_DATA_PATH=~/kb-data docker compose -f compose.nvidia.yaml up -d
|
KB_DATA_PATH=~/kb-data docker compose -f compose.nvidia.yaml up -d
|
||||||
|
|
||||||
# AMD GPU (ROCm)
|
|
||||||
KB_DATA_PATH=~/kb-data docker compose -f compose.rocm.yaml up -d
|
|
||||||
```
|
```
|
||||||
|
|
||||||
### Client
|
### Client
|
||||||
@@ -24,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.
|
||||||
@@ -50,7 +76,7 @@ The client embeds a `MinEngineVersion` (from `client/MIN_ENGINE_VERSION`) and wi
|
|||||||
./release-engine.sh --gitea --dry-run # preview without doing anything
|
./release-engine.sh --gitea --dry-run # preview without doing anything
|
||||||
```
|
```
|
||||||
|
|
||||||
Creates tag `engine-vX.Y.Z`, builds NVIDIA and ROCm Docker images, creates a Gitea/GitHub release, and pushes images to the registry.
|
Creates tag `engine-vX.Y.Z`, builds NVIDIA and CPU Docker images, creates a Gitea/GitHub release, and pushes images to the registry.
|
||||||
|
|
||||||
### Checking versions
|
### Checking versions
|
||||||
|
|
||||||
@@ -64,10 +90,14 @@ 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-rocm` — versioned
|
- `engine-v2.0.6-nvidia` / `engine-v2.0.6-cpu` — versioned
|
||||||
- `latest-nvidia` / `latest-rocm` — 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:
|
||||||
|
|
||||||
@@ -75,6 +105,15 @@ Override the registry and org via environment variables:
|
|||||||
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.
|
||||||
@@ -91,10 +130,10 @@ 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 |
|
||||||
## Future: ROCm runtime migration
|
| `POST` | `/bulk/tags` | Bulk add/remove tags by filter |
|
||||||
|
| `POST` | `/bulk/set-tags` | Bulk replace tags by filter |
|
||||||
The `onnxruntime-rocm` execution provider was removed from onnxruntime as of v1.23. AMD is pushing toward the **MIGraphX execution provider** as the replacement for ROCm GPU inference. When upgrading onnxruntime beyond v1.22, the ROCm Dockerfile will need to switch from `onnxruntime-rocm` to `onnxruntime` with the MIGraphX EP and install the `migraphx` runtime libraries instead.
|
|
||||||
|
|||||||
@@ -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,32 +20,32 @@ 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 collection/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 |
|
||||||
|
| `kb_delete` | Permanently delete a document by ID |
|
||||||
| `kb_status` | Engine health and statistics |
|
| `kb_status` | Engine health and statistics |
|
||||||
| `kb_jobs` | Ingestion queue status |
|
| `kb_jobs` | Ingestion queue status |
|
||||||
| `kb_upload_start` | Start a chunked file upload |
|
| `kb_upload_start` | Start a chunked file upload |
|
||||||
| `kb_upload_chunk` | Upload a base64-encoded file chunk |
|
| `kb_upload_chunk` | Upload a base64-encoded file chunk |
|
||||||
| `kb_upload_finish` | Finish upload and submit for ingestion |
|
| `kb_upload_finish` | Finish upload and submit for ingestion |
|
||||||
|
| `kb_bulk_delete` | Delete multiple documents matching a filter |
|
||||||
|
| `kb_bulk_tags` | Add/remove tags on multiple documents |
|
||||||
|
| `kb_bulk_set_tags` | Replace all tags on multiple documents |
|
||||||
|
|
||||||
## Collections
|
## Organising with tags
|
||||||
|
|
||||||
The MCP server supports **collections** — scoped document namespaces implemented via tag conventions. Use these to separate agent memory from user documents:
|
Use tags to separate agent data from user documents. For example, an agent can tag all its notes with `agent:mybot` and filter by that tag when searching. This is a naming convention — configure it in your agent's system prompt. No special server-side enforcement is needed.
|
||||||
|
|
||||||
- `documents` (default) — user-facing documents
|
Bulk tools accept filter-based selection (by tags, doc_type, ID list, or ID range) so agents can manage thousands of documents in a single call instead of looping. A safety threshold (default 70%, configurable via engine env var `KB_BULK_SAFETY_PERCENT`) prevents accidental mass operations unless `force: true` is set.
|
||||||
- `memory` — agent memory and preferences
|
|
||||||
- `workspace` — working context
|
|
||||||
|
|
||||||
Tools accept a `collection` parameter. The MCP server translates this to `collection:<name>` tags on the engine, and strips them from responses so agents see a clean `"collection": "memory"` field.
|
|
||||||
|
|
||||||
## MCP server configuration
|
## MCP server configuration
|
||||||
|
|
||||||
|
|||||||
@@ -12,9 +12,9 @@ Go CLI (kb) ──HTTP──▶ FastAPI Engine (Docker) ──▶ SQLite + GPU
|
|||||||
MCP Agents ──MCP/HTTP──▶ MCP Server (Docker) ──┘
|
MCP Agents ──MCP/HTTP──▶ MCP Server (Docker) ──┘
|
||||||
```
|
```
|
||||||
|
|
||||||
- **Engine**: Keeps the embedding model warm in memory. Handles search, ingestion, document management, and note mutation via REST API. Runs in Docker with NVIDIA GPU, AMD GPU (ROCm), or CPU-only support.
|
- **Engine**: Keeps the embedding model warm in memory. Handles search, ingestion, document management, and note mutation via REST API. Runs in Docker with NVIDIA GPU or CPU-only support.
|
||||||
- **Client**: Single static Go binary. No Python, no ML dependencies, instant startup. Talks to the engine over HTTP.
|
- **Client**: Single static Go binary. No Python, no ML dependencies, instant startup. Talks to the engine over HTTP.
|
||||||
- **MCP Server**: Exposes kb operations as native MCP tools over Streamable HTTP. Runs as a separate Docker container alongside the engine. Supports collections for scoping agent memory vs user documents.
|
- **MCP Server**: Exposes kb operations as native MCP tools over Streamable HTTP. Runs as a separate Docker container alongside the engine. Use tags to scope agent data from user documents.
|
||||||
- **Storage**: Single SQLite database with FTS5 (keyword search) and sqlite-vec (vector search). Portable via bind mount — just copy the data directory between hosts.
|
- **Storage**: Single SQLite database with FTS5 (keyword search) and sqlite-vec (vector search). Portable via bind mount — just copy the data directory between hosts.
|
||||||
|
|
||||||
## Quick start
|
## Quick start
|
||||||
@@ -33,19 +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
|
||||||
|
|
||||||
# AMD GPU (ROCm)
|
|
||||||
docker run -d --name kb-engine \
|
|
||||||
--device /dev/kfd --device /dev/dri \
|
|
||||||
--group-add video \
|
|
||||||
-p 8000:8000 \
|
|
||||||
-v ~/kb-data:/data \
|
|
||||||
-e KB_MODEL=all-MiniLM-L6-v2 \
|
|
||||||
-e KB_DEVICE=auto \
|
|
||||||
-e KB_API_KEY=your-secret-key \
|
|
||||||
--restart unless-stopped \
|
|
||||||
docker.dcglab.co.uk/dcg/kb/engine:latest-rocm
|
|
||||||
|
|
||||||
# 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 \
|
||||||
@@ -54,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:
|
||||||
@@ -63,9 +51,6 @@ Or use a compose file from the repo:
|
|||||||
# NVIDIA GPU
|
# NVIDIA GPU
|
||||||
KB_DATA_PATH=~/kb-data docker compose -f engine/compose.nvidia.yaml up -d
|
KB_DATA_PATH=~/kb-data docker compose -f engine/compose.nvidia.yaml up -d
|
||||||
|
|
||||||
# AMD GPU (ROCm)
|
|
||||||
KB_DATA_PATH=~/kb-data docker compose -f engine/compose.rocm.yaml up -d
|
|
||||||
|
|
||||||
# CPU only
|
# CPU only
|
||||||
KB_DATA_PATH=~/kb-data docker compose -f engine/compose.cpu.yaml up -d
|
KB_DATA_PATH=~/kb-data docker compose -f engine/compose.cpu.yaml up -d
|
||||||
```
|
```
|
||||||
@@ -126,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
|
||||||
@@ -137,24 +124,33 @@ 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
|
||||||
kb remove 3 --yes
|
kb remove 3 --yes
|
||||||
kb status
|
kb status
|
||||||
|
|
||||||
|
# Bulk operations
|
||||||
|
kb bulk-remove --tags "draft,old" --type note --yes
|
||||||
|
kb bulk-tag --type note --add "archived" --yes
|
||||||
|
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
|
||||||
@@ -169,11 +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_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:
|
||||||
@@ -186,13 +202,13 @@ rsync -a ~/kb-data/ user@target:/home/user/kb-data/
|
|||||||
KB_DATA_PATH=~/kb-data docker compose -f compose.nvidia.yaml up -d
|
KB_DATA_PATH=~/kb-data docker compose -f compose.nvidia.yaml up -d
|
||||||
```
|
```
|
||||||
|
|
||||||
Data is device-agnostic — you can ingest on NVIDIA and serve from AMD or CPU (or any combination) with the same data directory.
|
Data is device-agnostic — you can ingest on NVIDIA and serve from CPU (or vice versa) with the same data directory.
|
||||||
|
|
||||||
## MCP server (agent integration)
|
## MCP server (agent integration)
|
||||||
|
|
||||||
The MCP server exposes kb operations as native MCP tools over Streamable HTTP, so agents can search, add notes, upload files, and manage documents without shelling out to the CLI. Includes setup guides for Claude Code, VS Code, Cursor, Windsurf, and JetBrains IDEs.
|
The MCP server exposes kb operations as native MCP tools over Streamable HTTP, so agents can search, add notes, upload files, and manage documents without shelling out to the CLI. Includes setup guides for Claude Code, VS Code, Cursor, Windsurf, and JetBrains IDEs.
|
||||||
|
|
||||||
See **[MCP.md](MCP.md)** for full details — server setup, available tools, collections, configuration, and client examples.
|
See **[MCP.md](MCP.md)** for full details — server setup, available tools, tag-based organisation, configuration, and client examples.
|
||||||
|
|
||||||
## Agent skill
|
## Agent skill
|
||||||
|
|
||||||
|
|||||||
@@ -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,48 @@ 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
|
||||||
|
|
||||||
|
Operate on multiple documents at once using filter-based selection. Filters combine with AND logic.
|
||||||
|
|
||||||
|
**Filter flags (shared across all bulk commands):**
|
||||||
|
- `--tags tag1,tag2` — match documents with ALL specified tags
|
||||||
|
- `--type pdf|note|...` — match by document type
|
||||||
|
- `--ids 1,5,12` — match specific document IDs
|
||||||
|
- `--from-id N` — match documents with id >= N
|
||||||
|
- `--to-id N` — match documents with id <= N
|
||||||
|
- `--force` / `-f` — override safety threshold (blocks operations affecting >70% of all documents)
|
||||||
|
- `--yes` / `-y` — skip confirmation prompt
|
||||||
|
|
||||||
|
```bash
|
||||||
|
# Bulk delete
|
||||||
|
kb bulk-remove --tags "draft,old" --type note --yes # delete matching docs
|
||||||
|
kb bulk-remove --from-id 10 --to-id 50 --yes # delete by ID range
|
||||||
|
kb bulk-remove --ids "3,7,12" --yes # delete specific IDs
|
||||||
|
|
||||||
|
# Bulk tag add/remove
|
||||||
|
kb bulk-tag --tags "agent:mybot" --add "reviewed" --remove "pending" --yes
|
||||||
|
kb bulk-tag --type note --add "archived" --yes # tag all notes
|
||||||
|
|
||||||
|
# Bulk replace tags
|
||||||
|
kb bulk-set-tags --tags "old-scheme" --set "new-scheme,migrated" --yes
|
||||||
|
```
|
||||||
|
|
||||||
|
All bulk commands return a summary: matched count, succeeded count, failed count, and errors.
|
||||||
|
A safety threshold prevents accidentally affecting more than 70% of documents unless `--force` is used.
|
||||||
|
The threshold is configurable on the engine via `KB_BULK_SAFETY_PERCENT` (integer 0-100, default 70; 0 disables).
|
||||||
|
|
||||||
## Jobs (ingestion queue)
|
## Jobs (ingestion queue)
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
@@ -107,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,
|
||||||
@@ -115,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
|
||||||
}
|
}
|
||||||
```
|
```
|
||||||
|
|
||||||
@@ -172,12 +214,13 @@ For agent-to-agent integration, kb provides an MCP server alongside the CLI. The
|
|||||||
exposes the same operations as native MCP tools over Streamable HTTP transport, which agents
|
exposes the same operations as native MCP tools over Streamable HTTP transport, which agents
|
||||||
can connect to directly without subprocess overhead.
|
can connect to directly without subprocess overhead.
|
||||||
|
|
||||||
**MCP tools:** `kb_search`, `kb_addnote`, `kb_update_note`, `kb_get`, `kb_status`, `kb_jobs`,
|
**MCP tools:** `kb_search`, `kb_addnote`, `kb_update_note`, `kb_get`, `kb_delete`, `kb_status`,
|
||||||
`kb_upload_start`, `kb_upload_chunk`, `kb_upload_finish`.
|
`kb_jobs`, `kb_upload_start`, `kb_upload_chunk`, `kb_upload_finish`, `kb_bulk_delete`,
|
||||||
|
`kb_bulk_tags`, `kb_bulk_set_tags`.
|
||||||
|
|
||||||
The MCP server supports **collections** — scoped document namespaces (e.g. `memory`, `documents`,
|
Use tags to separate agent data from user documents (e.g. tag all agent notes with
|
||||||
`workspace`) implemented via tag conventions. This is the recommended way for agents to separate
|
`agent:mybot` and filter by that tag when searching). This convention is communicated
|
||||||
their memory from user documents.
|
via system prompt — no special server-side enforcement needed.
|
||||||
|
|
||||||
If the kb engine is already running via Docker Compose, add the MCP server by deploying the
|
If the kb engine is already running via Docker Compose, add the MCP server by deploying the
|
||||||
`kb-mcp` service from the same compose file. Agents connect to it on port 3000 (default).
|
`kb-mcp` service from the same compose file. Agents connect to it on port 3000 (default).
|
||||||
@@ -185,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 @@
|
|||||||
3.0.0
|
3.3.0
|
||||||
|
|||||||
+1
-1
@@ -1 +1 @@
|
|||||||
3.0.0
|
3.3.0
|
||||||
|
|||||||
+52
-4
@@ -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() {
|
||||||
|
if !wait || result.Duplicate {
|
||||||
output.PrintJSON([]interface{}{result.Raw})
|
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
|
||||||
}
|
}
|
||||||
|
if !wait || result.Duplicate {
|
||||||
results = append(results, result.Raw)
|
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
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
+25
-3
@@ -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"
|
||||||
@@ -27,16 +28,20 @@ var addnoteCmd = &cobra.Command{
|
|||||||
|
|
||||||
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() {
|
||||||
|
if !wait {
|
||||||
output.PrintJSON(result)
|
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
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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
|
||||||
|
}
|
||||||
@@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,186 @@
|
|||||||
|
package cmd
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bufio"
|
||||||
|
"fmt"
|
||||||
|
"os"
|
||||||
|
"strconv"
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
"github.com/kb-search/kb/internal/api"
|
||||||
|
"github.com/kb-search/kb/internal/output"
|
||||||
|
"github.com/spf13/cobra"
|
||||||
|
)
|
||||||
|
|
||||||
|
var bulkRemoveCmd = &cobra.Command{
|
||||||
|
Use: "bulk-remove",
|
||||||
|
Short: "Delete multiple documents matching a filter",
|
||||||
|
RunE: runBulkRemove,
|
||||||
|
}
|
||||||
|
|
||||||
|
func init() {
|
||||||
|
addBulkFilterFlags(bulkRemoveCmd)
|
||||||
|
rootCmd.AddCommand(bulkRemoveCmd)
|
||||||
|
}
|
||||||
|
|
||||||
|
func runBulkRemove(cmd *cobra.Command, args []string) error {
|
||||||
|
body, err := buildBulkBody(cmd)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
yes, _ := cmd.Flags().GetBool("yes")
|
||||||
|
if !yes {
|
||||||
|
desc := describeBulkFilter(cmd)
|
||||||
|
fmt.Printf("This will delete documents matching: %s\nProceed? [y/N] ", desc)
|
||||||
|
reader := bufio.NewReader(os.Stdin)
|
||||||
|
answer, _ := reader.ReadString('\n')
|
||||||
|
answer = strings.TrimSpace(strings.ToLower(answer))
|
||||||
|
if answer != "y" && answer != "yes" {
|
||||||
|
fmt.Println("Cancelled.")
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
client := api.NewClient()
|
||||||
|
resp, err := client.Post("/api/v1/bulk/delete", 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)
|
||||||
|
}
|
||||||
|
|
||||||
|
var result map[string]interface{}
|
||||||
|
if err := api.DecodeJSON(resp, &result); err != nil {
|
||||||
|
return fmt.Errorf("failed to decode response: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if output.IsJSON() {
|
||||||
|
output.PrintJSON(result)
|
||||||
|
} else {
|
||||||
|
printBulkResult("Deleted", result)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// Shared helpers for all bulk commands
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
func addBulkFilterFlags(cmd *cobra.Command) {
|
||||||
|
cmd.Flags().String("tags", "", "filter by tags (comma-separated)")
|
||||||
|
cmd.Flags().String("type", "", "filter by document type")
|
||||||
|
cmd.Flags().String("ids", "", "filter by document IDs (comma-separated)")
|
||||||
|
cmd.Flags().Int("from-id", 0, "filter by id >= value")
|
||||||
|
cmd.Flags().Int("to-id", 0, "filter by id <= value")
|
||||||
|
cmd.Flags().BoolP("force", "f", false, "override safety threshold")
|
||||||
|
cmd.Flags().BoolP("yes", "y", false, "skip confirmation prompt")
|
||||||
|
}
|
||||||
|
|
||||||
|
func buildBulkBody(cmd *cobra.Command) (map[string]interface{}, error) {
|
||||||
|
body := map[string]interface{}{}
|
||||||
|
|
||||||
|
tagsStr, _ := cmd.Flags().GetString("tags")
|
||||||
|
if tagsStr != "" {
|
||||||
|
body["tags"] = splitTags(tagsStr)
|
||||||
|
}
|
||||||
|
|
||||||
|
docType, _ := cmd.Flags().GetString("type")
|
||||||
|
if docType != "" {
|
||||||
|
body["doc_type"] = docType
|
||||||
|
}
|
||||||
|
|
||||||
|
idsStr, _ := cmd.Flags().GetString("ids")
|
||||||
|
if idsStr != "" {
|
||||||
|
ids, err := parseIntList(idsStr)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("invalid --ids: %w", err)
|
||||||
|
}
|
||||||
|
body["document_ids"] = ids
|
||||||
|
}
|
||||||
|
|
||||||
|
fromID, _ := cmd.Flags().GetInt("from-id")
|
||||||
|
if fromID > 0 {
|
||||||
|
body["from_id"] = fromID
|
||||||
|
}
|
||||||
|
|
||||||
|
toID, _ := cmd.Flags().GetInt("to-id")
|
||||||
|
if toID > 0 {
|
||||||
|
body["to_id"] = toID
|
||||||
|
}
|
||||||
|
|
||||||
|
force, _ := cmd.Flags().GetBool("force")
|
||||||
|
if force {
|
||||||
|
body["force"] = true
|
||||||
|
}
|
||||||
|
|
||||||
|
// Ensure at least one filter
|
||||||
|
hasFilter := tagsStr != "" || docType != "" || idsStr != "" || fromID > 0 || toID > 0
|
||||||
|
if !hasFilter {
|
||||||
|
return nil, fmt.Errorf("at least one filter is required (--tags, --type, --ids, --from-id, --to-id)")
|
||||||
|
}
|
||||||
|
|
||||||
|
return body, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func describeBulkFilter(cmd *cobra.Command) string {
|
||||||
|
var parts []string
|
||||||
|
|
||||||
|
tagsStr, _ := cmd.Flags().GetString("tags")
|
||||||
|
if tagsStr != "" {
|
||||||
|
parts = append(parts, fmt.Sprintf("tags=[%s]", tagsStr))
|
||||||
|
}
|
||||||
|
|
||||||
|
docType, _ := cmd.Flags().GetString("type")
|
||||||
|
if docType != "" {
|
||||||
|
parts = append(parts, fmt.Sprintf("type=%s", docType))
|
||||||
|
}
|
||||||
|
|
||||||
|
idsStr, _ := cmd.Flags().GetString("ids")
|
||||||
|
if idsStr != "" {
|
||||||
|
parts = append(parts, fmt.Sprintf("ids=[%s]", idsStr))
|
||||||
|
}
|
||||||
|
|
||||||
|
fromID, _ := cmd.Flags().GetInt("from-id")
|
||||||
|
if fromID > 0 {
|
||||||
|
parts = append(parts, fmt.Sprintf("from_id=%d", fromID))
|
||||||
|
}
|
||||||
|
|
||||||
|
toID, _ := cmd.Flags().GetInt("to-id")
|
||||||
|
if toID > 0 {
|
||||||
|
parts = append(parts, fmt.Sprintf("to_id=%d", toID))
|
||||||
|
}
|
||||||
|
|
||||||
|
return strings.Join(parts, " ")
|
||||||
|
}
|
||||||
|
|
||||||
|
func printBulkResult(action string, result map[string]interface{}) {
|
||||||
|
matched := int(result["matched"].(float64))
|
||||||
|
succeeded := int(result["succeeded"].(float64))
|
||||||
|
failed := int(result["failed"].(float64))
|
||||||
|
|
||||||
|
fmt.Printf("%s %d of %d documents", action, succeeded, matched)
|
||||||
|
if failed > 0 {
|
||||||
|
fmt.Printf(" (%d failed)", failed)
|
||||||
|
}
|
||||||
|
fmt.Println()
|
||||||
|
}
|
||||||
|
|
||||||
|
func parseIntList(s string) ([]int, error) {
|
||||||
|
var ids []int
|
||||||
|
for _, part := range strings.Split(s, ",") {
|
||||||
|
part = strings.TrimSpace(part)
|
||||||
|
if part == "" {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
id, err := strconv.Atoi(part)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("invalid ID %q: %w", part, err)
|
||||||
|
}
|
||||||
|
ids = append(ids, id)
|
||||||
|
}
|
||||||
|
return ids, nil
|
||||||
|
}
|
||||||
@@ -0,0 +1,73 @@
|
|||||||
|
package cmd
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bufio"
|
||||||
|
"fmt"
|
||||||
|
"os"
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
"github.com/kb-search/kb/internal/api"
|
||||||
|
"github.com/kb-search/kb/internal/output"
|
||||||
|
"github.com/spf13/cobra"
|
||||||
|
)
|
||||||
|
|
||||||
|
var bulkSetTagsCmd = &cobra.Command{
|
||||||
|
Use: "bulk-set-tags",
|
||||||
|
Short: "Replace all tags on multiple documents matching a filter",
|
||||||
|
RunE: runBulkSetTags,
|
||||||
|
}
|
||||||
|
|
||||||
|
func init() {
|
||||||
|
addBulkFilterFlags(bulkSetTagsCmd)
|
||||||
|
bulkSetTagsCmd.Flags().String("set", "", "replacement tags (comma-separated)")
|
||||||
|
rootCmd.AddCommand(bulkSetTagsCmd)
|
||||||
|
}
|
||||||
|
|
||||||
|
func runBulkSetTags(cmd *cobra.Command, args []string) error {
|
||||||
|
body, err := buildBulkBody(cmd)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
setStr, _ := cmd.Flags().GetString("set")
|
||||||
|
if setStr == "" {
|
||||||
|
return fmt.Errorf("--set is required (comma-separated list of replacement tags)")
|
||||||
|
}
|
||||||
|
body["new_tags"] = splitTags(setStr)
|
||||||
|
|
||||||
|
yes, _ := cmd.Flags().GetBool("yes")
|
||||||
|
if !yes {
|
||||||
|
desc := describeBulkFilter(cmd)
|
||||||
|
fmt.Printf("This will replace all tags with [%s] on documents matching: %s\nProceed? [y/N] ", setStr, desc)
|
||||||
|
reader := bufio.NewReader(os.Stdin)
|
||||||
|
answer, _ := reader.ReadString('\n')
|
||||||
|
answer = strings.TrimSpace(strings.ToLower(answer))
|
||||||
|
if answer != "y" && answer != "yes" {
|
||||||
|
fmt.Println("Cancelled.")
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
client := api.NewClient()
|
||||||
|
resp, err := client.Post("/api/v1/bulk/set-tags", 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)
|
||||||
|
}
|
||||||
|
|
||||||
|
var result map[string]interface{}
|
||||||
|
if err := api.DecodeJSON(resp, &result); err != nil {
|
||||||
|
return fmt.Errorf("failed to decode response: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if output.IsJSON() {
|
||||||
|
output.PrintJSON(result)
|
||||||
|
} else {
|
||||||
|
printBulkResult("Set tags on", result)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
@@ -0,0 +1,92 @@
|
|||||||
|
package cmd
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bufio"
|
||||||
|
"fmt"
|
||||||
|
"os"
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
"github.com/kb-search/kb/internal/api"
|
||||||
|
"github.com/kb-search/kb/internal/output"
|
||||||
|
"github.com/spf13/cobra"
|
||||||
|
)
|
||||||
|
|
||||||
|
var bulkTagCmd = &cobra.Command{
|
||||||
|
Use: "bulk-tag",
|
||||||
|
Short: "Add or remove tags on multiple documents matching a filter",
|
||||||
|
RunE: runBulkTag,
|
||||||
|
}
|
||||||
|
|
||||||
|
func init() {
|
||||||
|
addBulkFilterFlags(bulkTagCmd)
|
||||||
|
bulkTagCmd.Flags().String("add", "", "tags to add (comma-separated)")
|
||||||
|
bulkTagCmd.Flags().String("remove", "", "tags to remove (comma-separated)")
|
||||||
|
rootCmd.AddCommand(bulkTagCmd)
|
||||||
|
}
|
||||||
|
|
||||||
|
func runBulkTag(cmd *cobra.Command, args []string) error {
|
||||||
|
body, err := buildBulkBody(cmd)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
addStr, _ := cmd.Flags().GetString("add")
|
||||||
|
removeStr, _ := cmd.Flags().GetString("remove")
|
||||||
|
|
||||||
|
if addStr == "" && removeStr == "" {
|
||||||
|
return fmt.Errorf("specify --add and/or --remove")
|
||||||
|
}
|
||||||
|
|
||||||
|
if addStr != "" {
|
||||||
|
body["add"] = splitTags(addStr)
|
||||||
|
}
|
||||||
|
if removeStr != "" {
|
||||||
|
body["remove"] = splitTags(removeStr)
|
||||||
|
}
|
||||||
|
|
||||||
|
yes, _ := cmd.Flags().GetBool("yes")
|
||||||
|
if !yes {
|
||||||
|
desc := describeBulkFilter(cmd)
|
||||||
|
action := ""
|
||||||
|
if addStr != "" {
|
||||||
|
action += fmt.Sprintf("add=[%s]", addStr)
|
||||||
|
}
|
||||||
|
if removeStr != "" {
|
||||||
|
if action != "" {
|
||||||
|
action += " "
|
||||||
|
}
|
||||||
|
action += fmt.Sprintf("remove=[%s]", removeStr)
|
||||||
|
}
|
||||||
|
fmt.Printf("This will update tags (%s) on documents matching: %s\nProceed? [y/N] ", action, desc)
|
||||||
|
reader := bufio.NewReader(os.Stdin)
|
||||||
|
answer, _ := reader.ReadString('\n')
|
||||||
|
answer = strings.TrimSpace(strings.ToLower(answer))
|
||||||
|
if answer != "y" && answer != "yes" {
|
||||||
|
fmt.Println("Cancelled.")
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
client := api.NewClient()
|
||||||
|
resp, err := client.Post("/api/v1/bulk/tags", 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)
|
||||||
|
}
|
||||||
|
|
||||||
|
var result map[string]interface{}
|
||||||
|
if err := api.DecodeJSON(resp, &result); err != nil {
|
||||||
|
return fmt.Errorf("failed to decode response: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if output.IsJSON() {
|
||||||
|
output.PrintJSON(result)
|
||||||
|
} else {
|
||||||
|
printBulkResult("Tagged", result)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
@@ -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
|
||||||
`)
|
`)
|
||||||
|
|||||||
@@ -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
|
||||||
|
}
|
||||||
+9
-2
@@ -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)
|
||||||
@@ -48,6 +54,7 @@ func runInfo(cmd *cobra.Command, args []string) error {
|
|||||||
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"`
|
||||||
|
ChunkCount int `json:"chunk_count"`
|
||||||
Chunks []struct {
|
Chunks []struct {
|
||||||
ID int `json:"id"`
|
ID int `json:"id"`
|
||||||
Page interface{} `json:"page"`
|
Page interface{} `json:"page"`
|
||||||
@@ -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)
|
||||||
|
|
||||||
|
|||||||
+14
-2
@@ -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 {
|
||||||
@@ -62,6 +73,7 @@ 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"`
|
||||||
|
Filename string `json:"original_filename"`
|
||||||
Type string `json:"doc_type"`
|
Type string `json:"doc_type"`
|
||||||
Tags []string `json:"tags"`
|
Tags []string `json:"tags"`
|
||||||
}
|
}
|
||||||
@@ -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
|
||||||
|
|||||||
@@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
+72
-3
@@ -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 {
|
||||||
|
Reranked bool `json:"reranked"`
|
||||||
Results []struct {
|
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 {
|
||||||
|
|||||||
@@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -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
|
||||||
|
}
|
||||||
+3
-2
@@ -43,6 +43,7 @@ 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
|
||||||
|
|||||||
@@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -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."
|
||||||
|
}
|
||||||
|
]
|
||||||
|
}
|
||||||
@@ -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 · 2026-07-07 · 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 1–3 keep 75% retrieval weight, 4–10 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 (0–1). qmd adds a top-rank bonus (+0.05 for #1, +0.02 for #2–3 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 20–30 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, ~1–2s 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
@@ -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"
|
||||||
|
|||||||
@@ -13,15 +13,24 @@ 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 pip install -e .
|
||||||
uv pip install -e . && \
|
|
||||||
uv pip install --no-deps onnxruntime-gpu
|
|
||||||
|
|
||||||
ENV PATH="/app/.venv/bin:$PATH"
|
ENV PATH="/app/.venv/bin:$PATH"
|
||||||
ENV VIRTUAL_ENV="/app/.venv"
|
ENV VIRTUAL_ENV="/app/.venv"
|
||||||
|
|||||||
@@ -1,68 +0,0 @@
|
|||||||
# Stage 1: Build — install Python deps with dev tools available
|
|
||||||
FROM rocm/dev-ubuntu-24.04:6.4-complete AS builder
|
|
||||||
|
|
||||||
ENV DEBIAN_FRONTEND=noninteractive
|
|
||||||
|
|
||||||
RUN apt-get update && apt-get install -y --no-install-recommends \
|
|
||||||
python3.12 python3.12-venv python3.12-dev python3-pip \
|
|
||||||
libpoppler-cpp-dev poppler-utils \
|
|
||||||
build-essential curl \
|
|
||||||
&& rm -rf /var/lib/apt/lists/*
|
|
||||||
|
|
||||||
COPY --from=ghcr.io/astral-sh/uv:latest /uv /usr/local/bin/uv
|
|
||||||
|
|
||||||
WORKDIR /app
|
|
||||||
|
|
||||||
COPY pyproject.toml ./
|
|
||||||
COPY kb/ kb/
|
|
||||||
COPY main.py ./
|
|
||||||
COPY VERSION ./
|
|
||||||
|
|
||||||
RUN uv venv .venv && \
|
|
||||||
. .venv/bin/activate && \
|
|
||||||
uv pip install -e . && \
|
|
||||||
uv pip install --no-deps onnxruntime-rocm
|
|
||||||
|
|
||||||
# Stage 2: Runtime — minimal ROCm runtime libs only
|
|
||||||
FROM ubuntu:24.04
|
|
||||||
|
|
||||||
ENV DEBIAN_FRONTEND=noninteractive
|
|
||||||
|
|
||||||
# Add ROCm apt repository
|
|
||||||
RUN apt-get update && apt-get install -y --no-install-recommends \
|
|
||||||
ca-certificates curl gnupg \
|
|
||||||
&& mkdir -p /etc/apt/keyrings \
|
|
||||||
&& curl -fsSL https://repo.radeon.com/rocm/rocm.gpg.key \
|
|
||||||
| gpg --dearmor -o /etc/apt/keyrings/rocm.gpg \
|
|
||||||
&& echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/rocm.gpg] https://repo.radeon.com/rocm/apt/6.4.1 noble main" \
|
|
||||||
> /etc/apt/sources.list.d/rocm.list \
|
|
||||||
&& printf 'Package: *\nPin: release o=repo.radeon.com\nPin-Priority: 600\n' \
|
|
||||||
> /etc/apt/preferences.d/rocm-pin-600 \
|
|
||||||
&& apt-get update && apt-get install -y --no-install-recommends \
|
|
||||||
python3.12 python3.12-venv \
|
|
||||||
libpoppler-cpp0t64 poppler-utils \
|
|
||||||
libgl1 libglib2.0-0 \
|
|
||||||
rocm-hip-runtime \
|
|
||||||
rocm-hip-libraries \
|
|
||||||
miopen-hip \
|
|
||||||
&& rm -rf /var/lib/apt/lists/*
|
|
||||||
|
|
||||||
WORKDIR /app
|
|
||||||
|
|
||||||
# Copy built venv and application from builder
|
|
||||||
COPY --from=builder /app/.venv .venv
|
|
||||||
COPY --from=builder /app/kb kb
|
|
||||||
COPY --from=builder /app/main.py .
|
|
||||||
COPY --from=builder /app/pyproject.toml .
|
|
||||||
COPY --from=builder /app/VERSION .
|
|
||||||
|
|
||||||
ENV PATH="/app/.venv/bin:$PATH"
|
|
||||||
ENV VIRTUAL_ENV="/app/.venv"
|
|
||||||
ENV KB_DEVICE=auto
|
|
||||||
ENV KB_INGEST_DEVICE=auto
|
|
||||||
ENV KB_DATA_DIR=/data
|
|
||||||
|
|
||||||
EXPOSE 8000
|
|
||||||
VOLUME ["/data"]
|
|
||||||
|
|
||||||
CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000"]
|
|
||||||
+1
-1
@@ -1 +1 @@
|
|||||||
3.0.1
|
3.3.0
|
||||||
|
|||||||
@@ -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
|
||||||
|
|
||||||
|
|||||||
@@ -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
|
||||||
|
|
||||||
|
|||||||
@@ -1,38 +0,0 @@
|
|||||||
services:
|
|
||||||
kb-engine:
|
|
||||||
build:
|
|
||||||
context: .
|
|
||||||
dockerfile: Dockerfile.rocm
|
|
||||||
devices:
|
|
||||||
- "/dev/kfd"
|
|
||||||
- "/dev/dri"
|
|
||||||
group_add:
|
|
||||||
- "video"
|
|
||||||
ports:
|
|
||||||
- "${KB_PORT:-8000}:8000"
|
|
||||||
volumes:
|
|
||||||
- ${KB_DATA_PATH:-./data}:/data
|
|
||||||
environment:
|
|
||||||
- KB_MODEL=${KB_MODEL:-all-MiniLM-L6-v2}
|
|
||||||
- KB_DEVICE=${KB_DEVICE:-auto}
|
|
||||||
- KB_INGEST_DEVICE=${KB_INGEST_DEVICE:-auto}
|
|
||||||
- KB_API_KEY=${KB_API_KEY:-}
|
|
||||||
- KB_SEARCH_THRESHOLD=${KB_SEARCH_THRESHOLD:-0.01}
|
|
||||||
- HF_HUB_OFFLINE=${HF_HUB_OFFLINE:-}
|
|
||||||
restart: unless-stopped
|
|
||||||
|
|
||||||
kb-mcp:
|
|
||||||
build:
|
|
||||||
context: ../mcp
|
|
||||||
dockerfile: Dockerfile
|
|
||||||
ports:
|
|
||||||
- "${KB_MCP_PORT:-3000}:3000"
|
|
||||||
environment:
|
|
||||||
- KB_ENGINE_URL=http://kb-engine:8000
|
|
||||||
- KB_API_KEY=${KB_API_KEY:-}
|
|
||||||
- KB_MCP_API_KEY=${KB_MCP_API_KEY:-}
|
|
||||||
# Comma-separated IPs/FQDNs allowed to connect remotely (e.g. 192.168.1.50,kb.example.com)
|
|
||||||
- KB_MCP_ALLOWED_HOSTS=${KB_MCP_ALLOWED_HOSTS:-}
|
|
||||||
depends_on:
|
|
||||||
- kb-engine
|
|
||||||
restart: unless-stopped
|
|
||||||
@@ -20,6 +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.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"))
|
||||||
|
|
||||||
|
|||||||
@@ -74,6 +74,7 @@ def get_connection(db_path: str) -> sqlite3.Connection:
|
|||||||
conn.enable_load_extension(False)
|
conn.enable_load_extension(False)
|
||||||
conn.row_factory = sqlite3.Row
|
conn.row_factory = sqlite3.Row
|
||||||
conn.execute("PRAGMA journal_mode=WAL")
|
conn.execute("PRAGMA journal_mode=WAL")
|
||||||
|
conn.execute("PRAGMA busy_timeout=5000")
|
||||||
conn.execute("PRAGMA foreign_keys=ON")
|
conn.execute("PRAGMA foreign_keys=ON")
|
||||||
return conn
|
return conn
|
||||||
|
|
||||||
@@ -189,6 +190,16 @@ def init_schema(conn: sqlite3.Connection, embedding_dim: int) -> None:
|
|||||||
if "updated_at" not in doc_cols:
|
if "updated_at" not in doc_cols:
|
||||||
conn.execute("ALTER TABLE documents ADD COLUMN updated_at TEXT")
|
conn.execute("ALTER TABLE documents ADD COLUMN updated_at TEXT")
|
||||||
|
|
||||||
|
# Migrate: add job_type to jobs if missing (bulk operations)
|
||||||
|
job_cols = {row[1] for row in conn.execute("PRAGMA table_info(jobs)").fetchall()}
|
||||||
|
if "job_type" not in job_cols:
|
||||||
|
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()
|
||||||
|
|
||||||
|
|
||||||
@@ -329,6 +340,92 @@ def untag_document(conn: sqlite3.Connection, document_id: int, tag_names: list[s
|
|||||||
conn.commit()
|
conn.commit()
|
||||||
|
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# Bulk operation helpers
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
def resolve_bulk_selection(
|
||||||
|
conn: sqlite3.Connection,
|
||||||
|
document_ids: list[int] | None = None,
|
||||||
|
tags: list[str] | None = None,
|
||||||
|
doc_type: str | None = None,
|
||||||
|
from_id: int | None = None,
|
||||||
|
to_id: int | None = None,
|
||||||
|
) -> list[int]:
|
||||||
|
"""Return document IDs matching the bulk selection filter.
|
||||||
|
|
||||||
|
Filters combine with AND logic. At least one filter must be provided.
|
||||||
|
"""
|
||||||
|
sql = "SELECT DISTINCT d.id FROM documents d"
|
||||||
|
joins: list[str] = []
|
||||||
|
where: list[str] = []
|
||||||
|
params: list = []
|
||||||
|
|
||||||
|
if tags:
|
||||||
|
for i, tag in enumerate(tags):
|
||||||
|
joins.append(f"JOIN document_tags dt{i} ON d.id = dt{i}.document_id")
|
||||||
|
joins.append(f"JOIN tags t{i} ON dt{i}.tag_id = t{i}.id")
|
||||||
|
where.append(f"t{i}.name = ?")
|
||||||
|
params.append(tag)
|
||||||
|
|
||||||
|
if doc_type:
|
||||||
|
where.append("d.doc_type = ?")
|
||||||
|
params.append(doc_type)
|
||||||
|
|
||||||
|
if document_ids:
|
||||||
|
placeholders = ",".join("?" for _ in document_ids)
|
||||||
|
where.append(f"d.id IN ({placeholders})")
|
||||||
|
params.extend(document_ids)
|
||||||
|
|
||||||
|
if from_id is not None:
|
||||||
|
where.append("d.id >= ?")
|
||||||
|
params.append(from_id)
|
||||||
|
|
||||||
|
if to_id is not None:
|
||||||
|
where.append("d.id <= ?")
|
||||||
|
params.append(to_id)
|
||||||
|
|
||||||
|
if joins:
|
||||||
|
sql += " " + " ".join(joins)
|
||||||
|
if where:
|
||||||
|
sql += " WHERE " + " AND ".join(where)
|
||||||
|
|
||||||
|
rows = conn.execute(sql, params).fetchall()
|
||||||
|
return [row["id"] for row in rows]
|
||||||
|
|
||||||
|
|
||||||
|
def create_bulk_job(
|
||||||
|
conn: sqlite3.Connection,
|
||||||
|
job_type: str,
|
||||||
|
filters_json: str,
|
||||||
|
matched: int,
|
||||||
|
succeeded: int,
|
||||||
|
failed: int,
|
||||||
|
errors_json: str = "[]",
|
||||||
|
) -> int:
|
||||||
|
"""Create an audit log entry for a bulk operation and return its id."""
|
||||||
|
cur = conn.execute(
|
||||||
|
"""INSERT INTO jobs(filename, status, job_type, document_id, chunk_count, error, completed_at)
|
||||||
|
VALUES (?, ?, ?, ?, ?, ?, current_timestamp)""",
|
||||||
|
(
|
||||||
|
filters_json,
|
||||||
|
"done" if failed == 0 else "partial_failure",
|
||||||
|
job_type,
|
||||||
|
matched,
|
||||||
|
succeeded,
|
||||||
|
errors_json if failed > 0 else None,
|
||||||
|
),
|
||||||
|
)
|
||||||
|
conn.commit()
|
||||||
|
return cur.lastrowid
|
||||||
|
|
||||||
|
|
||||||
|
def count_documents(conn: sqlite3.Connection) -> int:
|
||||||
|
"""Return total number of documents in the database."""
|
||||||
|
row = conn.execute("SELECT COUNT(*) AS cnt FROM documents").fetchone()
|
||||||
|
return row["cnt"]
|
||||||
|
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
# ---------------------------------------------------------------------------
|
||||||
# Vec table management
|
# Vec table management
|
||||||
# ---------------------------------------------------------------------------
|
# ---------------------------------------------------------------------------
|
||||||
|
|||||||
@@ -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
|
||||||
@@ -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"),
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -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,
|
||||||
|
|||||||
@@ -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
|
||||||
@@ -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()
|
||||||
@@ -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()
|
||||||
@@ -0,0 +1,281 @@
|
|||||||
|
"""Bulk operation endpoints — delete, tag, and set-tags on multiple documents."""
|
||||||
|
|
||||||
|
import json
|
||||||
|
import logging
|
||||||
|
from pathlib import Path
|
||||||
|
from typing import Optional
|
||||||
|
|
||||||
|
from fastapi import HTTPException
|
||||||
|
from pydantic import BaseModel, model_validator
|
||||||
|
|
||||||
|
from main import app
|
||||||
|
from kb.config import cfg
|
||||||
|
from kb.database import (
|
||||||
|
get_connection,
|
||||||
|
resolve_bulk_selection,
|
||||||
|
count_documents,
|
||||||
|
create_bulk_job,
|
||||||
|
tag_document,
|
||||||
|
untag_document,
|
||||||
|
)
|
||||||
|
|
||||||
|
logger = logging.getLogger("kb.routes.bulk")
|
||||||
|
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# Request models
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
class BulkSelectionRequest(BaseModel):
|
||||||
|
document_ids: Optional[list[int]] = None
|
||||||
|
tags: Optional[list[str]] = None
|
||||||
|
doc_type: Optional[str] = None
|
||||||
|
from_id: Optional[int] = None
|
||||||
|
to_id: Optional[int] = None
|
||||||
|
force: bool = False
|
||||||
|
|
||||||
|
@model_validator(mode="after")
|
||||||
|
def require_at_least_one_filter(self):
|
||||||
|
if not any([self.document_ids, self.tags, self.doc_type,
|
||||||
|
self.from_id is not None, self.to_id is not None]):
|
||||||
|
raise ValueError("At least one selection filter is required")
|
||||||
|
return self
|
||||||
|
|
||||||
|
|
||||||
|
class BulkDeleteRequest(BulkSelectionRequest):
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
class BulkTagsRequest(BulkSelectionRequest):
|
||||||
|
add: Optional[list[str]] = None
|
||||||
|
remove: Optional[list[str]] = None
|
||||||
|
|
||||||
|
@model_validator(mode="after")
|
||||||
|
def require_add_or_remove(self):
|
||||||
|
if not self.add and not self.remove:
|
||||||
|
raise ValueError("At least one of 'add' or 'remove' is required")
|
||||||
|
return self
|
||||||
|
|
||||||
|
|
||||||
|
class BulkSetTagsRequest(BulkSelectionRequest):
|
||||||
|
new_tags: list[str]
|
||||||
|
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# Shared helpers
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
def _check_safety_threshold(matched: int, total: int, force: bool) -> None:
|
||||||
|
"""Raise 409 if the operation would affect too many documents."""
|
||||||
|
threshold = cfg.bulk_safety_percent
|
||||||
|
if threshold <= 0 or force or total == 0:
|
||||||
|
return
|
||||||
|
percent = (matched / total) * 100
|
||||||
|
if percent > threshold:
|
||||||
|
raise HTTPException(
|
||||||
|
status_code=409,
|
||||||
|
detail={
|
||||||
|
"error": "safety_threshold_exceeded",
|
||||||
|
"message": (
|
||||||
|
f"Operation would affect {matched} of {total} documents "
|
||||||
|
f"({percent:.1f}%). Exceeds safety threshold of {threshold}%. "
|
||||||
|
f"Use force: true to proceed."
|
||||||
|
),
|
||||||
|
"matched": matched,
|
||||||
|
"total": total,
|
||||||
|
"percent": round(percent, 1),
|
||||||
|
"threshold": threshold,
|
||||||
|
},
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _filters_dict(req: BulkSelectionRequest) -> str:
|
||||||
|
"""Build a JSON string of the selection filter for audit logging."""
|
||||||
|
d = {}
|
||||||
|
if req.document_ids:
|
||||||
|
d["document_ids"] = req.document_ids
|
||||||
|
if req.tags:
|
||||||
|
d["tags"] = req.tags
|
||||||
|
if req.doc_type:
|
||||||
|
d["doc_type"] = req.doc_type
|
||||||
|
if req.from_id is not None:
|
||||||
|
d["from_id"] = req.from_id
|
||||||
|
if req.to_id is not None:
|
||||||
|
d["to_id"] = req.to_id
|
||||||
|
return json.dumps(d)
|
||||||
|
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# Endpoints
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
@app.post("/api/v1/bulk/delete")
|
||||||
|
async def bulk_delete(req: BulkDeleteRequest):
|
||||||
|
conn = get_connection(cfg.db_path)
|
||||||
|
try:
|
||||||
|
doc_ids = resolve_bulk_selection(
|
||||||
|
conn, req.document_ids, req.tags, req.doc_type, req.from_id, req.to_id,
|
||||||
|
)
|
||||||
|
total = count_documents(conn)
|
||||||
|
_check_safety_threshold(len(doc_ids), total, req.force)
|
||||||
|
|
||||||
|
succeeded = 0
|
||||||
|
failed = 0
|
||||||
|
errors = []
|
||||||
|
stored_files: list[str] = []
|
||||||
|
|
||||||
|
for doc_id in doc_ids:
|
||||||
|
try:
|
||||||
|
doc = conn.execute(
|
||||||
|
"SELECT id, stored_path FROM documents WHERE id = ?", (doc_id,)
|
||||||
|
).fetchone()
|
||||||
|
if not doc:
|
||||||
|
failed += 1
|
||||||
|
errors.append({"document_id": doc_id, "error": "not found"})
|
||||||
|
continue
|
||||||
|
|
||||||
|
if doc["stored_path"]:
|
||||||
|
stored_files.append(doc["stored_path"])
|
||||||
|
|
||||||
|
# Delete embeddings
|
||||||
|
chunk_ids = conn.execute(
|
||||||
|
"SELECT id FROM chunks WHERE document_id = ?", (doc_id,)
|
||||||
|
).fetchall()
|
||||||
|
for row in chunk_ids:
|
||||||
|
conn.execute("DELETE FROM chunks_vec WHERE chunk_id = ?", (row["id"],))
|
||||||
|
|
||||||
|
# Delete document (cascades to chunks, document_tags)
|
||||||
|
conn.execute("DELETE FROM documents WHERE id = ?", (doc_id,))
|
||||||
|
succeeded += 1
|
||||||
|
except Exception as exc:
|
||||||
|
failed += 1
|
||||||
|
errors.append({"document_id": doc_id, "error": str(exc)})
|
||||||
|
|
||||||
|
conn.commit()
|
||||||
|
|
||||||
|
# Best-effort file cleanup after commit
|
||||||
|
for path in stored_files:
|
||||||
|
try:
|
||||||
|
f = Path(path)
|
||||||
|
if f.exists():
|
||||||
|
f.unlink()
|
||||||
|
except OSError as exc:
|
||||||
|
logger.warning("Failed to delete stored file %s: %s", path, exc)
|
||||||
|
|
||||||
|
errors_json = json.dumps(errors) if errors else "[]"
|
||||||
|
job_id = create_bulk_job(
|
||||||
|
conn, "bulk_delete", _filters_dict(req),
|
||||||
|
len(doc_ids), succeeded, failed, errors_json,
|
||||||
|
)
|
||||||
|
|
||||||
|
return {
|
||||||
|
"job_id": job_id,
|
||||||
|
"status": "done" if failed == 0 else "partial_failure",
|
||||||
|
"matched": len(doc_ids),
|
||||||
|
"succeeded": succeeded,
|
||||||
|
"failed": failed,
|
||||||
|
"errors": errors,
|
||||||
|
}
|
||||||
|
finally:
|
||||||
|
conn.close()
|
||||||
|
|
||||||
|
|
||||||
|
@app.post("/api/v1/bulk/tags")
|
||||||
|
async def bulk_tags(req: BulkTagsRequest):
|
||||||
|
conn = get_connection(cfg.db_path)
|
||||||
|
try:
|
||||||
|
doc_ids = resolve_bulk_selection(
|
||||||
|
conn, req.document_ids, req.tags, req.doc_type, req.from_id, req.to_id,
|
||||||
|
)
|
||||||
|
total = count_documents(conn)
|
||||||
|
_check_safety_threshold(len(doc_ids), total, req.force)
|
||||||
|
|
||||||
|
succeeded = 0
|
||||||
|
failed = 0
|
||||||
|
errors = []
|
||||||
|
|
||||||
|
for doc_id in doc_ids:
|
||||||
|
try:
|
||||||
|
if req.add:
|
||||||
|
tag_document(conn, doc_id, req.add)
|
||||||
|
if req.remove:
|
||||||
|
untag_document(conn, doc_id, req.remove)
|
||||||
|
conn.execute(
|
||||||
|
"UPDATE documents SET updated_at = current_timestamp WHERE id = ?",
|
||||||
|
(doc_id,),
|
||||||
|
)
|
||||||
|
succeeded += 1
|
||||||
|
except Exception as exc:
|
||||||
|
failed += 1
|
||||||
|
errors.append({"document_id": doc_id, "error": str(exc)})
|
||||||
|
|
||||||
|
conn.commit()
|
||||||
|
|
||||||
|
errors_json = json.dumps(errors) if errors else "[]"
|
||||||
|
job_id = create_bulk_job(
|
||||||
|
conn, "bulk_tags", _filters_dict(req),
|
||||||
|
len(doc_ids), succeeded, failed, errors_json,
|
||||||
|
)
|
||||||
|
|
||||||
|
return {
|
||||||
|
"job_id": job_id,
|
||||||
|
"status": "done" if failed == 0 else "partial_failure",
|
||||||
|
"matched": len(doc_ids),
|
||||||
|
"succeeded": succeeded,
|
||||||
|
"failed": failed,
|
||||||
|
"errors": errors,
|
||||||
|
}
|
||||||
|
finally:
|
||||||
|
conn.close()
|
||||||
|
|
||||||
|
|
||||||
|
@app.post("/api/v1/bulk/set-tags")
|
||||||
|
async def bulk_set_tags(req: BulkSetTagsRequest):
|
||||||
|
conn = get_connection(cfg.db_path)
|
||||||
|
try:
|
||||||
|
doc_ids = resolve_bulk_selection(
|
||||||
|
conn, req.document_ids, req.tags, req.doc_type, req.from_id, req.to_id,
|
||||||
|
)
|
||||||
|
total = count_documents(conn)
|
||||||
|
_check_safety_threshold(len(doc_ids), total, req.force)
|
||||||
|
|
||||||
|
succeeded = 0
|
||||||
|
failed = 0
|
||||||
|
errors = []
|
||||||
|
|
||||||
|
for doc_id in doc_ids:
|
||||||
|
try:
|
||||||
|
# Remove all existing tags
|
||||||
|
conn.execute(
|
||||||
|
"DELETE FROM document_tags WHERE document_id = ?", (doc_id,)
|
||||||
|
)
|
||||||
|
# Apply new tag set
|
||||||
|
if req.new_tags:
|
||||||
|
tag_document(conn, doc_id, req.new_tags)
|
||||||
|
conn.execute(
|
||||||
|
"UPDATE documents SET updated_at = current_timestamp WHERE id = ?",
|
||||||
|
(doc_id,),
|
||||||
|
)
|
||||||
|
succeeded += 1
|
||||||
|
except Exception as exc:
|
||||||
|
failed += 1
|
||||||
|
errors.append({"document_id": doc_id, "error": str(exc)})
|
||||||
|
|
||||||
|
conn.commit()
|
||||||
|
|
||||||
|
errors_json = json.dumps(errors) if errors else "[]"
|
||||||
|
job_id = create_bulk_job(
|
||||||
|
conn, "bulk_set_tags", _filters_dict(req),
|
||||||
|
len(doc_ids), succeeded, failed, errors_json,
|
||||||
|
)
|
||||||
|
|
||||||
|
return {
|
||||||
|
"job_id": job_id,
|
||||||
|
"status": "done" if failed == 0 else "partial_failure",
|
||||||
|
"matched": len(doc_ids),
|
||||||
|
"succeeded": succeeded,
|
||||||
|
"failed": failed,
|
||||||
|
"errors": errors,
|
||||||
|
}
|
||||||
|
finally:
|
||||||
|
conn.close()
|
||||||
@@ -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,6 +147,11 @@ 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.")
|
||||||
|
|
||||||
|
chunk_count = conn.execute(
|
||||||
|
"SELECT COUNT(*) AS n FROM chunks WHERE document_id = ?", (doc_id,)
|
||||||
|
).fetchone()["n"]
|
||||||
|
chunks = []
|
||||||
|
if include_chunks:
|
||||||
chunks = conn.execute(
|
chunks = conn.execute(
|
||||||
"SELECT * FROM chunks WHERE document_id = ? ORDER BY chunk_index",
|
"SELECT * FROM chunks WHERE document_id = ? ORDER BY chunk_index",
|
||||||
(doc_id,),
|
(doc_id,),
|
||||||
@@ -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:
|
||||||
|
|||||||
@@ -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 []
|
||||||
|
|||||||
@@ -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:
|
||||||
|
|||||||
@@ -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()
|
||||||
|
|||||||
@@ -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
@@ -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
|
||||||
|
|||||||
@@ -16,7 +16,8 @@ def stage_file(staging_dir: Path, filename: str, content: bytes) -> Path:
|
|||||||
The path to the newly created staged file.
|
The path to the newly created staged file.
|
||||||
"""
|
"""
|
||||||
staging_dir.mkdir(parents=True, exist_ok=True)
|
staging_dir.mkdir(parents=True, exist_ok=True)
|
||||||
dest = staging_dir / f"{uuid.uuid4()}_{filename}"
|
safe_filename = filename.replace("/", "_").replace("\\", "_")
|
||||||
|
dest = staging_dir / f"{uuid.uuid4()}_{safe_filename}"
|
||||||
dest.write_bytes(content)
|
dest.write_bytes(content)
|
||||||
logger.debug("Staged file: %s (%d bytes)", dest, len(content))
|
logger.debug("Staged file: %s (%d bytes)", dest, len(content))
|
||||||
return dest
|
return dest
|
||||||
@@ -31,7 +32,8 @@ def stage_note(staging_dir: Path, title: str, text: str) -> Path:
|
|||||||
The path to the newly created staged note file.
|
The path to the newly created staged note file.
|
||||||
"""
|
"""
|
||||||
staging_dir.mkdir(parents=True, exist_ok=True)
|
staging_dir.mkdir(parents=True, exist_ok=True)
|
||||||
dest = staging_dir / f"{uuid.uuid4()}_{title}.note"
|
safe_title = title.replace("/", "_").replace("\\", "_")
|
||||||
|
dest = staging_dir / f"{uuid.uuid4()}_{safe_title}.note"
|
||||||
dest.write_text(text, encoding="utf-8")
|
dest.write_text(text, encoding="utf-8")
|
||||||
logger.debug("Staged note: %s (%d chars)", dest, len(text))
|
logger.debug("Staged note: %s (%d chars)", dest, len(text))
|
||||||
return dest
|
return dest
|
||||||
|
|||||||
+9
-1
@@ -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}")
|
||||||
|
|
||||||
|
|||||||
+13
-1
@@ -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())
|
||||||
|
|
||||||
@@ -62,7 +74,7 @@ async def lifespan(app: FastAPI):
|
|||||||
app = FastAPI(title="kb-engine", version=__version__, lifespan=lifespan)
|
app = FastAPI(title="kb-engine", version=__version__, lifespan=lifespan)
|
||||||
|
|
||||||
# Import routes after app is created
|
# Import routes after app is created
|
||||||
from kb.routes import health, search, jobs, documents, tags, status, reindex, auth, notes # noqa: E402, F401
|
from kb.routes import health, search, jobs, documents, tags, status, reindex, auth, notes, bulk # noqa: E402, F401
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
import uvicorn
|
import uvicorn
|
||||||
|
|||||||
@@ -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))
|
||||||
@@ -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",
|
||||||
|
}]
|
||||||
@@ -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") == []
|
||||||
@@ -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)
|
||||||
@@ -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
|
||||||
@@ -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)
|
||||||
@@ -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"] == {}
|
||||||
+93
-1
@@ -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()
|
||||||
@@ -106,6 +111,93 @@ def update_tags(doc_id: int, add: list[str] | None = None,
|
|||||||
return r.json()
|
return r.json()
|
||||||
|
|
||||||
|
|
||||||
|
def delete_document(doc_id: int) -> dict:
|
||||||
|
with _client() as c:
|
||||||
|
r = c.delete(f"/api/v1/documents/{doc_id}")
|
||||||
|
r.raise_for_status()
|
||||||
|
return r.json()
|
||||||
|
|
||||||
|
|
||||||
|
def _bulk_body(
|
||||||
|
document_ids: list[int] | None = None,
|
||||||
|
tags: list[str] | None = None,
|
||||||
|
doc_type: str | None = None,
|
||||||
|
from_id: int | None = None,
|
||||||
|
to_id: int | None = None,
|
||||||
|
force: bool = False,
|
||||||
|
**extra,
|
||||||
|
) -> dict:
|
||||||
|
body: dict = {}
|
||||||
|
if document_ids:
|
||||||
|
body["document_ids"] = document_ids
|
||||||
|
if tags:
|
||||||
|
body["tags"] = tags
|
||||||
|
if doc_type:
|
||||||
|
body["doc_type"] = doc_type
|
||||||
|
if from_id is not None:
|
||||||
|
body["from_id"] = from_id
|
||||||
|
if to_id is not None:
|
||||||
|
body["to_id"] = to_id
|
||||||
|
if force:
|
||||||
|
body["force"] = True
|
||||||
|
body.update(extra)
|
||||||
|
return body
|
||||||
|
|
||||||
|
|
||||||
|
def bulk_delete(
|
||||||
|
document_ids: list[int] | None = None,
|
||||||
|
tags: list[str] | None = None,
|
||||||
|
doc_type: str | None = None,
|
||||||
|
from_id: int | None = None,
|
||||||
|
to_id: int | None = None,
|
||||||
|
force: bool = False,
|
||||||
|
) -> dict:
|
||||||
|
body = _bulk_body(document_ids, tags, doc_type, from_id, to_id, force)
|
||||||
|
with _client() as c:
|
||||||
|
r = c.post("/api/v1/bulk/delete", json=body)
|
||||||
|
r.raise_for_status()
|
||||||
|
return r.json()
|
||||||
|
|
||||||
|
|
||||||
|
def bulk_tags(
|
||||||
|
document_ids: list[int] | None = None,
|
||||||
|
tags: list[str] | None = None,
|
||||||
|
doc_type: str | None = None,
|
||||||
|
from_id: int | None = None,
|
||||||
|
to_id: int | None = None,
|
||||||
|
add: list[str] | None = None,
|
||||||
|
remove: list[str] | None = None,
|
||||||
|
force: bool = False,
|
||||||
|
) -> dict:
|
||||||
|
extra = {}
|
||||||
|
if add:
|
||||||
|
extra["add"] = add
|
||||||
|
if remove:
|
||||||
|
extra["remove"] = remove
|
||||||
|
body = _bulk_body(document_ids, tags, doc_type, from_id, to_id, force, **extra)
|
||||||
|
with _client() as c:
|
||||||
|
r = c.post("/api/v1/bulk/tags", json=body)
|
||||||
|
r.raise_for_status()
|
||||||
|
return r.json()
|
||||||
|
|
||||||
|
|
||||||
|
def bulk_set_tags(
|
||||||
|
document_ids: list[int] | None = None,
|
||||||
|
tags: list[str] | None = None,
|
||||||
|
doc_type: str | None = None,
|
||||||
|
from_id: int | None = None,
|
||||||
|
to_id: int | None = None,
|
||||||
|
new_tags: list[str] | None = None,
|
||||||
|
force: bool = False,
|
||||||
|
) -> dict:
|
||||||
|
extra = {"new_tags": new_tags or []}
|
||||||
|
body = _bulk_body(document_ids, tags, doc_type, from_id, to_id, force, **extra)
|
||||||
|
with _client() as c:
|
||||||
|
r = c.post("/api/v1/bulk/set-tags", json=body)
|
||||||
|
r.raise_for_status()
|
||||||
|
return r.json()
|
||||||
|
|
||||||
|
|
||||||
def upload_file(filename: str, file_bytes: bytes,
|
def upload_file(filename: str, file_bytes: bytes,
|
||||||
tags: list[str] | None = None) -> dict:
|
tags: list[str] | None = None) -> dict:
|
||||||
fields: dict = {}
|
fields: dict = {}
|
||||||
|
|||||||
+175
-100
@@ -20,68 +20,6 @@ import uploads
|
|||||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
|
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
|
||||||
logger = logging.getLogger("kb.mcp")
|
logger = logging.getLogger("kb.mcp")
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
# Collection helpers
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
|
|
||||||
COLLECTION_TAG_PREFIX = "collection:"
|
|
||||||
DEFAULT_COLLECTION = "documents"
|
|
||||||
|
|
||||||
|
|
||||||
def _collection_tag(collection: str | None) -> str:
|
|
||||||
return f"{COLLECTION_TAG_PREFIX}{collection or DEFAULT_COLLECTION}"
|
|
||||||
|
|
||||||
|
|
||||||
def _strip_collection_tags(tags: list[str]) -> tuple[str | None, list[str]]:
|
|
||||||
"""Split tags into (collection, remaining_tags)."""
|
|
||||||
collection = None
|
|
||||||
remaining = []
|
|
||||||
for t in tags:
|
|
||||||
if t.startswith(COLLECTION_TAG_PREFIX):
|
|
||||||
collection = t[len(COLLECTION_TAG_PREFIX):]
|
|
||||||
else:
|
|
||||||
remaining.append(t)
|
|
||||||
return collection, remaining
|
|
||||||
|
|
||||||
|
|
||||||
def _process_document(doc: dict) -> dict:
|
|
||||||
"""Strip collection tags from a document dict and add collection field."""
|
|
||||||
tags = doc.get("tags", [])
|
|
||||||
collection, clean_tags = _strip_collection_tags(tags)
|
|
||||||
doc["tags"] = clean_tags
|
|
||||||
doc["collection"] = collection
|
|
||||||
return doc
|
|
||||||
|
|
||||||
|
|
||||||
def _process_search_results(results: list[dict]) -> list[dict]:
|
|
||||||
"""Strip collection tags from search result dicts."""
|
|
||||||
for r in results:
|
|
||||||
if "tags" in r:
|
|
||||||
collection, clean_tags = _strip_collection_tags(r["tags"])
|
|
||||||
r["tags"] = clean_tags
|
|
||||||
r["collection"] = collection
|
|
||||||
if "document" in r and "tags" in r["document"]:
|
|
||||||
collection, clean_tags = _strip_collection_tags(r["document"]["tags"])
|
|
||||||
r["document"]["tags"] = clean_tags
|
|
||||||
r["document"]["collection"] = collection
|
|
||||||
return results
|
|
||||||
|
|
||||||
|
|
||||||
async def _ensure_exclusive_collection(doc_id: int, collection: str) -> None:
|
|
||||||
"""Remove existing collection tags and apply the new one."""
|
|
||||||
doc = engine.get_document(doc_id)
|
|
||||||
existing_collection_tags = [
|
|
||||||
t for t in doc.get("tags", [])
|
|
||||||
if t.startswith(COLLECTION_TAG_PREFIX)
|
|
||||||
]
|
|
||||||
new_tag = _collection_tag(collection)
|
|
||||||
if existing_collection_tags == [new_tag]:
|
|
||||||
return
|
|
||||||
if existing_collection_tags:
|
|
||||||
engine.update_tags(doc_id, remove=existing_collection_tags)
|
|
||||||
engine.update_tags(doc_id, add=[new_tag])
|
|
||||||
|
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
# ---------------------------------------------------------------------------
|
||||||
# Transport security — DNS rebinding protection with configurable allowed hosts
|
# Transport security — DNS rebinding protection with configurable allowed hosts
|
||||||
# ---------------------------------------------------------------------------
|
# ---------------------------------------------------------------------------
|
||||||
@@ -106,8 +44,16 @@ _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. This server requires Bearer token "
|
"kb_search uses dense vector embeddings (semantic similarity) fused with "
|
||||||
|
"BM25 full-text ranking, so it finds conceptually related content even "
|
||||||
|
"when the exact words don't match — agents can ask natural-language "
|
||||||
|
"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 "
|
"authentication — all requests are authenticated via the Authorization "
|
||||||
"header at the HTTP transport layer."
|
"header at the HTTP transport layer."
|
||||||
),
|
),
|
||||||
@@ -121,50 +67,67 @@ async def kb_search(
|
|||||||
top: int = 10,
|
top: int = 10,
|
||||||
tags: list[str] | None = None,
|
tags: list[str] | None = None,
|
||||||
doc_type: str | None = None,
|
doc_type: str | None = None,
|
||||||
collection: 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",
|
||||||
collection: Filter by collection name (e.g. "documents", "memory", "workspace").
|
"code", "data").
|
||||||
fts_only: If true, use only full-text search (no vector similarity).
|
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.
|
||||||
"""
|
"""
|
||||||
search_tags = list(tags) if tags else []
|
|
||||||
if collection:
|
|
||||||
search_tags.append(_collection_tag(collection))
|
|
||||||
|
|
||||||
result = engine.search(
|
result = engine.search(
|
||||||
query=query,
|
query=query,
|
||||||
top=top,
|
top=top,
|
||||||
tags=search_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", [])
|
||||||
processed = _process_search_results(results_list)
|
return json.dumps(results_list, indent=2)
|
||||||
return json.dumps(processed, indent=2)
|
|
||||||
|
|
||||||
|
|
||||||
@mcp.tool()
|
@mcp.tool()
|
||||||
async def kb_addnote(
|
async def kb_addnote(
|
||||||
text: str,
|
text: str,
|
||||||
collection: str | None = None,
|
|
||||||
tags: list[str] | None = None,
|
tags: list[str] | None = None,
|
||||||
title: str | None = None,
|
title: str | None = None,
|
||||||
) -> str:
|
) -> str:
|
||||||
@@ -175,15 +138,10 @@ async def kb_addnote(
|
|||||||
|
|
||||||
Args:
|
Args:
|
||||||
text: The note text content.
|
text: The note text content.
|
||||||
collection: Collection to add the note to (default "documents").
|
tags: Tags to apply to the note.
|
||||||
Standard collections: "documents", "memory", "workspace".
|
|
||||||
tags: Additional tags to apply to the note.
|
|
||||||
title: Optional title (auto-derived from first line if omitted).
|
title: Optional title (auto-derived from first line if omitted).
|
||||||
"""
|
"""
|
||||||
all_tags = list(tags) if tags else []
|
result = engine.add_note(text=text, tags=tags or None, title=title)
|
||||||
all_tags.append(_collection_tag(collection))
|
|
||||||
|
|
||||||
result = engine.add_note(text=text, tags=all_tags, title=title)
|
|
||||||
return json.dumps(result, indent=2)
|
return json.dumps(result, indent=2)
|
||||||
|
|
||||||
|
|
||||||
@@ -203,7 +161,7 @@ async def kb_update_note(
|
|||||||
text: The new text content for the note.
|
text: The new text content for the note.
|
||||||
"""
|
"""
|
||||||
result = engine.update_note(document_id, text)
|
result = engine.update_note(document_id, text)
|
||||||
return json.dumps(_process_document(result), indent=2)
|
return json.dumps(result, indent=2)
|
||||||
|
|
||||||
|
|
||||||
@mcp.tool()
|
@mcp.tool()
|
||||||
@@ -222,14 +180,14 @@ async def kb_get(
|
|||||||
"""
|
"""
|
||||||
if document_id is not None:
|
if document_id is not None:
|
||||||
result = engine.get_document(document_id)
|
result = engine.get_document(document_id)
|
||||||
return json.dumps(_process_document(result), indent=2)
|
return json.dumps(result, indent=2)
|
||||||
elif source_path is not None:
|
elif source_path is not None:
|
||||||
docs = engine.list_documents()
|
docs = engine.list_documents()
|
||||||
matches = [d for d in docs if d.get("source_path") == source_path]
|
matches = [d for d in docs if d.get("source_path") == source_path]
|
||||||
if not matches:
|
if not matches:
|
||||||
return json.dumps({"error": "No document found with that source_path"})
|
return json.dumps({"error": "No document found with that source_path"})
|
||||||
doc = engine.get_document(matches[0]["id"])
|
doc = engine.get_document(matches[0]["id"])
|
||||||
return json.dumps(_process_document(doc), indent=2)
|
return json.dumps(doc, indent=2)
|
||||||
else:
|
else:
|
||||||
return json.dumps({"error": "Provide either document_id or source_path"})
|
return json.dumps({"error": "Provide either document_id or source_path"})
|
||||||
|
|
||||||
@@ -262,12 +220,27 @@ async def kb_jobs(
|
|||||||
return json.dumps(result, indent=2)
|
return json.dumps(result, indent=2)
|
||||||
|
|
||||||
|
|
||||||
|
@mcp.tool()
|
||||||
|
async def kb_delete(
|
||||||
|
document_id: int,
|
||||||
|
) -> str:
|
||||||
|
"""Permanently delete a document from the knowledge base.
|
||||||
|
|
||||||
|
Removes the document and all associated data (chunks, embeddings, tags,
|
||||||
|
stored files). This action cannot be undone.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
document_id: The ID of the document to delete.
|
||||||
|
"""
|
||||||
|
result = engine.delete_document(document_id)
|
||||||
|
return json.dumps(result, indent=2)
|
||||||
|
|
||||||
|
|
||||||
@mcp.tool()
|
@mcp.tool()
|
||||||
async def kb_upload_start(
|
async def kb_upload_start(
|
||||||
filename: str,
|
filename: str,
|
||||||
total_size: int,
|
total_size: int,
|
||||||
tags: list[str] | None = None,
|
tags: list[str] | None = None,
|
||||||
collection: str | None = None,
|
|
||||||
) -> str:
|
) -> str:
|
||||||
"""Start a chunked file upload to the knowledge base.
|
"""Start a chunked file upload to the knowledge base.
|
||||||
|
|
||||||
@@ -277,7 +250,7 @@ async def kb_upload_start(
|
|||||||
3. Call kb_upload_finish to submit the file for ingestion
|
3. Call kb_upload_finish to submit the file for ingestion
|
||||||
|
|
||||||
Example for a 3MB file:
|
Example for a 3MB file:
|
||||||
upload = kb_upload_start(filename="report.pdf", total_size=3145728, collection="documents")
|
upload = kb_upload_start(filename="report.pdf", total_size=3145728, tags=["project:x"])
|
||||||
kb_upload_chunk(upload_id=upload["upload_id"], data="<base64 chunk 0>", chunk_index=0)
|
kb_upload_chunk(upload_id=upload["upload_id"], data="<base64 chunk 0>", chunk_index=0)
|
||||||
kb_upload_chunk(upload_id=upload["upload_id"], data="<base64 chunk 1>", chunk_index=1)
|
kb_upload_chunk(upload_id=upload["upload_id"], data="<base64 chunk 1>", chunk_index=1)
|
||||||
kb_upload_chunk(upload_id=upload["upload_id"], data="<base64 chunk 2>", chunk_index=2)
|
kb_upload_chunk(upload_id=upload["upload_id"], data="<base64 chunk 2>", chunk_index=2)
|
||||||
@@ -286,13 +259,9 @@ async def kb_upload_start(
|
|||||||
Args:
|
Args:
|
||||||
filename: Original filename (used for type detection).
|
filename: Original filename (used for type detection).
|
||||||
total_size: Total file size in bytes.
|
total_size: Total file size in bytes.
|
||||||
tags: Additional tags to apply.
|
tags: Tags to apply to the uploaded document.
|
||||||
collection: Collection name (default "documents").
|
|
||||||
"""
|
"""
|
||||||
all_tags = list(tags) if tags else []
|
upload_id = uploads.start_upload(filename, total_size, tags or [])
|
||||||
all_tags.append(_collection_tag(collection))
|
|
||||||
|
|
||||||
upload_id = uploads.start_upload(filename, total_size, all_tags)
|
|
||||||
return json.dumps({"upload_id": upload_id})
|
return json.dumps({"upload_id": upload_id})
|
||||||
|
|
||||||
|
|
||||||
@@ -338,6 +307,112 @@ async def kb_upload_finish(
|
|||||||
return json.dumps({"error": str(e)})
|
return json.dumps({"error": str(e)})
|
||||||
|
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# Bulk operation tools
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
|
||||||
|
@mcp.tool()
|
||||||
|
async def kb_bulk_delete(
|
||||||
|
document_ids: list[int] | None = None,
|
||||||
|
tags: list[str] | None = None,
|
||||||
|
doc_type: str | None = None,
|
||||||
|
from_id: int | None = None,
|
||||||
|
to_id: int | None = None,
|
||||||
|
force: bool = False,
|
||||||
|
) -> str:
|
||||||
|
"""Permanently delete multiple documents matching a filter.
|
||||||
|
|
||||||
|
Removes matched documents and all associated data (chunks, embeddings, tags,
|
||||||
|
stored files). This action cannot be undone.
|
||||||
|
|
||||||
|
Selection filters combine with AND logic — at least one is required.
|
||||||
|
|
||||||
|
A safety threshold applies: if the operation would affect more than 70% of
|
||||||
|
all documents, it is rejected unless force=true.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
document_ids: Delete documents with these specific IDs.
|
||||||
|
tags: Delete documents that have ALL of these tags (selection filter).
|
||||||
|
doc_type: Delete documents of this type (e.g. "note", "pdf").
|
||||||
|
from_id: Delete documents with id >= this value.
|
||||||
|
to_id: Delete documents with id <= this value.
|
||||||
|
force: Override the safety threshold if it would block the operation.
|
||||||
|
"""
|
||||||
|
result = engine.bulk_delete(
|
||||||
|
document_ids=document_ids, tags=tags, doc_type=doc_type,
|
||||||
|
from_id=from_id, to_id=to_id, force=force,
|
||||||
|
)
|
||||||
|
return json.dumps(result, indent=2)
|
||||||
|
|
||||||
|
|
||||||
|
@mcp.tool()
|
||||||
|
async def kb_bulk_tags(
|
||||||
|
document_ids: list[int] | None = None,
|
||||||
|
tags: list[str] | None = None,
|
||||||
|
doc_type: str | None = None,
|
||||||
|
from_id: int | None = None,
|
||||||
|
to_id: int | None = None,
|
||||||
|
add: list[str] | None = None,
|
||||||
|
remove: list[str] | None = None,
|
||||||
|
force: bool = False,
|
||||||
|
) -> str:
|
||||||
|
"""Add and/or remove tags on multiple documents matching a filter.
|
||||||
|
|
||||||
|
Selection filters combine with AND logic — at least one is required.
|
||||||
|
Note: the 'tags' parameter is a SELECTION FILTER (which documents to target),
|
||||||
|
while 'add' and 'remove' specify the TAG CHANGES to apply to those documents.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
document_ids: Target documents with these specific IDs.
|
||||||
|
tags: Target documents that have ALL of these tags (selection filter).
|
||||||
|
doc_type: Target documents of this type.
|
||||||
|
from_id: Target documents with id >= this value.
|
||||||
|
to_id: Target documents with id <= this value.
|
||||||
|
add: Tags to add to matched documents.
|
||||||
|
remove: Tags to remove from matched documents.
|
||||||
|
force: Override the safety threshold if it would block the operation.
|
||||||
|
"""
|
||||||
|
result = engine.bulk_tags(
|
||||||
|
document_ids=document_ids, tags=tags, doc_type=doc_type,
|
||||||
|
from_id=from_id, to_id=to_id, add=add, remove=remove, force=force,
|
||||||
|
)
|
||||||
|
return json.dumps(result, indent=2)
|
||||||
|
|
||||||
|
|
||||||
|
@mcp.tool()
|
||||||
|
async def kb_bulk_set_tags(
|
||||||
|
document_ids: list[int] | None = None,
|
||||||
|
tags: list[str] | None = None,
|
||||||
|
doc_type: str | None = None,
|
||||||
|
from_id: int | None = None,
|
||||||
|
to_id: int | None = None,
|
||||||
|
new_tags: list[str] | None = None,
|
||||||
|
force: bool = False,
|
||||||
|
) -> str:
|
||||||
|
"""Replace all tags on multiple documents with a new set.
|
||||||
|
|
||||||
|
Removes ALL existing tags from matched documents, then applies the new tag set.
|
||||||
|
Selection filters combine with AND logic — at least one is required.
|
||||||
|
Note: the 'tags' parameter is a SELECTION FILTER (which documents to target),
|
||||||
|
while 'new_tags' is the REPLACEMENT tag set to apply.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
document_ids: Target documents with these specific IDs.
|
||||||
|
tags: Target documents that have ALL of these tags (selection filter).
|
||||||
|
doc_type: Target documents of this type.
|
||||||
|
from_id: Target documents with id >= this value.
|
||||||
|
to_id: Target documents with id <= this value.
|
||||||
|
new_tags: The replacement tag set to apply to all matched documents.
|
||||||
|
force: Override the safety threshold if it would block the operation.
|
||||||
|
"""
|
||||||
|
result = engine.bulk_set_tags(
|
||||||
|
document_ids=document_ids, tags=tags, doc_type=doc_type,
|
||||||
|
from_id=from_id, to_id=to_id, new_tags=new_tags, force=force,
|
||||||
|
)
|
||||||
|
return json.dumps(result, indent=2)
|
||||||
|
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
# ---------------------------------------------------------------------------
|
||||||
# Auth middleware
|
# Auth middleware
|
||||||
# ---------------------------------------------------------------------------
|
# ---------------------------------------------------------------------------
|
||||||
|
|||||||
@@ -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
|
||||||
|
```
|
||||||
@@ -0,0 +1,194 @@
|
|||||||
|
## Context
|
||||||
|
|
||||||
|
The engine API (`engine/kb/routes/`) provides single-document operations for delete (`DELETE /api/v1/documents/{id}`) and tag management (`PUT /api/v1/documents/{id}/tags`). The MCP server (`mcp/server.py`) wraps these and adds a "collection" abstraction via `collection:`-prefixed tags — ~70 lines of helpers and translation logic that only the MCP layer understands.
|
||||||
|
|
||||||
|
The database is SQLite with WAL mode, FTS5 for full-text search, and sqlite-vec for embeddings. Foreign keys with `ON DELETE CASCADE` handle chunk cleanup when documents are deleted. Stored files on disk must be cleaned up separately.
|
||||||
|
|
||||||
|
## Goals / Non-Goals
|
||||||
|
|
||||||
|
**Goals:**
|
||||||
|
- Bulk delete, bulk tag add/remove, and bulk set-tags (replace) via engine API, MCP tools, and CLI
|
||||||
|
- Filter-based selection: by tag, doc_type, ID list, and ID range
|
||||||
|
- Safety threshold to prevent accidental mass operations
|
||||||
|
- Audit trail via jobs table
|
||||||
|
- Remove collection abstraction from MCP server
|
||||||
|
|
||||||
|
**Non-Goals:**
|
||||||
|
- Async/queued bulk operations (SQLite handles thousands of rows synchronously in <1s)
|
||||||
|
- Bulk document retrieval or bulk note creation
|
||||||
|
- Undo/recycle bin for bulk deletes
|
||||||
|
- Adding collection concept to engine or CLI (collections are being removed, not moved)
|
||||||
|
|
||||||
|
## Decisions
|
||||||
|
|
||||||
|
### 1. Common selection filter for all bulk endpoints
|
||||||
|
|
||||||
|
All three bulk endpoints accept the same selection body:
|
||||||
|
|
||||||
|
```json
|
||||||
|
{
|
||||||
|
"document_ids": [1, 5, 12],
|
||||||
|
"tags": ["agent:mybot", "draft"],
|
||||||
|
"doc_type": "note",
|
||||||
|
"from_id": 10,
|
||||||
|
"to_id": 50
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|
||||||
|
Filters combine with AND logic. At least one filter is required — the engine rejects requests with no selection criteria (400).
|
||||||
|
|
||||||
|
**Selection SQL generation**: A shared helper in `database.py` builds the WHERE clause from the filter. The `tags` filter uses the same JOIN pattern as `list_documents` (all specified tags must match). The `document_ids` filter uses `IN (?)`. The `from_id`/`to_id` filter uses `id >= ? AND id <= ?`.
|
||||||
|
|
||||||
|
**Alternative considered**: Separate endpoints per filter type. Rejected — combinable filters are more powerful and the SQL generation is straightforward.
|
||||||
|
|
||||||
|
### 2. Safety threshold with configurable percentage
|
||||||
|
|
||||||
|
Before executing, the engine counts matched documents and total documents. If `matched / total > threshold`, the request is rejected:
|
||||||
|
|
||||||
|
```
|
||||||
|
HTTP 409 Conflict
|
||||||
|
{
|
||||||
|
"error": "safety_threshold_exceeded",
|
||||||
|
"message": "Operation would affect 750 of 1000 documents (75.0%). Exceeds safety threshold of 70%. Use force: true to proceed.",
|
||||||
|
"matched": 750,
|
||||||
|
"total": 1000,
|
||||||
|
"percent": 75.0,
|
||||||
|
"threshold": 70
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|
||||||
|
- Default threshold: 70% (env var `KB_BULK_SAFETY_PERCENT`, integer 0-100)
|
||||||
|
- Override per-request: `"force": true` in the request body
|
||||||
|
- Threshold of 0 effectively disables the safety check
|
||||||
|
- CLI maps this to `--force` / `-f` flag
|
||||||
|
|
||||||
|
The check is a SELECT COUNT before the operation — minimal overhead.
|
||||||
|
|
||||||
|
**Alternative considered**: Dry-run mode (preview what would be affected, then confirm). Rejected — adds a two-step flow that doesn't help LLM callers (they'd just always confirm) and the safety threshold covers the dangerous case.
|
||||||
|
|
||||||
|
### 3. Synchronous execution with audit logging
|
||||||
|
|
||||||
|
Bulk operations execute synchronously and return a summary response:
|
||||||
|
|
||||||
|
```json
|
||||||
|
{
|
||||||
|
"job_id": 42,
|
||||||
|
"status": "done",
|
||||||
|
"matched": 750,
|
||||||
|
"succeeded": 748,
|
||||||
|
"failed": 2,
|
||||||
|
"errors": [
|
||||||
|
{"document_id": 42, "error": "file locked"},
|
||||||
|
{"document_id": 99, "error": "not found"}
|
||||||
|
]
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|
||||||
|
A job record is created in the `jobs` table with a new `bulk_delete` / `bulk_tags` / `bulk_set_tags` status type. This requires extending the jobs table:
|
||||||
|
|
||||||
|
- Add `job_type` column: `"ingest"` (default, for existing jobs) or `"bulk_delete"` / `"bulk_tags"` / `"bulk_set_tags"`
|
||||||
|
- The job's `filename` field stores a JSON summary of the selection filter for auditability
|
||||||
|
- `document_id` field stores the count of affected documents
|
||||||
|
- `error` field stores JSON array of individual errors if any
|
||||||
|
|
||||||
|
**Alternative considered**: Full async with job polling. Rejected — SQLite bulk operations are fast enough synchronously and async would require extra polling calls (defeating the purpose of reducing token usage).
|
||||||
|
|
||||||
|
### 4. Bulk delete implementation
|
||||||
|
|
||||||
|
For each matched document:
|
||||||
|
1. Collect chunk IDs
|
||||||
|
2. Delete embeddings from `chunks_vec`
|
||||||
|
3. Delete the document row (cascades to chunks, document_tags)
|
||||||
|
4. Delete stored file from disk
|
||||||
|
|
||||||
|
This follows the same logic as the existing `delete_document` endpoint but batched in a single transaction (except file deletion, which happens after commit). If a file deletion fails, the document is still counted as succeeded (the DB record is gone) but a warning is logged.
|
||||||
|
|
||||||
|
The operation processes documents within a single SQLite transaction for atomicity of the DB changes. File deletions happen post-commit and are best-effort.
|
||||||
|
|
||||||
|
### 5. Bulk tags implementation
|
||||||
|
|
||||||
|
Two distinct operations:
|
||||||
|
|
||||||
|
**`POST /api/v1/bulk/tags`** — Add and/or remove tags:
|
||||||
|
```json
|
||||||
|
{
|
||||||
|
"add": ["reviewed", "approved"],
|
||||||
|
"remove": ["draft"],
|
||||||
|
...selection filters...
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|
||||||
|
**`POST /api/v1/bulk/set-tags`** — Replace all tags:
|
||||||
|
```json
|
||||||
|
{
|
||||||
|
"tags": ["final", "approved"],
|
||||||
|
...selection filters...
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|
||||||
|
The `set-tags` operation removes all existing tags from matched documents, then applies the new set. This is useful for cleaning up tag clutter or migrating tagging schemes.
|
||||||
|
|
||||||
|
Both update `updated_at` on affected documents.
|
||||||
|
|
||||||
|
### 6. Remove collection abstraction from MCP
|
||||||
|
|
||||||
|
Remove from `mcp/server.py`:
|
||||||
|
- Constants: `COLLECTION_TAG_PREFIX`, `DEFAULT_COLLECTION`
|
||||||
|
- Functions: `_collection_tag`, `_strip_collection_tags`, `_process_document`, `_process_search_results`, `_ensure_exclusive_collection`
|
||||||
|
- Tool: `kb_set_collection` (entire tool removed)
|
||||||
|
- Parameters: `collection` from `kb_search`, `kb_addnote`, `kb_upload_start`
|
||||||
|
|
||||||
|
The `_process_document` and `_process_search_results` calls in remaining tools are removed — documents are returned as-is from the engine, with all tags visible.
|
||||||
|
|
||||||
|
Users/agents that need namespace isolation use a tag convention (e.g. `agent:claude-code`) communicated via system prompt or tool instructions.
|
||||||
|
|
||||||
|
### 7. Engine bulk route module
|
||||||
|
|
||||||
|
New file: `engine/kb/routes/bulk.py`
|
||||||
|
|
||||||
|
Three endpoints sharing common infrastructure:
|
||||||
|
- `_resolve_selection(conn, filters)` → list of document IDs + count
|
||||||
|
- `_check_safety_threshold(matched, total, force)` → raises HTTPException if exceeded
|
||||||
|
- `_log_bulk_job(conn, job_type, filters, matched, succeeded, failed, errors)` → job_id
|
||||||
|
|
||||||
|
### 8. MCP bulk tools
|
||||||
|
|
||||||
|
Three new tools in `mcp/server.py`, thin wrappers calling new `engine.py` methods:
|
||||||
|
|
||||||
|
- `kb_bulk_delete(document_ids?, tags?, doc_type?, from_id?, to_id?, force?)` → str (JSON)
|
||||||
|
- `kb_bulk_tags(document_ids?, tags?, doc_type?, from_id?, to_id?, add?, remove?, force?)` → str (JSON)
|
||||||
|
- `kb_bulk_set_tags(document_ids?, tags?, doc_type?, from_id?, to_id?, new_tags?, force?)` → str (JSON)
|
||||||
|
|
||||||
|
Note: The `tags` parameter on bulk tools serves as a **selection filter** (which documents to target), while `add`/`remove` (on bulk_tags) and `new_tags` (on bulk_set_tags) are the **operation** (what to do to the tags). Tool descriptions must make this distinction clear.
|
||||||
|
|
||||||
|
### 9. CLI bulk commands
|
||||||
|
|
||||||
|
Three new commands under `client/cmd/`:
|
||||||
|
|
||||||
|
```
|
||||||
|
kb bulk-remove --tags "draft,old" --type note --force --yes
|
||||||
|
kb bulk-tag --tags "agent:mybot" --add "reviewed" --remove "pending" --yes
|
||||||
|
kb bulk-set-tags --ids "1,5,12" --tags "clean,final" --yes
|
||||||
|
```
|
||||||
|
|
||||||
|
Filter flags (shared): `--tags`, `--type`, `--ids` (comma-separated), `--from-id`, `--to-id`, `--force`
|
||||||
|
Confirmation: `--yes` / `-y` to skip interactive prompt.
|
||||||
|
|
||||||
|
Without `--yes`, the CLI first shows the match count and asks for confirmation:
|
||||||
|
|
||||||
|
```
|
||||||
|
This will delete 47 documents matching: tags=[draft,old] type=note
|
||||||
|
Proceed? [y/N]
|
||||||
|
```
|
||||||
|
|
||||||
|
### 10. Engine config for safety threshold
|
||||||
|
|
||||||
|
New env var: `KB_BULK_SAFETY_PERCENT` (integer, default 70). Added to `engine/kb/config.py`.
|
||||||
|
|
||||||
|
## Risks / Trade-offs
|
||||||
|
|
||||||
|
- **[Bulk delete is irreversible]** → Safety threshold mitigates accidental mass deletion. CLI requires interactive confirmation. No undo mechanism — this is deliberate to keep the system simple.
|
||||||
|
- **[Naming collision: `tags` as filter vs operation]** → The `tags` parameter in bulk_tags selects documents, while `add`/`remove` specifies the tag changes. Clear naming and tool descriptions mitigate confusion. Engine request model uses the same field name as the existing list/search filter.
|
||||||
|
- **[SQLite lock during large bulk ops]** → A single transaction deleting 5000 documents will hold a write lock. With WAL mode, readers are not blocked. The lock duration should be under a few seconds for typical workloads.
|
||||||
|
- **[Breaking change: collection removal]** → Any MCP client relying on `collection` parameters will break. Since collections were only recently added and are not widely deployed, this is acceptable. Existing `collection:*` tags in the database remain as regular tags — they still work as filters, just without special treatment.
|
||||||
|
- **[Jobs table overload]** → Bulk operations add a new job type to a table designed for ingestion jobs. The schema change is minimal (one new column) and the audit trail value outweighs the mixing of concerns.
|
||||||
@@ -0,0 +1,91 @@
|
|||||||
|
## Why
|
||||||
|
|
||||||
|
Bulk operations on documents (delete, tag, retag) currently require one API/MCP call per document. When an LLM manages hundreds or thousands of documents, this means hundreds of tool calls — burning tokens, adding latency, and creating fragile multi-step flows that can fail partway through.
|
||||||
|
|
||||||
|
Additionally, the "collection" abstraction in the MCP server adds complexity without real benefit. Collections are implemented as `collection:`-prefixed tags, but this convention is only enforced in the MCP layer — the CLI and engine don't know about it. This creates inconsistency and extra code. Tags alone, with a naming convention communicated via system prompt or configuration, achieve the same namespace isolation more simply and uniformly.
|
||||||
|
|
||||||
|
## What Changes
|
||||||
|
|
||||||
|
### 1. Remove collections from MCP server
|
||||||
|
|
||||||
|
Strip all collection logic from `mcp/server.py`:
|
||||||
|
- Remove `COLLECTION_TAG_PREFIX`, `DEFAULT_COLLECTION`, and all collection helper functions
|
||||||
|
- Remove `collection` parameter from `kb_search`, `kb_addnote`, `kb_upload_start`
|
||||||
|
- Remove `kb_set_collection` tool entirely
|
||||||
|
- Remove `_process_document` / `_process_search_results` collection-tag stripping
|
||||||
|
- Update MCP server instructions to explain tag-based namespace convention
|
||||||
|
|
||||||
|
### 2. Add bulk engine endpoints
|
||||||
|
|
||||||
|
Three new endpoints in the engine API:
|
||||||
|
|
||||||
|
- **POST /api/v1/bulk/delete** — Delete multiple documents matching a filter
|
||||||
|
- **POST /api/v1/bulk/tags** — Add/remove tags on multiple documents matching a filter
|
||||||
|
- **POST /api/v1/bulk/set-tags** — Replace all tags on multiple documents matching a filter
|
||||||
|
|
||||||
|
All accept a common **selection filter** (combinable with AND logic):
|
||||||
|
- `document_ids` — explicit list of IDs
|
||||||
|
- `tags` — documents matching ALL specified tags
|
||||||
|
- `doc_type` — documents of this type
|
||||||
|
- `from_id` / `to_id` — ID range (inclusive)
|
||||||
|
|
||||||
|
At least one selection criterion is required.
|
||||||
|
|
||||||
|
**Safety threshold**: If the operation would affect more than N% of all documents (default 70%, configurable via `KB_BULK_SAFETY_PERCENT` env var), the request is rejected with a 409 response showing what would be affected. The caller must re-send with `force: true` to proceed.
|
||||||
|
|
||||||
|
**Response model**: Synchronous execution with summary response. The operation is logged to the jobs table for audit trail:
|
||||||
|
|
||||||
|
```json
|
||||||
|
{
|
||||||
|
"job_id": 42,
|
||||||
|
"status": "done",
|
||||||
|
"matched": 750,
|
||||||
|
"succeeded": 748,
|
||||||
|
"failed": 2,
|
||||||
|
"errors": [
|
||||||
|
{"document_id": 42, "error": "file locked"},
|
||||||
|
{"document_id": 99, "error": "not found"}
|
||||||
|
]
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|
||||||
|
### 3. Add bulk MCP tools
|
||||||
|
|
||||||
|
Expose the bulk engine endpoints as MCP tools:
|
||||||
|
- `kb_bulk_delete` — bulk delete with filter selection
|
||||||
|
- `kb_bulk_tags` — bulk add/remove tags with filter selection
|
||||||
|
- `kb_bulk_set_tags` — bulk replace tags with filter selection
|
||||||
|
|
||||||
|
These are thin wrappers around the engine bulk endpoints — no collection translation, no special logic.
|
||||||
|
|
||||||
|
### 4. Add bulk CLI commands
|
||||||
|
|
||||||
|
- `kb bulk-remove` — bulk delete with `--tags`, `--type`, `--ids`, `--from-id`, `--to-id`, `--force` flags
|
||||||
|
- `kb bulk-tag` — bulk tag/untag with `--add`, `--remove`, and the same filter flags
|
||||||
|
- `kb bulk-set-tags` — bulk replace tags with `--tags` (new tags) and the same filter flags
|
||||||
|
|
||||||
|
All show a confirmation prompt with match count before executing (unless `--yes`).
|
||||||
|
|
||||||
|
## Capabilities
|
||||||
|
|
||||||
|
### New Capabilities
|
||||||
|
|
||||||
|
- `bulk-operations`: Engine endpoints, MCP tools, and CLI commands for bulk delete, tag, and set-tags operations with filter-based selection and safety threshold.
|
||||||
|
|
||||||
|
### Modified Capabilities
|
||||||
|
|
||||||
|
- `mcp-document-management`: Remove `kb_set_collection` tool. Remove `collection` parameter from all tools.
|
||||||
|
|
||||||
|
### Removed Capabilities
|
||||||
|
|
||||||
|
- `mcp-collections`: The collection abstraction (collection helpers, collection parameters, collection tag stripping) is removed from the MCP server entirely.
|
||||||
|
|
||||||
|
## Impact
|
||||||
|
|
||||||
|
- **Engine API** (`engine/kb/routes/`): New `bulk.py` route module with 3 endpoints. New `bulk` job type in jobs table.
|
||||||
|
- **Engine database** (`engine/kb/database.py`): Helper functions for bulk selection queries and bulk delete/tag operations.
|
||||||
|
- **MCP server** (`mcp/server.py`): Remove ~70 lines of collection logic. Add 3 bulk tool definitions. Remove `collection` param from `kb_search`, `kb_addnote`, `kb_upload_start`. Remove `kb_set_collection`.
|
||||||
|
- **MCP engine client** (`mcp/engine.py`): Add bulk operation methods. Remove no longer needed code.
|
||||||
|
- **CLI** (`client/cmd/`): New `bulk_remove.go`, `bulk_tag.go`, `bulk_set_tags.go` command files.
|
||||||
|
- **CLI API client** (`client/internal/api/`): Add `Post` with JSON body support if not present.
|
||||||
|
- **Breaking changes**: `kb_set_collection` MCP tool removed. `collection` parameter removed from `kb_search`, `kb_addnote`, `kb_upload_start` MCP tools. Any MCP clients using collections will need to switch to tags.
|
||||||
+230
@@ -0,0 +1,230 @@
|
|||||||
|
## ADDED Requirements
|
||||||
|
|
||||||
|
### Requirement: Common selection filter
|
||||||
|
|
||||||
|
All bulk engine endpoints SHALL accept a JSON body with the following optional selection fields, combined with AND logic:
|
||||||
|
|
||||||
|
- `document_ids` (list of int) — match documents with these specific IDs
|
||||||
|
- `tags` (list of str) — match documents that have ALL specified tags
|
||||||
|
- `doc_type` (str) — match documents with this document type
|
||||||
|
- `from_id` (int) — match documents with id >= this value
|
||||||
|
- `to_id` (int) — match documents with id <= this value
|
||||||
|
|
||||||
|
At least one selection field MUST be present. If no selection fields are provided, the endpoint SHALL return 400 Bad Request.
|
||||||
|
|
||||||
|
#### Scenario: Filter by tags and doc_type
|
||||||
|
|
||||||
|
- **WHEN** a bulk endpoint receives `{"tags": ["draft"], "doc_type": "note"}`
|
||||||
|
- **THEN** it SHALL match only documents that have the tag "draft" AND have doc_type "note"
|
||||||
|
|
||||||
|
#### Scenario: Filter by ID range
|
||||||
|
|
||||||
|
- **WHEN** a bulk endpoint receives `{"from_id": 10, "to_id": 50}`
|
||||||
|
- **THEN** it SHALL match documents with id >= 10 AND id <= 50
|
||||||
|
|
||||||
|
#### Scenario: Filter by explicit IDs
|
||||||
|
|
||||||
|
- **WHEN** a bulk endpoint receives `{"document_ids": [1, 5, 12]}`
|
||||||
|
- **THEN** it SHALL match only documents with those specific IDs
|
||||||
|
|
||||||
|
#### Scenario: Combined filters
|
||||||
|
|
||||||
|
- **WHEN** a bulk endpoint receives `{"tags": ["agent:mybot"], "doc_type": "note", "from_id": 100}`
|
||||||
|
- **THEN** it SHALL match documents satisfying ALL three criteria
|
||||||
|
|
||||||
|
#### Scenario: No selection fields provided
|
||||||
|
|
||||||
|
- **WHEN** a bulk endpoint receives `{}` or `{"force": true}` with no selection fields
|
||||||
|
- **THEN** it SHALL return 400 Bad Request
|
||||||
|
|
||||||
|
### Requirement: Safety threshold
|
||||||
|
|
||||||
|
All bulk endpoints SHALL enforce a safety threshold. Before executing, the engine SHALL count the matched documents and the total documents in the database. If `matched / total * 100` exceeds the configured threshold, the request SHALL be rejected with 409 Conflict.
|
||||||
|
|
||||||
|
The response SHALL include: `error` ("safety_threshold_exceeded"), `message` (human-readable), `matched` (int), `total` (int), `percent` (float), and `threshold` (int).
|
||||||
|
|
||||||
|
The threshold SHALL default to 70 and be configurable via the `KB_BULK_SAFETY_PERCENT` environment variable (integer 0-100). A value of 0 disables the check.
|
||||||
|
|
||||||
|
The caller MAY override the threshold by including `"force": true` in the request body.
|
||||||
|
|
||||||
|
#### Scenario: Threshold exceeded
|
||||||
|
|
||||||
|
- **GIVEN** 1000 total documents and `KB_BULK_SAFETY_PERCENT` is 70
|
||||||
|
- **WHEN** a bulk endpoint matches 750 documents (75%) without `force: true`
|
||||||
|
- **THEN** it SHALL return 409 with `matched: 750`, `total: 1000`, `percent: 75.0`, `threshold: 70`
|
||||||
|
|
||||||
|
#### Scenario: Threshold not exceeded
|
||||||
|
|
||||||
|
- **GIVEN** 1000 total documents and `KB_BULK_SAFETY_PERCENT` is 70
|
||||||
|
- **WHEN** a bulk endpoint matches 500 documents (50%) without `force: true`
|
||||||
|
- **THEN** the operation SHALL proceed normally
|
||||||
|
|
||||||
|
#### Scenario: Force override
|
||||||
|
|
||||||
|
- **GIVEN** 1000 total documents and a match of 900 (90%)
|
||||||
|
- **WHEN** the request includes `"force": true`
|
||||||
|
- **THEN** the operation SHALL proceed regardless of threshold
|
||||||
|
|
||||||
|
#### Scenario: Zero threshold
|
||||||
|
|
||||||
|
- **GIVEN** `KB_BULK_SAFETY_PERCENT` is 0
|
||||||
|
- **THEN** the safety check SHALL be effectively disabled for all operations
|
||||||
|
|
||||||
|
### Requirement: Synchronous response with audit log
|
||||||
|
|
||||||
|
All bulk endpoints SHALL execute synchronously and return a JSON response with:
|
||||||
|
|
||||||
|
- `job_id` (int) — ID of the audit log entry in the jobs table
|
||||||
|
- `status` (str) — "done" or "partial_failure"
|
||||||
|
- `matched` (int) — number of documents that matched the selection
|
||||||
|
- `succeeded` (int) — number of documents successfully processed
|
||||||
|
- `failed` (int) — number of documents that failed
|
||||||
|
- `errors` (list) — array of `{"document_id": int, "error": str}` for each failure (empty on full success)
|
||||||
|
|
||||||
|
A job record SHALL be created in the jobs table with `job_type` set to the operation type. The `filename` field SHALL store a JSON representation of the selection filter. The `error` field SHALL store a JSON array of individual errors if any occurred.
|
||||||
|
|
||||||
|
#### Scenario: Full success
|
||||||
|
|
||||||
|
- **WHEN** a bulk operation matches 50 documents and all succeed
|
||||||
|
- **THEN** the response SHALL have `status: "done"`, `matched: 50`, `succeeded: 50`, `failed: 0`, `errors: []`
|
||||||
|
|
||||||
|
#### Scenario: Partial failure
|
||||||
|
|
||||||
|
- **WHEN** a bulk operation matches 50 documents but 2 fail
|
||||||
|
- **THEN** the response SHALL have `status: "partial_failure"`, `matched: 50`, `succeeded: 48`, `failed: 2`, and `errors` listing the 2 failures
|
||||||
|
|
||||||
|
### Requirement: Bulk delete endpoint
|
||||||
|
|
||||||
|
The engine SHALL expose `POST /api/v1/bulk/delete` which permanently deletes all documents matching the selection filter. For each matched document, it SHALL delete embeddings from `chunks_vec`, delete the document row (cascading to chunks and document_tags), and delete any stored file from disk.
|
||||||
|
|
||||||
|
Database deletions SHALL be performed within a single transaction. File deletions SHALL occur after the transaction commits and SHALL be best-effort (failures logged but not counted as document failures).
|
||||||
|
|
||||||
|
#### Scenario: Bulk delete by tag
|
||||||
|
|
||||||
|
- **WHEN** `POST /api/v1/bulk/delete` receives `{"tags": ["old", "draft"]}`
|
||||||
|
- **THEN** all documents with both tags "old" and "draft" SHALL be deleted
|
||||||
|
- **AND** their chunks, embeddings, tag associations, and stored files SHALL be removed
|
||||||
|
|
||||||
|
#### Scenario: Bulk delete with no matches
|
||||||
|
|
||||||
|
- **WHEN** `POST /api/v1/bulk/delete` receives a filter that matches 0 documents
|
||||||
|
- **THEN** the response SHALL have `matched: 0`, `succeeded: 0`, `failed: 0`
|
||||||
|
|
||||||
|
### Requirement: Bulk tags endpoint
|
||||||
|
|
||||||
|
The engine SHALL expose `POST /api/v1/bulk/tags` which adds and/or removes tags on all documents matching the selection filter. The request body SHALL include the selection filter plus:
|
||||||
|
|
||||||
|
- `add` (list of str, optional) — tags to add
|
||||||
|
- `remove` (list of str, optional) — tags to remove
|
||||||
|
|
||||||
|
At least one of `add` or `remove` MUST be present. The endpoint SHALL return 400 if neither is provided.
|
||||||
|
|
||||||
|
The endpoint SHALL update `updated_at` on all affected documents.
|
||||||
|
|
||||||
|
#### Scenario: Add and remove tags in one call
|
||||||
|
|
||||||
|
- **WHEN** `POST /api/v1/bulk/tags` receives `{"tags": ["agent:mybot"], "add": ["reviewed"], "remove": ["pending"]}`
|
||||||
|
- **THEN** all documents tagged "agent:mybot" SHALL have "reviewed" added and "pending" removed
|
||||||
|
|
||||||
|
### Requirement: Bulk set-tags endpoint
|
||||||
|
|
||||||
|
The engine SHALL expose `POST /api/v1/bulk/set-tags` which replaces all tags on matched documents with a new set. The request body SHALL include the selection filter plus:
|
||||||
|
|
||||||
|
- `new_tags` (list of str) — the replacement tag set
|
||||||
|
|
||||||
|
The endpoint SHALL remove all existing tag associations from matched documents, then apply the new set. It SHALL update `updated_at` on all affected documents.
|
||||||
|
|
||||||
|
#### Scenario: Replace all tags
|
||||||
|
|
||||||
|
- **WHEN** `POST /api/v1/bulk/set-tags` receives `{"doc_type": "note", "new_tags": ["clean", "final"]}`
|
||||||
|
- **THEN** all notes SHALL have their existing tags removed and replaced with "clean" and "final"
|
||||||
|
|
||||||
|
### Requirement: Jobs table extension
|
||||||
|
|
||||||
|
The jobs table SHALL be extended with a `job_type` column (TEXT, default "ingest") to distinguish ingestion jobs from bulk operation audit entries. Valid values: "ingest", "bulk_delete", "bulk_tags", "bulk_set_tags".
|
||||||
|
|
||||||
|
Existing jobs SHALL default to `job_type = "ingest"`. The existing jobs list endpoint and CLI `kb jobs` command SHALL continue to work unchanged.
|
||||||
|
|
||||||
|
#### Scenario: Migration adds column
|
||||||
|
|
||||||
|
- **GIVEN** an existing database without the `job_type` column
|
||||||
|
- **WHEN** the engine starts
|
||||||
|
- **THEN** the column SHALL be added with default value "ingest"
|
||||||
|
|
||||||
|
### Requirement: Engine config for safety threshold
|
||||||
|
|
||||||
|
The engine `Config` class SHALL read `KB_BULK_SAFETY_PERCENT` from the environment as an integer (default 70, range 0-100). This value SHALL be used as the default safety threshold for all bulk endpoints.
|
||||||
|
|
||||||
|
### Requirement: MCP bulk delete tool
|
||||||
|
|
||||||
|
The MCP server SHALL expose a `kb_bulk_delete` tool with parameters: `document_ids` (optional list of int), `tags` (optional list of str), `doc_type` (optional str), `from_id` (optional int), `to_id` (optional int), `force` (optional bool).
|
||||||
|
|
||||||
|
The tool SHALL call `POST /api/v1/bulk/delete` on the engine via the engine client and return the JSON response.
|
||||||
|
|
||||||
|
The tool description SHALL clearly state that `tags` is a selection filter (which documents to delete), not tags to delete.
|
||||||
|
|
||||||
|
#### Scenario: MCP bulk delete by tag
|
||||||
|
|
||||||
|
- **WHEN** `kb_bulk_delete(tags=["old"])` is called
|
||||||
|
- **THEN** the engine client SHALL send `POST /api/v1/bulk/delete` with `{"tags": ["old"]}`
|
||||||
|
- **AND** the tool SHALL return the engine's JSON response
|
||||||
|
|
||||||
|
### Requirement: MCP bulk tags tool
|
||||||
|
|
||||||
|
The MCP server SHALL expose a `kb_bulk_tags` tool with parameters: `document_ids`, `tags`, `doc_type`, `from_id`, `to_id` (selection filters), plus `add` (optional list of str), `remove` (optional list of str), and `force` (optional bool).
|
||||||
|
|
||||||
|
The tool description SHALL clearly distinguish `tags` (selection filter) from `add`/`remove` (tag changes to apply).
|
||||||
|
|
||||||
|
#### Scenario: MCP bulk tag update
|
||||||
|
|
||||||
|
- **WHEN** `kb_bulk_tags(tags=["agent:mybot"], add=["reviewed"], remove=["draft"])` is called
|
||||||
|
- **THEN** the engine client SHALL send the appropriate `POST /api/v1/bulk/tags` request
|
||||||
|
|
||||||
|
### Requirement: MCP bulk set-tags tool
|
||||||
|
|
||||||
|
The MCP server SHALL expose a `kb_bulk_set_tags` tool with parameters: `document_ids`, `tags`, `doc_type`, `from_id`, `to_id` (selection filters), plus `new_tags` (list of str) and `force` (optional bool).
|
||||||
|
|
||||||
|
#### Scenario: MCP bulk set tags
|
||||||
|
|
||||||
|
- **WHEN** `kb_bulk_set_tags(doc_type="note", new_tags=["clean"])` is called
|
||||||
|
- **THEN** the engine client SHALL send `POST /api/v1/bulk/set-tags` with `{"doc_type": "note", "new_tags": ["clean"]}`
|
||||||
|
|
||||||
|
### Requirement: MCP engine client bulk methods
|
||||||
|
|
||||||
|
The MCP engine client (`mcp/engine.py`) SHALL provide three new methods:
|
||||||
|
|
||||||
|
- `bulk_delete(document_ids?, tags?, doc_type?, from_id?, to_id?, force?)` → dict
|
||||||
|
- `bulk_tags(document_ids?, tags?, doc_type?, from_id?, to_id?, add?, remove?, force?)` → dict
|
||||||
|
- `bulk_set_tags(document_ids?, tags?, doc_type?, from_id?, to_id?, new_tags?, force?)` → dict
|
||||||
|
|
||||||
|
Each SHALL send a POST request to the corresponding `/api/v1/bulk/*` endpoint with the parameters as a JSON body. Each SHALL raise on non-2xx status codes, consistent with existing methods.
|
||||||
|
|
||||||
|
### Requirement: CLI bulk-remove command
|
||||||
|
|
||||||
|
The CLI SHALL expose a `kb bulk-remove` command with flags: `--tags` (comma-separated), `--type`, `--ids` (comma-separated), `--from-id`, `--to-id`, `--force`/`-f`, `--yes`/`-y`.
|
||||||
|
|
||||||
|
Without `--yes`, the CLI SHALL first display the match count and ask for interactive confirmation before proceeding.
|
||||||
|
|
||||||
|
The command SHALL call `POST /api/v1/bulk/delete` with the constructed filter.
|
||||||
|
|
||||||
|
#### Scenario: CLI bulk remove with confirmation
|
||||||
|
|
||||||
|
- **WHEN** `kb bulk-remove --tags "draft,old" --type note` is run without `--yes`
|
||||||
|
- **THEN** the CLI SHALL display "This will delete N documents matching: tags=[draft,old] type=note" and prompt "Proceed? [y/N]"
|
||||||
|
|
||||||
|
#### Scenario: CLI bulk remove with --yes
|
||||||
|
|
||||||
|
- **WHEN** `kb bulk-remove --tags "draft" --yes` is run
|
||||||
|
- **THEN** the CLI SHALL proceed without prompting
|
||||||
|
|
||||||
|
### Requirement: CLI bulk-tag command
|
||||||
|
|
||||||
|
The CLI SHALL expose a `kb bulk-tag` command with the same filter flags as `bulk-remove`, plus `--add` and `--remove` (comma-separated tag lists).
|
||||||
|
|
||||||
|
The command SHALL call `POST /api/v1/bulk/tags` with the constructed filter and tag changes.
|
||||||
|
|
||||||
|
### Requirement: CLI bulk-set-tags command
|
||||||
|
|
||||||
|
The CLI SHALL expose a `kb bulk-set-tags` command with the filter flags, plus `--set` (comma-separated list of replacement tags).
|
||||||
|
|
||||||
|
The command SHALL call `POST /api/v1/bulk/set-tags` with the constructed filter and `new_tags`.
|
||||||
+55
@@ -0,0 +1,55 @@
|
|||||||
|
## REMOVED Requirements
|
||||||
|
|
||||||
|
### Requirement: Collection abstraction in MCP server
|
||||||
|
|
||||||
|
The MCP server SHALL NOT maintain any collection abstraction. The following SHALL be removed:
|
||||||
|
|
||||||
|
- Constants: `COLLECTION_TAG_PREFIX`, `DEFAULT_COLLECTION`
|
||||||
|
- Functions: `_collection_tag`, `_strip_collection_tags`, `_process_document`, `_process_search_results`, `_ensure_exclusive_collection`
|
||||||
|
- Tool: `kb_set_collection` (entire tool)
|
||||||
|
- Parameters: `collection` from `kb_search`, `kb_addnote`, `kb_upload_start`
|
||||||
|
|
||||||
|
Documents SHALL be returned as-is from the engine with all tags visible. No tag stripping or collection field injection SHALL occur.
|
||||||
|
|
||||||
|
#### Scenario: Search results show all tags
|
||||||
|
|
||||||
|
- **WHEN** `kb_search` is called and a result has tags `["agent:mybot", "collection:documents", "draft"]`
|
||||||
|
- **THEN** all three tags SHALL be returned as-is — no stripping of `collection:*` tags
|
||||||
|
|
||||||
|
#### Scenario: kb_set_collection no longer exists
|
||||||
|
|
||||||
|
- **WHEN** an MCP client attempts to call `kb_set_collection`
|
||||||
|
- **THEN** the tool SHALL not be found (removed)
|
||||||
|
|
||||||
|
## MODIFIED Requirements
|
||||||
|
|
||||||
|
### Requirement: kb_search without collection parameter
|
||||||
|
|
||||||
|
The `kb_search` MCP tool SHALL accept `tags` (optional list of str) for filtering but SHALL NOT accept a `collection` parameter. Callers that previously used `collection="memory"` SHALL instead use `tags=["collection:memory"]` or whatever tag convention they prefer.
|
||||||
|
|
||||||
|
#### Scenario: Filter by tag instead of collection
|
||||||
|
|
||||||
|
- **WHEN** `kb_search(query="test", tags=["agent:mybot"])` is called
|
||||||
|
- **THEN** results SHALL be filtered to documents tagged "agent:mybot"
|
||||||
|
- **AND** no collection field SHALL be present in the response
|
||||||
|
|
||||||
|
### Requirement: kb_addnote without collection parameter
|
||||||
|
|
||||||
|
The `kb_addnote` MCP tool SHALL accept `tags` (optional list of str) but SHALL NOT accept a `collection` parameter. The tool SHALL NOT automatically apply any default collection tag — only explicitly provided tags are applied.
|
||||||
|
|
||||||
|
#### Scenario: Add note with explicit tags
|
||||||
|
|
||||||
|
- **WHEN** `kb_addnote(text="hello", tags=["agent:mybot", "memory"])` is called
|
||||||
|
- **THEN** the note SHALL be created with exactly those two tags — no `collection:documents` tag added
|
||||||
|
|
||||||
|
### Requirement: kb_upload_start without collection parameter
|
||||||
|
|
||||||
|
The `kb_upload_start` MCP tool SHALL accept `tags` (optional list of str) but SHALL NOT accept a `collection` parameter. The tool SHALL NOT automatically apply any default collection tag.
|
||||||
|
|
||||||
|
### Requirement: kb_update_note without collection processing
|
||||||
|
|
||||||
|
The `kb_update_note` MCP tool SHALL return the document as-is from the engine without passing it through `_process_document`. All tags SHALL be visible in the response.
|
||||||
|
|
||||||
|
### Requirement: kb_get without collection processing
|
||||||
|
|
||||||
|
The `kb_get` MCP tool SHALL return documents as-is from the engine without passing through `_process_document`. All tags SHALL be visible in the response. No `collection` field SHALL be injected.
|
||||||
@@ -0,0 +1,45 @@
|
|||||||
|
## 1. Remove collections from MCP server
|
||||||
|
|
||||||
|
- [x] 1.1 Remove collection constants and helper functions from `mcp/server.py` (`COLLECTION_TAG_PREFIX`, `DEFAULT_COLLECTION`, `_collection_tag`, `_strip_collection_tags`, `_process_document`, `_process_search_results`, `_ensure_exclusive_collection`)
|
||||||
|
- [x] 1.2 Remove `collection` parameter from `kb_search`, `kb_addnote`, `kb_upload_start` tools
|
||||||
|
- [x] 1.3 Remove `kb_set_collection` tool entirely
|
||||||
|
- [x] 1.4 Remove `_process_document` / `_process_search_results` calls from `kb_get`, `kb_update_note`, `kb_search`
|
||||||
|
- [x] 1.5 Update MCP server instructions text to reflect tags-only approach
|
||||||
|
|
||||||
|
## 2. Engine bulk infrastructure
|
||||||
|
|
||||||
|
- [x] 2.1 Add `bulk_safety_percent` to `Config` class in `engine/kb/config.py` (env var `KB_BULK_SAFETY_PERCENT`, default 70)
|
||||||
|
- [x] 2.2 Add `job_type` column migration to `database.py` `init_schema` (TEXT, default "ingest")
|
||||||
|
- [x] 2.3 Add `resolve_bulk_selection(conn, document_ids, tags, doc_type, from_id, to_id)` helper to `database.py` — returns list of matching document IDs
|
||||||
|
- [x] 2.4 Add `create_bulk_job(conn, job_type, filters_json, matched, succeeded, failed, errors_json)` helper to `database.py`
|
||||||
|
|
||||||
|
## 3. Engine bulk endpoints
|
||||||
|
|
||||||
|
- [x] 3.1 Create `engine/kb/routes/bulk.py` with shared Pydantic request model (`BulkSelectionRequest` with selection fields + `force` bool)
|
||||||
|
- [x] 3.2 Add `_check_safety_threshold` helper that returns 409 if threshold exceeded
|
||||||
|
- [x] 3.3 Implement `POST /api/v1/bulk/delete` — resolve selection, check threshold, delete documents in transaction, clean up files, log job, return summary
|
||||||
|
- [x] 3.4 Implement `POST /api/v1/bulk/tags` — resolve selection, check threshold, add/remove tags on matched docs, log job, return summary
|
||||||
|
- [x] 3.5 Implement `POST /api/v1/bulk/set-tags` — resolve selection, check threshold, clear and replace tags on matched docs, log job, return summary
|
||||||
|
- [x] 3.6 Import bulk routes in engine app startup (add to `engine/kb/routes/__init__.py` or `main.py`)
|
||||||
|
|
||||||
|
## 4. MCP bulk tools
|
||||||
|
|
||||||
|
- [x] 4.1 Add `bulk_delete`, `bulk_tags`, `bulk_set_tags` methods to `mcp/engine.py`
|
||||||
|
- [x] 4.2 Add `kb_bulk_delete` tool to `mcp/server.py`
|
||||||
|
- [x] 4.3 Add `kb_bulk_tags` tool to `mcp/server.py`
|
||||||
|
- [x] 4.4 Add `kb_bulk_set_tags` tool to `mcp/server.py`
|
||||||
|
|
||||||
|
## 5. CLI bulk commands
|
||||||
|
|
||||||
|
- [x] 5.1 Create `client/cmd/bulk_remove.go` — `kb bulk-remove` with filter flags, confirmation prompt, JSON output support
|
||||||
|
- [x] 5.2 Create `client/cmd/bulk_tag.go` — `kb bulk-tag` with filter flags + `--add`/`--remove`, confirmation prompt
|
||||||
|
- [x] 5.3 Create `client/cmd/bulk_set_tags.go` — `kb bulk-set-tags` with filter flags + `--set`, confirmation prompt
|
||||||
|
|
||||||
|
## 6. Verification
|
||||||
|
|
||||||
|
- [x] 6.1 Test collection removal: verify `kb_search`, `kb_addnote`, `kb_get`, `kb_update_note`, `kb_upload_start` work without collection params
|
||||||
|
- [x] 6.2 Test bulk delete via engine API: filter by tags, by IDs, by range, safety threshold trigger and force override
|
||||||
|
- [x] 6.3 Test bulk tags and bulk set-tags via engine API
|
||||||
|
- [x] 6.4 Test MCP bulk tools against running engine
|
||||||
|
- [x] 6.5 Test CLI bulk commands against running engine
|
||||||
|
- [x] 6.6 Test audit trail: verify bulk jobs appear in `kb jobs` output
|
||||||
@@ -0,0 +1,2 @@
|
|||||||
|
schema: spec-driven
|
||||||
|
created: 2026-04-04
|
||||||
@@ -0,0 +1,2 @@
|
|||||||
|
schema: spec-driven
|
||||||
|
created: 2026-04-04
|
||||||
@@ -0,0 +1,39 @@
|
|||||||
|
## Context
|
||||||
|
|
||||||
|
The MCP server (`mcp/server.py`) exposes KB operations as tools for LLM clients. Collections are an abstraction over tags — internally stored with a `collection:` prefix. The server already has helpers for managing collection tags (`_collection_tag`, `_ensure_exclusive_collection`, `_strip_collection_tags`) and the engine client (`mcp/engine.py`) already has an `update_tags()` method.
|
||||||
|
|
||||||
|
Document deletion is supported by the engine API at `DELETE /api/v1/documents/{doc_id}` but has no corresponding engine client method or MCP tool.
|
||||||
|
|
||||||
|
## Goals / Non-Goals
|
||||||
|
|
||||||
|
**Goals:**
|
||||||
|
- Expose collection assignment for existing documents via MCP (`kb_set_collection`)
|
||||||
|
- Expose document deletion via MCP (`kb_delete`)
|
||||||
|
- Follow existing patterns in `server.py` and `engine.py`
|
||||||
|
|
||||||
|
**Non-Goals:**
|
||||||
|
- Bulk operations (multi-document collection assignment or deletion)
|
||||||
|
- Tag management beyond collections (direct tag add/remove via MCP)
|
||||||
|
- Undo/recycle bin for deleted documents
|
||||||
|
- Changes to the engine API layer — all endpoints already exist
|
||||||
|
|
||||||
|
## Decisions
|
||||||
|
|
||||||
|
### 1. Reuse `_ensure_exclusive_collection` for kb_set_collection
|
||||||
|
|
||||||
|
The server already has `_ensure_exclusive_collection(doc_id, collection)` which removes any existing `collection:*` tags and applies the new one. The `kb_set_collection` tool will use this directly when a collection is provided, and manually remove collection tags when clearing.
|
||||||
|
|
||||||
|
**Alternative considered**: Exposing raw tag add/remove to the LLM. Rejected because it leaks the `collection:` prefix implementation detail and the LLM could create inconsistent state (multiple collections on one document).
|
||||||
|
|
||||||
|
### 2. New `engine.delete_document()` method for kb_delete
|
||||||
|
|
||||||
|
Add a simple `delete_document(doc_id)` to `mcp/engine.py` that calls `DELETE /api/v1/documents/{doc_id}`. This follows the same pattern as all other engine client methods.
|
||||||
|
|
||||||
|
### 3. Return confirmation with document metadata on delete
|
||||||
|
|
||||||
|
`kb_delete` will return the response from the engine API which includes `{"status": "deleted", "document_id": ..., "title": ...}`. This gives the LLM confirmation of what was deleted without needing a separate get call.
|
||||||
|
|
||||||
|
## Risks / Trade-offs
|
||||||
|
|
||||||
|
- **[Accidental deletion]** → The LLM could delete the wrong document. Mitigation: the tool requires an explicit `document_id`, and the response includes the title so the LLM can verify. No bulk delete is exposed.
|
||||||
|
- **[Collection cleared unexpectedly]** → Passing `collection=None` to `kb_set_collection` removes collection assignment. Mitigation: the parameter description will make this behavior explicit.
|
||||||
@@ -0,0 +1,25 @@
|
|||||||
|
## Why
|
||||||
|
|
||||||
|
LLMs using the KB MCP server can create notes in collections and search by collection, but cannot assign existing documents to a collection or delete documents. This forces users to drop out to the HTTP API for routine document management. Both operations are fully supported at the database and HTTP API layers but aren't wired through to MCP tools.
|
||||||
|
|
||||||
|
## What Changes
|
||||||
|
|
||||||
|
- Add `kb_set_collection` MCP tool — assigns, changes, or removes the collection on an existing document by manipulating `collection:` prefixed tags via the existing `engine.update_tags()` method.
|
||||||
|
- Add `kb_delete` MCP tool — deletes a document by ID, calling the existing `DELETE /api/v1/documents/{doc_id}` endpoint via a new `engine.delete_document()` method.
|
||||||
|
|
||||||
|
## Capabilities
|
||||||
|
|
||||||
|
### New Capabilities
|
||||||
|
|
||||||
|
- `mcp-document-management`: MCP tools for modifying and deleting existing documents (kb_set_collection, kb_delete).
|
||||||
|
|
||||||
|
### Modified Capabilities
|
||||||
|
|
||||||
|
_(none — the engine API endpoints already exist; this change only adds MCP tool wrappers)_
|
||||||
|
|
||||||
|
## Impact
|
||||||
|
|
||||||
|
- **MCP server** (`mcp/server.py`): Two new tool registrations.
|
||||||
|
- **MCP engine client** (`mcp/engine.py`): One new method (`delete_document`). The `update_tags` method already exists and will be reused.
|
||||||
|
- **Engine API**: No changes — `DELETE /api/v1/documents/{doc_id}` and `PUT /api/v1/documents/{doc_id}/tags` already exist.
|
||||||
|
- **Breaking changes**: None. Additive only.
|
||||||
+61
@@ -0,0 +1,61 @@
|
|||||||
|
## ADDED Requirements
|
||||||
|
|
||||||
|
### Requirement: Set collection on existing document via MCP
|
||||||
|
|
||||||
|
The MCP server SHALL expose a `kb_set_collection` tool that assigns or changes the collection of an existing document. The tool SHALL accept a `document_id` (required) and `collection` (optional string). When `collection` is provided, the tool SHALL ensure the document belongs to exactly that collection by removing any existing `collection:*` tags and adding the new one. When `collection` is omitted or null, the tool SHALL remove all `collection:*` tags from the document, leaving it unassigned.
|
||||||
|
|
||||||
|
The tool SHALL return the updated document with the `collection` field and cleaned tags (collection tags stripped), consistent with other MCP tool responses.
|
||||||
|
|
||||||
|
#### Scenario: Assign untagged document to a collection
|
||||||
|
|
||||||
|
- **WHEN** `kb_set_collection` is called with `document_id=42` and `collection="workspace"`
|
||||||
|
- **THEN** the document SHALL have the tag `collection:workspace` added
|
||||||
|
- **AND** the response SHALL include `"collection": "workspace"`
|
||||||
|
|
||||||
|
#### Scenario: Change document from one collection to another
|
||||||
|
|
||||||
|
- **WHEN** `kb_set_collection` is called with `document_id=42` and `collection="memory"` on a document currently in collection "documents"
|
||||||
|
- **THEN** the tag `collection:documents` SHALL be removed and `collection:memory` SHALL be added
|
||||||
|
- **AND** the response SHALL include `"collection": "memory"`
|
||||||
|
|
||||||
|
#### Scenario: Remove document from all collections
|
||||||
|
|
||||||
|
- **WHEN** `kb_set_collection` is called with `document_id=42` and no `collection` parameter
|
||||||
|
- **THEN** all `collection:*` tags SHALL be removed from the document
|
||||||
|
- **AND** the response SHALL include `"collection": null`
|
||||||
|
|
||||||
|
#### Scenario: Document not found
|
||||||
|
|
||||||
|
- **WHEN** `kb_set_collection` is called with a `document_id` that does not exist
|
||||||
|
- **THEN** the tool SHALL return an error response indicating the document was not found
|
||||||
|
|
||||||
|
### Requirement: Delete document via MCP
|
||||||
|
|
||||||
|
The MCP server SHALL expose a `kb_delete` tool that permanently deletes a document from the knowledge base. The tool SHALL accept a `document_id` (required integer). Deletion SHALL remove the document, its chunks, embeddings, tags, and any stored file on disk.
|
||||||
|
|
||||||
|
The tool SHALL return a confirmation response including the deleted document's ID and title.
|
||||||
|
|
||||||
|
#### Scenario: Successful deletion
|
||||||
|
|
||||||
|
- **WHEN** `kb_delete` is called with `document_id=42`
|
||||||
|
- **THEN** the document, its chunks, embeddings, tag associations, and stored file SHALL be deleted
|
||||||
|
- **AND** the response SHALL include `"status": "deleted"`, the `document_id`, and the document `title`
|
||||||
|
|
||||||
|
#### Scenario: Document not found
|
||||||
|
|
||||||
|
- **WHEN** `kb_delete` is called with a `document_id` that does not exist
|
||||||
|
- **THEN** the tool SHALL return an error response indicating the document was not found
|
||||||
|
|
||||||
|
### Requirement: Engine client delete method
|
||||||
|
|
||||||
|
The MCP engine client (`mcp/engine.py`) SHALL provide a `delete_document(doc_id)` method that sends a `DELETE` request to `/api/v1/documents/{doc_id}` and returns the JSON response. The method SHALL raise on non-2xx status codes, consistent with other engine client methods.
|
||||||
|
|
||||||
|
#### Scenario: Successful engine client delete call
|
||||||
|
|
||||||
|
- **WHEN** `delete_document(42)` is called and the engine API returns 200
|
||||||
|
- **THEN** the method SHALL return the parsed JSON response
|
||||||
|
|
||||||
|
#### Scenario: Engine client delete for missing document
|
||||||
|
|
||||||
|
- **WHEN** `delete_document(999)` is called and the engine API returns 404
|
||||||
|
- **THEN** the method SHALL raise an `httpx.HTTPStatusError`
|
||||||
@@ -0,0 +1,12 @@
|
|||||||
|
## 1. Engine Client
|
||||||
|
|
||||||
|
- [x] 1.1 Add `delete_document(doc_id)` method to `mcp/engine.py`
|
||||||
|
|
||||||
|
## 2. MCP Tools
|
||||||
|
|
||||||
|
- [x] 2.1 Add `kb_set_collection` tool to `mcp/server.py`
|
||||||
|
- [x] 2.2 Add `kb_delete` tool to `mcp/server.py`
|
||||||
|
|
||||||
|
## 3. Verification
|
||||||
|
|
||||||
|
- [x] 3.1 Test kb_set_collection and kb_delete against running engine
|
||||||
@@ -0,0 +1,2 @@
|
|||||||
|
schema: spec-driven
|
||||||
|
created: 2026-04-06
|
||||||
@@ -0,0 +1,37 @@
|
|||||||
|
## Context
|
||||||
|
|
||||||
|
The project currently ships three Docker image variants: CPU, NVIDIA, and AMD ROCm. The ROCm variant requires a 4.2GB pre-built torch wheel, a multi-stage Dockerfile with ROCm-specific runtime libraries, and additional build/push steps in the release pipeline. ROCm support is less tested and adds disproportionate complexity relative to its usage.
|
||||||
|
|
||||||
|
## Goals / Non-Goals
|
||||||
|
|
||||||
|
**Goals:**
|
||||||
|
- Remove all ROCm-specific files (Dockerfile, compose file, torch wheel)
|
||||||
|
- Remove ROCm build/push from the release pipeline
|
||||||
|
- Update all documentation to reflect CPU + NVIDIA only
|
||||||
|
- Update the docker-deployment spec to remove ROCm requirements
|
||||||
|
|
||||||
|
**Non-Goals:**
|
||||||
|
- Changing any engine application code (it is already GPU-vendor-agnostic via PyTorch)
|
||||||
|
- Modifying the CPU or NVIDIA Dockerfiles (beyond what's already in-flight)
|
||||||
|
- Providing a migration path for ROCm users (they can stay on 3.2.x or use CPU mode)
|
||||||
|
|
||||||
|
## Decisions
|
||||||
|
|
||||||
|
**1. Delete ROCm files outright rather than deprecating**
|
||||||
|
|
||||||
|
Remove `Dockerfile.rocm`, `compose.rocm.yaml`, and `assets/` immediately rather than marking them deprecated. There are no downstream consumers that depend on automated ROCm builds — anyone needing AMD support can pin to the last ROCm-supporting release.
|
||||||
|
|
||||||
|
*Alternative considered*: Keep files but stop publishing images. Rejected — dead code is confusing and still requires maintenance awareness.
|
||||||
|
|
||||||
|
**2. Leave archived openspec changes untouched**
|
||||||
|
|
||||||
|
Archived changes under `openspec/changes/archive/` contain historical ROCm references. These are historical records and should not be modified.
|
||||||
|
|
||||||
|
**3. Update GPU-vendor-agnostic requirement to reflect NVIDIA-only scope**
|
||||||
|
|
||||||
|
The existing spec requirement "Application code is GPU-vendor-agnostic" remains true at the code level (PyTorch abstracts GPU vendors), but the project no longer provides or tests ROCm images. The spec should be simplified to reflect that only NVIDIA and CPU are supported deployment targets.
|
||||||
|
|
||||||
|
## Risks / Trade-offs
|
||||||
|
|
||||||
|
- **[Breaking change for AMD users]** → Users on AMD GPUs must stay on 3.2.x or use CPU mode. Mitigated by the fact that ROCm support was already "less tested" per the original design risk assessment.
|
||||||
|
- **[Future re-addition harder]** → If ROCm support is needed later, the Dockerfile and compose file would need to be recreated. Mitigated by git history preserving the removed files.
|
||||||
@@ -0,0 +1,29 @@
|
|||||||
|
## Why
|
||||||
|
|
||||||
|
AMD ROCm support adds significant complexity and maintenance burden to the project — the ROCm torch wheel alone is 4.2GB, the Dockerfile requires a multi-stage build with ROCm-specific runtime libraries, and the release pipeline must build/push additional images. The final container is >20Gb. ROCm support is less tested and less commonly used than CPU or NVIDIA. Removing it keeps the project focused and manageable.
|
||||||
|
|
||||||
|
## What Changes
|
||||||
|
|
||||||
|
- **BREAKING**: Remove AMD ROCm Docker image (`Dockerfile.rocm`) and compose file (`compose.rocm.yaml`)
|
||||||
|
- **BREAKING**: Remove ROCm image build/push/release-notes from the engine release script
|
||||||
|
- Remove pre-built ROCm torch wheel from `assets/`
|
||||||
|
- Remove all AMD/ROCm references from user-facing docs (README, DEVELOPER)
|
||||||
|
- Update docker-deployment spec to reflect CPU + NVIDIA only
|
||||||
|
|
||||||
|
## Capabilities
|
||||||
|
|
||||||
|
### New Capabilities
|
||||||
|
|
||||||
|
_(none)_
|
||||||
|
|
||||||
|
### Modified Capabilities
|
||||||
|
|
||||||
|
- `docker-deployment`: Remove AMD ROCm Docker image requirement and all ROCm-specific scenarios. Deployment now covers CPU and NVIDIA only.
|
||||||
|
|
||||||
|
## Impact
|
||||||
|
|
||||||
|
- **Docker images**: ROCm image variant no longer published
|
||||||
|
- **Users**: Anyone running KB on AMD GPUs will need to stay on the last version with ROCm support (3.2.x) or switch to CPU mode
|
||||||
|
- **Release pipeline**: `release-engine.sh` simplified — only CPU and NVIDIA images
|
||||||
|
- **Repository size**: ~4.2GB reduction by removing the torch wheel from `assets/`
|
||||||
|
- **Docs**: README and DEVELOPER updated to remove AMD quick-start and build instructions
|
||||||
+76
@@ -0,0 +1,76 @@
|
|||||||
|
## REMOVED Requirements
|
||||||
|
|
||||||
|
### Requirement: AMD ROCm Docker image
|
||||||
|
|
||||||
|
**Reason**: AMD ROCm support removed to reduce project complexity and binary size. The ROCm torch wheel is 4.2GB and the variant is less tested than CPU or NVIDIA.
|
||||||
|
|
||||||
|
**Migration**: Users on AMD GPUs should stay on engine v3.2.x or switch to CPU mode (`KB_DEVICE=cpu`).
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## MODIFIED Requirements
|
||||||
|
|
||||||
|
### Requirement: Application code is GPU-vendor-agnostic
|
||||||
|
|
||||||
|
The Python engine code SHALL NOT reference CUDA directly. GPU abstraction SHALL be handled at the Docker image level (base image selection and pip package choice). The same application code SHALL run on both NVIDIA and CPU images without modification.
|
||||||
|
|
||||||
|
#### Scenario: Same engine code on both platforms
|
||||||
|
- **WHEN** the engine starts on an NVIDIA image and a CPU image with identical configuration
|
||||||
|
- **THEN** both SHALL load the model, accept requests, and return identical search results for the same query and data
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
### Requirement: Compose files for deployment
|
||||||
|
|
||||||
|
The project SHALL provide Docker Compose files for single-command deployment. Compose files SHALL use `build:` context for local development. Release notes SHALL document the versioned image tag for users pulling pre-built images.
|
||||||
|
|
||||||
|
#### Scenario: Start NVIDIA deployment
|
||||||
|
- **WHEN** an admin runs `docker compose -f compose.nvidia.yaml up -d`
|
||||||
|
- **THEN** the engine SHALL start with GPU access, bind-mount the data directory, and be reachable on the configured port
|
||||||
|
|
||||||
|
#### Scenario: Automatic restart
|
||||||
|
- **WHEN** the engine process crashes or the host reboots
|
||||||
|
- **THEN** Docker SHALL automatically restart the container (restart policy `unless-stopped`)
|
||||||
|
|
||||||
|
#### Scenario: Configure via environment
|
||||||
|
- **WHEN** an admin sets environment variables in the compose file (KB_MODEL, KB_API_KEY, KB_DEVICE, KB_MCP_ALLOWED_HOSTS, etc.)
|
||||||
|
- **THEN** the engine and MCP server SHALL use those values
|
||||||
|
|
||||||
|
#### Scenario: Pre-built image deployment
|
||||||
|
- **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`)
|
||||||
|
|
||||||
|
#### Scenario: MCP allowed hosts in Compose
|
||||||
|
- **WHEN** the kb-mcp service is defined in a Compose file
|
||||||
|
- **THEN** the environment block SHALL include `KB_MCP_ALLOWED_HOSTS` with a comment explaining its format and purpose
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
### Requirement: Bind-mount data directory
|
||||||
|
|
||||||
|
The engine SHALL store all persistent state (SQLite database, HF model cache, staging directory) under a single configurable data directory. This directory SHALL be mounted from the host via bind mount.
|
||||||
|
|
||||||
|
#### Scenario: Data directory structure
|
||||||
|
- **WHEN** the engine starts for the first time
|
||||||
|
- **THEN** it SHALL create the following structure under the data directory:
|
||||||
|
- `kb.db` — SQLite database
|
||||||
|
- `hf_cache/` — HuggingFace model cache
|
||||||
|
- `staging/` — temporary files for queued ingestion jobs
|
||||||
|
|
||||||
|
#### Scenario: Portable data across hosts
|
||||||
|
- **WHEN** an admin copies the data directory from Host A to Host B and starts the engine with the same bind mount path
|
||||||
|
- **THEN** the engine SHALL start successfully and serve all previously ingested documents without reprocessing
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
### Requirement: CPU-only fallback
|
||||||
|
|
||||||
|
The Dockerfiles SHALL produce images that work without GPU access. If no GPU is available, the engine SHALL fall back to CPU for all operations.
|
||||||
|
|
||||||
|
#### Scenario: No GPU available
|
||||||
|
- **WHEN** the container starts without GPU passthrough (no `--gpus`)
|
||||||
|
- **THEN** the engine SHALL detect no GPU, load the model on CPU, and log a warning that GPU acceleration is unavailable
|
||||||
|
|
||||||
|
#### Scenario: Explicit CPU mode
|
||||||
|
- **WHEN** `KB_DEVICE=cpu` and `KB_INGEST_DEVICE=cpu` are set in the environment
|
||||||
|
- **THEN** the engine SHALL use CPU regardless of GPU availability
|
||||||
@@ -0,0 +1,20 @@
|
|||||||
|
## 1. Delete ROCm files
|
||||||
|
|
||||||
|
- [x] 1.1 Delete `engine/Dockerfile.rocm`
|
||||||
|
- [x] 1.2 Delete `engine/compose.rocm.yaml`
|
||||||
|
- [x] 1.3 Delete `assets/` directory (ROCm torch wheel)
|
||||||
|
|
||||||
|
## 2. Update release pipeline
|
||||||
|
|
||||||
|
- [x] 2.1 Remove ROCm image build, tag, and push from `release-engine.sh`
|
||||||
|
- [x] 2.2 Remove ROCm entries from release notes output in `release-engine.sh`
|
||||||
|
|
||||||
|
## 3. Update documentation
|
||||||
|
|
||||||
|
- [x] 3.1 Remove AMD GPU quick-start section and ROCm references from `README.md`
|
||||||
|
- [x] 3.2 Remove ROCm build instructions and `compose.rocm.yaml` references from `DEVELOPER.md`
|
||||||
|
- [x] 3.3 Remove `onnxruntime-rocm` migration note from `DEVELOPER.md`
|
||||||
|
|
||||||
|
## 4. Update specs
|
||||||
|
|
||||||
|
- [x] 4.1 Update `openspec/specs/docker-deployment/spec.md` — remove AMD ROCm requirement, remove ROCm scenarios, update GPU-agnostic requirement to CPU + NVIDIA scope
|
||||||
@@ -0,0 +1,43 @@
|
|||||||
|
# Agent-Side Search Patterns
|
||||||
|
|
||||||
|
## Purpose
|
||||||
|
|
||||||
|
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
|
||||||
|
|
||||||
|
### Requirement: Query expansion guidance in tool description
|
||||||
|
|
||||||
|
The `kb_search` MCP tool description SHALL include guidance on query expansion as a recommended pattern for complex queries.
|
||||||
|
|
||||||
|
#### Scenario: Tool description includes expansion pattern
|
||||||
|
- **WHEN** an agent reads the `kb_search` tool description
|
||||||
|
- **THEN** the description SHALL include guidance such as: "For complex queries, consider expanding into 2-3 variant phrasings and calling this tool multiple times, then deduplicating results by chunk_id"
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
### Requirement: Reranking guidance in tool description
|
||||||
|
|
||||||
|
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 covers server-side reranking
|
||||||
|
- **WHEN** an agent reads the `kb_search` tool description
|
||||||
|
- **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 external LLM dependency
|
||||||
|
|
||||||
|
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 external LLM dependency
|
||||||
|
- **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
|
||||||
|
|
||||||
|
#### 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`
|
||||||
@@ -0,0 +1,230 @@
|
|||||||
|
## ADDED Requirements
|
||||||
|
|
||||||
|
### Requirement: Common selection filter
|
||||||
|
|
||||||
|
All bulk engine endpoints SHALL accept a JSON body with the following optional selection fields, combined with AND logic:
|
||||||
|
|
||||||
|
- `document_ids` (list of int) — match documents with these specific IDs
|
||||||
|
- `tags` (list of str) — match documents that have ALL specified tags
|
||||||
|
- `doc_type` (str) — match documents with this document type
|
||||||
|
- `from_id` (int) — match documents with id >= this value
|
||||||
|
- `to_id` (int) — match documents with id <= this value
|
||||||
|
|
||||||
|
At least one selection field MUST be present. If no selection fields are provided, the endpoint SHALL return 400 Bad Request.
|
||||||
|
|
||||||
|
#### Scenario: Filter by tags and doc_type
|
||||||
|
|
||||||
|
- **WHEN** a bulk endpoint receives `{"tags": ["draft"], "doc_type": "note"}`
|
||||||
|
- **THEN** it SHALL match only documents that have the tag "draft" AND have doc_type "note"
|
||||||
|
|
||||||
|
#### Scenario: Filter by ID range
|
||||||
|
|
||||||
|
- **WHEN** a bulk endpoint receives `{"from_id": 10, "to_id": 50}`
|
||||||
|
- **THEN** it SHALL match documents with id >= 10 AND id <= 50
|
||||||
|
|
||||||
|
#### Scenario: Filter by explicit IDs
|
||||||
|
|
||||||
|
- **WHEN** a bulk endpoint receives `{"document_ids": [1, 5, 12]}`
|
||||||
|
- **THEN** it SHALL match only documents with those specific IDs
|
||||||
|
|
||||||
|
#### Scenario: Combined filters
|
||||||
|
|
||||||
|
- **WHEN** a bulk endpoint receives `{"tags": ["agent:mybot"], "doc_type": "note", "from_id": 100}`
|
||||||
|
- **THEN** it SHALL match documents satisfying ALL three criteria
|
||||||
|
|
||||||
|
#### Scenario: No selection fields provided
|
||||||
|
|
||||||
|
- **WHEN** a bulk endpoint receives `{}` or `{"force": true}` with no selection fields
|
||||||
|
- **THEN** it SHALL return 400 Bad Request
|
||||||
|
|
||||||
|
### Requirement: Safety threshold
|
||||||
|
|
||||||
|
All bulk endpoints SHALL enforce a safety threshold. Before executing, the engine SHALL count the matched documents and the total documents in the database. If `matched / total * 100` exceeds the configured threshold, the request SHALL be rejected with 409 Conflict.
|
||||||
|
|
||||||
|
The response SHALL include: `error` ("safety_threshold_exceeded"), `message` (human-readable), `matched` (int), `total` (int), `percent` (float), and `threshold` (int).
|
||||||
|
|
||||||
|
The threshold SHALL default to 70 and be configurable via the `KB_BULK_SAFETY_PERCENT` environment variable (integer 0-100). A value of 0 disables the check.
|
||||||
|
|
||||||
|
The caller MAY override the threshold by including `"force": true` in the request body.
|
||||||
|
|
||||||
|
#### Scenario: Threshold exceeded
|
||||||
|
|
||||||
|
- **GIVEN** 1000 total documents and `KB_BULK_SAFETY_PERCENT` is 70
|
||||||
|
- **WHEN** a bulk endpoint matches 750 documents (75%) without `force: true`
|
||||||
|
- **THEN** it SHALL return 409 with `matched: 750`, `total: 1000`, `percent: 75.0`, `threshold: 70`
|
||||||
|
|
||||||
|
#### Scenario: Threshold not exceeded
|
||||||
|
|
||||||
|
- **GIVEN** 1000 total documents and `KB_BULK_SAFETY_PERCENT` is 70
|
||||||
|
- **WHEN** a bulk endpoint matches 500 documents (50%) without `force: true`
|
||||||
|
- **THEN** the operation SHALL proceed normally
|
||||||
|
|
||||||
|
#### Scenario: Force override
|
||||||
|
|
||||||
|
- **GIVEN** 1000 total documents and a match of 900 (90%)
|
||||||
|
- **WHEN** the request includes `"force": true`
|
||||||
|
- **THEN** the operation SHALL proceed regardless of threshold
|
||||||
|
|
||||||
|
#### Scenario: Zero threshold
|
||||||
|
|
||||||
|
- **GIVEN** `KB_BULK_SAFETY_PERCENT` is 0
|
||||||
|
- **THEN** the safety check SHALL be effectively disabled for all operations
|
||||||
|
|
||||||
|
### Requirement: Synchronous response with audit log
|
||||||
|
|
||||||
|
All bulk endpoints SHALL execute synchronously and return a JSON response with:
|
||||||
|
|
||||||
|
- `job_id` (int) — ID of the audit log entry in the jobs table
|
||||||
|
- `status` (str) — "done" or "partial_failure"
|
||||||
|
- `matched` (int) — number of documents that matched the selection
|
||||||
|
- `succeeded` (int) — number of documents successfully processed
|
||||||
|
- `failed` (int) — number of documents that failed
|
||||||
|
- `errors` (list) — array of `{"document_id": int, "error": str}` for each failure (empty on full success)
|
||||||
|
|
||||||
|
A job record SHALL be created in the jobs table with `job_type` set to the operation type. The `filename` field SHALL store a JSON representation of the selection filter. The `error` field SHALL store a JSON array of individual errors if any occurred.
|
||||||
|
|
||||||
|
#### Scenario: Full success
|
||||||
|
|
||||||
|
- **WHEN** a bulk operation matches 50 documents and all succeed
|
||||||
|
- **THEN** the response SHALL have `status: "done"`, `matched: 50`, `succeeded: 50`, `failed: 0`, `errors: []`
|
||||||
|
|
||||||
|
#### Scenario: Partial failure
|
||||||
|
|
||||||
|
- **WHEN** a bulk operation matches 50 documents but 2 fail
|
||||||
|
- **THEN** the response SHALL have `status: "partial_failure"`, `matched: 50`, `succeeded: 48`, `failed: 2`, and `errors` listing the 2 failures
|
||||||
|
|
||||||
|
### Requirement: Bulk delete endpoint
|
||||||
|
|
||||||
|
The engine SHALL expose `POST /api/v1/bulk/delete` which permanently deletes all documents matching the selection filter. For each matched document, it SHALL delete embeddings from `chunks_vec`, delete the document row (cascading to chunks and document_tags), and delete any stored file from disk.
|
||||||
|
|
||||||
|
Database deletions SHALL be performed within a single transaction. File deletions SHALL occur after the transaction commits and SHALL be best-effort (failures logged but not counted as document failures).
|
||||||
|
|
||||||
|
#### Scenario: Bulk delete by tag
|
||||||
|
|
||||||
|
- **WHEN** `POST /api/v1/bulk/delete` receives `{"tags": ["old", "draft"]}`
|
||||||
|
- **THEN** all documents with both tags "old" and "draft" SHALL be deleted
|
||||||
|
- **AND** their chunks, embeddings, tag associations, and stored files SHALL be removed
|
||||||
|
|
||||||
|
#### Scenario: Bulk delete with no matches
|
||||||
|
|
||||||
|
- **WHEN** `POST /api/v1/bulk/delete` receives a filter that matches 0 documents
|
||||||
|
- **THEN** the response SHALL have `matched: 0`, `succeeded: 0`, `failed: 0`
|
||||||
|
|
||||||
|
### Requirement: Bulk tags endpoint
|
||||||
|
|
||||||
|
The engine SHALL expose `POST /api/v1/bulk/tags` which adds and/or removes tags on all documents matching the selection filter. The request body SHALL include the selection filter plus:
|
||||||
|
|
||||||
|
- `add` (list of str, optional) — tags to add
|
||||||
|
- `remove` (list of str, optional) — tags to remove
|
||||||
|
|
||||||
|
At least one of `add` or `remove` MUST be present. The endpoint SHALL return 400 if neither is provided.
|
||||||
|
|
||||||
|
The endpoint SHALL update `updated_at` on all affected documents.
|
||||||
|
|
||||||
|
#### Scenario: Add and remove tags in one call
|
||||||
|
|
||||||
|
- **WHEN** `POST /api/v1/bulk/tags` receives `{"tags": ["agent:mybot"], "add": ["reviewed"], "remove": ["pending"]}`
|
||||||
|
- **THEN** all documents tagged "agent:mybot" SHALL have "reviewed" added and "pending" removed
|
||||||
|
|
||||||
|
### Requirement: Bulk set-tags endpoint
|
||||||
|
|
||||||
|
The engine SHALL expose `POST /api/v1/bulk/set-tags` which replaces all tags on matched documents with a new set. The request body SHALL include the selection filter plus:
|
||||||
|
|
||||||
|
- `new_tags` (list of str) — the replacement tag set
|
||||||
|
|
||||||
|
The endpoint SHALL remove all existing tag associations from matched documents, then apply the new set. It SHALL update `updated_at` on all affected documents.
|
||||||
|
|
||||||
|
#### Scenario: Replace all tags
|
||||||
|
|
||||||
|
- **WHEN** `POST /api/v1/bulk/set-tags` receives `{"doc_type": "note", "new_tags": ["clean", "final"]}`
|
||||||
|
- **THEN** all notes SHALL have their existing tags removed and replaced with "clean" and "final"
|
||||||
|
|
||||||
|
### Requirement: Jobs table extension
|
||||||
|
|
||||||
|
The jobs table SHALL be extended with a `job_type` column (TEXT, default "ingest") to distinguish ingestion jobs from bulk operation audit entries. Valid values: "ingest", "bulk_delete", "bulk_tags", "bulk_set_tags".
|
||||||
|
|
||||||
|
Existing jobs SHALL default to `job_type = "ingest"`. The existing jobs list endpoint and CLI `kb jobs` command SHALL continue to work unchanged.
|
||||||
|
|
||||||
|
#### Scenario: Migration adds column
|
||||||
|
|
||||||
|
- **GIVEN** an existing database without the `job_type` column
|
||||||
|
- **WHEN** the engine starts
|
||||||
|
- **THEN** the column SHALL be added with default value "ingest"
|
||||||
|
|
||||||
|
### Requirement: Engine config for safety threshold
|
||||||
|
|
||||||
|
The engine `Config` class SHALL read `KB_BULK_SAFETY_PERCENT` from the environment as an integer (default 70, range 0-100). This value SHALL be used as the default safety threshold for all bulk endpoints.
|
||||||
|
|
||||||
|
### Requirement: MCP bulk delete tool
|
||||||
|
|
||||||
|
The MCP server SHALL expose a `kb_bulk_delete` tool with parameters: `document_ids` (optional list of int), `tags` (optional list of str), `doc_type` (optional str), `from_id` (optional int), `to_id` (optional int), `force` (optional bool).
|
||||||
|
|
||||||
|
The tool SHALL call `POST /api/v1/bulk/delete` on the engine via the engine client and return the JSON response.
|
||||||
|
|
||||||
|
The tool description SHALL clearly state that `tags` is a selection filter (which documents to delete), not tags to delete.
|
||||||
|
|
||||||
|
#### Scenario: MCP bulk delete by tag
|
||||||
|
|
||||||
|
- **WHEN** `kb_bulk_delete(tags=["old"])` is called
|
||||||
|
- **THEN** the engine client SHALL send `POST /api/v1/bulk/delete` with `{"tags": ["old"]}`
|
||||||
|
- **AND** the tool SHALL return the engine's JSON response
|
||||||
|
|
||||||
|
### Requirement: MCP bulk tags tool
|
||||||
|
|
||||||
|
The MCP server SHALL expose a `kb_bulk_tags` tool with parameters: `document_ids`, `tags`, `doc_type`, `from_id`, `to_id` (selection filters), plus `add` (optional list of str), `remove` (optional list of str), and `force` (optional bool).
|
||||||
|
|
||||||
|
The tool description SHALL clearly distinguish `tags` (selection filter) from `add`/`remove` (tag changes to apply).
|
||||||
|
|
||||||
|
#### Scenario: MCP bulk tag update
|
||||||
|
|
||||||
|
- **WHEN** `kb_bulk_tags(tags=["agent:mybot"], add=["reviewed"], remove=["draft"])` is called
|
||||||
|
- **THEN** the engine client SHALL send the appropriate `POST /api/v1/bulk/tags` request
|
||||||
|
|
||||||
|
### Requirement: MCP bulk set-tags tool
|
||||||
|
|
||||||
|
The MCP server SHALL expose a `kb_bulk_set_tags` tool with parameters: `document_ids`, `tags`, `doc_type`, `from_id`, `to_id` (selection filters), plus `new_tags` (list of str) and `force` (optional bool).
|
||||||
|
|
||||||
|
#### Scenario: MCP bulk set tags
|
||||||
|
|
||||||
|
- **WHEN** `kb_bulk_set_tags(doc_type="note", new_tags=["clean"])` is called
|
||||||
|
- **THEN** the engine client SHALL send `POST /api/v1/bulk/set-tags` with `{"doc_type": "note", "new_tags": ["clean"]}`
|
||||||
|
|
||||||
|
### Requirement: MCP engine client bulk methods
|
||||||
|
|
||||||
|
The MCP engine client (`mcp/engine.py`) SHALL provide three new methods:
|
||||||
|
|
||||||
|
- `bulk_delete(document_ids?, tags?, doc_type?, from_id?, to_id?, force?)` → dict
|
||||||
|
- `bulk_tags(document_ids?, tags?, doc_type?, from_id?, to_id?, add?, remove?, force?)` → dict
|
||||||
|
- `bulk_set_tags(document_ids?, tags?, doc_type?, from_id?, to_id?, new_tags?, force?)` → dict
|
||||||
|
|
||||||
|
Each SHALL send a POST request to the corresponding `/api/v1/bulk/*` endpoint with the parameters as a JSON body. Each SHALL raise on non-2xx status codes, consistent with existing methods.
|
||||||
|
|
||||||
|
### Requirement: CLI bulk-remove command
|
||||||
|
|
||||||
|
The CLI SHALL expose a `kb bulk-remove` command with flags: `--tags` (comma-separated), `--type`, `--ids` (comma-separated), `--from-id`, `--to-id`, `--force`/`-f`, `--yes`/`-y`.
|
||||||
|
|
||||||
|
Without `--yes`, the CLI SHALL first display the match count and ask for interactive confirmation before proceeding.
|
||||||
|
|
||||||
|
The command SHALL call `POST /api/v1/bulk/delete` with the constructed filter.
|
||||||
|
|
||||||
|
#### Scenario: CLI bulk remove with confirmation
|
||||||
|
|
||||||
|
- **WHEN** `kb bulk-remove --tags "draft,old" --type note` is run without `--yes`
|
||||||
|
- **THEN** the CLI SHALL display "This will delete N documents matching: tags=[draft,old] type=note" and prompt "Proceed? [y/N]"
|
||||||
|
|
||||||
|
#### Scenario: CLI bulk remove with --yes
|
||||||
|
|
||||||
|
- **WHEN** `kb bulk-remove --tags "draft" --yes` is run
|
||||||
|
- **THEN** the CLI SHALL proceed without prompting
|
||||||
|
|
||||||
|
### Requirement: CLI bulk-tag command
|
||||||
|
|
||||||
|
The CLI SHALL expose a `kb bulk-tag` command with the same filter flags as `bulk-remove`, plus `--add` and `--remove` (comma-separated tag lists).
|
||||||
|
|
||||||
|
The command SHALL call `POST /api/v1/bulk/tags` with the constructed filter and tag changes.
|
||||||
|
|
||||||
|
### Requirement: CLI bulk-set-tags command
|
||||||
|
|
||||||
|
The CLI SHALL expose a `kb bulk-set-tags` command with the filter flags, plus `--set` (comma-separated list of replacement tags).
|
||||||
|
|
||||||
|
The command SHALL call `POST /api/v1/bulk/set-tags` with the constructed filter and `new_tags`.
|
||||||
@@ -10,7 +10,7 @@ DEVELOPER.md SHALL contain instructions for building both the engine and client
|
|||||||
|
|
||||||
#### Scenario: Engine build from source
|
#### Scenario: Engine build from source
|
||||||
- **WHEN** a developer reads DEVELOPER.md
|
- **WHEN** a developer reads DEVELOPER.md
|
||||||
- **THEN** it SHALL include instructions for starting the engine from source using compose files (both NVIDIA and ROCm)
|
- **THEN** it SHALL include instructions for starting the engine from source using compose files (NVIDIA and CPU)
|
||||||
|
|
||||||
#### Scenario: Client build from source
|
#### Scenario: Client build from source
|
||||||
- **WHEN** a developer reads DEVELOPER.md
|
- **WHEN** a developer reads DEVELOPER.md
|
||||||
@@ -31,13 +31,6 @@ DEVELOPER.md SHALL document the release process for both client and engine, incl
|
|||||||
- **WHEN** a developer reads DEVELOPER.md
|
- **WHEN** a developer reads DEVELOPER.md
|
||||||
- **THEN** it SHALL include how to check client and engine versions
|
- **THEN** it SHALL include how to check client and engine versions
|
||||||
|
|
||||||
### Requirement: DEVELOPER.md contains developer notes
|
|
||||||
DEVELOPER.md SHALL include any forward-looking developer notes such as migration plans or technical debt items.
|
|
||||||
|
|
||||||
#### Scenario: ROCm migration note
|
|
||||||
- **WHEN** a developer reads DEVELOPER.md
|
|
||||||
- **THEN** it SHALL include the ROCm runtime migration note about onnxruntime and MIGraphX
|
|
||||||
|
|
||||||
### Requirement: README.md excludes developer-only content
|
### Requirement: README.md excludes developer-only content
|
||||||
README.md SHALL NOT contain build-from-source instructions, release processes, or developer-only notes.
|
README.md SHALL NOT contain build-from-source instructions, release processes, or developer-only notes.
|
||||||
|
|
||||||
@@ -49,10 +42,6 @@ README.md SHALL NOT contain build-from-source instructions, release processes, o
|
|||||||
- **WHEN** a user reads README.md
|
- **WHEN** a user reads README.md
|
||||||
- **THEN** there SHALL be no "Building and releasing" section
|
- **THEN** there SHALL be no "Building and releasing" section
|
||||||
|
|
||||||
#### Scenario: No developer notes in README
|
|
||||||
- **WHEN** a user reads README.md
|
|
||||||
- **THEN** there SHALL be no "Future: ROCm runtime migration" section
|
|
||||||
|
|
||||||
### Requirement: README.md cross-references DEVELOPER.md
|
### Requirement: README.md cross-references DEVELOPER.md
|
||||||
README.md SHALL include a link to DEVELOPER.md for users who want to build from source or contribute.
|
README.md SHALL include a link to DEVELOPER.md for users who want to build from source or contribute.
|
||||||
|
|
||||||
|
|||||||
@@ -2,7 +2,7 @@
|
|||||||
|
|
||||||
## Purpose
|
## Purpose
|
||||||
|
|
||||||
Docker deployment provides containerized packaging of the knowledge base engine with GPU support for NVIDIA and AMD platforms, along with Compose files for single-command deployment.
|
Docker deployment provides containerized packaging of the knowledge base engine with GPU support for NVIDIA, along with Compose files for single-command deployment.
|
||||||
|
|
||||||
## Requirements
|
## Requirements
|
||||||
|
|
||||||
@@ -20,26 +20,12 @@ The project SHALL provide a `Dockerfile.nvidia` that builds the engine on an NVI
|
|||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
### Requirement: AMD ROCm Docker image
|
|
||||||
|
|
||||||
The project SHALL provide a `Dockerfile.rocm` that builds the engine on an AMD ROCm base image with GPU support for PyTorch and ONNX Runtime.
|
|
||||||
|
|
||||||
#### Scenario: Build ROCm image
|
|
||||||
- **WHEN** an admin runs `docker compose -f compose.rocm.yaml build`
|
|
||||||
- **THEN** the build SHALL produce a working image with ROCm runtime, PyTorch with ROCm support, onnxruntime-rocm, and all engine dependencies
|
|
||||||
|
|
||||||
#### Scenario: GPU access in ROCm container
|
|
||||||
- **WHEN** the ROCm container starts with `--device=/dev/kfd --device=/dev/dri`
|
|
||||||
- **THEN** `torch.cuda.is_available()` SHALL return True (via HIP) and the engine SHALL load the embedding model on GPU
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
### Requirement: Application code is GPU-vendor-agnostic
|
### Requirement: Application code is GPU-vendor-agnostic
|
||||||
|
|
||||||
The Python engine code SHALL NOT reference CUDA or ROCm directly. GPU vendor abstraction SHALL be handled entirely at the Docker image level (base image selection and pip package choice). The same application code SHALL run on both NVIDIA and AMD images without modification.
|
The Python engine code SHALL NOT reference CUDA directly. GPU abstraction SHALL be handled at the Docker image level (base image selection and pip package choice). The same application code SHALL run on both NVIDIA and CPU images without modification.
|
||||||
|
|
||||||
#### Scenario: Same engine code on both platforms
|
#### Scenario: Same engine code on both platforms
|
||||||
- **WHEN** the engine starts on an NVIDIA image and an AMD image with identical configuration
|
- **WHEN** the engine starts on an NVIDIA image and a CPU image with identical configuration
|
||||||
- **THEN** both SHALL load the model, accept requests, and return identical search results for the same query and data
|
- **THEN** both SHALL load the model, accept requests, and return identical search results for the same query and data
|
||||||
|
|
||||||
---
|
---
|
||||||
@@ -59,10 +45,6 @@ The engine SHALL store all persistent state (SQLite database, HF model cache, st
|
|||||||
- **WHEN** an admin copies the data directory from Host A to Host B and starts the engine with the same bind mount path
|
- **WHEN** an admin copies the data directory from Host A to Host B and starts the engine with the same bind mount path
|
||||||
- **THEN** the engine SHALL start successfully and serve all previously ingested documents without reprocessing
|
- **THEN** the engine SHALL start successfully and serve all previously ingested documents without reprocessing
|
||||||
|
|
||||||
#### Scenario: Portable data across GPU vendors
|
|
||||||
- **WHEN** an admin moves the data directory from an NVIDIA host to an AMD host (same model name)
|
|
||||||
- **THEN** the engine SHALL start successfully. Embeddings in the database remain valid (they are model-specific, not GPU-vendor-specific)
|
|
||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
### Requirement: Compose files for deployment
|
### Requirement: Compose files for deployment
|
||||||
@@ -73,21 +55,51 @@ The project SHALL provide Docker Compose files for single-command deployment. Co
|
|||||||
- **WHEN** an admin runs `docker compose -f compose.nvidia.yaml up -d`
|
- **WHEN** an admin runs `docker compose -f compose.nvidia.yaml up -d`
|
||||||
- **THEN** the engine SHALL start with GPU access, bind-mount the data directory, and be reachable on the configured port
|
- **THEN** the engine SHALL start with GPU access, bind-mount the data directory, and be reachable on the configured port
|
||||||
|
|
||||||
#### Scenario: Start ROCm deployment
|
|
||||||
- **WHEN** an admin runs `docker compose -f compose.rocm.yaml up -d`
|
|
||||||
- **THEN** the engine SHALL start with GPU access via ROCm device passthrough, bind-mount the data directory, and be reachable on the configured port
|
|
||||||
|
|
||||||
#### Scenario: Automatic restart
|
#### Scenario: Automatic restart
|
||||||
- **WHEN** the engine process crashes or the host reboots
|
- **WHEN** the engine process crashes or the host reboots
|
||||||
- **THEN** Docker SHALL automatically restart the container (restart policy `unless-stopped`)
|
- **THEN** Docker SHALL automatically restart the container (restart policy `unless-stopped`)
|
||||||
|
|
||||||
#### Scenario: Configure via environment
|
#### Scenario: Configure via environment
|
||||||
- **WHEN** an admin sets environment variables in the compose file (KB_MODEL, KB_API_KEY, KB_DEVICE, etc.)
|
- **WHEN** an admin sets environment variables in the compose file (KB_MODEL, KB_API_KEY, KB_DEVICE, KB_MCP_ALLOWED_HOSTS, etc.)
|
||||||
- **THEN** the engine SHALL use those values
|
- **THEN** the engine and MCP server SHALL use those values
|
||||||
|
|
||||||
#### 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
|
||||||
|
- **WHEN** the kb-mcp service is defined in a Compose file
|
||||||
|
- **THEN** the environment block SHALL include `KB_MCP_ALLOWED_HOSTS` with a comment explaining its format and purpose
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
### Requirement: Configurable MCP allowed hosts
|
||||||
|
|
||||||
|
The MCP server SHALL accept a `KB_MCP_ALLOWED_HOSTS` environment variable containing a comma-separated list of additional hosts (IP addresses or FQDNs) that are permitted to connect. The server SHALL always allow `127.0.0.1`, `localhost`, and `[::1]` regardless of this setting. DNS rebinding protection SHALL always be enabled.
|
||||||
|
|
||||||
|
#### Scenario: Remote client connects with allowed host
|
||||||
|
- **WHEN** `KB_MCP_ALLOWED_HOSTS` is set to `192.168.1.50` and a client connects with `Host: 192.168.1.50:3000`
|
||||||
|
- **THEN** the server SHALL accept the request and process it normally
|
||||||
|
|
||||||
|
#### Scenario: Remote client connects with disallowed host
|
||||||
|
- **WHEN** `KB_MCP_ALLOWED_HOSTS` is set to `192.168.1.50` and a client connects with `Host: 10.0.0.99:3000`
|
||||||
|
- **THEN** the server SHALL return HTTP 421 "Invalid Host header"
|
||||||
|
|
||||||
|
#### Scenario: Multiple allowed hosts
|
||||||
|
- **WHEN** `KB_MCP_ALLOWED_HOSTS` is set to `192.168.1.50,kb.example.com`
|
||||||
|
- **THEN** the server SHALL accept requests with `Host` matching either `192.168.1.50` or `kb.example.com` on any port
|
||||||
|
|
||||||
|
#### Scenario: Variable unset or empty
|
||||||
|
- **WHEN** `KB_MCP_ALLOWED_HOSTS` is unset or empty
|
||||||
|
- **THEN** the server SHALL allow only localhost addresses (`127.0.0.1`, `localhost`, `[::1]`) with any port
|
||||||
|
|
||||||
|
#### Scenario: Localhost always allowed
|
||||||
|
- **WHEN** `KB_MCP_ALLOWED_HOSTS` is set to `192.168.1.50`
|
||||||
|
- **THEN** the server SHALL still accept requests with `Host: localhost:3000` or `Host: 127.0.0.1:3000`
|
||||||
|
|
||||||
|
#### Scenario: Allowed origins derived from allowed hosts
|
||||||
|
- **WHEN** `KB_MCP_ALLOWED_HOSTS` includes `192.168.1.50`
|
||||||
|
- **THEN** the server SHALL accept `Origin: http://192.168.1.50:3000` (and any port) in addition to localhost origins
|
||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
@@ -96,7 +108,7 @@ The project SHALL provide Docker Compose files for single-command deployment. Co
|
|||||||
The Dockerfiles SHALL produce images that work without GPU access. If no GPU is available, the engine SHALL fall back to CPU for all operations.
|
The Dockerfiles SHALL produce images that work without GPU access. If no GPU is available, the engine SHALL fall back to CPU for all operations.
|
||||||
|
|
||||||
#### Scenario: No GPU available
|
#### Scenario: No GPU available
|
||||||
- **WHEN** the container starts without GPU passthrough (no `--gpus`, no `/dev/kfd`)
|
- **WHEN** the container starts without GPU passthrough (no `--gpus`)
|
||||||
- **THEN** the engine SHALL detect no GPU, load the model on CPU, and log a warning that GPU acceleration is unavailable
|
- **THEN** the engine SHALL detect no GPU, load the model on CPU, and log a warning that GPU acceleration is unavailable
|
||||||
|
|
||||||
#### Scenario: Explicit CPU mode
|
#### Scenario: Explicit CPU mode
|
||||||
|
|||||||
@@ -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
|
||||||
@@ -150,15 +176,19 @@ The engine SHALL provide endpoints to list, inspect, remove, and download origin
|
|||||||
|
|
||||||
#### Scenario: List documents
|
#### Scenario: List documents
|
||||||
- **WHEN** a client sends `GET /api/v1/documents`
|
- **WHEN** a client sends `GET /api/v1/documents`
|
||||||
- **THEN** the engine SHALL return a JSON array of documents with id, title, doc_type, tags, chunk_count, and created_at
|
- **THEN** the engine SHALL return a JSON array of documents with id, title, doc_type, tags, chunk_count, created_at, and updated_at
|
||||||
|
|
||||||
#### Scenario: List documents with filters
|
#### Scenario: List documents with filters
|
||||||
- **WHEN** a client sends `GET /api/v1/documents?type=pdf&tags=manual`
|
- **WHEN** a client sends `GET /api/v1/documents?type=pdf&tags=manual`
|
||||||
- **THEN** the engine SHALL return only documents matching all specified filters
|
- **THEN** the engine SHALL return only documents matching all specified filters
|
||||||
|
|
||||||
|
#### Scenario: List documents sorted by most recent
|
||||||
|
- **WHEN** a client requests documents sorted by date
|
||||||
|
- **THEN** the engine SHALL use `COALESCE(updated_at, created_at)` for ordering, so un-mutated documents sort by creation time and mutated documents sort by their last update
|
||||||
|
|
||||||
#### Scenario: Get document details
|
#### Scenario: Get document details
|
||||||
- **WHEN** a client sends `GET /api/v1/documents/{id}`
|
- **WHEN** a client sends `GET /api/v1/documents/{id}`
|
||||||
- **THEN** the engine SHALL return the full document record including all chunks, their text content, and whether the original file is available (`has_file: true/false`)
|
- **THEN** the engine SHALL return the full document record including all chunks, their text content, `updated_at`, and whether the original file is available (`has_file: true/false`)
|
||||||
|
|
||||||
#### Scenario: Download original file
|
#### Scenario: Download original file
|
||||||
- **WHEN** a client sends `GET /api/v1/documents/{id}/file`
|
- **WHEN** a client sends `GET /api/v1/documents/{id}/file`
|
||||||
@@ -174,13 +204,45 @@ The engine SHALL provide endpoints to list, inspect, remove, and download origin
|
|||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
|
### Requirement: Note mutation endpoint
|
||||||
|
|
||||||
|
The engine SHALL provide a `PATCH /api/v1/notes/{id}` endpoint for updating existing notes in place. See the `note-mutation` spec for full details.
|
||||||
|
|
||||||
|
#### Scenario: Note update endpoint exists
|
||||||
|
- **WHEN** a client sends `PATCH /api/v1/notes/42` with body `{"text": "new content"}`
|
||||||
|
- **THEN** the engine SHALL process the update synchronously and return the updated document
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
### Requirement: Document updated_at tracking
|
||||||
|
|
||||||
|
The engine SHALL track when documents are modified via an `updated_at` column. This column SHALL be NULL for documents that have never been updated.
|
||||||
|
|
||||||
|
#### Scenario: New document has no updated_at
|
||||||
|
- **WHEN** a document is first ingested
|
||||||
|
- **THEN** `updated_at` SHALL be NULL and `created_at` SHALL be set to the ingestion timestamp
|
||||||
|
|
||||||
|
#### Scenario: Note update sets updated_at
|
||||||
|
- **WHEN** a note is updated via `PATCH /api/v1/notes/{id}`
|
||||||
|
- **THEN** `updated_at` SHALL be set to the current timestamp
|
||||||
|
|
||||||
|
#### Scenario: Tag change sets updated_at
|
||||||
|
- **WHEN** tags are modified via `PUT /api/v1/documents/{id}/tags`
|
||||||
|
- **THEN** `updated_at` SHALL be set to the current timestamp
|
||||||
|
|
||||||
|
#### Scenario: Schema migration for updated_at
|
||||||
|
- **WHEN** the engine starts against a v2 database without an `updated_at` column
|
||||||
|
- **THEN** the engine SHALL automatically add `ALTER TABLE documents ADD COLUMN updated_at TEXT` and all existing documents SHALL have `updated_at = NULL`
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
### Requirement: Tag management
|
### Requirement: Tag management
|
||||||
|
|
||||||
The engine SHALL provide endpoints to list all tags and manage tags on documents.
|
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"]}`
|
||||||
@@ -192,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`
|
||||||
|
|||||||
@@ -265,17 +265,43 @@ The client SHALL provide a `kb reindex` command that triggers re-embedding of al
|
|||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
|
### Requirement: Update note command
|
||||||
|
|
||||||
|
The client SHALL provide a `kb updatenote <id> <text>` command that updates an existing note's content via the engine's `PATCH /api/v1/notes/{id}` endpoint.
|
||||||
|
|
||||||
|
#### Scenario: Update a note
|
||||||
|
- **WHEN** the user runs `kb updatenote 42 "Updated note content"`
|
||||||
|
- **THEN** the client SHALL send `PATCH /api/v1/notes/42` with body `{"text": "Updated note content"}` and display the result
|
||||||
|
|
||||||
|
#### Scenario: Update a note with JSON output
|
||||||
|
- **WHEN** the user runs `kb updatenote 42 "new content" --format json`
|
||||||
|
- **THEN** the client SHALL output the raw JSON response from the engine
|
||||||
|
|
||||||
|
#### Scenario: Update a non-existent document
|
||||||
|
- **WHEN** the user runs `kb updatenote 999 "text"` and the engine returns HTTP 404
|
||||||
|
- **THEN** the client SHALL display an error indicating the document was not found and exit with a non-zero code
|
||||||
|
|
||||||
|
#### Scenario: Update a non-note document
|
||||||
|
- **WHEN** the user runs `kb updatenote 42 "text"` and the engine returns HTTP 422
|
||||||
|
- **THEN** the client SHALL display an error indicating that only notes can be updated and exit with a non-zero code
|
||||||
|
|
||||||
|
#### Scenario: Missing arguments
|
||||||
|
- **WHEN** the user runs `kb updatenote` or `kb updatenote 42` with insufficient arguments
|
||||||
|
- **THEN** the client SHALL display usage help indicating that both document ID and text are required
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
### Requirement: Engine version compatibility check
|
### Requirement: Engine version compatibility check
|
||||||
|
|
||||||
The client SHALL verify that the connected engine meets a minimum version requirement before executing any API command. The minimum required engine version SHALL be embedded in the client binary at build time. If the engine version is below the minimum, the client SHALL print an error message and exit with a non-zero code. There SHALL be no flag to skip or suppress this check.
|
The client SHALL verify that the connected engine meets a minimum version requirement before executing any API command. The minimum required engine version SHALL be embedded in the client binary at build time. If the engine version is below the minimum, the client SHALL print an error message and exit with a non-zero code. There SHALL be no flag to skip or suppress this check.
|
||||||
|
|
||||||
#### Scenario: Compatible engine version
|
#### Scenario: Compatible engine version
|
||||||
- **WHEN** the client connects to an engine reporting version `2.1.5` and `MinEngineVersion` is `2.1.0`
|
- **WHEN** the client connects to an engine reporting version `3.0.0` and `MinEngineVersion` is `3.0.0`
|
||||||
- **THEN** the client SHALL proceed with the command normally
|
- **THEN** the client SHALL proceed with the command normally
|
||||||
|
|
||||||
#### Scenario: Incompatible engine version
|
#### Scenario: Incompatible engine version
|
||||||
- **WHEN** the client connects to an engine reporting version `2.0.3` and `MinEngineVersion` is `2.1.0`
|
- **WHEN** the client connects to an engine reporting version `2.1.0` and `MinEngineVersion` is `3.0.0`
|
||||||
- **THEN** the client SHALL print to stderr: `Error: kb client vX.Y.Z requires engine v2.1.0+ (connected engine is v2.0.3)` followed by an upgrade hint, and exit with code 1
|
- **THEN** the client SHALL print to stderr: `Error: kb client vX.Y.Z requires engine v3.0.0+ (connected engine is v2.1.0)` followed by an upgrade hint, and exit with code 1
|
||||||
|
|
||||||
#### Scenario: Engine unreachable during version check
|
#### Scenario: Engine unreachable during version check
|
||||||
- **WHEN** the client cannot reach the engine's `/api/v1/status` endpoint
|
- **WHEN** the client cannot reach the engine's `/api/v1/status` endpoint
|
||||||
|
|||||||
@@ -0,0 +1,198 @@
|
|||||||
|
# MCP Server
|
||||||
|
|
||||||
|
## Purpose
|
||||||
|
|
||||||
|
The MCP server provides a Model Context Protocol interface to the kb engine, exposing knowledge base operations as native MCP tools over Streamable HTTP transport. It runs as a separate Docker container alongside the engine, translating MCP tool calls into engine HTTP API calls.
|
||||||
|
|
||||||
|
## Requirements
|
||||||
|
|
||||||
|
### Requirement: MCP server transport and deployment
|
||||||
|
|
||||||
|
The MCP server SHALL expose tools via Streamable HTTP transport. It SHALL run as a Docker container, configured to connect to the kb engine's HTTP API. It SHALL read `KB_ENGINE_URL` and `KB_API_KEY` from environment variables to connect to the engine.
|
||||||
|
|
||||||
|
#### Scenario: MCP server starts and connects to engine
|
||||||
|
- **WHEN** the MCP server container starts with `KB_ENGINE_URL=http://engine:8000` and `KB_API_KEY=secret`
|
||||||
|
- **THEN** it SHALL begin accepting MCP connections over Streamable HTTP and use the configured URL and API key for all engine API calls
|
||||||
|
|
||||||
|
#### Scenario: Engine unreachable at startup
|
||||||
|
- **WHEN** the MCP server starts but cannot reach the engine at `KB_ENGINE_URL`
|
||||||
|
- **THEN** it SHALL start and accept connections, but tool calls SHALL return errors indicating the engine is unreachable
|
||||||
|
|
||||||
|
#### Scenario: Docker Compose deployment
|
||||||
|
- **WHEN** the MCP server is deployed via Docker Compose alongside the engine
|
||||||
|
- **THEN** it SHALL connect to the engine via the Docker network using the service name (e.g. `http://engine:8000`)
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
### Requirement: MCP server authentication
|
||||||
|
|
||||||
|
The MCP server SHALL require Bearer token authentication from calling agents via the `KB_MCP_API_KEY` environment variable. This is independent of the engine's `KB_API_KEY`.
|
||||||
|
|
||||||
|
#### Scenario: Valid MCP API key
|
||||||
|
- **WHEN** `KB_MCP_API_KEY` is set and a calling agent provides a matching Bearer token
|
||||||
|
- **THEN** the MCP server SHALL process the request normally
|
||||||
|
|
||||||
|
#### Scenario: Missing MCP API key when required
|
||||||
|
- **WHEN** `KB_MCP_API_KEY` is set and a calling agent connects without a Bearer token
|
||||||
|
- **THEN** the MCP server SHALL reject the connection with an authentication error
|
||||||
|
|
||||||
|
#### Scenario: Invalid MCP API key
|
||||||
|
- **WHEN** `KB_MCP_API_KEY` is set and a calling agent provides a non-matching Bearer token
|
||||||
|
- **THEN** the MCP server SHALL reject the connection with an authentication error
|
||||||
|
|
||||||
|
#### Scenario: MCP auth disabled
|
||||||
|
- **WHEN** `KB_MCP_API_KEY` is not set
|
||||||
|
- **THEN** the MCP server SHALL accept all connections without authentication
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
### Requirement: Search tool
|
||||||
|
|
||||||
|
The MCP server SHALL expose a `kb_search` tool that queries the knowledge base via the engine's search API.
|
||||||
|
|
||||||
|
#### Scenario: Basic search
|
||||||
|
- **WHEN** an agent calls `kb_search` with `{"query": "pension revaluation", "top": 5}`
|
||||||
|
- **THEN** the MCP server SHALL POST to the engine's `/api/v1/search` endpoint and return the results with chunk text, scores, document metadata, and tags
|
||||||
|
|
||||||
|
#### Scenario: Search with tag filter
|
||||||
|
- **WHEN** an agent calls `kb_search` with `{"query": "email preferences", "tags": ["agent:mybot"]}`
|
||||||
|
- **THEN** the MCP server SHALL include the tags in the filter and POST to the engine's search endpoint
|
||||||
|
|
||||||
|
#### Scenario: Search with mode override
|
||||||
|
- **WHEN** an agent calls `kb_search` with `{"query": "error log", "fts_only": true}`
|
||||||
|
- **THEN** the MCP server SHALL pass `fts_only: true` to the engine search endpoint
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
### Requirement: Add note tool
|
||||||
|
|
||||||
|
The MCP server SHALL expose a `kb_addnote` tool that submits a text note to the engine for ingestion.
|
||||||
|
|
||||||
|
#### Scenario: Add a note
|
||||||
|
- **WHEN** an agent calls `kb_addnote` with `{"text": "User prefers concise responses"}`
|
||||||
|
- **THEN** the MCP server SHALL submit the note to the engine's `POST /api/v1/jobs` endpoint and return the job ID
|
||||||
|
|
||||||
|
#### Scenario: Add a note with tags
|
||||||
|
- **WHEN** an agent calls `kb_addnote` with `{"text": "User prefers concise responses", "tags": ["agent:mybot", "feedback"]}`
|
||||||
|
- **THEN** the MCP server SHALL submit the note with exactly those tags to the engine
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
### Requirement: Chunked file upload tools
|
||||||
|
|
||||||
|
The MCP server SHALL expose a three-step chunked file upload pattern for transferring files from remote agents to the engine.
|
||||||
|
|
||||||
|
#### Scenario: Start an upload
|
||||||
|
- **WHEN** an agent calls `kb_upload_start` with `{"filename": "report.pdf", "total_size": 5242880, "tags": ["insurance"]}`
|
||||||
|
- **THEN** the MCP server SHALL create a staging entry, generate a UUID `upload_id`, and return `{"upload_id": "<uuid>"}`
|
||||||
|
|
||||||
|
#### Scenario: Upload a chunk
|
||||||
|
- **WHEN** an agent calls `kb_upload_chunk` with `{"upload_id": "<uuid>", "data": "<base64-encoded-data>", "chunk_index": 0}`
|
||||||
|
- **THEN** the MCP server SHALL decode the base64 data and write it to the staging area for the given upload
|
||||||
|
|
||||||
|
#### Scenario: Upload multiple chunks in sequence
|
||||||
|
- **WHEN** an agent calls `kb_upload_chunk` multiple times with sequential `chunk_index` values for the same `upload_id`
|
||||||
|
- **THEN** the MCP server SHALL store each chunk and track the sequence
|
||||||
|
|
||||||
|
#### Scenario: Finish an upload
|
||||||
|
- **WHEN** an agent calls `kb_upload_finish` with `{"upload_id": "<uuid>"}`
|
||||||
|
- **THEN** the MCP server SHALL reassemble the chunks in order, forward the complete file as a multipart upload to the engine's `POST /api/v1/jobs` endpoint with the tags from `kb_upload_start`, and return the job ID
|
||||||
|
|
||||||
|
#### Scenario: Upload with invalid upload_id
|
||||||
|
- **WHEN** an agent calls `kb_upload_chunk` or `kb_upload_finish` with an `upload_id` that does not exist
|
||||||
|
- **THEN** the MCP server SHALL return an error indicating the upload ID is not found
|
||||||
|
|
||||||
|
#### Scenario: Abandoned upload cleanup
|
||||||
|
- **WHEN** an agent starts an upload but does not call `kb_upload_finish` within 10 minutes
|
||||||
|
- **THEN** the MCP server SHALL clean up the staged chunks and remove the upload tracking entry
|
||||||
|
|
||||||
|
#### Scenario: MCP server restart during upload
|
||||||
|
- **WHEN** the MCP server container restarts while an upload is in progress
|
||||||
|
- **THEN** the in-progress upload SHALL be lost and the agent SHALL need to restart from `kb_upload_start`
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
### Requirement: Update note tool
|
||||||
|
|
||||||
|
The MCP server SHALL expose a `kb_update_note` tool that updates an existing note in place via the engine's note mutation endpoint.
|
||||||
|
|
||||||
|
#### Scenario: Update an existing note
|
||||||
|
- **WHEN** an agent calls `kb_update_note` with `{"document_id": 42, "text": "Updated preference: user prefers bullet points"}`
|
||||||
|
- **THEN** the MCP server SHALL send `PATCH /api/v1/notes/42` to the engine and return the updated document
|
||||||
|
|
||||||
|
#### Scenario: Update a non-existent document
|
||||||
|
- **WHEN** an agent calls `kb_update_note` with a `document_id` that does not exist
|
||||||
|
- **THEN** the MCP server SHALL return an error indicating the document was not found
|
||||||
|
|
||||||
|
#### Scenario: Update a non-note document
|
||||||
|
- **WHEN** an agent calls `kb_update_note` with a `document_id` that refers to a PDF
|
||||||
|
- **THEN** the MCP server SHALL return an error indicating that only notes can be updated
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
### Requirement: Get document tool
|
||||||
|
|
||||||
|
The MCP server SHALL expose a `kb_get` tool that retrieves document details from the engine.
|
||||||
|
|
||||||
|
#### Scenario: Get by document ID
|
||||||
|
- **WHEN** an agent calls `kb_get` with `{"document_id": 42}`
|
||||||
|
- **THEN** the MCP server SHALL fetch `GET /api/v1/documents/42` and return the document details with chunks
|
||||||
|
|
||||||
|
#### Scenario: Get by source path
|
||||||
|
- **WHEN** an agent calls `kb_get` with `{"source_path": "memory/feedback_testing.md"}`
|
||||||
|
- **THEN** the MCP server SHALL query the engine's documents endpoint filtered by source path and return matching documents
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
### Requirement: Status tool
|
||||||
|
|
||||||
|
The MCP server SHALL expose a `kb_status` tool that returns engine health and statistics.
|
||||||
|
|
||||||
|
#### Scenario: Get engine status
|
||||||
|
- **WHEN** an agent calls `kb_status` with no parameters
|
||||||
|
- **THEN** the MCP server SHALL fetch `GET /api/v1/status` and return engine version, model info, device info, document counts, and queue state
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
### Requirement: Jobs tool
|
||||||
|
|
||||||
|
The MCP server SHALL expose a `kb_jobs` tool that returns ingestion job status.
|
||||||
|
|
||||||
|
#### Scenario: List recent jobs
|
||||||
|
- **WHEN** an agent calls `kb_jobs` with no parameters
|
||||||
|
- **THEN** the MCP server SHALL fetch `GET /api/v1/jobs` and return the list of recent jobs
|
||||||
|
|
||||||
|
#### Scenario: Filter jobs by status
|
||||||
|
- **WHEN** an agent calls `kb_jobs` with `{"status": "failed"}`
|
||||||
|
- **THEN** the MCP server SHALL fetch `GET /api/v1/jobs?status=failed` and return matching jobs
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
### Requirement: Delete document tool
|
||||||
|
|
||||||
|
The MCP server SHALL expose a `kb_delete` tool that permanently deletes a document from the knowledge base. The tool SHALL accept a `document_id` (required integer). Deletion SHALL remove the document, its chunks, embeddings, tags, and any stored file on disk.
|
||||||
|
|
||||||
|
The tool SHALL return a confirmation response including the deleted document's ID and title.
|
||||||
|
|
||||||
|
#### Scenario: Successful deletion
|
||||||
|
- **WHEN** `kb_delete` is called with `document_id=42`
|
||||||
|
- **THEN** the document, its chunks, embeddings, tag associations, and stored file SHALL be deleted
|
||||||
|
- **AND** the response SHALL include `"status": "deleted"`, the `document_id`, and the document `title`
|
||||||
|
|
||||||
|
#### Scenario: Document not found
|
||||||
|
- **WHEN** `kb_delete` is called with a `document_id` that does not exist
|
||||||
|
- **THEN** the tool SHALL return an error response indicating the document was not found
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
### Requirement: Tags-only document organisation
|
||||||
|
|
||||||
|
The MCP server SHALL NOT maintain any collection abstraction. Documents SHALL be returned as-is from the engine with all tags visible. No tag stripping or collection field injection SHALL occur. Namespace isolation (e.g. separating agent memory from user documents) is achieved via tag conventions communicated through system prompts or tool descriptions.
|
||||||
|
|
||||||
|
#### Scenario: Search results show all tags
|
||||||
|
- **WHEN** `kb_search` is called and a result has tags `["agent:mybot", "collection:documents", "draft"]`
|
||||||
|
- **THEN** all three tags SHALL be returned as-is — no stripping of `collection:*` tags
|
||||||
|
|
||||||
|
#### Scenario: Add note with explicit tags only
|
||||||
|
- **WHEN** `kb_addnote(text="hello", tags=["agent:mybot", "memory"])` is called
|
||||||
|
- **THEN** the note SHALL be created with exactly those two tags — no default tags added
|
||||||
@@ -0,0 +1,43 @@
|
|||||||
|
# Note Mutation
|
||||||
|
|
||||||
|
## Purpose
|
||||||
|
|
||||||
|
Note mutation allows existing notes to be updated in place without requiring delete and re-add, preserving document identity (ID, creation timestamp) while updating content, embeddings, and the full-text index.
|
||||||
|
|
||||||
|
## Requirements
|
||||||
|
|
||||||
|
### Requirement: Note update endpoint
|
||||||
|
|
||||||
|
The engine SHALL provide a `PATCH /api/v1/notes/{id}` endpoint that accepts new text for an existing note, re-chunks and re-embeds it, and returns the updated document.
|
||||||
|
|
||||||
|
#### Scenario: Update an existing note
|
||||||
|
- **WHEN** a client sends `PATCH /api/v1/notes/42` with body `{"text": "Updated note content"}`
|
||||||
|
- **THEN** the engine SHALL delete existing chunks and embeddings for document 42, run the new text through the note chunking pipeline, generate embeddings for each chunk, insert new chunks and embeddings, update the document's `content_hash` and `updated_at`, and return the updated document with HTTP 200
|
||||||
|
|
||||||
|
#### Scenario: Update preserves document identity
|
||||||
|
- **WHEN** a note is updated via PATCH
|
||||||
|
- **THEN** the document SHALL retain its original `id` and `created_at` values, and `updated_at` SHALL be set to the current timestamp
|
||||||
|
|
||||||
|
#### Scenario: Update with long text that produces multiple chunks
|
||||||
|
- **WHEN** a client sends `PATCH /api/v1/notes/42` with text longer than the embedding model's token window
|
||||||
|
- **THEN** the engine SHALL chunk the text using the same note chunking pipeline as ingestion, producing multiple chunks, and embed each chunk separately
|
||||||
|
|
||||||
|
#### Scenario: Update a non-existent document
|
||||||
|
- **WHEN** a client sends `PATCH /api/v1/notes/999` and document 999 does not exist
|
||||||
|
- **THEN** the engine SHALL return HTTP 404
|
||||||
|
|
||||||
|
#### Scenario: Update a non-note document
|
||||||
|
- **WHEN** a client sends `PATCH /api/v1/notes/42` and document 42 has `doc_type = 'pdf'`
|
||||||
|
- **THEN** the engine SHALL return HTTP 422 with an error indicating that only notes can be updated via this endpoint
|
||||||
|
|
||||||
|
#### Scenario: Embedding failure during update
|
||||||
|
- **WHEN** a client sends `PATCH /api/v1/notes/42` but the embedding step fails
|
||||||
|
- **THEN** the engine SHALL roll back the entire transaction, preserving the original note content, chunks, and embeddings, and return HTTP 500
|
||||||
|
|
||||||
|
#### Scenario: FTS5 index updated on note mutation
|
||||||
|
- **WHEN** a note is updated via PATCH
|
||||||
|
- **THEN** the FTS5 virtual table SHALL be updated via the existing chunk triggers (`chunks_ad` for deletes, `chunks_ai` for inserts), keeping the full-text index consistent with the new content
|
||||||
|
|
||||||
|
#### Scenario: Tags preserved on update
|
||||||
|
- **WHEN** a note with tags `["feedback", "collection:memory"]` is updated via PATCH
|
||||||
|
- **THEN** the document's tags SHALL be unchanged — only the text content, chunks, and embeddings are replaced
|
||||||
+70
-20
@@ -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
|
||||||
@@ -151,15 +210,12 @@ fi
|
|||||||
echo "==> Building Docker engine images ($VERSION)"
|
echo "==> Building Docker engine images ($VERSION)"
|
||||||
|
|
||||||
NVIDIA_IMAGE="${IMAGE_BASE}/engine:${DOCKER_TAG}-nvidia"
|
NVIDIA_IMAGE="${IMAGE_BASE}/engine:${DOCKER_TAG}-nvidia"
|
||||||
ROCM_IMAGE="${IMAGE_BASE}/engine:${DOCKER_TAG}-rocm"
|
|
||||||
CPU_IMAGE="${IMAGE_BASE}/engine:${DOCKER_TAG}-cpu"
|
CPU_IMAGE="${IMAGE_BASE}/engine:${DOCKER_TAG}-cpu"
|
||||||
NVIDIA_LATEST="${IMAGE_BASE}/engine:latest-nvidia"
|
NVIDIA_LATEST="${IMAGE_BASE}/engine:latest-nvidia"
|
||||||
ROCM_LATEST="${IMAGE_BASE}/engine:latest-rocm"
|
|
||||||
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 "$ROCM_IMAGE" -t "$ROCM_LATEST" -f "$ENGINE_DIR/Dockerfile.rocm" "$ENGINE_DIR"
|
run docker build --provenance=false --sbom=false -t "$CPU_IMAGE" -t "$CPU_LATEST" -f "$ENGINE_DIR/Dockerfile.cpu" "$ENGINE_DIR"
|
||||||
run docker build -t "$CPU_IMAGE" -t "$CPU_LATEST" -f "$ENGINE_DIR/Dockerfile.cpu" "$ENGINE_DIR"
|
|
||||||
|
|
||||||
echo ""
|
echo ""
|
||||||
|
|
||||||
@@ -174,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
|
||||||
@@ -207,9 +263,6 @@ RELEASE_NOTES="## Docker images
|
|||||||
# NVIDIA GPU
|
# NVIDIA GPU
|
||||||
docker pull ${NVIDIA_IMAGE}
|
docker pull ${NVIDIA_IMAGE}
|
||||||
|
|
||||||
# AMD GPU (ROCm)
|
|
||||||
docker pull ${ROCM_IMAGE}
|
|
||||||
|
|
||||||
# CPU only
|
# CPU only
|
||||||
docker pull ${CPU_IMAGE}
|
docker pull ${CPU_IMAGE}
|
||||||
\`\`\`
|
\`\`\`
|
||||||
@@ -239,16 +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 "$ROCM_IMAGE"
|
push_image "$CPU_IMAGE"
|
||||||
run docker push "$ROCM_LATEST"
|
push_image "$CPU_LATEST"
|
||||||
run docker push "$CPU_IMAGE"
|
|
||||||
run docker push "$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 ""
|
||||||
@@ -256,7 +307,6 @@ echo "==> Release $GIT_TAG complete!"
|
|||||||
echo ""
|
echo ""
|
||||||
echo " Images:"
|
echo " Images:"
|
||||||
echo " $NVIDIA_IMAGE"
|
echo " $NVIDIA_IMAGE"
|
||||||
echo " $ROCM_IMAGE"
|
|
||||||
echo " $CPU_IMAGE"
|
echo " $CPU_IMAGE"
|
||||||
if [[ -n "${MCP_IMAGE:-}" ]]; then
|
if [[ -n "${MCP_IMAGE:-}" ]]; then
|
||||||
echo " $MCP_IMAGE"
|
echo " $MCP_IMAGE"
|
||||||
|
|||||||
Reference in New Issue
Block a user