jarvis-ai-assistant 0.1.126__py3-none-any.whl → 0.1.129__py3-none-any.whl

jarvis/jarvis_codebase/main.py

@@ -4,6 +4,8 @@ import numpy as np
 import faiss
 from typing import List, Tuple, Optional, Dict
 
+from yaspin import yaspin
+
 from jarvis.jarvis_platform.registry import PlatformRegistry
 import concurrent.futures
 from concurrent.futures import ThreadPoolExecutor
@@ -21,7 +23,11 @@ from jarvis.jarvis_utils.utils import  get_file_md5, init_env, user_confirm
 
 class CodeBase:
     def __init__(self, root_dir: str):
-        init_env()
+        with yaspin(text="正在初始化环境...", color="cyan") as spinner:
+            init_env()
+            spinner.text = "环境初始化完成"
+            spinner.ok("✅")
+            
         self.root_dir = root_dir
         os.chdir(self.root_dir)
         self.thread_count = get_thread_count()
@@ -29,22 +35,28 @@ class CodeBase:
         self.index = None
             
         # 初始化数据目录
-        self.data_dir = os.path.join(self.root_dir, ".jarvis/codebase")
-        self.cache_dir = os.path.join(self.data_dir, "cache")
-        if not os.path.exists(self.cache_dir):
-            os.makedirs(self.cache_dir)
-            
+        with yaspin(text="正在初始化数据目录...", color="cyan") as spinner:
+            self.data_dir = os.path.join(self.root_dir, ".jarvis/codebase")
+            self.cache_dir = os.path.join(self.data_dir, "cache")
+            if not os.path.exists(self.cache_dir):
+                os.makedirs(self.cache_dir)
+            spinner.text = "数据目录初始化完成"
+            spinner.ok("✅")
+            
+        with yaspin("正在初始化嵌入模型...", color="cyan") as spinner:
         # 初始化嵌入模型
-        try:
-            self.embedding_model = load_embedding_model()
-            test_text = """This is a test text"""
-            self.embedding_model.encode([test_text], 
-                                     convert_to_tensor=True,
-                                     normalize_embeddings=True)
-            PrettyOutput.print("模型加载成功", output_type=OutputType.SUCCESS)
-        except Exception as e:
-            PrettyOutput.print(f"加载模型失败: {str(e)}", output_type=OutputType.ERROR)
-            raise
+            try:
+                self.embedding_model = load_embedding_model()
+                test_text = """This is a test text"""
+                self.embedding_model.encode([test_text], 
+                                        convert_to_tensor=True,
+                                        normalize_embeddings=True)
+                spinner.text = "嵌入模型初始化完成"
+                spinner.ok("✅")
+            except Exception as e:
+                spinner.text = "嵌入模型初始化失败"
+                spinner.fail("❌")
+                raise
             
         self.vector_dim = self.embedding_model.get_sentence_embedding_dimension()
         self.git_file_list = self.get_git_file_list()
@@ -55,7 +67,8 @@ class CodeBase:
         self.file_paths = []
         
         # 加载所有缓存文件
-        self._load_all_cache()
+        with spinner.hidden():
+            self._load_all_cache()
 
     def get_git_file_list(self):
         """Get the list of files in the git repository, excluding the .jarvis-codebase directory"""
@@ -72,10 +85,6 @@ class CodeBase:
 
     def make_description(self, file_path: str, content: str) -> str:
         model = PlatformRegistry.get_global_platform_registry().get_cheap_platform()
