Add reranking, RRF fusion, bench harness, tag contexts, and data ingestion
Implements five of the six enhancements from docs/kb-enhancements-proposal.htm, closing the retrieval-quality gap identified in the qmd review. - Cross-encoder reranking: new kb/reranker.py loads an optional reranking model at startup (KB_RERANK_ENABLED, KB_RERANKER_MODEL, KB_RERANK_CANDIDATES). Search degrades gracefully to plain hybrid retrieval when the model is absent. Exposed via a "rerank" block in /status, a rerank flag on search, and --no-rerank in the CLI. - RRF rank fusion: FTS and vector lists now merge by reciprocal rank fusion with a top-rank bonus, replacing the old score blend. Scores are comparable across queries. - Bench harness and explain traces: kb bench runs a query fixture against each backend and reports precision@k, recall and MRR. --explain returns a per-result score breakdown. - Tag context descriptions: tags carry an optional one-line description (kb tag-describe), returned as tag_contexts with search results. Adds a tags.description column migration. - Structured data ingestion: .json/.yaml/.toml files ingest as text via the new "data" doc type, pretty-printing minified JSON before chunking. Query expansion (proposal item 5) is deliberately left out pending bench results. Requires engine v3.3.0. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -21,6 +21,9 @@ services:
|
||||
- KB_INGEST_DEVICE=${KB_INGEST_DEVICE:-auto}
|
||||
- KB_API_KEY=${KB_API_KEY:-}
|
||||
- KB_SEARCH_THRESHOLD=${KB_SEARCH_THRESHOLD:-0.01}
|
||||
- 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
|
||||
|
||||
|
||||
@@ -20,6 +20,9 @@ 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.host = os.environ.get("KB_HOST", "0.0.0.0")
|
||||
self.port = int(os.environ.get("KB_PORT", "8000"))
|
||||
|
||||
@@ -195,6 +195,11 @@ def init_schema(conn: sqlite3.Connection, embedding_dim: int) -> None:
|
||||
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()
|
||||
|
||||
|
||||
|
||||
@@ -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"),
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
|
||||
|
||||
+166
-16
@@ -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:
|
||||
@@ -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],
|
||||
@@ -311,6 +457,10 @@ def _enrich(
|
||||
"source_path": row[7],
|
||||
"created_at": row[8],
|
||||
"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
|
||||
|
||||
@@ -124,6 +124,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}")
|
||||
|
||||
|
||||
@@ -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())
|
||||
|
||||
|
||||
@@ -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,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,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"] == {}
|
||||
Reference in New Issue
Block a user