cfc/cfc-langgraph/src/llm/client.py

import os
import time
import functools
from langchain_openai import ChatOpenAI
from app.monitoring import llm_calls_total, llm_duration_seconds, llm_tokens_total


def _common_kwargs() -> dict:
    return {
        "api_key": os.getenv("LLM_API_KEY", os.getenv("OPENAI_API_KEY", "")),
        "base_url": os.getenv("LLM_BASE_URL", os.getenv("OPENAI_BASE_URL", "https://api.openai.com/v1")),
        "timeout": 180,
        "max_retries": 1,
    }


def _wrap_llm_ainvoke(llm: ChatOpenAI, model_name: str) -> ChatOpenAI:
    original_ainvoke = llm.ainvoke

    @functools.wraps(original_ainvoke)
    async def monitored_ainvoke(*args, **kwargs):
        llm_calls_total.labels(model=model_name, status="started").inc()
        start = time.perf_counter()
        try:
            result = await original_ainvoke(*args, **kwargs)
            elapsed = time.perf_counter() - start
            llm_duration_seconds.labels(model=model_name).observe(elapsed)
            llm_calls_total.labels(model=model_name, status="success").inc()
            try:
                usage = getattr(result, "response_metadata", {}).get("token_usage", {})
                if usage:
                    prompt_tokens = usage.get("prompt_tokens", 0)
                    completion_tokens = usage.get("completion_tokens", 0)
                    total_tokens = usage.get("total_tokens", prompt_tokens + completion_tokens)
                    llm_tokens_total.labels(model=model_name, type="prompt").inc(prompt_tokens)
                    llm_tokens_total.labels(model=model_name, type="completion").inc(completion_tokens)
                    llm_tokens_total.labels(model=model_name, type="total").inc(total_tokens)
            except Exception:
                pass
            return result
        except Exception as e:
            llm_duration_seconds.labels(model=model_name).observe(time.perf_counter() - start)
            llm_calls_total.labels(model=model_name, status="error").inc()
            raise

    try:
        # pydantic v2 models block arbitrary attribute assignment; fall back gracefully
        llm.ainvoke = monitored_ainvoke
    except (ValueError, AttributeError):
        pass
    return llm


def get_llm() -> ChatOpenAI:
    # 优先 LLM_* 系列（与 .env.production 一致），兼容 OPENAI_* 系列
    # timeout=180/max_retries=1：问卷推理可达 30-60s，重试叠加会拖垮 gunicorn worker
    llm = ChatOpenAI(
        model=os.getenv("LLM_MODEL", os.getenv("OPENAI_MODEL", "gpt-4o-mini")),
        temperature=0.7,
        **_common_kwargs(),
    )
    return _wrap_llm_ainvoke(llm, llm.model_name)


def get_vision_llm() -> ChatOpenAI:
    # 视觉模型（舌诊等图像分析），复用 LLM 网关与 key，仅切换模型名
    llm = ChatOpenAI(
        model=os.getenv("LLM_VISION_MODEL", "glm-5"),
        temperature=0.3,
        **_common_kwargs(),
    )
    return _wrap_llm_ainvoke(llm, llm.model_name)