jarvis-ai-assistant 0.1.108__py3-none-any.whl → 0.1.109__py3-none-any.whl

jarvis/utils.py

@@ -4,7 +4,7 @@ import time
 import os
 from enum import Enum
 from datetime import datetime
-from typing import Any, Dict, Optional
+from typing import Any, Dict, List, Optional
 import colorama
 from colorama import Fore, Style as ColoramaStyle
 import numpy as np
@@ -12,13 +12,14 @@ 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 tqdm import tqdm
 from transformers import AutoModelForSequenceClassification, AutoTokenizer
 import torch
 import yaml
 import faiss
 from pygments.lexers import guess_lexer
 from pygments.util import ClassNotFound
-
+import psutil
 from rich.console import Console
 from rich.theme import Theme
 from rich.panel import Panel
@@ -163,7 +164,7 @@ class PrettyOutput:
         return formatted
 
     @staticmethod
-    def print(text: str, output_type: OutputType, timestamp: bool = True, lang: Optional[str] = None):
+    def print(text: str, output_type: OutputType, timestamp: bool = True, lang: Optional[str] = None, traceback: bool = False):
         """Print formatted output using rich console"""
         # Get formatted header
         lang = lang if lang is not None else PrettyOutput._detect_language(text, default_lang='markdown')
@@ -176,7 +177,7 @@ class PrettyOutput:
         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:
+        if traceback or output_type == OutputType.ERROR:
             console.print_exception()
 
     @staticmethod
@@ -370,6 +371,7 @@ def find_git_root(dir="."):
     return ret
 
 def has_uncommitted_changes():
+    os.system("git add .")
     # Check working directory changes
     working_changes = os.popen("git diff --exit-code").read().strip() != ""
     # Check staged changes
@@ -398,6 +400,26 @@ def load_embedding_model():
     
     return embedding_model
 
+def load_tokenizer():
+    """Load tokenizer"""
+    model_name = "gpt2"
+    cache_dir = os.path.expanduser("~/.cache/huggingface/hub")
+    
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_name,
+            cache_dir=cache_dir,
+            local_files_only=True
+        )
+    except Exception as e:
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_name,
+            cache_dir=cache_dir,
+            local_files_only=False
+        )
+    
+    return tokenizer
+
 def load_rerank_model():
     """Load reranking model"""
     model_name = "BAAI/bge-reranker-v2-m3"
@@ -443,21 +465,21 @@ 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_length = get_max_context_length()
     threshold = max_length * 0.8
-    total_chars = 0
+    total_tokens = 0
     
     for file_path in files:
         try:
             with open(file_path, 'r', encoding='utf-8') as f:
                 content = f.read()
-                total_chars += len(content)
+                total_tokens += get_context_token_count(content)
                 
-                if total_chars > threshold:
+                if total_tokens > threshold:
                     return True
         except Exception as e:
             PrettyOutput.print(f"Failed to read file {file_path}: {e}", OutputType.WARNING)
             continue
             
-    return total_chars > threshold
+    return total_tokens > threshold
 
 
 
@@ -586,6 +608,72 @@ def get_file_line_count(filename: str) -> int:
     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 initialized: {torch.cuda.get_device_name(0)}\n"
+                f"Device Memory: {gpu_mem / 1024**3:.1f}GB\n"
+                f"Shared Memory: {config['shared_memory'] / 1024**3:.1f}GB", 
+                output_type=OutputType.SUCCESS
+            )
+        else:
+            PrettyOutput.print("No GPU available, using CPU mode", output_type=OutputType.WARNING)
+    except Exception as e:
+        PrettyOutput.print(f"GPU initialization failed: {str(e)}", output_type=OutputType.WARNING)
+        
+    return config
+
+
+def get_embedding(embedding_model: Any, text: str) -> np.ndarray:
+    """Get the vector representation of the text"""
+    embedding = embedding_model.encode(text, 
+                                        normalize_embeddings=True,
+                                        show_progress_bar=False)
+    return np.array(embedding, dtype=np.float32)
+
+def get_embedding_batch(embedding_model: Any, texts: List[str]) -> np.ndarray:
+    """Get embeddings for a batch of texts efficiently"""
+    try:
+        all_vectors = []
+        for text in texts:
+            vectors = get_embedding_with_chunks(embedding_model, text)
+            all_vectors.extend(vectors)
+        return np.vstack(all_vectors)
+    except Exception as e:
+        PrettyOutput.print(f"Batch embedding failed: {str(e)}", OutputType.ERROR)
+        return np.zeros((0, embedding_model.get_sentence_embedding_dimension()), dtype=np.float32)
+
+
+    
 def get_max_context_length():
     return int(os.getenv('JARVIS_MAX_CONTEXT_LENGTH', '131072'))  # 默认128k
     
