jarvis-ai-assistant 0.1.125__py3-none-any.whl → 0.1.128__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,17 +85,13 @@ 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"""Please analyze the following code file and generate a detailed description. The description should include:
-1. Overall file functionality description
-2. description for each global variable, function, type definition, class, method, and other code elements
-
-Please use concise and professional language, emphasizing technical functionality to facilitate subsequent code retrieval.
-File path: {file_path}
-Code content:
+        prompt = f"""请分析以下代码文件并生成详细描述。描述应包含：
+1. 文件整体功能描述
+2. 对每个全局变量、函数、类型定义、类、方法和其他代码元素的描述
+
+请使用简洁专业的语言，强调技术功能，以便于后续代码检索。
+文件路径: {file_path}
+代码内容:
 {content}
 """
         response = model.chat_until_success(prompt)
@@ -114,50 +123,52 @@ Code content:
 
     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"""
@@ -320,7 +331,7 @@ Content: {content}
             ids = []
             self.file_paths = []  # Reset the file path list
             
-            for i, (file_path, data) in enumerate(self.vector_cache.items()):
+            for i, ( file_path, data) in enumerate(self.vector_cache.items()):
                 if "vector" not in data:
                     PrettyOutput.print(f"无效的缓存数据 {file_path}: 缺少向量", 
                                      output_type=OutputType.WARNING)
@@ -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,79 +552,79 @@ 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"""
-        prompt = f"""As a code analysis expert, please help identify the most relevant files for the given query using chain-of-thought reasoning.
+        prompt = f"""作为一名代码分析专家，请使用链式思维推理帮助识别与给定查询最相关的文件。
 
-Query: {query}
+查询: {query}
 
-Available files:
+可用文件:
 {''.join(files_info)}
 
-Think through this step by step:
-1. First, analyze the query to identify key requirements and technical concepts
-2. For each file:
-   - Examine its path and content
-   - Assess how it relates to the query's requirements
-   - Consider both direct and indirect relationships
-   - Rate its relevance (high/medium/low)
-3. Select only files with clear relevance to the query
-4. Order files by relevance, with most relevant first
-
-Please output your selection in YAML format:
+请按以下步骤思考：
+1. 首先，分析查询以识别关键需求和技术概念
+2. 对于每个文件：
+   - 检查其路径和内容
+   - 评估其与查询需求的关系
+   - 考虑直接和间接关系
+   - 评估其相关性（高/中/低）
+3. 仅选择与查询明确相关的文件
+4. 按相关性排序，最相关的文件在前
+
+请以YAML格式输出您的选择：
 <FILES>
 - file: path/to/most/relevant.py
   reason: xxxxxxxxxx
@@ -620,17 +632,16 @@ Please output your selection in YAML format:
   reason: yyyyyyyyyy
 </FILES>
 
-Important:
-- Only include files that are truly relevant
-- Exclude files with weak or unclear connections
-- Focus on implementation rather than test files
-- Consider both file paths and content
-- Only output the file paths, no other text
+重要提示：
+- 仅包含真正相关的文件
+- 排除连接不明确或较弱的文件
+- 重点关注实现文件而非测试文件
+- 同时考虑文件路径和内容
+- 仅输出文件路径，不要包含其他文本
 """
 
         # 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,30 +667,28 @@ Important:
             List[str]: The query variants list
         """
         model = PlatformRegistry.get_global_platform_registry().get_normal_platform()
-        model.set_suppress_output(True)
-        prompt = f"""Please generate 10 different expressions optimized for vector search based on the following query. Each expression should:
-
-1. Focus on key technical concepts and terminology
-2. Use clear and specific language
-3. Include important contextual terms
-4. Avoid general or ambiguous words
-5. Maintain semantic similarity with original query
-6. Be suitable for embedding-based search
-
-Original query: 
+        prompt = f"""请基于以下查询生成10个针对向量搜索优化的不同表达。每个表达应满足：
+1. 聚焦关键技术概念和术语
+2. 使用清晰明确的语言
+3. 包含重要的上下文术语
+4. 避免使用通用或模糊的词语
+5. 保持与原始查询的语义相似性
+6. 适合基于嵌入的搜索
+
+原始查询: 
 {query}
 
-Example transformations:
-Query: "How to handle user login?"
-Output format:
+示例转换：
+查询: "如何处理用户登录？"
+输出格式:
 <QUESTION>
-- user authentication implementation and flow
-- login system architecture and components
-- credential validation and session management
+- 用户认证的实现与流程
+- 登录系统架构与组件
+- 凭证验证与会话管理
 - ...
 </QUESTION>
 
-Please provide 10 search-optimized expressions in the specified format.
+请以指定格式提供10个搜索优化的表达。
 """
         response = model.chat_until_success(prompt)
         
