auto-coder 1.0.0__py3-none-any.whl → 2.0.1__py3-none-any.whl

autocoder/common/llms/factory.py

@@ -0,0 +1,142 @@
+import byzerllm
+from typing import Union, List, Optional
+
+from .schema import LLMModel
+
+
+class LLMFactory:
+    """LLM 实例化工厂"""
+    
+    @staticmethod
+    def create_llm(model: LLMModel, product_mode: str) -> Union[byzerllm.ByzerLLM, byzerllm.SimpleByzerLLM]:
+        """
+        根据模型配置和产品模式创建 LLM 实例
+        
+        Args:
+            model: LLM 模型配置
+            product_mode: 产品模式 ("pro" 或 "lite")
+            
+        Returns:
+            配置好的 byzerllm 客户端实例
+            
+        Raises:
+            ValueError: 当模型未配置 API 密钥时
+        """
+        # 检查是否有 API 密钥
+        api_key = model.get_api_key()
+        if not api_key:
+            raise ValueError(f"Model {model.name} has no API key configured")
+        
+        if product_mode == "pro":
+            # Pro 模式：使用 ByzerLLM.from_default_model
+            return byzerllm.ByzerLLM.from_default_model(model.name)
+        
+        elif product_mode == "lite":
+            # Lite 模式：使用 SimpleByzerLLM
+            target_llm = byzerllm.SimpleByzerLLM(default_model_name=model.name)
+            
+            # 部署配置
+            infer_params = {
+                "saas.base_url": model.base_url,
+                "saas.api_key": api_key,
+                "saas.model": model.model_name,
+                "saas.is_reasoning": model.is_reasoning,
+                "saas.max_output_tokens": model.max_output_tokens
+            }
+            
+            target_llm.deploy(
+                model_path="",
+                pretrained_model_type=model.model_type,
+                udf_name=model.name,
+                infer_params=infer_params
+            )
+            
+            return target_llm
+        
+        else:
+            raise ValueError(f"Unsupported product mode: {product_mode}")
+    
+    @staticmethod
+    def create_llm_from_names(model_names: str, models: List[LLMModel], product_mode: str) -> Optional[Union[byzerllm.ByzerLLM, byzerllm.SimpleByzerLLM]]:
+        """
+        根据模型名称列表创建 LLM 实例（返回第一个可用的）
+        
+        Args:
+            model_names: 逗号分隔的模型名称列表
+            models: 可用的模型列表
+            product_mode: 产品模式
+            
+        Returns:
+            第一个成功创建的 LLM 实例，如果都失败则返回 None
+            
+        Raises:
+            ValueError: 当所有模型都无法创建时，包含详细的错误信息
+        """
+        # 处理模型名称列表
+        if "," in model_names:
+            names = [name.strip() for name in model_names.split(",")]
+        else:
+            names = [model_names.strip()]
+        
+        # 收集错误信息
+        errors = []
+        
+        # 查找并创建第一个可用的模型
+        for name in names:
+            # 在模型列表中查找
+            model = next((m for m in models if m.name == name), None)
+            if not model:
+                errors.append(f"Model '{name}' not found")
+                continue
+                
+            try:
+                return LLMFactory.create_llm(model, product_mode)
+            except ValueError as e:
+                errors.append(f"Model '{name}': {str(e)}")
+                continue
+        
+        # 如果所有模型都失败了，抛出包含所有错误信息的异常
+        if errors:
+            error_msg = f"Failed to create LLM instance for models: {model_names}\n" + "\n".join(f"  - {error}" for error in errors)
+            raise ValueError(error_msg)
+        
+        # 如果没有找到任何模型（理论上不应该到这里）
+        raise ValueError(f"No available models found: {model_names}")
+    
+    @staticmethod
+    def get_llm_names(llm: Union[byzerllm.ByzerLLM, byzerllm.SimpleByzerLLM, str, List], 
+                      target_model_type: Optional[str] = None) -> List[str]:
+        """
+        获取 LLM 的模型名称列表
+        
+        Args:
+            llm: LLM 实例或模型名称
+            target_model_type: 目标模型类型（可选）
+            
+        Returns:
+            模型名称列表
+        """
+        if target_model_type is None:
+            if isinstance(llm, list):
+                return [_llm.default_model_name for _llm in llm if hasattr(_llm, 'default_model_name') and _llm.default_model_name]
+            elif isinstance(llm, str):
+                return llm.split(",") if llm else []
+            return [llm.default_model_name] if hasattr(llm, 'default_model_name') and llm.default_model_name else []
+        
+        # 处理有 target_model_type 的情况
+        if not isinstance(llm, (str, list)) and hasattr(llm, 'get_sub_client'):
+            llms = llm.get_sub_client(target_model_type)
+            
+            if llms is None:
+                # 回退到默认处理
+                if isinstance(llm, list):
+                    return [_llm.default_model_name for _llm in llm if hasattr(_llm, 'default_model_name') and _llm.default_model_name]
+                return [llm.default_model_name] if hasattr(llm, 'default_model_name') and llm.default_model_name else []
+            elif isinstance(llms, list):
+                return [_llm.default_model_name for _llm in llms if hasattr(_llm, 'default_model_name') and _llm.default_model_name]
+            elif isinstance(llms, str) and llms:
+                return llms.split(",")
+            else:
+                return [llms.default_model_name] if hasattr(llms, 'default_model_name') and llms.default_model_name else []
+        
+        return [] 

