82 lines
2.5 KiB
Python
82 lines
2.5 KiB
Python
from uuid import UUID
|
|
|
|
from fastapi import APIRouter, Depends, HTTPException, Query
|
|
from sqlalchemy import func, select
|
|
from sqlalchemy.orm import Session
|
|
|
|
from app.db.session import get_db
|
|
from app.models import RetrievalRun
|
|
from app.schemas.catalog import Page
|
|
from app.schemas.hybrid import (
|
|
ChunkKeywordRequest,
|
|
ChunkKeywordResponse,
|
|
HybridRequest,
|
|
HybridResponse,
|
|
)
|
|
from app.schemas.retrieval import RetrievalRequest, RetrievalResponse
|
|
from app.schemas.vectors import VectorRetrievalResponse
|
|
from app.services.auth import require_admin
|
|
from app.services.embedding_runtime import get_embedding_provider
|
|
from app.services.hybrid import retrieve_chunks, retrieve_hybrid
|
|
from app.services.retrieval import retrieve
|
|
|
|
router = APIRouter(prefix="/retrieval", tags=["retrieval"])
|
|
|
|
RunResponse = RetrievalResponse | VectorRetrievalResponse | ChunkKeywordResponse | HybridResponse
|
|
|
|
|
|
@router.post("", response_model=RetrievalResponse)
|
|
def keyword_retrieval(payload: RetrievalRequest, db: Session = Depends(get_db)):
|
|
"""只检索已核验知识；保存原始检索结果供复核，不调用大模型。"""
|
|
return retrieve(db, payload)
|
|
|
|
|
|
@router.post("/chunks", response_model=ChunkKeywordResponse)
|
|
def chunk_keyword_retrieval(payload: ChunkKeywordRequest, db: Session = Depends(get_db)):
|
|
return retrieve_chunks(db, payload)
|
|
|
|
|
|
@router.post("/hybrid", response_model=HybridResponse)
|
|
def hybrid_retrieval(
|
|
payload: HybridRequest,
|
|
db: Session = Depends(get_db),
|
|
provider=Depends(get_embedding_provider),
|
|
):
|
|
return retrieve_hybrid(db, payload, provider)
|
|
|
|
|
|
@router.get(
|
|
"/runs",
|
|
response_model=Page[RunResponse],
|
|
dependencies=[Depends(require_admin)],
|
|
)
|
|
def list_runs(
|
|
page: int = Query(1, ge=1),
|
|
page_size: int = Query(20, ge=1, le=100),
|
|
db: Session = Depends(get_db),
|
|
):
|
|
rows = db.scalars(
|
|
select(RetrievalRun)
|
|
.order_by(RetrievalRun.created_at.desc(), RetrievalRun.id)
|
|
.offset((page - 1) * page_size)
|
|
.limit(page_size)
|
|
)
|
|
return {
|
|
"items": [row.response for row in rows],
|
|
"total": db.scalar(select(func.count()).select_from(RetrievalRun)),
|
|
"page": page,
|
|
"page_size": page_size,
|
|
}
|
|
|
|
|
|
@router.get(
|
|
"/runs/{run_id}",
|
|
response_model=RunResponse,
|
|
dependencies=[Depends(require_admin)],
|
|
)
|
|
def read_run(run_id: UUID, db: Session = Depends(get_db)):
|
|
row = db.get(RetrievalRun, str(run_id))
|
|
if row is None:
|
|
raise HTTPException(404, "检索记录不存在")
|
|
return row.response
|