auto-coder 0.1.399__py3-none-any.whl → 0.1.400__py3-none-any.whl

autocoder/common/tokens/counter.py

@@ -0,0 +1,211 @@
+import os
+import concurrent.futures
+from typing import List, Dict, Optional, Union, Callable
+from pathlib import Path
+import time
+import re
+
+from autocoder.rag.variable_holder import VariableHolder
+from .models import TokenResult, DirectoryTokenResult
+from .file_detector import FileTypeDetector
+from .filters import FileFilter
+
+
+class TokenCounter:
+    """Token 计数器，用于统计文件和目录的 token 数量"""
+    
+    def __init__(self, 
+                 timeout: int = 30, 
+                 parallel: bool = True, 
+                 max_workers: int = 4):
+        """
+        初始化 Token 计数器
+        
+        Args:
+            timeout: 单文件处理超时时间（秒）
+            parallel: 是否并行处理
+            max_workers: 最大工作线程数
+        """
+        self.timeout = timeout
+        self.parallel = parallel
+        self.max_workers = max_workers
+        
+        # 确保 tokenizer 已经加载
+        if VariableHolder.TOKENIZER_MODEL is None:
+            raise RuntimeError("Tokenizer model not initialized. Please call load_tokenizer() first.")
+    
+    def count_file(self, file_path: str) -> TokenResult:
+        """
+        统计单个文件的 token 数量
+        
+        Args:
+            file_path: 文件路径
+            
+        Returns:
+            TokenResult: 统计结果
+        """
+        try:
+            if not os.path.isfile(file_path):
+                return TokenResult(
+                    file_path=file_path,
+                    token_count=0,
+                    char_count=0,
+                    line_count=0,
+                    success=False,
+                    error="File does not exist"
+                )
+            
+            # 检查是否为文本文件
+            if not FileTypeDetector.is_text_file(file_path):
+                return TokenResult(
+                    file_path=file_path,
+                    token_count=0,
+                    char_count=0,
+                    line_count=0,
+                    success=False,
+                    error="Not a text file"
+                )
+            
+            # 检测文件编码
+            encoding = FileTypeDetector.detect_encoding(file_path)
+            
+            # 读取文件内容
+            with open(file_path, 'r', encoding=encoding, errors='replace') as f:
+                content = f.read()
+            
+            # 统计行数
+            line_count = content.count('\n') + (0 if content == "" or content.endswith('\n') else 1)
+            
+            # 统计字符数
+            char_count = len(content)
+            
+            # 统计 token 数量
+            tokens = VariableHolder.TOKENIZER_MODEL.encode(content)
+            token_count = len(tokens)
+            
+            return TokenResult(
+                file_path=file_path,
+                token_count=token_count,
+                char_count=char_count,
+                line_count=line_count
+            )
+        except Exception as e:
+            return TokenResult(
+                file_path=file_path,
+                token_count=0,
+                char_count=0,
+                line_count=0,
+                success=False,
+                error=str(e)
+            )
+    
+    def count_files(self, file_paths: List[str]) -> List[TokenResult]:
+        """
+        批量统计多个文件的 token 数量
+        
+        Args:
+            file_paths: 文件路径列表
+            
+        Returns:
+            List[TokenResult]: 统计结果列表
+        """
+        if not self.parallel or len(file_paths) <= 1:
+            return [self.count_file(file_path) for file_path in file_paths]
+        
+        results = []
+        with concurrent.futures.ThreadPoolExecutor(max_workers=self.max_workers) as executor:
+            future_to_file = {
+                executor.submit(self.count_file, file_path): file_path 
+                for file_path in file_paths
+            }
+            
+            for future in concurrent.futures.as_completed(future_to_file):
+                results.append(future.result())
+                
+        return results
+    
+    def count_directory(self, 
+                        dir_path: str, 
+                        pattern: str = None,
+                        exclude_pattern: str = None,
+                        recursive: bool = True,
+                        max_depth: int = None) -> DirectoryTokenResult:
+        """
+        统计目录中所有文件的 token 数量
+        
+        Args:
+            dir_path: 目录路径
+            pattern: 文件名匹配模式（正则表达式）
+            exclude_pattern: 排除的文件名模式（正则表达式）
+            recursive: 是否递归处理子目录
+            max_depth: 最大递归深度
+            
+        Returns:
+            DirectoryTokenResult: 目录统计结果
+        """
+        if not os.path.isdir(dir_path):
+            return DirectoryTokenResult(
+                directory_path=dir_path,
+                total_tokens=0,
+                file_count=0,
+                skipped_count=0,
+                files=[],
+                errors=["Directory does not exist"]
+            )
+        
+        # 创建文件过滤器
+        patterns = [pattern] if pattern else []
+        exclude_patterns = [exclude_pattern] if exclude_pattern else []
+        file_filter = FileFilter(patterns=patterns, exclude_patterns=exclude_patterns)
+        
+        # 收集所有匹配的文件
+        all_files = []
+        skipped_count = 0
+        
+        for root, dirs, files in os.walk(dir_path):
+            # 检查递归深度
+            if max_depth is not None:
+                current_depth = root[len(dir_path):].count(os.sep)
+                if current_depth >= max_depth:
+                    dirs.clear()  # 不再递归子目录
+            
+            for file in files:
+                file_path = os.path.join(root, file)
+                if file_filter.matches(file_path):
+                    all_files.append(file_path)
+                else:
+                    skipped_count += 1
+            
+            if not recursive:
+                break  # 不递归处理子目录
+        
+        # 统计所有文件
+        file_results = self.count_files(all_files)
+        
+        # 计算总 token 数
+        total_tokens = sum(result.token_count for result in file_results if result.success)
+        
+        # 收集错误
+        errors = [
+            f"{result.file_path}: {result.error}" 
+            for result in file_results if not result.success
+        ]
+        
+        return DirectoryTokenResult(
+            directory_path=dir_path,
+            total_tokens=total_tokens,
+            file_count=len(file_results),
+            skipped_count=skipped_count,
+            files=file_results,
+            errors=errors
+        )
+    
+    def set_tokenizer(self, tokenizer_name: str) -> None:
+        """
+        更改 tokenizer（目前不支持，仅为接口预留）
+        
+        Args:
+            tokenizer_name: tokenizer 名称
+        """
+        # 目前仅支持默认的 tokenizer
+        pass