autocoder/common/llms/manager.py

@@ -0,0 +1,264 @@
+import byzerllm
+from typing import Union, List, Optional, Dict
+
+from .schema import LLMModel
+from .registry import ModelRegistry
+from .factory import LLMFactory
+from .pricing import PricingManager
+
+
+class LLMManager:
+    """大语言模型管理器，统一管理所有 LLM 的配置、实例化和操作"""
+    
+    def __init__(self, models_json_path: Optional[str] = None):
+        """
+        初始化 LLM 管理器
+        
+        Args:
+            models_json_path: 模型配置文件路径，默认为 ~/.auto-coder/keys/models.json
+        """
+        self.registry = ModelRegistry(models_json_path)
+        self.pricing = PricingManager(self.registry)
+        # 立即加载模型以初始化缓存
+        self.registry.load()
+    
+    def get_model(self, name: str) -> Optional[LLMModel]:
+        """
+        获取模型配置信息
+        
+        Args:
+            name: 模型名称
+            
+        Returns:
+            LLMModel 对象，如果不存在则返回 None
+        """
+        return self.registry.get(name)
+    
+    def get_all_models(self) -> Dict[str, LLMModel]:
+        """获取所有模型配置"""
+        return self.registry.get_all()
+    
+    def get_model_info(self, name: str, product_mode: str) -> Optional[Dict]:
+        """
+        获取模型信息（兼容旧版本接口）
+        
+        Args:
+            name: 模型名称
+            product_mode: 产品模式（pro 或 lite）
+            
+        Returns:
+            模型信息字典，如果不存在则返回 None
+        """
+        if product_mode == "pro":
+            # Pro 模式下直接返回 None（与原实现保持一致）
+            return None
+        
+        model = self.get_model(name)
+        if not model:
+            return None
+        
+        # 转换为字典格式，保持兼容性
+        return {
+            "name": model.name,
+            "description": model.description,
+            "model_name": model.model_name,
+            "model_type": model.model_type,
+            "base_url": model.base_url,
+            "provider": model.provider,
+            "api_key": model.get_api_key(),
+            "is_reasoning": model.is_reasoning,
+            "input_price": model.input_price,
+            "output_price": model.output_price,
+            "max_output_tokens": model.max_output_tokens,
+            "context_window": model.context_window
+        }
+    
+    def get_single_llm(self, model_names: str, product_mode: str) -> Optional[Union[byzerllm.ByzerLLM, byzerllm.SimpleByzerLLM]]:
+        """
+        根据模型名称和产品模式获取 LLM 实例
+        
+        Args:
+            model_names: 逗号分隔的模型名称列表
+            product_mode: 产品模式（pro 或 lite）
+            
+        Returns:
+            配置好的 LLM 实例
+            
+        Raises:
+            ValueError: 当所有指定的模型都无法创建时，包含详细的错误信息
+        """
+        models = list(self.registry.get_all().values())
+        return LLMFactory.create_llm_from_names(model_names, models, product_mode)
+    
+    def get_llm_names(self, llm: Union[byzerllm.ByzerLLM, byzerllm.SimpleByzerLLM, str, List],
+                      target_model_type: Optional[str] = None) -> List[str]:
+        """
+        获取 LLM 的模型名称列表
+        
+        Args:
+            llm: LLM 实例或模型名称
+            target_model_type: 目标模型类型（可选）
+            
+        Returns:
+            模型名称列表
+        """
+        return LLMFactory.get_llm_names(llm, target_model_type)
+    
+    def add_models(self, models: List[Union[LLMModel, Dict]]) -> None:
+        """
+        批量添加模型
+        
+        Args:
+            models: 模型列表，可以是 LLMModel 对象或字典
+        """
+        for model_data in models:
+            if isinstance(model_data, dict):
+                # 如果是字典，转换为 LLMModel，设置 base_keys_dir
+                model_data_with_keys_dir = model_data.copy()
+                model_data_with_keys_dir["base_keys_dir"] = str(self.registry.models_json_path.parent)
+                model = LLMModel(**model_data_with_keys_dir)
+            else:
+                model = model_data
+                # 确保已有的 LLMModel 对象也设置了 base_keys_dir
+                if not model.base_keys_dir:
+                    model.base_keys_dir = str(self.registry.models_json_path.parent)
+            
+            # 保存 api_key 的副本（如果有的话）
+            api_key_to_save = model.api_key
+            
+            # 如果有 api_key，设置 api_key_path（如果没有的话）
+            if api_key_to_save and not model.api_key_path:
+                model.api_key_path = model.name.replace("/", "_")
+            
+            # 清除内存中的 api_key，以便正确保存到文件
+            model.api_key = None
+            
+            # 首先添加模型配置
+            self.registry.add_or_update(model)
+            
+            # 如果有 api_key，再保存密钥
+            if api_key_to_save and model.api_key_path:
+                # 使用文件锁保存密钥到文件
+                api_key_file = self.registry.models_json_path.parent / model.api_key_path
+                self.registry._ensure_dir()
+                
+                # 导入文件锁上下文管理器
+                from filelock import FileLock
+                
+                with FileLock(self.registry._get_lock_path(api_key_file), timeout=5):
+                    with open(api_key_file, 'w', encoding='utf-8') as f:
+                        f.write(api_key_to_save.strip())
+    
+    def update_model_with_api_key(self, name: str, api_key: str) -> bool:
+        """
+        更新模型的 API 密钥
+        
+        Args:
+            name: 模型名称
+            api_key: API 密钥
+            
+        Returns:
+            是否成功更新
+        """
+        return self.registry.save_api_key(name, api_key)
+    
+    def update_model(self, name: str, model_data: Dict) -> Optional[LLMModel]:
+        """
+        更新模型信息
+        
+        Args:
+            name: 要更新的模型名称
+            model_data: 包含新信息的字典
+            
+        Returns:
+            更新后的模型，如果未找到则返回 None
+        """
+        model = self.registry.get(name)
+        if not model:
+            return None
+        
+        # 更新字段
+        if "description" in model_data:
+            model.description = model_data["description"]
+        if "model_name" in model_data:
+            model.model_name = model_data["model_name"]
+        if "model_type" in model_data:
+            model.model_type = model_data["model_type"]
+        if "base_url" in model_data:
+            model.base_url = model_data["base_url"]
+        if "provider" in model_data:
+            model.provider = model_data["provider"]
+        if "is_reasoning" in model_data:
+            model.is_reasoning = model_data["is_reasoning"]
+        if "input_price" in model_data:
+            model.input_price = float(model_data["input_price"])
+        if "output_price" in model_data:
+            model.output_price = float(model_data["output_price"])
+        if "max_output_tokens" in model_data:
+            model.max_output_tokens = int(model_data["max_output_tokens"])
+        if "context_window" in model_data:
+            model.context_window = int(model_data["context_window"])
+        
+        # 保存更新
+        self.registry.add_or_update(model)
+        
+        # 如果提供了 API 密钥，单独处理
+        if "api_key" in model_data and model_data["api_key"]:
+            self.update_model_with_api_key(name, model_data["api_key"])
+        
+        return model
+    
+    def has_key(self, name: str) -> bool:
+        """
+        检查模型是否已配置密钥
+        
+        Args:
+            name: 模型名称
+            
+        Returns:
+            是否已配置密钥
+        """
+        model = self.get_model(name)
+        return model.has_api_key if model else False
+    
+    def check_model_exists(self, name: str) -> bool:
+        """
+        检查模型是否存在
+        
+        Args:
+            name: 模型名称
+            
+        Returns:
+            是否存在
+        """
+        return self.get_model(name) is not None
+    
+    # 价格管理相关方法
+    def update_input_price(self, name: str, price: float) -> bool:
+        """更新模型输入价格"""
+        return self.pricing.update_input_price(name, price)
+    
+    def update_output_price(self, name: str, price: float) -> bool:
+        """更新模型输出价格"""
+        return self.pricing.update_output_price(name, price)
+    
+    def update_price(self, name: str, input_price: Optional[float] = None,
+                    output_price: Optional[float] = None) -> bool:
+        """同时更新输入和输出价格"""
+        return self.pricing.update_prices(name, input_price, output_price)
+    
+    def get_cost_estimate(self, name: str, input_tokens: int, output_tokens: int) -> Optional[float]:
+        """估算使用成本"""
+        return self.pricing.get_cost_estimate(name, input_tokens, output_tokens)
+    
+    def remove_model(self, name: str) -> bool:
+        """
+        删除模型
+        
+        Args:
+            name: 模型名称
+            
+        Returns:
+            是否成功删除
+        """
+        return self.registry.remove_model(name) 