-        if self.thread_count > 1:
-            model.set_suppress_output(True)
-        else:
-            PrettyOutput.print(f"为 {file_path} 生成描述 ...", output_type=OutputType.PROGRESS)
         prompt = f"""请分析以下代码文件并生成详细描述。描述应包含：
 1. 文件整体功能描述
 2. 对每个全局变量、函数、类型定义、类、方法和其他代码元素的描述
@@ -114,50 +123,52 @@ class CodeBase:
 
     def _load_all_cache(self):
         """Load all cache files"""
-        try:
-            # 清空现有缓存和文件路径
-            self.vector_cache = {}
-            self.file_paths = []
-            vectors = []
-            
-            for cache_file in os.listdir(self.cache_dir):
-                if not cache_file.endswith('.cache'):
-                    continue
-                    
-                cache_path = os.path.join(self.cache_dir, cache_file)
-                try:
-                    with lzma.open(cache_path, 'rb') as f:
-                        cache_data = pickle.load(f)
-                        file_path = cache_data["path"]
-                        self.vector_cache[file_path] = cache_data
-                        self.file_paths.append(file_path)
-                        vectors.append(cache_data["vector"])
-                except Exception as e:
-                    PrettyOutput.print(f"加载缓存文件 {cache_file} 失败: {str(e)}", 
-                                     output_type=OutputType.WARNING)
-                    continue
-            
-            if vectors:
-                # 重建索引
-                vectors_array = np.vstack(vectors)
-                hnsw_index = faiss.IndexHNSWFlat(self.vector_dim, 16)
-                hnsw_index.hnsw.efConstruction = 40
-                hnsw_index.hnsw.efSearch = 16
-                self.index = faiss.IndexIDMap(hnsw_index)
-                self.index.add_with_ids(vectors_array, np.array(range(len(vectors)))) # type: ignore
+        with yaspin(text="正在加载缓存文件...", color="cyan") as spinner:
+            try:
+                # 清空现有缓存和文件路径
+                self.vector_cache = {}
+                self.file_paths = []
+                vectors = []
                 
-                PrettyOutput.print(f"加载 {len(self.vector_cache)} 个向量缓存并重建索引", 
-                                 output_type=OutputType.INFO)
-            else:
-                self.index = None
-                PrettyOutput.print("没有找到有效的缓存文件", output_type=OutputType.WARNING)
+                for cache_file in os.listdir(self.cache_dir):
+                    if not cache_file.endswith('.cache'):
+                        continue
+                        
+                    cache_path = os.path.join(self.cache_dir, cache_file)
+                    try:
+                        with lzma.open(cache_path, 'rb') as f:
+                            cache_data = pickle.load(f)
+                            file_path = cache_data["path"]
+                            self.vector_cache[file_path] = cache_data
+                            self.file_paths.append(file_path)
+                            vectors.append(cache_data["vector"])
+                            spinner.write(f"✅ 加载缓存文件成功 {file_path}")
+                    except Exception as e:
+                        spinner.write(f"❌ 加载缓存文件失败 {cache_file} {str(e)}")
+                        continue
                 
-        except Exception as e:
-            PrettyOutput.print(f"加载缓存目录失败: {str(e)}", 
-                             output_type=OutputType.WARNING)
-            self.vector_cache = {}
-            self.file_paths = []
-            self.index = None
+                if vectors:
+                    # 重建索引
+                    vectors_array = np.vstack(vectors)
+                    hnsw_index = faiss.IndexHNSWFlat(self.vector_dim, 16)
+                    hnsw_index.hnsw.efConstruction = 40
+                    hnsw_index.hnsw.efSearch = 16
+                    self.index = faiss.IndexIDMap(hnsw_index)
+                    self.index.add_with_ids(vectors_array, np.array(range(len(vectors)))) # type: ignore
+                    
+                    spinner.text = f"加载 {len(self.vector_cache)} 个向量缓存并重建索引"
+                    spinner.ok("✅")
+                else:
+                    self.index = None
+                    spinner.text = "没有找到有效的缓存文件"
+                    spinner.ok("✅")
+                    
+            except Exception as e:
+                spinner.text = f"加载缓存目录失败: {str(e)}"
+                spinner.fail("❌")
+                self.vector_cache = {}
+                self.file_paths = []
+                self.index = None
 
     def cache_vector(self, file_path: str, vector: np.ndarray, description: str):
         """Cache the vector representation of a file"""
@@ -450,7 +461,6 @@ Content: {content}
                 # If force is True, continue directly
                 if not force:
                     if not user_confirm("重建索引?", False):
-                        PrettyOutput.print("取消重建索引", output_type=OutputType.INFO)
                         return
                 
                 # Clean deleted files
@@ -464,7 +474,7 @@ Content: {content}
                 files_to_process = new_files + modified_files
                 processed_files = []
                 
-                with tqdm(total=len(files_to_process), desc="Processing files") as pbar:
+                with yaspin(text="正在处理文件...", color="cyan") as spinner:
                     # Use a thread pool to process files
                     with ThreadPoolExecutor(max_workers=self.thread_count) as executor:
                         # Submit all tasks
@@ -480,16 +490,18 @@ Content: {content}
                                 result = future.result()
                                 if result:
                                     processed_files.append(result)
+                                    spinner.write(f"✅ 处理文件成功 {file}")
                             except Exception as e:
-                                PrettyOutput.print(f"Failed to process file {file}: {str(e)}", 
-                                                output_type=OutputType.ERROR)
-                            pbar.update(1)
+                                spinner.write(f"❌ 处理文件失败 {file}: {str(e)}")
+                                
+                    spinner.text = f"处理完成"
+                    spinner.ok("✅")
 
                 if processed_files:
-                    PrettyOutput.print("重建向量数据库...", output_type=OutputType.INFO)
-                    self.gen_vector_db_from_cache()
-                    PrettyOutput.print(f"成功生成了 {len(processed_files)} 个文件的索引", 
-                                    output_type=OutputType.SUCCESS)
+                    with yaspin(text="重建向量数据库...", color="cyan") as spinner:
+                        self.gen_vector_db_from_cache()
+                        spinner.text = f"成功生成了 {len(processed_files)} 个文件的索引"
+                        spinner.ok("✅")
             else:
                 PrettyOutput.print("没有检测到文件变化, 不需要重建索引", output_type=OutputType.INFO)
                 
@@ -540,58 +552,58 @@ Content: {content}
         """
         if not initial_results:
             return []
