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

autocoder/common/llms/tests/test_registry.py

@@ -0,0 +1,364 @@
+import unittest
+import tempfile
+import json
+from pathlib import Path
+
+from ..registry import ModelRegistry
+from ..schema import LLMModel
+
+
+class TestModelRegistry(unittest.TestCase):
+    """ModelRegistry 单元测试"""
+    
+    def setUp(self):
+        """设置测试环境"""
+        self.temp_dir = tempfile.mkdtemp()
+        self.models_json = Path(self.temp_dir) / "models.json"
+        self.registry = ModelRegistry(str(self.models_json))
+    
+    def tearDown(self):
+        """清理测试环境"""
+        import shutil
+        shutil.rmtree(self.temp_dir, ignore_errors=True)
+    
+    def test_load_default_models(self):
+        """测试加载默认模型"""
+        # 第一次加载会创建默认配置
+        models = self.registry.load()
+        
+        # 验证默认模型已加载
+        self.assertGreater(len(models), 0)
+        
+        # 检查一些默认模型
+        model_names = [m.name for m in models]
+        self.assertIn("deepseek/v3", model_names)
+        self.assertIn("deepseek/r1", model_names)
+        
+        # 验证配置文件已创建
+        self.assertTrue(self.models_json.exists())
+    
+    def test_save_and_load(self):
+        """测试保存和加载"""
+        # 创建测试模型
+        test_models = [
+            LLMModel(
+                name="test/save1",
+                model_name="test-save-1",
+                model_type="saas/openai",
+                base_url="https://test.com/v1"
+            )
+        ]
+        
+        # 保存
+        self.registry.save(test_models)
+        
+        # 重新加载
+        loaded_models = self.registry.load()
+        
+        # 验证（应该包含默认模型和自定义模型）
+        loaded_names = [m.name for m in loaded_models]
+        self.assertIn("test/save1", loaded_names)
+        # 也应该包含默认模型
+        self.assertIn("deepseek/v3", loaded_names)
+    
+    def test_get_model(self):
+        """测试获取单个模型"""
+        # 获取存在的默认模型
+        model = self.registry.get("deepseek/v3")
+        self.assertIsNotNone(model)
+        if model:
+            self.assertEqual(model.name, "deepseek/v3")
+        
+        # 获取不存在的模型
+        model = self.registry.get("non/existent")
+        self.assertIsNone(model)
+    
+    def test_get_all_models(self):
+        """测试获取所有模型"""
+        models = self.registry.get_all()
+        self.assertIsInstance(models, dict)
+        self.assertGreater(len(models), 0)
+        
+        # 应该包含默认模型
+        self.assertIn("deepseek/v3", models)
+        self.assertIn("deepseek/r1", models)
+    
+    def test_add_or_update(self):
+        """测试添加或更新模型"""
+        # 添加新模型
+        new_model = LLMModel(
+            name="test/new",
+            model_name="test-new",
+            model_type="saas/openai",
+            base_url="https://new.com/v1"
+        )
+        self.registry.add_or_update(new_model)
+        
+        # 验证已添加
+        model = self.registry.get("test/new")
+        self.assertIsNotNone(model)
+        if model:
+            self.assertEqual(model.name, "test/new")
+        
+        # 更新模型
+        new_model.input_price = 5.0
+        self.registry.add_or_update(new_model)
+        
+        # 验证已更新
+        model = self.registry.get("test/new")
+        self.assertIsNotNone(model)
+        if model:
+            self.assertEqual(model.input_price, 5.0)
+    
+    def test_save_api_key(self):
+        """测试保存 API 密钥"""
+        # 保存密钥
+        result = self.registry.save_api_key("deepseek/v3", "test-api-key-123")
+        self.assertTrue(result)
+        
+        # 获取模型查看其 api_key_path
+        model = self.registry.get("deepseek/v3")
+        self.assertIsNotNone(model)
+        if model:
+            expected_key_path = model.api_key_path
+            self.assertIsNotNone(expected_key_path)
+            
+            if expected_key_path:
+                # 验证密钥文件已创建
+                key_file = self.models_json.parent / expected_key_path
+                self.assertTrue(key_file.exists())
+                self.assertEqual(key_file.read_text().strip(), "test-api-key-123")
+                
+                # 验证模型的 api_key_path 设置正确
+                self.assertEqual(model.api_key_path, expected_key_path)
+    
+    def test_save_api_key_nonexistent_model(self):
+        """测试为不存在的模型保存 API 密钥"""
+        result = self.registry.save_api_key("non/existent", "test-key")
+        self.assertFalse(result)
+    
+    def test_backward_compatibility(self):
+        """测试向后兼容性"""
+        # 创建没有新字段的旧格式数据
+        old_format_models = [
+            {
+                "name": "old/model",
+                "model_name": "old-model",
+                "model_type": "saas/openai",
+                "base_url": "https://old.com/v1",
+                "is_reasoning": False,
+                "input_price": 1.0,
+                "output_price": 2.0
+                # 没有 context_window, max_output_tokens, provider
+            }
+        ]
+        
+        # 写入文件
+        with open(self.models_json, 'w') as f:
+            json.dump(old_format_models, f)
+        
+        # 加载
+        models = self.registry.load()
+        
+        # 验证默认值已添加
+        old_model = next((m for m in models if m.name == "old/model"), None)
+        self.assertIsNotNone(old_model)
+        if old_model:
+            self.assertEqual(old_model.context_window, 32768)  # 默认值
+            self.assertEqual(old_model.max_output_tokens, 8096)  # 默认值
+            self.assertIsNone(old_model.provider)  # 默认值
+    
+    def test_no_cache_behavior(self):
+        """测试无缓存行为，每次都重新加载"""
+        # 第一次获取模型
+        model1 = self.registry.get("deepseek/v3")
+        self.assertIsNotNone(model1)
+        
+        # 添加一个新模型
+        new_model = LLMModel(
+            name="test/nocache",
+            model_name="test-nocache",
+            model_type="saas/openai",
+            base_url="https://test.com/v1"
+        )
+        self.registry.add_or_update(new_model)
+        
+        # 重新获取，应该能看到新模型
+        model2 = self.registry.get("test/nocache")
+        self.assertIsNotNone(model2)
+        if model2:
+            self.assertEqual(model2.name, "test/nocache")
+        
+        # 再次获取所有模型，应该包含新模型
+        all_models = self.registry.get_all()
+        self.assertIn("test/nocache", all_models)
+    
+    def test_file_corruption_handling(self):
+        """测试文件损坏的处理"""
+        # 写入无效的 JSON
+        with open(self.models_json, 'w') as f:
+            f.write("invalid json content")
+        
+        # 应该能正常加载（会重新创建默认配置）
+        models = self.registry.load()
+        self.assertGreater(len(models), 0)
+        
+        # 验证默认模型存在
+        model_names = [m.name for m in models]
+        self.assertIn("deepseek/v3", model_names)
+        
+        # 验证文件已被修复
+        self.assertTrue(self.models_json.exists())
+        with open(self.models_json, 'r') as f:
+            data = json.load(f)
+            self.assertIsInstance(data, list)
+    
+    def test_remove_model(self):
+        """测试删除模型"""
+        # 添加一个自定义模型
+        custom_model = LLMModel(
+            name="test/removable",
+            model_name="test-removable",
+            model_type="saas/openai",
+            base_url="https://test.com/v1"
+        )
+        self.registry.add_or_update(custom_model)
+        
+        # 验证模型已添加
+        self.assertTrue(self.registry.get("test/removable") is not None)
+        
+        # 删除模型
+        result = self.registry.remove_model("test/removable")
+        self.assertTrue(result)
+        
+        # 验证模型已被删除
+        self.assertIsNone(self.registry.get("test/removable"))
+    
+    def test_remove_default_model(self):
+        """测试删除默认模型（应该失败）"""
+        # 尝试删除默认模型
+        result = self.registry.remove_model("deepseek/v3")
+        self.assertFalse(result)
+        
+        # 验证默认模型仍然存在
+        self.assertIsNotNone(self.registry.get("deepseek/v3"))
+    
+    def test_remove_nonexistent_model(self):
+        """测试删除不存在的模型"""
+        result = self.registry.remove_model("non/existent")
+        self.assertFalse(result)
+    
+    def test_remove_model_with_api_key(self):
+        """测试删除带有 API 密钥的模型"""
+        # 添加带密钥的模型
+        custom_model = LLMModel(
+            name="test/with_key",
+            model_name="test-with-key",
+            model_type="saas/openai",
+            base_url="https://test.com/v1",
+            api_key="test-key-123"
+        )
+        self.registry.add_or_update(custom_model)
+        
+        # 保存密钥
+        self.registry.save_api_key("test/with_key", "test-key-123")
+        
+        # 获取模型并检查密钥文件是否存在
+        model = self.registry.get("test/with_key")
+        self.assertIsNotNone(model)
+        if model and model.api_key_path:
+            key_file = self.models_json.parent / model.api_key_path
+            self.assertTrue(key_file.exists())
+            
+            # 删除模型
+            result = self.registry.remove_model("test/with_key")
+            self.assertTrue(result)
+            
+            # 验证模型已被删除
+            self.assertIsNone(self.registry.get("test/with_key"))
+            
+            # 验证密钥文件也被删除（如果可能的话）
+            # 注意：这里不强制要求删除，因为可能有权限问题
+
+    def test_api_key_path_generation(self):
+        """测试 API 密钥路径自动生成逻辑"""
+        # 创建一个只有 api_key 没有 api_key_path 的模型数据
+        model_data = [
+            {
+                "name": "test/auto_path",
+                "model_name": "test-auto-path",
+                "model_type": "saas/openai",
+                "base_url": "https://test.com/v1",
+                "api_key": "test-key-456"
+                # 注意：没有 api_key_path
+            }
+        ]
+        
+        # 写入文件
+        with open(self.models_json, 'w') as f:
+            json.dump(model_data, f)
+        
+        # 创建密钥文件（按预期的路径）
+        expected_key_path = "test_auto_path"  # 模拟应该生成的路径
+        key_file = self.models_json.parent / expected_key_path
+        key_file.write_text("test-key-456")
+        
+        # 加载模型
+        models = self.registry.load()
+        
+        # 获取测试模型
+        test_model = next((m for m in models if m.name == "test/auto_path"), None)
+        self.assertIsNotNone(test_model)
+        
+        if test_model:
+            # 验证 api_key_path 已经被自动生成
+            self.assertEqual(test_model.api_key_path, expected_key_path)
+            
+            # 验证能够获取到 API 密钥
+            self.assertEqual(test_model.get_api_key(), "test-key-456")
+            
+            # 验证 has_api_key 返回 True
+            self.assertTrue(test_model.has_api_key)
+    
+    def test_api_key_path_priority(self):
+        """测试当 api_key_path 存在时，优先使用它而不是自动生成"""
+        # 创建一个既有 api_key 又有 api_key_path 的模型数据
+        model_data = [
+            {
+                "name": "test/priority",
+                "model_name": "test-priority",
+                "model_type": "saas/openai",
+                "base_url": "https://test.com/v1",
+                "api_key": "memory-key-789",
+                "api_key_path": "custom_key_path"
+            }
+        ]
+        
+        # 写入文件
+        with open(self.models_json, 'w') as f:
+            json.dump(model_data, f)
+        
+        # 创建密钥文件（使用自定义路径）
+        key_file = self.models_json.parent / "custom_key_path"
+        key_file.write_text("file-key-789")
+        
+        # 加载模型
+        models = self.registry.load()
+        
+        # 获取测试模型
+        test_model = next((m for m in models if m.name == "test/priority"), None)
+        self.assertIsNotNone(test_model)
+        
+        if test_model:
+            # 验证使用了指定的 api_key_path
+            self.assertEqual(test_model.api_key_path, "custom_key_path")
+            
+            # 验证从文件读取的密钥覆盖了内存中的密钥
+            self.assertEqual(test_model.get_api_key(), "file-key-789")
+            
+            # 验证 has_api_key 返回 True
+            self.assertTrue(test_model.has_api_key)
+
+
+if __name__ == '__main__':
+    unittest.main() 

