ForcePilot/backend/package/yuxi/models/chat.py

import os
import traceback

from openai import AsyncOpenAI
from tenacity import before_sleep_log, retry, retry_if_exception_type, stop_after_attempt, wait_exponential

from yuxi import config
from yuxi.services.model_cache import is_v2_spec_format
from yuxi.utils import logger


def split_model_spec(model_spec, sep="/"):
    """
    将 provider/model 形式的字符串拆分为 (provider, model)
    """
    if not model_spec or not isinstance(model_spec, str):
        return "", ""
    if not sep:
        return model_spec, ""
    try:
        provider, model_name = model_spec.split(sep, 1)
        return provider, model_name
    except ValueError:
        return model_spec, ""


class OpenAIBase:
    def __init__(self, api_key, base_url, model_name, **kwargs):
        self.api_key = api_key
        self.base_url = base_url
        self.client = AsyncOpenAI(api_key=api_key, base_url=base_url)
        self.model_name = model_name
        self.info = kwargs

    @retry(
        stop=stop_after_attempt(3),
        wait=wait_exponential(multiplier=1, min=1, max=10),
        retry=retry_if_exception_type((Exception,)),
        before_sleep=before_sleep_log(logger, log_level="WARNING"),
        reraise=True,
    )
    async def call(self, message, stream=False):
        if isinstance(message, str):
            messages = [{"role": "user", "content": message}]
        else:
            messages = message

        try:
            if stream:
                response = self._stream_response(messages)
            else:
                response = await self._get_response(messages)

        except Exception as e:
            err = (
                f"Error streaming response: {e}, URL: {self.base_url}, "
                f"API Key: {self.api_key[:5]}***, Model: {self.model_name}"
            )
            logger.error(err)
            raise Exception(err)

        return response

    async def _stream_response(self, messages):
        response = await self.client.chat.completions.create(
            model=self.model_name,
            messages=messages,
            stream=True,
        )
        async for chunk in response:
            if len(chunk.choices) > 0:
                yield chunk.choices[0].delta

    async def _get_response(self, messages):
        response = await self.client.chat.completions.create(
            model=self.model_name,
            messages=messages,
            stream=False,
        )
        return response.choices[0].message

    async def get_models(self):
        try:
            return await self.client.models.list(extra_query={"type": "text"})
        except Exception as e:
            logger.error(f"Error getting models: {e}")
            return []


class OpenModel(OpenAIBase):
    def __init__(self, model_name=None):
        model_name = model_name or "gpt-4o-mini"
        api_key = os.getenv("OPENAI_API_KEY")
        base_url = os.getenv("OPENAI_API_BASE")
        super().__init__(api_key=api_key, base_url=base_url, model_name=model_name)


class GeneralResponse:
    def __init__(self, content):
        self.content = content
        self.is_full = False


def select_model_v2(spec: str) -> OpenAIBase:
    """根据 v2 spec（provider_id:model_id）选择聊天模型。

    v2 spec 格式使用冒号分隔，如: siliconflow-cn:deepseek-ai/DeepSeek-V4-Flash
    数据来源为数据库中的 model_providers 表，通过全局缓存访问。
    """
    from yuxi.services.model_cache import model_cache

    info = model_cache.get_model_info(spec)
    if not info:
        raise ValueError(f"Unknown v2 model spec: {spec}")

    if info.model_type != "chat":
        raise ValueError(f"Model {spec} is not a chat model (type={info.model_type})")

    logger.info(f"Selecting v2 model: {spec} (provider_type={info.provider_type})")

    return OpenAIBase(
        api_key=info.api_key,
        base_url=info.base_url,
        model_name=info.model_id,
    )


def select_model(model_provider=None, model_name=None, model_spec=None):
    """根据模型提供者选择模型"""
    if model_spec and is_v2_spec_format(model_spec):
        from yuxi.services.model_cache import model_cache

        if model_cache.is_v2_spec(model_spec):
            return select_model_v2(model_spec)

        available = model_cache.get_all_specs("chat")
        available_ids = [s.spec for s in available[:10]]
        raise ValueError(f"未找到 V2 模型: '{model_spec}'。可用聊天模型 ({len(available)}): {available_ids}")

    logger.warning(
        f"旧版本的模型选择逻辑已废弃，建议尽快迁移至新的模型配置；"
        f"当前模型选择参数: provider={model_provider}, model_name={model_name}, spec={model_spec}"
    )

    if model_spec:
        spec_provider, spec_model_name = split_model_spec(model_spec)
        model_provider = model_provider or spec_provider
        model_name = model_name or spec_model_name

    if model_provider is None or not model_name:
        default_provider, default_model = split_model_spec(getattr(config, "default_model", ""))
        model_provider = model_provider or default_provider
        model_name = model_name or default_model

    assert model_provider, "Model provider not specified"

    model_info = config.model_names.get(model_provider)
    if not model_info:
        raise ValueError(f"Unknown model provider: {model_provider}")

    model_name = model_name or model_info.default

    if not model_name:
        raise ValueError(f"Model name not specified for provider {model_provider}")

    logger.info(f"Selecting model from `{model_provider}` with `{model_name}`")

    if model_provider == "openai":
        return OpenModel(model_name)

    # 其他模型，默认使用OpenAIBase
    try:
        model = OpenAIBase(
            api_key=os.environ.get(model_info.env, model_info.env),
            base_url=model_info.base_url,
            model_name=model_name,
        )
        return model
    except Exception as e:
        raise ValueError(f"Model provider {model_provider} load failed, {e} \n {traceback.format_exc()}")


async def test_chat_model_status_by_spec(spec: str) -> dict:
    """根据 full spec 测试聊天模型状态（自动识别 V1/V2）。

    V1 spec 格式: provider/model_name（斜杠分隔）
    V2 spec 格式: provider_id:model_id（冒号分隔）
    """
    try:
        logger.debug(f"Testing model status by spec: {spec}")
        model = select_model(model_spec=spec)

        test_messages = [{"role": "user", "content": "Say 1"}]
        response = await model.call(test_messages, stream=False)

        if response and response.content:
            return {"spec": spec, "status": "available", "message": "连接正常"}
        else:
            return {"spec": spec, "status": "unavailable", "message": "响应无效"}

    except Exception as e:
        logger.error(f"测试模型状态失败 {spec}: {e}")
        return {"spec": spec, "status": "error", "message": str(e)}


if __name__ == "__main__":
    pass