api_model.py

#!/usr/bin/env python3
"""
API模型接口 - 支持通过API调用商用模型进行代码修复
"""

import json
import logging
import requests
import time
from typing import Dict, Any, Optional
from model_integration import BaseCodeFixModel

logger = logging.getLogger(__name__)

class APICodeFixModel(BaseCodeFixModel):
    """通过API调用的代码修复模型"""
    
    def __init__(self, api_key: str, model_name: str = "gpt-3.5-turbo", 
                 base_url: str = "https://api.openai.com/v1", **kwargs):
        """
        初始化API模型
        
        Args:
            api_key: API密钥
            model_name: 模型名称 (如 gpt-3.5-turbo, gpt-4, claude-3-sonnet 等)
            base_url: API基础URL
            **kwargs: 其他配置参数
        """
        self.api_key = api_key
        self.model_name = model_name
        self.base_url = base_url.rstrip('/')
        self.config = kwargs
        
        # API配置
        self.max_tokens = kwargs.get('max_tokens', 1024)
        self.temperature = kwargs.get('temperature', 0.1)
        self.timeout = kwargs.get('timeout', 30)
        self.max_retries = kwargs.get('max_retries', 3)
        
        # 请求头
        self.headers = {
            'Content-Type': 'application/json',
            'Authorization': f'Bearer {self.api_key}'
        }
        
        # 支持的API提供商配置
        self.provider_configs = {
            'openai': {
                'chat_endpoint': '/chat/completions',
                'model_field': 'model',
                'messages_field': 'messages'
            },
            'anthropic': {
                'chat_endpoint': '/messages',
                'model_field': 'model', 
                'messages_field': 'messages'
            },
            'custom': {
                'chat_endpoint': '',  # 你的API端点已经包含在base_url中
                'model_field': 'model',
                'messages_field': 'input'  # 你的API使用input字段
            }
        }
        
        # 自动检测API提供商
        self.provider = self._detect_provider()
        
    def _detect_provider(self) -> str:
        """根据base_url自动检测API提供商"""
        if 'openai.com' in self.base_url:
            return 'openai'
        elif 'anthropic.com' in self.base_url:
            return 'anthropic'
        else:
            return 'custom'  # 默认使用OpenAI兼容格式
    
    def load_model(self) -> bool:
        """测试API连接"""
        try:
            logger.info(f"测试API连接: {self.base_url}")
            logger.info(f"使用模型: {self.model_name}")
            logger.info(f"API提供商: {self.provider}")
            
            # 发送一个简单的测试请求
            test_response = self._make_api_request("测试连接", "请回复'连接成功'")
            
            if test_response and "成功" in test_response:
                logger.info("✅ API连接测试成功")
                return True
            else:
                logger.info("✅ API连接正常（测试响应已收到）")
                return True
                
        except Exception as e:
            logger.error(f"❌ API连接失败: {e}")
            return False
    
    def fix_code(self, buggy_code: str, instruction: str, bug_type: str) -> str:
        """使用API修复代码"""
        try:
            prompt = self._build_prompt(buggy_code, instruction, bug_type)
            
            logger.info(f"正在修复 {bug_type} 类型的bug...")
            
            # 调用API
            response = self._make_api_request(prompt, "代码修复")
            
            if response:
                # 提取修复后的代码
                fixed_code = self._extract_code_from_response(response, buggy_code)
                return fixed_code
            else:
                logger.error("API调用失败，返回原代码")
                return buggy_code
                
        except Exception as e:
            logger.error(f"❌ 代码修复失败: {e}")
            return buggy_code
    
    def _make_api_request(self, prompt: str, task_type: str = "代码修复") -> Optional[str]:
        """发送API请求"""
        config = self.provider_configs[self.provider]
        url = f"{self.base_url}{config['chat_endpoint']}"
        
        # 构建请求数据
        if self.provider == 'anthropic':
            data = {
                config['model_field']: self.model_name,
                config['messages_field']: [
                    {"role": "user", "content": prompt}
                ],
                "max_tokens": self.max_tokens,
                "temperature": self.temperature
            }
            # Anthropic需要特殊的请求头
            headers = self.headers.copy()
            headers['anthropic-version'] = '2023-06-01'
        elif self.provider == 'custom':
            # 你的自定义API格式
            data = {
                config['model_field']: self.model_name,
                config['messages_field']: prompt,  # 直接使用prompt作为input
                "instructions": "You are an expert at Python programming and code debugging. Please fix the bug in the given code and return only the corrected code.",
                "max_tokens": self.max_tokens,
                "temperature": self.temperature
            }
            headers = self.headers
        else:
            # OpenAI格式（也适用于大多数兼容API）
            data = {
                config['model_field']: self.model_name,
                config['messages_field']: [
                    {"role": "system", "content": "你是一个专业的代码修复助手。请仔细分析代码中的bug并提供准确的修复方案。"},
                    {"role": "user", "content": prompt}
                ],
                "max_tokens": self.max_tokens,
                "temperature": self.temperature
            }
            headers = self.headers
        
        # 重试机制
        for attempt in range(self.max_retries):
            try:
                logger.debug(f"发送API请求 (尝试 {attempt + 1}/{self.max_retries})")
                
                response = requests.post(
                    url, 
                    headers=headers,
                    json=data,
                    timeout=self.timeout
                )
                
                if response.status_code == 200:
                    result = response.json()
                    
                    # 提取响应内容
                    if self.provider == 'anthropic':
                        if 'content' in result and result['content']:
                            return result['content'][0]['text']
                    elif self.provider == 'custom':
                        # 你的自定义API响应格式
                        try:
                            content = result['output'][0]['content'][0]['text']
                            return content
                        except (KeyError, IndexError):
                            # 备用提取方式
                            content = result.get('content', '')
                            if content:
                                return content
                    else:
                        # OpenAI格式
                        if 'choices' in result and result['choices']:
                            return result['choices'][0]['message']['content']
                    
                    logger.error(f"API响应格式异常: {result}")
                    return None
                    
                else:
                    logger.error(f"API请求失败 (状态码: {response.status_code}): {response.text}")
                    
                    # 如果是认证错误，不重试
                    if response.status_code in [401, 403]:
                        break
                        
            except requests.exceptions.Timeout:
                logger.warning(f"API请求超时 (尝试 {attempt + 1}/{self.max_retries})")
            except requests.exceptions.RequestException as e:
                logger.warning(f"API请求异常 (尝试 {attempt + 1}/{self.max_retries}): {e}")
            except Exception as e:
                logger.error(f"未知错误 (尝试 {attempt + 1}/{self.max_retries}): {e}")
            
            # 重试前等待
            if attempt < self.max_retries - 1:
                time.sleep(2 ** attempt)  # 指数退避
        
        return None
    
    def _build_prompt(self, buggy_code: str, instruction: str, bug_type: str) -> str:
        """构建修复代码的提示词"""
        prompt = f"""请修复以下代码中的bug。

**任务描述**: {instruction}

**Bug类型**: {bug_type}

**有问题的代码**:
```python
{buggy_code}
```

**要求**:
1. 仔细分析代码中的bug
2. 提供修复后的完整代码
3. 只返回修复后的代码，不要包含额外的解释
4. 保持原有的代码结构和风格
5. 确保修复后的代码能正确执行

**修复后的代码**:
```python"""
        
        return prompt
    
    def _extract_code_from_response(self, response: str, original_code: str) -> str:
        """从API响应中提取代码 - 专门针对HumanEvalFix格式"""
        try:
            extracted_code = None
            
            # 尝试提取代码块
            if "```python" in response:
                # 提取python代码块
                start = response.find("```python") + len("```python")
                end = response.find("```", start)
                if end != -1:
                    code = response[start:end].strip()
                    if code:
                        extracted_code = code
            
            elif "```" in response:
                # 提取普通代码块
                parts = response.split("```")
                for i, part in enumerate(parts):
                    if i % 2 == 1:  # 奇数索引是代码块
                        code = part.strip()
                        # 移除可能的语言标识符
                        if code.startswith(('python\n', 'py\n')):
                            code = '\n'.join(code.split('\n')[1:])
                        if code:
                            extracted_code = code
                            break
            
            # 如果提取到了代码，进行后处理
            if extracted_code:
                # 处理完整函数定义的情况
                if extracted_code.startswith(('def ', 'from ', 'import ')):
                    lines = extracted_code.split('\n')
                    function_body_lines = []
                    in_function = False
                    
                    for line in lines:
                        # 跳过import语句
                        if line.strip().startswith(('from ', 'import ')):
                            continue
                        
                        # 找到函数定义
                        if line.strip().startswith('def '):
                            in_function = True
                            continue
                        
                        # 收集函数体
                        if in_function:
                            # 保持原有的缩进格式，只移除函数定义级别的缩进
                            if line.startswith('    '):
                                # 保留函数体内的缩进结构
                                function_body_lines.append('    ' + line[4:])
                            elif line.strip() == '':
                                function_body_lines.append('')
                            else:
                                # 如果遇到不是函数体的行，停止
                                break
                    
                    if function_body_lines:
                        # 移除末尾的空行
                        while function_body_lines and not function_body_lines[-1].strip():
                            function_body_lines.pop()
                        
                        if function_body_lines:
                            return '\n'.join(function_body_lines)
                
                # 如果不是完整函数定义，直接返回
                return extracted_code
            
            # 如果没有代码块标记，尝试直接提取
            lines = response.strip().split('\n')
            code_lines = []
            in_code = False
            
            for line in lines:
                # 跳过明显的解释性文本
                if any(keyword in line.lower() for keyword in ['fix', 'bug', 'error', 'problem', 'here is', 'the issue']):
                    if not in_code:
                        continue
                
                # 检查是否是代码行
                stripped = line.strip()
                if stripped and (
                    stripped.startswith(('def ', 'class ', 'import ', 'from ', 'if ', 'for ', 'while ', 'try:', 'except', 'return')) or
                    '=' in stripped or
                    stripped.endswith(':') or
                    in_code
                ):
                    in_code = True
                    code_lines.append(line)
                elif in_code and not stripped:
                    code_lines.append(line)  # 保留空行
                elif in_code and stripped:
                    # 可能是代码的一部分
                    code_lines.append(line)
            
            if code_lines:
                extracted_code = '\n'.join(code_lines).strip()
                if extracted_code and len(extracted_code) > 10:  # 基本长度检查
                    return extracted_code
            
            # 如果都失败了，返回原代码
            logger.warning("无法从API响应中提取有效代码，返回原代码")
            return original_code
            
        except Exception as e:
            logger.error(f"提取代码时出错: {e}")
            return original_code
    
    def get_model_info(self) -> Dict[str, Any]:
        """获取模型信息"""
        return {
            "model_name": f"API-{self.model_name}",
            "provider": self.provider,
            "base_url": self.base_url,
            "config": {
                "max_tokens": self.max_tokens,
                "temperature": self.temperature,
                "timeout": self.timeout
            }
        }

# 便捷的预设配置
class OpenAIModel(APICodeFixModel):
    """OpenAI API模型"""
    def __init__(self, api_key: str, model_name: str = "gpt-3.5-turbo", **kwargs):
        super().__init__(
            api_key=api_key,
            model_name=model_name,
            base_url="https://api.openai.com/v1",
            **kwargs
        )

class ClaudeModel(APICodeFixModel):
    """Anthropic Claude API模型"""
    def __init__(self, api_key: str, model_name: str = "claude-3-sonnet-20240229", **kwargs):
        super().__init__(
            api_key=api_key,
            model_name=model_name,
            base_url="https://api.anthropic.com/v1",
            **kwargs
        )

class CustomAPIModel(APICodeFixModel):
    """自定义API模型（使用你的API格式）"""
    def __init__(self, api_key: str = "sk-8t4sGAakvPVKfFLn9801056499284a66B31aC07b1f9907F3", 
                 model_name: str = "gpt-4o-mini", **kwargs):
        super().__init__(
            api_key=api_key,
            model_name=model_name,
            base_url="https://api.apiyi.com/v1/responses",
            **kwargs
        )