jarvis-ai-assistant 0.1.76__py3-none-any.whl → 0.1.77__py3-none-any.whl

jarvis/__init__.py

@@ -1,3 +1,3 @@
 """Jarvis AI Assistant"""
 
-__version__ = "0.1.76" 
+__version__ = "0.1.77" 

jarvis/agent.py

@@ -8,7 +8,7 @@ import json
 
 from .models.registry import PlatformRegistry
 from .tools import ToolRegistry
-from .utils import PrettyOutput, OutputType, get_multiline_input, while_success
+from .utils import PrettyOutput, OutputType, get_multiline_input, load_embedding_model, while_success
 import os
 from datetime import datetime
 from prompt_toolkit import prompt
@@ -37,15 +37,15 @@ class Agent:
         
         # 初始化嵌入模型
         try:
-            os.environ["TOKENIZERS_PARALLELISM"] = "false"
-            PrettyOutput.print(f"正在加载嵌入模型: {self.embedding_model_name}...", OutputType.INFO)
-            self.embedding_model = SentenceTransformer(self.embedding_model_name)
+            self.embedding_model = load_embedding_model(self.embedding_model_name)
             
             # 预热模型并获取正确的维度
             test_text = "这是一段测试文本，用于确保模型完全加载。"
-            test_embedding = self.embedding_model.encode(test_text, 
-                                                      convert_to_tensor=True,
-                                                      normalize_embeddings=True)
+            test_embedding = self.embedding_model.encode(
+                test_text, 
+                convert_to_tensor=True,
+                normalize_embeddings=True
+            )
             self.embedding_dimension = len(test_embedding)
             PrettyOutput.print("嵌入模型加载完成", OutputType.SUCCESS)
             
@@ -291,37 +291,6 @@ class Agent:
         self.prompt = summary_prompt
         return self._call_model(self.prompt)
 
-    def choose_tools(self, user_input: str) -> List[Dict]:
-        """根据用户输入选择工具"""
-        PrettyOutput.print("选择工具...", OutputType.PLANNING)
-        tools = self.tool_registry.get_all_tools()
-        prompt = f"""你是一个工具选择专家，请根据用户输入选择合适的工具，返回可能使用到的工具的名称。以下是可用工具：
-"""
-        for tool in tools:
-            prompt += f"- {tool['name']}: {tool['description']}\n"
-        prompt += f"用户输入: {user_input}\n"
-        prompt += f"请返回可能使用到的工具的名称，如果无法确定，请返回空列表。"
-        prompt += f"返回的格式为：\n"
-        prompt += f"<TOOL_CHOICE_START>\n"
-        prompt += f"tool_name1\n"
-        prompt += f"tool_name2\n"
-        prompt += f"<TOOL_CHOICE_END>\n"
-        model = PlatformRegistry.get_global_platform()
-        model.set_suppress_output(True)
-        try:
-            response = model.chat(prompt)
-            response = response.replace("<TOOL_CHOICE_START>", "").replace("<TOOL_CHOICE_END>", "")
-            tools_name = response.split("\n")
-            choosed_tools = []
-            for tool_name in tools_name:
-                for tool in tools:
-                    if tool['name'] == tool_name:
-                        choosed_tools.append(tool)
-                        break
-            return choosed_tools
-        except Exception as e:
-            PrettyOutput.print(f"工具选择失败: {str(e)}", OutputType.ERROR)
-            return []
 
     def run(self, user_input: str, file_list: Optional[List[str]] = None, keep_history: bool = False) -> str:
         """处理用户输入并返回响应，返回任务总结报告
@@ -350,7 +319,7 @@ class Agent:
             tools_prompt = ""
 
             # 选择工具
-            tools = self.choose_tools(user_input)
+            tools = self.tool_registry.get_all_tools()
             if tools:
                 tools_prompt += "可用工具:\n"
                 for tool in tools:

jarvis/jarvis_codebase/main.py

@@ -7,7 +7,7 @@ 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
+from jarvis.utils import OutputType, PrettyOutput, find_git_root, load_embedding_model
 from jarvis.utils import load_env_from_file
 import argparse
 from sentence_transformers import SentenceTransformer
@@ -43,9 +43,7 @@ class CodeBase:
             
         # 初始化嵌入模型，使用系统默认缓存目录
         try:
-            os.environ["TOKENIZERS_PARALLELISM"] = "false"
-            PrettyOutput.print("正在加载/下载模型，请稍候...", output_type=OutputType.INFO)
-            self.embedding_model = SentenceTransformer(self.embedding_model_name)
+            self.embedding_model = load_embedding_model(self.embedding_model_name)
             
             # 强制完全加载所有模型组件
             test_text = """

