from fastapi import FastAPI, UploadFile, File, Form, HTTPException
from fastapi.responses import HTMLResponse
from pydantic import BaseModel
from typing import Optional
from pathlib import Path
import shutil
from .config import CONFIG
from .embed import embed_one
from .ingest import ingest_file
from . import db
app = FastAPI(title="KB Search")
class SearchRequest(BaseModel):
query: str
category: Optional[str] = None
k: int = 10
expand_neighbors: bool = True
neighbor_window: int = 1
rerank: bool = True
rerank_pool: int = 25
@app.post("/search")
def search(req: SearchRequest):
query_emb = embed_one(req.query)
# Fetch larger pool when reranking
initial_k = max(req.rerank_pool, req.k) if req.rerank else req.k
if req.expand_neighbors:
results = db.search_with_neighbors(
query_emb, k=initial_k, category=req.category,
neighbor_window=req.neighbor_window
)
else:
results = db.search(query_emb, k=initial_k, category=req.category)
rerank_used = False
if req.rerank and results:
try:
from .rerank import rerank as do_rerank
results = do_rerank(req.query, results, top_k=req.k)
rerank_used = True
except Exception as e:
print(f"Rerank failed, falling back to vector: {e}", flush=True)
results = results[:req.k]
else:
results = results[:req.k]
return {
"query": req.query,
"count": len(results),
"rerank_used": rerank_used,
"results": results,
}
@app.get("/stats")
def stats():
return db.get_stats()
@app.post("/upload")
async def upload(file: UploadFile = File(...), category: str = Form("uploads")):
if not file.filename:
raise HTTPException(400, "No filename")
safe_name = "".join(c for c in file.filename if c.isalnum() or c in "._- ")
target_dir = Path(CONFIG.corpus_root) / category
target_dir.mkdir(parents=True, exist_ok=True)
target = target_dir / safe_name
with open(target, "wb") as f:
shutil.copyfileobj(file.file, f)
result = ingest_file(target)
return result
@app.get("/upload", response_class=HTMLResponse)
def upload_page():
return """
KB Upload
KB Upload
"""
SEARCH_UI_HTML = """\
KB Search
"""
@app.get("/", response_class=HTMLResponse)
def search_ui():
return SEARCH_UI_HTML