autocoder/common/llms/pricing.py

@@ -0,0 +1,121 @@
+from typing import Optional
+
+from .schema import LLMModel
+from .registry import ModelRegistry
+
+
+class PricingManager:
+    """模型价格管理器"""
+    
+    def __init__(self, registry: ModelRegistry):
+        self.registry = registry
+    
+    def update_input_price(self, name: str, price: float) -> bool:
+        """
+        更新模型输入价格
+        
+        Args:
+            name: 模型名称
+            price: 新的输入价格（美元/百万tokens），必须大于等于0
+            
+        Returns:
+            bool: 是否成功更新
+            
+        Raises:
+            ValueError: 如果价格为负数
+        """
+        if price < 0:
+            raise ValueError("价格不能为负数")
+        
+        model = self.registry.get(name)
+        if not model:
+            return False
+        
+        model.input_price = float(price)
+        self.registry.add_or_update(model)
+        return True
+    
+    def update_output_price(self, name: str, price: float) -> bool:
+        """
+        更新模型输出价格
+        
+        Args:
+            name: 模型名称
+            price: 新的输出价格（美元/百万tokens），必须大于等于0
+            
+        Returns:
+            bool: 是否成功更新
+            
+        Raises:
+            ValueError: 如果价格为负数
+        """
+        if price < 0:
+            raise ValueError("价格不能为负数")
+        
+        model = self.registry.get(name)
+        if not model:
+            return False
+        
+        model.output_price = float(price)
+        self.registry.add_or_update(model)
+        return True
+    
+    def update_prices(self, name: str, input_price: Optional[float] = None, 
+                     output_price: Optional[float] = None) -> bool:
+        """
+        同时更新输入和输出价格
+        
+        Args:
+            name: 模型名称
+            input_price: 新的输入价格（可选）
+            output_price: 新的输出价格（可选）
+            
+        Returns:
+            bool: 是否成功更新
+            
+        Raises:
+            ValueError: 如果任一价格为负数
+        """
+        if input_price is not None and input_price < 0:
+            raise ValueError("输入价格不能为负数")
+        if output_price is not None and output_price < 0:
+            raise ValueError("输出价格不能为负数")
+        
+        model = self.registry.get(name)
+        if not model:
+            return False
+        
+        updated = False
+        if input_price is not None:
+            model.input_price = float(input_price)
+            updated = True
+        if output_price is not None:
+            model.output_price = float(output_price)
+            updated = True
+        
+        if updated:
+            self.registry.add_or_update(model)
+        
+        return updated
+    
+    def get_cost_estimate(self, name: str, input_tokens: int, output_tokens: int) -> Optional[float]:
+        """
+        估算使用成本
+        
+        Args:
+            name: 模型名称
+            input_tokens: 输入 token 数量
+            output_tokens: 输出 token 数量
+            
+        Returns:
+            估算的总成本（美元），如果模型不存在则返回 None
+        """
+        model = self.registry.get(name)
+        if not model:
+            return None
+        
+        # 计算成本：价格单位是 美元/百万tokens
+        input_cost = (input_tokens / 1_000_000) * model.input_price
+        output_cost = (output_tokens / 1_000_000) * model.output_price
+        
+        return input_cost + output_cost