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 service.search_service import query_data
from utils.deepseek_utils import generate_summary_query

router = APIRouter()

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


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


@router.get("/chat", response_model=ResponseWrapper)
async def chat(query, datasetIds):
    res = query_data(query, datasetIds.split(','))
    chat_res = generate_summary_query(query, res)
    data = {'results': res, 'chat_res': chat_res}
    return ResponseWrapper(
        status_code=200,
        detail="success",
        data=data
    )


@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
    )