@@ -733,57 +742,70 @@ Please provide 10 search-optimized expressions in the specified format.
 
     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]
-
-            results = self.pick_results(query_variants, [path for path, _, _ in results])
+                with spinner.hidden():
+                    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'''
-            PrettyOutput.print(output, output_type=OutputType.INFO, lang="markdown")
+                output = "Found related files:\n"
+                for file in results:
+                    output += f'''- {file['file']} ({file['reason']})\n'''
 
-            
-            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"""
@@ -794,85 +816,86 @@ Please provide 10 search-optimized expressions in the specified format.
             return [], ""
         
         prompt = f"""
-# 🤖 Role Definition
-You are a code analysis expert who provides comprehensive and accurate answers about codebases.
-
-# 🎯 Core Responsibilities
-- Analyze code files thoroughly
-- Explain technical concepts clearly
-- Provide relevant code examples
-- Identify missing information
-- Answer in user's language
-
-# 📋 Response Requirements
-## Content Quality
-- Focus on implementation details
-- Be technically precise
-- Include relevant code snippets
-- Indicate any missing information
-- Use professional terminology
-
-## Response Format
-- question: [Restate the question]
+# 🤖 角色定义
+您是一位代码分析专家，能够提供关于代码库的全面且准确的回答。
+
+# 🎯 核心职责
+- 深入分析代码文件
+- 清晰解释技术概念
+- 提供相关代码示例
+- 识别缺失的信息
+- 使用用户的语言进行回答
+
+# 📋 回答要求
+## 内容质量
+- 关注实现细节
+- 保持技术准确性
+- 包含相关代码片段
+- 指出任何缺失的信息
+- 使用专业术语
+
+## 回答格式
+- question: [重述问题]
   answer: |
-    [Detailed technical answer with:
-    - Implementation details
-    - Code examples (if relevant)
-    - Missing information (if any)
-    - Related technical concepts]
+    [详细的技术回答，包含：
+    - 实现细节
+    - 代码示例（如果相关）
+    - 缺失的信息（如果有）
+    - 相关技术概念]
 
-- question: [Follow-up question if needed]
+- question: [如果需要，提出后续问题]
   answer: |
-    [Additional technical details]
+    [额外的技术细节]
 
-# 🔍 Analysis Context
-Question: {query}
+# 🔍 分析上下文
+问题: {query}
 
-Relevant Code Files (by relevance):
+相关代码文件（按相关性排序）:
 """
 
-        # Add context with length control
-        available_count = self.max_token_count - get_context_token_count(prompt) - 1000  # Reserve space for answer
-        current_count = 0
-        
-        for path in files_from_codebase:
-            try:
-                content = open(path["file"], "r", encoding="utf-8").read()
-                file_content = f"""
-## File: {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 += """
-# ❗ Important Rules
-1. Always base answers on provided code
-2. Use technical precision
-3. Include code examples when relevant
-4. Indicate any missing information
-5. Maintain professional language
-6. Answer in user's language
-"""
-
-        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"""
@@ -973,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()