jarvis-ai-assistant 0.1.102__py3-none-any.whl → 0.1.104__py3-none-any.whl

jarvis/utils.py

@@ -1,32 +1,68 @@
 import hashlib
 from pathlib import Path
-import sys
 import time
 import os
 from enum import Enum
 from datetime import datetime
-from typing import Any
+from typing import Any, Dict, Optional
 import colorama
 from colorama import Fore, Style as ColoramaStyle
 import numpy as np
 from prompt_toolkit import PromptSession
 from prompt_toolkit.styles import Style as PromptStyle
 from prompt_toolkit.formatted_text import FormattedText
+from sentence_transformers import SentenceTransformer
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
+import torch
 import yaml
+import faiss
+from pygments.lexers import guess_lexer
+from pygments.util import ClassNotFound
+
+from rich.console import Console
+from rich.theme import Theme
+from rich.panel import Panel
+from rich.text import Text
+from rich.traceback import install as install_rich_traceback
+from rich.syntax import Syntax
+
+from prompt_toolkit.completion import Completer, Completion, PathCompleter
+from prompt_toolkit.document import Document
+from fuzzywuzzy import fuzz
 
 # 初始化colorama
 colorama.init()
 
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
-os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
 
 current_agent = []
 
+# Install rich traceback handler
+install_rich_traceback()
+
+# Create console with custom theme
+custom_theme = Theme({
+    "info": "yellow",
+    "warning": "yellow",
+    "error": "red",
+    "success": "green",
+    "system": "cyan",
+    "code": "green",
+    "result": "blue",
+    "planning": "magenta",
+    "progress": "white",
+    "debug": "blue",
+    "user": "green",
+    "tool": "yellow",
+})
+
+console = Console(theme=custom_theme)
+
 def add_agent(agent_name: str):
     current_agent.append(agent_name)
 
-def get_current_agent():
-    return current_agent[-1] if current_agent else "No Agent"
+def get_agent_list():
+    return ']['.join(current_agent) if current_agent else "No Agent"
 
 def delete_current_agent():
     current_agent.pop()
@@ -46,25 +82,9 @@ class OutputType(Enum):
     TOOL = "tool"         # Tool call
 
 class PrettyOutput:
-    """美化输出类"""
+    """Pretty output using rich"""
     
-    # 颜色方案 - 只使用前景色
-    COLORS = {
-        OutputType.SYSTEM: Fore.CYAN,      # Cyan - AI assistant
-        OutputType.CODE: Fore.GREEN,       # Green - Code
-        OutputType.RESULT: Fore.BLUE,      # Blue - Result
-        OutputType.ERROR: Fore.RED,        # Red - Error
-        OutputType.INFO: Fore.YELLOW,      # Yellow - Prompt
-        OutputType.PLANNING: Fore.MAGENTA, # Magenta - Planning
-        OutputType.PROGRESS: Fore.WHITE,   # White - Progress
-        OutputType.SUCCESS: Fore.GREEN,    # Green - Success
-        OutputType.WARNING: Fore.YELLOW,   # Yellow - Warning
-        OutputType.DEBUG: Fore.BLUE,       # Blue - Debug
-        OutputType.USER: Fore.GREEN,       # Green - User
-        OutputType.TOOL: Fore.YELLOW,      # Yellow - Tool
-    }
-    
-    # 图标方案
+    # Icons for different output types
     ICONS = {
         OutputType.SYSTEM: "🤖",    # Robot - AI assistant
         OutputType.CODE: "📝",      # Notebook - Code
@@ -79,67 +99,106 @@ class PrettyOutput:
         OutputType.USER: "👤",      # User - User
         OutputType.TOOL: "🔧",      # Wrench - Tool
     }
-    
-    # 前缀方案
-    PREFIXES = {
-        OutputType.SYSTEM: "Assistant",
-        OutputType.CODE: "Code",
-        OutputType.RESULT: "Result",
-        OutputType.ERROR: "Error",
-        OutputType.INFO: "Info",
-        OutputType.PLANNING: "Plan",
-        OutputType.PROGRESS: "Progress",
-        OutputType.SUCCESS: "Success",
-        OutputType.WARNING: "Warning",
-        OutputType.DEBUG: "Debug",
-        OutputType.USER: "User",
-        OutputType.TOOL: "Tool",
+
+    # Common language mapping dictionary
+    _lang_map = {
+        'Python': 'python',
+        'JavaScript': 'javascript',
+        'TypeScript': 'typescript',
+        'Java': 'java',
+        'C++': 'cpp',
+        'C#': 'csharp',
+        'Ruby': 'ruby',
+        'PHP': 'php',
+        'Go': 'go',
+        'Rust': 'rust',
+        'Bash': 'bash',
+        'HTML': 'html',
+        'CSS': 'css',
+        'SQL': 'sql',
+        'R': 'r',
+        'Kotlin': 'kotlin',
+        'Swift': 'swift',
+        'Scala': 'scala',
+        'Perl': 'perl',
+        'Lua': 'lua',
+        'YAML': 'yaml',
+        'JSON': 'json',
+        'XML': 'xml',
+        'Markdown': 'markdown',
+        'Text': 'text',
+        'Shell': 'bash',
+        'Dockerfile': 'dockerfile',
+        'Makefile': 'makefile',
+        'INI': 'ini',
+        'TOML': 'toml',
     }
 
     @staticmethod
-    def format(text: str, output_type: OutputType, timestamp: bool = True) -> str:
-        """Format output text"""
-        color = PrettyOutput.COLORS.get(output_type, "")
-        icon = PrettyOutput.ICONS.get(output_type, "")
-        prefix = PrettyOutput.PREFIXES.get(output_type, "")
-        
-        # 添加时间戳 - 使用白色
-        time_str = f"{Fore.BLUE}[{get_current_agent()}]{ColoramaStyle.RESET_ALL}{Fore.WHITE}[{datetime.now().strftime('%H:%M:%S')}]{ColoramaStyle.RESET_ALL} " if timestamp else ""
+    def _detect_language(text: str, default_lang: str = 'markdown') -> str:
+        """Helper method to detect language and map it to syntax highlighting name"""
+        try:
+            lexer = guess_lexer(text)
+            detected_lang = lexer.name
+            return PrettyOutput._lang_map.get(detected_lang, default_lang)
+        except ClassNotFound:
+            return default_lang
+        except Exception:
+            return default_lang
+
+    @staticmethod
+    def format(text: str, output_type: OutputType, timestamp: bool = True) -> Text:
+        """Format output text using rich Text"""
+        # Create rich Text object
+        formatted = Text()
         
-        # 格式化输出
-        formatted_text = f"{time_str}{color}{icon} {prefix}: {text}{ColoramaStyle.RESET_ALL}"
+        # Add timestamp and agent info
+        if timestamp:
+            formatted.append(f"[{datetime.now().strftime('%H:%M:%S')}] ", style="white")
+        formatted.append(f"[{get_agent_list()}]", style="blue")
+        # Add icon
+        icon = PrettyOutput.ICONS.get(output_type, "")
+        formatted.append(f"{icon} ", style=output_type.value)
         
-        return formatted_text
+        return formatted
 
     @staticmethod
-    def print(text: str, output_type: OutputType, timestamp: bool = True):
-        """Print formatted output"""
-        print(PrettyOutput.format(text, output_type, timestamp))
+    def print(text: str, output_type: OutputType, timestamp: bool = True, lang: Optional[str] = None):
+        """Print formatted output using rich console"""
+        # Get formatted header
+        lang = lang if lang is not None else PrettyOutput._detect_language(text, default_lang='markdown')
+        header = PrettyOutput.format("", output_type, timestamp)
+
+        content = Syntax(text, lang, theme="monokai")
+            
+        # Print panel with appropriate border style
+        border_style = "red" if output_type == OutputType.ERROR else output_type.value
+        console.print(Panel(content, border_style=border_style, title=header, title_align="left", highlight=True))
+        
+        # Print stack trace for errors
         if output_type == OutputType.ERROR:
-            import traceback
-            PrettyOutput.print(f"Error trace: {traceback.format_exc()}", OutputType.INFO)
+            console.print_exception()
 
     @staticmethod
     def section(title: str, output_type: OutputType = OutputType.INFO):
-        """Print paragraph title with separator"""
-        width = 100
-        color = PrettyOutput.COLORS.get(output_type, "")
-        print(f"\n{color}" + "=" * width + f"{ColoramaStyle.RESET_ALL}")
-        PrettyOutput.print(title.center(width - 25), output_type, timestamp=False)
-        print(f"{color}" + "=" * width + f"{ColoramaStyle.RESET_ALL}\n")
+        """Print section title in a panel"""
+        panel = Panel(
+            Text(title, style=output_type.value, justify="center"),
+            border_style=output_type.value
+        )
+        console.print()
+        console.print(panel)
+        console.print()
 
     @staticmethod
     def print_stream(text: str):
-        """Print stream output, no line break"""
-        color = PrettyOutput.COLORS.get(OutputType.SYSTEM, "")
-        sys.stdout.write(f"{color}{text}{ColoramaStyle.RESET_ALL}")
-        sys.stdout.flush()
+        """Print stream output without line break"""
+        console.print(text, style="system", end="")
 
     @staticmethod
     def print_stream_end():
-        """Stream output end, print line break"""
-        sys.stdout.write("\n")
-        sys.stdout.flush()
+        """End stream output with line break"""
+        console.print()
 
 def get_single_line_input(tip: str) -> str:
     """Get single line input, support direction key, history function, etc."""
@@ -149,14 +208,87 @@ def get_single_line_input(tip: str) -> str:
     })
     return session.prompt(f"{tip}", style=style)
 