jarvis/rag/main.py

@@ -0,0 +1,483 @@
+import os
+import hashlib
+import numpy as np
+import faiss
+from typing import List, Tuple, Optional, Dict
+from sentence_transformers import SentenceTransformer
+import pickle
+from jarvis.utils import OutputType, PrettyOutput, find_git_root, load_embedding_model
+from jarvis.utils import load_env_from_file
+import tiktoken
+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
+
+@dataclass
+class Document:
+    """文档类，用于存储文档内容和元数据"""
+    content: str  # 文档内容
+    metadata: Dict  # 元数据(文件路径、位置等)
+
+class FileProcessor:
+    """文件处理器基类"""
+    @staticmethod
+    def can_handle(file_path: str) -> bool:
+        """判断是否可以处理该文件"""
+        raise NotImplementedError
+        
+    @staticmethod
+    def extract_text(file_path: str) -> str:
+        """提取文件文本内容"""
+        raise NotImplementedError
+
+class TextFileProcessor(FileProcessor):
+    """文本文件处理器"""
+    ENCODINGS = ['utf-8', 'gbk', 'gb2312', 'latin1']
+    SAMPLE_SIZE = 8192  # 读取前8KB来检测编码
+    
+    @staticmethod
+    def can_handle(file_path: str) -> bool:
+        """判断文件是否为文本文件，通过尝试解码来判断"""
+        try:
+            # 读取文件开头的一小部分来检测
+            with open(file_path, 'rb') as f:
+                sample = f.read(TextFileProcessor.SAMPLE_SIZE)
+                
+            # 检查是否包含空字节（通常表示二进制文件）
+            if b'\x00' in sample:
+                return False
+                
+            # 检查是否包含过多的非打印字符（通常表示二进制文件）
+            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:  # 如果非打印字符超过30%，认为是二进制文件
+                return False
+                
+            # 尝试用不同编码解码
+            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:
+        """提取文本内容，使用检测到的正确编码"""
+        detected_encoding = None
+        try:
+            # 首先尝试检测编码
+            with open(file_path, 'rb') as f:
+                raw_data = f.read()
+                
+            # 尝试不同的编码
+            for encoding in TextFileProcessor.ENCODINGS:
+                try:
+                    raw_data.decode(encoding)
+                    detected_encoding = encoding
+                    break
+                except UnicodeDecodeError:
+                    continue
+                    
+            if not detected_encoding:
+                raise UnicodeDecodeError(f"无法用支持的编码解码文件: {file_path}")
+                
+            # 使用检测到的编码读取文件
+            with open(file_path, 'r', encoding=detected_encoding, errors='replace') as f:
+                content = f.read()
+                
+            # 规范化Unicode字符
+            import unicodedata
+            content = unicodedata.normalize('NFKC', content)
+            
+            return content
+            
+        except Exception as e:
+            raise Exception(f"读取文件失败: {str(e)}")
+
+class PDFProcessor(FileProcessor):
+    """PDF文件处理器"""
+    @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:
+            for page in doc:
+                text_parts.append(page.get_text())
+        return "\n".join(text_parts)
+
+class DocxProcessor(FileProcessor):
+    """DOCX文件处理器"""
+    @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):
+        """初始化RAG工具
+        
+        Args:
+            root_dir: 项目根目录
+        """
+        load_env_from_file()
+        self.root_dir = root_dir
+        os.chdir(self.root_dir)
+        
+        # 初始化配置
+        self.min_paragraph_length = int(os.environ.get("JARVIS_MIN_PARAGRAPH_LENGTH", "50"))  # 最小段落长度
+        self.max_paragraph_length = int(os.environ.get("JARVIS_MAX_PARAGRAPH_LENGTH", "1000"))  # 最大段落长度
+        self.embedding_model_name = os.environ.get("JARVIS_EMBEDDING_MODEL", "BAAI/bge-large-zh-v1.5")
+        
+        # 初始化数据目录
+        self.data_dir = os.path.join(self.root_dir, ".jarvis-rag")
+        if not os.path.exists(self.data_dir):
+            os.makedirs(self.data_dir)
+            
+        # 初始化嵌入模型
+        try:
+            self.embedding_model = load_embedding_model(self.embedding_model_name)
+            self.vector_dim = self.embedding_model.get_sentence_embedding_dimension()
+            PrettyOutput.print("模型加载完成", output_type=OutputType.SUCCESS)
+        except Exception as e:
+            PrettyOutput.print(f"加载模型失败: {str(e)}", output_type=OutputType.ERROR)
+            raise
+
+        # 初始化缓存和索引
+        self.cache_path = os.path.join(self.data_dir, "cache.pkl")
+        self.documents: List[Document] = []
+        self.index = None
+        
+        # 加载缓存
+        self._load_cache()
+
+        # 注册文件处理器
+        self.file_processors = [
+            TextFileProcessor(),
+            PDFProcessor(),
+            DocxProcessor()
+        ]
+
+    def _load_cache(self):
+        """加载缓存数据"""
+        if os.path.exists(self.cache_path):
+            try:
+                with open(self.cache_path, 'rb') as f:
+                    cache_data = pickle.load(f)
+                    self.documents = cache_data["documents"]
+                    vectors = cache_data["vectors"]
+                    
+                # 重建索引
+                self._build_index(vectors)
+                PrettyOutput.print(f"加载了 {len(self.documents)} 个文档片段", 
+                                output_type=OutputType.INFO)
+            except Exception as e:
+                PrettyOutput.print(f"加载缓存失败: {str(e)}", 
+                                output_type=OutputType.WARNING)
+                self.documents = []
+                self.index = None
+
+    def _save_cache(self, vectors: np.ndarray):
+        """保存缓存数据"""
+        try:
+            cache_data = {
+                "documents": self.documents,
+                "vectors": vectors
+            }
+            with open(self.cache_path, 'wb') as f:
+                pickle.dump(cache_data, f)
+            PrettyOutput.print(f"保存了 {len(self.documents)} 个文档片段", 
+                            output_type=OutputType.INFO)
+        except Exception as e:
+            PrettyOutput.print(f"保存缓存失败: {str(e)}", 
+                            output_type=OutputType.ERROR)
+
+    def _build_index(self, vectors: np.ndarray):
+        """构建FAISS索引"""
+        # 创建HNSW索引
+        hnsw_index = faiss.IndexHNSWFlat(self.vector_dim, 16)
+        hnsw_index.hnsw.efConstruction = 40
+        hnsw_index.hnsw.efSearch = 16
+        
+        # 用IndexIDMap包装HNSW索引
+        self.index = faiss.IndexIDMap(hnsw_index)
+        
+        # 添加向量到索引
+        if vectors.shape[0] > 0:
+            self.index.add_with_ids(vectors, np.arange(vectors.shape[0]))
+        else:
+            self.index = None
+
+    def _split_text(self, text: str) -> List[str]:
+        """将文本分割成段落
+        
+        Args:
+            text: 要分割的文本
+            
+        Returns:
+            分割后的段落列表
+        """
+        # 首先按空行分割
+        paragraphs = []
+        current_paragraph = []
+        
+        for line in text.split('\n'):
+            line = line.strip()
+            if not line:  # 空行表示段落结束
+                if current_paragraph:
+                    paragraph_text = ' '.join(current_paragraph)
+                    if len(paragraph_text) >= self.min_paragraph_length:
+                        paragraphs.append(paragraph_text)
+                    current_paragraph = []
+            else:
+                current_paragraph.append(line)
+        
+        # 处理最后一个段落
+        if current_paragraph:
+            paragraph_text = ' '.join(current_paragraph)
+            if len(paragraph_text) >= self.min_paragraph_length:
+                paragraphs.append(paragraph_text)
+        
+        # 处理过长的段落
+        final_paragraphs = []
+        for paragraph in paragraphs:
+            if len(paragraph) <= self.max_paragraph_length:
+                final_paragraphs.append(paragraph)
+            else:
+                # 按句子分割过长的段落
+                sentences = []
+                current_sentence = []
+                
+                # 中文句子结束标记
+                sentence_ends = {'。', '！', '？', '…', '.', '!', '?'}
+                
+                for char in paragraph:
+                    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)
+                
+                # 组合句子成适当长度的段落
+                current_chunk = []
+                current_length = 0
+                
+                for sentence in sentences:
+                    sentence_length = len(sentence)
+                    if current_length + sentence_length > self.max_paragraph_length:
+                        if current_chunk:
+                            final_paragraphs.append(''.join(current_chunk))
+                        current_chunk = [sentence]
+                        current_length = sentence_length
+                    else:
+                        current_chunk.append(sentence)
+                        current_length += sentence_length
+                
+                # 处理最后一个chunk
+                if current_chunk:
+                    final_paragraphs.append(''.join(current_chunk))
+        
+        # 过滤掉太短的段落
+        final_paragraphs = [p for p in final_paragraphs if len(p) >= self.min_paragraph_length]
+        
+        return final_paragraphs
+
+    def _get_embedding(self, text: str) -> np.ndarray:
+        """获取文本的向量表示"""
+        embedding = self.embedding_model.encode(text, 
+                                            normalize_embeddings=True,
+                                            show_progress_bar=False)
+        return np.array(embedding, dtype=np.float32)
+
+    def _process_file(self, file_path: str) -> List[Document]:
+        """处理单个文件
+        
+        Args:
+            file_path: 文件路径
+            
+        Returns:
+            文档对象列表
+        """
+        try:
+            # 查找合适的处理器
+            processor = None
+            for p in self.file_processors:
+                if p.can_handle(file_path):
+                    processor = p
+                    break
+                    
+            if not processor:
+                PrettyOutput.print(f"跳过不支持的文件: {file_path}", 
+                                output_type=OutputType.WARNING)
+                return []
+            
+            # 提取文本内容
+            content = processor.extract_text(file_path)
+            if not content.strip():
+                PrettyOutput.print(f"文件内容为空: {file_path}", 
+                                output_type=OutputType.WARNING)
+                return []
+            
+                
+            # 分割文本
+            chunks = self._split_text(content)
+            
+            # 创建文档对象
+            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)
+                    }
+                )
+                documents.append(doc)
+                
+            return documents
+            
+        except Exception as e:
+            PrettyOutput.print(f"处理文件失败 {file_path}: {str(e)}", 
+                            output_type=OutputType.ERROR)
+            return []
+
+    def build_index(self):
+        """构建文档索引"""
+        # 获取所有文件
+        all_files = []
+        for root, _, files in os.walk(self.root_dir):
+            if any(ignored in root for ignored in ['.jarvis-rag', '.git', '__pycache__', 'node_modules']):
+                continue
+            for file in files:
+                file_path = os.path.join(root, file)
+                # 跳过大文件
+                if os.path.getsize(file_path) > 10 * 1024 * 1024:  # 10MB
+                    PrettyOutput.print(f"跳过大文件: {file_path}", 
+                                    output_type=OutputType.WARNING)
+                    continue
+                all_files.append(file_path)
+
+        # 处理所有文件
+        self.documents = []
+        for file_path in tqdm(all_files, desc="处理文件"):
+            docs = self._process_file(file_path)
+            self.documents.extend(docs)
+
+        # 获取所有文档的向量表示
+        vectors = []
+        for doc in tqdm(self.documents, desc="生成向量"):
+            vector = self._get_embedding(doc.content)
+            vectors.append(vector)
+
+        if vectors:
+            vectors = np.vstack(vectors)
+            # 构建索引
+            self._build_index(vectors)
+            # 保存缓存
+            self._save_cache(vectors)
+            
+        PrettyOutput.print(f"成功索引了 {len(self.documents)} 个文档片段", 
+                        output_type=OutputType.SUCCESS)
+
+    def search(self, query: str, top_k: int = 5) -> List[Tuple[Document, float]]:
+        """搜索相关文档
+        
+        Args:
+            query: 查询文本
+            top_k: 返回结果数量
+            
+        Returns:
+            文档和相似度得分的列表
+        """
+        if not self.index:
+            raise ValueError("索引未构建，请先调用build_index()")
+            
+        # 获取查询的向量表示
+        query_vector = self._get_embedding(query)
+        query_vector = query_vector.reshape(1, -1)
+        
+        # 搜索最相似的向量
+        distances, indices = self.index.search(query_vector, top_k)
+        
+        # 返回结果
+        results = []
+        for idx, distance in zip(indices[0], distances[0]):
+            if idx == -1:  # FAISS返回-1表示无效结果
+                continue
+            similarity = 1.0 / (1.0 + float(distance))
+            results.append((self.documents[idx], similarity))
+                
+        return results
+
+def main():
+    """命令行入口"""
+    import argparse
+    import sys
+    
+    # 设置标准输出编码为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='RAG工具')
+    parser.add_argument('--dir', type=str, default=os.getcwd(), help='项目根目录')
+    parser.add_argument('--build', action='store_true', help='构建索引')
+    parser.add_argument('--query', type=str, help='搜索查询')
+    parser.add_argument('--top-k', type=int, default=5, help='返回结果数量')
+    
+    args = parser.parse_args()
+    
+    try:
+        rag = RAGTool(args.dir)
+        
+        if args.build:
+            rag.build_index()
+            
+        if args.query:
+            results = rag.search(args.query, args.top_k)
+            
+            if not results:
+                PrettyOutput.print("未找到相关内容", output_type=OutputType.WARNING)
+                return
+                
+            PrettyOutput.print("\n搜索结果:", output_type=OutputType.INFO)
+            for doc, score in results:
+                PrettyOutput.print("\n" + "="*50, output_type=OutputType.INFO)
+                PrettyOutput.print(f"文件: {doc.metadata['file_path']}", output_type=OutputType.INFO)
+                PrettyOutput.print(f"相似度: {score:.3f}", output_type=OutputType.INFO)
+                PrettyOutput.print(f"片段 {doc.metadata['chunk_index'] + 1}/{doc.metadata['total_chunks']}", 
+                                output_type=OutputType.INFO)
+                PrettyOutput.print("\n内容:", output_type=OutputType.INFO)
+                # 确保内容是UTF-8编码
+                content = doc.content.encode('utf-8', errors='replace').decode('utf-8')
+                PrettyOutput.print(content, output_type=OutputType.INFO)
+                
+    except Exception as e:
+        PrettyOutput.print(f"执行失败: {str(e)}", output_type=OutputType.ERROR)
+        return 1
+
+if __name__ == "__main__":
+    main()

