from server.chat.report_chat import report_chat
from configs import VECTOR_SEARCH_TOP_K, SCORE_THRESHOLD, MAX_TOKENS, LLM_MODELS, DEFAULT_REPORT_BASE
import json
import asyncio
from typing import List, Tuple, Any
from pydantic import BaseModel, Field


class ReportKnowledgeInput(BaseModel):
    location: str = Field(description="The report related query to be searched")


async def report_knowledgebase_search_iter(query: str) -> tuple[str | Any, list[Any] | Any]:
    response = await report_chat(query=query,
                                 fileName=None,
                                 knowledge_base_name=DEFAULT_REPORT_BASE,
                                 model_name=LLM_MODELS[0],
                                 temperature=0.01,
                                 history=[],
                                 top_k=VECTOR_SEARCH_TOP_K,
                                 max_tokens=MAX_TOKENS,
                                 prompt_name="default",
                                 score_threshold=SCORE_THRESHOLD,
                                 stream=False)

    contents = ""
    docs = []
    async for data in response.body_iterator:  # 这里的data是一个json字符串
        data = json.loads(data)
        print("data>>>>>", data)
        contents = data["answer"] if "answer" in data.keys() else contents
        docs = data["docs"] if "docs" in data.keys() else docs
    return contents, docs


def report_knowledgebase_search(query: str) -> tuple[str | Any, list[Any] | Any]:
    return asyncio.run(report_knowledgebase_search_iter(query))


if __name__ == "__main__":
    result = report_knowledgebase_search("大数据男女比例")
    print("答案:", result)