trade-assistant/backend/app/ai/providers/openai.py

from typing import Dict, Any, Optional
import json
from app.ai.base import AIProvider


SYSTEM_PROMPTS = {
    "translate": "You are a professional translator specialized in foreign trade and e-commerce. "
                 "Accurately translate business terms like MOQ, FOB, CIF, lead time, etc. "
                 "Return ONLY the translated text, no explanations.",
    "reply": "You are an experienced foreign trade sales expert. Write professional, "
             "clear business replies. Be concise but warm. Include relevant details "
             "naturally. Return ONLY the reply text, no explanations.",
    "marketing": "You are a creative copywriter for international trade. Write compelling "
                 "marketing content that drives action. Adapt to the target audience's culture. "
                 "Return ONLY the copy, no explanations.",
    "extract": "You extract structured data from text. Return ONLY valid JSON matching the requested schema.",
    "chat": "你是 TradeMate（外贸小助手）的 AI 助手。你的职责是帮助外贸从业者解答关于本工具使用的问题，以及提供外贸业务建议。\n"
            "你可以回答的问题包括：\n"
            "- 功能介绍：翻译、客户管理、产品管理、报价单、营销文案、WhatsApp 集成等\n"
            "- 使用帮助：如何添加客户、如何生成报价单、如何导出数据等\n"
            "- 外贸知识：贸易术语（FOB、CIF 等）、谈判技巧、跟进策略等\n\n"
            "回答要求：\n"
            "- 简洁扼要，用中文回答\n"
            "- 涉及操作步骤时用数字列表说明\n"
            "- 不确定的问题不要编造，直接说需要查证\n"
            "- 语气友好专业",
}