autocoder/common/mcp_tools/__init__.py

@@ -0,0 +1,62 @@
+"""
+MCP Tools Module - A comprehensive Model Context Protocol (MCP) management system
+
+This module provides a complete set of tools for managing MCP servers, including
+server connections, tool execution, resource access, and server installation.
+"""
+
+from .hub import McpHub
+from .server import McpServer, get_mcp_server
+from .executor import McpExecutor, McpToolCall, McpResourceAccess
+from .installer import McpServerInstaller
+from .types import (
+    McpRequest, McpInstallRequest, McpRemoveRequest, McpListRequest,
+    McpListRunningRequest, McpRefreshRequest, McpServerInfoRequest,
+    McpResponse, ServerInfo, InstallResult, RemoveResult, ListResult,
+    ListRunningResult, RefreshResult, QueryResult, ErrorResult, ServerConfig,
+    StringResult, ExternalServerInfo, McpExternalServer, MarketplaceAddRequest,
+    MarketplaceAddResult, MarketplaceUpdateRequest, MarketplaceUpdateResult,
+    MarketplaceMCPServerItem
+)
+
+__all__ = [
+    # Main classes
+    'McpHub',
+    'McpServer',
+    'McpExecutor',
+    'McpServerInstaller',
+    'get_mcp_server',
+    
+    # Tool and resource classes
+    'McpToolCall',
+    'McpResourceAccess',
+    
+    # Request/Response types
+    'McpRequest',
+    'McpInstallRequest',
+    'McpRemoveRequest',
+    'McpListRequest',
+    'McpListRunningRequest',
+    'McpRefreshRequest',
+    'McpServerInfoRequest',
+    'McpResponse',
+    
+    # Result types
+    'ServerInfo',
+    'InstallResult',
+    'RemoveResult',
+    'ListResult',
+    'ListRunningResult',
+    'RefreshResult',
+    'QueryResult',
+    'ErrorResult',
+    'ServerConfig',
+    'StringResult',
+    'ExternalServerInfo',
+    'McpExternalServer',
+    'MarketplaceAddRequest',
+    'MarketplaceAddResult',
+    'MarketplaceUpdateRequest',
+    'MarketplaceUpdateResult',
+    'MarketplaceMCPServerItem',
+] 

