RAG/api/search.py

from concurrent.futures import ThreadPoolExecutor

from fastapi import APIRouter

from schemas import ResponseWrapper
from schemas.schemas import ContentParam, DatasetParam
from service.content_service import get_contents, add_contents, get_content
from service.dataset_service import get_datasets, add_datasets
from utils.embedding_utils import get_embedding_content_data

router = APIRouter()

# 创建线程池执行器
executor = ThreadPoolExecutor(max_workers=10)


@router.get("/query", response_model=ResponseWrapper)
async def query_keyword(query, datasetIds):
    print(query)
    print(datasetIds)
    embedding_res = get_embedding_content_data(query, datasetIds.split(','))
    res = {'results': embedding_res}
    return ResponseWrapper(
        status_code=200,
        detail="success",
        data=res
    )


@router.get("/content/list", response_model=ResponseWrapper)
async def content_list(page: int = 1, pageSize: int = 10, datasetId: int = None):
    data = get_contents(page, pageSize, datasetId)
    return ResponseWrapper(
        status_code=200,
        detail="success",
        data=data
    )


@router.post("/content/add", response_model=ResponseWrapper)
async def add_content(content_param: ContentParam):
    res = add_contents(content_param)
    return ResponseWrapper(
        status_code=200,
        detail="success",
        data=res
    )


@router.get("/dataset/list", response_model=ResponseWrapper)
async def dataset_list():
    data = get_datasets()
    return ResponseWrapper(
        status_code=200,
        detail="success",
        data=data
    )

@router.post("/dataset/add", response_model=ResponseWrapper)
async def get_dataset(dataset_param : DatasetParam):
    res = add_datasets(dataset_param)
    return ResponseWrapper(
        status_code=200,
        detail="success",
        data=res
    )


@router.get("/content/get", response_model=ResponseWrapper)
async def content_get(docId):
    data = get_content(docId)
    return ResponseWrapper(
        status_code=200,
        detail="success",
        data=data
    )