176 lines
4.9 KiB
Python
176 lines
4.9 KiB
Python
#!/usr/bin/env python3
|
|
|
|
import json
|
|
from pathlib import Path
|
|
from typing import Any, List, Optional
|
|
|
|
import numpy as np
|
|
import faiss
|
|
from fastapi import FastAPI, HTTPException
|
|
from pydantic import BaseModel
|
|
from sentence_transformers import SentenceTransformer
|
|
|
|
|
|
# ============================================================
|
|
# Paths
|
|
# ============================================================
|
|
|
|
BASE_PATH = Path(__file__).resolve().parents[2]
|
|
KNOWLEDGE_DIR = BASE_PATH / "var" / "knowledge"
|
|
|
|
CHUNK_INDEX_PATH = KNOWLEDGE_DIR / "vector.index"
|
|
CHUNK_MAP_PATH = KNOWLEDGE_DIR / "vector.index.meta.json"
|
|
|
|
TAG_INDEX_PATH = KNOWLEDGE_DIR / "vector_tags.index"
|
|
TAG_MAP_PATH = KNOWLEDGE_DIR / "vector_tags.index.meta.json"
|
|
|
|
INDEX_META_PATH = KNOWLEDGE_DIR / "index_meta.json"
|
|
|
|
|
|
# ============================================================
|
|
# FastAPI
|
|
# ============================================================
|
|
|
|
app = FastAPI()
|
|
|
|
model: Optional[SentenceTransformer] = None
|
|
chunk_index = None
|
|
chunk_ids: Optional[List[Any]] = None
|
|
tag_index = None
|
|
tag_ids: Optional[List[Any]] = None
|
|
loaded_embedding_model_name: Optional[str] = None
|
|
|
|
|
|
# ============================================================
|
|
# Models
|
|
# ============================================================
|
|
|
|
class SearchRequest(BaseModel):
|
|
query: str
|
|
limit: int = 8
|
|
|
|
|
|
# ============================================================
|
|
# Loader
|
|
# ============================================================
|
|
|
|
def load_all():
|
|
global model, chunk_index, chunk_ids, tag_index, tag_ids, loaded_embedding_model_name
|
|
|
|
if not INDEX_META_PATH.exists():
|
|
raise RuntimeError("index_meta.json not found")
|
|
|
|
meta = json.loads(INDEX_META_PATH.read_text())
|
|
embedding_model_name = meta.get("embedding_model")
|
|
|
|
if not embedding_model_name:
|
|
raise RuntimeError("embedding_model missing in index_meta.json")
|
|
|
|
# Reload model only if changed
|
|
if model is None or embedding_model_name != loaded_embedding_model_name:
|
|
print(f"[Reload] Loading embedding model: {embedding_model_name}")
|
|
model = SentenceTransformer(embedding_model_name)
|
|
loaded_embedding_model_name = embedding_model_name
|
|
|
|
# Reload chunk index
|
|
if CHUNK_INDEX_PATH.exists() and CHUNK_MAP_PATH.exists():
|
|
print("[Reload] Loading chunk index")
|
|
chunk_index = faiss.read_index(str(CHUNK_INDEX_PATH))
|
|
chunk_ids = json.loads(CHUNK_MAP_PATH.read_text())
|
|
else:
|
|
chunk_index = None
|
|
chunk_ids = None
|
|
|
|
# Reload tag index
|
|
if TAG_INDEX_PATH.exists() and TAG_MAP_PATH.exists():
|
|
print("[Reload] Loading tag index")
|
|
tag_index = faiss.read_index(str(TAG_INDEX_PATH))
|
|
tag_ids = json.loads(TAG_MAP_PATH.read_text())
|
|
else:
|
|
tag_index = None
|
|
tag_ids = None
|
|
|
|
print("[Reload] Completed")
|
|
|
|
|
|
# ============================================================
|
|
# Startup
|
|
# ============================================================
|
|
|
|
@app.on_event("startup")
|
|
def startup_event():
|
|
load_all()
|
|
print("[VectorService] Ready")
|
|
|
|
|
|
# ============================================================
|
|
# Endpoints
|
|
# ============================================================
|
|
|
|
@app.get("/health")
|
|
def health():
|
|
return {
|
|
"status": "ok",
|
|
"chunk_index_loaded": chunk_index is not None,
|
|
"tag_index_loaded": tag_index is not None,
|
|
"model_loaded": model is not None,
|
|
}
|
|
|
|
|
|
@app.post("/reload")
|
|
def reload():
|
|
try:
|
|
load_all()
|
|
return {"status": "reloaded"}
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
|
|
@app.post("/search-chunks")
|
|
def search_chunks(req: SearchRequest):
|
|
if chunk_index is None or chunk_ids is None:
|
|
raise HTTPException(status_code=503, detail="Chunk index not available")
|
|
|
|
query_vec = model.encode([req.query], normalize_embeddings=True)
|
|
query_vec = np.array(query_vec).astype("float32")
|
|
|
|
scores, indices = chunk_index.search(query_vec, req.limit)
|
|
|
|
results = []
|
|
for score, idx in zip(scores[0], indices[0]):
|
|
if idx == -1:
|
|
continue
|
|
if idx < 0 or idx >= len(chunk_ids):
|
|
continue
|
|
|
|
results.append({
|
|
"chunk_id": chunk_ids[idx],
|
|
"score": float(score),
|
|
})
|
|
|
|
return results
|
|
|
|
|
|
@app.post("/search-tags")
|
|
def search_tags(req: SearchRequest):
|
|
if tag_index is None or tag_ids is None:
|
|
raise HTTPException(status_code=503, detail="Tag index not available")
|
|
|
|
query_vec = model.encode([req.query], normalize_embeddings=True)
|
|
query_vec = np.array(query_vec).astype("float32")
|
|
|
|
scores, indices = tag_index.search(query_vec, req.limit)
|
|
|
|
results = []
|
|
for score, idx in zip(scores[0], indices[0]):
|
|
if idx == -1:
|
|
continue
|
|
if idx < 0 or idx >= len(tag_ids):
|
|
continue
|
|
|
|
results.append({
|
|
"chunk_id": tag_ids[idx],
|
|
"score": float(score),
|
|
})
|
|
|
|
return results |