auto-coder 0.1.209__py3-none-any.whl → 0.1.212__py3-none-any.whl

autocoder/common/buildin_tokenizer.py

@@ -0,0 +1,37 @@
+import pkg_resources
+from tokenizers import Tokenizer
+from typing import Optional
+
+class BuildinTokenizer:
+    _instance: Optional['BuildinTokenizer'] = None
+    _tokenizer: Optional[Tokenizer] = None
+
+    def __new__(cls) -> 'BuildinTokenizer':
+        if cls._instance is None:
+            cls._instance = super(BuildinTokenizer, cls).__new__(cls)
+            cls._instance._initialize()
+        return cls._instance
+
+    def _initialize(self) -> None:
+        try:
+            tokenizer_path = pkg_resources.resource_filename(
+                "autocoder", "data/tokenizer.json"
+            )
+        except FileNotFoundError:
+            tokenizer_path = None
+            
+        if tokenizer_path:
+            self._tokenizer = Tokenizer.from_file(tokenizer_path)
+        else:
+            raise ValueError("Cannot find tokenizer.json file in package data directory")
+
+    def count_tokens(self, text: str) -> int:
+        if not self._tokenizer:
+            raise ValueError("Tokenizer is not initialized")
+            
+        encoded = self._tokenizer.encode('{"role":"user","content":"' + text + '"}')
+        return len(encoded.ids)
+
+    @property
+    def tokenizer(self) -> Optional[Tokenizer]:
+        return self._tokenizer

autocoder/common/code_auto_generate.py

@@ -195,9 +195,7 @@ class CodeAutoGenerate:
             for result in results:
                 conversations_list.append(
                     conversations + [{"role": "assistant", "content": result}])
-        else:
-            results = []
-            conversations_list = []
+        else:            
             for _ in range(self.args.human_model_num):
                 v = self.llms[0].chat_oai(
                     conversations=conversations, llm_config=llm_config)

autocoder/common/code_auto_generate_diff.py

@@ -341,9 +341,7 @@ class CodeAutoGenerateDiff:
             for result in results:
                 conversations_list.append(
                     conversations + [{"role": "assistant", "content": result}])
-        else:
-            results = []
-            conversations_list = []
+        else:            
             for _ in range(self.args.human_model_num):
                 v = self.llms[0].chat_oai(
                     conversations=conversations, llm_config=llm_config)

autocoder/common/code_auto_generate_editblock.py

@@ -423,9 +423,7 @@ class CodeAutoGenerateEditBlock:
             for result in results:
                 conversations_list.append(
                     conversations + [{"role": "assistant", "content": result}])
-        else:
-            results = []
-            conversations_list = []
+        else:            
             for _ in range(self.args.human_model_num):
                 v = self.llms[0].chat_oai(
                     conversations=conversations, llm_config=llm_config)

autocoder/common/code_auto_generate_strict_diff.py

@@ -311,9 +311,7 @@ class CodeAutoGenerateStrictDiff:
             for result in results:
                 conversations_list.append(
                     conversations + [{"role": "assistant", "content": result}])
-        else:
-            results = []
-            conversations_list = []
+        else:            
             for _ in range(self.args.human_model_num):
                 v = self.llms[0].chat_oai(
                     conversations=conversations, llm_config=llm_config)

autocoder/common/code_modification_ranker.py