+def make_choice_input(tip: str, choices: list) -> str:
+    """Get choice input, support direction key, history function, etc."""
+    session = PromptSession(history=None)
+    style = PromptStyle.from_dict({
+        'prompt': 'ansicyan',
+    })
+    return session.prompt(f"{tip}", style=style)
+
+class FileCompleter(Completer):
+    """Custom completer for file paths with fuzzy matching."""
+    def __init__(self):
+        self.path_completer = PathCompleter()
+        
+    def get_completions(self, document: Document, complete_event):
+        text = document.text_before_cursor
+        cursor_pos = document.cursor_position
+        
+        # Find all @ positions in text
+        at_positions = [i for i, char in enumerate(text) if char == '@']
+        
+        if not at_positions:
+            return
+            
+        # Get the last @ position
+        current_at_pos = at_positions[-1]
+        
+        # If cursor is not after the last @, don't complete
+        if cursor_pos <= current_at_pos:
+            return
+            
+        # Check if there's a space after @
+        text_after_at = text[current_at_pos + 1:cursor_pos]
+        if ' ' in text_after_at:
+            return
+            
+        # Get the text after the current @
+        file_path = text_after_at.strip()
+        
+        # Get all possible files from current directory
+        all_files = []
+        for root, _, files in os.walk('.'):
+            for f in files:
+                path = os.path.join(root, f)
+                # Remove ./ from the beginning
+                path = path[2:] if path.startswith('./') else path
+                all_files.append(path)
+        
+        # If no input after @, show all files
+        # Otherwise use fuzzy matching
+        if not file_path:
+            scored_files = [(path, 100) for path in all_files]
+        else:
+            scored_files = [
+                (path, fuzz.ratio(file_path.lower(), path.lower()))
+                for path in all_files
+            ]
+            scored_files.sort(key=lambda x: x[1], reverse=True)
+        
+        # Return completions for files
+        for path, score in scored_files:
+            if not file_path or score > 30:  # Show all if no input, otherwise filter by score
+                completion = Completion(
+                    text=path,
+                    start_position=-len(file_path),
+                    display=f"{path}" if not file_path else f"{path} ({score}%)",
+                    display_meta="File"
+                )
+                yield completion
+
 def get_multiline_input(tip: str) -> str:
-    """Get multi-line input, support direction key, history function, etc."""
-    print(f"{Fore.GREEN}{tip}{ColoramaStyle.RESET_ALL}")
+    """Get multi-line input, support direction key, history function, and file completion.
     
-    # 创建输入会话，启用历史记录
-    session = PromptSession(history=None)  # 使用默认历史记录
+    Args:
+        tip: The prompt tip to display
+        
+    Returns:
+        str: The entered text
+    """
+    print(f"{Fore.GREEN}{tip}{ColoramaStyle.RESET_ALL}")
     
-    # 定义提示符样式
+    # Define prompt style
     style = PromptStyle.from_dict({
         'prompt': 'ansicyan',
     })
