auto-coder 0.1.287__py3-none-any.whl → 0.1.289__py3-none-any.whl

autocoder/plugins/token_helper_plugin.py

@@ -0,0 +1,343 @@
+"""
+Token Helper Plugin for Chat Auto Coder.
+Provides token counting functionality for files and projects.
+"""
+
+import os
+from typing import Any, Callable, Dict, List, Optional, Tuple
+from dataclasses import dataclass
+from collections import defaultdict
+
+from autocoder.plugins import Plugin, PluginManager
+from autocoder.rag.token_counter import count_tokens
+from autocoder.suffixproject import SuffixProject
+from autocoder.common import AutoCoderArgs, SourceCode
+from loguru import logger
+import json
+
+
+@dataclass
+class TokenCount:
+    """Represents token count information for a file."""
+    filename: str
+    tokens: int
+    relative_path: Optional[str] = None
+    file_size: Optional[int] = None
+
+
+class TokenHelperPlugin(Plugin):
+    """Token helper plugin for the Chat Auto Coder."""
+
+    name = "token_helper"
+    description = "Token helper plugin providing token counting for files and projects"
+    version = "0.1.0"
+
+    def __init__(self, manager: PluginManager, config: Optional[Dict[str, Any]] = None, config_path: Optional[str] = None):
+        """Initialize the Token helper plugin."""
+        super().__init__(manager, config, config_path)
+        self.token_counts = {}
+        self.project_dir = os.getcwd()        
+        self.base_persist_dir = os.path.join(".auto-coder", "plugins", "chat-auto-coder")        
+        self.auto_coder_config = {}
+        self.exclude_files = []
+
+    def load_auto_coder_config(self):
+        memory_path = os.path.join(self.base_persist_dir, "memory.json")
+        if os.path.exists(memory_path):
+            with open(memory_path, "r", encoding="utf-8") as f:
+                _memory = json.load(f)
+                self.auto_coder_config = _memory.get("conf",{})   
+                self.exclude_files = _memory.get("exclude_files",[])
+
+    def initialize(self) -> bool:
+        """Initialize the plugin.
+
+        Returns:
+            True if initialization was successful
+        """
+        self.load_auto_coder_config()
+        print(f"[{self.name}] Token Helper plugin initialized")
+        return True
+
+    def get_commands(self) -> Dict[str, Tuple[Callable, str]]:
+        """Get commands provided by this plugin.
+
+        Returns:
+            A dictionary of command name to handler and description
+        """
+        return {
+            "token/count": (self.count_tokens_in_project, "Count tokens in all project files"),
+            "token/top": (self.show_top_token_files, "Show top N files by token count"),
+            "token/file": (self.count_tokens_in_file, "Count tokens in a specific file"),
+            "token/summary": (self.show_token_summary, "Show token count summary for the project"),
+        }
+
+    def get_completions(self) -> Dict[str, List[str]]:
+        """Get completions provided by this plugin.
+
+        Returns:
+            A dictionary mapping command prefixes to completion options
+        """
+        completions = {
+            "/token/count": [],
+            "/token/top": ["5", "10", "20", "50", "100"],
+            "/token/file": [],
+            "/token/summary": [],
+        }
+        return completions
+
+    def count_tokens_in_project(self, args: str) -> None:
+        """Count tokens in all project files.
+        
+        Args:
+            args: Optional arguments including project directory and file types
+        """
+        args_dict = self._parse_args(args)
+        project_dir = args_dict.get("dir", self.auto_coder_config.get("project_dir", os.getcwd()))
+        project_type = args_dict.get("type", self.auto_coder_config.get("project_type", ".py"))
+        exclude = args_dict.get("exclude_files", [])
+        
+        self.project_dir = project_dir
+        print(f"Counting tokens in project: {project_dir}")
+        print(f"File types: {project_type}")
+        
+        try:
+            # Create AutoCoderArgs with necessary parameters
+            coder_args = AutoCoderArgs(
+                source_dir=project_dir,
+                project_type=project_type,
+                exclude_files=[f"regex://{exclude}"] + self.exclude_files if exclude else self.exclude_files
+            )
+            
+            # Use SuffixProject to get all source files
+            project = SuffixProject(args=coder_args)
+            
+            total_tokens = 0
+            file_count = 0
+            self.token_counts = {}
+            
+            print("Scanning files and counting tokens...")
+            
+            for source_code in project.get_source_codes():
+                file_count += 1
+                if file_count % 10 == 0:
+                    print(f"Processed {file_count} files...")
+                
+                tokens = count_tokens(source_code.source_code)
+                file_path = source_code.module_name
+                relative_path = os.path.relpath(file_path, project_dir)
+                file_size = len(source_code.source_code)
+                
+                self.token_counts[file_path] = TokenCount(
+                    filename=file_path,
+                    tokens=tokens,
+                    relative_path=relative_path,
+                    file_size=file_size
+                )
+                
+                total_tokens += tokens
+            
+            print(f"\nToken count complete!")
+            print(f"Total files: {file_count}")
+            print(f"Total tokens: {total_tokens:,}")
+            print(f"Use /token/top N to see the top N files by token count")
+            print(f"Use /token/summary to see a summary by file type")
+            
+        except Exception as e:
+            logger.error(f"Error counting tokens in project: {str(e)}")
+            print(f"Error: {str(e)}")
+
+    def _parse_args(self, args: str) -> Dict[str, str]:
+        """Parse command arguments.
+        
+        Args:
+            args: Command arguments string. Supports both:
+                 - Key=value format: dir=. type=.py,.java
+                 - Command line format: --dir . --type .py,.java
+            
+        Returns:
+            Dictionary of parsed arguments
+        """
+        result = {}
+        if not args:
+            return result
+            
+        # Try using argparse first
+        try:
+            import argparse
+            import shlex
+            
+            # Create parser with arguments expected by the plugin
+            parser = argparse.ArgumentParser(description='Token counter options')
+            parser.add_argument('--dir', '-d', help='Project directory')
+            parser.add_argument('--type', '-t', help='File types (comma separated)')
+            parser.add_argument('--exclude', '-e', help='Exclude pattern')
+            parser.add_argument('--verbose', '-v', action='store_true', help='Verbose output')
+            
+            # Parse with shlex to handle quoted arguments properly
+            parsed_args, unknown = parser.parse_known_args(shlex.split(args))
+            
+            # Convert namespace to dictionary, ignoring None values
+            for key, value in vars(parsed_args).items():
+                if value is not None:
+                    result[key] = value
+                    
+            # Handle any unknown arguments as key=value pairs
+            for arg in unknown:
+                if '=' in arg:
+                    key, value = arg.split('=', 1)
+                    # Strip leading dashes if present
+                    if key.startswith('--'):
+                        key = key[2:]
+                    elif key.startswith('-'):
+                        key = key[1:]
+                    result[key] = value
+            
+            return result
+            
+        except Exception as e:
+            # Fallback to original implementation if argparse fails
+            logger.debug(f"Argparse failed, using fallback parser: {str(e)}")
+            
+            parts = args.split()
+            i = 0
+            while i < len(parts):
+                part = parts[i]
+                
+                # Handle key=value format
+                if "=" in part:
+                    key, value = part.split("=", 1)
+                    # Strip leading dashes if present
+                    if key.startswith("--"):
+                        key = key[2:]
+                    elif key.startswith("-"):
+                        key = key[1:]
+                    result[key] = value
+                    i += 1
+                    continue
+                    
+                # Handle --key value or -key value format
+                if part.startswith("--"):
+                    key = part[2:]
+                    # Check if there's a value following this key
+                    if i + 1 < len(parts) and not parts[i + 1].startswith("-"):
+                        result[key] = parts[i + 1]
+                        i += 2
+                    else:
+                        # Flag option without value
+                        result[key] = "true"
+                        i += 1
+                elif part.startswith("-"):
+                    key = part[1:]
+                    # Check if there's a value following this key
+                    if i + 1 < len(parts) and not parts[i + 1].startswith("-"):
+                        result[key] = parts[i + 1]
+                        i += 2
+                    else:
+                        # Flag option without value
+                        result[key] = "true"
+                        i += 1
+                else:
+                    # Standalone argument without a key
+                    i += 1
+            
+            return result
+
+    def show_top_token_files(self, args: str) -> None:
+        """Show top N files by token count.
+        
+        Args:
+            args: Number of files to show
+        """
+        if not self.token_counts:
+            print("No token count data available. Run /token/count first.")
+            return
+            
+        try:
+            n = int(args.strip()) if args.strip() else 10
+        except ValueError:
+            print(f"Invalid value: {args}. Using default of 10.")
+            n = 10
+            
+        print(f"\nTop {n} files by token count:")
+        print(f"{'Tokens':<10} {'Size (bytes)':<15} {'File'}")
+        print(f"{'-'*10} {'-'*15} {'-'*50}")
+        
+        sorted_files = sorted(
+            self.token_counts.values(), 
+            key=lambda x: x.tokens, 
+            reverse=True
+        )
+        
+        for i, token_count in enumerate(sorted_files[:n], 1):
+            relative_path = token_count.relative_path or token_count.filename
+            print(f"{token_count.tokens:<10,} {token_count.file_size:<15,} {relative_path}")
+
+    def count_tokens_in_file(self, args: str) -> None:
+        """Count tokens in a specific file.
+        
+        Args:
+            args: Path to the file
+        """
+        if not args:
+            print("Please specify a file path.")
+            return
+            
+        file_path = args.strip()
+        
+        try:
+            with open(file_path, 'r', encoding='utf-8') as f:
+                content = f.read()
+                
+            tokens = count_tokens(content)
+            print(f"\nFile: {file_path}")
+            print(f"Tokens: {tokens:,}")
+            print(f"File size: {len(content):,} bytes")
+            print(f"Avg bytes per token: {len(content)/tokens:.2f}")
+            
+        except Exception as e:
+            print(f"Error counting tokens in file: {str(e)}")
+
+    def show_token_summary(self, args: str) -> None:
+        """Show token count summary by file type.
+        
+        Args:
+            args: Optional arguments
+        """
+        if not self.token_counts:
+            print("No token count data available. Run /token/count first.")
+            return
+            
+        by_extension = defaultdict(lambda: {"files": 0, "tokens": 0, "size": 0})
+        
+        for token_count in self.token_counts.values():
+            filename = token_count.filename
+            ext = os.path.splitext(filename)[1].lower() or "no_extension"
+            
+            by_extension[ext]["files"] += 1
+            by_extension[ext]["tokens"] += token_count.tokens
+            by_extension[ext]["size"] += token_count.file_size or 0
+            
+        total_tokens = sum(data["tokens"] for data in by_extension.values())
+        total_files = sum(data["files"] for data in by_extension.values())
+        total_size = sum(data["size"] for data in by_extension.values())
+        
+        print(f"\nToken count summary by file type:")
+        print(f"{'Extension':<12} {'Files':<8} {'Tokens':<12} {'% of Total':<12} {'Size (KB)':<12}")
+        print(f"{'-'*12} {'-'*8} {'-'*12} {'-'*12} {'-'*12}")
+        
+        for ext, data in sorted(by_extension.items(), key=lambda x: x[1]["tokens"], reverse=True):
+            percent = (data["tokens"] / total_tokens * 100) if total_tokens > 0 else 0
+            size_kb = data["size"] / 1024
+            print(f"{ext:<12} {data['files']:<8} {data['tokens']:<12,} {percent:<12.2f} {size_kb:<12.2f}")
+            
+        print(f"\nTotal Files: {total_files:,}")
+        print(f"Total Tokens: {total_tokens:,}")
+        print(f"Total Size: {total_size/1024/1024:.2f} MB")
+        
+        if self.project_dir:
+            print(f"Project Directory: {self.project_dir}")
+
+    def shutdown(self) -> None:
+        """Shutdown the plugin."""
+        print(f"[{self.name}] Token Helper plugin shutdown") 

autocoder/plugins/utils.py

@@ -0,0 +1,9 @@
+import json
+
+def load_json_file(file_path: str) -> dict:
+    with open(file_path, 'r') as f:
+        return json.load(f)
+
+def save_json_file(file_path: str, data: dict):
+    with open(file_path, 'w') as f:
+        json.dump(data, f, ensure_ascii=False, indent=4)

autocoder/rag/long_context_rag.py

@@ -38,7 +38,7 @@ from pydantic import BaseModel
 from byzerllm.utils.types import SingleOutputMeta
 from autocoder.rag.lang import get_message_with_format_and_newline
 from autocoder.rag.qa_conversation_strategy import get_qa_strategy
-
+from autocoder.rag.searchable import SearchableResults
 try:
     from autocoder_pro.rag.llm_compute import LLMComputeEngine
     pro_version = version("auto-coder-pro")
@@ -257,7 +257,7 @@ class LongContextRAG:
         请根据提供的文档内容、用户对话历史以及最后一个问题，提取并总结文档中与问题相关的重要信息。
         如果文档中没有相关信息，请回复"该文档中没有与问题相关的信息"。
         提取的信息尽量保持和原文中的一样，并且只输出这些信息。
-        """    
+        """
 
     def _get_document_retriever_class(self):
         """Get the document retriever class based on configuration."""
