jarvis-ai-assistant 0.1.124__py3-none-any.whl → 0.1.125__py3-none-any.whl

jarvis/jarvis_utils/methodology.py

@@ -0,0 +1,128 @@
+"""
+Methodology Management Module
+This module provides utilities for loading and searching methodologies.
+It includes functions for:
+- Creating methodology embeddings
+- Loading and processing methodology data
+- Building and searching methodology index
+- Generating methodology prompts
+"""
+import os
+import yaml
+import numpy as np
+import faiss
+from typing import Dict, Any, List
+from jarvis.jarvis_utils.output import PrettyOutput, OutputType
+from jarvis.jarvis_utils.embedding import load_embedding_model
+from jarvis.jarvis_utils.config import dont_use_local_model
+def _create_methodology_embedding(embedding_model: Any, methodology_text: str) -> np.ndarray:
+    """
+    Create embedding vector for methodology text.
+    
+    Args:
+        embedding_model: The embedding model to use
+        methodology_text: The text to create embedding for
+        
+    Returns:
+        np.ndarray: The embedding vector
+    """
+    try:
+        # Truncate long text
+        max_length = 512
+        text = ' '.join(methodology_text.split()[:max_length])
+        
+        # 使用sentence_transformers模型获取嵌入向量
+        embedding = embedding_model.encode([text], 
+                                          convert_to_tensor=True,
+                                          normalize_embeddings=True)
+        vector = np.array(embedding.cpu().numpy(), dtype=np.float32)
+        return vector[0]  # Return first vector, because we only encoded one text
+    except Exception as e:
+        PrettyOutput.print(f"创建方法论嵌入向量失败: {str(e)}", OutputType.ERROR)
+        return np.zeros(1536, dtype=np.float32)
+def make_methodology_prompt(data: Dict[str, str]) -> str:
+    """
+    Generate a formatted prompt from methodology data.
+    
+    Args:
+        data: Dictionary of methodology data
+        
+    Returns:
+        str: Formatted prompt string
+    """
+    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
+def load_methodology(user_input: str) -> str:
+    """
+    Load methodology and build vector index for similarity search.
+    
+    Args:
+        user_input: The input text to search methodologies for
+        
+    Returns:
+        str: Relevant methodology prompt or empty string if no methodology found
+    """
+    PrettyOutput.print("加载方法论...", OutputType.PROGRESS)
+    user_jarvis_methodology = os.path.expanduser("~/.jarvis/methodology")
+    if not os.path.exists(user_jarvis_methodology):
+        return ""
+    
+    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: List[Dict[str, str]] = []
+        vectors: List[np.ndarray] = []
+        ids: List[int] = []
+        # Get embedding model
+        embedding_model = load_embedding_model()
+        
+        # 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"检索方法论...", 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"加载方法论失败: {str(e)}", OutputType.ERROR)
+        return ""

jarvis/jarvis_utils/output.py

