auto-coder 0.1.215__py3-none-any.whl → 0.1.218__py3-none-any.whl

autocoder/command_args.py

@@ -171,7 +171,7 @@ def parse_args(input_args: Optional[List[str]] = None) -> AutoCoderArgs:
     parser.add_argument("--rag_url", default="", help="")
     parser.add_argument("--rag_params_max_tokens", default=4096, help="")
     parser.add_argument(
-        "--rag_type", default="storage", help="RAG type, default is storage"
+        "--rag_type", default="simple", help="RAG type, default is simple"
     )
 
     parser.add_argument(
@@ -502,8 +502,8 @@ def parse_args(input_args: Optional[List[str]] = None) -> AutoCoderArgs:
     chat_parser.add_argument("--rag_params_max_tokens", default=4096, help="")
     chat_parser.add_argument(
         "--rag_type",
-        default="storage",
-        help="RAG type (simple/storage), default is storage",
+        default="simple",
+        help="RAG type (simple/storage), default is simple",
     )
     chat_parser.add_argument("--target_file", default="./output.txt", help="")
     chat_parser.add_argument(
@@ -661,7 +661,7 @@ def parse_args(input_args: Optional[List[str]] = None) -> AutoCoderArgs:
     auto_tool_parser.add_argument(
         "--rag_params_max_tokens", default=4096, help="")
     auto_tool_parser.add_argument(
-        "--rag_type", default="storage", help="RAG type, default is storage"
+        "--rag_type", default="simple", help="RAG type, default is simple"
     )
     auto_tool_parser.add_argument(
         "--target_file", default="./output.txt", help="")
@@ -728,7 +728,7 @@ def parse_args(input_args: Optional[List[str]] = None) -> AutoCoderArgs:
     planner_parser.add_argument(
         "--rag_params_max_tokens", default=4096, help="")
     planner_parser.add_argument(
-        "--rag_type", default="storage", help="RAG type, default is storage"
+        "--rag_type", default="simple", help="RAG type, default is simple"
     )
     planner_parser.add_argument(
         "--target_file", default="./output.txt", help="")

autocoder/common/anything2img.py

@@ -0,0 +1,196 @@
+import os
+from typing import List, Optional, Dict, Any, Tuple
+from PIL import Image
+import fitz  # PyMuPDF
+import byzerllm
+from autocoder.common import AutoCoderArgs
+from loguru import logger
+import pydantic
+from docx import Document
+from spire.doc import Document
+from spire.doc import ImageType
+from PIL import Image
+from concurrent.futures import ThreadPoolExecutor, as_completed
+class ImageInfo(pydantic.BaseModel):
+    """
+    图片信息
+    """
+    coordinates: List[float] = pydantic.Field(..., description="图片坐标 [x1,y1,x2,y2]")
+    text: Optional[str] = pydantic.Field(None, description="图片描述")
+
+class Page(pydantic.BaseModel):
+    """
+    页面信息，包含文本和图片
+    """
+    text: str = pydantic.Field(..., description="页面文本内容")
+    images: List[ImageInfo] = pydantic.Field(default_factory=list, description="页面中的图片信息")
+    width: int = pydantic.Field(..., description="页面宽度")
+    height: int = pydantic.Field(..., description="页面高度")
+
+class Anything2Img:
+    def __init__(
+        self,
+        llm: byzerllm.ByzerLLM,
+        args: AutoCoderArgs,        
+        keep_conversion: bool = False,                
+    ):
+        self.llm = llm
+        self.vl_model = llm.get_sub_client("vl_model")
+        self.args = args
+        self.output_dir = args.output
+        os.makedirs(self.output_dir, exist_ok=True)
+        self.keep_conversion = keep_conversion            
+
+    @byzerllm.prompt()
+    def analyze_image(self, image_path: str) -> str:
+        """
+        {{ image }}
+        图片中一般包含文字，图片，图表。分析图片，返回该图片包含的文本内容以及图片位置信息。
+        请遵循以下格式返回：
+
+        ```json
+        {
+            "text": "页面的文本内容",
+            "images": [
+                {
+                    "coordinates": [x1, y1, x2, y2],
+                    "text": "对图片的描述"                    
+                }
+            ],
+            "width": 页面宽度,
+            "height": 页面高度
+        }
+        ```
+
+        注意：
+        1. 其中x1,y1是左上角坐标，x2,y2是右下角坐标，使用绝对坐标，也就是图片的像素坐标。
+        2. 文本内容应保持原有的段落格式
+        3. width和height是页面宽度，高度,要求整数类型
+        4. 格局图片中文本和图片的位置关系，在文本中使用 <image_placeholder> 来表示图片。
+        """
+        image = byzerllm.Image.load_image_from_path(image_path)
+        return {"image": image}
+
+    def convert_pdf(self, file_path: str) -> List[str]:
+        """转换PDF文件为图片列表"""
+        pdf_document = fitz.open(file_path)
+        image_paths = []                             
+        try:
+            # 分别保存每一页
+            for page_num in range(len(pdf_document)):
+                page = pdf_document[page_num]
+                pix = page.get_pixmap()
+                basename = os.path.basename(file_path).replace(" ", "_")
+                image_path = os.path.join(self.output_dir, f"{basename}_page{page_num + 1}.png")
+                pix.save(image_path)
+                image_paths.append(image_path)
+        finally:
+            # 确保PDF文档关闭
+            pdf_document.close()
+        return image_paths
+
+    def convert_docx(self, file_path: str) -> List[str]:
+        """使用 Spire.Doc 将 Word 文档直接转换为图片"""
+        # 创建 Spire.Doc 文档对象
+        doc = Document()
+        doc.LoadFromFile(file_path)
+
+        # 设置图片保存选项
+        image_paths = []
+        try:
+            # 将每一页保存为图片
+            for i in range(doc.GetPageCount()):
+                imageStream = doc.SaveImageToStreams(i, ImageType.Bitmap)
+                basename = os.path.basename(file_path).replace(" ", "_")
+                image_path = os.path.join(self.output_dir, f"{basename}_page{i + 1}.png")
+                with open(image_path, 'wb') as imageFile:
+                    imageFile.write(imageStream.ToArray())
+                image_paths.append(image_path)
+        finally:
+            # 确保文档关闭
+            doc.Close()
+
+        return image_paths
+
+    def convert(self, file_path: str) -> List[str]:
+        """根据文件类型选择合适的转换方法"""
+        file_path = os.path.abspath(file_path)
+        if file_path.lower().endswith('.pdf'):
+            return self.convert_pdf(file_path)
+        elif file_path.lower().endswith('.docx'):
+            return self.convert_docx(file_path)
+        else:
+            raise ValueError(f"Unsupported file format: {file_path}")
+
+    def to_markdown(self, file_path: str, size: int = -1, max_workers: int = 10) -> str:
+        """
+        将文档转换为Markdown格式
+
+        Args:
+            file_path: 文件路径
+            size: 转换的页数，-1表示全部
+            max_workers: 并行度，控制同时分析图片的线程数
+        """
+        # 创建 _images 目录
+        images_dir = os.path.join(self.output_dir, "_images")
+        os.makedirs(images_dir, exist_ok=True)
+        
+        # 转换文档为图片
+        if size == -1:
+            image_paths = self.convert(file_path)
+        else:
+            image_paths = self.convert(file_path)[0:size]
+        
+        pages: List[Page] = []
+        # 使用线程池并行分析图片
+        with ThreadPoolExecutor(max_workers=max_workers) as executor:
+            futures = {
+                executor.submit(
+                    self.analyze_image.with_llm(self.vl_model).with_return_type(Page).run,
+                    image_path
+                ): image_path for image_path in image_paths
+            }
+            
+            for future in as_completed(futures):
+                image_path = futures[future]
+                try:
+                    result = future.result()
+                    pages.append(result)
+                    logger.info(f"Analyzed {image_path}")
+                except Exception as e:
+                    logger.error(f"Failed to analyze {image_path}: {str(e)}")
+        
+        # 生成Markdown内容
+        markdown_content = []
+        
+        # 遍历每个页面和对应的图片路径
+        for page, image_path in zip(pages, image_paths):
+            # 处理页面中的每个图片
+            for img in page.images:                                
+                # 打开原始图片
+                original_image = Image.open(image_path)                
+                
+                # 获得坐标
+                x1 = img.coordinates[0]
+                y1 = img.coordinates[1]
+                x2 = img.coordinates[2]
+                y2 = img.coordinates[3]
+                
+                # 截取图片
+                cropped_image = original_image.crop((x1, y1, x2, y2))
+                
+                # 保存截取后的图片
+                cropped_image_path = os.path.join(images_dir, f"cropped_{os.path.basename(image_path)}")
+                cropped_image.save(cropped_image_path)
+                
+                # 将图片路径转换为Markdown格式
+                image_markdown = f"![{img.text}]({cropped_image_path})"
+                
+                # 替换文本中的<image_placeholder>为实际的图片Markdown
+                page.text = page.text.replace("<image_placeholder>", image_markdown, 1)
+            
+            # 将处理后的页面文本添加到Markdown内容中
+            markdown_content.append(page.text)
+        
+        # 将所有页面内容合并为一个Markdown文档
+        return '\n\n'.join(markdown_content)

autocoder/common/code_auto_generate.py

@@ -52,12 +52,15 @@ class CodeAutoGenerate:
 
         {%- if content %}
         下面是一些文件路径以及每个文件对应的源码：
-
+        <files>
         {{ content }}
+        </files>
         {%- endif %}
 
         {%- if context %}
+        <extra_context>
         {{ context }}
+        </extra_context>
         {%- endif %}
 
         下面是用户的需求：

autocoder/common/code_auto_generate_diff.py

@@ -131,12 +131,15 @@ class CodeAutoGenerateDiff:
 
         {%- if content %}
         下面是一些文件路径以及每个文件对应的源码：
-
+        <files>
         {{ content }}
+        </files>
         {%- endif %}
 
         {%- if context %}
+        <extra_context>
         {{ context }}
+        </extra_context>
         {%- endif %}
 
         下面是用户的需求：
@@ -265,12 +268,15 @@ class CodeAutoGenerateDiff:
 
         {%- if content %}
         下面是一些文件路径以及每个文件对应的源码：
-
+        <files>
         {{ content }}
+        </files>
         {%- endif %}
 
         {%- if context %}
+        <extra_context>
         {{ context }}
+        </extra_context>
         {%- endif %}
 
         下面是用户的需求：

autocoder/common/code_auto_generate_editblock.py

@@ -182,12 +182,15 @@ class CodeAutoGenerateEditBlock:
 
         {%- if content %}
         下面是一些文件路径以及每个文件对应的源码：
-
+        <files>
         {{ content }}
+        </files>
         {%- endif %}
 
         {%- if context %}
+        <extra_context>
         {{ context }}
+        </extra_context>
         {%- endif %}     
 
         下面是用户的需求：
@@ -343,12 +346,15 @@ class CodeAutoGenerateEditBlock:
 
         {%- if content %}
         下面是一些文件路径以及每个文件对应的源码：
-
+        <files>
         {{ content }}
+        </files>
         {%- endif %}
 
         {%- if context %}
+        <extra_context>
         {{ context }}
+        </extra_context>
         {%- endif %}     
 
         下面是用户的需求：

autocoder/common/code_auto_generate_strict_diff.py

@@ -114,12 +114,15 @@ class CodeAutoGenerateStrictDiff:
 
         {%- if content %}
         下面是一些文件路径以及每个文件对应的源码：
-
+        <files>
         {{ content }}
+        </files>
         {%- endif %}
 
         {%- if context %}
+        <extra_context>
         {{ context }}
+        </extra_context>
         {%- endif %}
 
         下面是用户的需求：
@@ -235,12 +238,15 @@ class CodeAutoGenerateStrictDiff:
 
         {%- if content %}
         下面是一些文件路径以及每个文件对应的源码：
-
+        <files>
         {{ content }}
+        </files>
         {%- endif %}
 
         {%- if context %}
+        <extra_context>
         {{ context }}
+        </extra_context>
         {%- endif %}
 
         下面是用户的需求：

autocoder/common/command_completer.py

@@ -14,6 +14,13 @@ COMMANDS = {
     },
     "/coding": {"/apply": {}, "/next": {}},
     "/chat": {"/new": {}, "/review": {}, "/no_context": {}},
+    "/mcp": {
+        "/add": "",         
+        "/remove": "",
+        "/list": "",
+        "/list_running": "",
+        "/refresh": ""
+    },
     "/lib": {
         "/add": "",
         "/remove": "",

autocoder/common/mcp_hub.py

@@ -1,6 +1,7 @@
 import os
 import json
 import asyncio
+import aiohttp
 from datetime import datetime
 from typing import Dict, List, Optional, Any, Set, Optional
 from pathlib import Path
@@ -10,6 +11,8 @@ from mcp import ClientSession
 from mcp.client.stdio import stdio_client, StdioServerParameters
 import mcp.types as mcp_types
 from loguru import logger
+import time
+
 
 class McpTool(BaseModel):
     """Represents an MCP tool configuration"""
@@ -60,6 +63,24 @@ class McpConnection:
         )
 
 
+MCP_PERPLEXITY_SERVER = '''
+{
+    "perplexity": {
+        "command": "python",
+        "args": [
+            "-m", "autocoder.common.mcp_servers.mcp_server_perplexity"
+        ],
+        "env": {
+            "PERPLEXITY_API_KEY": "{{PERPLEXITY_API_KEY}}"
+        }
+    }
+}
+'''
+
+MCP_BUILD_IN_SERVERS = {
+    "perplexity": json.loads(MCP_PERPLEXITY_SERVER)["perplexity"]
+}
+
 class McpHub:
     """
     Manages MCP server connections and interactions.
@@ -98,6 +119,45 @@ class McpHub:
         with open(self.settings_path, "w") as f:
             json.dump(default_settings, f, indent=2)
 
+    async def add_server_config(self, name: str, config:Dict[str,Any]) -> None:
+        """
+        Add or update MCP server configuration in settings file.
+
+        Args:
+            server_name_or_config: Name of the server or server configuration dictionary
+        """
+        try:
+            settings = self._read_settings()            
+            settings["mcpServers"][name] = config
+            with open(self.settings_path, "w") as f:
+                json.dump(settings, f, indent=2, ensure_ascii=False)
+            await self.initialize()
+            logger.info(f"Added/updated MCP server config: {name}")
+        except Exception as e:
+            logger.error(f"Failed to add MCP server config: {e}")
+            raise
+
+    async def remove_server_config(self, name: str) -> None:
+        """
+        Remove MCP server configuration from settings file.
+
+        Args:
+            name: Name of the server to remove
+        """
+        try:
+            settings = self._read_settings()
+            if name in settings["mcpServers"]:
+                del settings["mcpServers"][name]
+                with open(self.settings_path, "w") as f:
+                    json.dump(settings, f, indent=2, ensure_ascii=False)
+                logger.info(f"Removed MCP server config: {name}")
+                await self.initialize()
+            else:
+                logger.warning(f"MCP server {name} not found in settings")
+        except Exception as e:
+            logger.error(f"Failed to remove MCP server config: {e}")
+            raise
+
     async def initialize(self):
         """Initialize MCP server connections from settings"""
         try:
@@ -128,14 +188,15 @@ class McpHub:
             server_params = StdioServerParameters(
                 command=config["command"],
                 args=config.get("args", []),
-                env={**config.get("env", {}), "PATH": os.environ.get("PATH", "")},
+                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__()                 
+                session = await ClientSession(transport[0], transport[1]).__aenter__()
                 await session.initialize()
 
                 # Store connection with transport manager
@@ -148,9 +209,9 @@ class McpHub:
                 server.resources = await self._fetch_resources(name)
                 server.resource_templates = await self._fetch_resource_templates(name)
 
-            except Exception as e:                                
+            except Exception as e:
                 # Clean up transport if session initialization fails
-                
+
                 await transport_manager.__aexit__(None, None, None)
                 raise
 
@@ -179,6 +240,19 @@ class McpHub:
                 if name in self.connections:
                     del self.connections[name]
 
+    async def refresh_server_connection(self, name: str) -> None:
+        """
+        Refresh a server connection
+        """
+        try:
+            config = self._read_settings()
+            await self.delete_connection(name)
+            await self.connect_to_server(name, config.get("mcpServers", {}).get(name, {}))
+        except Exception as e:
+            logger.error(f"Failed to refresh MCP server {name}: {e}")
+            raise
+        
+
     async def update_server_connections(self, new_servers: Dict[str, Any]) -> None:
         """
         Update server connections based on new configuration
@@ -204,7 +278,8 @@ class McpHub:
                 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}")
+                    logger.info(
+                        f"Reconnected MCP server with updated config: {name}")
 
         finally:
             self.is_connecting = False
@@ -247,7 +322,7 @@ class McpHub:
                 for resource in response.resources
             ]
         except Exception as e:
-            logger.error(f"Failed to fetch resources for {server_name}: {e}")
+            logger.warning(f"Failed to fetch resources for {server_name}: {e}")
             return []
 
     async def _fetch_resource_templates(
@@ -284,7 +359,8 @@ class McpHub:
                 for template in response.resourceTemplates
             ]
         except Exception as e:
-            logger.error(f"Failed to fetch resource templates for {server_name}: {e}")
+            logger.warning(
+                f"Failed to fetch resource templates for {server_name}: {e}")
             return []
 
     def _read_settings(self) -> dict: