vent_agent_2.0/agents/model_factory.py

"""模型实例构建。

三个要点：
1. **思考级别烧进实例**——级别是构造参数，运行时改级别等于换模型，需要
   重建 agent（见 agents/agent_registry.py 的失效重建）。
2. **必须显式声明上下文窗口**——内置 SummarizationMiddleware 的
   ``trigger=("fraction", x)`` 依赖模型 ``profile["max_input_tokens"]``；
   不声明则框架在首次触发压缩时直接报错，而不是降级。
3. **不信任进程环境**——用 ``config`` 读 .env 的值，而不是让 SDK 自己去
   捡 ``OPENAI_API_KEY`` 等环境变量，避免被外部环境污染。
"""

from __future__ import annotations

import functools
from typing import Any

import config
from obs import get_logger

logger = get_logger(__name__)

# 已知模型的上下文窗口（token）。来源是各家的公开文档；命中不了就用兜底值。
KNOWN_CONTEXT_WINDOWS: dict[str, int] = {
    "deepseek-flash": 128_000,
    "deepseek-v4-flash": 128_000,
    "deepseek-v4-pro": 128_000,
    "deepseek-chat": 64_000,
    "deepseek-reasoner": 64_000,
    "gpt-4o": 128_000,
    "gpt-4o-mini": 128_000,
}

# 兜底窗口：按最保守的假设给，宁可早点压缩也不要超窗被服务端截断
FALLBACK_CONTEXT_WINDOW = 64_000

VALID_THINKING_LEVELS = ("off", "high", "highest")


def context_window(model_name: str | None = None) -> int:
    """解析模型的上下文窗口。

    ``CONTEXT_LIMIT`` 显式配置时优先——这是给「模型真实窗口比文档小」或
    「想留更多余量」的场景留的手动闸。
    """
    if config.CONTEXT_LIMIT > 0:
        return config.CONTEXT_LIMIT
    name = model_name or config.MODEL
    if name in KNOWN_CONTEXT_WINDOWS:
        return KNOWN_CONTEXT_WINDOWS[name]
    # 前缀匹配：deepseek-xxx-v5 这种没登记过的新版本，沿用同族窗口
    for known, window in KNOWN_CONTEXT_WINDOWS.items():
        family = known.split("-")[0]
        if name.startswith(family):
            return window
    logger.warning(
        "模型 %s 未登记上下文窗口，按兜底值 %d 处理（可用 CONTEXT_LIMIT 显式覆盖）",
        name, FALLBACK_CONTEXT_WINDOW,
    )
    return FALLBACK_CONTEXT_WINDOW


def normalize_thinking(level: str | None) -> str:
    """把思考级别归一到合法值；非法值回落 off 并告警。"""
    value = (level or config.THINKING_LEVEL or "off").strip().lower()
    if value not in VALID_THINKING_LEVELS:
        logger.warning("非法思考级别 %r，回落 off", level)
        return "off"
    return value


def _thinking_payload(level: str) -> dict[str, Any] | None:
    """思考级别 → 透传给 provider 的 extra_body。

    ``off`` 返回 None（**什么都不发**）：不同 provider 对「关闭思考」的
    表达方式不一致，发错字段会被服务端拒；不发才是真正中立的默认。
    """
    if level == "off":
        return None
    if level == "high":
        return {"thinking": {"type": "enabled"}, "reasoning_effort": "high"}
    return {"thinking": {"type": "enabled"}, "reasoning_effort": "highest"}


def build_model(
    model_name: str | None = None,
    thinking: str | None = None,
    timeout: int | None = None,
    context_limit: int | None = None,
) -> Any:
    """构造一个聊天模型实例。

    deepseek 系列优先用 ``langchain-deepseek``（它解析 ``reasoning_content``，
    官方 OpenAI 兼容层会把它丢掉）；其余走 ``init_chat_model`` 的 openai 通道。
    """
    name = model_name or config.MODEL
    level = normalize_thinking(thinking)
    window = context_limit or context_window(name)
    extra_body = _thinking_payload(level)

    common: dict[str, Any] = {
        "model": name,
        "api_key": config.OPENAI_API_KEY or "sk-placeholder",
        "timeout": timeout or config.OPENAI_TIMEOUT,
        "max_retries": 2,
        # 框架靠这个字段决定何时压缩上下文，必须显式给
        "profile": {"max_input_tokens": window},
    }
    if config.OPENAI_BASE_URL:
        common["base_url"] = config.OPENAI_BASE_URL
    if extra_body:
        common["extra_body"] = extra_body

    if not config.OPENAI_API_KEY:
        logger.warning("OPENAI_API_KEY 未配置，模型实例无法真正发起请求（仅可用于结构验证）")

    model = _instantiate(name, common)
    logger.info(
        "模型就绪 model=%s thinking=%s window=%d base_url=%s",
        name, level, window, config.OPENAI_BASE_URL or "(默认)",
    )
    return model


def _instantiate(name: str, common: dict[str, Any]) -> Any:
    """按模型名选择集成包。导入失败时逐级回落，并说明回落原因。"""
    if name.lower().startswith("deepseek"):
        try:
            from langchain_deepseek import ChatDeepSeek

            return ChatDeepSeek(**common)
        except ImportError:
            logger.warning("langchain-deepseek 未安装，回落到 ChatOpenAI（会丢失 reasoning_content）")

    try:
        from langchain_openai import ChatOpenAI

        return ChatOpenAI(**common)
    except ImportError:
        pass

    from langchain.chat_models import init_chat_model

    init_kwargs = {
        "api_key": common["api_key"],
        "timeout": common["timeout"],
        "max_retries": common["max_retries"],
    }
    if config.OPENAI_BASE_URL:
        init_kwargs["base_url"] = config.OPENAI_BASE_URL
    return init_chat_model(f"openai:{name}", **init_kwargs)


# ============================================================
# 进程级缓存
# ============================================================
# 没有「摘要模型」这一档：压缩用的 SummarizationMiddleware 由 deepagents 内部
# 构造，拿到的就是传进 create_deep_agent 的那个实例，没有注入另一个模型的入口。
# 所以思考级别同样作用于压缩调用；要单独给压缩配便宜模型，得先能替换掉那个
# 中间件（同名冲突见 AGENTS.md 不变式 5）。

@functools.lru_cache(maxsize=8)
def _cached_main(model_name: str, thinking: str) -> Any:
    return build_model(model_name, thinking)


def get_model(model_name: str | None = None, thinking: str | None = None) -> Any:
    """取主模型实例（按 model+thinking 缓存，同配置复用同一实例）。"""
    name = model_name or config.MODEL
    level = normalize_thinking(thinking)
    return _cached_main(name, level)


def clear_cache() -> None:
    """丢弃模型实例缓存（思考级别切换后调用）。"""
    _cached_main.cache_clear()


def describe() -> dict[str, Any]:
    """当前模型配置摘要，供 GET /api/model 展示。"""
    name = config.MODEL
    return {
        "model": name,
        "thinking": normalize_thinking(None),
        "thinking_levels": list(VALID_THINKING_LEVELS),
        "context_window": context_window(name),
        "context_limit_override": config.CONTEXT_LIMIT,
        "base_url": config.OPENAI_BASE_URL or "(默认)",
        "configured": bool(config.OPENAI_API_KEY),
    }