import aiohttp
from huibiao_framework.client.data_model.model_resp_vo import ModelBaseRespVo
from huibiao_framework.client.llm.abstrast_llm_client import AbstractLlmModelClient
from huibiao_framework.client.llm.schema import HuizeQwen32bAwqDto, HuizeQwen32bAwqVo
from huibiao_framework.config.config import LlmModelNameConstant


class HuiZeQwen32bQwqClient(AbstractLlmModelClient):
    """
    慧泽Qwen-32B模型客户端
    url: http://vllm-qwen-32b.model.hhht.ctcdn.cn:9080/common/query
    request:
        {
        "Action": "NormalChat",
        "DoSample": true,
        "Messages": [
                {
                    "content": "请将下面这段英文翻译成中文：请将下面这段英文翻译成中文：I am a test。",
                    "role": "user"
                }
            ]
        }
    response:
        {
            "code": 0,
            "result": {
                "Output": "我是一个测试。",
                "TokenProbs": [
                    1.0
                ]
            },
            "message": "success"
        }
    """

    def __init__(self, session: aiohttp.ClientSession, url: str = None):
        super().__init__(
            client_name=LlmModelNameConstant.HuiZeQwen32bAwq, session=session, url=url
        )

    def cal_content_len(self, dto: HuizeQwen32bAwqDto):
        return [len(i.content) for i in dto.Messages]

    async def query(
        self, dto: HuizeQwen32bAwqDto, session_id: str = ""
    ) -> ModelBaseRespVo[HuizeQwen32bAwqVo]:
        resp_json = await self.post(dto=dto, session_id=session_id)
        return ModelBaseRespVo[HuizeQwen32bAwqVo](**resp_json)

    async def simple_query(
        self, content, session_id: str = "", do_sample: bool = None, top_k: int = None
    ) -> str:
        dto = HuizeQwen32bAwqDto()

        dto.DoSample = do_sample
        dto.TopK = top_k

        dto.Messages.append(HuizeQwen32bAwqDto.Message(content=content))
        vo: ModelBaseRespVo[HuizeQwen32bAwqVo] = await self.query(dto, session_id)
        return vo.result.Output
