AgentCoreService/pqai_agent/toolkit/image_describer.py

import diskcache

from pqai_agent import chat_service
from pqai_agent.chat_service import VOLCENGINE_MODEL_DOUBAO_1_5_VISION_PRO
from pqai_agent.logging_service import logger
from pqai_agent.toolkit.base import BaseToolkit
from pqai_agent.toolkit.function_tool import FunctionTool


class ImageDescriber(BaseToolkit):
    def __init__(self, cache_dir: str = None):
        self.model = VOLCENGINE_MODEL_DOUBAO_1_5_VISION_PRO
        self.llm_client = chat_service.OpenAICompatible.create_client(self.model)
        if not cache_dir:
            cache_dir = 'image_descriptions_cache'
        self.cache = diskcache.Cache(cache_dir, size_limit=100*1024*1024)
        super().__init__()

    def analyse_image(self, image_url: str):
        """Takes an image URL as input and returns a detailed description of the image.

        Args:
            image_url (str): The URL of the image to be described.
        Returns:
            str: A detailed description of the image.
        """
        if image_url in self.cache:
            logger.debug(f"Cache hit for image URL: {image_url}")
            return self.cache[image_url]

        system_prompt = "你是一位图像分析专家。请提供输入图像的详细描述，包括图像中的文本内容（如果存在）"

        messages = [
            {'role': 'system', 'content': system_prompt},
            {'role': 'user', 'content': [
                {
                    'type': 'image_url',
                    'image_url': image_url
                }
            ]}
        ]
        response = self.llm_client.chat.completions.create(messages=messages, model=self.model)
        response_content = response.choices[0].message.content
        logger.debug(f"ImageDescriber response: {response_content}")
        self.cache[image_url] = response_content
        return response_content

    def get_tools(self):
        return [FunctionTool(self.analyse_image)]