@@ -500,6 +500,9 @@ class LongContextRAG:
         except json.JSONDecodeError:
             pass
 
+        if not only_contexts and extra_request_params.get("only_contexts", False):
+            only_contexts = True
+
         logger.info(f"Query: {query} only_contexts: {only_contexts}")
         start_time = time.time()
 
@@ -593,10 +596,19 @@ class LongContextRAG:
             )
 
             if only_contexts:
-                final_docs = []
-                for doc in relevant_docs:
-                    final_docs.append(doc.model_dump())
-                return [json.dumps(final_docs, ensure_ascii=False)], []
+                try:
+                    searcher = SearchableResults()
+                    result = searcher.reorder(docs=relevant_docs)
+                    yield (json.dumps(result.model_dump(), ensure_ascii=False), SingleOutputMeta(input_tokens_count=rag_stat.recall_stat.total_input_tokens + rag_stat.chunk_stat.total_input_tokens,
+                                                                                                 generated_tokens_count=rag_stat.recall_stat.total_generated_tokens +
+                                                                                                 rag_stat.chunk_stat.total_generated_tokens,
+                                                                                                 ))
+                except Exception as e:
+                    yield (str(e), SingleOutputMeta(input_tokens_count=rag_stat.recall_stat.total_input_tokens + rag_stat.chunk_stat.total_input_tokens,
+                                                    generated_tokens_count=rag_stat.recall_stat.total_generated_tokens +
+                                                    rag_stat.chunk_stat.total_generated_tokens,
+                                                    ))
+                return
 
             if not relevant_docs:
                 yield ("没有找到可以回答你问题的相关文档", SingleOutputMeta(input_tokens_count=rag_stat.recall_stat.total_input_tokens + rag_stat.chunk_stat.total_input_tokens,
@@ -816,12 +828,13 @@ class LongContextRAG:
 
                 self._print_rag_stats(rag_stat)
             else:
-                
-                qa_strategy = get_qa_strategy(self.args.rag_qa_conversation_strategy)
+
+                qa_strategy = get_qa_strategy(
+                    self.args.rag_qa_conversation_strategy)
                 new_conversations = qa_strategy.create_conversation(
                     documents=[doc.source_code for doc in relevant_docs],
                     conversations=conversations
-                )                
+                )
 
                 chunks = target_llm.stream_chat_oai(
                     conversations=new_conversations,

autocoder/rag/relevant_utils.py

@@ -19,7 +19,7 @@ class TaskTiming(BaseModel):
     
 class FilterDoc(BaseModel):
     source_code: SourceCode
-    relevance: DocRelevance
+    relevance: Optional[DocRelevance]
     task_timing: TaskTiming
 
 

autocoder/rag/searchable.py

@@ -0,0 +1,58 @@
+import json
+from collections import Counter
+from typing import Dict, List, Any, Optional, Tuple, Set
+from pydantic import BaseModel
+from autocoder.rag.relevant_utils import FilterDoc
+
+
+class FileOccurrence(BaseModel):
+    """Represents a file and its occurrence count in search results"""
+    file_path: str
+    count: int
+    score: float = 0.0  # Optional relevance score
+
+class FileResult(BaseModel):
+    files: List[FileOccurrence]
+
+class SearchableResults:
+    """Class to process and organize search results by file frequency"""
+    
+    def __init__(self):
+        """Initialize the SearchableResults instance"""
+        pass
+    
+    def extract_original_docs(self, docs: List[FilterDoc]) -> List[str]:
+        """Extract all original_docs from a list of document metadata"""
+        all_files = []
+        
+        for doc in docs:
+            # Extract from metadata if available
+            metadata = doc.source_code.metadata
+            if "original_docs" in metadata:
+                all_files.extend(metadata["original_docs"])
+            # Also include the module_name from source_code as a fallback
+            else:
+                all_files.append(doc.source_code.module_name)
+        
+        return all_files
+    
+    def count_file_occurrences(self, files: List[str]) -> List[FileOccurrence]:
+        """Count occurrences of each file and return sorted list"""
+        # Count occurrences
+        counter = Counter(files)
+        
+        # Convert to FileOccurrence objects
+        occurrences = [
+            FileOccurrence(file_path=file_path, count=count)
+            for file_path, count in counter.items()
+        ]
+        
+        # Sort by count (descending)
+        return sorted(occurrences, key=lambda x: x.count, reverse=True)
+    
+    def reorder(self, docs: List[FilterDoc]) -> List[FileOccurrence]:
+        """Process search results to extract and rank files by occurrence (main entry point)"""
+        all_files = self.extract_original_docs(docs)
+        return FileResult(files=self.count_file_occurrences(all_files))
+    
+    

autocoder/version.py

@@ -1 +1 @@
-__version__ = "0.1.287"
+__version__ = "0.1.289"