jarvis/tools/coder.py

@@ -7,21 +7,21 @@ class CoderTool:
     """代码修改工具"""
     
     name = "coder"
-    description = "用于自动修改和生成代码的工具"
+    description = "分析并修改现有代码，用于实现新功能、修复bug、重构代码等。能理解代码上下文并进行精确的代码编辑。"
     parameters = {
         "feature": {
             "type": "string",
-            "description": "要实现的功能描述",
+            "description": "要实现的功能描述或需要修改的内容，例如：'添加日志功能'、'修复内存泄漏'、'优化性能'等",
             "required": True
         },
         "dir": {
             "type": "string", 
-            "description": "项目根目录",
+            "description": "项目根目录，默认为当前目录",
             "required": False
         },
         "language": {
             "type": "string",
-            "description": "编程语言",
+            "description": "项目的主要编程语言，默认为python",
             "required": False
         }
     }

jarvis/tools/search.py

@@ -8,33 +8,58 @@ from urllib.parse import quote
 def bing_search(query):
     try:
         with sync_playwright() as p:
-            browser = p.chromium.launch()
-            page = browser.new_page()
-            page.goto(
-                f"https://www.bing.com/search?form=QBRE&q={quote(query)}&cc=US"
+            # 启动浏览器时设置参数
+            browser = p.chromium.launch(
+                headless=True,  # 无头模式
+                args=['--disable-gpu', '--no-sandbox', '--disable-dev-shm-usage']
             )
-
-            page.wait_for_selector("#b_results", timeout=10000)
             
+            # 创建新页面并设置超时
+            page = browser.new_page(
+                user_agent='Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/121.0.0.0 Safari/537.36',
+                viewport={'width': 1920, 'height': 1080}
+            )
+            
+            # 设置页面超时
+            page.set_default_timeout(60000)
+            
+            # 访问搜索页面
+            url = f"https://www.bing.com/search?q={quote(query)}&form=QBLH&sp=-1"
+            page.goto(url, wait_until="networkidle")
+            
+            # 等待搜索结果加载
+            page.wait_for_selector("#b_results", state="visible", timeout=30000)
+            
+            # 等待一下以确保结果完全加载
+            page.wait_for_timeout(1000)
+            
+            # 提取搜索结果
             summaries = page.evaluate("""() => {
-                const liElements = Array.from(
-                    document.querySelectorAll("#b_results > .b_algo")
-                );
-                return liElements.map((li) => {
-                    const abstractElement = li.querySelector(".b_caption > p");
-                    const linkElement = li.querySelector("a");
-                    const href = linkElement.getAttribute("href");
-                    const title = linkElement.textContent;
-                    const abstract = abstractElement ? abstractElement.textContent : "";
-                    return { href, title, abstract };
-                });
+                const results = [];
+                const elements = document.querySelectorAll("#b_results > .b_algo");
+                
+                for (const el of elements) {
+                    const titleEl = el.querySelector("h2");
+                    const linkEl = titleEl ? titleEl.querySelector("a") : null;
+                    const abstractEl = el.querySelector(".b_caption p");
+                    
+                    if (linkEl) {
+                        results.push({
+                            title: titleEl.innerText.trim(),
+                            href: linkEl.href,
+                            abstract: abstractEl ? abstractEl.innerText.trim() : ""
+                        });
+                    }
+                }
+                return results;
             }""")
             
             browser.close()
-            print(summaries)
             return summaries
+            
     except Exception as error:
-        print("An error occurred:", error)
+        PrettyOutput.print(f"搜索出错: {str(error)}", OutputType.ERROR)
+        return None
 
 class SearchTool:
     name = "search"
@@ -158,4 +183,46 @@ class SearchTool:
             return {
                 "success": False,
                 "error": f"搜索失败: {str(e)}"
-            } 
+            }
+
+def main():
+    """命令行直接运行搜索工具"""
+    import argparse
+    import sys
+    
+    parser = argparse.ArgumentParser(description='Bing搜索工具')
+    parser.add_argument('query', help='搜索关键词')
+    parser.add_argument('--max', type=int, default=5, help='最大结果数量(默认5)')
+    parser.add_argument('--url-only', action='store_true', help='只显示URL')
+    args = parser.parse_args()
+    
+    try:
+        PrettyOutput.print(f"正在搜索: {args.query}", OutputType.INFO)
+        
+        results = bing_search(args.query)
+        
+        if not results:
+            PrettyOutput.print("未找到搜索结果", OutputType.WARNING)
+            sys.exit(1)
+            
+        PrettyOutput.print(f"\n找到 {len(results)} 条结果:", OutputType.INFO)
+        
+        for i, result in enumerate(results[:args.max], 1):
+            PrettyOutput.print(f"\n{'-'*50}", OutputType.INFO)
+            if args.url_only:
+                PrettyOutput.print(f"{i}. {result['href']}", OutputType.INFO)
+            else:
+                PrettyOutput.print(f"{i}. {result['title']}", OutputType.INFO)
+                PrettyOutput.print(f"链接: {result['href']}", OutputType.INFO)
+                if result['abstract']:
+                    PrettyOutput.print(f"摘要: {result['abstract']}", OutputType.INFO)
+                    
+    except KeyboardInterrupt:
+        PrettyOutput.print("\n搜索已取消", OutputType.WARNING)
+        sys.exit(1)
+    except Exception as e:
+        PrettyOutput.print(f"执行出错: {str(e)}", OutputType.ERROR)
+        sys.exit(1)
+
+if __name__ == "__main__":
+    main() 

