RAG/api/search.py

import json
from cgitb import reset
from concurrent.futures import ThreadPoolExecutor
from typing import List

from fastapi import APIRouter, BackgroundTasks

from schemas import ResponseWrapper
from schemas.schemas import Query, ContentData
from tools_v1 import query_keyword_summary_results, query_keyword_content_results
from utils.deepseek_utils import get_keywords
from utils.json_parse_utils import process_texts_concurrently

router = APIRouter()

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


@router.post("/query", response_model=ResponseWrapper)
async def query_keyword(query: Query):
    keywords = get_keywords(query.text)['keywords']
    print(keywords)
    summary_res = query_keyword_summary_results(keywords)
    content_res = query_keyword_content_results(keywords)
    res = {'summary_results': summary_res, 'content_results': content_res}
    return ResponseWrapper(
        status_code=200,
        detail="success",
        data=res
    )


@router.post("/add/data", response_model=ResponseWrapper)
async def query_keyword(content_list: List[ContentData]):
    param = []
    for content in content_list:
        param.append({'body_text': content.body_text})
    print(json.dumps(param, ensure_ascii=False))

    # 将处理任务提交给后台线程池
    executor.submit(process_texts_concurrently, param)

    return ResponseWrapper(
        status_code=200,
        detail="success",
        data="正在后台处理中"
    )

# @router.post("/query/keyword/content", response_model=ResponseWrapper)
# async def query_keyword(query: Query):
#     res = query_keyword_content_results(query.text)
#     return ResponseWrapper(
#         status_code=200,
#         detail="success",
#         data=res
#     )

# @router.post("/query/embedding", response_model=ResponseWrapper)
# async def query_keyword(query: Query):
#     res = query_embedding_results(query.text)
#     return ResponseWrapper(
#         status_code=200,
#         detail="success",
#         data=res
#     )