steve
6dfc13be1d
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>
82 lines
2.3 KiB
Python
82 lines
2.3 KiB
Python
"""Engine entry point — FastAPI server with eager model loading."""
|
|
|
|
import asyncio
|
|
import logging
|
|
import os
|
|
from contextlib import asynccontextmanager
|
|
from pathlib import Path
|
|
|
|
from fastapi import FastAPI
|
|
|
|
_version_file = Path(__file__).parent / "VERSION"
|
|
__version__ = _version_file.read_text().strip() if _version_file.exists() else "dev"
|
|
|
|
from kb.config import cfg
|
|
from kb.embeddings import load_model
|
|
from kb.database import get_connection, init_schema
|
|
from kb.worker import ingestion_worker
|
|
|
|
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
|
|
log = logging.getLogger("kb.engine")
|
|
|
|
# Track readiness for health endpoint
|
|
ready = False
|
|
worker_task: asyncio.Task | None = None
|
|
|
|
|
|
@asynccontextmanager
|
|
async def lifespan(app: FastAPI):
|
|
global ready, worker_task
|
|
|
|
# Set HF cache before any model imports
|
|
os.environ["HF_HOME"] = str(cfg.hf_cache)
|
|
|
|
log.info("Starting engine...")
|
|
cfg.ensure_dirs()
|
|
|
|
# Initialise database
|
|
conn = get_connection(cfg.db_path)
|
|
model_dim = load_model(cfg.model, cfg.device)
|
|
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())
|
|
|
|
ready = True
|
|
log.info("Engine ready — model: %s, device: %s", cfg.model, cfg.device)
|
|
|
|
yield
|
|
|
|
# Shutdown
|
|
ready = False
|
|
if worker_task:
|
|
worker_task.cancel()
|
|
try:
|
|
await worker_task
|
|
except asyncio.CancelledError:
|
|
pass
|
|
log.info("Engine stopped.")
|
|
|
|
|
|
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, notes, bulk # noqa: E402, F401
|
|
|
|
if __name__ == "__main__":
|
|
import uvicorn
|
|
uvicorn.run("main:app", host=cfg.host, port=cfg.port, log_level="info")
|