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

jarvis/jarvis_codebase/main.py

@@ -0,0 +1,871 @@
+import hashlib
+import os
+import numpy as np
+import faiss
+from typing import List, Tuple, Optional, Dict
+
+import yaml
+from jarvis.models.registry import PlatformRegistry
+import concurrent.futures
+from threading import Lock
+from concurrent.futures import ThreadPoolExecutor
+from jarvis.utils import OutputType, PrettyOutput, find_git_root, get_file_md5, get_max_context_length, get_single_line_input, get_thread_count, load_embedding_model, load_rerank_model, user_confirm
+from jarvis.utils import init_env
+import argparse
+import pickle
+import lzma  # 添加 lzma 导入
+from tqdm import tqdm
+import re
+
+class CodeBase:
+    def __init__(self, root_dir: str):
+        init_env()
+        self.root_dir = root_dir
+        os.chdir(self.root_dir)
+        self.thread_count = get_thread_count()
+        self.max_context_length = get_max_context_length()
+        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)
+            
+        # 初始化嵌入模型
+        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("Model loaded successfully", output_type=OutputType.SUCCESS)
+        except Exception as e:
+            PrettyOutput.print(f"Failed to load model: {str(e)}", output_type=OutputType.ERROR)
+            raise
+            
+        self.vector_dim = self.embedding_model.get_sentence_embedding_dimension()
+        self.git_file_list = self.get_git_file_list()
+        self.platform_registry = PlatformRegistry.get_global_platform_registry()
+        
+        # 初始化缓存和索引
+        self.vector_cache = {}
+        self.file_paths = []
+        
+        # 加载所有缓存文件
+        self._load_all_cache()
+
+    def get_git_file_list(self):
+        """Get the list of files in the git repository, excluding the .jarvis-codebase directory"""
+        files = os.popen("git ls-files").read().splitlines()
+        # Filter out files in the .jarvis-codebase directory
+        return [f for f in files if not f.startswith(".jarvis")]
+
+    def is_text_file(self, file_path: str):
+        with open(file_path, "r", encoding="utf-8") as f:
+            try:
+                f.read()
+                return True
+            except UnicodeDecodeError:
+                return False
+
+    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"Make description for {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:
+{content}
+"""
+        response = model.chat_until_success(prompt)
+        return response
+
+    def export(self):
+        """Export the current index data to standard output"""
+        for file_path, data in self.vector_cache.items():
+            print(f"## {file_path}")
+            print(f"- path: {file_path}")
+            print(f"- description: {data['description']}")
+    
+    def _get_cache_path(self, file_path: str) -> str:
+        """Get cache file path for a source file
+        
+        Args:
+            file_path: Source file path
+            
+        Returns:
+            str: Cache file path
+        """
+        # 处理文件路径：
+        # 1. 移除开头的 ./ 或 /
+        # 2. 将 / 替换为 --
+        # 3. 添加 .cache 后缀
+        clean_path = file_path.lstrip('./').lstrip('/')
+        cache_name = clean_path.replace('/', '--') + '.cache'
+        return os.path.join(self.cache_dir, cache_name)
+
+    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"Failed to load cache file {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
+                
+                PrettyOutput.print(f"Loaded {len(self.vector_cache)} vector cache and rebuilt index", 
+                                 output_type=OutputType.INFO)
+            else:
+                self.index = None
+                PrettyOutput.print("No valid cache files found", output_type=OutputType.WARNING)
+                
+        except Exception as e:
+            PrettyOutput.print(f"Failed to load cache directory: {str(e)}", 
+                             output_type=OutputType.WARNING)
+            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"""
+        try:
+            with open(file_path, "rb") as f:
+                file_md5 = hashlib.md5(f.read()).hexdigest()
+        except Exception as e:
+            PrettyOutput.print(f"Failed to calculate MD5 for {file_path}: {str(e)}", 
+                              output_type=OutputType.ERROR)
+            file_md5 = ""
+        
+        # 准备缓存数据
+        cache_data = {
+            "path": file_path,  # 保存文件路径
+            "md5": file_md5,    # 保存文件MD5
+            "description": description,  # 保存文件描述
+            "vector": vector    # 保存向量
+        }
+        
+        # 更新内存缓存
+        self.vector_cache[file_path] = cache_data
+        
+        # 保存到单独的缓存文件
+        cache_path = self._get_cache_path(file_path)
+        try:
+            with lzma.open(cache_path, 'wb') as f:
+                pickle.dump(cache_data, f, protocol=pickle.HIGHEST_PROTOCOL)
+        except Exception as e:
+            PrettyOutput.print(f"Failed to save cache for {file_path}: {str(e)}", 
+                             output_type=OutputType.ERROR)
+
+    def get_cached_vector(self, file_path: str, description: str) -> Optional[np.ndarray]:
+        """Get the vector representation of a file from the cache"""
+        if file_path not in self.vector_cache:
+            return None
+        
+        # Check if the file has been modified
+        try:
+            with open(file_path, "rb") as f:
+                current_md5 = hashlib.md5(f.read()).hexdigest()
+        except Exception as e:
+            PrettyOutput.print(f"Failed to calculate MD5 for {file_path}: {str(e)}", 
+                              output_type=OutputType.ERROR)
+            return None
+        
+        cached_data = self.vector_cache[file_path]
+        if cached_data["md5"] != current_md5:
+            return None
+        
+        # Check if the description has changed
+        if cached_data["description"] != description:
+            return None
+        
+        return cached_data["vector"]
+
+    def get_embedding(self, text: str) -> np.ndarray:
+        """Use the transformers model to get the vector representation of text"""
+        # Truncate long text
+        max_length = 512  # Or other suitable length
+        text = ' '.join(text.split()[:max_length])
+        
+        # Get the embedding vector
+        embedding = self.embedding_model.encode(text, 
+                                                 normalize_embeddings=True,  # L2 normalization
+                                                 show_progress_bar=False)
+        vector = np.array(embedding, dtype=np.float32)
+        return vector
+
+    def vectorize_file(self, file_path: str, description: str) -> np.ndarray:
+        """Vectorize the file content and description"""
+        try:
+            # Try to get the vector from the cache first
+            cached_vector = self.get_cached_vector(file_path, description)
+            if cached_vector is not None:
+                return cached_vector
+                
+            # Read the file content and combine information
+            content = open(file_path, "r", encoding="utf-8").read()[:self.max_context_length]  # Limit the file content length
+            
+            # Combine file information, including file content
+            combined_text = f"""
+File path: {file_path}
+Description: {description}
+Content: {content}
+"""
+            vector = self.get_embedding(combined_text)
+            
+            # Save to cache
+            self.cache_vector(file_path, vector, description)
+            return vector
+        except Exception as e:
+            PrettyOutput.print(f"Error vectorizing file {file_path}: {str(e)}", 
+                             output_type=OutputType.ERROR)
+            return np.zeros(self.vector_dim, dtype=np.float32) # type: ignore
+
+    def clean_cache(self) -> bool:
+        """Clean expired cache records"""
+        try:
+            files_to_delete = []
+            for file_path in list(self.vector_cache.keys()):
+                if not os.path.exists(file_path):
+                    files_to_delete.append(file_path)
+                    cache_path = self._get_cache_path(file_path)
+                    try:
+                        os.remove(cache_path)
+                    except Exception:
+                        pass
+                        
+            for file_path in files_to_delete:
+                del self.vector_cache[file_path]
+                if file_path in self.file_paths:
+                    self.file_paths.remove(file_path)
+                    
+            return bool(files_to_delete)
+            
+        except Exception as e:
+            PrettyOutput.print(f"Failed to clean cache: {str(e)}", 
+                             output_type=OutputType.ERROR)
+            return False
+
+    def process_file(self, file_path: str):
+        """Process a single file"""
+        try:
+            # Skip non-existent files
+            if not os.path.exists(file_path):
+                return None
+                
+            if not self.is_text_file(file_path):
+                return None
+            
+            md5 = get_file_md5(file_path)
+
+            content = open(file_path, "r", encoding="utf-8").read()
+            
+            # Check if the file has already been processed and the content has not changed
+            if file_path in self.vector_cache:
+                if self.vector_cache[file_path].get("md5") == md5:
+                    return None
+                    
+            description = self.make_description(file_path, content)  # Pass the truncated content
+            vector = self.vectorize_file(file_path, description)
+            
+            # Save to cache, using the actual file path as the key
+            self.vector_cache[file_path] = {
+                "vector": vector,
+                "description": description,
+                "md5": md5
+            }
+            
+            return file_path
+            
+        except Exception as e:
+            PrettyOutput.print(f"Failed to process file {file_path}: {str(e)}", 
+                             output_type=OutputType.ERROR)
+            return None
+
+    def build_index(self):
+        """Build a faiss index from the vector cache"""
+        try:
+            if not self.vector_cache:
+                self.index = None
+                return
+
+            # Create the underlying HNSW index
+            hnsw_index = faiss.IndexHNSWFlat(self.vector_dim, 16)
+            hnsw_index.hnsw.efConstruction = 40
+            hnsw_index.hnsw.efSearch = 16
+            
+            # Wrap the HNSW index with IndexIDMap
+            self.index = faiss.IndexIDMap(hnsw_index)
+            
+            vectors = []
+            ids = []
+            self.file_paths = []  # Reset the file path list
+            
+            for i, (file_path, data) in enumerate(self.vector_cache.items()):
+                if "vector" not in data:
+                    PrettyOutput.print(f"Invalid cache data for {file_path}: missing vector", 
+                                     output_type=OutputType.WARNING)
+                    continue
+                    
+                vector = data["vector"]
+                if not isinstance(vector, np.ndarray):
+                    PrettyOutput.print(f"Invalid vector type for {file_path}: {type(vector)}", 
+                                     output_type=OutputType.WARNING)
+                    continue
+                    
+                vectors.append(vector.reshape(1, -1))
+                ids.append(i)
+                self.file_paths.append(file_path)
+                
+            if vectors:
+                vectors = np.vstack(vectors)
+                if len(vectors) != len(ids):
+                    PrettyOutput.print(f"Vector count mismatch: {len(vectors)} vectors vs {len(ids)} ids", 
+                                     output_type=OutputType.ERROR)
+                    self.index = None
+                    return
+                    
+                try:
+                    self.index.add_with_ids(vectors, np.array(ids)) # type: ignore
+                    PrettyOutput.print(f"Successfully built index with {len(vectors)} vectors", 
+                                     output_type=OutputType.SUCCESS)
+                except Exception as e:
+                    PrettyOutput.print(f"Failed to add vectors to index: {str(e)}", 
+                                     output_type=OutputType.ERROR)
+                    self.index = None
+            else:
+                PrettyOutput.print("No valid vectors found, index not built", 
+                                 output_type=OutputType.WARNING)
+                self.index = None
+                
+        except Exception as e:
+            PrettyOutput.print(f"Failed to build index: {str(e)}", 
+                             output_type=OutputType.ERROR)
+            self.index = None
+
+    def gen_vector_db_from_cache(self):
+        """Generate a vector database from the cache"""
+        self.build_index()
+        self._load_all_cache()
+
+
+    def generate_codebase(self, force: bool = False):
+        """Generate the codebase index
+        Args:
+            force: Whether to force rebuild the index, without asking the user
+        """
+        try:
+            # Update the git file list
+            self.git_file_list = self.get_git_file_list()
+            
+            # Check file changes
+            PrettyOutput.print("Check file changes...", output_type=OutputType.INFO)
+            changes_detected = False
+            new_files = []
+            modified_files = []
+            deleted_files = []
+            
+            # Check deleted files
+            files_to_delete = []
+            for file_path in list(self.vector_cache.keys()):
+                if file_path not in self.git_file_list:
+                    deleted_files.append(file_path)
+                    files_to_delete.append(file_path)
+                    changes_detected = True
+            # Check new and modified files
+            from rich.progress import Progress
+            with Progress() as progress:
+                task = progress.add_task("Check file status", total=len(self.git_file_list))
+                for file_path in self.git_file_list:
+                    if not os.path.exists(file_path) or not self.is_text_file(file_path):
+                        progress.advance(task)
+                        continue
+                    
+                    try:
+                        current_md5 = get_file_md5(file_path)
+                        
+                        if file_path not in self.vector_cache:
+                            new_files.append(file_path)
+                            changes_detected = True
+                        elif self.vector_cache[file_path].get("md5") != current_md5:
+                            modified_files.append(file_path)
+                            changes_detected = True
+                    except Exception as e:
+                        PrettyOutput.print(f"Failed to check file {file_path}: {str(e)}", 
+                                         output_type=OutputType.ERROR)
+                    progress.advance(task)
+            
+            # If changes are detected, display changes and ask the user
+            if changes_detected:
+                output_lines = ["Detected the following changes:"]
+                if new_files:
+                    output_lines.append("New files:")
+                    output_lines.extend(f"  {f}" for f in new_files)
+                if modified_files:
+                    output_lines.append("Modified files:")
+                    output_lines.extend(f"  {f}" for f in modified_files)
+                if deleted_files:
+                    output_lines.append("Deleted files:")
+                    output_lines.extend(f"  {f}" for f in deleted_files)
+                
+                PrettyOutput.print("\n".join(output_lines), output_type=OutputType.WARNING)
+
+                # If force is True, continue directly
+                if not force:
+                    if not user_confirm("Rebuild the index?", False):
+                        PrettyOutput.print("Cancel rebuilding the index", output_type=OutputType.INFO)
+                        return
+                
+                # Clean deleted files
+                for file_path in files_to_delete:
+                    del self.vector_cache[file_path]
+                if files_to_delete:
+                    PrettyOutput.print(f"Cleaned the cache of {len(files_to_delete)} files", 
+                                     output_type=OutputType.INFO)
+                
+                # Process new and modified files
+                files_to_process = new_files + modified_files
+                processed_files = []
+                
+                with tqdm(total=len(files_to_process), desc="Processing files") as pbar:
+                    # Use a thread pool to process files
+                    with ThreadPoolExecutor(max_workers=self.thread_count) as executor:
+                        # Submit all tasks
+                        future_to_file = {
+                            executor.submit(self.process_file, file): file 
+                            for file in files_to_process
+                        }
+                        
+                        # Process completed tasks
+                        for future in concurrent.futures.as_completed(future_to_file):
+                            file = future_to_file[future]
+                            try:
+                                result = future.result()
+                                if result:
+                                    processed_files.append(result)
+                            except Exception as e:
+                                PrettyOutput.print(f"Failed to process file {file}: {str(e)}", 
+                                                output_type=OutputType.ERROR)
+                            pbar.update(1)
+
+                if processed_files:
+                    PrettyOutput.print("Rebuilding the vector database...", output_type=OutputType.INFO)
+                    self.gen_vector_db_from_cache()
+                    PrettyOutput.print(f"Successfully generated the index for {len(processed_files)} files", 
+                                    output_type=OutputType.SUCCESS)
+            else:
+                PrettyOutput.print("No file changes detected, no need to rebuild the index", output_type=OutputType.INFO)
+                
+        except Exception as e:
+            # Try to save the cache when an exception occurs
+            try:
+                self._load_all_cache()
+            except Exception as save_error:
+                PrettyOutput.print(f"Failed to save cache: {str(save_error)}", 
+                                output_type=OutputType.ERROR)
+            raise e  # Re-raise the original exception
+
+
+    def _text_search_score(self, content: str, keywords: List[str]) -> float:
+        """Calculate the matching score between the text content and the keywords
+        
+        Args:
+            content: Text content
+            keywords: List of keywords
+            
+        Returns:
+            float: Matching score (0-1)
+        """
+        if not keywords:
+            return 0.0
+            
+        content = content.lower()
+        matched_keywords = set()
+        
+        for keyword in keywords:
+            keyword = keyword.lower()
+            if keyword in content:
+                matched_keywords.add(keyword)
+                
+        # Calculate the matching score
+        score = len(matched_keywords) / len(keywords)
+        return score
+
+    def pick_results(self, query: str, initial_results: List[str]) -> List[str]:
+        """Use a large model to pick the search results
+        
+        Args:
+            query: Search query
+            initial_results: Initial results list of file paths
+            
+        Returns:
+            List[str]: The picked results list, each item is a file path
+        """
+        if not initial_results:
+            return []
+            
+        try:
+            PrettyOutput.print(f"Picking results for query: {query}", output_type=OutputType.INFO)
+            
+            # Maximum content length per batch
+            max_batch_length = self.max_context_length - 1000  # Reserve space for prompt
+            max_file_length = max_batch_length // 3  # Limit individual file size
+            
+            # Process files in batches
+            all_selected_files = set()
+            current_batch = []
+            current_length = 0
+            
+            for path in initial_results:
+                try:
+                    content = open(path, "r", encoding="utf-8").read()
+                    # Truncate large files
+                    if len(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"
+                    file_length = len(file_info)
+                    
+                    # If adding this file would exceed batch limit
+                    if current_length + file_length > max_batch_length:
+                        # Process current batch
+                        if current_batch:
+                            selected = self._process_batch(query, current_batch)
+                            all_selected_files.update(selected)
+                        # Start new batch
+                        current_batch = [file_info]
+                        current_length = file_length
+                    else:
+                        current_batch.append(file_info)
+                        current_length += file_length
+                        
+                except Exception as e:
+                    PrettyOutput.print(f"Failed to read file {path}: {str(e)}", OutputType.ERROR)
+                    continue
+            
+            # Process final batch
+            if current_batch:
+                selected = self._process_batch(query, current_batch)
+                all_selected_files.update(selected)
+            
+            # Convert set to list and maintain original order
+            final_results = [path for path in initial_results if path in all_selected_files]
+            return final_results
+
+        except Exception as e:
+            PrettyOutput.print(f"Failed to pick: {str(e)}", OutputType.ERROR)
+            return initial_results
+            
+    def _process_batch(self, query: str, files_info: List[str]) -> List[str]:
+        """Process a batch of files
+        
+        Args:
+            query: Search query
+            files_info: List of file information strings
+            
+        Returns:
+            List[str]: Selected file paths from this batch
+        """
+        prompt = f"""Please analyze the following code files and determine which files are most relevant to the given query. Consider file paths and code content to make your judgment.
+
+Query: {query}
+
+Available files:
+{''.join(files_info)}
+
+Please output a YAML list of relevant file paths, ordered by relevance (most relevant first). Only include files that are truly relevant to the query.
+Output format:
+<FILES>
+- path/to/file1.py
+- path/to/file2.py
+</FILES>
+
+Note: Only include files that have a strong connection to the query."""
+
+        # Use a large model to evaluate
+        model = PlatformRegistry.get_global_platform_registry().get_normal_platform()
+        response = model.chat_until_success(prompt)
+
+        # Parse the response
+        import yaml
+        files_match = re.search(r'<FILES>\n(.*?)</FILES>', response, re.DOTALL)
+        if not files_match:
+            return []
+
+        # Extract the file list
+        try:
+            selected_files = yaml.safe_load(files_match.group(1))
+            return selected_files if selected_files else []
+        except Exception as e:
+            PrettyOutput.print(f"Failed to parse response: {str(e)}", OutputType.ERROR)
+            return []
+
+    def _generate_query_variants(self, query: str) -> List[str]:
+        """Generate different expressions of the query
+        
+        Args:
+            query: Original query
+            
+        Returns:
+            List[str]: The query variants list
+        """
+        model = PlatformRegistry.get_global_platform_registry().get_normal_platform()
+        prompt = f"""Please generate 3 different expressions based on the following query, each expression should fully convey the meaning of the original query. These expressions will be used for code search, maintain professionalism and accuracy.
+Original query: {query}
+
+Please output 3 expressions directly, separated by two line breaks, without numbering or other markers.
+"""
+        variants = model.chat_until_success(prompt).strip().split('\n\n')
+        variants.append(query)  # Add the original query
+        return variants
+
+    def _vector_search(self, query_variants: List[str], top_k: int) -> Dict[str, Tuple[str, float, str]]:
+        """Use vector search to find related files
+        
+        Args:
+            query_variants: The query variants list
+            top_k: The number of results to return
+            
+        Returns:
+            Dict[str, Tuple[str, float, str]]: The mapping from file path to (file path, score, description)
+        """
+        results = {}
+        for query in query_variants:
+            query_vector = self.get_embedding(query)
+            query_vector = query_vector.reshape(1, -1)
+            
+            distances, indices = self.index.search(query_vector, top_k) # type: ignore
+            
+            for i, distance in zip(indices[0], distances[0]):
+                if i == -1:
+                    continue
+                    
+                similarity = 1.0 / (1.0 + float(distance))
+                file_path = self.file_paths[i]
+                # Use the highest similarity score
+                if file_path not in results:
+                    if similarity > 0.5:
+                        data = self.vector_cache[file_path]
+                        results[file_path] = (file_path, similarity, data["description"])
+        
+        return results
+
+
+    def search_similar(self, query: str, top_k: int = 30) -> List[str]:
+        """Search related files"""
+        try:
+            self.generate_codebase()
+            if self.index is None:
+                return []            
+            # Generate the query variants
+            query_variants = self._generate_query_variants(query)
+            
+            # Perform vector search
+            vector_results = self._vector_search(query_variants, top_k)            
+
+            results = list(vector_results.values())
+            results.sort(key=lambda x: x[1], reverse=True)
+
+            # Take the top top_k results for reordering
+            initial_results = results[:top_k]
+            
+            # If no results are found, return directly
+            if not initial_results:
+                return []
+            
+            # Filter low-scoring results
+            initial_results = [(path, score, desc) for path, score, desc in initial_results if score >= 0.5]
+
+            message = "Found related files:\n"
+            for path, score, _ in initial_results:
+                message += f"File: {path} Similarity: {score:.3f}\n"
+            PrettyOutput.print(message.rstrip(), output_type=OutputType.INFO, lang="markdown")
+                
+            # Reorder the preliminary results
+            return self.pick_results(query, [path for path, _, _ in initial_results])
+            
+        except Exception as e:
+            PrettyOutput.print(f"Failed to search: {str(e)}", output_type=OutputType.ERROR)
+            return []
+
+    def ask_codebase(self, query: str, top_k: int=20) -> str:
+        """Query the codebase"""
+        results = self.search_similar(query, top_k)
+        if not results:
+            PrettyOutput.print("No related files found", output_type=OutputType.WARNING)
+            return ""
+        
+        message = "Found related files:\n"
+        for path in results:
+            message += f"File: {path}\n"
+        PrettyOutput.print(message.rstrip(), output_type=OutputType.SUCCESS, lang="markdown")
+        
+        prompt = f"""You are a code expert, please answer the user's question based on the following file information:
+"""
+        for path in results:
+            try:
+                if len(prompt) > self.max_context_length:
+                    PrettyOutput.print(f"Avoid context overflow, discard low-related file: {path}", OutputType.WARNING)
+                    continue
+                content = open(path, "r", encoding="utf-8").read()
+                prompt += f"""
+File path: {path}
+File content:
+{content}
+========================================
+"""
+            except Exception as e:
+                PrettyOutput.print(f"Failed to read file {path}: {str(e)}", 
+                                 output_type=OutputType.ERROR)
+                continue
+                
+        prompt += f"""
+User question: {query}
+
+Please answer the user's question in Chinese using professional language. If the provided file content is insufficient to answer the user's question, please inform the user. Never make up information.
+"""
+        model = PlatformRegistry.get_global_platform_registry().get_codegen_platform()
+        response = model.chat_until_success(prompt)
+        return response
+
+    def is_index_generated(self) -> bool:
+        """Check if the index has been generated"""
+        try:
+            # 1. 检查基本条件
+            if not self.vector_cache or not self.file_paths:
+                return False
+                
+            if not hasattr(self, 'index') or self.index is None:
+                return False
+                
+            # 2. 检查索引是否可用
+            # 创建测试向量
+            test_vector = np.zeros((1, self.vector_dim), dtype=np.float32) # type: ignore
+            try:
+                self.index.search(test_vector, 1) # type: ignore
+            except Exception:
+                return False
+                
+            # 3. 验证向量缓存和文件路径的一致性
+            if len(self.vector_cache) != len(self.file_paths):
+                return False
+                
+            # 4. 验证所有缓存文件
+            for file_path in self.file_paths:
+                if file_path not in self.vector_cache:
+                    return False
+                    
+                cache_path = self._get_cache_path(file_path)
+                if not os.path.exists(cache_path):
+                    return False
+                    
+                cache_data = self.vector_cache[file_path]
+                if not isinstance(cache_data.get("vector"), np.ndarray):
+                    return False
+            
+            return True
+            
+        except Exception as e:
+            PrettyOutput.print(f"Error checking index status: {str(e)}", 
+                             output_type=OutputType.ERROR)
+            return False
+
+
+
+
+
+def main():
+
+    parser = argparse.ArgumentParser(description='Codebase management and search tool')
+    subparsers = parser.add_subparsers(dest='command', help='Available commands')
+
+    # Generate command
+    generate_parser = subparsers.add_parser('generate', help='Generate codebase index')
+    generate_parser.add_argument('--force', action='store_true', help='Force rebuild index')
+
+    # Search command
+    search_parser = subparsers.add_parser('search', help='Search similar code files')
+    search_parser.add_argument('query', type=str, help='Search query')
+    search_parser.add_argument('--top-k', type=int, default=20, help='Number of results to return (default: 20)')
+
+    # Ask command
+    ask_parser = subparsers.add_parser('ask', help='Ask a question about the codebase')
+    ask_parser.add_argument('question', type=str, help='Question to ask')
+    ask_parser.add_argument('--top-k', type=int, default=20, help='Number of results to use (default: 20)')
+
+    export_parser = subparsers.add_parser('export', help='Export current index data')
+    args = parser.parse_args()
+    
+    current_dir = find_git_root()
+    codebase = CodeBase(current_dir)
+
+    if args.command == 'export':
+        codebase.export()
+        return
+
+    # 如果没有生成索引，且不是生成命令，提示用户先生成索引
+    if not codebase.is_index_generated() and args.command != 'generate':
+        PrettyOutput.print("索引尚未生成，请先运行 'generate' 命令生成索引", output_type=OutputType.WARNING)
+        return
+
+    if args.command == 'generate':
+        try:
+            codebase.generate_codebase(force=args.force)
+            PrettyOutput.print("Codebase generation completed", output_type=OutputType.SUCCESS)
+        except Exception as e:
+            PrettyOutput.print(f"Error during codebase generation: {str(e)}", output_type=OutputType.ERROR)
+    
+    elif args.command == 'search':
+        results = codebase.search_similar(args.query, args.top_k)
+        if not results:
+            PrettyOutput.print("No similar files found", output_type=OutputType.WARNING)
+            return
+            
+        output = "Search Results:\n"
+        for path in results:
+            output += f"""- {path}\n"""
+        PrettyOutput.print(output, output_type=OutputType.INFO, lang="markdown")
+
+    elif args.command == 'ask':            
+        response = codebase.ask_codebase(args.question, args.top_k)
+        output = f"""Answer:\n{response}"""
+        PrettyOutput.print(output, output_type=OutputType.INFO)
+
+    else:
+        parser.print_help()
+
+
+if __name__ == "__main__":
+    exit(main())