jarvis/utils.py

@@ -9,6 +9,7 @@ from colorama import Fore, Style as ColoramaStyle
 from prompt_toolkit import PromptSession
 from prompt_toolkit.styles import Style as PromptStyle
 from prompt_toolkit.formatted_text import FormattedText
+from sentence_transformers import SentenceTransformer
 
 # 初始化colorama
 colorama.init()
@@ -206,4 +207,27 @@ def find_git_root(dir="."):
     os.chdir(dir)
     ret = os.popen("git rev-parse --show-toplevel").read().strip()
     os.chdir(curr_dir)
-    return ret
+    return ret
+
+def load_embedding_model(model_name: str):
+    os.environ["TOKENIZERS_PARALLELISM"] = "false"
+    PrettyOutput.print(f"正在加载嵌入模型: {model_name}...", OutputType.INFO)
+    try:
+        # 首先尝试离线加载
+        embedding_model = SentenceTransformer(
+            model_name,
+            device="cpu",
+            cache_folder=os.path.expanduser("~/.cache/huggingface/hub"),
+            local_files_only=True
+        )
+        PrettyOutput.print("使用本地缓存加载模型成功", OutputType.SUCCESS)
+    except Exception as local_error:
+        PrettyOutput.print(f"本地加载失败，尝试在线下载: {str(local_error)}", OutputType.WARNING)
+        # 如果离线加载失败，尝试在线下载
+        embedding_model = SentenceTransformer(
+            model_name,
+            device="cpu",
+            cache_folder=os.path.expanduser("~/.cache/huggingface/hub")
+        )
+        PrettyOutput.print("模型下载并加载成功", OutputType.SUCCESS)
+    return embedding_model