autocoder/common/tokens/file_detector.py

@@ -0,0 +1,105 @@
+import os
+import mimetypes
+from pathlib import Path
+
+
+class FileTypeDetector:
+    """文件类型检测器，用于判断文件类型和编码"""
+    
+    # 常见的文本文件MIME类型前缀
+    TEXT_MIME_PREFIXES = ('text/', 'application/json', 'application/xml', 'application/javascript')
+    
+    # 常见的文本文件扩展名
+    TEXT_EXTENSIONS = {
+        '.txt', '.md', '.py', '.js', '.jsx', '.ts', '.tsx', '.html', '.css', '.scss', '.sass',
+        '.json', '.xml', '.yaml', '.yml', '.ini', '.conf', '.sh', '.bash', '.zsh', '.c', '.cpp',
+        '.h', '.hpp', '.java', '.kt', '.rs', '.go', '.rb', '.php', '.pl', '.swift', '.dart',
+        '.vue', '.svelte', '.lua', '.r', '.sql', '.graphql', '.toml', '.csv'
+    }
+    
+    @staticmethod
+    def is_text_file(file_path: str) -> bool:
+        """
+        判断文件是否为文本文件
+        
+        Args:
+            file_path: 文件路径
+            
+        Returns:
+            bool: 是否为文本文件
+        """
+        # 首先通过扩展名判断
+        ext = os.path.splitext(file_path)[1].lower()
+        if ext in FileTypeDetector.TEXT_EXTENSIONS:
+            return True
+            
+        # 通过MIME类型判断
+        mime_type = FileTypeDetector.get_mime_type(file_path)
+        if any(mime_type.startswith(prefix) for prefix in FileTypeDetector.TEXT_MIME_PREFIXES):
+            return True
+            
+        # 通过文件内容判断
+        try:
+            with open(file_path, 'rb') as f:
+                # 读取前4KB进行判断
+                chunk = f.read(4096)
+                # 检查是否包含空字节（二进制文件通常包含空字节）
+                if b'\x00' in chunk:
+                    return False
+                # 尝试解码为UTF-8
+                try:
+                    chunk.decode('utf-8')
+                    return True
+                except UnicodeDecodeError:
+                    # 尝试其他常见编码
+                    for encoding in ['gbk', 'latin1', 'ascii']:
+                        try:
+                            chunk.decode(encoding)
+                            return True
+                        except UnicodeDecodeError:
+                            continue
+                    return False
+        except (IOError, OSError):
+            pass
+            
+        return False
+    
+    @staticmethod
+    def detect_encoding(file_path: str) -> str:
+        """
+        检测文件编码
+        
+        Args:
+            file_path: 文件路径
+            
+        Returns:
+            str: 文件编码，默认为utf-8
+        """
+        # 尝试常见编码
+        encodings = ['utf-8', 'gbk', 'latin1', 'ascii']
+        
+        for encoding in encodings:
+            try:
+                with open(file_path, 'r', encoding=encoding) as f:
+                    f.read(100)  # 尝试读取一小部分内容
+                    return encoding
+            except UnicodeDecodeError:
+                continue
+            except (IOError, OSError):
+                break
+                
+        return 'utf-8'  # 默认编码
+    
+    @staticmethod
+    def get_mime_type(file_path: str) -> str:
+        """
+        获取文件的MIME类型
+        
+        Args:
+            file_path: 文件路径
+            
+        Returns:
+            str: MIME类型
+        """
+        mime_type, _ = mimetypes.guess_type(file_path)
+        return mime_type or 'application/octet-stream'  # 默认为二进制类型