-            
-        try:
-            PrettyOutput.print(f"Picking results ...", output_type=OutputType.INFO)
-            
-            # Maximum content length per batch
-            max_batch_length = self.max_token_count - 1000  # Reserve space for prompt
-            max_file_length = max_batch_length // 3  # Limit individual file size
-            
-            # Process files in batches
-            all_selected_files = []
-            current_batch = []
-            current_token_count = 0
-            
-            for path in initial_results:
-                try:
-                    content = open(path, "r", encoding="utf-8").read()
-                    # Truncate large files
-                    if get_context_token_count(content) > max_file_length:
-                        PrettyOutput.print(f"Truncating large file: {path}", OutputType.WARNING)
-                        content = content[:max_file_length] + "\n... (content truncated)"
-                    
-                    file_info = f"File: {path}\nContent: {content}\n\n"
-                    tokens_count = get_context_token_count(file_info)
-                    
-                    # If adding this file would exceed batch limit
-                    if current_token_count + tokens_count > max_batch_length:
-                        # Process current batch
-                        if current_batch:
-                            selected = self._process_batch('\n'.join(query), current_batch)
-                            all_selected_files.extend(selected)
-                        # Start new batch
-                        current_batch = [file_info]
-                        current_token_count = tokens_count
-                    else:
-                        current_batch.append(file_info)
-                        current_token_count += tokens_count
+        with yaspin(text="正在筛选结果...", color="cyan") as spinner:
+            try:
+                # Maximum content length per batch
+                max_batch_length = self.max_token_count - 1000  # Reserve space for prompt
+                max_file_length = max_batch_length // 3  # Limit individual file size
+                
+                # Process files in batches
+                all_selected_files = []
+                current_batch = []
+                current_token_count = 0
+                
+                for path in initial_results:
+                    try:
+                        content = open(path, "r", encoding="utf-8").read()
+                        # Truncate large files
+                        if get_context_token_count(content) > max_file_length:
+                            spinner.write(f"❌ 截断大文件: {path}")
+                            content = content[:max_file_length] + "\n... (content truncated)"
                         