autocoder/common/{mcp_tools.py → mcp_tools/executor.py}

@@ -1,56 +1,56 @@
+"""
+MCP Tools Executor Module
+
+This module provides the McpExecutor class for executing MCP tools and
+managing tool interactions. It handles tool calls, resource access,
+and result formatting.
+"""
+
 from typing import Dict, Any, List, Optional, Union
-from byzerllm import prompt
 import json
-import byzerllm
 import re
-from pydantic import BaseModel, Field
 from loguru import logger
-import os
-import time
+from byzerllm import prompt
+import byzerllm
+
+from .types import McpToolCall, McpResourceAccess
+from .hub import McpHub
+
 try:    
     import mcp.types as mcp_types    
 except ImportError:
     mcp_types = None
     logger.error("mcp is not installed(which requires python>=3.11), please install it by `pip install mcp`")
-from autocoder.common.mcp_hub import McpHub    
-
-
-class McpToolCall(BaseModel):
-    server_name: str = Field(..., description="The name of the MCP server")
-    tool_name: str = Field(..., description="The name of the tool to call")
-    arguments: Dict[str, Any] = Field(
-        default_factory=dict, description="The arguments to pass to the tool")
-
-
-class McpResourceAccess(BaseModel):
-    server_name: str = Field(..., description="The name of the MCP server")
-    uri: str = Field(..., description="The URI of the resource to access")
 
 
 class McpExecutor:
