26 lines
895 B
Python
26 lines
895 B
Python
from fastapi import APIRouter, HTTPException, Query
|
|
from pydantic import BaseModel
|
|
import numpy as np
|
|
from app.core.index_manager import UserIndexManager
|
|
from app.core.embedding import embedder
|
|
from app.core.config import settings
|
|
|
|
router = APIRouter()
|
|
index_manager = UserIndexManager()
|
|
|
|
class QueryRequest(BaseModel):
|
|
query: str
|
|
|
|
@router.post("/search")
|
|
def search_docs(request: QueryRequest, user_id: str = Query(..., description="用户ID")):
|
|
try:
|
|
query_vector = embedder.encode([request.query])
|
|
index = index_manager.get_index(user_id)
|
|
D, I = index.search(query_vector, settings.TOP_K)
|
|
return {
|
|
"user_id": user_id,
|
|
"query": request.query,
|
|
"results": [{"id": int(idx), "score": float(dist)} for dist, idx in zip(D, I)]
|
|
}
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=str(e)) |