-                except Exception as e:
-                    PrettyOutput.print(f"读取 {path} 失败: {str(e)}", OutputType.ERROR)
-                    continue
-            
-            # Process final batch
-            if current_batch:
-                selected = self._process_batch('\n'.join(query), current_batch)
-                all_selected_files.extend(selected)
-            
-            # Convert set to list and maintain original order
-            return all_selected_files
+                        file_info = f"File: {path}\nContent: {content}\n\n"
+                        tokens_count = get_context_token_count(file_info)
+                        
+                        # If adding this file would exceed batch limit
+                        if current_token_count + tokens_count > max_batch_length:
+                            # Process current batch
+                            if current_batch:
+                                selected = self._process_batch('\n'.join(query), current_batch)
+                                all_selected_files.extend(selected)
+                            # Start new batch
+                            current_batch = [file_info]
+                            current_token_count = tokens_count
+                        else:
+                            current_batch.append(file_info)
+                            current_token_count += tokens_count
+                            
+                    except Exception as e:
+                        spinner.write(f"❌ 读取 {path} 失败: {str(e)}")
+                        continue
+                
+                # Process final batch
+                if current_batch:
+                    selected = self._process_batch('\n'.join(query), current_batch)
+                    all_selected_files.extend(selected)
+                
+                spinner.write("✅ 结果筛选完成")
+                # Convert set to list and maintain original order
+                return all_selected_files
 
-        except Exception as e:
-            PrettyOutput.print(f"选择失败: {str(e)}", OutputType.ERROR)
-            return [{"file": f, "reason": "" } for f in initial_results]
+            except Exception as e:
+                spinner.text = f"选择失败: {str(e)}"
+                spinner.fail("❌")
+                return [{"file": f, "reason": "" } for f in initial_results]
             
     def _process_batch(self, query: str, files_info: List[str]) -> List[Dict[str, str]]:
         """Process a batch of files"""
@@ -630,7 +642,6 @@ Content: {content}
 
         # Use a large model to evaluate
         model = PlatformRegistry.get_global_platform_registry().get_normal_platform()
-        model.set_suppress_output(True)
         response = model.chat_until_success(prompt)
 
         # Parse the response
@@ -656,7 +667,6 @@ Content: {content}
             List[str]: The query variants list
         """
         model = PlatformRegistry.get_global_platform_registry().get_normal_platform()
-        model.set_suppress_output(True)
         prompt = f"""请基于以下查询生成10个针对向量搜索优化的不同表达。每个表达应满足：
 1. 聚焦关键技术概念和术语
 2. 使用清晰明确的语言
@@ -732,57 +742,70 @@ Content: {content}
 
     def search_similar(self, query: str, top_k: int = 30) -> List[Dict[str, str]]:
         """Search related files with optimized retrieval"""
-        try:
-            self.generate_codebase()
-            if self.index is None:
-                return []
-                
-            # Generate query variants for better coverage
-            query_variants = self._generate_query_variants(query)
-            
-            # Collect results from all variants
-            all_results = []
-            seen_files = set()
-            
-            for variant in query_variants:
-                # Get vector for each variant
-                query_vector = get_embedding(self.embedding_model, variant)
-                query_vector = query_vector.reshape(1, -1)
+        with yaspin(text="正在搜索相关文件...", color="cyan") as spinner:
+            try:
+                with spinner.hidden():
+                    self.generate_codebase()
+                if self.index is None:
+                    spinner.text = "没有找到有效的缓存文件"
+                    spinner.ok("✅")
+                    return []
+                    
+                # Generate query variants for better coverage
+                spinner.text = "生成查询变体..."
+                query_variants = self._generate_query_variants(query)
+                spinner.write("✅ 查询变体生成完成")
                 
-                # Search with current variant
-                initial_k = min(top_k * 2, len(self.file_paths))
-                distances, indices = self.index.search(query_vector, initial_k) # type: ignore
+                # Collect results from all variants
+                spinner.text = "收集结果..."
+                all_results = []
+                seen_files = set()
                 
