85 lines
2.6 KiB
Python
85 lines
2.6 KiB
Python
from fastapi import FastAPI, HTTPException
|
|
|
|
from config import get_settings
|
|
from faiss_manager import faiss_manager
|
|
from models import EmbeddingInput, IndexDeleteInput, IndexInsertInput, IndexSearchInput, SearchInput
|
|
|
|
settings = get_settings()
|
|
|
|
app = FastAPI(
|
|
title="FAISS 服务",
|
|
description="向量插入 + 相似搜索 + 多索引管理",
|
|
version="1.0.0",
|
|
)
|
|
|
|
|
|
def log_business(message: str):
|
|
if settings.ENABLE_REQUEST_LOGS:
|
|
print(message, flush=True)
|
|
|
|
|
|
@app.post("/insert")
|
|
async def insert(data: EmbeddingInput):
|
|
try:
|
|
vector_id = faiss_manager.insert(data.embedding)
|
|
log_business(f"[faiss] global insert id={vector_id}")
|
|
return {"id": vector_id}
|
|
except Exception as exc:
|
|
raise HTTPException(status_code=400, detail=str(exc))
|
|
|
|
|
|
@app.post("/search")
|
|
async def search(data: SearchInput):
|
|
try:
|
|
result = faiss_manager.search(data.embedding, data.k)
|
|
log_business(f"[faiss] global search ids={result['ids']} similarity_scores={result['similarity_scores']}")
|
|
return result
|
|
except Exception as exc:
|
|
raise HTTPException(status_code=400, detail=str(exc))
|
|
|
|
|
|
@app.post("/create")
|
|
async def create():
|
|
try:
|
|
idx = faiss_manager.create()
|
|
log_business(f"[faiss] create idx={idx}")
|
|
return {"idx": idx}
|
|
except Exception as exc:
|
|
raise HTTPException(status_code=400, detail=str(exc))
|
|
|
|
|
|
@app.post("/insert_idx")
|
|
async def insert_idx(data: IndexInsertInput):
|
|
try:
|
|
vector_id = faiss_manager.insert_idx(data.idx, data.id, data.embedding)
|
|
log_business(f"[faiss] insert_idx idx={data.idx} id={vector_id}")
|
|
return {"id": vector_id, "idx": data.idx}
|
|
except Exception as exc:
|
|
raise HTTPException(status_code=400, detail=str(exc))
|
|
|
|
|
|
@app.post("/search_idx")
|
|
async def search_idx(data: IndexSearchInput):
|
|
try:
|
|
result = faiss_manager.search_idx(data.idx, data.embedding, data.k)
|
|
log_business(f"[faiss] search_idx idx={data.idx} ids={result['ids']} similarity_scores={result['similarity_scores']}")
|
|
return result
|
|
except Exception as exc:
|
|
raise HTTPException(status_code=400, detail=str(exc))
|
|
|
|
|
|
@app.post("/del")
|
|
async def delete_idx(data: IndexDeleteInput):
|
|
try:
|
|
faiss_manager.delete(data.idx)
|
|
log_business(f"[faiss] delete idx={data.idx}")
|
|
return {"status": "success", "idx": data.idx}
|
|
except Exception as exc:
|
|
raise HTTPException(status_code=400, detail=str(exc))
|
|
|
|
|
|
@app.post("/persist")
|
|
async def persist():
|
|
faiss_manager.persist()
|
|
return {"status": "success", "message": "索引已持久化"}
|