BrandAnalysis/agent/glm.py

from config import load_config
from utils import image_to_base
from zhipuai import ZhipuAI
class Glm:
    _instance = None
    def __new__(cls):
        if not cls._instance:
            cls._instance = super(Glm, cls).__new__(cls)
            cls._instance._initialized = False
        return cls._instance
    
    def __init__(self):
        if not self._initialized:
            self.cfg = load_config()['glm']
            self.client = ZhipuAI(api_key=self.cfg["api_key"])
            self._initialized = True
            
        self.model_name = "glm-4-plus"
    
    def text_response(self, query):
        """文字问答"""
        resonse = self.client.chat.completions.create(
            model=self.model_name,
            messages=[
                {"role": "user", "content": query}
            ],
            response_format= {
                'type': 'json_object'
            }
        )
        return resonse.choices[0].message
    
    def image_response(self, query):
        """单图像问答"""
        response = self.client.chat.completions.create(
            model="glm-4v-plus-0111",
            messages=[
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "text",
                            "text": query
                        },
                        {
                            "type": "image_url",
                            "image_url": {"url": "http://h2.appsimg.com/a.appsimg.com/upload/merchandise/pdcvis/613214/2024/1120/88/9b9027dd-95b7-4024-b71e-fb7cbfde16a1.jpg"}
                        }
                    ]
                }
            ]
        )
        
        return response.choices[0].message
    
    def multi_epoch_image_response(self, logo_path):
        """多轮图像问答"""
        image_base = image_to_base(logo_path)
        response = self.client.chat.completions.create(
            model="glm-4v-plus-0111",
            messages=[
                {
                    "content": [
                        {
                            "image_url": {"url": image_base},
                            "type": "image_url"
                        },
                        {
                            "text": "这是李宁的logo",
                            "type": "text"
                        }
                    ],
                    "role": "user"
                },
                {
                    "content": [
                        {
                            "image_url": {"url": "http://h2.appsimg.com/a.appsimg.com/upload/merchandise/pdcvis/2023/04/14/79/d75fa3db-3bec-45c4-b7a6-54332d42e373.jpg"},
                            "type": "image_url"
                        },
                        {
                            "text": "第二张图像中的产品是否包含logo，如果包含的话是否是第一张的logo？",
                            "type": "text"
                        }
                    ],
                    "role": "user"
                }
            ]
        )
        
        return response.choices[0].message
    
    def set_modelname(self, modelname):
        self.model_name = modelname
    
if __name__ == '__main__':
    glm = Glm()
    response = glm.multi_epoch_image_response("./logo/lining.jpg")
    print(response)