autocoder/common/tokens/filters.py

@@ -0,0 +1,111 @@
+import os
+import re
+from pathlib import Path
+from typing import List, Optional, Tuple
+
+
+class FileFilter:
+    """文件过滤器，用于过滤需要统计的文件"""
+    
+    def __init__(self, 
+                 patterns: List[str] = None, 
+                 exclude_patterns: List[str] = None,
+                 min_size: int = None, 
+                 max_size: int = None,
+                 only_text_files: bool = True):
+        """
+        初始化文件过滤器
+        
+        Args:
+            patterns: 包含的文件名模式（正则表达式）
+            exclude_patterns: 排除的文件名模式（正则表达式）
+            min_size: 最小文件大小（字节）
+            max_size: 最大文件大小（字节）
+            only_text_files: 是否只包含文本文件
+        """
+        self.patterns = []
+        self.exclude_patterns = []
+        self.min_size = min_size
+        self.max_size = max_size
+        self.only_text_files = only_text_files
+        
+        if patterns:
+            for pattern in patterns:
+                self.add_pattern(pattern)
+                
+        if exclude_patterns:
+            for pattern in exclude_patterns:
+                self.add_exclude_pattern(pattern)
+    
+    def add_pattern(self, pattern: str) -> None:
+        """
+        添加包含的文件名模式
+        
+        Args:
+            pattern: 正则表达式模式
+        """
+        try:
+            self.patterns.append(re.compile(pattern))
+        except re.error:
+            raise ValueError(f"Invalid regex pattern: {pattern}")
+    
+    def add_exclude_pattern(self, pattern: str) -> None:
+        """
+        添加排除的文件名模式
+        
+        Args:
+            pattern: 正则表达式模式
+        """
+        try:
+            self.exclude_patterns.append(re.compile(pattern))
+        except re.error:
+            raise ValueError(f"Invalid regex pattern: {pattern}")
+    
+    def set_size_range(self, min_size: Optional[int] = None, max_size: Optional[int] = None) -> None:
+        """
+        设置文件大小范围
+        
+        Args:
+            min_size: 最小文件大小（字节）
+            max_size: 最大文件大小（字节）
+        """
+        self.min_size = min_size
+        self.max_size = max_size
+    
+    def matches(self, file_path: str) -> bool:
+        """
+        检查文件是否匹配过滤条件
+        
+        Args:
+            file_path: 文件路径
+            
+        Returns:
+            bool: 是否匹配
+        """
+        # 检查文件是否存在
+        if not os.path.isfile(file_path):
+            return False
+            
+        # 检查文件大小
+        if self.min_size is not None or self.max_size is not None:
+            size = os.path.getsize(file_path)
+            if self.min_size is not None and size < self.min_size:
+                return False
+            if self.max_size is not None and size > self.max_size:
+                return False
+        
+        # 检查是否匹配排除模式
+        for pattern in self.exclude_patterns:
+            if pattern.search(file_path):
+                return False
+        
+        # 如果没有包含模式，则默认匹配所有文件
+        if not self.patterns:
+            return True
+            
+        # 检查是否匹配包含模式
+        for pattern in self.patterns:
+            if pattern.search(file_path):
+                return True
+                
+        return False

autocoder/common/tokens/models.py

@@ -0,0 +1,28 @@
+from typing import List, Dict, Optional
+from dataclasses import dataclass
+
+
+@dataclass
+class TokenResult:
+    """单个文件的 token 统计结果"""
+    file_path: str
+    token_count: int
+    char_count: int
+    line_count: int
+    success: bool = True
+    error: Optional[str] = None
+
+
+@dataclass
+class DirectoryTokenResult:
+    """目录的 token 统计结果"""
+    directory_path: str
+    total_tokens: int
+    file_count: int
+    skipped_count: int
+    files: List[TokenResult]
+    errors: List[str] = None
+    
+    def __post_init__(self):
+        if self.errors is None:
+            self.errors = []