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.data_utils import add_data
from utils.deepseek_utils import get_keywords
from utils.embedding_utils import get_embedding_content_data

router = APIRouter()

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


@router.post("/query", response_model=ResponseWrapper)
async def query_keyword(query: Query):
    print(query.text)
    keywords = get_keywords(query.text)['keywords']
    print(keywords)
    summary_res = query_keyword_summary_results(keywords)
    content_res = query_keyword_content_results(keywords)
    embedding_res = get_embedding_content_data(query.text)
    res = {'summary_results': summary_res, 'content_results': content_res, 'embedding_results': embedding_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]):
    res_list = []
    for content in content_list:
        if content.body_text:
            print(content.body_text)
            res = add_data(content.body_text)
            res_list.append(res)
    return ResponseWrapper(
        status_code=200,
        detail="success",
        data=res_list
    )

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