{jarvis_ai_assistant-0.1.76.dist-info → jarvis_ai_assistant-0.1.77.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: jarvis-ai-assistant
-Version: 0.1.76
+Version: 0.1.77
 Summary: Jarvis: An AI assistant that uses tools to interact with the system
 Home-page: https://github.com/skyfireitdiy/Jarvis
 Author: skyfire
@@ -35,6 +35,7 @@ Classifier: Programming Language :: Python :: 3.8
 Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
+Classifier: Operating System :: POSIX :: Linux
 Requires-Python: >=3.8
 Description-Content-Type: text/markdown
 License-File: LICENSE
@@ -48,6 +49,10 @@ Requires-Dist: numpy>=1.24.0
 Requires-Dist: faiss-cpu>=1.8.0
 Requires-Dist: sentence-transformers>=2.2.2
 Requires-Dist: bs4>=0.0.1
+Requires-Dist: PyMuPDF>=1.21.0
+Requires-Dist: python-docx>=0.8.11
+Requires-Dist: tiktoken>=0.3.0
+Requires-Dist: tqdm>=4.65.0
 Provides-Extra: dev
 Requires-Dist: pytest; extra == "dev"
 Requires-Dist: black; extra == "dev"
@@ -124,6 +129,7 @@ Jarvis supports configuration through environment variables that can be set in t
 |---------|------|--------|------|
 | JARVIS_PLATFORM | AI platform to use, supports kimi/openai/ai8 etc | kimi | Yes |
 | JARVIS_MODEL | Model name to use | - | No |
+| JARVIS_THREAD_COUNT | Number of threads for parallel processing | 10 | No |
 | JARVIS_CODEGEN_PLATFORM | AI platform for code generation | Same as JARVIS_PLATFORM | No |
 | JARVIS_CODEGEN_MODEL | Model name for code generation | Same as JARVIS_MODEL | No |
 | JARVIS_CHEAP_PLATFORM | AI platform for cheap operations | Same as JARVIS_PLATFORM | No |
@@ -141,36 +147,47 @@ Jarvis supports configuration through environment variables that can be set in t
 
 ## 🎯 Usage
 
-### Basic Usage
+### Main Assistant
 ```bash
 jarvis
 ```
 
-
-### With Specific Model
+### Code Generation
 ```bash
-jarvis -p kimi  # Use Kimi platform
-jarvis -p openai  # Use OpenAI platform
-```
-
-### Code Modification
-```bash
-jarvis-coder --feature "Add new feature"  # Modify code to add new feature
+jarvis-coder
 ```
 
 ### Codebase Search
 ```bash
-jarvis-codebase --search "database connection"  # Search codebase
+# Generate codebase index
+jarvis-codebase --generate
+
+# Search similar code
+jarvis-codebase --search "your search query"
+
+# Ask questions about codebase
+jarvis-codebase --ask "your question"
 ```
 
-### Codebase Question
+### Document Analysis (RAG)
 ```bash
-jarvis-codebase --ask "How to use the database?"  # Ask about codebase
+# Build document index
+jarvis-rag --dir /path/to/documents --build
+
+# Search documents
+jarvis-rag --query "your search query"
 ```
 
-### Keep Chat History
+### Search Tool
 ```bash
-jarvis --keep-history  # Don't delete chat session after completion
+# Basic search
+jarvis-search "your query"
+
+# Show only URLs
+jarvis-search "your query" --url-only
+
+# Limit results
+jarvis-search "your query" --max 3
 ```
 
 ## 🛠️ Tools

{jarvis_ai_assistant-0.1.76.dist-info → jarvis_ai_assistant-0.1.77.dist-info}/RECORD

@@ -1,9 +1,9 @@
-jarvis/__init__.py,sha256=mqMOFwEAV4m463m3kW2Lz8KGIkWz1vJScSEYFQi4jsc,50
-jarvis/agent.py,sha256=10YxntRu9CWp-xZt7PmUYbzj_knwGM0QqyRbM4FaKSk,21473
+jarvis/__init__.py,sha256=WZAfXLVIclhkO-zAnoGnZ9Gc0F0C_LUd0MwIpGBFrdw,50
+jarvis/agent.py,sha256=lUYH9wijs4wwbXMxqZBtOkEfL0jHLv9rw82r9EjPpXs,19743
 jarvis/main.py,sha256=7EcSlxa5JFFXBujzKDWdNtwX6axLhFFdJMc2GxTjfdk,6295
-jarvis/utils.py,sha256=vZV8sHj0ggZy4Rb8RxIujQhRWgeNEomhqVl4WXmpq7c,7498
+jarvis/utils.py,sha256=TKKO7vbWLJ80n3JqL30HpbW-GuFfJTgR0GqBrsPPiwk,8566
 jarvis/jarvis_codebase/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-jarvis/jarvis_codebase/main.py,sha256=G8ADdTgjZTxEjvP11IBbjHiNWCvuTkFiLkDKq_8UpI0,26947
+jarvis/jarvis_codebase/main.py,sha256=aGhfH9JTW-CX97OJErzxztJGhL_CH6lXeabfwZxHHxQ,26807
 jarvis/jarvis_coder/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 jarvis/jarvis_coder/main.py,sha256=4YhQ20MD8ntdD5hjmjJbBHYEDPDZ2aPIgAviJCvtOXE,24458
 jarvis/models/__init__.py,sha256=mrOt67nselz_H1gX9wdAO4y2DY5WPXzABqJbr5Des8k,63
@@ -13,21 +13,23 @@ jarvis/models/kimi.py,sha256=1iTB0Z_WOmCML3Ufsge6jmeKOYvccr7I5lS3JUXymU4,17611
 jarvis/models/openai.py,sha256=ayaBWAN5VexMcKVrjEPDNB-Q9wx0sCV9Z4BCrvwYJ9w,4315
 jarvis/models/oyi.py,sha256=X2c5SWDIuQDCCFBcEKbzIWEz3I34eOAi0d1XAFgxlpw,15001
 jarvis/models/registry.py,sha256=YpooKSpk5pSWfb5cBDz5wRfPK-abb9uuUZr4WBejqwI,7762
+jarvis/rag/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+jarvis/rag/main.py,sha256=sgtOzmqi9cO9DcCenSXSad5Xr7MI0U7Ji2GtQL-YaO4,18161
 jarvis/tools/__init__.py,sha256=7Rqyj5hBAv5cWDVr5T9ZTZASO7ssBHeQNm2_4ZARdkA,72
 jarvis/tools/base.py,sha256=EGRGbdfbLXDLwtyoWdvp9rlxNX7bzc20t0Vc2VkwIEY,652
 jarvis/tools/codebase_qa.py,sha256=LsowsgL7HBmdBwa7zXcYi_OkwOok4qbnzYWYsuZxHtU,2413
-jarvis/tools/coder.py,sha256=ZJfPInKms4Hj3-eQlBwamVsvZ-2nlZ-4jsqJ-tJc6mg,2040
+jarvis/tools/coder.py,sha256=kmotT2Klsug44S51QoSW9DzkxLzcF-XonyYAEoWZV6c,2295
 jarvis/tools/file_ops.py,sha256=h8g0eT9UvlJf4kt0DLXvdSsjcPj7x19lxWdDApeDfpg,3842
 jarvis/tools/generator.py,sha256=vVP3eN5cCDpRXf_fn0skETkPXAW1XZFWx9pt2_ahK48,5999
 jarvis/tools/methodology.py,sha256=UG6s5VYRcd9wrKX4cg6f7zJhet5AIcthFGMOAdevBiw,5175
 jarvis/tools/registry.py,sha256=MeTYNdZNRdhlgABviVxzbDPSgLpwDp2Nx2dGzedRu8U,7212
-jarvis/tools/search.py,sha256=1EqOVvLhg2Csh-i03-XeCrusbyfmH69FZ8khwZt8Tow,6131
+jarvis/tools/search.py,sha256=3FX5oAkPiKD4Bvu8gBpKLQveOgl_K5jkvxRvkRPoH3M,8787
 jarvis/tools/shell.py,sha256=UPKshPyOaUwTngresUw-ot1jHjQIb4wCY5nkJqa38lU,2520
 jarvis/tools/sub_agent.py,sha256=rEtAmSVY2ZjFOZEKr5m5wpACOQIiM9Zr_3dT92FhXYU,2621
 jarvis/tools/webpage.py,sha256=d3w3Jcjcu1ESciezTkz3n3Zf-rp_l91PrVoDEZnckOo,2391
-jarvis_ai_assistant-0.1.76.dist-info/LICENSE,sha256=AGgVgQmTqFvaztRtCAXsAMryUymB18gZif7_l2e1XOg,1063
-jarvis_ai_assistant-0.1.76.dist-info/METADATA,sha256=X1sBPb59GZyR4UjtnFntEQY4SzIe0q_KOgmYt62hMcU,12399
-jarvis_ai_assistant-0.1.76.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-jarvis_ai_assistant-0.1.76.dist-info/entry_points.txt,sha256=QNUeqmUJd7nHufel2FO7cRttS1uKFfnbIyObv8eVyOY,140
-jarvis_ai_assistant-0.1.76.dist-info/top_level.txt,sha256=1BOxyWfzOP_ZXj8rVTDnNCJ92bBGB0rwq8N1PCpoMIs,7
-jarvis_ai_assistant-0.1.76.dist-info/RECORD,,
+jarvis_ai_assistant-0.1.77.dist-info/LICENSE,sha256=AGgVgQmTqFvaztRtCAXsAMryUymB18gZif7_l2e1XOg,1063
+jarvis_ai_assistant-0.1.77.dist-info/METADATA,sha256=Fo61TgkIMKSfekG4kV5RR_kfMDAgEpl5D2_2jl8geqU,12736
+jarvis_ai_assistant-0.1.77.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+jarvis_ai_assistant-0.1.77.dist-info/entry_points.txt,sha256=iEvZ0rn-muMxWCNH5QEvw_mTTy_EVhyQyxDxWXqQQVo,174
+jarvis_ai_assistant-0.1.77.dist-info/top_level.txt,sha256=1BOxyWfzOP_ZXj8rVTDnNCJ92bBGB0rwq8N1PCpoMIs,7
+jarvis_ai_assistant-0.1.77.dist-info/RECORD,,

{jarvis_ai_assistant-0.1.76.dist-info → jarvis_ai_assistant-0.1.77.dist-info}/entry_points.txt

@@ -2,3 +2,4 @@
 jarvis = jarvis.main:main
 jarvis-codebase = jarvis.jarvis_codebase.main:main
 jarvis-coder = jarvis.jarvis_coder.main:main
+jarvis-rag = jarvis.rag.main:main