+    """
+    Executes MCP tools and manages tool interactions.
+    
+    This class provides a high-level interface for executing MCP tools,
+    accessing resources, and formatting results.
+    """
+    
     def __init__(self, mcp_hub: McpHub, llm: byzerllm.ByzerLLM):
         self.mcp_hub = mcp_hub
         self.llm = llm
 
-    def get_server_names(self) -> List[str]:
+    def get_server_names(self) -> str:
         """
         Get the names of all connected MCP servers.
 
         Returns:
-            List of server names
-            
+            Comma-separated string of server names or "(None running currently)"
         """
         server_names = [server.name for server in self.mcp_hub.get_servers()]
         return ",".join(server_names) or "(None running currently)"
 
     def get_connected_servers_info(self) -> str:
-        """Generate formatted information about connected MCP servers
-
-        Args:
-            mcp_hub: McpHub instance to get server information from
+        """
+        Generate formatted information about connected MCP servers
 
         Returns:
-            Formatted string with server details
+            Formatted string with server details including tools, resources, and templates
         """
         servers = self.mcp_hub.get_servers()
         if not servers:
@@ -108,11 +108,9 @@ class McpExecutor:
             info.append(server_info)
 
         return "\n\n".join(info)
-    
-    
 
     @byzerllm.prompt()
-    def mcp_prompt(self,query:str) -> str:
+    def mcp_prompt(self, query: str) -> str:
         """    
         TOOL USE
 