@@ -164,28 +296,34 @@ def get_multiline_input(tip: str) -> str:
     lines = []
     try:
         while True:
-            # 设置提示符
+            # Set prompt
             prompt = FormattedText([
                 ('class:prompt', '... ' if lines else '>>> ')
             ])
             
-            # 获取输入
+            # Create new session with new completer for each line
+            session = PromptSession(
+                history=None,  # Use default history
+                completer=FileCompleter()  # New completer instance for each line
+            )
+            
+            # Get input with completion support
             line = session.prompt(
                 prompt,
                 style=style,
             ).strip()
             
-            # 空行处理
+            # Handle empty line
             if not line:
-                if not lines:  # 第一行就输入空行
+                if not lines:  # First line is empty
                     return ""
-                break  # 结束多行输入
+                break  # End multi-line input
                 
             lines.append(line)
             
     except KeyboardInterrupt:
-        PrettyOutput.print("\nInput cancelled", OutputType.INFO)
-        return "__interrupt__"
+        PrettyOutput.print("Input cancelled", OutputType.INFO)
+        return ""
     
     return "\n".join(lines)
 
@@ -239,6 +377,73 @@ def find_git_root(dir="."):
     os.chdir(curr_dir)
     return ret
 
+def has_uncommitted_changes():
+    # Check working directory changes
+    working_changes = os.popen("git diff --exit-code").read().strip() != ""
+    # Check staged changes
+    staged_changes = os.popen("git diff --cached --exit-code").read().strip() != ""
+    return working_changes or staged_changes
+
+def load_embedding_model():
+    model_name = "BAAI/bge-m3"
+    cache_dir = os.path.expanduser("~/.cache/huggingface/hub")
+    
+
+    try:
+        # Load model
+        embedding_model = SentenceTransformer(
+            model_name,
+            cache_folder=cache_dir,
+            local_files_only=True
+        )
+    except Exception as e:
+        # Load model
+        embedding_model = SentenceTransformer(
+            model_name,
+            cache_folder=cache_dir,
+            local_files_only=False
+        )
+    
+    return embedding_model
+
+def load_rerank_model():
+    """Load reranking model"""
+    model_name = "BAAI/bge-reranker-v2-m3"
+    cache_dir = os.path.expanduser("~/.cache/huggingface/hub")
+    
+    PrettyOutput.print(f"Loading reranking model: {model_name}...", OutputType.INFO)
+    
+    try:
+        # Load model and tokenizer
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_name,
+            cache_dir=cache_dir,
+            local_files_only=True
+        )
+        model = AutoModelForSequenceClassification.from_pretrained(
+            model_name,
+            cache_dir=cache_dir,
+            local_files_only=True
+        )
+    except Exception as e:
+        # Load model and tokenizer
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_name,
+            cache_dir=cache_dir,
+            local_files_only=False
+        )
+        model = AutoModelForSequenceClassification.from_pretrained(
+            model_name,
+            cache_dir=cache_dir,
+            local_files_only=False
+        )
+    
+    # Use GPU if available
+    if torch.cuda.is_available():
+        model = model.cuda()
+    model.eval()
+    
+    return model, tokenizer
 
 def get_max_context_length():
     return int(os.getenv('JARVIS_MAX_CONTEXT_LENGTH', '131072'))  # 默认128k
@@ -270,6 +475,10 @@ def get_file_md5(filepath: str)->str:
     return hashlib.md5(open(filepath, "rb").read(100*1024*1024)).hexdigest()
 
 
+def dont_use_local_model():
+    return os.getenv('JARVIS_DONT_USE_LOCAL_MODEL', 'false') == 'true'
+
+
 def _create_methodology_embedding(embedding_model: Any, methodology_text: str) -> np.ndarray:
     """Create embedding vector for methodology text"""
     try:
@@ -290,19 +499,77 @@ def _create_methodology_embedding(embedding_model: Any, methodology_text: str) -
 
 def load_methodology(user_input: str) -> str:
     """Load methodology and build vector index"""
+    PrettyOutput.print("Loading methodology...", OutputType.PROGRESS)
     user_jarvis_methodology = os.path.expanduser("~/.jarvis/methodology")
     if not os.path.exists(user_jarvis_methodology):
         return ""
