# -*- coding: utf-8 -*-
# 
# @author: lyg
# @date: 2025-5-8
# @version: 1
# @description: 利用LLM将图片转为文本。

from langchain_core.prompts import HumanMessagePromptTemplate, ChatPromptTemplate
from langchain_core.messages import HumanMessage, SystemMessage
from langchain_core.output_parsers import JsonOutputParser
import json
import base64

from knowledgebase.llm import vision_llm


class ImageToText:
    def __init__(self):
        self.llm = vision_llm
        self.prompt = ChatPromptTemplate.from_messages([
            ("system", "你是一个资深软件工程师，请分析图片中的内容。"),
            (
                "user",
                [
                    {"type": "text", "text": "{msg}"},
                    {
                        "type": "image_url",
                        "image_url": {"url": "data:image/jpeg;base64,{image}"},
                    }
                ],
            )
        ])

    def gen_text_from_img(self, image: bytes) -> str:
        """
        从图片生成文本。

        :param image:  图片数据
        :return: 文本
        """
        image = base64.b64encode(image).decode()
        chain = self.prompt | self.llm
        resp = chain.invoke({"msg": "使用自然语言输出图片中的内容，不要做过多的解释。输出格式为纯文本。", "image": image})
        return resp.content