llm.py

import os, math, time
from typing import List, Dict, Any, Optional
from openai import OpenAI

def _norm(x):
    if x is None: return ""
    if isinstance(x, float) and math.isnan(x): return ""
    return str(x)

def _backoff(attempt):
    # 0.5, 1, 2, 4 ... (上限 8s)
    return min(0.5 * (2 ** attempt), 8.0)

class OpenAILLM:
    def __init__(self, model_chat: str = "gpt-4o-mini", model_translate: str = "gpt-4o-mini"):
        api_key = os.environ.get("OPENAI_API_KEY2")
        if not api_key:
            raise ValueError("環境変数 OPENAI_API_KEY2 が設定されていません。")
        self.client = OpenAI(api_key=api_key)
        self.model_chat = model_chat
        self.model_translate = model_translate
        self.last_usage = {"prompt_tokens":0, "completion_tokens":0, "total_tokens":0}

    def _update_usage(self, rsp):
        try:
            u = rsp.usage
            if u:
                self.last_usage = {
                    "prompt_tokens": getattr(u, "prompt_tokens", 0),
                    "completion_tokens": getattr(u, "completion_tokens", 0),
                    "total_tokens": getattr(u, "total_tokens", 0),
                }
        except Exception:
            pass

    def _chat(self, model, messages, temperature=0.2, max_retries=3):
        for i in range(max_retries+1):
            try:
                rsp = self.client.chat.completions.create(
                    model=model, messages=messages, temperature=temperature
                )
                self._update_usage(rsp)
                return rsp.choices[0].message.content.strip()
            except Exception as e:
                if i == max_retries:
                    raise
                time.sleep(_backoff(i))

    def generate_ceo_message(self, meta, kpi: Dict[str, float], esg_rows: List[Dict[str, Any]]) -> str:
        prompt = (
            "以下の企業情報・KPI・ESG指標をもとに、日本語で200字程度のCEOメッセージ草案を出力。"
            "事実ベース・簡潔・投資家向け。数値は丸め過ぎないこと。\n\n"
            f"企業情報: {meta.model_dump()}\nKPI: {kpi}\nESG: {esg_rows}\n"
        )
        return self._chat(self.model_chat, [{"role":"user","content":prompt}], temperature=0.2)

    def generate_risk_opportunity(self, meta, kpi: Dict[str, float], esg_rows: List[Dict[str, Any]]) -> str:
        prompt = (
            "以下に基づき主要なリスクと機会を150字程度で日本語要約。具体的観点を1-2点：\n\n"
            f"企業情報: {meta.model_dump()}\nKPI: {kpi}\nESG: {esg_rows}\n"
        )
        return self._chat(self.model_chat, [{"role":"user","content":prompt}], temperature=0.2)

    def translate_texts(self, texts: List[Any], target_lang: str = "en", glossary: Optional[Dict[str,str]] = None) -> List[str]:
        norm = [_norm(t) for t in texts]
        SEP = "\n<<<SEP>>>\n"
        rules = ""
        if glossary:
            rules = "用語統一ルール（厳守）:\n" + "\n".join([f"- {k} -> {v}" for k,v in glossary.items()])

        system = (
            "You are a precise financial/ESG translator. Preserve numbers and units. "
            "Follow the glossary strictly. Keep tone concise."
        )
        prompt = f"Translate the following into {target_lang}. Each part is separated by <<<SEP>>>.\n{rules}\n\n" + SEP.join(norm)

        txt = self._chat(self.model_translate, [{"role":"system","content":system},{"role":"user","content":prompt}], temperature=0.1)
        parts = [p.strip() for p in txt.split("<<<SEP>>>")]
        if len(parts) != len(norm):
            parts = [txt] + norm[1:]
            parts = parts[:len(norm)]
        return parts