@@ -0,0 +1,235 @@
+"""
+Output Formatting Module
+This module provides rich text formatting and display utilities for the Jarvis system.
+It includes:
+- OutputType enum for categorizing different types of output
+- PrettyOutput class for formatting and displaying styled output
+- Syntax highlighting support for various programming languages
+- Panel-based display for structured output
+"""
+from enum import Enum
+from datetime import datetime
+from typing import Optional
+from rich.panel import Panel
+from rich.box import HEAVY
+from rich.text import Text
+from rich.syntax import Syntax
+from rich.style import Style as RichStyle
+from pygments.lexers import guess_lexer
+from pygments.util import ClassNotFound
+from .globals import console, get_agent_list
+class OutputType(Enum):
+    """
+    Enumeration of output types for categorizing and styling different types of messages.
+    
+    Attributes:
+        SYSTEM: AI assistant message
+        CODE: Code related output
+        RESULT: Tool execution result
+        ERROR: Error information
+        INFO: System prompt
+        PLANNING: Task planning
+        PROGRESS: Execution progress
+        SUCCESS: Success information
+        WARNING: Warning information
+        DEBUG: Debug information
+        USER: User input
+        TOOL: Tool call
+    """
+    SYSTEM = "SYSTEM"
+    CODE = "CODE"
+    RESULT = "RESULT"
+    ERROR = "ERROR"
+    INFO = "INFO"
+    PLANNING = "PLANNING"
+    PROGRESS = "PROGRESS"
+    SUCCESS = "SUCCESS"
+    WARNING = "WARNING"
+    DEBUG = "DEBUG"
+    USER = "USER"
+    TOOL = "TOOL"
+class PrettyOutput:
+    """
+    Class for formatting and displaying rich text output using the rich library.
+    
+    Provides methods for:
+    - Formatting different types of output with appropriate styling
+    - Syntax highlighting for code blocks
+    - Panel-based display for structured content
+    - Stream output for progressive display
+    """
+    # Icons for different output types
+    _ICONS = {
+        OutputType.SYSTEM: "🤖",
+        OutputType.CODE: "📝",
+        OutputType.RESULT: "✨",
+        OutputType.ERROR: "❌",
+        OutputType.INFO: "ℹ️",
+        OutputType.PLANNING: "📋",
+        OutputType.PROGRESS: "⏳",
+        OutputType.SUCCESS: "✅",
+        OutputType.WARNING: "⚠️",
+        OutputType.DEBUG: "🔍",
+        OutputType.USER: "👤",
+        OutputType.TOOL: "🔧",
+    }
+    # Language mapping for syntax highlighting
+    _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 _detect_language(text: str, default_lang: str = 'markdown') -> str:
+        """
+        Detect the programming language of the given text.
+        
+        Args:
+            text: The text to analyze
+            default_lang: Default language if detection fails
+            
+        Returns:
+            str: Detected language name
+        """
+        try:
+            lexer = guess_lexer(text)
+            detected_lang = lexer.name
+            return PrettyOutput._lang_map.get(detected_lang, default_lang)
+        except (ClassNotFound, Exception):
+            return default_lang
+    @staticmethod
+    def _format(output_type: OutputType, timestamp: bool = True) -> Text:
+        """
+        Format the output header with timestamp and icon.
+        
+        Args:
+            output_type: Type of output
+            timestamp: Whether to include timestamp
+            
+        Returns:
+            Text: Formatted rich Text object
+        """
+        formatted = Text()
+        if timestamp:
+            formatted.append(f"[{datetime.now().strftime('%H:%M:%S')}][{output_type.value}]", style=output_type.value)
+        agent_info = get_agent_list()
+        if agent_info:
+            formatted.append(f"[{agent_info}]", style="blue")
+        icon = PrettyOutput._ICONS.get(output_type, "")
+        formatted.append(f" {icon} ", style=output_type.value)
+        return formatted
+    @staticmethod
+    def print(text: str, output_type: OutputType, timestamp: bool = True, lang: Optional[str] = None, traceback: bool = False):
+        """
+        Print formatted output with styling and syntax highlighting.
+        
+        Args:
+            text: The text content to print
+            output_type: The type of output (affects styling)
+            timestamp: Whether to show timestamp
+            lang: Language for syntax highlighting
+            traceback: Whether to show traceback for errors
+        """
+        styles = {
+            OutputType.SYSTEM: RichStyle(color="bright_cyan", bgcolor="#1a1a1a", frame=True, meta={"icon": "🤖"}),
+            OutputType.CODE: RichStyle(color="green", bgcolor="#1a1a1a", frame=True, meta={"icon": "📝"}),
+            OutputType.RESULT: RichStyle(color="bright_blue", bgcolor="#1a1a1a", frame=True, meta={"icon": "✨"}),
+            OutputType.ERROR: RichStyle(color="red", frame=True, bgcolor="dark_red", meta={"icon": "❌"}),
+            OutputType.INFO: RichStyle(color="gold1", frame=True, bgcolor="grey11", meta={"icon": "ℹ️"}),
+            OutputType.PLANNING: RichStyle(color="purple", bold=True, frame=True, meta={"icon": "📋"}),
+            OutputType.PROGRESS: RichStyle(color="white", encircle=True, frame=True, meta={"icon": "⏳"}),
+            OutputType.SUCCESS: RichStyle(color="bright_green", bold=True, strike=False, meta={"icon": "✅"}),
+            OutputType.WARNING: RichStyle(color="yellow", bold=True, blink2=True, bgcolor="dark_orange", meta={"icon": "⚠️"}),
+            OutputType.DEBUG: RichStyle(color="grey58", dim=True, conceal=True, meta={"icon": "🔍"}),
+            OutputType.USER: RichStyle(color="spring_green2", frame=True, meta={"icon": "👤"}),
+            OutputType.TOOL: RichStyle(color="dark_sea_green4", bgcolor="grey19", frame=True, meta={"icon": "🔧"}),
+        }
+        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", word_wrap=True)
+        panel = Panel(
+            content,
+            style=styles[output_type],
+            border_style=styles[output_type],
+            title=header,
+            title_align="left",
+            padding=(0, 0),
+            highlight=True,
+            box=HEAVY,
+        )
+        console.print(panel)
+        if traceback or output_type == OutputType.ERROR:
+            console.print_exception()
+    @staticmethod
+    def section(title: str, output_type: OutputType = OutputType.INFO):
+        """
+        Print a section title in a styled panel.
+        
+        Args:
+            title: The section title text
+            output_type: The type of output (affects styling)
+        """
+        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 without line break.
+        
+        Args:
+            text: The text to print
+        """
+        style = PrettyOutput._get_style(OutputType.SYSTEM)
+        console.print(text, style=style, end="")
+    @staticmethod
+    def print_stream_end():
+        """
+        End stream output with line break.
+        """
+        end_style = PrettyOutput._get_style(OutputType.SUCCESS)
+        console.print("\n", style=end_style)
+        console.file.flush()
+    @staticmethod
+    def _get_style(output_type: OutputType) -> RichStyle:
+        """
+        Get pre-defined RichStyle for output type.
+        
+        Args:
+            output_type: The output type to get style for
+            
+        Returns:
+            RichStyle: The corresponding style
+        """
+        return console.get_style(output_type.value)

jarvis/jarvis_utils/utils.py

@@ -0,0 +1,150 @@
+import os
+import time
+import hashlib
+from pathlib import Path
+from typing import Dict, List
+import psutil
+from jarvis.jarvis_utils.config import get_max_token_count
+from jarvis.jarvis_utils.embedding import get_context_token_count
+from jarvis.jarvis_utils.input import get_single_line_input
+from jarvis.jarvis_utils.output import PrettyOutput, OutputType
+def init_env():
+    """Initialize environment variables from ~/.jarvis/env file.
+    
+    Creates the .jarvis directory if it doesn't exist and loads environment variables
+    from the env file. Handles file reading errors gracefully.
+    """
+    jarvis_dir = Path.home() / ".jarvis"
+    env_file = jarvis_dir / "env"
+    
+    # Check if ~/.jarvis directory exists
+    if not jarvis_dir.exists():
+        jarvis_dir.mkdir(parents=True)
+    if env_file.exists():
+        try:
+            with open(env_file, "r", encoding="utf-8") as f:
+                for line in f:
+                    line = line.strip()
+                    if line and not line.startswith(("#", ";")):
+                        try:
+                            key, value = line.split("=", 1)
+                            os.environ[key.strip()] = value.strip().strip("'").strip('"')
+                        except ValueError:
+                            continue
+        except Exception as e:
+            PrettyOutput.print(f"警告: 读取 {env_file} 失败: {e}", OutputType.WARNING)
+def while_success(func, sleep_time: float = 0.1):
+    while True:
+        try:
+            return func()
+        except Exception as e:
+            PrettyOutput.print(f"执行失败: {str(e)}, 等待 {sleep_time}s...", OutputType.ERROR)
+            time.sleep(sleep_time)
+            continue
+def while_true(func, sleep_time: float = 0.1):
+    """Loop execution function, until the function returns True"""
+    while True:
+        ret = func()
+        if ret:
+            break
+        PrettyOutput.print(f"执行失败, 等待 {sleep_time}s...", OutputType.WARNING)
+        time.sleep(sleep_time)
+    return ret
+def get_file_md5(filepath: str)->str:    
+    """Calculate the MD5 hash of a file's content.
+    
+    Args:
+        filepath: Path to the file to hash
+        
+    Returns:
+        str: MD5 hash of the file's content
+    """
+    return hashlib.md5(open(filepath, "rb").read(100*1024*1024)).hexdigest()
+def user_confirm(tip: str, default: bool = True) -> bool:
+    """Prompt the user for confirmation with a yes/no question.
+    
+    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:
+    """Count the number of lines in a file.
+    
+    Args:
+        filename: Path to the file to count lines for
+        
+    Returns:
+        int: Number of lines in the file, 0 if file cannot be read
+    """
+    try:
+        return len(open(filename, "r", encoding="utf-8").readlines())
+    except Exception as e:
+        return 0
+def init_gpu_config() -> Dict:
+    """Initialize GPU configuration based on available hardware.
+    
+    Returns:
+        Dict: GPU configuration including memory sizes and availability
+    """
+    config = {
+        "has_gpu": False,
+        "shared_memory": 0,
+        "device_memory": 0,
+        "memory_fraction": 0.8  # 默认使用80%的可用内存
+    }
+    
+    try:
+        import torch
+        if torch.cuda.is_available():
+            # 获取GPU信息
+            gpu_mem = torch.cuda.get_device_properties(0).total_memory
+            config["has_gpu"] = True
+            config["device_memory"] = gpu_mem
+            
+            # 估算共享内存 (通常是系统内存的一部分)
+            system_memory = psutil.virtual_memory().total
+            config["shared_memory"] = min(system_memory * 0.5, gpu_mem * 2)  # 取系统内存的50%或GPU内存的2倍中的较小值
+            
+            # 设置CUDA内存分配
+            torch.cuda.set_per_process_memory_fraction(config["memory_fraction"])
+            torch.cuda.empty_cache()
+            
+            PrettyOutput.print(
+                f"GPU已初始化: {torch.cuda.get_device_name(0)}\n"
+                f"设备内存: {gpu_mem / 1024**3:.1f}GB\n"
+                f"共享内存: {config['shared_memory'] / 1024**3:.1f}GB", 
+                output_type=OutputType.SUCCESS
+            )
+        else:
+            PrettyOutput.print("没有GPU可用, 使用CPU模式", output_type=OutputType.WARNING)
+    except Exception as e:
+        PrettyOutput.print(f"GPU初始化失败: {str(e)}", output_type=OutputType.WARNING)
+        
+    return config
+
+
+def is_long_context(files: list) -> bool:
+    """Check if the file list belongs to a long context (total characters exceed 80% of the maximum context length)"""
+    max_token_count = get_max_token_count()
+    threshold = max_token_count * 0.8
+    total_tokens = 0
+    
+    for file_path in files:
+        try:
+            with open(file_path, 'r', encoding='utf-8') as f:
+                content = f.read()
+                total_tokens += get_context_token_count(content)
+                
+                if total_tokens > threshold:
+                    return True
+        except Exception as e:
+            PrettyOutput.print(f"读取文件 {file_path} 失败: {e}", OutputType.WARNING)
+            continue
+            
+    return total_tokens > threshold