CSPON-TSC/landsite_agent/config.py

from typing import Dict, Any
from pydantic import BaseModel

class ModelConfig(BaseModel):
    """AI模型配置类"""
    api_key: str
    api_base: str
    model_name: str
    temperature: float = 0
    max_tokens: int = 2000


# AI模型配置
model_list: Dict[str, Dict[str, Any]] = {
    "openai": {
        "api_key": "none",
        "api_base": "http://ac.zjugis.com:8511/v1",
        "model_name": "qwen2.5-instruct"
    },
    "zjstai": {
        "api_key": "none",
        "api_base": "http://172.27.27.20:20331/v1",
        "model_name": "DeepSeek-R1-Distill-Qwen-32B",
    }
}

# 新增全局默认模型类型开关
DEFAULT_MODEL_TYPE = "openai"
# # 线上服务器配置
# DEFAULT_MODEL_TYPE = "zjstai"

def get_model_config(model_type: str = "openai") -> ModelConfig:
    """
    获取模型配置
    :param model_type: 模型类型，默认为openai
    :return: 模型配置对象
    """
    if model_type not in model_list:
        raise ValueError(f"Unsupported model type: {model_type}")
    
    return ModelConfig(**model_list[model_type])