@@ -614,10 +702,7 @@ def get_min_paragraph_length() -> int:
     return int(os.getenv('JARVIS_MIN_PARAGRAPH_LENGTH', '50'))
 
 def get_max_paragraph_length() -> int:
-    return int(os.getenv('JARVIS_MAX_PARAGRAPH_LENGTH', '1000'))
-
-def get_context_window() -> int:
-    return int(os.getenv('JARVIS_CONTEXT_WINDOW', '5'))
+    return int(os.getenv('JARVIS_MAX_PARAGRAPH_LENGTH', '12800'))
 
 def get_shell_name() -> str:
     return os.getenv('SHELL', 'bash')
@@ -645,3 +730,57 @@ def get_cheap_platform_name() -> str:
 
 def get_cheap_model_name() -> str:
     return os.getenv('JARVIS_CHEAP_MODEL', os.getenv('JARVIS_MODEL', 'kimi'))
+
+def split_text_into_chunks(text: str, max_length: int = 512) -> List[str]:
+    """Split text into chunks with overlapping windows"""
+    chunks = []
+    start = 0
+    while start < len(text):
+        end = start + max_length
+        # Find the nearest sentence boundary
+        if end < len(text):
+            while end > start and text[end] not in {'.', '!', '?', '\n'}:
+                end -= 1
+            if end == start:  # No punctuation found, hard cut
+                end = start + max_length
+        chunk = text[start:end]
+        chunks.append(chunk)
+        # Overlap 20% of the window
+        start = end - int(max_length * 0.2)
+    return chunks
+
+def get_embedding_with_chunks(embedding_model: Any, text: str) -> List[np.ndarray]:
+    """Get embeddings for text chunks"""
+    chunks = split_text_into_chunks(text, 512)
+    if not chunks:
+        return []
+    
+    vectors = []
+    for chunk in chunks:
+        vector = get_embedding(embedding_model, chunk)
+        vectors.append(vector)
+    return vectors
+
+
+def get_context_token_count(text: str) -> int:
+    """Get the token count of the text using the tokenizer
+    
+    Args:
+        text: The input text to count tokens for
+        
+    Returns:
+        int: The number of tokens in the text
+    """
+    try:
+        # Use a fast tokenizer that's good at general text
+        tokenizer = load_tokenizer()
+        tokens = tokenizer.encode(text)
+        return len(tokens)
+        
+    except Exception as e:
+        PrettyOutput.print(f"Error counting tokens: {str(e)}", OutputType.WARNING)
+        # Fallback to rough character-based estimate
+        return len(text) // 4  # Rough estimate of 4 chars per token
+
+
+