+    
+    def make_methodology_prompt(data: Dict) -> str:
+        ret = """This is the standard methodology for handling previous problems, if the current task is similar, you can refer to it, if not,just ignore it:\n""" 
+        for key, value in data.items():
+            ret += f"Problem: {key}\nMethodology: {value}\n"
+        return ret
 
     try:
         with open(user_jarvis_methodology, "r", encoding="utf-8") as f:
             data = yaml.safe_load(f)
 
+        if dont_use_local_model():
+            return make_methodology_prompt(data)
+
+        # Reset data structure
+        methodology_data = []
+        vectors = []
+        ids = []
+
+        # Get embedding model
+        embedding_model = load_embedding_model()
         
-        ret = """This is the standard methodology for handling previous problems, if the current task is similar, you can refer to it:"""
-        for k, v in data.items():
-            ret += f"\n{k}: {v}\n"
-        return ret
+        # Create test embedding to get correct dimension
+        test_embedding = _create_methodology_embedding(embedding_model, "test")
+        embedding_dimension = len(test_embedding)
+
+        # Create embedding vector for each methodology
+        for i, (key, value) in enumerate(data.items()):
+            methodology_text = f"{key}\n{value}"
+            embedding = _create_methodology_embedding(embedding_model, methodology_text)
+            vectors.append(embedding)
+            ids.append(i)
+            methodology_data.append({"key": key, "value": value})
+
+        if vectors:
+            vectors_array = np.vstack(vectors)
+            # Use correct dimension from test embedding
+            hnsw_index = faiss.IndexHNSWFlat(embedding_dimension, 16)
+            hnsw_index.hnsw.efConstruction = 40
+            hnsw_index.hnsw.efSearch = 16
+            methodology_index = faiss.IndexIDMap(hnsw_index)
+            methodology_index.add_with_ids(vectors_array, np.array(ids)) # type: ignore
+            query_embedding = _create_methodology_embedding(embedding_model, user_input)
+            k = min(3, len(methodology_data))
+            PrettyOutput.print(f"Retrieving methodology...", OutputType.INFO)
+            distances, indices = methodology_index.search(
+                query_embedding.reshape(1, -1), k
+            ) # type: ignore
+
+            relevant_methodologies = {}
+            output_lines = []
+            for dist, idx in zip(distances[0], indices[0]):
+                if idx >= 0:
+                    similarity = 1.0 / (1.0 + float(dist))
+                    methodology = methodology_data[idx]
+                    output_lines.append(
+                        f"Methodology '{methodology['key']}' similarity: {similarity:.3f}"
+                    )
+                    if similarity >= 0.5:
+                        relevant_methodologies[methodology["key"]] = methodology["value"]
+            
+            if output_lines:
+                PrettyOutput.print("\n".join(output_lines), OutputType.INFO)
+                    
+            if relevant_methodologies:
+                return make_methodology_prompt(relevant_methodologies)
+        return make_methodology_prompt(data)
 
     except Exception as e:
         PrettyOutput.print(f"Error loading methodology: {str(e)}", OutputType.ERROR)
@@ -310,5 +577,28 @@ def load_methodology(user_input: str) -> str:
         PrettyOutput.print(f"Error trace: {traceback.format_exc()}", OutputType.INFO)
         return ""
     
-def no_embedding() -> bool:
-    return bool(os.environ.get("JARVIS_NO_EMBEDDING", ""))
+def is_auto_complete() -> bool:
+    return os.getenv('JARVIS_AUTO_COMPLETE', 'false') == 'true'
+
+def is_disable_codebase() -> bool:
+    return os.getenv('JARVIS_DISABLE_CODEBASE', 'false') == 'true'
+
+def user_confirm(tip: str, default: bool = True) -> bool:
+    """Prompt the user for confirmation.
+    
+    Args:
+        tip: The message to show to the user
+        default: The default response if user hits enter
+        
+    Returns:
+        bool: True if user confirmed, False otherwise
+    """
+    suffix = "[Y/n]" if default else "[y/N]"
+    ret = get_single_line_input(f"{tip} {suffix}: ")
+    return default if ret == "" else ret.lower() == "y"
+
+def get_file_line_count(filename: str) -> int:
+    try:
+        return len(open(filename, "r", encoding="utf-8").readlines())
+    except Exception as e:
+        return 0