modal_deploy/client.py

"""
PregoPal - MiniCPM-o API 客户端

用法：
    from modal_deploy.client import MiniCPMClient
    
    client = MiniCPMClient(base_url="https://your-app.modal.run")
    response = client.chat("今天孕妇可以吃什么？")
"""
import json
import logging
from typing import Optional
from dataclasses import dataclass, field

import requests

logger = logging.getLogger(__name__)


@dataclass
class ChatMessage:
    role: str
    content: str


@dataclass
class MultimodalContent:
    type: str  # "text" or "image_url"
    text: Optional[str] = None
    image_url: Optional[dict] = None


class MiniCPMClient:
    """封装 MiniCPM-o API 的客户端"""

    def __init__(
        self,
        base_url: str = "http://localhost:8080",
        timeout: int = 120,
        default_max_tokens: int = 1024,
        default_temperature: float = 0.7,
    ):
        self.base_url = base_url.rstrip("/")
        self.timeout = timeout
        self.default_max_tokens = default_max_tokens
        self.default_temperature = default_temperature

    def _post(self, path: str, body: dict) -> dict:
        url = f"{self.base_url}{path}"
        try:
            r = requests.post(url, json=body, timeout=self.timeout)
            r.raise_for_status()
            return r.json()
        except requests.exceptions.RequestException as e:
            logger.error(f"API request failed: {e}")
            raise

    def health(self) -> dict:
        """检查服务健康状态"""
        r = requests.get(f"{self.base_url}/health", timeout=10)
        return r.json()

    # ── 文本对话 ──────────────────────────────────────────

    def chat(
        self,
        messages: list[dict],
        system_prompt: Optional[str] = None,
        max_tokens: Optional[int] = None,
        temperature: Optional[float] = None,
        stream: bool = False,
    ) -> dict | str:
        """
        OpenAI 兼容的聊天接口
        
        Args:
            messages: 消息列表 [{"role": "user", "content": "..."}]
            system_prompt: 系统提示词
            max_tokens: 最大生成长度
            temperature: 温度参数
            stream: 是否流式
        
        Returns:
            dict 或流式字符串
        """
        msgs = []
        if system_prompt:
            msgs.append({"role": "system", "content": system_prompt})
        msgs.extend(messages)

        body = {
            "model": "MiniCPM-o-4_5",
            "messages": msgs,
            "max_tokens": max_tokens or self.default_max_tokens,
            "temperature": temperature or self.default_temperature,
            "stream": stream,
        }
        return self._post("/v1/chat/completions", body)

    def ask(self, prompt: str, system_prompt: Optional[str] = None) -> str:
        """简化调用：发送问题，返回文本回答"""
        result = self.chat(
            messages=[{"role": "user", "content": prompt}],
            system_prompt=system_prompt,
        )
        try:
            return result["choices"][0]["message"]["content"]
        except (KeyError, IndexError, TypeError):
            logger.error(f"Unexpected response format: {result}")
            return str(result)

    # ── 多模态（图片理解） ───────────────────────────────

    def chat_with_image(
        self,
        prompt: str,
        image_base64: str,
        image_format: str = "jpeg",
        max_tokens: Optional[int] = None,
        temperature: Optional[float] = None,
    ) -> dict:
        """
        发送图片 + 文字进行多模态对话
        
        Args:
            prompt: 文字提示
            image_base64: base64 编码的图片
            image_format: 图片格式 (jpeg, png, webp)
        """
        body = {
            "messages": [
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": prompt},
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:image/{image_format};base64,{image_base64}"
                            },
                        },
                    ],
                }
            ],
            "max_tokens": max_tokens or self.default_max_tokens,
            "temperature": temperature or self.default_temperature,
        }
        return self._post("/v1/chat/completions", body)

    def describe_image(self, image_base64: str, image_format: str = "jpeg") -> str:
        """简化调用：描述图片内容"""
        result = self.chat_with_image(
            prompt="请详细描述这张图片中的内容。",
            image_base64=image_base64,
            image_format=image_format,
        )
        return result.get("response", str(result))

    # ── 嵌入向量 ──────────────────────────────────────────

    def embed(self, texts: list[str]) -> list[list[float]]:
        """获取文本的嵌入向量"""
        body = {
            "model": "MiniCPM-o-4_5",
            "input": texts,
        }
        result = self._post("/v1/embeddings", body)
        return [item["embedding"] for item in result.get("data", [])]


# ── 预设提示词 (针对孕期陪护场景) ────────────────────────

SYSTEM_PROMPTS = {
    "diet_advisor": """你是一位专业的孕期营养顾问。请根据以下原则提供建议：
1. 优先考虑食材的易得性和家庭制作能力
2. 注意孕期各阶段的营养需求差异
3. 避免生冷、辛辣、高汞鱼类等孕期禁忌食物
4. 推荐富含叶酸、铁、钙、DHA 的食物
5. 回答简洁实用，给出具体可操作的菜品建议""",

    "meal_planner": """你是一位家庭营养师，专门为孕妇设计日常菜单。要求：
1. 每天设计早中晚三餐 + 2次加餐
2. 确保营养均衡，覆盖蛋白质、碳水、脂肪、维生素、矿物质
3. 考虑中国家庭的烹饪习惯和常见食材
4. 标注每餐的关键营养素
5. 给出简要的制作步骤""",

    "food_analyzer": """分析用户提供的食物/菜品描述，提取以下信息：
1. 菜名
2. 主要食材
3. 营养成分评估（热量、蛋白质、脂肪、碳水、关键微量元素）
4. 孕期适宜度（推荐/可食用/慎食/忌食）
5. 建议（如有）
请用 JSON 格式回复。""",

    "nutrition_summary": """根据用户提供的饮食记录，总结营养摄入情况：
1. 各大类营养素摄入评估
2. 与孕期推荐标准的对比
3. 需要补充或调整的方向
4. 具体的一周饮食改善建议
5. 用结构化格式输出，便于前端可视化""",
}