iov_data_analysis_agent/config/llm_config.py

# -*- coding: utf-8 -*-
"""
配置管理模块
"""

import os
from typing import Dict, Any
from dataclasses import dataclass, asdict


from dotenv import load_dotenv

load_dotenv()


@dataclass
class LLMConfig:
    """LLM配置"""

    provider: str = os.environ.get("LLM_PROVIDER", "openai")  # openai, gemini, etc.
    api_key: str = os.environ.get("OPENAI_API_KEY", "sk-c44i1hy64xgzwox6x08o4zug93frq6rgn84oqugf2pje1tg4")
    base_url: str = os.environ.get("OPENAI_BASE_URL", "https://api.xiaomimimo.com/v1")
    model: str = os.environ.get("OPENAI_MODEL", "mimo-v2-flash")
    temperature: float = 0.5
    max_tokens: int = 131072

    def __post_init__(self):
        """配置初始化后的处理"""
        if self.provider == "gemini":
            # 如果使用 Gemini，尝试从环境变量加载 Gemini 配置，或者使用默认的 Gemini 配置
            # 注意：如果 OPENAI_API_KEY 已设置且 GEMINI_API_KEY 未设置，可能会沿用 OpenAI 的 Key，
            # 但既然用户切换了 provider，通常会有配套的 Key。
            self.api_key = os.environ.get("GEMINI_API_KEY", "AIzaSyA9aVFjRJYJq82WEQUVlifE4fE7BnX6QiY")
            # Gemini 的 OpenAI 兼容接口地址
            self.base_url = os.environ.get("GEMINI_BASE_URL", "https://gemini.jeason.online")
            self.model = os.environ.get("GEMINI_MODEL", "gemini-2.5-flash")

    def to_dict(self) -> Dict[str, Any]:
        """转换为字典"""
        return asdict(self)

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "LLMConfig":
        """从字典创建配置"""
        return cls(**data)

    def validate(self) -> bool:
        """验证配置有效性"""
        if not self.api_key:
            raise ValueError("OPENAI_API_KEY is required")
        if not self.base_url:
            raise ValueError("OPENAI_BASE_URL is required")
        if not self.model:
            raise ValueError("OPENAI_MODEL is required")
        return True