-                # Process results
-                for idx, dist in zip(indices[0], distances[0]):
-                    if idx != -1:
-                        file_path = self.file_paths[idx]
-                        if file_path not in seen_files:
-                            similarity = 1.0 / (1.0 + float(dist))
-                            if similarity > 0.3:  # Lower threshold for better recall
-                                seen_files.add(file_path)
-                                all_results.append((file_path, similarity, self.vector_cache[file_path]["description"]))
-            
-            if not all_results:
-                return []
+                for variant in query_variants:
+                    # Get vector for each variant
+                    query_vector = get_embedding(self.embedding_model, variant)
+                    query_vector = query_vector.reshape(1, -1)
+                    
+                    # Search with current variant
+                    initial_k = min(top_k * 2, len(self.file_paths))
+                    distances, indices = self.index.search(query_vector, initial_k) # type: ignore
+                    
+                    # Process results
+                    for idx, dist in zip(indices[0], distances[0]):
+                        if idx != -1:
+                            file_path = self.file_paths[idx]
+                            if file_path not in seen_files:
+                                similarity = 1.0 / (1.0 + float(dist))
+                                if similarity > 0.3:  # Lower threshold for better recall
+                                    seen_files.add(file_path)
+                                    all_results.append((file_path, similarity, self.vector_cache[file_path]["description"]))
+                spinner.write("✅ 结果收集完成")
+                if not all_results:
+                    spinner.text = "没有找到相关文件"
+                    spinner.ok("✅")
+                    return []
+                    
+                spinner.text = "排序..."
+                # Sort by similarity and take top_k
+                all_results.sort(key=lambda x: x[1], reverse=True)
+                results = all_results[:top_k]
+                spinner.write("✅ 排序完成")
                 
-            # Sort by similarity and take top_k
-            all_results.sort(key=lambda x: x[1], reverse=True)
-            results = all_results[:top_k]
+                with spinner.hidden():
+                    results = self.pick_results(query_variants, [path for path, _, _ in results])
 
-            results = self.pick_results(query_variants, [path for path, _, _ in results])
+                output = "Found related files:\n"
+                for file in results:
+                    output += f'''- {file['file']} ({file['reason']})\n'''
 
-            output = "Found related files:\n"
-            for file in results:
-                output += f'''- {file['file']} ({file['reason']})\n'''
-            PrettyOutput.print(output, output_type=OutputType.INFO, lang="markdown")
-
-            
-            return results
-            
-        except Exception as e:
-            PrettyOutput.print(f"搜索失败: {str(e)}", output_type=OutputType.ERROR)
-            return []
+                spinner.text="结果输出完成"
+                spinner.ok("✅")
+                return results
+                
+            except Exception as e:
+                spinner.text = f"搜索失败: {str(e)}"
+                spinner.fail("❌")
+                return []
 
     def ask_codebase(self, query: str, top_k: int=20) -> Tuple[List[Dict[str, str]], str]:
         """Query the codebase with enhanced context building"""
