Agent/agent/prompts/wrapper.py

"""
Prompt Wrapper - 为 .prompt 文件提供 Prompt 实现

类似 Resonote 的 SimpleHPrompt，但增加了多模态支持
"""

import base64
from pathlib import Path
from typing import List, Dict, Any, Union, Optional
from agent.prompts.loader import load_prompt, get_message


class SimplePrompt:
    """
    通用的 Prompt 包装器

    特性：
    - 加载 .prompt 文件（YAML frontmatter + sections）
    - 支持参数替换（%variable%）
    - 支持多模态消息（图片）

    使用示例：
        # 纯文本
        prompt = SimplePrompt(Path("task.prompt"))
        messages = prompt.build_messages(text="内容")

        # 多模态（文本 + 图片）
        messages = prompt.build_messages(
            text="分析这张图片",
            images="path/to/image.png"  # 或 images=["img1.png", "img2.png"]
        )
    """

    def __init__(self, prompt_path: Union[Path, str]):
        """
        Args:
            prompt_path: .prompt 文件路径
        """
        self.prompt_path = Path(prompt_path) if isinstance(prompt_path, str) else prompt_path

        # 加载 .prompt 文件
        self.config, self._messages = load_prompt(self.prompt_path)

    def build_messages(self, **context) -> List[Dict[str, Any]]:
        """
        构造消息列表（支持多模态）

        Args:
            **context: 参数
                - 普通参数：用于替换 %variable%
                - images: 图片资源（可选）
                  - 单个图片：str 或 Path
                  - 多个图片：List[str | Path]
                  - 格式：文件路径或 base64 字符串

        Returns:
            消息列表，格式遵循 OpenAI API 规范

        Example:
            >>> messages = prompt.build_messages(
            ...     text="特征描述",
            ...     images="input/image.png"
            ... )
            [
                {"role": "system", "content": "..."},
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": "..."},
                        {"type": "image_url", "image_url": {"url": "data:image/png;base64,..."}}
                    ]
                }
            ]
        """
        # 提取图片资源（从 context 中移除，避免传入 get_message）
        images = context.pop('images', None)

        # 构建文本内容（支持参数替换）
        system_content = get_message(self._messages, 'system', **context)
        user_content = get_message(self._messages, 'user', **context)

        messages = []

        # 添加 system 消息
        if system_content:
            messages.append({"role": "system", "content": system_content})

        # 添加 user 消息（可能是多模态）
        if images:
            # 多模态消息
            user_message = {"role": "user", "content": []}

            # 添加文本部分
            if user_content:
                user_message["content"].append({
                    "type": "text",
                    "text": user_content
                })

            # 添加图片部分
            if isinstance(images, (list, tuple)):
                for img in images:
                    user_message["content"].append(self._build_image_content(img))
            else:
                user_message["content"].append(self._build_image_content(images))

            messages.append(user_message)
        else:
            # 纯文本消息
            if user_content:
                messages.append({"role": "user", "content": user_content})

        return messages

    def _build_image_content(self, image: Union[str, Path]) -> Dict[str, Any]:
        """
        构建图片内容部分（OpenAI 格式）

        Args:
            image: 图片路径或 base64 字符串

        Returns:
            {"type": "image_url", "image_url": {"url": "data:..."}}
        """
        # 如果已经是 base64 data URL，直接使用
        if isinstance(image, str) and image.startswith("data:"):
            return {
                "type": "image_url",
                "image_url": {"url": image}
            }

        # 否则，读取文件并转为 base64
        image_path = Path(image) if isinstance(image, str) else image

        # 推断 MIME type
        suffix = image_path.suffix.lower()
        mime_type_map = {
            '.png': 'image/png',
            '.jpg': 'image/jpeg',
            '.jpeg': 'image/jpeg',
            '.gif': 'image/gif',
            '.webp': 'image/webp'
        }
        mime_type = mime_type_map.get(suffix, 'image/png')

        # 读取并编码
        with open(image_path, 'rb') as f:
            image_data = base64.b64encode(f.read()).decode('utf-8')

        data_url = f"data:{mime_type};base64,{image_data}"

        return {
            "type": "image_url",
            "image_url": {"url": data_url}
        }


def create_prompt(prompt_path: Union[Path, str]) -> SimplePrompt:
    """
    工厂函数：创建 SimplePrompt 实例

    Args:
        prompt_path: .prompt 文件路径

    Returns:
        SimplePrompt 实例
    """
    return SimplePrompt(prompt_path)