import aiohttp

from huibiao_framework.client.llm import AbstractLlmModelClient
from huibiao_framework.client.llm.schema import MindleDs32bAwqDto, MindleDs32bAwqVo
from huibiao_framework.config.config import LlmModelNameConstant, ModelConfig


class MindleDs32bAwqClient(AbstractLlmModelClient):
    """
    curl --location --request POST 'http://mindie-ds-32b-awq.office.sh.ctyun.cn:9080/v1/chat/completions' \
    --header 'Content-Type: application/json' \
    --data-raw '{
        "model": "huize",
        "messages": [
            {
                "role": "user",
                "content": "你好"
            }
        ]
    }'

    响应
        {
        "id": "chatcmpl-ab718673ce220e43572b4055414a7362",
        "choices": [
            {
                "finish_reason": "stop",
                "index": 0,
                "message": {
                    "role": "assistant",
                    "content": "<think>\n\n</think>\n\n你好！很高兴见到你，有什么我可以帮忙的吗？无论是学习、工作还是生活中的问题，都可以告诉我哦！😊",
                    "function_call": null,
                    "tool_calls": []
                },
                "logprobs": null
            }
        ],
        "created": 1761207189,
        "model": "huize",
        "object": "chat.completion",
        "usage": {
            "prompt_tokens": 5,
            "total_tokens": 36,
            "completion_tokens": 31
        }
    }
    """

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

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

    async def query(
        self, dto: MindleDs32bAwqDto, session_id: str = ""
    ) -> MindleDs32bAwqVo:
        dto.model = ModelConfig.OPEN_AI_MODEL_TYPE
        resp_json = await self.post(dto=dto, session_id=session_id)
        return MindleDs32bAwqVo(**resp_json)

    async def simple_query(self, content, session_id: str = "", **kwargs) -> str:
        dto = MindleDs32bAwqDto()
        dto.messages.append(MindleDs32bAwqDto.Message(content=content))
        vo: MindleDs32bAwqVo = await self.query(dto=dto, session_id=session_id)
        assert vo.choices and vo.choices[0].message, f"响应数据为空:{vo.model_dump()}"
        return vo.choices[0].message.content