@@ -830,48 +853,49 @@ Content: {content}
 相关代码文件（按相关性排序）:
 """
 
-        # 添加上下文，控制长度
-        available_count = self.max_token_count - get_context_token_count(prompt) - 1000  # 为回答预留空间
-        current_count = 0
-        
-        for path in files_from_codebase:
-            try:
-                content = open(path["file"], "r", encoding="utf-8").read()
-                file_content = f"""
-## 文件: {path["file"]}
-```
-{content}
-```
----
-"""
-                if current_count + get_context_token_count(file_content) > available_count:
-                    PrettyOutput.print(
-                        "由于上下文长度限制, 一些文件被省略", 
-                        output_type=OutputType.WARNING
-                    )
-                    break
+        with yaspin(text="正在生成回答...", color="cyan") as spinner:
+            # 添加上下文，控制长度
+            spinner.text = "添加上下文..."
+            available_count = self.max_token_count - get_context_token_count(prompt) - 1000  # 为回答预留空间
+            current_count = 0
+            
+            for path in files_from_codebase:
+                try:
+                    content = open(path["file"], "r", encoding="utf-8").read()
+                    file_content = f"""
+    ## 文件: {path["file"]}
+    ```
+    {content}
+    ```
+    ---
+    """
+                    if current_count + get_context_token_count(file_content) > available_count:
+                        spinner.write("⚠️ 由于上下文长度限制, 一些文件被省略")
+                        break
+                        
+                    prompt += file_content
+                    current_count += get_context_token_count(file_content)
                     
-                prompt += file_content
-                current_count += get_context_token_count(file_content)
-                
-            except Exception as e:
-                PrettyOutput.print(f"读取 {path} 失败: {str(e)}", 
-                                output_type=OutputType.ERROR)
-                continue
-
-        prompt += """
-# ❗ 重要规则
-1. 始终基于提供的代码进行回答
-2. 保持技术准确性
-3. 在相关时包含代码示例
-4. 指出任何缺失的信息
-5. 保持专业语言
-6. 使用用户的语言进行回答
-"""
-
-        model = PlatformRegistry.get_global_platform_registry().get_thinking_platform()
+                except Exception as e:
+                    spinner.write(f"❌ 读取 {path} 失败: {str(e)}")
+                    continue
 
-        return files_from_codebase, model.chat_until_success(prompt)
+            prompt += """
+    # ❗ 重要规则
+    1. 始终基于提供的代码进行回答
+    2. 保持技术准确性
+    3. 在相关时包含代码示例
+    4. 指出任何缺失的信息
+    5. 保持专业语言
+    6. 使用用户的语言进行回答
+    """
+
+            model = PlatformRegistry.get_global_platform_registry().get_thinking_platform()
+            spinner.text = "生成回答..."
+            ret = files_from_codebase, model.chat_until_success(prompt)
+            spinner.text = "回答生成完成"
+            spinner.ok("✅")
+            return ret
 
     def is_index_generated(self) -> bool:
         """Check if the index has been generated"""
@@ -972,9 +996,12 @@ def main():
         PrettyOutput.print(output, output_type=OutputType.INFO, lang="markdown")
 
     elif args.command == 'ask':            
-        response = codebase.ask_codebase(args.question, args.top_k)
-        output = f"""{response}"""
-        PrettyOutput.print(output, output_type=OutputType.INFO)
+        files, answer = codebase.ask_codebase(args.question, args.top_k)
+        output = f"# 相关文件：\n"
+        for file in files:
+            output += f"""- {file['file']} ({file['reason']})\n"""
+        output += f"# 回答：\n{answer}"
+        PrettyOutput.print(output, output_type=OutputType.SYSTEM, lang="markdown")
 
     else:
         parser.print_help()

jarvis/jarvis_dev/main.py

@@ -2,6 +2,7 @@ from jarvis.jarvis_platform.registry import PlatformRegistry
 from jarvis.jarvis_multi_agent import MultiAgent, AgentConfig
 from jarvis.jarvis_tools.registry import ToolRegistry
 from jarvis.jarvis_utils.input import get_multiline_input
+from jarvis.jarvis_utils.output import OutputType, PrettyOutput
 from jarvis.jarvis_utils.utils import init_env
 
 # 定义每个角色的系统提示
@@ -891,13 +892,13 @@ def main():
                 break
                 
             result = dev_team.run("My requirement: " + user_input)
-            print("\nFinal Result:", result)
+            PrettyOutput.print(result, output_type=OutputType.SYSTEM)
             
         except KeyboardInterrupt:
-            print("\nExiting...")
+            PrettyOutput.print("Exiting...", output_type=OutputType.SYSTEM)
             break
         except Exception as e:
-            print(f"\nError: {str(e)}")
+            PrettyOutput.print(f"Error: {str(e)}", output_type=OutputType.SYSTEM)
             continue
 
 if __name__ == "__main__":