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"
    },
    "azure": {
        "api_key": "your-azure-api-key",
        "api_base": "https://your-azure-endpoint.openai.azure.com",
        "model_name": "gpt-35-turbo",
        "temperature": 0,
        "max_tokens": 2000
    },
    "local": {
        "api_key": "your-local-api-key",
        "api_base": "http://localhost:8000/v1",
        "model_name": "local-model",
        "temperature": 0,
        "max_tokens": 2000
    }
}

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