# -*- coding: utf-8 -*-
"""
Created on 2024-10-31 
---------
@summary:  智普文件抽取服务
---------
@author: Dzr
"""
import base64

from zhipuai import ZhipuAI


class ZhipuFileExtract:

    def __init__(self, api_key):
        self._client = ZhipuAI(
            api_key=api_key,
            base_url="https://open.bigmodel.cn/api/paas/v4"
        )

    def glm_4v_flash(self, image: bytes, text=None):
        if text is None:
            text = "这是一张字体图片,图片内容是什么？你只需要返回图中内容即可,不要解释,不要返回任何其它内容"

        img_base = base64.b64encode(image).decode('utf-8')
        response = self._client.chat.completions.create(
            model="glm-4v-flash",  # 填写需要调用的模型名称
            messages=[
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": img_base
                            }
                        },
                        {
                            "type": "text",
                            "text": text
                        }
                    ]
                }
            ]
        )
        return response.choices[0].message.content