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 images_response(self, query, image_url_list):
        """多图像单轮问答"""
        content = [
            {
                "type": "text",
                "text": query
            }
        ]
        for image_url in image_url_list:
            content.append(
                {
                    "type": "image_url",
                    "image_url": {"url": image_url}
                }
            )
        response = self.client.chat.completions.create(
            model="glm-4v-plus-0111",
            messages=[
                {
                    "role": "user",
                    "content": content
                }
            ],
            response_format={
                "type": "json_object"
            }
        )
        
        return response.choices[0].message
    
    def multi_images_similarity_judge(self, image_url_list1, image_url_list2, query):
        """多组图像中的产品图像相似度对比"""
        content1 = [
            {
                "text": '这是第一组图像，商品信息中的产品图像',
                "type": "text"
            }
        ]
        
        for image_url in image_url_list1:
            content1.append(
                {
                    "type": "image_url",
                    "image_url": {"url": image_url}
                }
            )
        
        content2 = [
            {
                "text": query,
                "type": "text"
            }
        ]
        
        for image_url in image_url_list2:
            content2.append(
                {
                    "type": "image_url",
                    "image_url": {"url": image_url}
                }
            )
            
        response = self.client.chat.completions.create(
            model=self.model_name,
            messages=[
                {
                    "content": content1,
                    "role": "user"
                },
                {
                    "content": content2,
                    "role": "user"
                }
            ],
            response_format = {
                "type": "json_object"
            }
        )
        
        return response.choices[0].message
    
    def image_response(self, query, image_url1, image_url2):
        """单图像问答"""
        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": image_url1}
                        },
                        {
                            "type": "image_url",
                            "image_url": {"url": image_url2}
                        }
                    ]
                }
            ],
            response_format = {
                "type": "json_object"
            }
        )
        
        return response.choices[0].message
    
    def multi_epoch_image_response(self, logo_path, image_url, query):
        """多轮图像问答"""
        image_base = image_to_base(logo_path)
        response = self.client.chat.completions.create(
            model=self.model_name,
            messages=[
                {
                    "content": [
                        {
                            "image_url": {"url": image_base},
                            "type": "image_url"
                        },
                        {
                            "text": "这是李宁的logo",
                            "type": "text"
                        }
                    ],
                    "role": "user"
                },
                {
                    "content": [
                        {
                            "image_url": {"url": image_url},
                            "type": "image_url"
                        },
                        {
                            "text": query,
                            "type": "text"
                        }
                    ],
                    "role": "user"
                }
            ],
            response_format = {
                "type": "json_object"
            }
        )
        
        return response.choices[0].message
    
    def web_search_in_chat(self, search_prompt, content):
        """网络搜索工具"""
        tools = [{
            "type": "web_search",
            "web_search": {
                "enable": True,
                "search_engine": "search_pro_sogou", # 选择搜索引擎
                "search_result": True,
                "search_prompt": search_prompt,
            }
        }]
        
        messages = [{
            "role": "user",
            "content": content
        }]
        
        response = self.client.chat.completions.create(
            model=self.model_name,
            messages=messages,
            tools=tools,
            response_format={
                "type": "json_object"
            }
        )
        
        return response.choices[0].message
        
        
    
    def set_modelname(self, modelname):
        self.model_name = modelname
    
if __name__ == '__main__':
    glm = Glm()
    glm.set_modelname("glm-4v-plus-0111")
    query = f"""
        你是一个产品图像分析助手，你的任务是判断第二张图像中的产品上是否包含logo，并与第一张logo图像做对比，判断第二张图像中的logo是否与第一张的一致。
        判断思路：
        1. 首先判断第二张图像中是否包含logo。
        2. 如果包含logo的话，再进行判断，是否与第一张的logo一致。
        3. 如果与第一张的logo不一致，请根据你的经验判断图像中的logo是什么品牌，如果不知道返回'未知'
        最终结果返回为以下给出的json格式
        输出结果示例：
        
        {{
            "is_contain_logo": false
        }}
        
        
        {{
            "is_contain_logo": true,
            "is_jugement_logo": true,
            "brand_name": "李宁"
        }}
    """
    response = glm.multi_epoch_image_response("./logo/lining.jpg", "http://h2.appsimg.com/a.appsimg.com/upload/merchandise/pdcvis/613214/2024/1120/88/9b9027dd-95b7-4024-b71e-fb7cbfde16a1.jpg", query)
    print(response.content)