Compare commits
47 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 | |||
| 0c124c4ab7 | |||
| da5b8435bc | |||
| e39e00a2c0 | |||
| d078af9ad3 | |||
| b3dce188e1 | |||
| 0dc3065979 | |||
| e7136a4a20 | |||
| adeba21712 | |||
| 2d179af557 | |||
| a6bab5e55e | |||
| c5191df9c0 | |||
| afbe270181 | |||
| 9e957f1a9a | |||
| bbe6a5e909 | |||
| 743102aee4 | |||
| 0f3b3be59f | |||
| 2fa2ac1134 | |||
| b2176c36ea | |||
| 5f9946efc9 | |||
| ea3d5707e1 | |||
| 7f4decee26 | |||
| 528a09ca90 | |||
| b04823e67b | |||
| 6a4bce4659 | |||
| 4590c124ad |
@@ -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
|
||||
@@ -1,2 +1,9 @@
|
||||
examples/
|
||||
.claude/
|
||||
__pycache__/
|
||||
engine/data/
|
||||
|
||||
TMP/
|
||||
.env
|
||||
.venv/
|
||||
test_mcp_client.py
|
||||
|
||||
+139
@@ -0,0 +1,139 @@
|
||||
# Developer Guide
|
||||
|
||||
Instructions for building from source, releasing, and contributing to kb.
|
||||
|
||||
## Building from source
|
||||
|
||||
### Engine
|
||||
|
||||
```bash
|
||||
cd engine
|
||||
|
||||
# NVIDIA GPU
|
||||
KB_DATA_PATH=~/kb-data docker compose -f compose.nvidia.yaml up -d
|
||||
```
|
||||
|
||||
### Client
|
||||
|
||||
```bash
|
||||
cd client
|
||||
make build # produces ./kb binary
|
||||
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
|
||||
|
||||
Client and engine are versioned independently via `client/VERSION` and `engine/VERSION`. Each has its own release script and git tag prefix.
|
||||
|
||||
### Release client
|
||||
|
||||
```bash
|
||||
./release-client.sh --gitea # patch bump, release via Gitea
|
||||
./release-client.sh --github --minor # minor bump, release via GitHub
|
||||
./release-client.sh --gitea --no-increment # release current version as-is
|
||||
./release-client.sh --gitea --dry-run # preview without doing anything
|
||||
```
|
||||
|
||||
Creates tag `client-vX.Y.Z`, builds Go binaries for all platforms, and creates a Gitea/GitHub release with binaries attached.
|
||||
|
||||
The client embeds a `MinEngineVersion` (from `client/MIN_ENGINE_VERSION`) and will hard-fail if the connected engine is too old.
|
||||
|
||||
### Release engine
|
||||
|
||||
```bash
|
||||
./release-engine.sh --gitea # patch bump, release via Gitea
|
||||
./release-engine.sh --github --minor # minor bump, release via GitHub
|
||||
./release-engine.sh --gitea --no-increment # release current version as-is
|
||||
./release-engine.sh --gitea --dry-run # preview without doing anything
|
||||
```
|
||||
|
||||
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
|
||||
|
||||
```bash
|
||||
# Client
|
||||
kb --version
|
||||
|
||||
# Engine
|
||||
curl http://localhost:8000/api/v1/status | jq .version
|
||||
```
|
||||
|
||||
### Docker images
|
||||
|
||||
Images are pushed to `docker.dcglab.co.uk/public/kb/engine` with tags:
|
||||
|
||||
- `engine-v2.0.6-nvidia` / `engine-v2.0.6-cpu` — versioned
|
||||
- `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:
|
||||
|
||||
```bash
|
||||
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
|
||||
|
||||
All endpoints are under `/api/v1/`. Requires `Authorization: Bearer <key>` header when `KB_API_KEY` is set.
|
||||
|
||||
| Method | Endpoint | Description |
|
||||
|---|---|---|
|
||||
| `GET` | `/health` | Health check (bypasses auth) |
|
||||
| `POST` | `/search` | Hybrid search (JSON body) |
|
||||
| `POST` | `/jobs` | Upload file/note for ingestion (multipart, returns 202 or 409 if duplicate) |
|
||||
| `GET` | `/jobs` | List ingestion jobs |
|
||||
| `GET` | `/jobs/{id}` | Job details |
|
||||
| `GET` | `/documents` | List documents |
|
||||
| `GET` | `/documents/{id}` | Document details with chunks |
|
||||
| `GET` | `/documents/{id}/file` | Download original file |
|
||||
| `DELETE` | `/documents/{id}` | Remove a document (and stored file) |
|
||||
| `PUT` | `/documents/{id}/tags` | Add/remove tags |
|
||||
| `GET` | `/tags` | List all tags (with descriptions) |
|
||||
| `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` | `/bulk/delete` | Bulk delete documents by filter |
|
||||
| `POST` | `/bulk/tags` | Bulk add/remove tags by filter |
|
||||
| `POST` | `/bulk/set-tags` | Bulk replace tags by filter |
|
||||
@@ -0,0 +1,174 @@
|
||||
# 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. `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
|
||||
|
||||
The compose files include a `kb-mcp` service alongside the engine. Set `KB_MCP_API_KEY` to require Bearer token auth from connecting agents:
|
||||
|
||||
```bash
|
||||
KB_API_KEY=your-engine-key KB_MCP_API_KEY=your-agent-key \
|
||||
docker compose -f engine/compose.nvidia.yaml up -d
|
||||
```
|
||||
|
||||
Or run the MCP server standalone:
|
||||
|
||||
```bash
|
||||
docker run -d --name kb-mcp \
|
||||
-p 3000:3000 \
|
||||
-e KB_ENGINE_URL=http://your-engine-host:8000 \
|
||||
-e KB_API_KEY=your-engine-key \
|
||||
-e KB_MCP_API_KEY=your-agent-key \
|
||||
--restart unless-stopped \
|
||||
docker.dcglab.co.uk/public/kb/mcp:latest
|
||||
```
|
||||
|
||||
## MCP tools
|
||||
|
||||
| Tool | Description |
|
||||
|---|---|
|
||||
| `kb_search` | Hybrid semantic (vector) + full-text search with tag/type filters |
|
||||
| `kb_addnote` | Add a text note (queued for async ingestion) |
|
||||
| `kb_update_note` | Update an existing note in place |
|
||||
| `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_jobs` | Ingestion queue status |
|
||||
| `kb_upload_start` | Start a chunked file upload |
|
||||
| `kb_upload_chunk` | Upload a base64-encoded file chunk |
|
||||
| `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 |
|
||||
|
||||
## Organising with tags
|
||||
|
||||
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.
|
||||
|
||||
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.
|
||||
|
||||
## MCP server configuration
|
||||
|
||||
| Variable | Default | Description |
|
||||
|---|---|---|
|
||||
| `KB_ENGINE_URL` | `http://localhost:8000` | Engine API URL |
|
||||
| `KB_API_KEY` | (none) | Engine API key |
|
||||
| `KB_MCP_API_KEY` | (none) | Bearer token required from agents (disabled if unset) |
|
||||
| `KB_MCP_PORT` | `3000` | Port to listen on |
|
||||
|
||||
## Connecting AI coding tools
|
||||
|
||||
The kb MCP server uses **Streamable HTTP** transport at `http://your-host:3000/mcp`. Below are configuration examples for popular AI coding tools.
|
||||
|
||||
### Claude Code (CLI / Desktop / Web)
|
||||
|
||||
Add the server to your project or user settings:
|
||||
|
||||
```bash
|
||||
claude mcp add kb-server --transport http http://localhost:3000/mcp
|
||||
```
|
||||
|
||||
Or add it manually to `.claude/settings.json` (project) or `~/.claude/settings.json` (global):
|
||||
|
||||
```json
|
||||
{
|
||||
"mcpServers": {
|
||||
"kb-server": {
|
||||
"type": "http",
|
||||
"url": "http://localhost:3000/mcp",
|
||||
"headers": {
|
||||
"Authorization": "Bearer your-agent-key"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
### VS Code (GitHub Copilot)
|
||||
|
||||
Add to your `.vscode/settings.json` (workspace) or user settings:
|
||||
|
||||
```json
|
||||
{
|
||||
"mcp": {
|
||||
"servers": {
|
||||
"kb-server": {
|
||||
"type": "http",
|
||||
"url": "http://localhost:3000/mcp",
|
||||
"headers": {
|
||||
"Authorization": "Bearer your-agent-key"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
Or add to `.vscode/mcp.json` in your workspace:
|
||||
|
||||
```json
|
||||
{
|
||||
"servers": {
|
||||
"kb-server": {
|
||||
"type": "http",
|
||||
"url": "http://localhost:3000/mcp",
|
||||
"headers": {
|
||||
"Authorization": "Bearer your-agent-key"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
### Cursor
|
||||
|
||||
Add to `.cursor/mcp.json` in your project root:
|
||||
|
||||
```json
|
||||
{
|
||||
"mcpServers": {
|
||||
"kb-server": {
|
||||
"type": "streamable-http",
|
||||
"url": "http://localhost:3000/mcp",
|
||||
"headers": {
|
||||
"Authorization": "Bearer your-agent-key"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
### Windsurf
|
||||
|
||||
Add to `~/.codeium/windsurf/mcp_config.json`:
|
||||
|
||||
```json
|
||||
{
|
||||
"mcpServers": {
|
||||
"kb-server": {
|
||||
"serverUrl": "http://localhost:3000/mcp",
|
||||
"headers": {
|
||||
"Authorization": "Bearer your-agent-key"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
### JetBrains IDEs (IntelliJ, WebStorm, PyCharm, etc.)
|
||||
|
||||
Add to `.junie/mcp.json` in your project root, or configure via **Settings > Tools > AI Assistant > MCP Servers**:
|
||||
|
||||
```json
|
||||
{
|
||||
"servers": {
|
||||
"kb-server": {
|
||||
"type": "http",
|
||||
"url": "http://localhost:3000/mcp",
|
||||
"headers": {
|
||||
"Authorization": "Bearer your-agent-key"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
@@ -2,33 +2,62 @@
|
||||
|
||||
Personal knowledge base with hybrid search (full-text + semantic vector search).
|
||||
|
||||
v2 uses a client-server architecture: a **FastAPI engine** running in Docker (with GPU acceleration) and a lightweight **Go CLI client** that talks to it over HTTP.
|
||||
Client-server architecture: a **FastAPI engine** running in Docker (with optional GPU acceleration), a lightweight **Go CLI client**, and an **MCP server** for native agent integration.
|
||||
|
||||
## Architecture
|
||||
|
||||
```
|
||||
Go CLI (kb) ──HTTP──▶ FastAPI Engine (Docker) ──▶ SQLite + GPU
|
||||
▲
|
||||
MCP Agents ──MCP/HTTP──▶ MCP Server (Docker) ──┘
|
||||
```
|
||||
|
||||
- **Engine**: Keeps the embedding model warm in GPU memory. Handles search, ingestion, and document management via REST API. Runs in Docker with NVIDIA or AMD GPU 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.
|
||||
- **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.
|
||||
|
||||
## Quick start
|
||||
|
||||
### 1. Start the engine
|
||||
|
||||
**From pre-built images** (recommended):
|
||||
|
||||
```bash
|
||||
cd engine
|
||||
|
||||
# NVIDIA GPU
|
||||
KB_DATA_PATH=~/kb-data docker compose -f compose.nvidia.yaml up -d
|
||||
docker run -d --name kb-engine \
|
||||
--gpus all \
|
||||
-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/public/kb/engine:latest-nvidia
|
||||
|
||||
# AMD GPU (ROCm)
|
||||
KB_DATA_PATH=~/kb-data docker compose -f compose.rocm.yaml up -d
|
||||
# CPU only (no GPU required — smaller image)
|
||||
docker run -d --name kb-engine \
|
||||
-p 8000:8000 \
|
||||
-v ~/kb-data:/data \
|
||||
-e KB_MODEL=all-MiniLM-L6-v2 \
|
||||
-e KB_API_KEY=your-secret-key \
|
||||
--restart unless-stopped \
|
||||
docker.dcglab.co.uk/public/kb/engine:latest-cpu
|
||||
```
|
||||
|
||||
The engine will download the embedding model on first start (~90MB) and load it onto the GPU. Check readiness:
|
||||
Or use a compose file from the repo:
|
||||
|
||||
```bash
|
||||
# NVIDIA GPU
|
||||
KB_DATA_PATH=~/kb-data docker compose -f engine/compose.nvidia.yaml up -d
|
||||
|
||||
# CPU only
|
||||
KB_DATA_PATH=~/kb-data docker compose -f engine/compose.cpu.yaml up -d
|
||||
```
|
||||
|
||||
See [DEVELOPER.md](DEVELOPER.md) to run the engine from source.
|
||||
|
||||
The engine will download the embedding model on first start (~90MB) and load it into memory (GPU or CPU). Check readiness:
|
||||
|
||||
```bash
|
||||
curl http://localhost:8000/api/v1/health
|
||||
@@ -37,18 +66,32 @@ curl http://localhost:8000/api/v1/health
|
||||
|
||||
### 2. Install the client
|
||||
|
||||
Build from source:
|
||||
**From a release** (recommended):
|
||||
|
||||
Check [releases](https://gitea.dcglab.co.uk/steve/kb/releases) for the latest client tag, then:
|
||||
|
||||
```bash
|
||||
cd client
|
||||
make build # produces ./kb binary
|
||||
# Set the version tag
|
||||
TAG=client-v3.0.0
|
||||
|
||||
# Linux (amd64)
|
||||
curl -L -o kb https://gitea.dcglab.co.uk/steve/kb/releases/download/${TAG}/kb-linux-amd64
|
||||
|
||||
# Linux (arm64)
|
||||
curl -L -o kb https://gitea.dcglab.co.uk/steve/kb/releases/download/${TAG}/kb-linux-arm64
|
||||
|
||||
# macOS (Apple Silicon)
|
||||
curl -L -o kb https://gitea.dcglab.co.uk/steve/kb/releases/download/${TAG}/kb-darwin-arm64
|
||||
|
||||
# macOS (Intel)
|
||||
curl -L -o kb https://gitea.dcglab.co.uk/steve/kb/releases/download/${TAG}/kb-darwin-amd64
|
||||
|
||||
# Then install
|
||||
chmod +x kb
|
||||
sudo mv kb /usr/local/bin/
|
||||
```
|
||||
|
||||
Or cross-compile for all platforms:
|
||||
|
||||
```bash
|
||||
make all # produces dist/kb-{os}-{arch} binaries
|
||||
```
|
||||
See [DEVELOPER.md](DEVELOPER.md) to build the client from source.
|
||||
|
||||
### 3. Configure the client
|
||||
|
||||
@@ -65,10 +108,15 @@ Override via environment variables (`KB_ENGINE_URL`, `KB_API_KEY`) or CLI flags
|
||||
### 4. Use it
|
||||
|
||||
```bash
|
||||
# Add documents (async — uploads and exits immediately)
|
||||
kb add ~/docs/manual.pdf --tags admin
|
||||
kb add ~/notes/ --recursive
|
||||
kb add --note "Always restart nginx after config changes" --tags ops
|
||||
# Add notes
|
||||
kb addnote "Always restart nginx after config changes"
|
||||
kb addnote "Server room is building 3, floor 2" --tags ops
|
||||
kb addnote "Deploy checklist" --wait
|
||||
|
||||
# Add files (async by default; --wait blocks until ingestion finishes)
|
||||
kb addfile ~/docs/manual.pdf --tags admin
|
||||
kb addfile ~/notes/ --recursive
|
||||
kb addfile ~/docs/manual.pdf --wait
|
||||
|
||||
# Check ingestion progress
|
||||
kb jobs
|
||||
@@ -76,20 +124,33 @@ kb jobs
|
||||
# Search
|
||||
kb search "how to install git"
|
||||
kb search "deploy process" --tags ops --type pdf
|
||||
kb find "vehicle handbook" --type pdf
|
||||
|
||||
# Update a note in place
|
||||
kb updatenote 42 "revised note content"
|
||||
|
||||
# Manage
|
||||
kb list
|
||||
kb info 1
|
||||
kb list --title handbook
|
||||
kb list --filename M38T_PHEV
|
||||
kb info 1 --no-chunks
|
||||
kb tags
|
||||
kb tag 1 --add important
|
||||
kb export 1 -o manual.pdf # download original file
|
||||
kb remove 3 --yes
|
||||
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
|
||||
|
||||
- **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.
|
||||
- **Search**: Hybrid retrieval combining BM25 keyword scoring (FTS5) and vector similarity (sqlite-vec), merged via Reciprocal Rank Fusion. Sub-100ms with a warm model.
|
||||
- **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 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.
|
||||
|
||||
## Engine configuration
|
||||
@@ -100,12 +161,34 @@ The engine is configured via environment variables (set in the compose file or v
|
||||
|---|---|---|
|
||||
| `KB_DATA_DIR` | `/data` | Data directory inside the container (bind-mounted) |
|
||||
| `KB_MODEL` | `all-MiniLM-L6-v2` | HuggingFace embedding model name |
|
||||
| `KB_DEVICE` | `auto` | Embedding device: `auto`, `cpu`, or `cuda` |
|
||||
| `KB_INGEST_DEVICE` | `auto` | Docling layout detection device |
|
||||
| `KB_DEVICE` | `auto` | Embedding/search 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_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_DATA_PATH` | `./data` | Host path for bind mount (compose variable) |
|
||||
| `KB_HOST` | `0.0.0.0` | Host to bind to |
|
||||
| `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) |
|
||||
|
||||
### 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
|
||||
|
||||
@@ -119,77 +202,14 @@ rsync -a ~/kb-data/ user@target:/home/user/kb-data/
|
||||
KB_DATA_PATH=~/kb-data docker compose -f compose.nvidia.yaml up -d
|
||||
```
|
||||
|
||||
Data is GPU-vendor-agnostic — you can ingest on NVIDIA and serve from AMD (or vice versa) 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.
|
||||
|
||||
## API reference
|
||||
## MCP server (agent integration)
|
||||
|
||||
All endpoints are under `/api/v1/`. Requires `Authorization: Bearer <key>` header when `KB_API_KEY` is set.
|
||||
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.
|
||||
|
||||
| Method | Endpoint | Description |
|
||||
|---|---|---|
|
||||
| `GET` | `/health` | Health check (bypasses auth) |
|
||||
| `POST` | `/search` | Hybrid search (JSON body) |
|
||||
| `POST` | `/jobs` | Upload file/note for ingestion (multipart, returns 202 or 409 if duplicate) |
|
||||
| `GET` | `/jobs` | List ingestion jobs |
|
||||
| `GET` | `/jobs/{id}` | Job details |
|
||||
| `GET` | `/documents` | List documents |
|
||||
| `GET` | `/documents/{id}` | Document details with chunks |
|
||||
| `DELETE` | `/documents/{id}` | Remove a document |
|
||||
| `PUT` | `/documents/{id}/tags` | Add/remove tags |
|
||||
| `GET` | `/tags` | List all tags |
|
||||
| `GET` | `/status` | Engine status, GPU info, DB stats |
|
||||
| `POST` | `/reindex` | Re-embed all chunks |
|
||||
See **[MCP.md](MCP.md)** for full details — server setup, available tools, tag-based organisation, configuration, and client examples.
|
||||
|
||||
## Building and releasing
|
||||
## Agent skill
|
||||
|
||||
Versioning is managed via `client/VERSION` and `engine/VERSION` files. The release script bumps these, builds all artifacts, tags, and publishes in one step.
|
||||
|
||||
### Release
|
||||
|
||||
```bash
|
||||
./release.sh --gitea # patch bump (e.g. 2.0.0 → 2.0.1), release via Gitea
|
||||
./release.sh --github --minor # minor bump (e.g. 2.0.1 → 2.1.0), release via GitHub
|
||||
./release.sh --gitea --major # major bump (e.g. 2.1.0 → 3.0.0)
|
||||
./release.sh --gitea --no-increment # release current version as-is
|
||||
./release.sh --gitea --dry-run # preview without doing anything
|
||||
```
|
||||
|
||||
The script will:
|
||||
|
||||
1. Bump the version in both `client/VERSION` and `engine/VERSION` (unless `--no-increment`)
|
||||
2. Build Go client binaries for all platforms (linux/darwin/windows, amd64/arm64)
|
||||
3. Build Docker engine images for NVIDIA and ROCm
|
||||
4. Commit the version bump, create an annotated git tag, and push
|
||||
5. Create a release (with client binaries attached) via `tea` or `gh`
|
||||
6. Push Docker images to the registry
|
||||
|
||||
### Checking versions
|
||||
|
||||
```bash
|
||||
# Client
|
||||
kb --version
|
||||
|
||||
# Engine
|
||||
curl http://localhost:8000/api/v1/status | jq .version
|
||||
```
|
||||
|
||||
### Docker images
|
||||
|
||||
Images are pushed to `docker.dcglab.co.uk/dcg/kb/engine` with tags:
|
||||
|
||||
- `v2.1.0-nvidia` / `v2.1.0-rocm` — versioned
|
||||
- `latest-nvidia` / `latest-rocm` — latest release
|
||||
|
||||
Override the registry and org via environment variables:
|
||||
|
||||
```bash
|
||||
REGISTRY=ghcr.io IMAGE_ORG=myorg ./release.sh --github
|
||||
```
|
||||
|
||||
## Future: ROCm runtime migration
|
||||
|
||||
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.
|
||||
|
||||
## Claude Code skill
|
||||
|
||||
This tool is designed to be wrapped as a Claude Code skill. See `SKILL.md` for the skill definition.
|
||||
If you are restricted from using MCP server, or you just prefer to utilise Agent SKILLS, please also see `SKILL.md` for the skill definition.
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
# kb-search skill
|
||||
|
||||
Search the user's personal knowledge base containing PDFs, markdown documents, code snippets, and text notes.
|
||||
Search, manage, and add to the user's personal knowledge base containing PDFs, Word docs, HTML, markdown, code files, and text notes.
|
||||
|
||||
## When to use
|
||||
|
||||
@@ -8,10 +8,18 @@ Search the user's personal knowledge base containing PDFs, markdown documents, c
|
||||
- User explicitly says "check my notes", "search kb", "look in my knowledge base", "what do my docs say about..."
|
||||
- User references documents or notes they've previously stored
|
||||
- User asks "how do I..." style questions that their knowledge base likely covers
|
||||
- User wants to save a note, add a file, or manage their knowledge base
|
||||
|
||||
## Available commands
|
||||
## Adding notes
|
||||
|
||||
### Search (primary)
|
||||
```bash
|
||||
kb addnote "remember to update DNS records" # add a note
|
||||
kb addnote "server room is building 3, floor 2" --tags ops # add a tagged note
|
||||
```
|
||||
|
||||
The note text must be a single quoted argument.
|
||||
|
||||
## Search (primary use case)
|
||||
|
||||
```bash
|
||||
kb search "<query>" --top 10 --format json
|
||||
@@ -20,25 +28,117 @@ kb search "<query>" --top 10 --format json
|
||||
Returns JSON with ranked results combining full-text and semantic search.
|
||||
|
||||
**Flags:**
|
||||
- `--top N` — number of results (default: 10)
|
||||
- `-n, --top N` — number of results (default: 10)
|
||||
- `--tags tag1,tag2` — filter by tags (AND logic)
|
||||
- `--type pdf|markdown|code|note` — filter by document type
|
||||
- `--format json|human` — output format (always use json)
|
||||
- `--type pdf|markdown|code|note|data` — filter by document type
|
||||
- `--format json|human` — output format (always use json for parsing)
|
||||
- `--fts-only` — keyword search only (skip semantic)
|
||||
- `--vec-only` — semantic search only (skip keyword)
|
||||
- `--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)
|
||||
|
||||
### Other useful commands
|
||||
## Adding files
|
||||
|
||||
```bash
|
||||
kb list --format json # List all documents
|
||||
kb list --type pdf --format json # List only PDFs
|
||||
kb tags --format json # List tags with counts
|
||||
kb info <doc_id> --format json # Document details
|
||||
kb status --format json # DB stats
|
||||
kb addfile report.pdf # single file
|
||||
kb addfile report.pdf --tags admin,reference # with tags
|
||||
kb addfile ~/docs/ --recursive # directory (recursive)
|
||||
kb addfile ~/docs/ --recursive --tags reference # directory with tags
|
||||
```
|
||||
|
||||
## Output format (search)
|
||||
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:**
|
||||
- `--tags tag1,tag2` — tags (comma-separated)
|
||||
- `-r, --recursive` — recursively add directory contents
|
||||
|
||||
## Document management
|
||||
|
||||
```bash
|
||||
kb list --format json # list all documents
|
||||
kb list --type pdf --format json # filter by type
|
||||
kb list --tags admin --format json # filter by tags
|
||||
kb info <doc_id> --format json # document details with chunks
|
||||
kb export <doc_id> -o file.pdf # download original file
|
||||
kb remove <doc_id> # remove (prompts for confirmation)
|
||||
kb remove <doc_id> --yes # remove without confirmation
|
||||
```
|
||||
|
||||
## Tag management
|
||||
|
||||
```bash
|
||||
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> --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)
|
||||
|
||||
```bash
|
||||
kb jobs --format json # list recent jobs
|
||||
kb jobs --status failed --format json # filter by status
|
||||
kb jobs <job_id> --format json # job details
|
||||
```
|
||||
|
||||
## Examples
|
||||
|
||||
```bash
|
||||
kb examples # show common usage examples
|
||||
```
|
||||
|
||||
## Engine status and maintenance
|
||||
|
||||
```bash
|
||||
kb status --format json # engine status, GPU info, DB stats
|
||||
kb reindex --yes # re-embed all chunks (skip confirmation)
|
||||
```
|
||||
|
||||
## Global flags
|
||||
|
||||
All commands support:
|
||||
- `--format json|human` — output format (always use `json` for machine parsing)
|
||||
- `--engine <url>` — engine API URL (default: http://localhost:8000)
|
||||
- `--api-key <key>` — API key for authentication
|
||||
|
||||
## Search output format
|
||||
|
||||
```json
|
||||
{
|
||||
@@ -46,27 +146,26 @@ kb status --format json # DB stats
|
||||
"results": [
|
||||
{
|
||||
"chunk_id": 1423,
|
||||
"document_id": 87,
|
||||
"score": 0.031,
|
||||
"score_breakdown": {"fts": 0.016, "vector": 0.015},
|
||||
"text": "To install the latest version of git from source...",
|
||||
"source": {
|
||||
"document_id": 42,
|
||||
"title": "Git Admin Guide",
|
||||
"path": "/home/user/docs/git-admin.pdf",
|
||||
"type": "pdf",
|
||||
"page": 12,
|
||||
"chunk_index": 3,
|
||||
"total_chunks": 28,
|
||||
"tags": ["git", "admin"]
|
||||
}
|
||||
"chunk_metadata": {"page": 12},
|
||||
"title": "Git Admin Guide",
|
||||
"doc_type": "pdf",
|
||||
"source_path": "/home/user/docs/git-admin.pdf",
|
||||
"created_at": "2026-03-15T10:30:00",
|
||||
"tags": ["git", "admin"],
|
||||
"tag_contexts": {"admin": "System administration guides"}
|
||||
}
|
||||
],
|
||||
"total_matches": 47,
|
||||
"returned": 10
|
||||
"returned": 10,
|
||||
"reranked": true
|
||||
}
|
||||
```
|
||||
|
||||
## How to answer
|
||||
## How to answer search queries
|
||||
|
||||
1. Run `kb search "<query>" --top 10 --format json`
|
||||
2. Read the returned chunks
|
||||
@@ -93,7 +192,7 @@ Query 2: kb search "git merge explanation" --top 5 --format json
|
||||
Query 3: kb search "git rebase vs merge" --top 5 --format json
|
||||
```
|
||||
|
||||
## Filtering
|
||||
## Filtering tips
|
||||
|
||||
Use filters when the question implies a specific domain:
|
||||
|
||||
@@ -101,10 +200,37 @@ Use filters when the question implies a specific domain:
|
||||
- From a specific topic → `--tags <topic>`
|
||||
- Check available tags first: `kb tags --format json`
|
||||
|
||||
## Updating notes
|
||||
|
||||
```bash
|
||||
kb updatenote 42 "revised note content" # update note by ID
|
||||
```
|
||||
|
||||
Updates the text of an existing note in place, preserving its ID, creation timestamp, and tags. Re-chunks and re-embeds the new text.
|
||||
|
||||
## MCP server (agent integration)
|
||||
|
||||
For agent-to-agent integration, kb provides an MCP server alongside the CLI. The MCP server
|
||||
exposes the same operations as native MCP tools over Streamable HTTP transport, which agents
|
||||
can connect to directly without subprocess overhead.
|
||||
|
||||
**MCP tools:** `kb_search`, `kb_addnote`, `kb_update_note`, `kb_get`, `kb_delete`, `kb_status`,
|
||||
`kb_jobs`, `kb_upload_start`, `kb_upload_chunk`, `kb_upload_finish`, `kb_bulk_delete`,
|
||||
`kb_bulk_tags`, `kb_bulk_set_tags`.
|
||||
|
||||
Use tags to separate agent data from user documents (e.g. tag all agent notes with
|
||||
`agent:mybot` and filter by that tag when searching). This convention is communicated
|
||||
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
|
||||
`kb-mcp` service from the same compose file. Agents connect to it on port 3000 (default).
|
||||
|
||||
## Important notes
|
||||
|
||||
- Always use `--format json` for machine parsing
|
||||
- The `score` field is relative, not absolute — compare scores within a result set
|
||||
- `source.page` is only present for PDF documents
|
||||
- `source.section_header` is only present for markdown documents with headers
|
||||
- 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.section_header` is only present for markdown documents with headers
|
||||
- Results are already ranked by relevance (hybrid FTS + vector search)
|
||||
- Duplicate files are detected at upload time (HTTP 409) — the client handles this gracefully
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
3.3.0
|
||||
+2
-1
@@ -1,5 +1,6 @@
|
||||
VERSION ?= $(shell cat VERSION 2>/dev/null || echo "dev")
|
||||
LDFLAGS := -ldflags "-s -w -X github.com/kb-search/kb/cmd.Version=$(VERSION)"
|
||||
MIN_ENGINE_VERSION ?= $(shell cat MIN_ENGINE_VERSION 2>/dev/null || echo "dev")
|
||||
LDFLAGS := -ldflags "-s -w -X github.com/kb-search/kb/cmd.Version=$(VERSION) -X github.com/kb-search/kb/cmd.MinEngineVersion=$(MIN_ENGINE_VERSION)"
|
||||
|
||||
PLATFORMS := linux/amd64 linux/arm64 darwin/amd64 darwin/arm64 windows/amd64
|
||||
|
||||
|
||||
+1
-1
@@ -1 +1 @@
|
||||
2.0.4
|
||||
3.3.0
|
||||
|
||||
+76
-86
@@ -6,7 +6,9 @@ import (
|
||||
"net/http"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sort"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/kb-search/kb/internal/api"
|
||||
"github.com/kb-search/kb/internal/output"
|
||||
@@ -37,95 +39,35 @@ var supportedExts = map[string]bool{
|
||||
".py": true,
|
||||
".sh": true,
|
||||
".go": true,
|
||||
".json": true,
|
||||
".yaml": true,
|
||||
".yml": true,
|
||||
".toml": true,
|
||||
}
|
||||
|
||||
var addCmd = &cobra.Command{
|
||||
Use: "add <path>",
|
||||
Short: "Add a document or directory to the knowledge base",
|
||||
Args: cobra.MaximumNArgs(1),
|
||||
RunE: runAdd,
|
||||
var addfileCmd = &cobra.Command{
|
||||
Use: "addfile <path>",
|
||||
Short: "Upload a file or directory to the knowledge base",
|
||||
Args: cobra.ExactArgs(1),
|
||||
RunE: runAddfile,
|
||||
}
|
||||
|
||||
func init() {
|
||||
addCmd.Flags().String("tags", "", "tags (comma-separated)")
|
||||
addCmd.Flags().String("type", "", "document type")
|
||||
addCmd.Flags().BoolP("recursive", "r", false, "recursively add directory contents")
|
||||
addCmd.Flags().String("note", "", "add a text note instead of a file")
|
||||
addCmd.Flags().String("title", "", "title for the note")
|
||||
rootCmd.AddCommand(addCmd)
|
||||
addfileCmd.Flags().String("tags", "", "tags (comma-separated)")
|
||||
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)
|
||||
}
|
||||
|
||||
func runAdd(cmd *cobra.Command, args []string) error {
|
||||
func runAddfile(cmd *cobra.Command, args []string) error {
|
||||
tags, _ := cmd.Flags().GetString("tags")
|
||||
docType, _ := cmd.Flags().GetString("type")
|
||||
recursive, _ := cmd.Flags().GetBool("recursive")
|
||||
note, _ := cmd.Flags().GetString("note")
|
||||
title, _ := cmd.Flags().GetString("title")
|
||||
wait, _ := cmd.Flags().GetBool("wait")
|
||||
timeout, _ := cmd.Flags().GetDuration("wait-timeout")
|
||||
|
||||
client := api.NewClient()
|
||||
|
||||
// Note mode
|
||||
if note != "" {
|
||||
fields := map[string]string{
|
||||
"note": note,
|
||||
}
|
||||
if title != "" {
|
||||
fields["title"] = title
|
||||
}
|
||||
if tags != "" {
|
||||
fields["tags"] = tags
|
||||
}
|
||||
if docType != "" {
|
||||
fields["type"] = docType
|
||||
}
|
||||
|
||||
resp, err := client.PostMultipart("/api/v1/jobs", fields, nil)
|
||||
if err != nil {
|
||||
fmt.Fprintln(os.Stderr, err)
|
||||
os.Exit(1)
|
||||
}
|
||||
|
||||
if resp.StatusCode == http.StatusConflict {
|
||||
var result 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 {
|
||||
if m, ok := result.(map[string]interface{}); ok {
|
||||
if docID, ok := m["document_id"].(float64); ok {
|
||||
fmt.Printf("Already imported: %s (doc ID: %.0f)\n", m["title"], docID)
|
||||
} else if jobID, ok := m["job_id"].(float64); ok {
|
||||
fmt.Printf("Already queued: %s (job ID: %.0f)\n", m["title"], jobID)
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
if err := api.CheckError(resp); err != nil {
|
||||
fmt.Fprintln(os.Stderr, err)
|
||||
os.Exit(1)
|
||||
}
|
||||
|
||||
var result 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 {
|
||||
fmt.Println("Queued: note")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
if len(args) == 0 {
|
||||
return fmt.Errorf("path argument is required (or use --note)")
|
||||
}
|
||||
|
||||
path := args[0]
|
||||
info, err := os.Stat(path)
|
||||
if err != nil {
|
||||
@@ -133,20 +75,44 @@ func runAdd(cmd *cobra.Command, args []string) error {
|
||||
}
|
||||
|
||||
if !info.IsDir() {
|
||||
// Validate extension
|
||||
ext := strings.ToLower(filepath.Ext(path))
|
||||
if !supportedExts[ext] {
|
||||
supported := make([]string, 0, len(supportedExts))
|
||||
for e := range supportedExts {
|
||||
supported = append(supported, e)
|
||||
}
|
||||
sort.Strings(supported)
|
||||
return fmt.Errorf("unsupported file type %q — supported: %s", ext, strings.Join(supported, ", "))
|
||||
}
|
||||
|
||||
// Single file upload
|
||||
result, err := uploadFile(client, path, tags, docType)
|
||||
result, err := uploadFile(client, path, tags)
|
||||
if err != nil {
|
||||
fmt.Fprintln(os.Stderr, err)
|
||||
os.Exit(1)
|
||||
}
|
||||
|
||||
if output.IsJSON() {
|
||||
if !wait || result.Duplicate {
|
||||
output.PrintJSON([]interface{}{result.Raw})
|
||||
}
|
||||
} else if result.Duplicate {
|
||||
fmt.Println(result.duplicateMsg())
|
||||
} else {
|
||||
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
|
||||
}
|
||||
|
||||
@@ -174,15 +140,18 @@ func runAdd(cmd *cobra.Command, args []string) error {
|
||||
}
|
||||
|
||||
var results []interface{}
|
||||
var pending []*uploadResult
|
||||
queued := 0
|
||||
duplicates := 0
|
||||
for _, f := range files {
|
||||
result, err := uploadFile(client, f, tags, docType)
|
||||
result, err := uploadFile(client, f, tags)
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "Error uploading %s: %v\n", f, err)
|
||||
continue
|
||||
}
|
||||
if !wait || result.Duplicate {
|
||||
results = append(results, result.Raw)
|
||||
}
|
||||
if result.Duplicate {
|
||||
duplicates++
|
||||
if !output.IsJSON() {
|
||||
@@ -190,11 +159,25 @@ func runAdd(cmd *cobra.Command, args []string) error {
|
||||
}
|
||||
} else {
|
||||
queued++
|
||||
pending = append(pending, result)
|
||||
if !output.IsJSON() {
|
||||
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() {
|
||||
output.PrintJSON(results)
|
||||
@@ -206,7 +189,7 @@ func runAdd(cmd *cobra.Command, args []string) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func uploadFile(client *api.Client, path, tags, docType string) (*uploadResult, error) {
|
||||
func uploadFile(client *api.Client, path, tags string) (*uploadResult, error) {
|
||||
f, err := os.Open(path)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("cannot open %s: %w", path, err)
|
||||
@@ -217,9 +200,6 @@ func uploadFile(client *api.Client, path, tags, docType string) (*uploadResult,
|
||||
if tags != "" {
|
||||
fields["tags"] = tags
|
||||
}
|
||||
if docType != "" {
|
||||
fields["type"] = docType
|
||||
}
|
||||
|
||||
upload := &api.FileUpload{
|
||||
FieldName: "file",
|
||||
@@ -258,9 +238,19 @@ func uploadFile(client *api.Client, path, tags, docType string) (*uploadResult,
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var result interface{}
|
||||
if err := api.DecodeJSON(resp, &result); err != nil {
|
||||
var raw json.RawMessage
|
||||
if err := api.DecodeJSON(resp, &raw); err != nil {
|
||||
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
|
||||
}
|
||||
|
||||
@@ -0,0 +1,110 @@
|
||||
package cmd
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
"os"
|
||||
"time"
|
||||
|
||||
"github.com/kb-search/kb/internal/api"
|
||||
"github.com/kb-search/kb/internal/output"
|
||||
"github.com/spf13/cobra"
|
||||
)
|
||||
|
||||
var addnoteCmd = &cobra.Command{
|
||||
Use: "addnote <text>",
|
||||
Short: "Add a text note to the knowledge base",
|
||||
Args: func(cmd *cobra.Command, args []string) error {
|
||||
if len(args) == 0 {
|
||||
return fmt.Errorf("requires a note text argument\n\n Usage: kb addnote \"your note text here\"")
|
||||
}
|
||||
if len(args) > 1 {
|
||||
return fmt.Errorf("accepts 1 arg but received %d — quote your note text, e.g. kb addnote \"your note text here\"", len(args))
|
||||
}
|
||||
return nil
|
||||
},
|
||||
RunE: runAddnote,
|
||||
}
|
||||
|
||||
func init() {
|
||||
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)
|
||||
}
|
||||
|
||||
func runAddnote(cmd *cobra.Command, args []string) error {
|
||||
tags, _ := cmd.Flags().GetString("tags")
|
||||
wait, _ := cmd.Flags().GetBool("wait")
|
||||
timeout, _ := cmd.Flags().GetDuration("wait-timeout")
|
||||
client := api.NewClient()
|
||||
return submitNote(client, args[0], tags, wait, timeout)
|
||||
}
|
||||
|
||||
func submitNote(client *api.Client, note, tags string, wait bool, timeout time.Duration) error {
|
||||
fields := map[string]string{
|
||||
"note": note,
|
||||
}
|
||||
if tags != "" {
|
||||
fields["tags"] = tags
|
||||
}
|
||||
|
||||
resp, err := client.PostMultipart("/api/v1/jobs", fields, nil)
|
||||
if err != nil {
|
||||
fmt.Fprintln(os.Stderr, err)
|
||||
os.Exit(1)
|
||||
}
|
||||
|
||||
if resp.StatusCode == http.StatusConflict {
|
||||
var result 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 {
|
||||
if m, ok := result.(map[string]interface{}); ok {
|
||||
if docID, ok := m["document_id"].(float64); ok {
|
||||
fmt.Printf("Already imported: %s (doc ID: %.0f)\n", m["title"], docID)
|
||||
} else if jobID, ok := m["job_id"].(float64); ok {
|
||||
fmt.Printf("Already queued: %s (job ID: %.0f)\n", m["title"], jobID)
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
if err := api.CheckError(resp); err != nil {
|
||||
fmt.Fprintln(os.Stderr, err)
|
||||
os.Exit(1)
|
||||
}
|
||||
|
||||
var result struct {
|
||||
JobID int `json:"job_id"`
|
||||
Status string `json:"status"`
|
||||
Filename string `json:"filename"`
|
||||
}
|
||||
if err := api.DecodeJSON(resp, &result); err != nil {
|
||||
return fmt.Errorf("failed to decode response: %w", err)
|
||||
}
|
||||
|
||||
if output.IsJSON() {
|
||||
if !wait {
|
||||
output.PrintJSON(result)
|
||||
}
|
||||
} else {
|
||||
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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -0,0 +1,44 @@
|
||||
package cmd
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
|
||||
"github.com/spf13/cobra"
|
||||
)
|
||||
|
||||
var examplesCmd = &cobra.Command{
|
||||
Use: "examples",
|
||||
Short: "Show common usage examples",
|
||||
Args: cobra.NoArgs,
|
||||
Run: func(cmd *cobra.Command, args []string) {
|
||||
fmt.Print(`Add notes:
|
||||
kb addnote "Remember to update DNS records"
|
||||
kb addnote "Server room is building 3" --tags ops
|
||||
kb addnote "Deploy checklist" --wait
|
||||
|
||||
Add files:
|
||||
kb addfile report.pdf
|
||||
kb addfile ~/docs/ --recursive --tags reference
|
||||
kb addfile report.pdf --wait
|
||||
|
||||
Search:
|
||||
kb search "how to restart nginx"
|
||||
kb search "deploy" --tags ops --top 5
|
||||
kb find "quarterly report" --type pdf
|
||||
|
||||
Update notes:
|
||||
kb updatenote 42 "revised note content"
|
||||
|
||||
Manage documents:
|
||||
kb list --type pdf
|
||||
kb list --filename report.pdf
|
||||
kb info 3 --no-chunks
|
||||
kb tag 3 --add important,ops
|
||||
kb remove 3 --yes
|
||||
`)
|
||||
},
|
||||
}
|
||||
|
||||
func init() {
|
||||
rootCmd.AddCommand(examplesCmd)
|
||||
}
|
||||
@@ -0,0 +1,74 @@
|
||||
package cmd
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"io"
|
||||
"mime"
|
||||
"os"
|
||||
"path/filepath"
|
||||
|
||||
"github.com/kb-search/kb/internal/api"
|
||||
"github.com/spf13/cobra"
|
||||
)
|
||||
|
||||
var exportCmd = &cobra.Command{
|
||||
Use: "export <id>",
|
||||
Short: "Download original document file",
|
||||
Args: cobra.ExactArgs(1),
|
||||
RunE: runExport,
|
||||
}
|
||||
|
||||
func init() {
|
||||
exportCmd.Flags().StringP("output", "o", "", "output file path (default: original filename to current directory)")
|
||||
rootCmd.AddCommand(exportCmd)
|
||||
}
|
||||
|
||||
func runExport(cmd *cobra.Command, args []string) error {
|
||||
client := api.NewClient()
|
||||
resp, err := client.Get("/api/v1/documents/" + args[0] + "/file")
|
||||
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)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
outPath, _ := cmd.Flags().GetString("output")
|
||||
|
||||
if outPath == "" {
|
||||
// Try to get filename from Content-Disposition header
|
||||
cd := resp.Header.Get("Content-Disposition")
|
||||
if cd != "" {
|
||||
_, params, err := mime.ParseMediaType(cd)
|
||||
if err == nil && params["filename"] != "" {
|
||||
outPath = params["filename"]
|
||||
}
|
||||
}
|
||||
if outPath == "" {
|
||||
outPath = "document-" + args[0]
|
||||
}
|
||||
}
|
||||
|
||||
if outPath == "-" {
|
||||
_, err := io.Copy(os.Stdout, resp.Body)
|
||||
return err
|
||||
}
|
||||
|
||||
outPath = filepath.Clean(outPath)
|
||||
f, err := os.Create(outPath)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to create output file: %w", err)
|
||||
}
|
||||
defer f.Close()
|
||||
|
||||
n, err := io.Copy(f, resp.Body)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to write file: %w", err)
|
||||
}
|
||||
|
||||
fmt.Fprintf(os.Stderr, "Saved %s (%d bytes)\n", outPath, n)
|
||||
return nil
|
||||
}
|
||||
@@ -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() {
|
||||
infoCmd.Flags().Bool("no-chunks", false, "return document metadata without chunk details")
|
||||
rootCmd.AddCommand(infoCmd)
|
||||
}
|
||||
|
||||
func runInfo(cmd *cobra.Command, args []string) error {
|
||||
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 {
|
||||
fmt.Fprintln(os.Stderr, err)
|
||||
os.Exit(1)
|
||||
@@ -48,6 +54,7 @@ func runInfo(cmd *cobra.Command, args []string) error {
|
||||
Tags []string `json:"tags"`
|
||||
CreatedAt string `json:"created_at"`
|
||||
UpdatedAt string `json:"updated_at"`
|
||||
ChunkCount int `json:"chunk_count"`
|
||||
Chunks []struct {
|
||||
ID int `json:"id"`
|
||||
Page interface{} `json:"page"`
|
||||
@@ -65,7 +72,7 @@ func runInfo(cmd *cobra.Command, args []string) error {
|
||||
{"Tags", joinStrings(doc.Tags)},
|
||||
{"Created", doc.CreatedAt},
|
||||
{"Updated", doc.UpdatedAt},
|
||||
{"Chunks", fmt.Sprintf("%d", len(doc.Chunks))},
|
||||
{"Chunks", fmt.Sprintf("%d", doc.ChunkCount)},
|
||||
}
|
||||
output.PrintKeyValue(pairs)
|
||||
|
||||
|
||||
+14
-2
@@ -13,18 +13,23 @@ import (
|
||||
var listCmd = &cobra.Command{
|
||||
Use: "list",
|
||||
Short: "List documents in the knowledge base",
|
||||
Args: cobra.NoArgs,
|
||||
RunE: runList,
|
||||
}
|
||||
|
||||
func init() {
|
||||
listCmd.Flags().String("type", "", "filter by document type")
|
||||
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)
|
||||
}
|
||||
|
||||
func runList(cmd *cobra.Command, args []string) error {
|
||||
docType, _ := cmd.Flags().GetString("type")
|
||||
tags, _ := cmd.Flags().GetString("tags")
|
||||
title, _ := cmd.Flags().GetString("title")
|
||||
filename, _ := cmd.Flags().GetString("filename")
|
||||
|
||||
params := url.Values{}
|
||||
if docType != "" {
|
||||
@@ -33,6 +38,12 @@ func runList(cmd *cobra.Command, args []string) error {
|
||||
if tags != "" {
|
||||
params.Set("tags", tags)
|
||||
}
|
||||
if title != "" {
|
||||
params.Set("title", title)
|
||||
}
|
||||
if filename != "" {
|
||||
params.Set("filename", filename)
|
||||
}
|
||||
|
||||
path := "/api/v1/documents"
|
||||
if len(params) > 0 {
|
||||
@@ -62,6 +73,7 @@ func runList(cmd *cobra.Command, args []string) error {
|
||||
var docs []struct {
|
||||
ID int `json:"id"`
|
||||
Title string `json:"title"`
|
||||
Filename string `json:"original_filename"`
|
||||
Type string `json:"doc_type"`
|
||||
Tags []string `json:"tags"`
|
||||
}
|
||||
@@ -74,10 +86,10 @@ func runList(cmd *cobra.Command, args []string) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
headers := []string{"ID", "TITLE", "TYPE", "TAGS"}
|
||||
headers := []string{"ID", "TITLE", "FILENAME", "TYPE", "TAGS"}
|
||||
var rows [][]string
|
||||
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)
|
||||
return nil
|
||||
|
||||
@@ -0,0 +1,83 @@
|
||||
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 reindexCmd = &cobra.Command{
|
||||
Use: "reindex",
|
||||
Short: "Re-embed all chunks with the current engine model",
|
||||
Args: cobra.NoArgs,
|
||||
RunE: runReindex,
|
||||
}
|
||||
|
||||
func init() {
|
||||
reindexCmd.Flags().BoolP("yes", "y", false, "skip confirmation prompt")
|
||||
rootCmd.AddCommand(reindexCmd)
|
||||
}
|
||||
|
||||
func runReindex(cmd *cobra.Command, args []string) error {
|
||||
yes, _ := cmd.Flags().GetBool("yes")
|
||||
|
||||
client := api.NewClient()
|
||||
|
||||
if !yes {
|
||||
// Fetch model name from engine status
|
||||
modelName := "current"
|
||||
statusResp, err := client.Get("/api/v1/status")
|
||||
if err == nil && api.CheckError(statusResp) == nil {
|
||||
var status struct {
|
||||
ModelName string `json:"model_name"`
|
||||
}
|
||||
if api.DecodeJSON(statusResp, &status) == nil && status.ModelName != "" {
|
||||
modelName = status.ModelName
|
||||
}
|
||||
}
|
||||
|
||||
fmt.Printf("Reindex all chunks? This will re-embed everything with the %s model. [y/N] ", modelName)
|
||||
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
|
||||
}
|
||||
}
|
||||
resp, err := client.Post("/api/v1/reindex", nil)
|
||||
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 struct {
|
||||
ChunksReindexed int `json:"chunks_reindexed"`
|
||||
Model string `json:"model"`
|
||||
}
|
||||
|
||||
if output.IsJSON() {
|
||||
var raw interface{}
|
||||
if err := api.DecodeJSON(resp, &raw); err != nil {
|
||||
fmt.Fprintln(os.Stderr, "Failed to parse response:", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
output.PrintJSON(raw)
|
||||
} else {
|
||||
if err := api.DecodeJSON(resp, &result); err != nil {
|
||||
fmt.Fprintln(os.Stderr, "Failed to parse response:", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
fmt.Printf("Reindexed %d chunks (model: %s)\n", result.ChunksReindexed, result.Model)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
+7
-2
@@ -4,6 +4,7 @@ import (
|
||||
"fmt"
|
||||
"os"
|
||||
|
||||
"github.com/kb-search/kb/internal/api"
|
||||
"github.com/kb-search/kb/internal/config"
|
||||
"github.com/spf13/cobra"
|
||||
)
|
||||
@@ -11,6 +12,9 @@ import (
|
||||
// Version is set at build time via -ldflags.
|
||||
var Version = "dev"
|
||||
|
||||
// MinEngineVersion is set at build time via -ldflags.
|
||||
var MinEngineVersion = "dev"
|
||||
|
||||
var (
|
||||
flagEngine string
|
||||
flagFormat string
|
||||
@@ -18,9 +22,9 @@ var (
|
||||
)
|
||||
|
||||
var rootCmd = &cobra.Command{
|
||||
Use: "kb",
|
||||
Use: "kb [command]",
|
||||
Short: "kb-search CLI client",
|
||||
Long: "A CLI client for the kb-search v2 engine API.",
|
||||
Long: "A CLI client for the kb-search v2 engine API.\nRun 'kb examples' for common usage patterns.",
|
||||
PersistentPreRunE: func(cmd *cobra.Command, args []string) error {
|
||||
if err := config.Load(); err != nil {
|
||||
return err
|
||||
@@ -31,6 +35,7 @@ var rootCmd = &cobra.Command{
|
||||
}
|
||||
|
||||
func init() {
|
||||
api.SetVersionInfo(Version, MinEngineVersion)
|
||||
rootCmd.Version = Version
|
||||
rootCmd.PersistentFlags().StringVar(&flagEngine, "engine", "", "engine API URL")
|
||||
rootCmd.PersistentFlags().StringVar(&flagFormat, "format", "", "output format (human|json)")
|
||||
|
||||
@@ -0,0 +1,82 @@
|
||||
package cmd
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestRootCmd_NoArgs_ShowsHelp(t *testing.T) {
|
||||
rootCmd.SetArgs([]string{})
|
||||
|
||||
var stdout bytes.Buffer
|
||||
rootCmd.SetOut(&stdout)
|
||||
|
||||
err := rootCmd.Execute()
|
||||
if err != nil {
|
||||
t.Fatalf("expected no error for zero args, got: %v", err)
|
||||
}
|
||||
|
||||
output := stdout.String()
|
||||
if !strings.Contains(output, "Available Commands") {
|
||||
t.Errorf("expected help output, got: %s", output)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRootCmd_UnknownCommand_ReturnsError(t *testing.T) {
|
||||
rootCmd.SetArgs([]string{"notacommand"})
|
||||
|
||||
var stderr bytes.Buffer
|
||||
rootCmd.SetErr(&stderr)
|
||||
|
||||
err := rootCmd.Execute()
|
||||
if err == nil {
|
||||
t.Fatal("expected error for unknown command, got nil")
|
||||
}
|
||||
|
||||
errMsg := err.Error()
|
||||
if !strings.Contains(errMsg, "unknown command") {
|
||||
t.Errorf("expected 'unknown command' error, got: %s", errMsg)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAddnoteCmd_NoArgs_ReturnsError(t *testing.T) {
|
||||
rootCmd.SetArgs([]string{"addnote"})
|
||||
|
||||
err := rootCmd.Execute()
|
||||
if err == nil {
|
||||
t.Fatal("expected error for addnote with no args, got nil")
|
||||
}
|
||||
|
||||
errMsg := err.Error()
|
||||
if !strings.Contains(errMsg, "requires a note text argument") {
|
||||
t.Errorf("expected 'requires a note text argument' error, got: %s", errMsg)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAddnoteCmd_TooManyArgs_ReturnsError(t *testing.T) {
|
||||
rootCmd.SetArgs([]string{"addnote", "hello", "world"})
|
||||
|
||||
err := rootCmd.Execute()
|
||||
if err == nil {
|
||||
t.Fatal("expected error for addnote with too many args, got nil")
|
||||
}
|
||||
|
||||
errMsg := err.Error()
|
||||
if !strings.Contains(errMsg, "quote your note text") {
|
||||
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)
|
||||
}
|
||||
}
|
||||
+84
-16
@@ -23,6 +23,8 @@ func init() {
|
||||
searchCmd.Flags().Bool("fts-only", false, "use full-text search only")
|
||||
searchCmd.Flags().Bool("vec-only", false, "use vector search only")
|
||||
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)
|
||||
}
|
||||
|
||||
@@ -33,16 +35,18 @@ func runSearch(cmd *cobra.Command, args []string) error {
|
||||
ftsOnly, _ := cmd.Flags().GetBool("fts-only")
|
||||
vecOnly, _ := cmd.Flags().GetBool("vec-only")
|
||||
threshold, _ := cmd.Flags().GetFloat64("threshold")
|
||||
explain, _ := cmd.Flags().GetBool("explain")
|
||||
noRerank, _ := cmd.Flags().GetBool("no-rerank")
|
||||
|
||||
body := map[string]interface{}{
|
||||
"query": args[0],
|
||||
"top": top,
|
||||
}
|
||||
if tags != "" {
|
||||
body["tags"] = tags
|
||||
body["tags"] = splitTags(tags)
|
||||
}
|
||||
if docType != "" {
|
||||
body["type"] = docType
|
||||
body["doc_type"] = docType
|
||||
}
|
||||
if ftsOnly {
|
||||
body["fts_only"] = true
|
||||
@@ -53,6 +57,12 @@ func runSearch(cmd *cobra.Command, args []string) error {
|
||||
if threshold > 0 {
|
||||
body["threshold"] = threshold
|
||||
}
|
||||
if explain {
|
||||
body["explain"] = true
|
||||
}
|
||||
if noRerank {
|
||||
body["rerank"] = false
|
||||
}
|
||||
|
||||
client := api.NewClient()
|
||||
resp, err := client.Post("/api/v1/search", body)
|
||||
@@ -66,16 +76,17 @@ func runSearch(cmd *cobra.Command, args []string) error {
|
||||
}
|
||||
|
||||
var result struct {
|
||||
Reranked bool `json:"reranked"`
|
||||
Results []struct {
|
||||
Score float64 `json:"score"`
|
||||
Document struct {
|
||||
DocumentID int64 `json:"document_id"`
|
||||
Title string `json:"title"`
|
||||
Type string `json:"doc_type"`
|
||||
DocType string `json:"doc_type"`
|
||||
Tags []string `json:"tags"`
|
||||
} `json:"document"`
|
||||
Page interface{} `json:"page"`
|
||||
Section string `json:"section"`
|
||||
TagContexts map[string]string `json:"tag_contexts"`
|
||||
ChunkMetadata map[string]interface{} `json:"chunk_metadata"`
|
||||
Text string `json:"text"`
|
||||
Explain map[string]interface{} `json:"explain"`
|
||||
} `json:"results"`
|
||||
}
|
||||
|
||||
@@ -97,32 +108,46 @@ func runSearch(cmd *cobra.Command, args []string) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
if result.Reranked {
|
||||
fmt.Println("(reranked)")
|
||||
}
|
||||
|
||||
for i, r := range result.Results {
|
||||
snippet := r.Text
|
||||
if len(snippet) > 200 {
|
||||
snippet = snippet[:200] + "..."
|
||||
}
|
||||
|
||||
fmt.Printf("\n%d. [%.4f] %s\n", i+1, r.Score, r.Document.Title)
|
||||
fmt.Printf("\n%d. [%.4f] %s (doc:%d)\n", i+1, r.Score, r.Title, r.DocumentID)
|
||||
|
||||
location := ""
|
||||
if r.Page != nil {
|
||||
location = fmt.Sprintf("Page %v", r.Page)
|
||||
if page, ok := r.ChunkMetadata["page"]; ok && page != nil {
|
||||
location = fmt.Sprintf("Page %v", page)
|
||||
}
|
||||
if r.Section != "" {
|
||||
if section, ok := r.ChunkMetadata["section_header"]; ok && section != nil {
|
||||
if s, ok := section.(string); ok && s != "" {
|
||||
if location != "" {
|
||||
location += " / "
|
||||
}
|
||||
location += r.Section
|
||||
location += s
|
||||
}
|
||||
}
|
||||
if location != "" {
|
||||
fmt.Printf(" Location: %s\n", location)
|
||||
}
|
||||
if r.Document.Type != "" {
|
||||
fmt.Printf(" Type: %s\n", r.Document.Type)
|
||||
if r.DocType != "" {
|
||||
fmt.Printf(" Type: %s\n", r.DocType)
|
||||
}
|
||||
if len(r.Document.Tags) > 0 {
|
||||
fmt.Printf(" Tags: %s\n", joinStrings(r.Document.Tags))
|
||||
if len(r.Tags) > 0 {
|
||||
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)
|
||||
}
|
||||
@@ -130,6 +155,49 @@ func runSearch(cmd *cobra.Command, args []string) error {
|
||||
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 {
|
||||
result := ""
|
||||
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 {
|
||||
Name string `json:"name"`
|
||||
Count int `json:"count"`
|
||||
Description string `json:"description"`
|
||||
}
|
||||
if err := api.DecodeJSON(resp, &tags); err != nil {
|
||||
return fmt.Errorf("failed to decode response: %w", err)
|
||||
@@ -53,10 +54,10 @@ func runTags(cmd *cobra.Command, args []string) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
headers := []string{"TAG", "COUNT"}
|
||||
headers := []string{"TAG", "COUNT", "DESCRIPTION"}
|
||||
var rows [][]string
|
||||
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)
|
||||
return nil
|
||||
|
||||
@@ -0,0 +1,61 @@
|
||||
package cmd
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"strconv"
|
||||
|
||||
"github.com/kb-search/kb/internal/api"
|
||||
"github.com/kb-search/kb/internal/output"
|
||||
"github.com/spf13/cobra"
|
||||
)
|
||||
|
||||
var updatenoteCmd = &cobra.Command{
|
||||
Use: "updatenote <id> <text>",
|
||||
Short: "Update an existing note's content",
|
||||
Args: func(cmd *cobra.Command, args []string) error {
|
||||
if len(args) < 2 {
|
||||
return fmt.Errorf("requires document ID and text arguments\n\n Usage: kb updatenote 42 \"updated note text\"")
|
||||
}
|
||||
if _, err := strconv.Atoi(args[0]); err != nil {
|
||||
return fmt.Errorf("document ID must be an integer, got %q", args[0])
|
||||
}
|
||||
return nil
|
||||
},
|
||||
RunE: runUpdatenote,
|
||||
}
|
||||
|
||||
func init() {
|
||||
rootCmd.AddCommand(updatenoteCmd)
|
||||
}
|
||||
|
||||
func runUpdatenote(cmd *cobra.Command, args []string) error {
|
||||
docID := args[0]
|
||||
text := args[1]
|
||||
|
||||
client := api.NewClient()
|
||||
|
||||
body := map[string]string{"text": text}
|
||||
resp, err := client.Patch(fmt.Sprintf("/api/v1/notes/%s", docID), 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 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 {
|
||||
fmt.Printf("Updated note %s\n", docID)
|
||||
}
|
||||
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)
|
||||
}
|
||||
}
|
||||
@@ -7,6 +7,9 @@ import (
|
||||
"io"
|
||||
"mime/multipart"
|
||||
"net/http"
|
||||
"os"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"github.com/kb-search/kb/internal/config"
|
||||
)
|
||||
@@ -18,11 +21,25 @@ type FileUpload struct {
|
||||
Reader io.Reader
|
||||
}
|
||||
|
||||
// Package-level version info, set once by cmd.init via SetVersionInfo.
|
||||
var (
|
||||
clientVersion string
|
||||
minEngineVersion string
|
||||
)
|
||||
|
||||
// SetVersionInfo configures the client and minimum engine version for compatibility checking.
|
||||
// Called once from cmd package initialization.
|
||||
func SetVersionInfo(cv, minEV string) {
|
||||
clientVersion = cv
|
||||
minEngineVersion = minEV
|
||||
}
|
||||
|
||||
// Client is an HTTP client for the kb-search engine API.
|
||||
type Client struct {
|
||||
baseURL string
|
||||
apiKey string
|
||||
httpClient *http.Client
|
||||
versionChecked bool
|
||||
}
|
||||
|
||||
// NewClient creates a Client from the current configuration.
|
||||
@@ -48,6 +65,7 @@ func (c *Client) newRequest(method, path string, body io.Reader) (*http.Request,
|
||||
}
|
||||
|
||||
func (c *Client) do(req *http.Request) (*http.Response, error) {
|
||||
c.checkEngineVersion()
|
||||
resp, err := c.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("Cannot reach engine at %s: %v", c.baseURL, err)
|
||||
@@ -55,6 +73,75 @@ func (c *Client) do(req *http.Request) (*http.Response, error) {
|
||||
return resp, nil
|
||||
}
|
||||
|
||||
func (c *Client) checkEngineVersion() {
|
||||
if c.versionChecked {
|
||||
return
|
||||
}
|
||||
c.versionChecked = true
|
||||
|
||||
minVer := minEngineVersion
|
||||
if minVer == "" || minVer == "dev" {
|
||||
return
|
||||
}
|
||||
|
||||
statusReq, err := c.newRequest(http.MethodGet, "/api/v1/status", nil)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
resp, err := c.httpClient.Do(statusReq)
|
||||
if err != nil {
|
||||
return // unreachable — let the actual request surface the error
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return // auth error or other issue — let the actual request surface it
|
||||
}
|
||||
|
||||
var status struct {
|
||||
Version string `json:"version"`
|
||||
}
|
||||
if err := json.NewDecoder(resp.Body).Decode(&status); err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
if !semverAtLeast(status.Version, minVer) {
|
||||
fmt.Fprintf(os.Stderr, "Error: kb client v%s requires engine v%s+ (connected engine is v%s)\nUpdate your engine image to engine-v%s or later.\n",
|
||||
clientVersion, minVer, status.Version, minVer)
|
||||
os.Exit(1)
|
||||
}
|
||||
}
|
||||
|
||||
// semverAtLeast returns true if version >= minimum, comparing major.minor.patch.
|
||||
func semverAtLeast(version, minimum string) bool {
|
||||
parse := func(s string) (int, int, int) {
|
||||
s = strings.TrimPrefix(s, "v")
|
||||
parts := strings.SplitN(s, ".", 3)
|
||||
var major, minor, patch int
|
||||
if len(parts) >= 1 {
|
||||
major, _ = strconv.Atoi(parts[0])
|
||||
}
|
||||
if len(parts) >= 2 {
|
||||
minor, _ = strconv.Atoi(parts[1])
|
||||
}
|
||||
if len(parts) >= 3 {
|
||||
patch, _ = strconv.Atoi(parts[2])
|
||||
}
|
||||
return major, minor, patch
|
||||
}
|
||||
|
||||
vMaj, vMin, vPat := parse(version)
|
||||
mMaj, mMin, mPat := parse(minimum)
|
||||
|
||||
if vMaj != mMaj {
|
||||
return vMaj > mMaj
|
||||
}
|
||||
if vMin != mMin {
|
||||
return vMin > mMin
|
||||
}
|
||||
return vPat >= mPat
|
||||
}
|
||||
|
||||
// Get performs a GET request to the given path.
|
||||
func (c *Client) Get(path string) (*http.Response, error) {
|
||||
req, err := c.newRequest(http.MethodGet, path, nil)
|
||||
@@ -134,6 +221,20 @@ func (c *Client) Put(path string, body interface{}) (*http.Response, error) {
|
||||
return c.do(req)
|
||||
}
|
||||
|
||||
// Patch performs a PATCH request with a JSON body.
|
||||
func (c *Client) Patch(path string, body interface{}) (*http.Response, error) {
|
||||
data, err := json.Marshal(body)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to marshal request body: %w", err)
|
||||
}
|
||||
req, err := c.newRequest(http.MethodPatch, path, bytes.NewReader(data))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
return c.do(req)
|
||||
}
|
||||
|
||||
// DecodeJSON reads the response body and decodes it into target.
|
||||
func DecodeJSON(resp *http.Response, target interface{}) error {
|
||||
defer resp.Body.Close()
|
||||
|
||||
@@ -0,0 +1,136 @@
|
||||
package api
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestSemverAtLeast(t *testing.T) {
|
||||
tests := []struct {
|
||||
version string
|
||||
minimum string
|
||||
expected bool
|
||||
}{
|
||||
{"2.1.0", "2.0.0", true},
|
||||
{"2.0.0", "2.0.0", true},
|
||||
{"2.0.5", "2.0.0", true},
|
||||
{"2.1.5", "2.1.0", true},
|
||||
{"2.0.9", "2.1.0", false},
|
||||
{"1.9.9", "2.0.0", false},
|
||||
{"3.0.0", "2.9.9", true},
|
||||
{"2.0.0", "2.0.1", false},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.version+">="+tt.minimum, func(t *testing.T) {
|
||||
got := semverAtLeast(tt.version, tt.minimum)
|
||||
if got != tt.expected {
|
||||
t.Errorf("semverAtLeast(%q, %q) = %v, want %v", tt.version, tt.minimum, got, tt.expected)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestCheckEngineVersion_Compatible(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
json.NewEncoder(w).Encode(map[string]string{"version": "2.1.0"})
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
clientVersion = "2.2.0"
|
||||
minEngineVersion = "2.1.0"
|
||||
defer func() { clientVersion = ""; minEngineVersion = "" }()
|
||||
|
||||
c := &Client{
|
||||
baseURL: srv.URL,
|
||||
httpClient: &http.Client{},
|
||||
}
|
||||
|
||||
// Should not panic or exit
|
||||
c.checkEngineVersion()
|
||||
|
||||
if !c.versionChecked {
|
||||
t.Error("versionChecked should be true after check")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCheckEngineVersion_SkipsWhenDev(t *testing.T) {
|
||||
clientVersion = "dev"
|
||||
minEngineVersion = "dev"
|
||||
defer func() { clientVersion = ""; minEngineVersion = "" }()
|
||||
|
||||
c := &Client{
|
||||
baseURL: "http://localhost:99999",
|
||||
httpClient: &http.Client{},
|
||||
}
|
||||
|
||||
// Should not attempt connection
|
||||
c.checkEngineVersion()
|
||||
|
||||
if !c.versionChecked {
|
||||
t.Error("versionChecked should be true after skipping")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCheckEngineVersion_SkipsWhenEmpty(t *testing.T) {
|
||||
clientVersion = "1.0.0"
|
||||
minEngineVersion = ""
|
||||
defer func() { clientVersion = ""; minEngineVersion = "" }()
|
||||
|
||||
c := &Client{
|
||||
baseURL: "http://localhost:99999",
|
||||
httpClient: &http.Client{},
|
||||
}
|
||||
|
||||
c.checkEngineVersion()
|
||||
|
||||
if !c.versionChecked {
|
||||
t.Error("versionChecked should be true after skipping")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCheckEngineVersion_SkipsWhenUnreachable(t *testing.T) {
|
||||
clientVersion = "2.0.0"
|
||||
minEngineVersion = "2.0.0"
|
||||
defer func() { clientVersion = ""; minEngineVersion = "" }()
|
||||
|
||||
c := &Client{
|
||||
baseURL: "http://localhost:99999",
|
||||
httpClient: &http.Client{},
|
||||
}
|
||||
|
||||
// Should not panic — just skip
|
||||
c.checkEngineVersion()
|
||||
|
||||
if !c.versionChecked {
|
||||
t.Error("versionChecked should be true even when unreachable")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCheckEngineVersion_CachedAfterFirstCall(t *testing.T) {
|
||||
callCount := 0
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
callCount++
|
||||
json.NewEncoder(w).Encode(map[string]string{"version": "2.1.0"})
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
clientVersion = "2.1.0"
|
||||
minEngineVersion = "2.0.0"
|
||||
defer func() { clientVersion = ""; minEngineVersion = "" }()
|
||||
|
||||
c := &Client{
|
||||
baseURL: srv.URL,
|
||||
httpClient: &http.Client{},
|
||||
}
|
||||
|
||||
c.checkEngineVersion()
|
||||
c.checkEngineVersion()
|
||||
c.checkEngineVersion()
|
||||
|
||||
if callCount != 1 {
|
||||
t.Errorf("expected 1 status call, got %d", callCount)
|
||||
}
|
||||
}
|
||||
@@ -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>
|
||||
@@ -0,0 +1,52 @@
|
||||
FROM ubuntu:24.04
|
||||
|
||||
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 \
|
||||
libgl1 libglib2.0-0 \
|
||||
build-essential curl \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
COPY --from=ghcr.io/astral-sh/uv:latest /uv /usr/local/bin/uv
|
||||
|
||||
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 kb/ kb/
|
||||
COPY main.py ./
|
||||
COPY VERSION ./
|
||||
|
||||
# Remaining dependencies resolve against the CPU torch already present.
|
||||
RUN . .venv/bin/activate && \
|
||||
uv pip install "sentence-transformers[onnx]" && \
|
||||
uv pip install -e .
|
||||
|
||||
ENV PATH="/app/.venv/bin:$PATH"
|
||||
ENV VIRTUAL_ENV="/app/.venv"
|
||||
ENV KB_DEVICE=cpu
|
||||
ENV KB_INGEST_DEVICE=cpu
|
||||
ENV KB_DATA_DIR=/data
|
||||
|
||||
EXPOSE 8000
|
||||
VOLUME ["/data"]
|
||||
|
||||
CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000"]
|
||||
@@ -13,15 +13,24 @@ COPY --from=ghcr.io/astral-sh/uv:latest /uv /usr/local/bin/uv
|
||||
|
||||
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 kb/ kb/
|
||||
COPY main.py ./
|
||||
COPY VERSION ./
|
||||
|
||||
RUN uv venv .venv && \
|
||||
. .venv/bin/activate && \
|
||||
uv pip install -e . && \
|
||||
uv pip install --no-deps onnxruntime-gpu
|
||||
RUN . .venv/bin/activate && \
|
||||
uv pip install -e .
|
||||
|
||||
ENV PATH="/app/.venv/bin:$PATH"
|
||||
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 @@
|
||||
2.0.4
|
||||
3.3.0
|
||||
|
||||
@@ -0,0 +1,34 @@
|
||||
services:
|
||||
kb-engine:
|
||||
build:
|
||||
context: .
|
||||
dockerfile: Dockerfile.cpu
|
||||
ports:
|
||||
- "${KB_PORT:-8000}:8000"
|
||||
volumes:
|
||||
- ${KB_DATA_PATH:-./data}:/data
|
||||
environment:
|
||||
- KB_MODEL=${KB_MODEL:-all-MiniLM-L6-v2}
|
||||
- KB_DEVICE=cpu
|
||||
- KB_INGEST_DEVICE=cpu
|
||||
- KB_API_KEY=${KB_API_KEY:-}
|
||||
- KB_SEARCH_THRESHOLD=${KB_SEARCH_THRESHOLD:-0.01}
|
||||
- KB_MIN_CHUNK_ALNUM=${KB_MIN_CHUNK_ALNUM:-3}
|
||||
- HF_HUB_OFFLINE=${HF_HUB_OFFLINE:-}
|
||||
restart: unless-stopped
|
||||
|
||||
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
|
||||
@@ -21,4 +21,25 @@ services:
|
||||
- KB_INGEST_DEVICE=${KB_INGEST_DEVICE:-auto}
|
||||
- KB_API_KEY=${KB_API_KEY:-}
|
||||
- KB_SEARCH_THRESHOLD=${KB_SEARCH_THRESHOLD:-0.01}
|
||||
- KB_MIN_CHUNK_ALNUM=${KB_MIN_CHUNK_ALNUM:-3}
|
||||
- KB_RERANK_ENABLED=${KB_RERANK_ENABLED:-true}
|
||||
- KB_RERANKER_MODEL=${KB_RERANKER_MODEL:-BAAI/bge-reranker-v2-m3}
|
||||
- KB_RERANK_CANDIDATES=${KB_RERANK_CANDIDATES:-40}
|
||||
- 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
|
||||
|
||||
@@ -1,21 +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}
|
||||
restart: unless-stopped
|
||||
Executable
+4
@@ -0,0 +1,4 @@
|
||||
#!/bin/bash
|
||||
|
||||
docker stop engine-kb-engine-1
|
||||
KB_MODEL=BAAI/bge-base-en-v1.5 KB_DATA_PATH=~/kb-data docker compose -f compose.nvidia.yaml up -d --build
|
||||
@@ -20,6 +20,11 @@ class Config:
|
||||
self.ingest_device = os.environ.get("KB_INGEST_DEVICE", "auto")
|
||||
self.api_key = os.environ.get("KB_API_KEY") or None
|
||||
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.port = int(os.environ.get("KB_PORT", "8000"))
|
||||
|
||||
@@ -35,10 +40,15 @@ class Config:
|
||||
def staging_dir(self) -> Path:
|
||||
return self.data_dir / "staging"
|
||||
|
||||
@property
|
||||
def documents_dir(self) -> Path:
|
||||
return self.data_dir / "documents"
|
||||
|
||||
def ensure_dirs(self):
|
||||
self.data_dir.mkdir(parents=True, exist_ok=True)
|
||||
self.hf_cache.mkdir(exist_ok=True)
|
||||
self.staging_dir.mkdir(exist_ok=True)
|
||||
self.documents_dir.mkdir(exist_ok=True)
|
||||
|
||||
|
||||
cfg = Config()
|
||||
|
||||
+180
-7
@@ -10,6 +10,60 @@ import struct
|
||||
from typing import Any, Optional
|
||||
|
||||
|
||||
def build_enriched_text(title: str, chunk_text: str, metadata: dict | None = None) -> str:
|
||||
"""Build enriched text by prepending document title and optional section header.
|
||||
|
||||
Format: "{title} > {section_header}\\n\\n{chunk_text}" or "{title}\\n\\n{chunk_text}".
|
||||
"""
|
||||
section_header = (metadata or {}).get("section_header")
|
||||
if section_header:
|
||||
return f"{title} > {section_header}\n\n{chunk_text}"
|
||||
return f"{title}\n\n{chunk_text}"
|
||||
|
||||
|
||||
def _backfill_enriched_text(conn: sqlite3.Connection) -> None:
|
||||
"""Backfill enriched_text for all existing chunks."""
|
||||
rows = conn.execute(
|
||||
"SELECT c.id, c.text, c.metadata, d.title "
|
||||
"FROM chunks c JOIN documents d ON c.document_id = d.id"
|
||||
).fetchall()
|
||||
for row in rows:
|
||||
metadata = json.loads(row["metadata"]) if row["metadata"] else None
|
||||
enriched = build_enriched_text(row["title"], row["text"], metadata)
|
||||
conn.execute("UPDATE chunks SET enriched_text = ? WHERE id = ?", (enriched, row["id"]))
|
||||
|
||||
|
||||
def _rebuild_fts(conn: sqlite3.Connection) -> None:
|
||||
"""Drop and recreate chunks_fts to index enriched_text, with updated triggers."""
|
||||
conn.executescript("""
|
||||
DROP TRIGGER IF EXISTS chunks_ai;
|
||||
DROP TRIGGER IF EXISTS chunks_ad;
|
||||
DROP TRIGGER IF EXISTS chunks_au;
|
||||
DROP TABLE IF EXISTS chunks_fts;
|
||||
|
||||
CREATE VIRTUAL TABLE chunks_fts USING fts5(
|
||||
text,
|
||||
content=chunks,
|
||||
content_rowid=id
|
||||
);
|
||||
|
||||
CREATE TRIGGER chunks_ai AFTER INSERT ON chunks BEGIN
|
||||
INSERT INTO chunks_fts(rowid, text) VALUES (new.id, new.enriched_text);
|
||||
END;
|
||||
|
||||
CREATE TRIGGER chunks_ad AFTER DELETE ON chunks BEGIN
|
||||
INSERT INTO chunks_fts(chunks_fts, rowid, text) VALUES ('delete', old.id, old.enriched_text);
|
||||
END;
|
||||
|
||||
CREATE TRIGGER chunks_au AFTER UPDATE ON chunks BEGIN
|
||||
INSERT INTO chunks_fts(chunks_fts, rowid, text) VALUES ('delete', old.id, old.enriched_text);
|
||||
INSERT INTO chunks_fts(rowid, text) VALUES (new.id, new.enriched_text);
|
||||
END;
|
||||
""")
|
||||
# Repopulate FTS from existing enriched_text
|
||||
conn.execute("INSERT INTO chunks_fts(rowid, text) SELECT id, enriched_text FROM chunks")
|
||||
|
||||
|
||||
def get_connection(db_path: str) -> sqlite3.Connection:
|
||||
"""Return a sqlite3 connection with WAL mode, Row factory, and foreign keys enabled."""
|
||||
import sqlite_vec
|
||||
@@ -20,6 +74,7 @@ def get_connection(db_path: str) -> sqlite3.Connection:
|
||||
conn.enable_load_extension(False)
|
||||
conn.row_factory = sqlite3.Row
|
||||
conn.execute("PRAGMA journal_mode=WAL")
|
||||
conn.execute("PRAGMA busy_timeout=5000")
|
||||
conn.execute("PRAGMA foreign_keys=ON")
|
||||
return conn
|
||||
|
||||
@@ -34,6 +89,8 @@ def init_schema(conn: sqlite3.Connection, embedding_dim: int) -> None:
|
||||
content_hash TEXT UNIQUE,
|
||||
doc_type TEXT,
|
||||
language TEXT,
|
||||
stored_path TEXT,
|
||||
original_filename TEXT,
|
||||
created_at TEXT DEFAULT current_timestamp
|
||||
);
|
||||
|
||||
@@ -42,6 +99,7 @@ def init_schema(conn: sqlite3.Connection, embedding_dim: int) -> None:
|
||||
document_id INTEGER REFERENCES documents(id) ON DELETE CASCADE,
|
||||
chunk_index INTEGER,
|
||||
text TEXT,
|
||||
enriched_text TEXT,
|
||||
token_count INTEGER,
|
||||
metadata TEXT DEFAULT '{{}}',
|
||||
UNIQUE(document_id, chunk_index)
|
||||
@@ -53,18 +111,18 @@ def init_schema(conn: sqlite3.Connection, embedding_dim: int) -> None:
|
||||
content_rowid=id
|
||||
);
|
||||
|
||||
-- Triggers to keep FTS index in sync with chunks table
|
||||
-- Triggers to keep FTS index in sync with chunks table (using enriched_text)
|
||||
CREATE TRIGGER IF NOT EXISTS chunks_ai AFTER INSERT ON chunks BEGIN
|
||||
INSERT INTO chunks_fts(rowid, text) VALUES (new.id, new.text);
|
||||
INSERT INTO chunks_fts(rowid, text) VALUES (new.id, new.enriched_text);
|
||||
END;
|
||||
|
||||
CREATE TRIGGER IF NOT EXISTS chunks_ad AFTER DELETE ON chunks BEGIN
|
||||
INSERT INTO chunks_fts(chunks_fts, rowid, text) VALUES ('delete', old.id, old.text);
|
||||
INSERT INTO chunks_fts(chunks_fts, rowid, text) VALUES ('delete', old.id, old.enriched_text);
|
||||
END;
|
||||
|
||||
CREATE TRIGGER IF NOT EXISTS chunks_au AFTER UPDATE ON chunks BEGIN
|
||||
INSERT INTO chunks_fts(chunks_fts, rowid, text) VALUES ('delete', old.id, old.text);
|
||||
INSERT INTO chunks_fts(rowid, text) VALUES (new.id, new.text);
|
||||
INSERT INTO chunks_fts(chunks_fts, rowid, text) VALUES ('delete', old.id, old.enriched_text);
|
||||
INSERT INTO chunks_fts(rowid, text) VALUES (new.id, new.enriched_text);
|
||||
END;
|
||||
|
||||
CREATE TABLE IF NOT EXISTS tags (
|
||||
@@ -114,6 +172,34 @@ def init_schema(conn: sqlite3.Connection, embedding_dim: int) -> None:
|
||||
if "content_hash" not in cols:
|
||||
conn.execute("ALTER TABLE jobs ADD COLUMN content_hash TEXT")
|
||||
|
||||
# Migrate: add stored_path and original_filename to documents if missing
|
||||
doc_cols = {row[1] for row in conn.execute("PRAGMA table_info(documents)").fetchall()}
|
||||
if "stored_path" not in doc_cols:
|
||||
conn.execute("ALTER TABLE documents ADD COLUMN stored_path TEXT")
|
||||
if "original_filename" not in doc_cols:
|
||||
conn.execute("ALTER TABLE documents ADD COLUMN original_filename TEXT")
|
||||
|
||||
# Migrate: add enriched_text to chunks and rebuild FTS to index it
|
||||
chunk_cols = {row[1] for row in conn.execute("PRAGMA table_info(chunks)").fetchall()}
|
||||
if "enriched_text" not in chunk_cols:
|
||||
conn.execute("ALTER TABLE chunks ADD COLUMN enriched_text TEXT")
|
||||
_backfill_enriched_text(conn)
|
||||
_rebuild_fts(conn)
|
||||
|
||||
# Migrate: add updated_at to documents if missing (v3.0.0)
|
||||
if "updated_at" not in doc_cols:
|
||||
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()
|
||||
|
||||
|
||||
@@ -196,6 +282,7 @@ def insert_chunk(
|
||||
document_id: int,
|
||||
chunk_index: int,
|
||||
text: str,
|
||||
enriched_text: str | None = None,
|
||||
token_count: Optional[int] = None,
|
||||
metadata: Any = None,
|
||||
) -> int:
|
||||
@@ -208,8 +295,8 @@ def insert_chunk(
|
||||
metadata_str = str(metadata)
|
||||
|
||||
cur = conn.execute(
|
||||
"INSERT INTO chunks(document_id, chunk_index, text, token_count, metadata) VALUES (?, ?, ?, ?, ?)",
|
||||
(document_id, chunk_index, text, token_count, metadata_str),
|
||||
"INSERT INTO chunks(document_id, chunk_index, text, enriched_text, token_count, metadata) VALUES (?, ?, ?, ?, ?, ?)",
|
||||
(document_id, chunk_index, text, enriched_text or text, token_count, metadata_str),
|
||||
)
|
||||
conn.commit()
|
||||
return cur.lastrowid
|
||||
@@ -253,6 +340,92 @@ def untag_document(conn: sqlite3.Connection, document_id: int, tag_names: list[s
|
||||
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
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@@ -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"),
|
||||
".sh": ("code", "bash"),
|
||||
".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
|
||||
HierarchicalChunker,
|
||||
)
|
||||
from kb.ingest.quality import has_minimum_content
|
||||
|
||||
|
||||
def _fixed_size_chunks(text: str, max_chars: int = 2000) -> list[str]:
|
||||
@@ -40,6 +41,7 @@ def _fixed_size_chunks(text: str, max_chars: int = 2000) -> list[str]:
|
||||
def chunk_document(
|
||||
file_path: Path,
|
||||
ingest_device: str = "cpu",
|
||||
min_chunk_alnum: int = 3,
|
||||
) -> list[dict]:
|
||||
"""Convert and chunk a PDF/DOCX/HTML document using Docling.
|
||||
|
||||
@@ -71,7 +73,7 @@ def chunk_document(
|
||||
chunks: list[dict] = []
|
||||
for idx, chunk in enumerate(raw_chunks):
|
||||
text = chunk.text.strip() if hasattr(chunk, "text") else str(chunk).strip()
|
||||
if not text:
|
||||
if not text or not has_minimum_content(text, min_chunk_alnum):
|
||||
continue
|
||||
|
||||
metadata: dict = {}
|
||||
@@ -98,6 +100,8 @@ def chunk_document(
|
||||
if not full_text and hasattr(doc, "text"):
|
||||
full_text = doc.text
|
||||
for idx, piece in enumerate(_fixed_size_chunks(full_text)):
|
||||
if not has_minimum_content(piece, min_chunk_alnum):
|
||||
continue
|
||||
chunks.append({
|
||||
"text": piece,
|
||||
"chunk_index": idx,
|
||||
|
||||
@@ -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()
|
||||
@@ -1 +1 @@
|
||||
from kb.routes import health, search, jobs, documents, tags, status, reindex, auth
|
||||
from kb.routes import health, search, jobs, documents, tags, status, reindex, auth, notes
|
||||
|
||||
@@ -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()
|
||||
@@ -1,26 +1,36 @@
|
||||
"""Document management endpoints — list, view, and delete documents."""
|
||||
|
||||
import json
|
||||
import logging
|
||||
import mimetypes
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
from fastapi import HTTPException, Query
|
||||
from pydantic import BaseModel, Field
|
||||
from fastapi.responses import FileResponse
|
||||
|
||||
from main import app
|
||||
from kb.config import cfg
|
||||
from kb.database import get_connection
|
||||
from kb.search import hybrid_search
|
||||
|
||||
logger = logging.getLogger("kb.routes.documents")
|
||||
|
||||
|
||||
@app.get("/api/v1/documents")
|
||||
async def list_documents(
|
||||
type: Optional[str] = Query(None),
|
||||
tags: Optional[str] = Query(None),
|
||||
title: Optional[str] = Query(None),
|
||||
filename: Optional[str] = Query(None),
|
||||
):
|
||||
conn = get_connection(cfg.db_path)
|
||||
try:
|
||||
sql = """
|
||||
SELECT d.id, d.title, d.doc_type,
|
||||
SELECT d.id, d.title, d.original_filename, d.source_path, d.doc_type,
|
||||
(SELECT COUNT(*) FROM chunks c WHERE c.document_id = d.id) AS chunk_count,
|
||||
d.created_at
|
||||
d.created_at, d.updated_at
|
||||
FROM documents d
|
||||
"""
|
||||
joins: list[str] = []
|
||||
@@ -31,6 +41,14 @@ async def list_documents(
|
||||
where.append("d.doc_type = ?")
|
||||
params.append(type)
|
||||
|
||||
if title:
|
||||
where.append("d.title LIKE ? COLLATE NOCASE")
|
||||
params.append(f"%{title}%")
|
||||
|
||||
if filename:
|
||||
where.append("d.original_filename LIKE ? COLLATE NOCASE")
|
||||
params.append(f"%{filename}%")
|
||||
|
||||
if tags:
|
||||
tag_list = [t.strip() for t in tags.split(",") if t.strip()]
|
||||
for i, tag in enumerate(tag_list):
|
||||
@@ -44,7 +62,7 @@ async def list_documents(
|
||||
if where:
|
||||
sql += " WHERE " + " AND ".join(where)
|
||||
|
||||
sql += " ORDER BY d.created_at DESC"
|
||||
sql += " ORDER BY COALESCE(d.updated_at, d.created_at) DESC"
|
||||
|
||||
rows = conn.execute(sql, params).fetchall()
|
||||
|
||||
@@ -64,10 +82,13 @@ async def list_documents(
|
||||
results.append({
|
||||
"id": row["id"],
|
||||
"title": row["title"],
|
||||
"original_filename": row["original_filename"],
|
||||
"source_path": row["source_path"],
|
||||
"doc_type": row["doc_type"],
|
||||
"tags": [t["name"] for t in tag_rows],
|
||||
"chunk_count": row["chunk_count"],
|
||||
"created_at": row["created_at"],
|
||||
"updated_at": row["updated_at"],
|
||||
})
|
||||
|
||||
return results
|
||||
@@ -75,8 +96,49 @@ async def list_documents(
|
||||
conn.close()
|
||||
|
||||
|
||||
class DocumentFindRequest(BaseModel):
|
||||
query: str
|
||||
top: int = Field(default=10, ge=1, le=100)
|
||||
tags: Optional[list[str]] = None
|
||||
doc_type: Optional[str] = None
|
||||
|
||||
|
||||
@app.post("/api/v1/documents/find")
|
||||
async def find_documents(req: DocumentFindRequest):
|
||||
"""Return document-level results aggregated from hybrid chunk search."""
|
||||
conn = get_connection(cfg.db_path)
|
||||
try:
|
||||
search_result = hybrid_search(
|
||||
conn,
|
||||
req.query,
|
||||
cfg,
|
||||
top=max(req.top * 10, 50),
|
||||
tags=req.tags,
|
||||
doc_type=req.doc_type,
|
||||
)
|
||||
documents: dict[int, dict] = {}
|
||||
for result in search_result["results"]:
|
||||
doc_id = result["document_id"]
|
||||
if doc_id not in documents:
|
||||
documents[doc_id] = {
|
||||
"document_id": doc_id,
|
||||
"title": result["title"],
|
||||
"doc_type": result["doc_type"],
|
||||
"source_path": result["source_path"],
|
||||
"original_filename": result["original_filename"],
|
||||
"tags": result["tags"],
|
||||
"score": result["score"],
|
||||
"hit_count": 0,
|
||||
"top_chunk": result["text"],
|
||||
}
|
||||
documents[doc_id]["hit_count"] += 1
|
||||
return list(documents.values())[: req.top]
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
@app.get("/api/v1/documents/{doc_id}")
|
||||
async def get_document(doc_id: int):
|
||||
async def get_document(doc_id: int, include_chunks: bool = Query(True)):
|
||||
conn = get_connection(cfg.db_path)
|
||||
try:
|
||||
doc = conn.execute(
|
||||
@@ -85,6 +147,11 @@ async def get_document(doc_id: int):
|
||||
if not doc:
|
||||
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(
|
||||
"SELECT * FROM chunks WHERE document_id = ? ORDER BY chunk_index",
|
||||
(doc_id,),
|
||||
@@ -100,21 +167,67 @@ async def get_document(doc_id: int):
|
||||
(doc_id,),
|
||||
).fetchall()
|
||||
|
||||
stored_path = doc["stored_path"]
|
||||
has_file = bool(stored_path and Path(stored_path).exists())
|
||||
|
||||
return {
|
||||
**dict(doc),
|
||||
"has_file": has_file,
|
||||
"tags": [t["name"] for t in tag_rows],
|
||||
"chunk_count": chunk_count,
|
||||
"chunks": [dict(c) for c in chunks],
|
||||
}
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
@app.get("/api/v1/documents/{doc_id}/file")
|
||||
async def download_document_file(doc_id: int):
|
||||
conn = get_connection(cfg.db_path)
|
||||
try:
|
||||
doc = conn.execute(
|
||||
"SELECT id, title, stored_path, original_filename FROM documents WHERE id = ?",
|
||||
(doc_id,),
|
||||
).fetchone()
|
||||
if not doc:
|
||||
raise HTTPException(status_code=404, detail="Document not found.")
|
||||
|
||||
stored_path = doc["stored_path"]
|
||||
if not stored_path:
|
||||
raise HTTPException(
|
||||
status_code=404,
|
||||
detail="Original file not available - ingested before document storage was enabled.",
|
||||
)
|
||||
|
||||
file_path = Path(stored_path)
|
||||
if not file_path.exists():
|
||||
raise HTTPException(
|
||||
status_code=404,
|
||||
detail="Stored file not found on disk.",
|
||||
)
|
||||
|
||||
original_filename = doc["original_filename"]
|
||||
if not original_filename:
|
||||
ext = file_path.suffix
|
||||
original_filename = (doc["title"] or "document") + ext
|
||||
|
||||
media_type = mimetypes.guess_type(original_filename)[0] or "application/octet-stream"
|
||||
|
||||
return FileResponse(
|
||||
path=str(file_path),
|
||||
media_type=media_type,
|
||||
filename=original_filename,
|
||||
)
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
@app.delete("/api/v1/documents/{doc_id}")
|
||||
async def delete_document(doc_id: int):
|
||||
conn = get_connection(cfg.db_path)
|
||||
try:
|
||||
doc = conn.execute(
|
||||
"SELECT id, title FROM documents WHERE id = ?", (doc_id,)
|
||||
"SELECT id, title, stored_path FROM documents WHERE id = ?", (doc_id,)
|
||||
).fetchone()
|
||||
if not doc:
|
||||
raise HTTPException(status_code=404, detail="Document not found.")
|
||||
@@ -134,6 +247,19 @@ async def delete_document(doc_id: int):
|
||||
conn.execute("DELETE FROM documents WHERE id = ?", (doc_id,))
|
||||
conn.commit()
|
||||
|
||||
# Delete stored file from disk
|
||||
stored_path = doc["stored_path"]
|
||||
if stored_path:
|
||||
try:
|
||||
file_path = Path(stored_path)
|
||||
if file_path.exists():
|
||||
file_path.unlink()
|
||||
logger.info("Deleted stored file: %s", stored_path)
|
||||
else:
|
||||
logger.warning("Stored file already missing: %s", stored_path)
|
||||
except OSError as exc:
|
||||
logger.warning("Failed to delete stored file %s: %s", stored_path, exc)
|
||||
|
||||
return {
|
||||
"status": "deleted",
|
||||
"document_id": doc_id,
|
||||
|
||||
@@ -10,6 +10,7 @@ from fastapi.responses import JSONResponse
|
||||
from main import app
|
||||
from kb.config import cfg
|
||||
from kb.database import get_connection, create_job, get_job, list_jobs, get_document_by_hash
|
||||
from kb.ingest.note import auto_title
|
||||
from kb.staging import stage_file, stage_note
|
||||
|
||||
|
||||
@@ -32,6 +33,7 @@ async def submit_job(
|
||||
content_hash = hashlib.sha256(content).hexdigest()
|
||||
filename = file.filename
|
||||
else:
|
||||
title = title or auto_title(note) or "note"
|
||||
content = note.encode("utf-8")
|
||||
content_hash = hashlib.sha256(content).hexdigest()
|
||||
filename = None
|
||||
@@ -48,7 +50,7 @@ async def submit_job(
|
||||
if file:
|
||||
staging_path = stage_file(cfg.staging_dir, file.filename, content)
|
||||
else:
|
||||
staging_path = stage_note(cfg.staging_dir, title or "note", note)
|
||||
staging_path = stage_note(cfg.staging_dir, title, note)
|
||||
filename = staging_path.name
|
||||
|
||||
tags_list = [t.strip() for t in tags.split(",") if t.strip()] if tags else []
|
||||
|
||||
@@ -0,0 +1,120 @@
|
||||
"""Note mutation endpoint — update existing notes in place."""
|
||||
|
||||
import hashlib
|
||||
import logging
|
||||
|
||||
from fastapi import HTTPException
|
||||
from pydantic import BaseModel
|
||||
|
||||
from main import app
|
||||
from kb.config import cfg
|
||||
from kb.database import (
|
||||
get_connection,
|
||||
build_enriched_text,
|
||||
insert_chunk,
|
||||
insert_embedding,
|
||||
)
|
||||
from kb.embeddings import embed_texts
|
||||
from kb.ingest.note import chunk_note
|
||||
|
||||
logger = logging.getLogger("kb.routes.notes")
|
||||
|
||||
|
||||
class NoteUpdateRequest(BaseModel):
|
||||
text: str
|
||||
|
||||
|
||||
@app.patch("/api/v1/notes/{doc_id}")
|
||||
async def update_note(doc_id: int, req: NoteUpdateRequest):
|
||||
conn = get_connection(cfg.db_path)
|
||||
try:
|
||||
doc = conn.execute(
|
||||
"SELECT id, title, doc_type FROM documents WHERE id = ?", (doc_id,)
|
||||
).fetchone()
|
||||
if not doc:
|
||||
raise HTTPException(status_code=404, detail="Document not found.")
|
||||
if doc["doc_type"] != "note":
|
||||
raise HTTPException(
|
||||
status_code=422,
|
||||
detail="Only notes can be updated via this endpoint.",
|
||||
)
|
||||
|
||||
title = doc["title"]
|
||||
|
||||
# Delete existing chunks and their 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"],))
|
||||
conn.execute("DELETE FROM chunks WHERE document_id = ?", (doc_id,))
|
||||
|
||||
# Run note chunking pipeline on new text
|
||||
chunks = chunk_note(req.text)
|
||||
chunk_texts = [c["text"] for c in chunks]
|
||||
chunk_metas = [
|
||||
{k: v for k, v in c.items() if k != "text"} or None for c in chunks
|
||||
]
|
||||
|
||||
enriched_texts = [
|
||||
build_enriched_text(title, ct, cm)
|
||||
for ct, cm in zip(chunk_texts, chunk_metas)
|
||||
]
|
||||
|
||||
# Embed — if this fails, the transaction rolls back
|
||||
vectors = embed_texts(enriched_texts)
|
||||
|
||||
for idx, (chunk_text, enriched, vector) in enumerate(
|
||||
zip(chunk_texts, enriched_texts, vectors)
|
||||
):
|
||||
chunk_id = insert_chunk(
|
||||
conn,
|
||||
document_id=doc_id,
|
||||
chunk_index=idx,
|
||||
text=chunk_text,
|
||||
enriched_text=enriched,
|
||||
metadata=chunk_metas[idx],
|
||||
)
|
||||
insert_embedding(conn, chunk_id, vector)
|
||||
|
||||
# Update content_hash and updated_at
|
||||
content_hash = hashlib.sha256(req.text.encode("utf-8")).hexdigest()
|
||||
conn.execute(
|
||||
"UPDATE documents SET content_hash = ?, updated_at = current_timestamp WHERE id = ?",
|
||||
(content_hash, doc_id),
|
||||
)
|
||||
conn.commit()
|
||||
|
||||
# Return updated document
|
||||
updated_doc = conn.execute(
|
||||
"SELECT * FROM documents WHERE id = ?", (doc_id,)
|
||||
).fetchone()
|
||||
|
||||
new_chunks = conn.execute(
|
||||
"SELECT * FROM chunks WHERE document_id = ? ORDER BY chunk_index",
|
||||
(doc_id,),
|
||||
).fetchall()
|
||||
|
||||
tag_rows = conn.execute(
|
||||
"""
|
||||
SELECT t.name FROM tags t
|
||||
JOIN document_tags dt ON t.id = dt.tag_id
|
||||
WHERE dt.document_id = ?
|
||||
ORDER BY t.name
|
||||
""",
|
||||
(doc_id,),
|
||||
).fetchall()
|
||||
|
||||
return {
|
||||
**dict(updated_doc),
|
||||
"tags": [t["name"] for t in tag_rows],
|
||||
"chunks": [dict(c) for c in new_chunks],
|
||||
}
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception:
|
||||
conn.rollback()
|
||||
logger.exception("Failed to update note %d", doc_id)
|
||||
raise HTTPException(status_code=500, detail="Failed to update note.")
|
||||
finally:
|
||||
conn.close()
|
||||
@@ -19,10 +19,10 @@ async def reindex():
|
||||
|
||||
conn = get_connection(cfg.db_path)
|
||||
try:
|
||||
# Fetch all chunks
|
||||
rows = conn.execute("SELECT id, text FROM chunks ORDER BY id").fetchall()
|
||||
# Fetch all chunks — use enriched_text for embedding (includes title context)
|
||||
rows = conn.execute("SELECT id, enriched_text FROM chunks ORDER BY id").fetchall()
|
||||
chunk_ids = [row["id"] for row in rows]
|
||||
chunk_texts = [row["text"] for row in rows]
|
||||
chunk_texts = [row["enriched_text"] or "" for row in rows]
|
||||
|
||||
logger.info("Reindexing %d chunks with model '%s'", len(chunk_ids), cfg.model)
|
||||
|
||||
|
||||
@@ -19,6 +19,8 @@ class SearchRequest(BaseModel):
|
||||
fts_only: bool = False
|
||||
vec_only: bool = False
|
||||
threshold: Optional[float] = None
|
||||
explain: bool = False
|
||||
rerank: Optional[bool] = None
|
||||
|
||||
|
||||
@app.post("/api/v1/search")
|
||||
@@ -35,6 +37,8 @@ async def search(req: SearchRequest):
|
||||
fts_only=req.fts_only,
|
||||
vec_only=req.vec_only,
|
||||
threshold=req.threshold,
|
||||
explain=req.explain,
|
||||
rerank=req.rerank,
|
||||
)
|
||||
return result
|
||||
except Exception as exc:
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
import os
|
||||
|
||||
from main import app, __version__
|
||||
from kb import reranker
|
||||
from kb.config import cfg
|
||||
from kb.database import get_connection
|
||||
from kb.embeddings import get_model_dim
|
||||
@@ -62,6 +63,12 @@ async def status():
|
||||
"queued": queue_stats.get("queued", 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:
|
||||
conn.close()
|
||||
|
||||
@@ -16,14 +16,45 @@ async def list_tags():
|
||||
try:
|
||||
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
|
||||
LEFT JOIN document_tags dt ON t.id = dt.tag_id
|
||||
GROUP BY t.id, t.name
|
||||
ORDER BY t.name
|
||||
"""
|
||||
).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:
|
||||
conn.close()
|
||||
|
||||
@@ -48,6 +79,13 @@ async def update_document_tags(doc_id: int, req: TagUpdateRequest):
|
||||
if req.remove:
|
||||
untag_document(conn, doc_id, req.remove)
|
||||
|
||||
if req.add or req.remove:
|
||||
conn.execute(
|
||||
"UPDATE documents SET updated_at = current_timestamp WHERE id = ?",
|
||||
(doc_id,),
|
||||
)
|
||||
conn.commit()
|
||||
|
||||
tag_rows = conn.execute(
|
||||
"""
|
||||
SELECT t.name FROM tags t
|
||||
|
||||
+168
-17
@@ -18,6 +18,8 @@ def hybrid_search(
|
||||
fts_only: bool = False,
|
||||
vec_only: bool = False,
|
||||
threshold: float | None = None,
|
||||
explain: bool = False,
|
||||
rerank: bool | None = None,
|
||||
) -> dict:
|
||||
"""Run hybrid search and return merged, enriched results.
|
||||
|
||||
@@ -31,11 +33,29 @@ def hybrid_search(
|
||||
fts_only: Only use FTS5 (skip vector search).
|
||||
vec_only: Only use vector search (skip FTS5).
|
||||
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:
|
||||
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
|
||||
if do_rerank:
|
||||
candidate_count = max(candidate_count, cfg.rerank_candidates)
|
||||
|
||||
fts_results: dict[int, float] = {}
|
||||
vec_results: dict[int, float] = {}
|
||||
@@ -49,10 +69,12 @@ def hybrid_search(
|
||||
# --- merge ---------------------------------------------------------------
|
||||
if fts_only:
|
||||
merged = sorted(fts_results.items(), key=lambda x: x[1], reverse=True)
|
||||
details = _single_arm_details("fts", fts_results)
|
||||
elif vec_only:
|
||||
merged = sorted(vec_results.items(), key=lambda x: x[1], reverse=True)
|
||||
details = _single_arm_details("vec", vec_results)
|
||||
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
|
||||
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]
|
||||
|
||||
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]
|
||||
|
||||
# --- enrich --------------------------------------------------------------
|
||||
results = _enrich(conn, merged)
|
||||
results = _enrich(conn, merged, details if explain else None)
|
||||
|
||||
return {
|
||||
"query": query,
|
||||
"results": results,
|
||||
"total_matches": total_matches,
|
||||
"returned": len(results),
|
||||
"reranked": reranked,
|
||||
}
|
||||
|
||||
|
||||
@@ -232,31 +268,135 @@ def _rrf_merge(
|
||||
fts_results: dict[int, float],
|
||||
vec_results: dict[int, float],
|
||||
k: int = 60,
|
||||
) -> list[tuple[int, float]]:
|
||||
) -> tuple[list[tuple[int, float]], dict[int, dict]]:
|
||||
"""Reciprocal Rank Fusion over two scored result sets.
|
||||
|
||||
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:
|
||||
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)
|
||||
vec_ranked = _rank_by_score(vec_results)
|
||||
|
||||
all_ids = set(fts_ranked) | set(vec_ranked)
|
||||
scores: list[tuple[int, float]] = []
|
||||
details: dict[int, dict] = {}
|
||||
|
||||
for chunk_id in all_ids:
|
||||
rrf = 0.0
|
||||
if chunk_id in fts_ranked:
|
||||
rrf += 1.0 / (k + fts_ranked[chunk_id])
|
||||
if chunk_id in vec_ranked:
|
||||
rrf += 1.0 / (k + vec_ranked[chunk_id])
|
||||
fts_rank = fts_ranked.get(chunk_id)
|
||||
vec_rank = vec_ranked.get(chunk_id)
|
||||
rrf_fts = 1.0 / (k + fts_rank) if fts_rank is not None else None
|
||||
rrf_vec = 1.0 / (k + vec_rank) if vec_rank is not None else None
|
||||
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.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]:
|
||||
@@ -268,8 +408,13 @@ def _rank_by_score(results: dict[int, float]) -> dict[int, int]:
|
||||
def _enrich(
|
||||
conn: sqlite3.Connection,
|
||||
merged: list[tuple[int, float]],
|
||||
details: dict[int, dict] | None = None,
|
||||
) -> 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] = []
|
||||
|
||||
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,
|
||||
d.id AS doc_id, d.title, d.doc_type, d.source_path,
|
||||
d.created_at
|
||||
d.created_at, d.original_filename
|
||||
FROM chunks c
|
||||
JOIN documents d ON c.document_id = d.id
|
||||
WHERE c.id = ?
|
||||
@@ -292,7 +437,7 @@ def _enrich(
|
||||
|
||||
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
|
||||
WHERE dt.document_id = ?
|
||||
ORDER BY t.name
|
||||
@@ -300,8 +445,9 @@ def _enrich(
|
||||
(row[4],), # doc_id
|
||||
).fetchall()
|
||||
|
||||
results.append({
|
||||
result = {
|
||||
"chunk_id": row[0],
|
||||
"document_id": row[4],
|
||||
"score": round(score, 6),
|
||||
"text": row[1],
|
||||
"chunk_index": row[2],
|
||||
@@ -310,7 +456,12 @@ def _enrich(
|
||||
"doc_type": row[6],
|
||||
"source_path": row[7],
|
||||
"created_at": row[8],
|
||||
"original_filename": row[9],
|
||||
"tags": [t[0] for t in tag_rows],
|
||||
})
|
||||
"tag_contexts": {t[0]: t[1] for t in tag_rows if t[1]},
|
||||
}
|
||||
if details is not None and row[0] in details:
|
||||
result["explain"] = details[row[0]]
|
||||
results.append(result)
|
||||
|
||||
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.
|
||||
"""
|
||||
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)
|
||||
logger.debug("Staged file: %s (%d bytes)", dest, len(content))
|
||||
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.
|
||||
"""
|
||||
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")
|
||||
logger.debug("Staged note: %s (%d chars)", dest, len(text))
|
||||
return dest
|
||||
|
||||
+52
-9
@@ -4,9 +4,11 @@ import asyncio
|
||||
import hashlib
|
||||
import json
|
||||
import logging
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
|
||||
from kb import config, database, embeddings, staging
|
||||
from kb.database import build_enriched_text
|
||||
from kb.ingest import detector
|
||||
|
||||
logger = logging.getLogger("kb.worker")
|
||||
@@ -111,7 +113,9 @@ def _process_job(job_row) -> tuple[str, int | None, int]:
|
||||
chunks = chunk_note(text)
|
||||
elif doc_type == "pdf":
|
||||
from kb.ingest.docling_pipeline import chunk_document
|
||||
chunks = chunk_document(staged_path, cfg.ingest_device)
|
||||
chunks = chunk_document(
|
||||
staged_path, cfg.ingest_device, cfg.min_chunk_alnum
|
||||
)
|
||||
elif doc_type == "markdown":
|
||||
text = staged_path.read_text(encoding="utf-8")
|
||||
from kb.ingest.markdown import chunk_markdown
|
||||
@@ -122,6 +126,12 @@ def _process_job(job_row) -> tuple[str, int | None, int]:
|
||||
_, language = detector.detect_type(Path(filename))
|
||||
from kb.ingest.code import chunk_code
|
||||
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:
|
||||
raise ValueError(f"Unsupported doc_type: {doc_type}")
|
||||
|
||||
@@ -145,20 +155,30 @@ def _process_job(job_row) -> tuple[str, int | None, int]:
|
||||
)
|
||||
|
||||
chunk_texts = [c if isinstance(c, str) else c["text"] for c in chunks]
|
||||
vectors = embeddings.embed_texts(chunk_texts)
|
||||
chunk_metas = []
|
||||
for idx, c in enumerate(chunks):
|
||||
if isinstance(c, str):
|
||||
chunk_metas.append(None)
|
||||
else:
|
||||
meta = {k: v for k, v in c.items() if k != "text"} or None
|
||||
chunk_metas.append(meta)
|
||||
|
||||
for idx, (chunk_text, vector) in enumerate(zip(chunk_texts, vectors)):
|
||||
metadata = None
|
||||
if not isinstance(chunks[idx], str):
|
||||
metadata = {
|
||||
k: v for k, v in chunks[idx].items() if k != "text"
|
||||
} or None
|
||||
enriched_texts = [
|
||||
build_enriched_text(title, ct, cm)
|
||||
for ct, cm in zip(chunk_texts, chunk_metas)
|
||||
]
|
||||
vectors = embeddings.embed_texts(enriched_texts)
|
||||
|
||||
for idx, (chunk_text, enriched, vector) in enumerate(
|
||||
zip(chunk_texts, enriched_texts, vectors)
|
||||
):
|
||||
chunk_id = database.insert_chunk(
|
||||
conn,
|
||||
document_id=doc_id,
|
||||
chunk_index=idx,
|
||||
text=chunk_text,
|
||||
metadata=metadata,
|
||||
enriched_text=enriched,
|
||||
metadata=chunk_metas[idx],
|
||||
)
|
||||
database.insert_embedding(conn, chunk_id, vector)
|
||||
|
||||
@@ -168,8 +188,31 @@ def _process_job(job_row) -> tuple[str, int | None, int]:
|
||||
database.tag_document(conn, doc_id, tags)
|
||||
|
||||
conn.commit()
|
||||
|
||||
# --- Move original file to persistent storage ---------------------
|
||||
ext = Path(filename).suffix or staged_path.suffix
|
||||
dest = cfg.documents_dir / f"{content_hash}{ext}"
|
||||
try:
|
||||
cfg.documents_dir.mkdir(parents=True, exist_ok=True)
|
||||
shutil.move(str(staged_path), str(dest))
|
||||
conn_update = database.get_connection(cfg.db_path)
|
||||
try:
|
||||
conn_update.execute(
|
||||
"UPDATE documents SET stored_path = ?, original_filename = ? WHERE id = ?",
|
||||
(str(dest), filename, doc_id),
|
||||
)
|
||||
conn_update.commit()
|
||||
finally:
|
||||
conn_update.close()
|
||||
logger.info("Stored original file: %s", dest)
|
||||
except Exception as exc:
|
||||
logger.warning("Failed to store original file: %s", exc)
|
||||
staging.cleanup(staged_path)
|
||||
|
||||
return ("done", doc_id, len(chunk_texts))
|
||||
|
||||
finally:
|
||||
conn.close()
|
||||
# Only clean up staging if the file is still there (not moved)
|
||||
if staged_path.exists():
|
||||
staging.cleanup(staged_path)
|
||||
|
||||
+13
-1
@@ -40,6 +40,18 @@ async def lifespan(app: FastAPI):
|
||||
init_schema(conn, model_dim)
|
||||
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
|
||||
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)
|
||||
|
||||
# Import routes after app is created
|
||||
from kb.routes import health, search, jobs, documents, tags, status, reindex, auth # noqa: E402, F401
|
||||
from kb.routes import health, search, jobs, documents, tags, status, reindex, auth, notes, bulk # noqa: E402, F401
|
||||
|
||||
if __name__ == "__main__":
|
||||
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,223 @@
|
||||
"""Tests for original document storage feature."""
|
||||
|
||||
import hashlib
|
||||
import shutil
|
||||
import sqlite3
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def data_dir(tmp_path):
|
||||
"""Create a temporary data directory with required subdirectories."""
|
||||
staging = tmp_path / "staging"
|
||||
staging.mkdir()
|
||||
documents = tmp_path / "documents"
|
||||
documents.mkdir()
|
||||
return tmp_path
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def db_conn(data_dir):
|
||||
"""Create an in-memory-style SQLite DB with the full schema."""
|
||||
db_path = data_dir / "kb.db"
|
||||
conn = sqlite3.connect(str(db_path))
|
||||
conn.row_factory = sqlite3.Row
|
||||
conn.execute("PRAGMA foreign_keys=ON")
|
||||
conn.executescript("""
|
||||
CREATE TABLE IF NOT EXISTS documents (
|
||||
id INTEGER PRIMARY KEY,
|
||||
title TEXT,
|
||||
source_path TEXT,
|
||||
content_hash TEXT UNIQUE,
|
||||
doc_type TEXT,
|
||||
language TEXT,
|
||||
stored_path TEXT,
|
||||
original_filename TEXT,
|
||||
created_at TEXT DEFAULT current_timestamp
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS chunks (
|
||||
id INTEGER PRIMARY KEY,
|
||||
document_id INTEGER REFERENCES documents(id) ON DELETE CASCADE,
|
||||
chunk_index INTEGER,
|
||||
text TEXT,
|
||||
token_count INTEGER,
|
||||
metadata TEXT DEFAULT '{}',
|
||||
UNIQUE(document_id, chunk_index)
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS tags (
|
||||
id INTEGER PRIMARY KEY,
|
||||
name TEXT UNIQUE COLLATE NOCASE
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS document_tags (
|
||||
document_id INTEGER REFERENCES documents(id) ON DELETE CASCADE,
|
||||
tag_id INTEGER REFERENCES tags(id) ON DELETE CASCADE,
|
||||
UNIQUE(document_id, tag_id)
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS jobs (
|
||||
id INTEGER PRIMARY KEY,
|
||||
filename TEXT,
|
||||
status TEXT DEFAULT 'queued',
|
||||
doc_type TEXT,
|
||||
tags_json TEXT DEFAULT '[]',
|
||||
title TEXT,
|
||||
error TEXT,
|
||||
document_id INTEGER,
|
||||
chunk_count INTEGER DEFAULT 0,
|
||||
staging_path TEXT,
|
||||
content_hash TEXT,
|
||||
created_at TEXT DEFAULT current_timestamp,
|
||||
completed_at TEXT
|
||||
);
|
||||
""")
|
||||
conn.commit()
|
||||
yield conn
|
||||
conn.close()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_pdf(data_dir):
|
||||
"""Create a fake PDF file in staging."""
|
||||
content = b"%PDF-1.4 fake pdf content for testing"
|
||||
staging = data_dir / "staging"
|
||||
path = staging / "test_upload.pdf"
|
||||
path.write_bytes(content)
|
||||
return path, content
|
||||
|
||||
|
||||
class TestWorkerFileStorage:
|
||||
"""Tests for worker moving files to persistent storage."""
|
||||
|
||||
def test_successful_ingestion_stores_file(self, data_dir, db_conn, sample_pdf):
|
||||
"""7.1 - Test successful ingestion stores file at expected path."""
|
||||
staged_path, content = sample_pdf
|
||||
content_hash = hashlib.sha256(content).hexdigest()
|
||||
documents_dir = data_dir / "documents"
|
||||
|
||||
expected_dest = documents_dir / f"{content_hash}.pdf"
|
||||
|
||||
# Simulate what the worker does: move file to documents dir
|
||||
shutil.move(str(staged_path), str(expected_dest))
|
||||
|
||||
assert expected_dest.exists()
|
||||
assert expected_dest.read_bytes() == content
|
||||
assert not staged_path.exists()
|
||||
|
||||
# Simulate DB update
|
||||
db_conn.execute(
|
||||
"INSERT INTO documents(title, source_path, content_hash, doc_type, stored_path, original_filename) "
|
||||
"VALUES (?, ?, ?, ?, ?, ?)",
|
||||
("Test PDF", str(staged_path), content_hash, "pdf", str(expected_dest), "test_upload.pdf"),
|
||||
)
|
||||
db_conn.commit()
|
||||
|
||||
row = db_conn.execute("SELECT stored_path, original_filename FROM documents WHERE content_hash = ?", (content_hash,)).fetchone()
|
||||
assert row["stored_path"] == str(expected_dest)
|
||||
assert row["original_filename"] == "test_upload.pdf"
|
||||
|
||||
def test_failed_ingestion_no_file_in_documents(self, data_dir, sample_pdf):
|
||||
"""7.2 - Test failed ingestion does not leave file in documents dir."""
|
||||
staged_path, _ = sample_pdf
|
||||
documents_dir = data_dir / "documents"
|
||||
|
||||
# Simulate failure: staging file gets cleaned up, nothing in documents dir
|
||||
staged_path.unlink()
|
||||
|
||||
assert len(list(documents_dir.iterdir())) == 0
|
||||
|
||||
def test_document_deletion_removes_stored_file(self, data_dir, db_conn, sample_pdf):
|
||||
"""7.4 - Test document deletion removes stored file."""
|
||||
staged_path, content = sample_pdf
|
||||
content_hash = hashlib.sha256(content).hexdigest()
|
||||
documents_dir = data_dir / "documents"
|
||||
|
||||
dest = documents_dir / f"{content_hash}.pdf"
|
||||
shutil.move(str(staged_path), str(dest))
|
||||
|
||||
db_conn.execute(
|
||||
"INSERT INTO documents(title, source_path, content_hash, doc_type, stored_path, original_filename) "
|
||||
"VALUES (?, ?, ?, ?, ?, ?)",
|
||||
("Test PDF", str(staged_path), content_hash, "pdf", str(dest), "test_upload.pdf"),
|
||||
)
|
||||
db_conn.commit()
|
||||
|
||||
# Simulate delete: remove from DB and disk
|
||||
doc = db_conn.execute("SELECT id, stored_path FROM documents WHERE content_hash = ?", (content_hash,)).fetchone()
|
||||
stored = Path(doc["stored_path"])
|
||||
db_conn.execute("DELETE FROM documents WHERE id = ?", (doc["id"],))
|
||||
db_conn.commit()
|
||||
|
||||
if stored.exists():
|
||||
stored.unlink()
|
||||
|
||||
assert not stored.exists()
|
||||
assert db_conn.execute("SELECT COUNT(*) FROM documents", ()).fetchone()[0] == 0
|
||||
|
||||
def test_download_404_for_document_without_stored_file(self, db_conn):
|
||||
"""7.5 - Test download returns 404 for documents without stored files."""
|
||||
db_conn.execute(
|
||||
"INSERT INTO documents(title, source_path, content_hash, doc_type) "
|
||||
"VALUES (?, ?, ?, ?)",
|
||||
("Old Doc", "/tmp/gone", "abc123", "pdf"),
|
||||
)
|
||||
db_conn.commit()
|
||||
|
||||
row = db_conn.execute("SELECT stored_path FROM documents WHERE content_hash = 'abc123'").fetchone()
|
||||
assert row["stored_path"] is None
|
||||
|
||||
|
||||
class TestFileDownloadEndpoint:
|
||||
"""Tests for the /api/v1/documents/{id}/file endpoint logic."""
|
||||
|
||||
def test_file_response_uses_original_filename(self, data_dir, db_conn, sample_pdf):
|
||||
"""7.3 - Test file download uses correct original filename."""
|
||||
staged_path, content = sample_pdf
|
||||
content_hash = hashlib.sha256(content).hexdigest()
|
||||
documents_dir = data_dir / "documents"
|
||||
|
||||
dest = documents_dir / f"{content_hash}.pdf"
|
||||
shutil.move(str(staged_path), str(dest))
|
||||
|
||||
db_conn.execute(
|
||||
"INSERT INTO documents(title, source_path, content_hash, doc_type, stored_path, original_filename) "
|
||||
"VALUES (?, ?, ?, ?, ?, ?)",
|
||||
("My Report", str(staged_path), content_hash, "pdf", str(dest), "quarterly_report.pdf"),
|
||||
)
|
||||
db_conn.commit()
|
||||
|
||||
doc = db_conn.execute("SELECT stored_path, original_filename, title FROM documents WHERE content_hash = ?", (content_hash,)).fetchone()
|
||||
|
||||
# Verify the original filename is preserved and different from title
|
||||
assert doc["original_filename"] == "quarterly_report.pdf"
|
||||
assert doc["title"] == "My Report"
|
||||
assert Path(doc["stored_path"]).exists()
|
||||
|
||||
def test_fallback_to_title_when_no_original_filename(self, data_dir, db_conn):
|
||||
"""Test that title+ext is used when original_filename is NULL."""
|
||||
documents_dir = data_dir / "documents"
|
||||
fake_file = documents_dir / "somehash.pdf"
|
||||
fake_file.write_bytes(b"fake")
|
||||
|
||||
db_conn.execute(
|
||||
"INSERT INTO documents(title, source_path, content_hash, doc_type, stored_path) "
|
||||
"VALUES (?, ?, ?, ?, ?)",
|
||||
("Engine Manual", "/tmp/old", "hash456", "pdf", str(fake_file)),
|
||||
)
|
||||
db_conn.commit()
|
||||
|
||||
doc = db_conn.execute("SELECT original_filename, title, stored_path FROM documents WHERE content_hash = 'hash456'").fetchone()
|
||||
|
||||
# When original_filename is NULL, the endpoint should fall back to title + ext
|
||||
original_filename = doc["original_filename"]
|
||||
if not original_filename:
|
||||
ext = Path(doc["stored_path"]).suffix
|
||||
original_filename = (doc["title"] or "document") + ext
|
||||
|
||||
assert original_filename == "Engine Manual.pdf"
|
||||
@@ -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"] == {}
|
||||
@@ -0,0 +1,17 @@
|
||||
FROM python:3.12-slim
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
COPY requirements.txt ./
|
||||
RUN pip install --no-cache-dir -r requirements.txt
|
||||
|
||||
COPY *.py ./
|
||||
|
||||
ENV KB_ENGINE_URL=http://engine:8000
|
||||
ENV KB_API_KEY=
|
||||
ENV KB_MCP_API_KEY=
|
||||
ENV KB_MCP_PORT=3000
|
||||
|
||||
EXPOSE 3000
|
||||
|
||||
CMD ["python", "server.py"]
|
||||
@@ -0,0 +1,17 @@
|
||||
"""Configuration from environment variables."""
|
||||
|
||||
import os
|
||||
|
||||
|
||||
KB_ENGINE_URL = os.environ.get("KB_ENGINE_URL", "http://localhost:8000")
|
||||
KB_API_KEY = os.environ.get("KB_API_KEY", "")
|
||||
KB_MCP_API_KEY = os.environ.get("KB_MCP_API_KEY", "")
|
||||
KB_MCP_PORT = int(os.environ.get("KB_MCP_PORT", "3000"))
|
||||
KB_MCP_ALLOWED_HOSTS = os.environ.get("KB_MCP_ALLOWED_HOSTS", "")
|
||||
|
||||
|
||||
def parse_allowed_hosts() -> list[str]:
|
||||
"""Parse KB_MCP_ALLOWED_HOSTS into a list of host strings."""
|
||||
if not KB_MCP_ALLOWED_HOSTS:
|
||||
return []
|
||||
return [h.strip() for h in KB_MCP_ALLOWED_HOSTS.split(",") if h.strip()]
|
||||
+213
@@ -0,0 +1,213 @@
|
||||
"""HTTP client for the kb engine API."""
|
||||
|
||||
import httpx
|
||||
|
||||
from config import KB_ENGINE_URL, KB_API_KEY
|
||||
|
||||
|
||||
def _auth_headers() -> dict[str, str]:
|
||||
h: dict[str, str] = {}
|
||||
if KB_API_KEY:
|
||||
h["Authorization"] = f"Bearer {KB_API_KEY}"
|
||||
return h
|
||||
|
||||
|
||||
def _client() -> httpx.Client:
|
||||
return httpx.Client(base_url=KB_ENGINE_URL, headers=_auth_headers(), timeout=60.0)
|
||||
|
||||
|
||||
def search(query: str, top: int = 10, tags: list[str] | None = None,
|
||||
doc_type: str | None = None, fts_only: bool = False,
|
||||
vec_only: bool = False, threshold: float | None = None,
|
||||
explain: bool = False, rerank: bool | None = None) -> dict:
|
||||
body: dict = {"query": query, "top": top}
|
||||
if tags:
|
||||
body["tags"] = tags
|
||||
if doc_type:
|
||||
body["doc_type"] = doc_type
|
||||
if fts_only:
|
||||
body["fts_only"] = True
|
||||
if vec_only:
|
||||
body["vec_only"] = True
|
||||
if threshold is not None:
|
||||
body["threshold"] = threshold
|
||||
if explain:
|
||||
body["explain"] = True
|
||||
if rerank is not None:
|
||||
body["rerank"] = rerank
|
||||
with _client() as c:
|
||||
r = c.post("/api/v1/search", json=body)
|
||||
r.raise_for_status()
|
||||
return r.json()
|
||||
|
||||
|
||||
def add_note(text: str, tags: list[str] | None = None,
|
||||
title: str | None = None) -> dict:
|
||||
fields = {"note": text}
|
||||
if tags:
|
||||
fields["tags"] = ",".join(tags)
|
||||
if title:
|
||||
fields["title"] = title
|
||||
with _client() as c:
|
||||
r = c.post("/api/v1/jobs", data=fields)
|
||||
r.raise_for_status()
|
||||
return r.json()
|
||||
|
||||
|
||||
def update_note(doc_id: int, text: str) -> dict:
|
||||
with _client() as c:
|
||||
r = c.patch(f"/api/v1/notes/{doc_id}", json={"text": text})
|
||||
r.raise_for_status()
|
||||
return r.json()
|
||||
|
||||
|
||||
def get_document(doc_id: int) -> dict:
|
||||
with _client() as c:
|
||||
r = c.get(f"/api/v1/documents/{doc_id}")
|
||||
r.raise_for_status()
|
||||
return r.json()
|
||||
|
||||
|
||||
def list_documents(doc_type: str | None = None,
|
||||
tags: str | None = None) -> list[dict]:
|
||||
params: dict = {}
|
||||
if doc_type:
|
||||
params["type"] = doc_type
|
||||
if tags:
|
||||
params["tags"] = tags
|
||||
with _client() as c:
|
||||
r = c.get("/api/v1/documents", params=params)
|
||||
r.raise_for_status()
|
||||
return r.json()
|
||||
|
||||
|
||||
def get_status() -> dict:
|
||||
with _client() as c:
|
||||
r = c.get("/api/v1/status")
|
||||
r.raise_for_status()
|
||||
return r.json()
|
||||
|
||||
|
||||
def list_jobs(status: str | None = None) -> list[dict]:
|
||||
params: dict = {}
|
||||
if status:
|
||||
params["status"] = status
|
||||
with _client() as c:
|
||||
r = c.get("/api/v1/jobs", params=params)
|
||||
r.raise_for_status()
|
||||
return r.json()
|
||||
|
||||
|
||||
def update_tags(doc_id: int, add: list[str] | None = None,
|
||||
remove: list[str] | None = None) -> dict:
|
||||
body: dict = {}
|
||||
if add:
|
||||
body["add"] = add
|
||||
if remove:
|
||||
body["remove"] = remove
|
||||
with _client() as c:
|
||||
r = c.put(f"/api/v1/documents/{doc_id}/tags", json=body)
|
||||
r.raise_for_status()
|
||||
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,
|
||||
tags: list[str] | None = None) -> dict:
|
||||
fields: dict = {}
|
||||
if tags:
|
||||
fields["tags"] = ",".join(tags)
|
||||
with _client() as c:
|
||||
r = c.post(
|
||||
"/api/v1/jobs",
|
||||
data=fields,
|
||||
files={"file": (filename, file_bytes)},
|
||||
)
|
||||
r.raise_for_status()
|
||||
return r.json()
|
||||
@@ -0,0 +1,4 @@
|
||||
mcp>=1.9.0
|
||||
httpx>=0.27
|
||||
uvicorn>=0.30
|
||||
starlette>=0.38
|
||||
+478
@@ -0,0 +1,478 @@
|
||||
"""kb MCP server — exposes knowledge base operations as MCP tools."""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
|
||||
from mcp.server.fastmcp import FastMCP
|
||||
from mcp.server.transport_security import TransportSecuritySettings
|
||||
from starlette.applications import Starlette
|
||||
from starlette.middleware import Middleware
|
||||
from starlette.middleware.base import BaseHTTPMiddleware
|
||||
from starlette.requests import Request
|
||||
from starlette.responses import JSONResponse
|
||||
from starlette.routing import Mount
|
||||
|
||||
import config
|
||||
import engine
|
||||
import uploads
|
||||
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
|
||||
logger = logging.getLogger("kb.mcp")
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Transport security — DNS rebinding protection with configurable allowed hosts
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_LOCALHOST_HOSTS = ["127.0.0.1:*", "localhost:*", "[::1]:*"]
|
||||
_LOCALHOST_ORIGINS = ["http://127.0.0.1:*", "http://localhost:*", "http://[::1]:*"]
|
||||
|
||||
_extra_hosts = config.parse_allowed_hosts()
|
||||
_allowed_hosts = _LOCALHOST_HOSTS + [f"{h}:*" for h in _extra_hosts]
|
||||
_allowed_origins = _LOCALHOST_ORIGINS + [f"http://{h}:*" for h in _extra_hosts]
|
||||
|
||||
_transport_security = TransportSecuritySettings(
|
||||
enable_dns_rebinding_protection=True,
|
||||
allowed_hosts=_allowed_hosts,
|
||||
allowed_origins=_allowed_origins,
|
||||
)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# FastMCP server
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
mcp = FastMCP(
|
||||
"kb",
|
||||
instructions=(
|
||||
"Knowledge base MCP server with hybrid semantic + full-text search. "
|
||||
"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 "
|
||||
"header at the HTTP transport layer."
|
||||
),
|
||||
transport_security=_transport_security,
|
||||
)
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
async def kb_search(
|
||||
query: str,
|
||||
top: int = 10,
|
||||
tags: list[str] | None = None,
|
||||
doc_type: str | None = None,
|
||||
fts_only: bool = False,
|
||||
explain: bool = False,
|
||||
rerank: bool | None = None,
|
||||
) -> str:
|
||||
"""Hybrid semantic (vector) + full-text search over the knowledge base.
|
||||
|
||||
Combines dense vector embeddings (semantic similarity — finds conceptually
|
||||
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:
|
||||
query: The search query — a natural language question or keywords.
|
||||
top: Maximum number of results to return (default 10).
|
||||
tags: Filter results to documents with ALL of these tags.
|
||||
doc_type: Filter by document type (e.g. "note", "pdf", "markdown",
|
||||
"code", "data").
|
||||
fts_only: Disable the vector/semantic component and use only BM25
|
||||
keyword matching. Default false (hybrid mode). Set true only when
|
||||
you need exact-string matching (e.g. an error code, identifier).
|
||||
explain: Include a per-result score breakdown (BM25 score/rank, vector
|
||||
similarity/rank, rank-fusion contributions, rerank blend) under an
|
||||
"explain" key. Useful for diagnosing why a result ranked where it did.
|
||||
rerank: Set false to skip server-side reranking for lower latency.
|
||||
Default (None) uses the engine's configured behaviour.
|
||||
|
||||
Tips for complex queries:
|
||||
- Consider expanding into 2-3 variant phrasings and calling this tool multiple
|
||||
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.
|
||||
- If the engine's reranker is disabled, you can still rerank the returned
|
||||
results yourself using your own judgement of relevance to the question.
|
||||
- Call kb_status to see which embedding model is in use and whether
|
||||
server-side reranking is active.
|
||||
"""
|
||||
result = engine.search(
|
||||
query=query,
|
||||
top=top,
|
||||
tags=tags or None,
|
||||
doc_type=doc_type,
|
||||
fts_only=fts_only,
|
||||
explain=explain,
|
||||
rerank=rerank,
|
||||
)
|
||||
|
||||
results_list = result if isinstance(result, list) else result.get("results", [])
|
||||
return json.dumps(results_list, indent=2)
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
async def kb_addnote(
|
||||
text: str,
|
||||
tags: list[str] | None = None,
|
||||
title: str | None = None,
|
||||
) -> str:
|
||||
"""Add a text note to the knowledge base for indexing and search.
|
||||
|
||||
The note is queued for ingestion — it will be chunked, embedded, and made
|
||||
searchable. Use kb_jobs to check ingestion status.
|
||||
|
||||
Args:
|
||||
text: The note text content.
|
||||
tags: Tags to apply to the note.
|
||||
title: Optional title (auto-derived from first line if omitted).
|
||||
"""
|
||||
result = engine.add_note(text=text, tags=tags or None, title=title)
|
||||
return json.dumps(result, indent=2)
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
async def kb_update_note(
|
||||
document_id: int,
|
||||
text: str,
|
||||
) -> str:
|
||||
"""Update an existing note's content in place.
|
||||
|
||||
Replaces the note text, re-chunks, and re-embeds while preserving the
|
||||
document ID, creation timestamp, and tags. Only works on documents with
|
||||
doc_type "note".
|
||||
|
||||
Args:
|
||||
document_id: The ID of the note document to update.
|
||||
text: The new text content for the note.
|
||||
"""
|
||||
result = engine.update_note(document_id, text)
|
||||
return json.dumps(result, indent=2)
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
async def kb_get(
|
||||
document_id: int | None = None,
|
||||
source_path: str | None = None,
|
||||
) -> str:
|
||||
"""Retrieve document details from the knowledge base.
|
||||
|
||||
Look up a document by its ID or source path. Returns full document metadata,
|
||||
tags, and chunk contents.
|
||||
|
||||
Args:
|
||||
document_id: The numeric document ID.
|
||||
source_path: The document's source path (alternative to document_id).
|
||||
"""
|
||||
if document_id is not None:
|
||||
result = engine.get_document(document_id)
|
||||
return json.dumps(result, indent=2)
|
||||
elif source_path is not None:
|
||||
docs = engine.list_documents()
|
||||
matches = [d for d in docs if d.get("source_path") == source_path]
|
||||
if not matches:
|
||||
return json.dumps({"error": "No document found with that source_path"})
|
||||
doc = engine.get_document(matches[0]["id"])
|
||||
return json.dumps(doc, indent=2)
|
||||
else:
|
||||
return json.dumps({"error": "Provide either document_id or source_path"})
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
async def kb_status() -> str:
|
||||
"""Get knowledge base engine status.
|
||||
|
||||
Returns engine version, embedding model info, device info, document counts,
|
||||
database size, and ingestion queue state.
|
||||
"""
|
||||
result = engine.get_status()
|
||||
result["authenticated"] = bool(config.KB_MCP_API_KEY)
|
||||
return json.dumps(result, indent=2)
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
async def kb_jobs(
|
||||
status: str | None = None,
|
||||
) -> str:
|
||||
"""List ingestion jobs and their status.
|
||||
|
||||
Returns recent jobs showing what has been queued, is processing, completed,
|
||||
or failed.
|
||||
|
||||
Args:
|
||||
status: Filter by job status ("queued", "processing", "done", "failed", "skipped").
|
||||
"""
|
||||
result = engine.list_jobs(status=status)
|
||||
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()
|
||||
async def kb_upload_start(
|
||||
filename: str,
|
||||
total_size: int,
|
||||
tags: list[str] | None = None,
|
||||
) -> str:
|
||||
"""Start a chunked file upload to the knowledge base.
|
||||
|
||||
Use this for uploading files from a remote agent. The upload process is:
|
||||
1. Call kb_upload_start to get an upload_id
|
||||
2. Call kb_upload_chunk repeatedly with base64-encoded file chunks (recommended ~1MB each)
|
||||
3. Call kb_upload_finish to submit the file for ingestion
|
||||
|
||||
Example for a 3MB file:
|
||||
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 1>", chunk_index=1)
|
||||
kb_upload_chunk(upload_id=upload["upload_id"], data="<base64 chunk 2>", chunk_index=2)
|
||||
result = kb_upload_finish(upload_id=upload["upload_id"])
|
||||
|
||||
Args:
|
||||
filename: Original filename (used for type detection).
|
||||
total_size: Total file size in bytes.
|
||||
tags: Tags to apply to the uploaded document.
|
||||
"""
|
||||
upload_id = uploads.start_upload(filename, total_size, tags or [])
|
||||
return json.dumps({"upload_id": upload_id})
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
async def kb_upload_chunk(
|
||||
upload_id: str,
|
||||
data: str,
|
||||
chunk_index: int,
|
||||
) -> str:
|
||||
"""Upload a base64-encoded chunk of a file.
|
||||
|
||||
Part of the chunked upload flow started by kb_upload_start.
|
||||
|
||||
Args:
|
||||
upload_id: The upload ID from kb_upload_start.
|
||||
data: Base64-encoded file data for this chunk.
|
||||
chunk_index: Zero-based index of this chunk.
|
||||
"""
|
||||
try:
|
||||
uploads.add_chunk(upload_id, data, chunk_index)
|
||||
return json.dumps({"status": "ok", "chunk_index": chunk_index})
|
||||
except KeyError as e:
|
||||
return json.dumps({"error": str(e)})
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
async def kb_upload_finish(
|
||||
upload_id: str,
|
||||
) -> str:
|
||||
"""Finish a chunked upload and submit the file for ingestion.
|
||||
|
||||
Reassembles all uploaded chunks and forwards the complete file to the
|
||||
engine for processing. Returns the ingestion job ID.
|
||||
|
||||
Args:
|
||||
upload_id: The upload ID from kb_upload_start.
|
||||
"""
|
||||
try:
|
||||
filename, file_bytes, tags = uploads.finish_upload(upload_id)
|
||||
result = engine.upload_file(filename, file_bytes, tags)
|
||||
return json.dumps(result, indent=2)
|
||||
except KeyError as 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
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class BearerAuthMiddleware(BaseHTTPMiddleware):
|
||||
async def dispatch(self, request: Request, call_next):
|
||||
if not config.KB_MCP_API_KEY:
|
||||
return await call_next(request)
|
||||
|
||||
auth_header = request.headers.get("authorization", "")
|
||||
if auth_header.startswith("Bearer ") and auth_header[7:] == config.KB_MCP_API_KEY:
|
||||
return await call_next(request)
|
||||
|
||||
return JSONResponse(
|
||||
status_code=401,
|
||||
content={"error": "Unauthorized"},
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# ASGI app assembly
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def create_app():
|
||||
"""Create the ASGI app with auth middleware wrapping the MCP server."""
|
||||
from contextlib import asynccontextmanager
|
||||
|
||||
mcp_app = mcp.streamable_http_app()
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app):
|
||||
uploads.start_cleanup_task()
|
||||
logger.info("Upload cleanup task started")
|
||||
# Delegate to the MCP app's lifespan if it has one
|
||||
if hasattr(mcp_app, 'router') and hasattr(mcp_app.router, 'lifespan_context'):
|
||||
async with mcp_app.router.lifespan_context(app):
|
||||
yield
|
||||
else:
|
||||
yield
|
||||
|
||||
app = Starlette(
|
||||
routes=[Mount("/", app=mcp_app)],
|
||||
middleware=[Middleware(BearerAuthMiddleware)],
|
||||
lifespan=lifespan,
|
||||
)
|
||||
return app
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Entry point
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
if __name__ == "__main__":
|
||||
import uvicorn
|
||||
|
||||
logger.info(
|
||||
"Starting kb MCP server on port %d, engine=%s",
|
||||
config.KB_MCP_PORT,
|
||||
config.KB_ENGINE_URL,
|
||||
)
|
||||
|
||||
app = create_app()
|
||||
uvicorn.run(app, host="0.0.0.0", port=config.KB_MCP_PORT)
|
||||
@@ -0,0 +1,96 @@
|
||||
"""Chunked upload staging management."""
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import logging
|
||||
import shutil
|
||||
import tempfile
|
||||
import time
|
||||
import uuid
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
|
||||
logger = logging.getLogger("kb.mcp.uploads")
|
||||
|
||||
UPLOAD_TIMEOUT_SECONDS = 600 # 10 minutes
|
||||
|
||||
|
||||
@dataclass
|
||||
class StagedUpload:
|
||||
upload_id: str
|
||||
filename: str
|
||||
total_size: int
|
||||
tags: list[str]
|
||||
staging_dir: Path
|
||||
created_at: float = field(default_factory=time.time)
|
||||
chunks: dict[int, Path] = field(default_factory=dict)
|
||||
|
||||
|
||||
_uploads: dict[str, StagedUpload] = {}
|
||||
_cleanup_task: asyncio.Task | None = None
|
||||
|
||||
|
||||
def start_upload(filename: str, total_size: int, tags: list[str]) -> str:
|
||||
upload_id = str(uuid.uuid4())
|
||||
staging_dir = Path(tempfile.mkdtemp(prefix=f"kb_upload_{upload_id[:8]}_"))
|
||||
_uploads[upload_id] = StagedUpload(
|
||||
upload_id=upload_id,
|
||||
filename=filename,
|
||||
total_size=total_size,
|
||||
tags=tags,
|
||||
staging_dir=staging_dir,
|
||||
)
|
||||
logger.info("Started upload %s for %s (%d bytes)", upload_id, filename, total_size)
|
||||
return upload_id
|
||||
|
||||
|
||||
def add_chunk(upload_id: str, data_b64: str, chunk_index: int) -> None:
|
||||
upload = _uploads.get(upload_id)
|
||||
if upload is None:
|
||||
raise KeyError(f"Upload ID not found: {upload_id}")
|
||||
chunk_bytes = base64.b64decode(data_b64)
|
||||
chunk_path = upload.staging_dir / f"chunk_{chunk_index:06d}"
|
||||
chunk_path.write_bytes(chunk_bytes)
|
||||
upload.chunks[chunk_index] = chunk_path
|
||||
logger.info("Added chunk %d to upload %s (%d bytes)", chunk_index, upload_id, len(chunk_bytes))
|
||||
|
||||
|
||||
def finish_upload(upload_id: str) -> tuple[str, bytes, list[str]]:
|
||||
"""Reassemble chunks and return (filename, file_bytes, tags)."""
|
||||
upload = _uploads.get(upload_id)
|
||||
if upload is None:
|
||||
raise KeyError(f"Upload ID not found: {upload_id}")
|
||||
try:
|
||||
parts = []
|
||||
for idx in sorted(upload.chunks.keys()):
|
||||
parts.append(upload.chunks[idx].read_bytes())
|
||||
file_bytes = b"".join(parts)
|
||||
return upload.filename, file_bytes, upload.tags
|
||||
finally:
|
||||
_cleanup_upload(upload_id)
|
||||
|
||||
|
||||
def _cleanup_upload(upload_id: str) -> None:
|
||||
upload = _uploads.pop(upload_id, None)
|
||||
if upload and upload.staging_dir.exists():
|
||||
shutil.rmtree(upload.staging_dir, ignore_errors=True)
|
||||
|
||||
|
||||
async def cleanup_abandoned_uploads() -> None:
|
||||
"""Background task that removes uploads older than the timeout."""
|
||||
while True:
|
||||
await asyncio.sleep(60)
|
||||
now = time.time()
|
||||
expired = [
|
||||
uid for uid, u in _uploads.items()
|
||||
if now - u.created_at > UPLOAD_TIMEOUT_SECONDS
|
||||
]
|
||||
for uid in expired:
|
||||
logger.warning("Cleaning up abandoned upload %s", uid)
|
||||
_cleanup_upload(uid)
|
||||
|
||||
|
||||
def start_cleanup_task() -> None:
|
||||
global _cleanup_task
|
||||
if _cleanup_task is None or _cleanup_task.done():
|
||||
_cleanup_task = asyncio.create_task(cleanup_abandoned_uploads())
|
||||
@@ -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,2 @@
|
||||
schema: spec-driven
|
||||
created: 2026-03-27
|
||||
@@ -0,0 +1,84 @@
|
||||
## Context
|
||||
|
||||
Currently, uploaded files pass through a staging directory and are deleted after the worker extracts chunks and embeddings. The `documents.source_path` column stores the (now-stale) staging path. Users who want the original file must re-source it externally. The data directory structure today is:
|
||||
|
||||
```
|
||||
/data/
|
||||
kb.db
|
||||
hf_cache/
|
||||
staging/ # temporary, cleaned after processing
|
||||
```
|
||||
|
||||
## Goals / Non-Goals
|
||||
|
||||
**Goals:**
|
||||
- Persist every successfully-ingested original file for the lifetime of the document
|
||||
- Serve the original file via API (`GET /api/v1/documents/{id}/file`)
|
||||
- Clean up stored files when a document is deleted
|
||||
- Work transparently with the existing Docker volume mount (`/data`)
|
||||
|
||||
**Non-Goals:**
|
||||
- Serving transformed/converted versions of documents (e.g. PDF→HTML)
|
||||
- De-duplicating file storage (same content hash = same row, so 1:1 is fine)
|
||||
- Compression or archival of stored files
|
||||
- Retroactive storage of files ingested before this change (they're already gone)
|
||||
|
||||
## Decisions
|
||||
|
||||
### 1. Storage layout: content-hash-based flat directory
|
||||
|
||||
Store files at `{data_dir}/documents/{content_hash}{ext}` (e.g. `documents/a1b2c3...d4.pdf`).
|
||||
|
||||
**Why over document-ID naming:** Content hash is available at staging time before the DB row exists, avoids race conditions, and makes dedup trivially safe (same hash = same file, overwrite is harmless). The hash is already computed for dedup checks.
|
||||
|
||||
**Why flat over nested:** The KB is a personal tool — expected scale is hundreds to low-thousands of documents. A flat directory is simpler and sufficient. If needed later, a `ab/cd/` prefix scheme is easy to add.
|
||||
|
||||
**Alternatives considered:**
|
||||
- *Store in SQLite as BLOBs*: Bloats the DB, complicates backups, and degrades WAL performance for large files. Rejected.
|
||||
- *Keep the staging path as-is*: Staging uses UUID prefixes which are meaningless; content-hash naming is deterministic and self-deduplicating.
|
||||
|
||||
### 2. Move file from staging to documents dir (not copy)
|
||||
|
||||
Use `shutil.move()` from staging to documents dir after successful ingestion, before `staging.cleanup()`. This avoids doubling disk usage during processing.
|
||||
|
||||
**Why not copy-then-delete:** Move is atomic on the same filesystem (which `/data/staging` and `/data/documents` share). Faster, no temporary disk spike.
|
||||
|
||||
### 3. New columns `stored_path` and `original_filename` on `documents` table
|
||||
|
||||
Add two nullable columns:
|
||||
- `stored_path TEXT` — permanent file location on disk
|
||||
- `original_filename TEXT` — the exact filename from the upload (e.g. `report.pdf`)
|
||||
|
||||
Both are nullable because existing documents (ingested before this change) won't have values.
|
||||
|
||||
**Why `original_filename` separate from `title`:** The `title` field can be user-overridden (e.g. "Engine Manual" instead of `report.pdf`). When serving the file for download, the `Content-Disposition` header should use the original filename so the downloaded file has the correct name and extension. The `original_filename` is sourced from `jobs.filename` which is already captured at upload time.
|
||||
|
||||
Keep `source_path` as-is for backward compatibility (it records what the staging path was). `stored_path` is the permanent location.
|
||||
|
||||
**Migration:** Two `ALTER TABLE` statements — safe additive migrations, no data rewrite needed.
|
||||
|
||||
### 4. File download endpoint returns the file directly
|
||||
|
||||
`GET /api/v1/documents/{id}/file` uses FastAPI's `FileResponse` with:
|
||||
- `media_type` derived from the file extension
|
||||
- `Content-Disposition: attachment; filename="{original_filename}"` (falls back to `{title}{ext}` if `original_filename` is NULL)
|
||||
- Returns 404 if `stored_path` is NULL or file is missing from disk
|
||||
|
||||
### 5. Delete cascades to file removal
|
||||
|
||||
When `DELETE /api/v1/documents/{id}` is called, delete the stored file from disk after the DB delete succeeds. If file removal fails (already gone, permissions), log a warning but don't fail the API call — the DB is the source of truth.
|
||||
|
||||
## Risks / Trade-offs
|
||||
|
||||
- **Disk usage increases** — every ingested file persists. For the personal-use scale this is expected and acceptable. Users manage this via document deletion.
|
||||
→ Mitigation: Document the storage behavior; `GET /api/v1/status` already shows DB size, could add documents-dir size later.
|
||||
|
||||
- **Pre-existing documents have no stored file** — `stored_path` will be NULL for documents ingested before this change.
|
||||
→ Mitigation: The download endpoint returns 404 with a clear message ("original file not available — ingested before document storage was enabled"). No attempt to backfill.
|
||||
|
||||
- **File-DB consistency** — crash between DB commit and file move could leave orphan staged files or missing stored files.
|
||||
→ Mitigation: Move file first, then commit DB. If DB commit fails, the file in documents dir is harmless (orphan cleanup can be added later). If move fails, the job fails and staged file remains for retry.
|
||||
|
||||
## Open Questions
|
||||
|
||||
None — the scope is straightforward enough to proceed.
|
||||
@@ -0,0 +1,30 @@
|
||||
## Why
|
||||
|
||||
The knowledge base currently discards original files after chunking and embedding. Once a document is ingested, only the extracted text chunks and vectors remain — the original PDF, markdown, or code file is deleted from staging. Users cannot retrieve the source document from the KB, which limits its usefulness as a document store and prevents use cases like re-processing with a different model or serving the original file to downstream tools.
|
||||
|
||||
## What Changes
|
||||
|
||||
- Add a persistent document storage directory (`{data_dir}/documents/`) alongside the SQLite database
|
||||
- After successful ingestion, copy the original file from staging to permanent storage instead of deleting it
|
||||
- Store the permanent file path in the `documents` table (`stored_path` column) and the original upload filename (`original_filename` column) so downloads use the correct name
|
||||
- Add an API endpoint to download the original file by document ID
|
||||
- Add a CLI command to export/retrieve the original document
|
||||
- **BREAKING**: Delete document now also removes the stored file from disk
|
||||
- Notes (text-only) are stored as `.note` files in the same directory for consistency
|
||||
|
||||
## Capabilities
|
||||
|
||||
### New Capabilities
|
||||
- `document-storage`: Persistent storage of original uploaded files on disk, lifecycle management (store on ingest, delete on document removal), and retrieval via API
|
||||
|
||||
### Modified Capabilities
|
||||
- `engine-api`: New endpoint `GET /api/v1/documents/{id}/file` to download the original file; delete endpoint must also clean up stored files; ingestion worker stores files instead of discarding them
|
||||
|
||||
## Impact
|
||||
|
||||
- **Engine config**: New `documents_dir` property on Config, new directory created at startup via `ensure_dirs()`
|
||||
- **Worker**: After successful chunking, move/copy file from staging to documents dir; update `source_path` → `stored_path` with permanent location
|
||||
- **Database schema**: Add `stored_path` and `original_filename` columns to `documents` table (migration for existing DBs)
|
||||
- **Routes**: New file-download endpoint; update delete handler to remove stored file
|
||||
- **Go client**: New `export` / `get-file` subcommand to download original documents
|
||||
- **Docker**: `documents/` directory lives inside the existing `/data` volume — no new mounts needed
|
||||
+83
@@ -0,0 +1,83 @@
|
||||
## ADDED Requirements
|
||||
|
||||
### Requirement: Persistent original file storage
|
||||
|
||||
The engine SHALL persistently store the original uploaded file on disk after successful ingestion. Files SHALL be stored at `{data_dir}/documents/{content_hash}{extension}` where `content_hash` is the SHA-256 hex digest already computed for dedup and `extension` is preserved from the original filename. The `documents` table SHALL record the stored file path in a `stored_path` column and the original upload filename in an `original_filename` column.
|
||||
|
||||
#### Scenario: File stored after successful ingestion
|
||||
- **WHEN** the background worker successfully processes an ingestion job for a PDF file
|
||||
- **THEN** the worker SHALL move the staged file to `{data_dir}/documents/{content_hash}.pdf`, store the permanent path in `documents.stored_path`, store the original filename in `documents.original_filename`, and delete the staging entry
|
||||
|
||||
#### Scenario: Note stored after successful ingestion
|
||||
- **WHEN** the background worker successfully processes an ingestion job for a text note
|
||||
- **THEN** the worker SHALL move the staged `.note` file to `{data_dir}/documents/{content_hash}.note` and store the permanent path in `documents.stored_path`
|
||||
|
||||
#### Scenario: Markdown file stored after successful ingestion
|
||||
- **WHEN** the background worker successfully processes an ingestion job for a markdown file
|
||||
- **THEN** the worker SHALL move the staged file to `{data_dir}/documents/{content_hash}.md` and store the permanent path in `documents.stored_path`
|
||||
|
||||
#### Scenario: Code file stored after successful ingestion
|
||||
- **WHEN** the background worker successfully processes an ingestion job for a code file (e.g. `.py`, `.go`)
|
||||
- **THEN** the worker SHALL move the staged file to `{data_dir}/documents/{content_hash}{original_extension}` and store the permanent path in `documents.stored_path`
|
||||
|
||||
#### Scenario: Documents directory created at startup
|
||||
- **WHEN** the engine starts up and calls `ensure_dirs()`
|
||||
- **THEN** the `{data_dir}/documents/` directory SHALL be created if it does not exist
|
||||
|
||||
#### Scenario: Ingestion failure does not store file
|
||||
- **WHEN** the background worker fails to process an ingestion job
|
||||
- **THEN** the staged file SHALL be cleaned up as before and no file SHALL be written to the documents directory
|
||||
|
||||
---
|
||||
|
||||
### Requirement: File retrieval via API
|
||||
|
||||
The engine SHALL serve the original stored file for any document that has a stored file on disk.
|
||||
|
||||
#### Scenario: Download original file
|
||||
- **WHEN** a client sends `GET /api/v1/documents/{id}/file` for a document with a stored file
|
||||
- **THEN** the engine SHALL return the file with appropriate `Content-Type` based on file extension and `Content-Disposition: attachment; filename="{original_filename}"` header, falling back to `{title}{ext}` if `original_filename` is NULL
|
||||
|
||||
#### Scenario: Download file for pre-existing document
|
||||
- **WHEN** a client sends `GET /api/v1/documents/{id}/file` for a document ingested before this feature was added (stored_path is NULL)
|
||||
- **THEN** the engine SHALL return HTTP 404 with `{"error": "Original file not available - ingested before document storage was enabled"}`
|
||||
|
||||
#### Scenario: Download file when file missing from disk
|
||||
- **WHEN** a client sends `GET /api/v1/documents/{id}/file` for a document whose `stored_path` is set but the file no longer exists on disk
|
||||
- **THEN** the engine SHALL return HTTP 404 with `{"error": "Stored file not found on disk"}`
|
||||
|
||||
#### Scenario: Download file for non-existent document
|
||||
- **WHEN** a client sends `GET /api/v1/documents/{id}/file` with a non-existent document ID
|
||||
- **THEN** the engine SHALL return HTTP 404 with `{"error": "Document not found"}`
|
||||
|
||||
---
|
||||
|
||||
### Requirement: File cleanup on document deletion
|
||||
|
||||
The engine SHALL remove the stored original file from disk when a document is deleted.
|
||||
|
||||
#### Scenario: Delete document with stored file
|
||||
- **WHEN** a client sends `DELETE /api/v1/documents/{id}` for a document with a stored file
|
||||
- **THEN** the engine SHALL delete the document from the database (cascading to chunks, embeddings, tags) AND delete the stored file from disk
|
||||
|
||||
#### Scenario: Delete document when stored file already missing
|
||||
- **WHEN** a client sends `DELETE /api/v1/documents/{id}` for a document whose stored file has been manually removed from disk
|
||||
- **THEN** the engine SHALL delete the document from the database successfully and log a warning about the missing file
|
||||
|
||||
#### Scenario: Delete document without stored file (pre-existing)
|
||||
- **WHEN** a client sends `DELETE /api/v1/documents/{id}` for a document with `stored_path` NULL
|
||||
- **THEN** the engine SHALL delete the document from the database without attempting file removal
|
||||
|
||||
---
|
||||
|
||||
### Requirement: Database schema migration for stored_path and original_filename
|
||||
|
||||
The engine SHALL add `stored_path` and `original_filename` columns to the `documents` table for tracking permanent file locations and original upload filenames.
|
||||
|
||||
#### Scenario: Fresh database initialization
|
||||
- **WHEN** the engine initializes a new database
|
||||
- **THEN** the `documents` table SHALL include `stored_path TEXT` and `original_filename TEXT` columns in its schema
|
||||
|
||||
#### Scenario: Existing database migration
|
||||
- **WHEN** the engine starts with a database created before this feature
|
||||
- **THEN** the engine SHALL add `stored_path TEXT` and `original_filename TEXT` to the `documents` table via `ALTER TABLE` if the columns do not exist
|
||||
@@ -0,0 +1,61 @@
|
||||
## MODIFIED Requirements
|
||||
|
||||
### 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), generate embeddings using the resident model, insert chunks and vectors into the database, and move the original file to persistent storage.
|
||||
|
||||
#### Scenario: Successful PDF ingestion
|
||||
- **WHEN** the background worker picks up a queued PDF job
|
||||
- **THEN** it SHALL update the job status to `processing`, run Docling conversion and chunking, embed all chunks, insert document and chunks into the database, move the staged file to `{data_dir}/documents/{content_hash}.pdf`, update `documents.stored_path` with the permanent path, store the original filename in `documents.original_filename`, update the job status to `done` with the resulting document_id and chunk count, and clean up the staging entry
|
||||
|
||||
#### Scenario: Ingestion failure
|
||||
- **WHEN** the background worker encounters an error during processing (e.g., corrupt PDF)
|
||||
- **THEN** it SHALL update the job status to `failed` with the error message, delete the staged file, and continue processing the next queued job
|
||||
|
||||
#### Scenario: Search during active ingestion
|
||||
- **WHEN** a search request arrives while the background worker is processing a job
|
||||
- **THEN** the search SHALL execute without blocking (SQLite WAL mode) and return results from already-ingested documents
|
||||
|
||||
---
|
||||
|
||||
### Requirement: Document management
|
||||
|
||||
The engine SHALL provide endpoints to list, inspect, remove, and download original files for ingested documents.
|
||||
|
||||
#### Scenario: List 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
|
||||
|
||||
#### Scenario: List documents with filters
|
||||
- **WHEN** a client sends `GET /api/v1/documents?type=pdf&tags=manual`
|
||||
- **THEN** the engine SHALL return only documents matching all specified filters
|
||||
|
||||
#### Scenario: Get document details
|
||||
- **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`)
|
||||
|
||||
#### Scenario: Download original file
|
||||
- **WHEN** a client sends `GET /api/v1/documents/{id}/file`
|
||||
- **THEN** the engine SHALL return the original file with appropriate Content-Type and `Content-Disposition: attachment; filename="{original_filename}"` headers, or HTTP 404 if the file is not available
|
||||
|
||||
#### Scenario: Remove a document
|
||||
- **WHEN** a client sends `DELETE /api/v1/documents/{id}`
|
||||
- **THEN** the engine SHALL delete the document, all its chunks, associated embeddings, tag associations, and the stored original file from disk, and return HTTP 200 with a confirmation
|
||||
|
||||
#### Scenario: Remove non-existent document
|
||||
- **WHEN** a client sends `DELETE /api/v1/documents/{id}` with a non-existent ID
|
||||
- **THEN** the engine SHALL return HTTP 404
|
||||
|
||||
---
|
||||
|
||||
### Requirement: Engine configuration via environment variables
|
||||
|
||||
The engine SHALL be configured via environment variables. No config file is read by the engine — all configuration comes from the environment (set via compose.yaml or Docker run).
|
||||
|
||||
#### Scenario: Default configuration
|
||||
- **WHEN** the engine starts with no environment variables set
|
||||
- **THEN** it SHALL use defaults: data directory `/data`, model `all-MiniLM-L6-v2`, device `auto`, no API key required. It SHALL create `staging/` and `documents/` subdirectories under the data directory.
|
||||
|
||||
#### Scenario: Custom model
|
||||
- **WHEN** `KB_MODEL` is set to `BAAI/bge-small-en-v1.5`
|
||||
- **THEN** the engine SHALL download and load that model instead of the default
|
||||
@@ -0,0 +1,38 @@
|
||||
## 1. Config and Schema
|
||||
|
||||
- [x] 1.1 Add `documents_dir` property to `Config` in `engine/kb/config.py` returning `{data_dir}/documents`
|
||||
- [x] 1.2 Add `documents_dir.mkdir()` to `Config.ensure_dirs()`
|
||||
- [x] 1.3 Add `stored_path TEXT` and `original_filename TEXT` columns to `documents` table in `init_schema()` (both CREATE TABLE and ALTER TABLE migration for existing DBs)
|
||||
|
||||
## 2. Worker — File Persistence
|
||||
|
||||
- [x] 2.1 In `worker._process_job()`, after successful DB commit, move staged file to `{documents_dir}/{content_hash}{ext}` using `shutil.move()`
|
||||
- [x] 2.2 Update `documents.stored_path` and `documents.original_filename` (from `jobs.filename`) after moving the file
|
||||
- [x] 2.3 Remove `staging.cleanup()` call for successful jobs (file is moved, not deleted); keep cleanup on failure path
|
||||
|
||||
## 3. API — File Download Endpoint
|
||||
|
||||
- [x] 3.1 Add `GET /api/v1/documents/{id}/file` route in `engine/kb/routes/documents.py` using FastAPI `FileResponse`
|
||||
- [x] 3.2 Return appropriate `Content-Type` from file extension and `Content-Disposition: attachment; filename="{original_filename}"` (fall back to `{title}{ext}` if NULL)
|
||||
- [x] 3.3 Handle 404 cases: document not found, `stored_path` is NULL, file missing from disk
|
||||
|
||||
## 4. API — Delete Cleanup
|
||||
|
||||
- [x] 4.1 Update `DELETE /api/v1/documents/{id}` in `engine/kb/routes/documents.py` to also delete the stored file from disk
|
||||
- [x] 4.2 Handle missing file gracefully (log warning, don't fail the request)
|
||||
|
||||
## 5. Document Details Enhancement
|
||||
|
||||
- [x] 5.1 Add `has_file` boolean to `GET /api/v1/documents/{id}` response based on `stored_path` presence and file existence on disk
|
||||
|
||||
## 6. Go Client
|
||||
|
||||
- [x] 6.1 Add `kb export <doc_id>` subcommand to the Go client that calls `GET /api/v1/documents/{id}/file` and writes to stdout or a specified output path
|
||||
|
||||
## 7. Testing
|
||||
|
||||
- [x] 7.1 Test successful ingestion stores file at expected path
|
||||
- [x] 7.2 Test failed ingestion does not leave file in documents dir
|
||||
- [x] 7.3 Test file download endpoint returns correct content and headers
|
||||
- [x] 7.4 Test document deletion removes stored file
|
||||
- [x] 7.5 Test download returns 404 for documents without stored files
|
||||
@@ -0,0 +1,2 @@
|
||||
schema: spec-driven
|
||||
created: 2026-03-29
|
||||
@@ -0,0 +1,23 @@
|
||||
## Context
|
||||
|
||||
The engine's `POST /api/v1/reindex` re-embeds all chunks synchronously and returns `{"chunks_reindexed": N, "model": "..."}`. The client has an established confirmation pattern in `remove.go` using `--yes`/`-y` flag.
|
||||
|
||||
## Goals / Non-Goals
|
||||
|
||||
**Goals:**
|
||||
- Add `kb reindex` with confirmation prompt matching `kb remove` pattern
|
||||
- Display human-readable and JSON output
|
||||
|
||||
**Non-Goals:**
|
||||
- Progress reporting during reindex (engine returns synchronously)
|
||||
- Model selection from the client (model is engine-side config)
|
||||
|
||||
## Decisions
|
||||
|
||||
### 1. Confirmation prompt before reindex
|
||||
|
||||
Reindex drops and rebuilds the vector table — destructive if interrupted. Use the same `[y/N]` prompt pattern as `kb remove`, skippable with `--yes`/`-y`.
|
||||
|
||||
### 2. Warn that it may take a while
|
||||
|
||||
The prompt should mention that reindex re-embeds all chunks, so the user knows it's not instant.
|
||||
@@ -0,0 +1,22 @@
|
||||
## Why
|
||||
|
||||
The engine exposes `POST /api/v1/reindex` but there's no client command for it. Users switching embedding models must use curl directly. Adding `kb reindex` with a confirmation prompt keeps it consistent with other destructive commands like `kb remove`.
|
||||
|
||||
## What Changes
|
||||
|
||||
- Add `kb reindex` command to the Go client with confirmation prompt (skip with `--yes`/`-y`)
|
||||
- Display reindex results (chunks reindexed, model used)
|
||||
|
||||
## Capabilities
|
||||
|
||||
### New Capabilities
|
||||
|
||||
(none)
|
||||
|
||||
### Modified Capabilities
|
||||
|
||||
- `go-client`: Add reindex command requirement
|
||||
|
||||
## Impact
|
||||
|
||||
- New file: `client/cmd/reindex.go`
|
||||
@@ -0,0 +1,25 @@
|
||||
## ADDED Requirements
|
||||
|
||||
### Requirement: Reindex command
|
||||
|
||||
The client SHALL provide a `kb reindex` command that triggers re-embedding of all chunks on the engine. The command SHALL prompt for confirmation before proceeding.
|
||||
|
||||
#### Scenario: Reindex with confirmation
|
||||
- **WHEN** the user runs `kb reindex`
|
||||
- **THEN** the client SHALL display a warning that all chunks will be re-embedded and prompt `Reindex all chunks? This will re-embed everything. [y/N]`. If confirmed, it SHALL POST to `/api/v1/reindex` and display the result.
|
||||
|
||||
#### Scenario: Reindex with skip confirmation
|
||||
- **WHEN** the user runs `kb reindex --yes`
|
||||
- **THEN** the client SHALL skip the confirmation prompt and POST to `/api/v1/reindex` immediately
|
||||
|
||||
#### Scenario: Reindex cancelled
|
||||
- **WHEN** the user runs `kb reindex` and responds with anything other than `y` or `yes`
|
||||
- **THEN** the client SHALL print `Cancelled.` and exit with code 0
|
||||
|
||||
#### Scenario: Reindex human output
|
||||
- **WHEN** the reindex completes successfully with default format
|
||||
- **THEN** the client SHALL print `Reindexed N chunks (model: <model_name>)`
|
||||
|
||||
#### Scenario: Reindex JSON output
|
||||
- **WHEN** the user runs `kb reindex --yes --format json`
|
||||
- **THEN** the client SHALL output the raw JSON response from the engine
|
||||
@@ -0,0 +1,5 @@
|
||||
## 1. Implementation
|
||||
|
||||
- [x] 1.1 Create `client/cmd/reindex.go` with `kb reindex` command, `--yes`/`-y` flag, confirmation prompt matching `remove.go` pattern
|
||||
- [x] 1.2 POST to `/api/v1/reindex`, handle human output (`Reindexed N chunks (model: ...)`) and JSON output
|
||||
- [x] 1.3 Verify build compiles and command appears in `kb --help`
|
||||
@@ -0,0 +1,2 @@
|
||||
schema: spec-driven
|
||||
created: 2026-03-29
|
||||
@@ -0,0 +1,43 @@
|
||||
## Context
|
||||
|
||||
The `add` command currently handles both file uploads and notes via a `--note` string flag. This creates confusing flag parsing and a muddled help screen. The engine already auto-detects file type from extension (`detector.py`) and rejects unsupported ones, so the client's `--type` flag is redundant.
|
||||
|
||||
## Goals / Non-Goals
|
||||
|
||||
**Goals:**
|
||||
- `kb "my note"` as the sole note entry path (replaces `kb add --note`)
|
||||
- `kb addfile <path>` as a file-only upload command (replaces `kb add`)
|
||||
- Client-side extension validation before uploading
|
||||
- Clean, unambiguous help text for both paths
|
||||
|
||||
**Non-Goals:**
|
||||
- Engine changes — type detection stays server-side
|
||||
- Backward compatibility shim for `kb add` — clean break
|
||||
- Client-side MIME type detection — extension check is sufficient
|
||||
|
||||
## Decisions
|
||||
|
||||
### Rename add → addfile, strip note/type flags
|
||||
|
||||
Rename the cobra command from `add` to `addfile`. Remove `--note`, `--title`, and `--type` flags. Keep `--tags`, `--recursive`. The command becomes purely about file uploads.
|
||||
|
||||
**Why not keep `add` as an alias?** Clean break is simpler. The old form was confusing — better to force a quick migration than maintain two paths.
|
||||
|
||||
### Extension validation on single file uploads
|
||||
|
||||
The `supportedExts` map already gates recursive walks. Apply the same check to single file uploads — reject with a clear error listing supported extensions. This gives instant feedback instead of a round-trip to the engine.
|
||||
|
||||
### Root command RunE for note shorthand
|
||||
|
||||
Use cobra's `Args: cobra.ArbitraryArgs` and `RunE` on the root command. When args are present and no subcommand matched, join all args into a single note string and submit. `--tags` flag on root for tagging notes. No `--title` — keep it minimal.
|
||||
|
||||
**Why join all args?** `kb remember to update dns` (unquoted) should work the same as `kb "remember to update dns"`.
|
||||
|
||||
### Reuse note submission logic via shared helper
|
||||
|
||||
Extract `submitNote` from the current `runAdd` so both the root command and any future callers use the same POST + duplicate-handling + output logic.
|
||||
|
||||
## Risks / Trade-offs
|
||||
|
||||
- **Breaking change** → Anyone with `kb add` in scripts needs to update to `kb addfile`. Acceptable for a personal tool.
|
||||
- **No `--type` override** → If a user ever needs to force a type, they'd have to go through the engine API directly. Low risk since the engine's auto-detection covers all supported formats.
|
||||
@@ -0,0 +1,34 @@
|
||||
## Why
|
||||
|
||||
Adding a note requires `kb add --note "my note"` — too much ceremony for what should be instant. The `--note` flag taking a string value also creates confusing flag parsing (e.g. `kb add --note --tags foo` parses `--tags` as the note value). Meanwhile, `kb add` tries to do two things (files and notes) which muddies its help text and UX.
|
||||
|
||||
Splitting these into distinct paths makes the CLI clearer:
|
||||
- **Notes**: `kb "my note"` — zero-friction, no subcommand needed
|
||||
- **Files**: `kb addfile report.pdf` — explicit, file-only command
|
||||
|
||||
## What Changes
|
||||
|
||||
- **Add `kb "text"` shorthand**: bare string arguments without a subcommand are treated as notes, submitted via `POST /api/v1/jobs`
|
||||
- **Rename `add` → `addfile`**: the command becomes file-only, no more `--note`/`--title` flags
|
||||
- **Drop `--type` flag**: the engine already auto-detects type from file extension (`detector.py`); the client doesn't need to override this
|
||||
- **Add client-side extension validation**: reject unsupported file extensions with a clear error before uploading, using the same extension set as recursive directory walks
|
||||
- **Update README**: document the new shorthand and renamed command
|
||||
- **BREAKING**: `kb add` no longer exists; `kb add --note` no longer exists
|
||||
|
||||
## Capabilities
|
||||
|
||||
### New Capabilities
|
||||
|
||||
_(none)_
|
||||
|
||||
### Modified Capabilities
|
||||
|
||||
- `go-client`: Rename `add` to `addfile`, remove `--note`/`--title`/`--type` flags, add extension validation for single file uploads, add implicit note shorthand on root command
|
||||
|
||||
## Impact
|
||||
|
||||
- `client/cmd/add.go` → renamed/refactored to `addfile` command, stripped of note logic, added extension check
|
||||
- `client/cmd/root.go` — bare args handling + `--tags` flag for note shorthand
|
||||
- `README.md` — updated usage examples
|
||||
- No engine changes — engine already detects type from extension and rejects unsupported files
|
||||
- Breaking change for any scripts using `kb add` or `kb add --note`
|
||||
@@ -0,0 +1,95 @@
|
||||
## ADDED Requirements
|
||||
|
||||
### Requirement: Implicit note shorthand
|
||||
|
||||
The client SHALL treat bare string arguments (with no subcommand) as an implicit note. `kb "my note"` SHALL behave identically to submitting a note via `POST /api/v1/jobs`. All persistent flags (`--format`, `--engine`, `--api-key`) and the root `--tags` flag SHALL work with the shorthand form.
|
||||
|
||||
#### Scenario: Quick note via bare argument
|
||||
- **WHEN** the user runs `kb "remember to update DNS"`
|
||||
- **THEN** the client SHALL submit the text as a note via `POST /api/v1/jobs` and print `Queued: note`
|
||||
|
||||
#### Scenario: Bare argument with tags
|
||||
- **WHEN** the user runs `kb "server room is building 3" --tags ops`
|
||||
- **THEN** the client SHALL submit the note with the specified tags
|
||||
|
||||
#### Scenario: Bare argument with JSON output
|
||||
- **WHEN** the user runs `kb "my note" --format json`
|
||||
- **THEN** the client SHALL output the raw JSON response from the engine
|
||||
|
||||
#### Scenario: Bare argument duplicate detection
|
||||
- **WHEN** the user runs `kb "my note"` and the engine returns HTTP 409
|
||||
- **THEN** the client SHALL handle the duplicate response identically to the previous `kb add --note` behaviour
|
||||
|
||||
#### Scenario: Multiple unquoted words
|
||||
- **WHEN** the user runs `kb remember to update dns` (without quotes)
|
||||
- **THEN** the client SHALL join all arguments into a single note string and submit it
|
||||
|
||||
#### Scenario: No interference with subcommands
|
||||
- **WHEN** the user runs `kb search "query"` or any other existing subcommand
|
||||
- **THEN** the client SHALL route to the subcommand as before — the implicit note shorthand SHALL NOT interfere
|
||||
|
||||
#### Scenario: No arguments
|
||||
- **WHEN** the user runs `kb` with no arguments
|
||||
- **THEN** the client SHALL display the help text
|
||||
|
||||
---
|
||||
|
||||
## MODIFIED Requirements
|
||||
|
||||
### Requirement: Add command (file and note ingestion)
|
||||
|
||||
The client SHALL provide a `kb addfile` command that uploads files to the engine for async ingestion. The command SHALL validate file extensions before uploading and reject unsupported types. The client SHALL handle duplicate rejection (HTTP 409) and display the existing document information. The command SHALL NOT handle notes — notes are submitted via the implicit note shorthand (`kb "text"`).
|
||||
|
||||
#### Scenario: Add a single file
|
||||
- **WHEN** the user runs `kb addfile report.pdf`
|
||||
- **THEN** the client SHALL validate the file extension, upload the file via `POST /api/v1/jobs` (multipart), print "Queued: report.pdf", and exit
|
||||
|
||||
#### Scenario: Add a file with tags
|
||||
- **WHEN** the user runs `kb addfile manual.pdf --tags car,maintenance`
|
||||
- **THEN** the client SHALL include the tags in the multipart upload metadata
|
||||
|
||||
#### Scenario: Add a directory recursively
|
||||
- **WHEN** the user runs `kb addfile ~/documents/ --recursive`
|
||||
- **THEN** the client SHALL discover all supported files in the directory tree, upload each one sequentially, and print "Queued: N files"
|
||||
|
||||
#### Scenario: Unsupported file extension
|
||||
- **WHEN** the user runs `kb addfile photo.jpg`
|
||||
- **THEN** the client SHALL print an error listing supported extensions and exit with a non-zero code without making any API call
|
||||
|
||||
#### Scenario: Duplicate file rejected (already ingested)
|
||||
- **WHEN** the user runs `kb addfile report.pdf` and the engine returns HTTP 409 with `{"error": "duplicate", "document_id": 42, "title": "report.pdf"}`
|
||||
- **THEN** the client SHALL print "Already imported: report.pdf (doc ID: 42)" and exit with code 0
|
||||
|
||||
#### Scenario: Duplicate file rejected (in-flight job)
|
||||
- **WHEN** the user runs `kb addfile report.pdf` and the engine returns HTTP 409 with `{"error": "duplicate", "job_id": 7, "title": "report.pdf"}`
|
||||
- **THEN** the client SHALL print "Already queued: report.pdf (job ID: 7)" and exit with code 0
|
||||
|
||||
#### Scenario: Duplicate file in recursive add
|
||||
- **WHEN** the user runs `kb addfile ~/documents/ --recursive` and some files are rejected as duplicates
|
||||
- **THEN** the client SHALL print the duplicate message for each rejected file, continue uploading remaining files, and include a summary (e.g., "Queued: 5 files, 2 duplicates skipped")
|
||||
|
||||
#### Scenario: Duplicate with JSON output
|
||||
- **WHEN** the user runs `kb addfile report.pdf --format json` and the engine returns HTTP 409
|
||||
- **THEN** the client SHALL output the raw JSON response from the engine including the document_id and title
|
||||
|
||||
#### Scenario: Add with JSON output
|
||||
- **WHEN** the user runs `kb addfile report.pdf --format json`
|
||||
- **THEN** the client SHALL output the JSON response from the engine including the job_id
|
||||
|
||||
#### Scenario: File not found
|
||||
- **WHEN** the user runs `kb addfile nonexistent.pdf`
|
||||
- **THEN** the client SHALL print an error and exit with a non-zero code without making any API call
|
||||
|
||||
#### Scenario: Upload failure
|
||||
- **WHEN** the upload fails (network error, engine returns 4xx/5xx other than 409)
|
||||
- **THEN** the client SHALL print the error and exit with a non-zero code
|
||||
|
||||
## REMOVED Requirements
|
||||
|
||||
### Requirement: Note ingestion via add command
|
||||
**Reason**: Notes are now submitted via the implicit note shorthand (`kb "text"`). The `--note` and `--title` flags on the add command are removed.
|
||||
**Migration**: Use `kb "my note"` or `kb "my note" --tags ops` instead of `kb add --note "my note" --tags ops`.
|
||||
|
||||
### Requirement: Document type override via add command
|
||||
**Reason**: The engine auto-detects document type from file extension (`detector.py`). The client `--type` flag is redundant.
|
||||
**Migration**: Remove `--type` from scripts. The engine handles type detection automatically.
|
||||
@@ -0,0 +1,19 @@
|
||||
## 1. Refactor note submission
|
||||
|
||||
- [x] 1.1 Extract note submission logic from `runAdd` into a shared `submitNote` helper (multipart POST, duplicate detection, output formatting)
|
||||
|
||||
## 2. Root command shorthand
|
||||
|
||||
- [x] 2.1 Add `Args: cobra.ArbitraryArgs` and `RunE` to the root command — join args into a note string, call `submitNote`; show help when no args
|
||||
- [x] 2.2 Add `--tags` flag on the root command for note tagging
|
||||
|
||||
## 3. Rename add → addfile
|
||||
|
||||
- [x] 3.1 Rename command from `add` to `addfile` (`Use: "addfile <path>"`)
|
||||
- [x] 3.2 Remove `--note`, `--title`, and `--type` flags from the command
|
||||
- [x] 3.3 Add extension validation for single file uploads — reject unsupported extensions with a clear error listing supported types
|
||||
|
||||
## 4. Documentation and verification
|
||||
|
||||
- [x] 4.1 Update README.md usage section: show `kb "text"` shorthand, rename `add` references to `addfile`
|
||||
- [x] 4.2 Verify build compiles, `kb --help` and `kb addfile --help` show expected output
|
||||
@@ -0,0 +1,2 @@
|
||||
schema: spec-driven
|
||||
created: 2026-03-28
|
||||
@@ -0,0 +1,93 @@
|
||||
## Context
|
||||
|
||||
Currently the project uses a single version number shared between client and engine, managed by `release.sh`. Both `client/VERSION` and `engine/VERSION` are always bumped to the same value. A single git tag `vX.Y.Z` is created, and a single Gitea release bundles Go client binaries and Docker engine image references. This means any change to either component forces a full release of both.
|
||||
|
||||
The client is a Go binary distributed as platform-specific downloads. The engine is a Python FastAPI server distributed as Docker images. They communicate over HTTP via `/api/v1/` endpoints. The engine already exposes its version via `GET /api/v1/status` → `{"version": "X.Y.Z", ...}`.
|
||||
|
||||
## Goals / Non-Goals
|
||||
|
||||
**Goals:**
|
||||
- Allow client and engine to have independent version numbers and release cadences
|
||||
- Provide a runtime compatibility check so users get a clear error when their client is too new for their engine
|
||||
- Split release tooling so each component can be released without touching the other
|
||||
|
||||
**Non-Goals:**
|
||||
- API versioning beyond the existing `/api/v1/` path prefix
|
||||
- Backward-compatible negotiation or feature detection (client either works or fails)
|
||||
- Automatic upgrades or update notifications
|
||||
- Version checking in the other direction (engine requiring minimum client)
|
||||
|
||||
## Decisions
|
||||
|
||||
### 1. Tag naming: `client-vX.Y.Z` and `engine-vX.Y.Z`
|
||||
|
||||
Prefix-style tags clearly identify which component a release belongs to and sort well in git tag listings.
|
||||
|
||||
**Why over path-style (`client/vX.Y.Z`):** Slashes in git tags can cause issues with some tooling and are less conventional. Prefix-style is simpler and widely used in monorepos.
|
||||
|
||||
**Why over separate repos:** The project is small and tightly coupled at the API level. A monorepo with prefixed tags keeps everything together while allowing independent releases.
|
||||
|
||||
### 2. Two release scripts: `release-client.sh` and `release-engine.sh`
|
||||
|
||||
Each script handles its own component end-to-end: version bump, build, tag, release, push.
|
||||
|
||||
**Why over a single script with flags:** Two simple scripts are easier to understand and maintain than one script with component-selection logic. Each script is ~100 lines instead of one ~200-line script with branching. The shared logic (version helpers, pre-flight checks) is minimal and acceptable to duplicate.
|
||||
|
||||
**Shared structure for both scripts:**
|
||||
1. Pre-flight checks (on main branch, tag doesn't exist)
|
||||
2. Version bump (reads/writes component's VERSION file only)
|
||||
3. Build artifacts (Go binaries or Docker images)
|
||||
4. Commit version bump, create prefixed tag, push
|
||||
5. Create Gitea release with assets
|
||||
6. (Engine only) Push Docker images
|
||||
|
||||
### 3. `MinEngineVersion` as a build-time constant in the Go client
|
||||
|
||||
The client embeds a `MinEngineVersion` string constant alongside the existing `Version` constant. It is set via `-ldflags` at build time, sourced from a `client/MIN_ENGINE_VERSION` file.
|
||||
|
||||
**Why a separate file over embedding in `VERSION`:** The two values have different lifecycles. `VERSION` changes every release; `MIN_ENGINE_VERSION` changes only when the client starts using a new engine feature. A separate file makes the intent clear.
|
||||
|
||||
**Why ldflags over hardcoding in Go source:** Consistent with how `Version` is already injected. The value lives in a plain text file that's easy to bump manually.
|
||||
|
||||
### 4. Compatibility check on every API call via the `Client` struct
|
||||
|
||||
The `api.Client` checks engine compatibility on its first HTTP call by hitting `GET /api/v1/status` and comparing the `version` field against `MinEngineVersion`. The result is cached on the `Client` instance — subsequent calls skip the check.
|
||||
|
||||
**Flow:**
|
||||
1. First call to any `Client` method (Get/Post/Delete/Put)
|
||||
2. Before the actual request, call `GET /api/v1/status`
|
||||
3. Parse `version` from response
|
||||
4. Compare against `MinEngineVersion` using semver major.minor.patch comparison
|
||||
5. If engine version < min: print error to stderr, `os.Exit(1)`
|
||||
6. If check passes: set `versionChecked = true`, proceed with original request
|
||||
7. If status endpoint unreachable: proceed with original request (connectivity error will surface on the actual call)
|
||||
|
||||
**Why hard fail, no skip flag:** This is a personal tool. If the client needs a newer engine, the user needs to update. A skip flag adds complexity for a scenario where the outcome (broken behavior) is worse than the error.
|
||||
|
||||
**Why check on first API call, not at startup:** The `PersistentPreRunE` in cobra runs before every command, but some future commands might not need the engine (e.g. `kb version`, `kb help`). Checking in the `Client` ensures we only check when actually contacting the engine.
|
||||
|
||||
**Why proceed when status endpoint is unreachable:** If we can't reach `/status`, the actual API call will also fail with a connection error. No point in double-failing. The compatibility check is for version mismatch, not connectivity.
|
||||
|
||||
### 5. Compose files: use `build:` context, not pinned image tags
|
||||
|
||||
The compose files currently use `build:` directives, not pre-built image references. Users who build locally don't need pinned tags — they're building from source. Users pulling pre-built images will reference the image tag directly in their own compose file or `docker run` command.
|
||||
|
||||
**Decision:** Leave compose files as-is. Release notes for engine releases will include the exact `docker pull` command with the versioned tag.
|
||||
|
||||
### 6. Semver comparison: major.minor.patch, no pre-release
|
||||
|
||||
Compare versions as three integers. No support for pre-release suffixes (`-rc1`, `-beta`) — the project doesn't use them. If `MinEngineVersion` is `2.1.0` and engine reports `2.1.5`, the check passes. If engine reports `2.0.9`, it fails.
|
||||
|
||||
## Risks / Trade-offs
|
||||
|
||||
- **Extra HTTP round-trip on first command** — One additional `GET /api/v1/status` call per client invocation. Negligible for a local-network tool.
|
||||
→ Mitigation: Cached after first check within the Client instance.
|
||||
|
||||
- **Developer must remember to bump `MIN_ENGINE_VERSION`** — When adding client code that depends on a new engine endpoint/field, the developer must manually update the file.
|
||||
→ Mitigation: This is a conscious decision point. The file's existence serves as a reminder. Could add a CI check later if needed.
|
||||
|
||||
- **Breaking change to git tag format** — Existing `v2.0.x` tags won't match the new `client-v*` / `engine-v*` convention. Old tags remain in history.
|
||||
→ Mitigation: No migration needed. Old tags stay as historical artifacts. New convention starts from the first independent release.
|
||||
|
||||
- **Two Gitea releases per coordinated release** — When both components change, two releases are created instead of one.
|
||||
→ Mitigation: Acceptable trade-off. Each release is self-contained with its own assets and notes.
|
||||
@@ -0,0 +1,32 @@
|
||||
## Why
|
||||
|
||||
Client and engine are currently locked to the same version number and released together via a single script. This means a client-only bug fix (e.g. output formatting) forces a full engine Docker image rebuild and push, and vice versa. Decoupling versions allows each component to be released independently on its own cadence, while a compatibility check ensures users don't run a client that requires engine features not yet deployed.
|
||||
|
||||
## What Changes
|
||||
|
||||
- **Separate version files** — `client/VERSION` and `engine/VERSION` may diverge (they already exist as separate files, but are currently always set to the same value)
|
||||
- **Split release script** — Replace single `release.sh` with `release-client.sh` (builds Go binaries, tags `client-vX.Y.Z`, creates release) and `release-engine.sh` (builds Docker images, tags `engine-vX.Y.Z`, creates release, pushes images)
|
||||
- **Client compatibility check** — Client embeds a `MinEngineVersion` constant (set at build time or in code). On every command that contacts the engine, the client calls `GET /api/v1/status`, compares the engine's reported version against `MinEngineVersion`, and hard-fails with an actionable error if the engine is too old. No skip flag, no warning — just a clear error with upgrade instructions.
|
||||
- **Tag naming convention** — `client-vX.Y.Z` and `engine-vX.Y.Z` replace the current `vX.Y.Z` tag format. **BREAKING** — existing tag format changes.
|
||||
|
||||
## Capabilities
|
||||
|
||||
### New Capabilities
|
||||
|
||||
(none)
|
||||
|
||||
### Modified Capabilities
|
||||
|
||||
- `go-client`: Add engine version compatibility check requirement (hard fail if engine version < MinEngineVersion)
|
||||
- `engine-api`: Status endpoint already returns `version` — no change needed, but delta spec documents the contract that the version field is required for compatibility checking
|
||||
- `docker-deployment`: Compose files pin engine image tag; release script changes affect image tagging
|
||||
|
||||
## Impact
|
||||
|
||||
- `release.sh` — replaced by `release-client.sh` + `release-engine.sh`
|
||||
- `client/cmd/root.go` — new `MinEngineVersion` constant
|
||||
- `client/internal/api/client.go` — version check on first API call
|
||||
- `client/Makefile` — may inject `MinEngineVersion` via ldflags alongside `Version`
|
||||
- Git tags — new naming convention (`client-v*`, `engine-v*`)
|
||||
- Gitea releases — two separate releases per independent release cycle
|
||||
- `engine/compose.nvidia.yaml`, `engine/compose.rocm.yaml` — add pinned image tag
|
||||
+25
@@ -0,0 +1,25 @@
|
||||
## MODIFIED Requirements
|
||||
|
||||
### 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: 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
|
||||
- **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, etc.)
|
||||
- **THEN** the engine 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`)
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user