class OpenAIProvider(AIProvider):
    def __init__(self, api_key: str, model: str = "gpt-4o", base_url: Optional[str] = None, http_client=None):
        try:
            from openai import AsyncOpenAI
        except ImportError:
            raise ImportError(
                "openai>=1.0 is required for OpenAIProvider. "
                "Install it with: pip install 'openai>=1.0'"
            )
        kwargs = {"api_key": api_key}
        if base_url:
            kwargs["base_url"] = base_url
        if http_client:
            kwargs["http_client"] = http_client
        self.client = AsyncOpenAI(**kwargs)
        self.model = model
        self._name = f"openai-{model}"
        self._pricing = {
            "gpt-4o": {"input": 0.01, "output": 0.03},
            "gpt-4o-mini": {"input": 0.0015, "output": 0.006},
        }
        self._cheap_model = "gpt-4o-mini" if model == "gpt-4o" else model

    async def translate(self, text: str, source_lang: Optional[str], target_lang: str, context: Optional[str] = None) -> Dict[str, Any]:
        from app.services.translation_quota import TranslationQuotaService
        from app.database import AsyncSessionLocal

        async with AsyncSessionLocal() as db:
            quota_svc = TranslationQuotaService(db)
            if not await quota_svc.check_quota("llm"):
                raise Exception("LLM translation quota exhausted or disabled")
            result = await self._do_translate(text, source_lang, target_lang, context)
            if result and result.get("translated_text"):
                await quota_svc.consume("llm", len(text))
                await db.commit()
            return result

    async def _do_translate(self, text: str, source_lang: Optional[str], target_lang: str, context: Optional[str] = None) -> Dict[str, Any]:
        system = SYSTEM_PROMPTS["translate"]
        if context:
            system += f"\nContext: this is about {context}"
        if source_lang and source_lang != "auto":
            system += f"\nSource language: {source_lang}"

        content = await self._call(system, f"Translate to {target_lang}:\n\n{text}", model=self._cheap_model)
        return {"translated_text": content, "provider": self.name, "model": self.model}

    async def reply(self, inquiry: str, context: Optional[Dict[str, Any]] = None, tone: str = "professional", preference_context: Optional[str] = None) -> Dict[str, Any]:
        system = SYSTEM_PROMPTS["reply"] + f"\nTone: {tone}"
        if preference_context:
            system += f"\nUser preference: {preference_context}"

        context_str = ""
        if context:
            if context.get("product"):
                context_str += f"Product: {context['product']}\n"
            if context.get("price"):
                context_str += f"Price: {context['price']}\n"
            if context.get("customer_history"):
                context_str += f"Customer history: {context['customer_history']}\n"
            if context.get("conversation_history"):
                context_str += f"Previous messages: {context['conversation_history']}\n"

        prompt = f"{context_str}\nCustomer inquiry:\n{inquiry}\n\nWrite a reply:"
        content = await self._call(system, prompt)
        return {"reply": content, "provider": self.name, "model": self.model}

    async def generate_marketing(self, product_info: Dict[str, Any], target: str, style: str = "professional", language: str = "en", preference_context: Optional[str] = None) -> Dict[str, Any]:
        system = SYSTEM_PROMPTS["marketing"] + f"\nStyle: {style}\nTarget audience: {target}\nLanguage: {language}"
        if preference_context:
            system += f"\nUser preference: {preference_context}"

        product_str = json.dumps(product_info, ensure_ascii=False, indent=2)
        prompt = f"Product information:\n{product_str}\n\nGenerate marketing copy:"
        content = await self._call(system, prompt)
        return {"content": content, "provider": self.name, "model": self.model}

    async def extract_info(self, text: str, schema: Dict[str, Any]) -> Dict[str, Any]:
        system = SYSTEM_PROMPTS["extract"]
        schema_str = json.dumps(schema, indent=2)
        prompt = f"Schema:\n{schema_str}\n\nText:\n{text}\n\nExtracted JSON:"
        try:
            content = await self._call(system, prompt, response_format={"type": "json_object"})
        except Exception:
            content = await self._call(system, prompt)
        try:
            data = json.loads(content)
            return {"data": data, "confidence": 0.9, "provider": self.name}
        except json.JSONDecodeError:
            return {"data": {}, "confidence": 0.0, "provider": self.name, "error": "parse_failed"}

    async def chat(self, message: str, history: list = None, system_prompt: str = None) -> Dict[str, Any]:
        system = system_prompt or SYSTEM_PROMPTS["chat"]
        messages = [{"role": "system", "content": system}]
        if history:
            for h in history[-10:]:
                messages.append(h)
        messages.append({"role": "user", "content": message})

        kwargs = {
            "model": self._cheap_model,
            "messages": messages,
            "max_tokens": 2000,
            "temperature": 0.7,
        }
        resp = await self.client.chat.completions.create(**kwargs)
        content = resp.choices[0].message.content or ""
        return {"reply": content, "provider": self.name, "model": self.model}

    async def _call(self, system: str, prompt: str, max_tokens: int = 3000, response_format: Optional[Dict] = None, model: Optional[str] = None) -> str:
        kwargs = {
            "model": model or self.model,
            "messages": [
                {"role": "system", "content": system},
                {"role": "user", "content": prompt},
            ],
            "max_tokens": max_tokens,
            "temperature": 0.7,
        }
        if response_format:
            kwargs["response_format"] = response_format

        resp = await self.client.chat.completions.create(**kwargs)
        content = resp.choices[0].message.content
        
        if content is None and hasattr(resp.choices[0].message, 'reasoning'):
            reasoning = resp.choices[0].message.reasoning
            if reasoning:
                import re
                final_output_patterns = [
                    r'Final Output Generation[:：]\s*(.+?)(?:\n\n|$)',
                    r'Final Output[:：]\s*(.+?)(?:\n\n|$)',
                    r'7\.\s*Final Output Generation[:：]\s*(.+?)(?:\n\n|$)',
                    r'翻译结果[：:]\s*(.+?)(?:\n\n|$)',
                    r'最终输出[：:]\s*(.+?)(?:\n\n|$)',
                ]
                for pattern in final_output_patterns:
                    match = re.search(pattern, reasoning, re.DOTALL)
                    if match:
                        content = match.group(1).strip()
                        break
                
                if content is None:
                    paragraphs = re.split(r'\n\n+', reasoning.strip())
                    if paragraphs:
                        for p in reversed(paragraphs):
                            p = p.strip()
                            if p and len(p) > 10:
                                if not p.startswith('步骤') and not p.startswith('Step'):
                                    content = p
                                    break
        
        if content is None and hasattr(resp.choices[0].message, 'reasoning'):
            reasoning = resp.choices[0].message.reasoning
            if reasoning:
                import re
                cleaned = re.sub(r'^步骤\d+[：:].*$', '', reasoning, flags=re.MULTILINE)
                cleaned = re.sub(r'^Step \d+[：:].*$', '', cleaned, flags=re.MULTILINE)
                cleaned = re.sub(r'\n+', '\n', cleaned).strip()
                if cleaned:
                    content = cleaned
        
        return content

    @property
    def name(self) -> str:
        return self._name

    @property
    def cost_per_1k_tokens(self) -> float:
        p = self._pricing.get(self.model, {"input": 0.01, "output": 0.03})
        return (p["input"] + p["output"]) / 2