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

jarvis/jarvis_rag/main.py

@@ -0,0 +1,818 @@
+import os
+import numpy as np
+import faiss
+from typing import List, Tuple, Optional, Dict
+import pickle
+from jarvis.utils import OutputType, PrettyOutput, get_file_md5, get_max_context_length, load_embedding_model, load_rerank_model
+from jarvis.utils import init_env
+from dataclasses import dataclass
+from tqdm import tqdm
+import fitz  # PyMuPDF for PDF files
+from docx import Document as DocxDocument  # python-docx for DOCX files
+from pathlib import Path
+from jarvis.models.registry import PlatformRegistry
+import shutil
+from datetime import datetime
+import lzma  # 添加 lzma 导入
+from concurrent.futures import ThreadPoolExecutor
+from threading import Lock
+
+@dataclass
+class Document:
+    """Document class, for storing document content and metadata"""
+    content: str  # Document content
+    metadata: Dict  # Metadata (file path, position, etc.)
+    md5: str = ""  # File MD5 value, for incremental update detection
+
+class FileProcessor:
+    """Base class for file processor"""
+    @staticmethod
+    def can_handle(file_path: str) -> bool:
+        """Determine if the file can be processed"""
+        raise NotImplementedError
+        
+    @staticmethod
+    def extract_text(file_path: str) -> str:
+        """Extract file text content"""
+        raise NotImplementedError
+
+class TextFileProcessor(FileProcessor):
+    """Text file processor"""
+    ENCODINGS = ['utf-8', 'gbk', 'gb2312', 'latin1']
+    SAMPLE_SIZE = 8192  # Read the first 8KB to detect encoding
+    
+    @staticmethod
+    def can_handle(file_path: str) -> bool:
+        """Determine if the file is a text file by trying to decode it"""
+        try:
+            # Read the first part of the file to detect encoding
+            with open(file_path, 'rb') as f:
+                sample = f.read(TextFileProcessor.SAMPLE_SIZE)
+                
+            # Check if it contains null bytes (usually represents a binary file)
+            if b'\x00' in sample:
+                return False
+                
+            # Check if it contains too many non-printable characters (usually represents a binary file)
+            non_printable = sum(1 for byte in sample if byte < 32 and byte not in (9, 10, 13))  # tab, newline, carriage return
+            if non_printable / len(sample) > 0.3:  # If non-printable characters exceed 30%, it is considered a binary file
+                return False
+                
+            # Try to decode with different encodings
+            for encoding in TextFileProcessor.ENCODINGS:
+                try:
+                    sample.decode(encoding)
+                    return True
+                except UnicodeDecodeError:
+                    continue
+                    
+            return False
+            
+        except Exception:
+            return False
+    
+    @staticmethod
+    def extract_text(file_path: str) -> str:
+        """Extract text content, using the detected correct encoding"""
+        detected_encoding = None
+        try:
+            # First try to detect encoding
+            with open(file_path, 'rb') as f:
+                raw_data = f.read()
+                
+            # Try different encodings
+            for encoding in TextFileProcessor.ENCODINGS:
+                try:
+                    raw_data.decode(encoding)
+                    detected_encoding = encoding
+                    break
+                except UnicodeDecodeError:
+                    continue
+                    
+            if not detected_encoding:
+                raise UnicodeDecodeError(f"Failed to decode file with supported encodings: {file_path}") # type: ignore
+                
+            # Use the detected encoding to read the file
+            with open(file_path, 'r', encoding=detected_encoding, errors='replace') as f:
+                content = f.read()
+                
+            # Normalize Unicode characters
+            import unicodedata
+            content = unicodedata.normalize('NFKC', content)
+            
+            return content
+            
+        except Exception as e:
+            raise Exception(f"Failed to read file: {str(e)}")
+
+class PDFProcessor(FileProcessor):
+    """PDF file processor"""
+    @staticmethod
+    def can_handle(file_path: str) -> bool:
+        return Path(file_path).suffix.lower() == '.pdf'
+    
+    @staticmethod
+    def extract_text(file_path: str) -> str:
+        text_parts = []
+        with fitz.open(file_path) as doc: # type: ignore
+            for page in doc:
+                text_parts.append(page.get_text()) # type: ignore
+        return "\n".join(text_parts)
+
+class DocxProcessor(FileProcessor):
+    """DOCX file processor"""
+    @staticmethod
+    def can_handle(file_path: str) -> bool:
+        return Path(file_path).suffix.lower() == '.docx'
+    
+    @staticmethod
+    def extract_text(file_path: str) -> str:
+        doc = DocxDocument(file_path)
+        return "\n".join([paragraph.text for paragraph in doc.paragraphs])
+
+class RAGTool:
+    def __init__(self, root_dir: str):
+        """Initialize RAG tool
+        
+        Args:
+            root_dir: Project root directory
+        """
+        init_env()
+        self.root_dir = root_dir
+        os.chdir(self.root_dir)
+        
+        # Initialize configuration
+        self.min_paragraph_length = int(os.environ.get("JARVIS_MIN_PARAGRAPH_LENGTH", "50"))  # Minimum paragraph length
+        self.max_paragraph_length = int(os.environ.get("JARVIS_MAX_PARAGRAPH_LENGTH", "1000"))  # Maximum paragraph length
+        self.context_window = int(os.environ.get("JARVIS_CONTEXT_WINDOW", "5"))  # Context window size, default前后各5个片段
+        self.max_context_length = int(get_max_context_length() * 0.8)
+        
+        # Initialize data directory
+        self.data_dir = os.path.join(self.root_dir, ".jarvis-rag")
+        if not os.path.exists(self.data_dir):
+            os.makedirs(self.data_dir)
+            
+        # Initialize embedding model
+        try:
+            self.embedding_model = load_embedding_model()
+            self.vector_dim = self.embedding_model.get_sentence_embedding_dimension()
+            PrettyOutput.print("Model loaded", output_type=OutputType.SUCCESS)
+        except Exception as e:
+            PrettyOutput.print(f"Failed to load model: {str(e)}", output_type=OutputType.ERROR)
+            raise
+
+        # Initialize cache and index
+        self.cache_path = os.path.join(self.data_dir, "cache.pkl")
+        self.documents: List[Document] = []
+        self.index = None  # IVF index for search
+        self.flat_index = None  # Store original vectors
+        self.file_md5_cache = {}  # Store file MD5 values
+        
+        # Load cache
+        self._load_cache()
+
+        # Register file processors
+        self.file_processors = [
+            TextFileProcessor(),
+            PDFProcessor(),
+            DocxProcessor()
+        ]
+
+        # Add thread related configuration
+        self.thread_count = int(os.environ.get("JARVIS_THREAD_COUNT", os.cpu_count() or 4))
+        self.vector_lock = Lock()  # Protect vector list concurrency
+
+    def _load_cache(self):
+        """Load cache data"""
+        if os.path.exists(self.cache_path):
+            try:
+                with lzma.open(self.cache_path, 'rb') as f:
+                    cache_data = pickle.load(f)
+                    self.documents = cache_data["documents"]
+                    vectors = cache_data["vectors"]
+                    self.file_md5_cache = cache_data.get("file_md5_cache", {})  # 加载MD5缓存
+                    
+                # 重建索引
+                if vectors is not None:
+                    self._build_index(vectors)
+                PrettyOutput.print(f"Loaded {len(self.documents)} document fragments", 
+                                output_type=OutputType.INFO)
+            except Exception as e:
+                PrettyOutput.print(f"Failed to load cache: {str(e)}", 
+                                output_type=OutputType.WARNING)
+                self.documents = []
+                self.index = None
+                self.flat_index = None
+                self.file_md5_cache = {}
+
+    def _save_cache(self, vectors: np.ndarray):
+        """Optimize cache saving"""
+        try:
+            cache_data = {
+                "version": "1.0",
+                "timestamp": datetime.now().isoformat(),
+                "documents": self.documents,
+                "vectors": vectors.copy() if vectors is not None else None,  # Create a copy of the array
+                "file_md5_cache": dict(self.file_md5_cache),  # Create a copy of the dictionary
+                "metadata": {
+                    "vector_dim": self.vector_dim,
+                    "total_docs": len(self.documents),
+                    "model_name": self.embedding_model.__class__.__name__
+                }
+            }
+            
+            # First serialize the data to a byte stream
+            data = pickle.dumps(cache_data, protocol=pickle.HIGHEST_PROTOCOL)
+            
+            # Then use LZMA to compress the byte stream
+            with lzma.open(self.cache_path, 'wb') as f:
+                f.write(data)
+            
+            # Create a backup
+            backup_path = f"{self.cache_path}.backup"
+            shutil.copy2(self.cache_path, backup_path)
+            
+            PrettyOutput.print(f"Cache saved: {len(self.documents)} document fragments", 
+                            output_type=OutputType.INFO)
+        except Exception as e:
+            PrettyOutput.print(f"Failed to save cache: {str(e)}", 
+                            output_type=OutputType.ERROR)
+            raise
+
+    def _build_index(self, vectors: np.ndarray):
+        """Build FAISS index"""
+        if vectors.shape[0] == 0:
+            self.index = None
+            self.flat_index = None
+            return
+            
+        # Create a flat index to store original vectors, for reconstruction
+        self.flat_index = faiss.IndexFlatIP(self.vector_dim)
+        self.flat_index.add(vectors) # type: ignore
+        
+        # Create an IVF index for fast search
+        nlist = max(4, int(vectors.shape[0] / 1000))  # 每1000个向量一个聚类中心
+        quantizer = faiss.IndexFlatIP(self.vector_dim)
+        self.index = faiss.IndexIVFFlat(quantizer, self.vector_dim, nlist, faiss.METRIC_INNER_PRODUCT)
+        
+        # Train and add vectors
+        self.index.train(vectors) # type: ignore
+        self.index.add(vectors) # type: ignore
+        # Set the number of clusters to probe during search
+        self.index.nprobe = min(nlist, 10)
+
+    def _split_text(self, text: str) -> List[str]:
+        """Use a more intelligent splitting strategy"""
+        # Add overlapping blocks to maintain context consistency
+        overlap_size = min(200, self.max_paragraph_length // 4)
+        
+        paragraphs = []
+        current_chunk = []
+        current_length = 0
+        
+        # First split by sentence
+        sentences = []
+        current_sentence = []
+        sentence_ends = {'。', '！', '？', '…', '.', '!', '?'}
+        
+        for char in text:
+            current_sentence.append(char)
+            if char in sentence_ends:
+                sentence = ''.join(current_sentence)
+                if sentence.strip():
+                    sentences.append(sentence)
+                current_sentence = []
+        
+        if current_sentence:
+            sentence = ''.join(current_sentence)
+            if sentence.strip():
+                sentences.append(sentence)
+        
+        # Build overlapping blocks based on sentences
+        for sentence in sentences:
+            if current_length + len(sentence) > self.max_paragraph_length:
+                if current_chunk:
+                    chunk_text = ' '.join(current_chunk)
+                    if len(chunk_text) >= self.min_paragraph_length:
+                        paragraphs.append(chunk_text)
+                        
+                    # Keep some content as overlap
+                    overlap_text = ' '.join(current_chunk[-2:])  # Keep the last two sentences
+                    current_chunk = []
+                    if overlap_text:
+                        current_chunk.append(overlap_text)
+                        current_length = len(overlap_text)
+                    else:
+                        current_length = 0
+                        
+            current_chunk.append(sentence)
+            current_length += len(sentence)
+        
+        # Process the last chunk
+        if current_chunk:
+            chunk_text = ' '.join(current_chunk)
+            if len(chunk_text) >= self.min_paragraph_length:
+                paragraphs.append(chunk_text)
+        
+        return paragraphs
+
+    def _get_embedding(self, text: str) -> np.ndarray:
+        """Get the vector representation of the text"""
+        embedding = self.embedding_model.encode(text, 
+                                            normalize_embeddings=True,
+                                            show_progress_bar=False)
+        return np.array(embedding, dtype=np.float32)
+
+    def _get_embedding_batch(self, texts: List[str]) -> np.ndarray:
+        """Get the vector representation of the text batch
+        
+        Args:
+            texts: Text list
+            
+        Returns:
+            np.ndarray: Vector representation array
+        """
+        try:
+            embeddings = self.embedding_model.encode(texts, 
+                                                normalize_embeddings=True,
+                                                show_progress_bar=False,
+                                                batch_size=32)  # Use batch processing to improve efficiency
+            return np.array(embeddings, dtype=np.float32)
+        except Exception as e:
+            PrettyOutput.print(f"Failed to get vector representation: {str(e)}", 
+                            output_type=OutputType.ERROR)
+            return np.zeros((len(texts), self.vector_dim), dtype=np.float32) # type: ignore
+
+    def _process_document_batch(self, documents: List[Document]) -> List[np.ndarray]:
+        """Process a batch of documents vectorization
+        
+        Args:
+            documents: Document list
+            
+        Returns:
+            List[np.ndarray]: Vector list
+        """
+        texts = []
+        for doc in documents:
+            # Combine document information
+            combined_text = f"""
+File: {doc.metadata['file_path']}
+Content: {doc.content}
+"""
+            texts.append(combined_text)
+            
+        return self._get_embedding_batch(texts) # type: ignore
+
+    def _process_file(self, file_path: str) -> List[Document]:
+        """Process a single file"""
+        try:
+            # Calculate file MD5
+            current_md5 = get_file_md5(file_path)
+            if not current_md5:
+                return []
+
+            # Check if the file needs to be reprocessed
+            if file_path in self.file_md5_cache and self.file_md5_cache[file_path] == current_md5:
+                return []
+
+            # Find the appropriate processor
+            processor = None
+            for p in self.file_processors:
+                if p.can_handle(file_path):
+                    processor = p
+                    break
+                    
+            if not processor:
+                # If no appropriate processor is found, return an empty document
+                return []
+            
+            # Extract text content
+            content = processor.extract_text(file_path)
+            if not content.strip():
+                return []
+            
+            # Split text
+            chunks = self._split_text(content)
+            
+            # Create document objects
+            documents = []
+            for i, chunk in enumerate(chunks):
+                doc = Document(
+                    content=chunk,
+                    metadata={
+                        "file_path": file_path,
+                        "file_type": Path(file_path).suffix.lower(),
+                        "chunk_index": i,
+                        "total_chunks": len(chunks)
+                    },
+                    md5=current_md5
+                )
+                documents.append(doc)
+            
+            # Update MD5 cache
+            self.file_md5_cache[file_path] = current_md5
+            return documents
+            
+        except Exception as e:
+            PrettyOutput.print(f"Failed to process file {file_path}: {str(e)}", 
+                            output_type=OutputType.ERROR)
+            return []
+
+    def build_index(self, dir: str):
+        """Build document index"""
+        # Get all files
+        all_files = []
+        for root, _, files in os.walk(dir):
+            if any(ignored in root for ignored in ['.git', '__pycache__', 'node_modules']) or \
+               any(part.startswith('.jarvis-') for part in root.split(os.sep)):
+                continue
+            for file in files:
+                if file.startswith('.jarvis-'):
+                    continue
+                    
+                file_path = os.path.join(root, file)
+                if os.path.getsize(file_path) > 100 * 1024 * 1024:  # 100MB
+                    PrettyOutput.print(f"Skip large file: {file_path}", 
+                                    output_type=OutputType.WARNING)
+                    continue
+                all_files.append(file_path)
+
+        # Clean up cache for deleted files
+        deleted_files = set(self.file_md5_cache.keys()) - set(all_files)
+        for file_path in deleted_files:
+            del self.file_md5_cache[file_path]
+            # Remove related documents
+            self.documents = [doc for doc in self.documents if doc.metadata['file_path'] != file_path]
+
+        # Check file changes
+        files_to_process = []
+        unchanged_files = []
+        
+        with tqdm(total=len(all_files), desc="Check file status") as pbar:
+            for file_path in all_files:
+                current_md5 = get_file_md5(file_path)
+                if current_md5:  # Only process files that can successfully calculate MD5
+                    if file_path in self.file_md5_cache and self.file_md5_cache[file_path] == current_md5:
+                        # File未变化，记录但不重新处理
+                        unchanged_files.append(file_path)
+                    else:
+                        # New file or modified file
+                        files_to_process.append(file_path)
+                pbar.update(1)
+
+        # Keep documents for unchanged files
+        unchanged_documents = [doc for doc in self.documents 
+                            if doc.metadata['file_path'] in unchanged_files]
+
+        # Process new files and modified files
+        new_documents = []
+        if files_to_process:
+            with tqdm(total=len(files_to_process), desc="Process files") as pbar:
+                for file_path in files_to_process:
+                    try:
+                        docs = self._process_file(file_path)
+                        if len(docs) > 0:
+                            new_documents.extend(docs)
+                    except Exception as e:
+                        PrettyOutput.print(f"Failed to process file {file_path}: {str(e)}", 
+                                        output_type=OutputType.ERROR)
+                    pbar.update(1)
+
+        # Update document list
+        self.documents = unchanged_documents + new_documents
+
+        if not self.documents:
+            PrettyOutput.print("No documents to process", output_type=OutputType.WARNING)
+            return
+
+        # Only vectorize new documents
+        if new_documents:
+            PrettyOutput.print(f"Start processing {len(new_documents)} new documents", 
+                            output_type=OutputType.INFO)
+            
+            # Use thread pool to process vectorization
+            batch_size = 32
+            new_vectors = []
+            
+            with tqdm(total=len(new_documents), desc="Generating vectors") as pbar:
+                with ThreadPoolExecutor(max_workers=self.thread_count) as executor:
+                    for i in range(0, len(new_documents), batch_size):
+                        batch = new_documents[i:i + batch_size]
+                        future = executor.submit(self._process_document_batch, batch)
+                        batch_vectors = future.result()
+                        
+                        with self.vector_lock:
+                            new_vectors.extend(batch_vectors)
+                        
+                        pbar.update(len(batch))
+
+            # Merge new and old vectors
+            if self.flat_index is not None:
+                # Get vectors for unchanged documents
+                unchanged_vectors = []
+                for doc in unchanged_documents:
+                    # Get vectors from existing index
+                    doc_idx = next((i for i, d in enumerate(self.documents) 
+                                if d.metadata['file_path'] == doc.metadata['file_path']), None)
+                    if doc_idx is not None:
+                        # Reconstruct vectors from flat index
+                        vector = np.zeros((1, self.vector_dim), dtype=np.float32) # type: ignore
+                        self.flat_index.reconstruct(doc_idx, vector.ravel())
+                        unchanged_vectors.append(vector)
+                
+                if unchanged_vectors:
+                    unchanged_vectors = np.vstack(unchanged_vectors)
+                    vectors = np.vstack([unchanged_vectors, np.vstack(new_vectors)])
+                else:
+                    vectors = np.vstack(new_vectors)
+            else:
+                vectors = np.vstack(new_vectors)
+
+            # Build index
+            self._build_index(vectors)
+            # Save cache
+            self._save_cache(vectors)
+        
+        PrettyOutput.print(f"Successfully indexed {len(self.documents)} document fragments (Added/Modified: {len(new_documents)}, Unchanged: {len(unchanged_documents)})", 
+                        output_type=OutputType.SUCCESS)
+
+    def search(self, query: str, top_k: int = 30) -> List[Tuple[Document, float]]:
+        """Optimize search strategy"""
+        if not self.index:
+            PrettyOutput.print("Index not built, building...", output_type=OutputType.INFO)
+            self.build_index(self.root_dir)
+        
+        # Implement MMR (Maximal Marginal Relevance) to increase result diversity
+        def mmr(query_vec, doc_vecs, doc_ids, lambda_param=0.5, n_docs=top_k):
+            selected = []
+            selected_ids = []
+            
+            while len(selected) < n_docs and len(doc_ids) > 0:
+                best_score = -1
+                best_idx = -1
+                
+                for i, (doc_vec, doc_id) in enumerate(zip(doc_vecs, doc_ids)):
+                    # Calculate similarity with query
+                    query_sim = float(np.dot(query_vec, doc_vec))
+                    
+                    # Calculate maximum similarity with selected documents
+                    if selected:
+                        doc_sims = [float(np.dot(doc_vec, selected_doc)) for selected_doc in selected]
+                        max_doc_sim = max(doc_sims)
+                    else:
+                        max_doc_sim = 0
+                    
+                    # MMR score
+                    score = lambda_param * query_sim - (1 - lambda_param) * max_doc_sim
+                    
+                    if score > best_score:
+                        best_score = score
+                        best_idx = i
+                
+                if best_idx == -1:
+                    break
+                
+                selected.append(doc_vecs[best_idx])
+                selected_ids.append(doc_ids[best_idx])
+                doc_vecs = np.delete(doc_vecs, best_idx, axis=0)
+                doc_ids = np.delete(doc_ids, best_idx)
+            
+            return selected_ids
+        
+        # Get query vector
+        query_vector = self._get_embedding(query)
+        query_vector = query_vector.reshape(1, -1)
+        
+        # Initial search more results for MMR
+        initial_k = min(top_k * 2, len(self.documents))
+        distances, indices = self.index.search(query_vector, initial_k) # type: ignore
+        
+        # Get valid results
+        valid_indices = indices[0][indices[0] != -1]
+        valid_vectors = np.vstack([self._get_embedding(self.documents[idx].content) for idx in valid_indices])
+        
+        # Apply MMR
+        final_indices = mmr(query_vector[0], valid_vectors, valid_indices, n_docs=top_k)
+        
+        # Build results
+        results = []
+        for idx in final_indices:
+            doc = self.documents[idx]
+            similarity = 1.0 / (1.0 + float(distances[0][np.where(indices[0] == idx)[0][0]]))
+            results.append((doc, similarity))
+        
+        return results
+
+    def _rerank_results(self, query: str, initial_results: List[Tuple[Document, float]]) -> List[Tuple[Document, float]]:
+        """Use rerank model to rerank search results"""
+        try:
+            import torch
+            model, tokenizer = load_rerank_model()
+            
+            # Prepare data
+            pairs = []
+            for doc, _ in initial_results:
+                # Combine document information
+                doc_content = f"""
+File: {doc.metadata['file_path']}
+Content: {doc.content}
+"""
+                pairs.append([query, doc_content])
+                
+            # Score each document pair
+            scores = []
+            batch_size = 8
+            
+            with torch.no_grad():
+                for i in range(0, len(pairs), batch_size):
+                    batch_pairs = pairs[i:i + batch_size]
+                    encoded = tokenizer(
+                        batch_pairs,
+                        padding=True,
+                        truncation=True,
+                        max_length=512,
+                        return_tensors='pt'
+                    )
+                    
+                    if torch.cuda.is_available():
+                        encoded = {k: v.cuda() for k, v in encoded.items()}
+                        
+                    outputs = model(**encoded)
+                    batch_scores = outputs.logits.squeeze(-1).cpu().numpy()
+                    scores.extend(batch_scores.tolist())
+            
+            # Normalize scores to 0-1 range
+            if scores:
+                min_score = min(scores)
+                max_score = max(scores)
+                if max_score > min_score:
+                    scores = [(s - min_score) / (max_score - min_score) for s in scores]
+            
+            # Combine scores with documents and sort
+            scored_results = []
+            for (doc, _), score in zip(initial_results, scores):
+                if score >= 0.5:  # Only keep results with a score greater than 0.5
+                    scored_results.append((doc, float(score)))
+                    
+            # Sort by score in descending order
+            scored_results.sort(key=lambda x: x[1], reverse=True)
+            
+            return scored_results
+            
+        except Exception as e:
+            PrettyOutput.print(f"Failed to rerank, using original sorting: {str(e)}", output_type=OutputType.WARNING)
+            return initial_results
+
+    def is_index_built(self):
+        """Check if index is built"""
+        return self.index is not None
+
+    def query(self, query: str) -> List[Document]:
+        """Query related documents
+        
+        Args:
+            query: Query text
+            
+        Returns:
+            List[Document]: Related documents, including context
+        """
+        results = self.search(query)
+        return [doc for doc, _ in results]
+
+    def ask(self, question: str) -> Optional[str]:
+        """Ask about documents
+        
+        Args:
+            question: User question
+            
+        Returns:
+            Model answer, return None if failed
+        """
+        try:
+            # Search related document fragments
+            results = self.query(question)
+            if not results:
+                return None
+            
+            # Display found document fragments
+            for doc in results:
+                output = f"""File: {doc.metadata['file_path']}\n"""
+                output += f"""Fragment {doc.metadata['chunk_index'] + 1}/{doc.metadata['total_chunks']}\n"""
+                output += f"""Content:\n{doc.content}\n"""
+                PrettyOutput.print(output, output_type=OutputType.INFO, lang="markdown")
+
+            # Build base prompt
+            base_prompt = f"""Please answer the user's question based on the following document fragments. If the document content is not sufficient to answer the question completely, please clearly indicate.
+
+User question: {question}
+
+Related document fragments:
+"""
+            end_prompt = "\nPlease provide an accurate and concise answer. If the document content is not sufficient to answer the question completely, please clearly indicate."
+            
+            # Calculate the maximum length that can be used for document content
+            # Leave some space for the model's answer
+            available_length = self.max_context_length - len(base_prompt) - len(end_prompt) - 500
+            
+            # Build context, while controlling the total length
+            context = []
+            current_length = 0
+            
+            for doc in results:
+                # Calculate the length of this document fragment's content
+                doc_content = f"""
+Source file: {doc.metadata['file_path']}
+Content:
+{doc.content}
+---
+"""
+                content_length = len(doc_content)
+                
+                # If adding this fragment would exceed the limit, stop adding
+                if current_length + content_length > available_length:
+                    PrettyOutput.print("Due to context length limit, some related document fragments were omitted", 
+                                    output_type=OutputType.WARNING)
+                    break
+                    
+                context.append(doc_content)
+                current_length += content_length
+
+            # Build complete prompt
+            prompt = base_prompt + ''.join(context) + end_prompt
+            
+            # Get model instance and generate answer
+            model = PlatformRegistry.get_global_platform_registry().get_normal_platform()
+            response = model.chat_until_success(prompt)
+            
+            return response
+            
+        except Exception as e:
+            PrettyOutput.print(f"Failed to answer: {str(e)}", output_type=OutputType.ERROR)
+            return None
+
+def main():
+    """Main function"""
+    import argparse
+    import sys
+    
+    # Set standard output encoding to UTF-8
+    if sys.stdout.encoding != 'utf-8':
+        import codecs
+        sys.stdout = codecs.getwriter('utf-8')(sys.stdout.buffer, 'strict')
+        sys.stderr = codecs.getwriter('utf-8')(sys.stderr.buffer, 'strict')
+    
+    parser = argparse.ArgumentParser(description='Document retrieval and analysis tool')
+    parser.add_argument('--dir', type=str, help='Directory to process')
+    parser.add_argument('--build', action='store_true', help='Build document index')
+    parser.add_argument('--search', type=str, help='Search document content')
+    parser.add_argument('--ask', type=str, help='Ask about documents')
+    args = parser.parse_args()
+
+    try:
+        current_dir = os.getcwd()
+        rag = RAGTool(current_dir)
+
+        if not args.dir:
+            args.dir = current_dir
+
+        if args.dir and args.build:
+            PrettyOutput.print(f"Processing directory: {args.dir}", output_type=OutputType.INFO)
+            rag.build_index(args.dir)
+            return 0
+
+        if args.search or args.ask:
+
+            if args.search:
+                results = rag.query(args.search)
+                if not results:
+                    PrettyOutput.print("No related content found", output_type=OutputType.WARNING)
+                    return 1
+                    
+                for doc in results:
+                    output = f"""File: {doc.metadata['file_path']}\n"""
+                    output += f"""Fragment {doc.metadata['chunk_index'] + 1}/{doc.metadata['total_chunks']}\n"""
+                    output += f"""Content:\n{doc.content}\n"""
+                    PrettyOutput.print(output, output_type=OutputType.INFO, lang="markdown")
+                return 0
+
+            if args.ask:
+                # Call ask method
+                response = rag.ask(args.ask)
+                if not response:
+                    PrettyOutput.print("Failed to get answer", output_type=OutputType.WARNING)
+                    return 1
+                    
+                # Display answer
+                output = f"""Answer:\n{response}"""
+                PrettyOutput.print(output, output_type=OutputType.INFO)
+                return 0
+
+        PrettyOutput.print("Please specify operation parameters. Use -h to view help.", output_type=OutputType.WARNING)
+        return 1
+
+    except Exception as e:
+        PrettyOutput.print(f"Failed to execute: {str(e)}", output_type=OutputType.ERROR)
+        return 1
+
+if __name__ == "__main__":
+    main()