CSPON-TSC/landsite_agent/config.py

from typing import Dict, Any
from pydantic import BaseModel
import os
from dotenv import load_dotenv

# 加载config.env文件
load_dotenv("config.env")

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.72:1025/v1",
        "model_name": "DeepSeek-R1-Distill-Qwen-32B",
    }
}

# 从环境变量获取默认模型类型，如果未设置则使用"openai"作为默认值
DEFAULT_MODEL_TYPE = os.getenv("DEFAULT_MODEL_TYPE", "openai")

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])