{jarvis_ai_assistant-0.1.108.dist-info → jarvis_ai_assistant-0.1.109.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: jarvis-ai-assistant
-Version: 0.1.108
+Version: 0.1.109
 Summary: Jarvis: An AI assistant that uses tools to interact with the system
 Home-page: https://github.com/skyfireitdiy/Jarvis
 Author: skyfire
@@ -58,8 +58,9 @@ Requires-Dist: yaspin>=2.5.0
 Requires-Dist: rich>=13.3.1
 Requires-Dist: pygments>=2.15.0
 Requires-Dist: fuzzywuzzy>=0.18.0
-Requires-Dist: python-Levenshtein>=0.26.1
-Requires-Dist: jedi>=0.19.0
+Requires-Dist: python-Levenshtein>=0.25.0
+Requires-Dist: jedi>=0.17.2
+Requires-Dist: psutil>=7.0.0
 Provides-Extra: dev
 Requires-Dist: pytest; extra == "dev"
 Requires-Dist: black; extra == "dev"
@@ -161,6 +162,14 @@ Jarvis supports configuration through environment variables that can be set in t
 | OYI_API_KEY | API key for OYI platform | - | Required for OYI |
 | OLLAMA_API_BASE | Base URL for Ollama API | http://localhost:11434 | No |
 
+## Minimal Configuration (Example with OpenAI-compatible Interface)
+
+```bash
+JARVIS_PLATFORM=openai
+JARVIS_MODEL=deepseek-chat
+OPENAI_API_KEY=your_openai_api_key
+OPENAI_API_BASE=https://api.deepseek.com/v1
+```
 
 ## 🎯 Usage
 

{jarvis_ai_assistant-0.1.108.dist-info → jarvis_ai_assistant-0.1.109.dist-info}/RECORD

@@ -1,13 +1,13 @@
-jarvis/__init__.py,sha256=hCWwxdx-HHSIwE8vf6Yfn-___pFPgINrCmQcUL8j_XI,51
-jarvis/agent.py,sha256=7FDrJc2_JlY9u7TRfeHKZRQ0PrY04r-0w4H64eGcbUM,22626
-jarvis/utils.py,sha256=0w1rYsSovS7vgbHNdfdzpo3zSb3y-KWM7RvYMqBhDnM,22086
+jarvis/__init__.py,sha256=CB3MYpgrHjKZT-LJIoysQYNgLyXYOJSKmxopMqJsAmU,51
+jarvis/agent.py,sha256=x6LmdW5k50nhAK6hdb8V1Vvkt64szwrnaewAto8murg,22690
+jarvis/utils.py,sha256=_GQY1UlxzhKM0DwdnuTUeSmaae8dV83SffsjMPPXFM0,26889
 jarvis/jarvis_code_agent/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-jarvis/jarvis_code_agent/code_agent.py,sha256=DDKqM4m235Ryd8aMHbHKeC0ULdNLF6JPW__FKbI-Gk0,10748
-jarvis/jarvis_code_agent/file_select.py,sha256=KNxalhepCM2e-V__ca8ErmbXSXHP_1xmd0UEVWUXic8,8083
-jarvis/jarvis_code_agent/patch.py,sha256=bbNB8k8mebjPVsNdI8aT3oOyjLyAhUQbKmX54tyziDk,4034
-jarvis/jarvis_code_agent/relevant_files.py,sha256=PxSKQyHfCe6878bDqP6XyQd_jwcvNK4a9YKTfpLImRI,6160
+jarvis/jarvis_code_agent/code_agent.py,sha256=nigsmCK6D2z0dFU_1HFNYEvXr3lWdl0rm6p4VgiOk6o,5980
+jarvis/jarvis_code_agent/file_select.py,sha256=1kOVRLPS1GZcDyGpCW9hOPbfCEwF8f0-qVPaRZPHzoM,8154
+jarvis/jarvis_code_agent/patch.py,sha256=bOhegGKs4JEmJJOZfUlmwzGI6kakMyi2Q62HADJ7Npk,4594
+jarvis/jarvis_code_agent/relevant_files.py,sha256=Q4nI45zuyWt5aKuc4OR7-a6UbOXOym3oEzQJvqxkF8Q,946
 jarvis/jarvis_codebase/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-jarvis/jarvis_codebase/main.py,sha256=UfFLwOZQoOKSKnpk_Ive1azEYO4-Cm8fICxULGAvqUU,39889
+jarvis/jarvis_codebase/main.py,sha256=7SwL0QPeojsC95SIlbKjq9FcvKnvPBbtH1kv3SjPUBY,39701
 jarvis/jarvis_lsp/base.py,sha256=_7pdbMKjdtYBW0DsRbjIodDHM3J7df-YgXHejN_WIrU,4490
 jarvis/jarvis_lsp/cpp.py,sha256=F7Zo3BErkvtWS1_H9zQO83pX_FUmnijux-2SjhWzKCE,4985
 jarvis/jarvis_lsp/go.py,sha256=p8LULiFdq4qjDYQzXFlzH0-FQZ3IyfiwN_sbO9i0L_A,5310
@@ -25,7 +25,7 @@ jarvis/jarvis_platform/registry.py,sha256=9QLoihcnkYckrCzgNnlTqaLn_z_HMhaxMSyUNb
 jarvis/jarvis_platform_manager/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 jarvis/jarvis_platform_manager/main.py,sha256=17607aNAStqJ1sOQLTGi6Tnv-cIQme_r5YvbB_S3enc,4985
 jarvis/jarvis_rag/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-jarvis/jarvis_rag/main.py,sha256=3BpRz4XrdSLUmH1IENXoKAFqjMxckwtWrGrhdUxOsNw,34626
+jarvis/jarvis_rag/main.py,sha256=ZdZZnOyQX8q1mNUeMymb9qg1l77JnG6WBz_KDov0NLQ,31366
 jarvis/jarvis_smart_shell/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 jarvis/jarvis_smart_shell/main.py,sha256=VdUR-x932OccEwU0pcQM_pb_I4yfrAutE3hfm6jf5es,3955
 jarvis/jarvis_tools/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -51,12 +51,12 @@ jarvis/jarvis_tools/methodology.py,sha256=RFqcVjKuj8ESGmNYcQz_HyphsitDvF3XtqgGaq
 jarvis/jarvis_tools/rag.py,sha256=2fQHqc4bw8JM-OxGTsHobLIOTo8Mip3rdtJCmAoY8XU,4952
 jarvis/jarvis_tools/read_code.py,sha256=5DGmeXTgumAiG0RP1xB4sF4NdmBm5BEGjRRlIBzjGnQ,4002
 jarvis/jarvis_tools/read_webpage.py,sha256=JCReSXhkDHDkQ606sZYIKG1Itlprjpmu1sSbF-Ed-jI,2478
-jarvis/jarvis_tools/registry.py,sha256=mkAQ1NDPwDy1ESAaAFnzSxAJRrhce3NO3E7cWkb-quA,11732
+jarvis/jarvis_tools/registry.py,sha256=mYx7q78BQ9El68c6DPbbeMyRp0Dr9GnZoozjhSViLPo,11859
 jarvis/jarvis_tools/search.py,sha256=PLSSNETyajpqDoStCTfkoy-D41IMNudTuVzonMlT6Aw,9225
 jarvis/jarvis_tools/select_code_files.py,sha256=bjJGwCNw0Ue_8jW60K1gcy1rUgKqoHihicu5SS58WNk,1890
-jarvis_ai_assistant-0.1.108.dist-info/LICENSE,sha256=AGgVgQmTqFvaztRtCAXsAMryUymB18gZif7_l2e1XOg,1063
-jarvis_ai_assistant-0.1.108.dist-info/METADATA,sha256=OWN4NKRNvFiqIc7mL9XN4-a488mjwbV_TCUX5bhyFkQ,14153
-jarvis_ai_assistant-0.1.108.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-jarvis_ai_assistant-0.1.108.dist-info/entry_points.txt,sha256=UYj4FYvOH8jJ0GgCJTA_TAmJ3wvikos-hUVbCwt_KOc,480
-jarvis_ai_assistant-0.1.108.dist-info/top_level.txt,sha256=1BOxyWfzOP_ZXj8rVTDnNCJ92bBGB0rwq8N1PCpoMIs,7
-jarvis_ai_assistant-0.1.108.dist-info/RECORD,,
+jarvis_ai_assistant-0.1.109.dist-info/LICENSE,sha256=AGgVgQmTqFvaztRtCAXsAMryUymB18gZif7_l2e1XOg,1063
+jarvis_ai_assistant-0.1.109.dist-info/METADATA,sha256=pxorU-ZxDXj6dNQOZfIh285QvRa9hUxXmKGVbgotBTc,14392
+jarvis_ai_assistant-0.1.109.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+jarvis_ai_assistant-0.1.109.dist-info/entry_points.txt,sha256=UYj4FYvOH8jJ0GgCJTA_TAmJ3wvikos-hUVbCwt_KOc,480
+jarvis_ai_assistant-0.1.109.dist-info/top_level.txt,sha256=1BOxyWfzOP_ZXj8rVTDnNCJ92bBGB0rwq8N1PCpoMIs,7
+jarvis_ai_assistant-0.1.109.dist-info/RECORD,,