api_client.py

import os
import openai
from typing import Dict
from dotenv import load_dotenv
from models import AgentConfig

# 加载环境变量
load_dotenv()


class APIClient:
    """API客户端管理类"""
    
    @staticmethod
    def get_qwen_client(api_key: str):
        """获取Qwen API客户端"""
        # 使用OpenAI兼容接口
        client = openai.OpenAI(
            api_key=api_key,
            base_url="https://dashscope.aliyuncs.com/compatible-mode/v1"
        )
        return client

    @staticmethod
    def get_deepseek_client(api_key: str):
        """获取DeepSeek API客户端"""
        client = openai.OpenAI(
            api_key=api_key,
            base_url="https://api.deepseek.com"
        )
        return client


class ModelManager:
    """模型管理类"""
    
    def __init__(self):
        self.agents = self._initialize_agents()
    
    def _initialize_agents(self) -> Dict[str, AgentConfig]:
        """初始化智能体配置"""
        agents = {}
        
        # 初始化Qwen-Max (作为A讨论者)
        qwen_api_key = os.getenv("QWEN_API_KEY")
        if qwen_api_key:
            agents["A"] = AgentConfig("A", qwen_api_key, "qwen-max")
            # 初始化Qwen-Turbo (作为C裁判)
            agents["C"] = AgentConfig("C", qwen_api_key, "qwen-turbo")
        
        # 初始化DeepSeek (作为B讨论者)
        deepseek_api_key = os.getenv("DEEPSEEK_API_KEY")
        if deepseek_api_key:
            agents["B"] = AgentConfig("B", deepseek_api_key, "deepseek-chat")
        
        return agents
    
    def get_agent(self, name: str) -> AgentConfig:
        """获取指定名称的智能体配置"""
        if name not in self.agents:
            raise ValueError(f"未知的智能体: {name}")
        return self.agents[name]
    
    def call_model(self, agent_config: AgentConfig, prompt: str) -> str:
        """调用模型生成响应"""
        if "qwen" in agent_config.model:
            client = APIClient.get_qwen_client(agent_config.api_key)
        elif "deepseek" in agent_config.model:
            client = APIClient.get_deepseek_client(agent_config.api_key)
        else:
            raise ValueError(f"不支持的模型: {agent_config.model}")
        
        try:
            response = client.chat.completions.create(
                model=agent_config.model,
                messages=[
                    {"role": "user", "content": prompt}
                ],
                temperature=0.7,
                max_tokens=4096
            )
            
            return response.choices[0].message.content.strip()
            
        except Exception as e:
            print(f"调用模型 {agent_config.model} 时出错: {str(e)}")
            return f"错误: {str(e)}"