AgentCoreService/pqai_agent/toolkit/image_describer.py

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):
        self.model = VOLCENGINE_MODEL_DOUBAO_1_5_VISION_PRO
        self.llm_client = chat_service.OpenAICompatible.create_client(self.model)
        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.
        """
        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}")
        return response_content

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