@@ -45,12 +45,15 @@ class CodeModificationRanker:
         }
         ```
 
-        注意，只输出前面要求的 Json 格式就好，不要输出其他内容，Json 需要使用 ```json ```包裹。
+        注意：        
+        1. 只输出前面要求的 Json 格式就好，不要输出其他内容，Json 需要使用 ```json ```包裹        
         '''
         
 
     def rank_modifications(self, generate_result: CodeGenerateResult) -> CodeGenerateResult:
         import time
+        from collections import defaultdict
+
         start_time = time.time()
         
         # 如果只有一个候选，直接返回
@@ -72,29 +75,44 @@ class CodeModificationRanker:
                     ) for _ in range(generate_times)
                 ]
                 
-                # Process results as they complete
+                # Collect all results
+                results = []
                 for future in as_completed(futures):
                     try:
                         v = future.result()
-                        # If we get a valid result, use it and cancel other tasks
-                        for f in futures:
-                            f.cancel()
-                            
-                        elapsed = time.time() - start_time
-                        logger.info(f"Ranking completed in {elapsed:.2f}s, best candidate index: {v.rank_result[0]}")
-                        
-                        rerank_contents = [generate_result.contents[i] for i in v.rank_result]
-                        rerank_conversations = [generate_result.conversations[i] for i in v.rank_result]
-                        return CodeGenerateResult(contents=rerank_contents,conversations=rerank_conversations)
+                        results.append(v.rank_result)
                     except Exception as e:
                         logger.warning(f"Ranking request failed: {str(e)}")
                         logger.debug(traceback.format_exc())
                         continue
+                
+                if not results:
+                    raise Exception("All ranking requests failed")
+                
+                # Calculate scores for each candidate
+                candidate_scores = defaultdict(float)
+                for rank_result in results:
+                    for idx, candidate_id in enumerate(rank_result):
+                        # Score is 1/(position + 1) since position starts from 0
+                        candidate_scores[candidate_id] += 1.0 / (idx + 1)
+                
+                # Sort candidates by score in descending order
+                sorted_candidates = sorted(candidate_scores.keys(), 
+                                         key=lambda x: candidate_scores[x], 
+                                         reverse=True)
+                
+                elapsed = time.time() - start_time
+                # Format scores for logging
+                score_details = ", ".join([f"candidate {i}: {candidate_scores[i]:.2f}" for i in sorted_candidates])
+                logger.info(f"Ranking completed in {elapsed:.2f}s, best candidate index: {sorted_candidates[0]}, scores: {score_details}")
+                
+                rerank_contents = [generate_result.contents[i] for i in sorted_candidates]
+                rerank_conversations = [generate_result.conversations[i] for i in sorted_candidates]
+                return CodeGenerateResult(contents=rerank_contents,conversations=rerank_conversations)
+                
         except Exception as e:
             logger.error(f"Ranking process failed: {str(e)}")
             logger.debug(traceback.format_exc())
-        
-        # If all requests failed, use the original codes
-        elapsed = time.time() - start_time
-        logger.warning(f"All ranking requests failed in {elapsed:.2f}s, using original order")
-        return generate_result
+            elapsed = time.time() - start_time
+            logger.warning(f"Ranking failed in {elapsed:.2f}s, using original order")
+            return generate_result

autocoder/common/mcp_hub.py

@@ -0,0 +1,326 @@
+import os
+import json
+import asyncio
+from datetime import datetime
+from typing import Dict, List, Optional, Any, Set, Optional
+from pathlib import Path
+from pydantic import BaseModel, Field
+
+from mcp import ClientSession
+from mcp.client.stdio import stdio_client, StdioServerParameters
+import mcp.types as mcp_types
+from loguru import logger
+
+class McpTool(BaseModel):
+    """Represents an MCP tool configuration"""
+
+    name: str
+    description: Optional[str] = None
+    input_schema: dict = Field(default_factory=dict)
+
+
+class McpResource(BaseModel):
+    """Represents an MCP resource configuration"""
+
+    uri: str
+    name: str
+    description: Optional[str] = None
+    mime_type: Optional[str] = None
+
+
+class McpResourceTemplate(BaseModel):
+    """Represents an MCP resource template"""
+
+    uri_template: str
+    name: str
+    description: Optional[str] = None
+    mime_type: Optional[str] = None
+
+
+class McpServer(BaseModel):
+    """Represents an MCP server configuration and status"""
+
+    name: str
+    config: str  # JSON string of server config
+    status: str = "disconnected"  # connected, disconnected, connecting
+    error: Optional[str] = None
+    tools: List[McpTool] = Field(default_factory=list)
+    resources: List[McpResource] = Field(default_factory=list)
+    resource_templates: List[McpResourceTemplate] = Field(default_factory=list)
+
+
+class McpConnection:
+    """Represents an active MCP server connection"""
+
+    def __init__(self, server: McpServer, session: ClientSession, transport_manager):
+        self.server = server
+        self.session = session
+        self.transport_manager = (
+            transport_manager  # Will hold transport context manager
+        )
+
+
+class McpHub:
+    """
+    Manages MCP server connections and interactions.
+    Similar to the TypeScript McpHub but adapted for Python/asyncio.
+    """
+
+    _instance = None
+
+    def __new__(cls, settings_path: Optional[str] = None):
+        if cls._instance is None:
+            cls._instance = super(McpHub, cls).__new__(cls)
+            cls._instance._initialized = False
+        return cls._instance
+
+    def __init__(self, settings_path: Optional[str] = None):
+        if self._initialized:
+            return
+        """Initialize the MCP Hub with a path to settings file"""
+        if settings_path is None:
+            self.settings_path = Path.home() / ".auto-coder" / "mcp" / "settings.json"
+        else:
+            self.settings_path = Path(settings_path)
+        self.connections: Dict[str, McpConnection] = {}
+        self.is_connecting = False
+
+        # Ensure settings directory exists
+        self.settings_path.parent.mkdir(parents=True, exist_ok=True)
+        if not self.settings_path.exists():
+            self._write_default_settings()
+
+        self._initialized = True
+
+    def _write_default_settings(self):
+        """Write default MCP settings file"""
+        default_settings = {"mcpServers": {}}
+        with open(self.settings_path, "w") as f:
+            json.dump(default_settings, f, indent=2)
+
+    async def initialize(self):
+        """Initialize MCP server connections from settings"""
+        try:
+            config = self._read_settings()
+            await self.update_server_connections(config.get("mcpServers", {}))
+        except Exception as e:
+            logger.error(f"Failed to initialize MCP servers: {e}")
+            raise
+
+    def get_servers(self) -> List[McpServer]:
+        """Get list of all configured servers"""
+        return [conn.server for conn in self.connections.values()]
+
+    async def connect_to_server(self, name: str, config: dict) -> None:
+        """
+        Establish connection to an MCP server with proper resource management
+        """
+        # Remove existing connection if present
+        if name in self.connections:
+            await self.delete_connection(name)
+
+        try:
+            server = McpServer(
+                name=name, config=json.dumps(config), status="connecting"
+            )
+
+            # Setup transport parameters
+            server_params = StdioServerParameters(
+                command=config["command"],
+                args=config.get("args", []),
+                env={**config.get("env", {}), "PATH": os.environ.get("PATH", "")},
+            )
+
+            # Create transport using context manager
+            transport_manager = stdio_client(server_params)
+            transport = await transport_manager.__aenter__()
+            try:
+                session = await ClientSession(transport[0], transport[1]).__aenter__()                 
+                await session.initialize()
+
+                # Store connection with transport manager
+                connection = McpConnection(server, session, transport_manager)
+                self.connections[name] = connection
+
+                # Update server status and fetch capabilities
+                server.status = "connected"
+                server.tools = await self._fetch_tools(name)
+                server.resources = await self._fetch_resources(name)
+                server.resource_templates = await self._fetch_resource_templates(name)
+
+            except Exception as e:                                
+                # Clean up transport if session initialization fails
+                
+                await transport_manager.__aexit__(None, None, None)
+                raise
+
+        except Exception as e:
+            error_msg = str(e)
+            logger.error(f"Failed to connect to server {name}: {error_msg}")
+            if name in self.connections:
+                self.connections[name].server.status = "disconnected"
+                self.connections[name].server.error = error_msg
+            raise
+
+    async def delete_connection(self, name: str) -> None:
+        """
+        Close and remove a server connection with proper cleanup
+        """
+        if name in self.connections:
+            try:
+                connection = self.connections[name]
+                # Clean up in reverse order of creation
+                await connection.session.__aexit__(None, None, None)
+                await connection.transport_manager.__aexit__(None, None, None)
+                del self.connections[name]
+            except Exception as e:
+                logger.error(f"Error closing connection to {name}: {e}")
+                # Continue with deletion even if cleanup fails
+                if name in self.connections:
+                    del self.connections[name]
+
+    async def update_server_connections(self, new_servers: Dict[str, Any]) -> None:
+        """
+        Update server connections based on new configuration
+        """
+        self.is_connecting = True
+        try:
+            current_names = set(self.connections.keys())
+            new_names = set(new_servers.keys())
+
+            # Remove deleted servers
+            for name in current_names - new_names:
+                await self.delete_connection(name)
+                logger.info(f"Deleted MCP server: {name}")
+
+            # Add or update servers
+            for name, config in new_servers.items():
+                current_conn = self.connections.get(name)
+
+                if not current_conn:
+                    # New server
+                    await self.connect_to_server(name, config)
+                    logger.info(f"Connected to new MCP server: {name}")
+                elif current_conn.server.config != json.dumps(config):
+                    # Updated configuration
+                    await self.connect_to_server(name, config)
+                    logger.info(f"Reconnected MCP server with updated config: {name}")
+
+        finally:
+            self.is_connecting = False
+
+    async def _fetch_tools(self, server_name: str) -> List[McpTool]:
+        """Fetch available tools from server"""
+        try:
+            connection = self.connections.get(server_name)
+            if not connection:
+                return []
+
+            response = await connection.session.list_tools()
+            return [
+                McpTool(
+                    name=tool.name,
+                    description=tool.description,
+                    input_schema=tool.inputSchema,
+                )
+                for tool in response.tools
+            ]
+        except Exception as e:
+            logger.error(f"Failed to fetch tools for {server_name}: {e}")
+            return []
+
+    async def _fetch_resources(self, server_name: str) -> List[McpResource]:
+        """Fetch available resources from server"""
+        try:
+            connection = self.connections.get(server_name)
+            if not connection:
+                return []
+
+            response = await connection.session.list_resources()
+            return [
+                McpResource(
+                    uri=resource.uri,
+                    name=resource.name,
+                    description=resource.description,
+                    mime_type=resource.mimeType,
+                )
+                for resource in response.resources
+            ]
+        except Exception as e:
+            logger.error(f"Failed to fetch resources for {server_name}: {e}")
+            return []
+
+    async def _fetch_resource_templates(
+        self, server_name: str
+    ) -> List[McpResourceTemplate]:
+        """Fetch available resource templates from server"""
+        try:
+            connection = self.connections.get(server_name)
+            if not connection:
+                return []
+
+            #       return await self.send_request(
+            #     types.ClientRequest(
+            #         types.PingRequest(
+            #             method="ping",
+            #         )
+            #     ),
+            #     types.EmptyResult,
+            # )
+
+            response = await connection.session.send_request(
+                mcp_types.ClientRequest(mcp_types.ListResourceTemplatesRequest(
+                    method="resources/templates/list",
+                )),
+                mcp_types.ListResourceTemplatesResult,
+            )
+            return [
+                McpResourceTemplate(
+                    uri_template=template.uriTemplate,
+                    name=template.name,
+                    description=template.description,
+                    mime_type=template.mimeType,
+                )
+                for template in response.resourceTemplates
+            ]
+        except Exception as e:
+            logger.error(f"Failed to fetch resource templates for {server_name}: {e}")
+            return []
+
+    def _read_settings(self) -> dict:
+        """Read MCP settings file"""
+        try:
+            with open(self.settings_path) as f:
+                return json.load(f)
+        except Exception as e:
+            logger.error(f"Failed to read MCP settings: {e}")
+            return {"mcpServers": {}}
+
+    async def call_tool(
+        self, server_name: str, tool_name: str, tool_arguments: Optional[Dict] = None
+    ) -> mcp_types.CallToolResult:
+        """
+        Call an MCP tool with arguments
+        """
+        connection = self.connections.get(server_name)
+        if not connection:
+            raise ValueError(f"No connection found for server: {server_name}")
+
+        return await connection.session.call_tool(tool_name, tool_arguments or {})
+
+    async def read_resource(self, server_name: str, uri: str) -> mcp_types.ReadResourceResult:
+        """
+        Read an MCP resource
+        """
+        connection = self.connections.get(server_name)
+        if not connection:
+            raise ValueError(f"No connection found for server: {server_name}")
+
+        return await connection.session.read_resource(uri)
+
+    async def shutdown(self):
+        """
+        Clean shutdown of all connections
+        """
+        for name in list(self.connections.keys()):
+            await self.delete_connection(name)

autocoder/common/mcp_server.py

@@ -0,0 +1,83 @@
+import asyncio
+from asyncio import Queue as AsyncQueue
+import threading
+from typing import List, Dict, Any, Optional
+from dataclasses import dataclass
+import byzerllm
+from autocoder.common.mcp_hub import McpHub
+from autocoder.common.mcp_tools import McpExecutor
+
+@dataclass 
+class McpRequest:
+    query: str
+    model: Optional[str] = None
+    
+@dataclass
+class McpResponse:
+    result: str
+    error: Optional[str] = None
+
+class McpServer:
+    def __init__(self):
+        self._request_queue = AsyncQueue()
+        self._response_queue = AsyncQueue()
+        self._running = False
+        self._task = None
+        self._loop = None
+        
+    def start(self):
+        if self._running:
+            return
+            
+        self._running = True
+        self._loop = asyncio.new_event_loop()
+        threading.Thread(target=self._run_event_loop, daemon=True).start()
+        
+    def stop(self):
+        if self._running:
+            self._running = False            
+            if self._loop:
+                self._loop.stop()
+                self._loop.close()
+                
+    def _run_event_loop(self):
+        asyncio.set_event_loop(self._loop)
+        self._task = self._loop.create_task(self._process_request())
+        self._loop.run_forever()
+                
+    async def _process_request(self):
+        hub = McpHub()
+        await hub.initialize()
+        
+        while self._running:
+            try:
+                request = await self._request_queue.get()
+                if request is None:
+                    break
+                    
+                llm = byzerllm.ByzerLLM.from_default_model(model=request.model)
+                mcp_executor = McpExecutor(hub, llm)
+                conversations = [{"role": "user", "content": request.query}]
+                _, results = await mcp_executor.run(conversations)
+                results_str = "\n\n".join(mcp_executor.format_mcp_result(result) for result in results)
+                await self._response_queue.put(McpResponse(result=results_str))
+            except Exception as e:
+                await self._response_queue.put(McpResponse(result="", error=str(e)))
+        
+    def send_request(self, request: McpRequest) -> McpResponse:
+        async def _send():
+            await self._request_queue.put(request)
+            return await self._response_queue.get()
+            
+        future = asyncio.run_coroutine_threadsafe(_send(), self._loop)
+        return future.result()
+
+# Global MCP server instance
+_mcp_server = None
+
+def get_mcp_server():
+    global _mcp_server
+    if _mcp_server is None:
+        _mcp_server = McpServer()
+        _mcp_server.start()
+    return _mcp_server