api/retrieval.py

from fastapi import APIRouter
from pydantic import BaseModel
from retriever.vectordb import search_documents

router = APIRouter()

# 요청받을 데이터 포맷 정의
class SearchRequest(BaseModel):
    query: str
    top_k: int = 5  # 기본적으로 상위 5개 검색

# 응답 데이터 포맷 정의
class SearchResponse(BaseModel):
    documents: list

@router.post("/search", response_model=SearchResponse)
async def search(request: SearchRequest):
    docs = search_documents(request.query, top_k=request.top_k)
    return {"documents": docs}