@@ -221,7 +219,7 @@ class McpExecutor:
             "server_names": self.get_server_names()
         }
 
-    async def extract_mcp_calls(self,content: str) -> List[Union[McpToolCall, McpResourceAccess]]:
+    async def extract_mcp_calls(self, content: str) -> List[Union[McpToolCall, McpResourceAccess]]:
         """
         Extract MCP tool calls and resource accesses from content.
 
@@ -269,17 +267,26 @@ class McpExecutor:
 
         return results
 
-    async def run(self, conversations: List[Dict[str, Any]]):                
+    async def run(self, conversations: List[Dict[str, Any]]):
+        """
+        Run MCP tool execution based on conversations.
+        
+        Args:
+            conversations: List of conversation messages
+            
+        Returns:
+            Tuple of (new_conversations, tool_results)
+        """
         all_tool_results = []
 
         query = conversations[-1]["content"]
-        new_conversations = conversations[0:-1]+[{
+        new_conversations = conversations[0:-1] + [{
             "role": "user",
             "content": self.mcp_prompt.prompt(query=query)
         }]
 
         v = self.llm.chat_oai(conversations=new_conversations)
-        content = v[0].output  
+        content = str(v[0].output)  
 
         tools = await self.extract_mcp_calls(content)
 
@@ -287,7 +294,7 @@ class McpExecutor:
             results = await self.execute_mcp_tools(tools)
             all_tool_results += results            
             
-        return new_conversations,all_tool_results
+        return new_conversations, all_tool_results
 
     def format_mcp_result(self, result: Any) -> str:
         """
@@ -301,18 +308,20 @@ class McpExecutor:
         """
         if result is None:
             return "(No result)"
-        if isinstance(result, mcp_types.CallToolResult):
-            if len(result.content) ==1:
-                return result.content[0].text
+        
+        if mcp_types and isinstance(result, mcp_types.CallToolResult):
+            if len(result.content) == 1:
+                content_item = result.content[0]
+                if hasattr(content_item, 'text'):
+                    return content_item.text
                                                                       
-        return json.dumps(result.model_dump(), indent=2, ensure_ascii=False)
+        return json.dumps(result.model_dump() if hasattr(result, 'model_dump') else result, indent=2, ensure_ascii=False)
 
     async def execute_mcp_tools(self, tools: List[Union[McpToolCall, McpResourceAccess]]) -> List[Any]:
         """
         Execute MCP tools and return results in order.
 
         Args:
-            mcp_hub: McpHub instance to execute tools
             tools: List of McpToolCall and McpResourceAccess objects
 
         Returns:
@@ -332,4 +341,4 @@ class McpExecutor:
             except Exception as e:
                 logger.error(f"Failed to execute MCP tool {tool}: {e}")
                 results.append(None)
-        return results
+        return results