jarvis-ai-assistant 0.1.86__tar.gz → 0.1.87__tar.gz

{jarvis_ai_assistant-0.1.86/src/jarvis_ai_assistant.egg-info → jarvis_ai_assistant-0.1.87}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: jarvis-ai-assistant
-Version: 0.1.86
+Version: 0.1.87
 Summary: Jarvis: An AI assistant that uses tools to interact with the system
 Home-page: https://github.com/skyfireitdiy/Jarvis
 Author: skyfire
@@ -134,7 +134,6 @@ Jarvis supports configuration through environment variables that can be set in t
 | 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 |
 | JARVIS_CHEAP_MODEL | Model name for cheap operations | Same as JARVIS_MODEL | No |
-| JARVIS_EMBEDDING_MODEL | Embedding model for code analysis | BAAI/bge-large-zh-v1.5 | No |
 | OPENAI_API_KEY | API key for OpenAI platform | - | Required for OpenAI |
 | OPENAI_API_BASE | Base URL for OpenAI API | https://api.deepseek.com | No |
 | OPENAI_MODEL_NAME | Model name for OpenAI | deepseek-chat | No |

{jarvis_ai_assistant-0.1.86 → jarvis_ai_assistant-0.1.87}/README.md

@@ -70,7 +70,6 @@ Jarvis supports configuration through environment variables that can be set in t
 | 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 |
 | JARVIS_CHEAP_MODEL | Model name for cheap operations | Same as JARVIS_MODEL | No |
-| JARVIS_EMBEDDING_MODEL | Embedding model for code analysis | BAAI/bge-large-zh-v1.5 | No |
 | OPENAI_API_KEY | API key for OpenAI platform | - | Required for OpenAI |
 | OPENAI_API_BASE | Base URL for OpenAI API | https://api.deepseek.com | No |
 | OPENAI_MODEL_NAME | Model name for OpenAI | deepseek-chat | No |

{jarvis_ai_assistant-0.1.86 → jarvis_ai_assistant-0.1.87}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "jarvis-ai-assistant"
-version = "0.1.86"
+version = "0.1.87"
 description = "Jarvis: An AI assistant that uses tools to interact with the system"
 readme = "README.md"
 authors = [{ name = "Your Name", email = "your.email@example.com" }]

{jarvis_ai_assistant-0.1.86 → jarvis_ai_assistant-0.1.87}/setup.py

@@ -2,7 +2,7 @@ from setuptools import setup, find_packages
 
 setup(
     name="jarvis-ai-assistant",
-    version="0.1.86",
+    version="0.1.87",
     author="skyfire",
     author_email="skyfireitdiy@hotmail.com",
     description="An AI assistant that uses various tools to interact with the system",

{jarvis_ai_assistant-0.1.86 → jarvis_ai_assistant-0.1.87}/src/jarvis/__init__.py

@@ -1,3 +1,3 @@
 """Jarvis AI Assistant"""
 
-__version__ = "0.1.86" 
+__version__ = "0.1.87" 

{jarvis_ai_assistant-0.1.86 → jarvis_ai_assistant-0.1.87}/src/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, get_max_context_length, get_thread_count, load_embedding_model
+from jarvis.utils import OutputType, PrettyOutput, find_git_root, get_max_context_length, get_thread_count, load_embedding_model, load_rerank_model
 from jarvis.utils import load_env_from_file
 import argparse
 from sentence_transformers import SentenceTransformer
@@ -98,7 +98,7 @@ class CodeBase:
 5. 关键业务逻辑和处理流程
 6. 特殊功能点和亮点特性
 
-请用简洁专业的语言描述，突出代码的技术特征和功能特点，以便后续进行相似代码检索。
+请用简洁专业的语言描述，突出代码的技术特征和功能特点，以便后续进行关联代码检索。
 
 文件路径：{file_path}
 代码内容：
@@ -196,10 +196,15 @@ class CodeBase:
             if cached_vector is not None:
                 return cached_vector
                 
-            # 组合文件信息
+            # 读取文件内容并组合信息
+            with open(file_path, "r", encoding="utf-8") as f:
+                content = f.read()[:self.max_context_length]  # 限制文件内容长度
+            
+            # 组合文件信息，包含文件内容
             combined_text = f"""
 文件路径: {file_path}
 文件描述: {description}
+文件内容: {content}
 """
             vector = self.get_embedding(combined_text)
             
@@ -398,116 +403,126 @@ class CodeBase:
             PrettyOutput.print("没有检测到文件变更，无需重建索引", output_type=OutputType.INFO)
 
     def rerank_results(self, query: str, initial_results: List[Tuple[str, float, str]]) -> List[Tuple[str, float, str]]:
-        """使用大模型对搜索结果重新排序"""
+        """使用 BAAI/bge-reranker-v2-m3 对搜索结果重新排序"""
         if not initial_results:
             return []
-
-        model = PlatformRegistry.get_global_platform_registry().get_codegen_platform()
-        # model.set_suppress_output(True)
-
-        # 构建重排序的prompt
-        prompt = f"""请根据用户的查询，对以下代码文件进行相关性排序。对每个文件给出0-100的相关性分数，分数越高表示越相关。
-只需要输出每个文件的分数，格式为：
-<RERANK>
-文件路径: 分数
-文件路径: 分数
-</RERANK>
-
-用户查询: {query}
-
-待评估文件:
-"""
-        for path, _, desc in initial_results:
-            prompt += f"""
-文件: {path}
-描述: {desc}
----
-"""
-        
-        response = model.chat(prompt)
-        
-        # 提取<RERANK>和</RERANK>之间的内容
-        start_tag = "<RERANK>"
-        end_tag = "</RERANK>"
-        if start_tag in response and end_tag in response:
-            response = response[response.find(start_tag) + len(start_tag):response.find(end_tag)]
-        
-        # 解析响应，提取文件路径和分数
-        scored_results = []
-        for line in response.split('\n'):
-            if ':' not in line:
-                continue
-            try:
-                file_path, score_str = line.split(':', 1)
-                file_path = file_path.strip()
-                score = float(score_str.strip()) / 100.0  # 转换为0-1范围
-                # 只保留相关度大于等于0.7的结果
-                if score >= 0.7:
-                    # 找到对应的原始描述
-                    desc = next((desc for p, _, desc in initial_results if p == file_path), "")
-                    scored_results.append((file_path, score, desc))
-            except:
-                continue
-        
-        # 按分数降序排序
-        return sorted(scored_results, key=lambda x: x[1], reverse=True)
-        
+            
+        try:
+            import torch
+            
+            # 加载模型和分词器
+            model, tokenizer = load_rerank_model()
+            
+            # 准备数据 - 加入文件内容进行更准确的重排序
+            pairs = []
+            for path, _, desc in initial_results:
+                try:
+                    with open(path, "r", encoding="utf-8") as f:
+                        content = f.read()[:512]  # 限制内容长度
+                    # 组合文件路径、描述和内容
+                    doc_content = f"文件: {path}\n描述: {desc}\n内容: {content}"
+                    pairs.append([query, doc_content])
+                except Exception as e:
+                    PrettyOutput.print(f"读取文件失败 {path}: {str(e)}", 
+                                    output_type=OutputType.ERROR)
+                    doc_content = f"文件: {path}\n描述: {desc}"
+                    pairs.append([query, doc_content])
+            
+            # 使用更大的batch size提高处理速度
+            batch_size = 16  # 根据GPU显存调整
+            
+            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)
+                    # 修改这里：直接使用 outputs.logits 作为分数
+                    batch_scores = outputs.logits.squeeze(-1).cpu().numpy()
+            
+            # 归一化分数到 0-1 范围
+            if batch_scores:
+                min_score = min(batch_scores)
+                max_score = max(batch_scores)
+                if max_score > min_score:
+                    batch_scores = [(s - min_score) / (max_score - min_score) for s in batch_scores]
+            
+            # 将分数与原始结果组合并排序
+            scored_results = []
+            for (path, _, desc), score in zip(initial_results, batch_scores):
+                if score >= 0.5:  # 只保留相关度大于 0.5 的结果
+                    scored_results.append((path, float(score), desc))
+                    
+            # 按分数降序排序
+            scored_results.sort(key=lambda x: x[1], reverse=True)
+            
+            return scored_results
+            
+        except Exception as e:
+            PrettyOutput.print(f"重排序失败，使用原始排序: {str(e)}", output_type=OutputType.WARNING)
+            return initial_results
 
     def search_similar(self, query: str, top_k: int = 30) -> List[Tuple[str, float, str]]:
-        """搜索相似文件"""
-        model = PlatformRegistry.get_global_platform_registry().get_normal_platform()
-        model.set_suppress_output(True)
-        
-
-        prompt = f"""请根据以下查询，生成意思完全相同的另一个表述。这个表述将用于代码搜索，所以要保持专业性和准确性。
+        """搜索关联文件"""
+        try:
+            # 生成多个查询变体以提高召回率
+            model = PlatformRegistry.get_global_platform_registry().get_normal_platform()
+            model.set_suppress_output(True)
+            
+            prompt = f"""请根据以下查询，生成3个不同的表述，每个表述都要完整表达原始查询的意思。这些表述将用于代码搜索，要保持专业性和准确性。
 原始查询: {query}
 
-请直接输出新表述，不要有编号或其他标记。
+请直接输出3个表述，用换行分隔，不要有编号或其他标记。
 """
-        
-        query = model.chat(prompt)
+            query_variants = model.chat(prompt).strip().split('\n')
+            query_variants.append(query)  # 添加原始查询
             
-
-
-        PrettyOutput.print(f"查询:  {query}", output_type=OutputType.INFO)
-        
-        # 为每个查询获取相似文件
-        q_vector = self.get_embedding(query)
-        q_vector = q_vector.reshape(1, -1)
-            
-        distances, indices = self.index.search(q_vector, top_k)
-
-        PrettyOutput.print(f"查询 {query} 的结果: ", output_type=OutputType.INFO)
-
-        initial_results = []
-        
-        for i, distance in zip(indices[0], distances[0]):
-            if i == -1:  # faiss返回-1表示无效结果
-                continue
+            # 对每个查询变体进行搜索
+            all_results = {}
+            for q in query_variants:
+                q_vector = self.get_embedding(q)
+                q_vector = q_vector.reshape(1, -1)
                 
-            similarity = 1.0 / (1.0 + float(distance))
-            # 只保留相似度大于等于0.5的结果
-            if similarity >= 0.5:
-                PrettyOutput.print(f"  {self.file_paths[i]} : 距离 {distance:.3f}, 相似度 {similarity:.3f}", 
-                                 output_type=OutputType.INFO)
+                distances, indices = self.index.search(q_vector, top_k)
                 
-                file_path = self.file_paths[i]
-                data = self.vector_cache[file_path]
-                initial_results.append((file_path, similarity, data["description"]))
-
-        if not initial_results:
-            PrettyOutput.print("没有找到相似度大于0.5的文件", output_type=OutputType.WARNING)
+                for i, distance in zip(indices[0], distances[0]):
+                    if i == -1:
+                        continue
+                        
+                    similarity = 1.0 / (1.0 + float(distance))
+                    if similarity >= 0.5:
+                        file_path = self.file_paths[i]
+                        # 使用最高的相似度分数
+                        if file_path not in all_results or similarity > all_results[file_path][1]:
+                            data = self.vector_cache[file_path]
+                            all_results[file_path] = (file_path, similarity, data["description"])
+            
+            # 转换为列表并排序
+            results = list(all_results.values())
+            results.sort(key=lambda x: x[1], reverse=True)
+            
+            return results[:top_k]
+            
+        except Exception as e:
+            PrettyOutput.print(f"搜索失败: {str(e)}", output_type=OutputType.ERROR)
             return []
 
-        # 使用大模型重新排序
-        PrettyOutput.print("使用大模型重新排序...", output_type=OutputType.INFO)
-        reranked_results = self.rerank_results(query, initial_results)
-        
-        return reranked_results
-
     def ask_codebase(self, query: str, top_k: int=20) -> str:
         """查询代码库"""
         results = self.search_similar(query, top_k)
+        if not results:
+            PrettyOutput.print("没有找到关联的文件", output_type=OutputType.WARNING)
+            return ""
+        
         PrettyOutput.print(f"找到的关联文件: ", output_type=OutputType.SUCCESS)
         for path, score, _ in results:
             PrettyOutput.print(f"文件: {path} 关联度: {score:.3f}", 

{jarvis_ai_assistant-0.1.86 → jarvis_ai_assistant-0.1.87}/src/jarvis/jarvis_coder/main.py

@@ -127,27 +127,55 @@ class JarvisCoder:
 修改格式说明：
 1. 每个修改块格式如下：
 <PATCH>
->>>>>> path/to/file
-要替换的内容
+> path/to/file
+def old_function():
+    print("old code")
+    return False
 =======
-新的内容
->>>>>>
+def old_function():
+    print("new code")
+    return True
 </PATCH>
 
 2. 如果是新文件或者替换整个文件内容，格式如下：
 <PATCH>
->>>>>> path/to/new/file
+> src/new_module.py
 =======
-新文件的完整内容
->>>>>>
+from typing import List
+
+def new_function():
+    return "This is a new file"
 </PATCH>
 
 3. 如果要删除文件中的某一段，格式如下：
 <PATCH>
->>>>>> path/to/file
-要删除的内容
+> path/to/file
+    # 这是要删除的注释
+    deprecated_code = True
+    if deprecated_code:
+        print("old feature")
+=======
+</PATCH>
+
+4. 如果要修改导入语句，格式如下：
+<PATCH>
+> src/main.py
+from old_module import old_class
+=======
+from new_module import new_class
+</PATCH>
+
+5. 如果要修改类定义，格式如下：
+<PATCH>
+> src/models.py
+class OldModel:
+    def __init__(self):
+        self.value = 0
 =======
->>>>>>
+class OldModel:
+    def __init__(self):
+        self.value = 1
+        self.name = "new"
 </PATCH>
 
 文件列表如下：
@@ -170,7 +198,7 @@ class JarvisCoder:
 3、要替换的内容，一定要与文件内容完全一致，不要有任何多余或者缺失的内容
 4、每个patch不超过20行，超出20行，请生成多个patch
 """
-        
+
         success, response = self._call_model_with_retry(self.main_model, prompt)
         if not success:
             return []
@@ -204,7 +232,7 @@ class JarvisCoder:
                     continue
                     
                 # 获取文件路径
-                file_path_match = re.search(r'>>>>>> (.*)', lines[0])
+                file_path_match = re.search(r'> (.*)', lines[0])
                 if not file_path_match:
                     error_info.append(f"无法解析文件路径: {lines[0]}")
                     return False, "\n".join(error_info)
@@ -220,7 +248,7 @@ class JarvisCoder:
                     return False, "\n".join(error_info)
                 
                 old_content = parts[0]
-                new_content = parts[1].split(">>>>>>")[0]
+                new_content = parts[1].split("</PATCH>")[0]
                 
                 # 处理新文件
                 if not old_content:
@@ -348,19 +376,59 @@ class JarvisCoder:
                     }
                 else:
                     self._revert_changes()
-                    return {
-                        "success": False,
-                        "stdout": "",
-                        "stderr": "修改被用户取消，文件未发生任何变化",
-                        "error": UserWarning("用户取消修改")
-                    }
+                    
+                    # 让用户输入调整意见
+                    user_feedback = get_multiline_input("""
+请提供修改建议，帮助生成更好的补丁：
+1. 修改的位置是否正确？
+2. 修改的内容是否合适？
+3. 是否有遗漏的修改？
+4. 其他调整建议？
+
+请输入调整意见(直接回车跳过):""")
+                    
+                    if not user_feedback:
+                        return {
+                            "success": False,
+                            "stdout": "",
+                            "stderr": "修改被用户取消，文件未发生任何变化",
+                            "error": UserWarning("用户取消修改")
+                        }
+                    
+                    retry_prompt = f"""补丁被用户拒绝，请根据用户意见重新生成补丁：
+
+用户意见：
+{user_feedback}
+
+请重新生成补丁，确保：
+1. 按照用户意见调整修改内容
+2. 准确定位要修改的代码位置
+3. 正确处理代码缩进
+4. 考虑代码上下文
+"""
+                    patches = self._remake_patch(retry_prompt)
+                    continue
             else:
-                PrettyOutput.print(f"补丁应用失败，请求重新生成: {error_info}", OutputType.WARNING)
-                retry_prompt = f"""补丁应用失败，请根据以下错误信息重新生成补丁：
+                PrettyOutput.print(f"补丁应用失败: {error_info}", OutputType.WARNING)
+                
+                # 让用户输入补充信息
+                user_info = get_multiline_input("""
+补丁应用失败。请提供更多信息来帮助修复问题：
+1. 是否需要调整代码位置？
+2. 是否有特殊的格式要求？
+3. 是否需要考虑其他文件的依赖？
+4. 其他补充说明？
+
+请输入补充信息(直接回车跳过):""")
+                
+                retry_prompt = f"""补丁应用失败，请根据以下信息重新生成补丁：
 
 错误信息：
 {error_info}
 
+用户补充信息：
+{user_info if user_info else "用户未提供补充信息"}
+
 请确保：
 1. 准确定位要修改的代码位置
 2. 正确处理代码缩进

{jarvis_ai_assistant-0.1.86 → jarvis_ai_assistant-0.1.87}/src/jarvis/jarvis_rag/main.py

@@ -5,7 +5,7 @@ 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, get_max_context_length, load_embedding_model
+from jarvis.utils import OutputType, PrettyOutput, find_git_root, get_max_context_length, load_embedding_model, load_rerank_model
 from jarvis.utils import load_env_from_file
 import tiktoken
 from dataclasses import dataclass
@@ -14,6 +14,8 @@ 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
 
 @dataclass
 class Document:
@@ -193,15 +195,30 @@ class RAGTool:
                 self.index = None
 
     def _save_cache(self, vectors: np.ndarray):
-        """保存缓存数据"""
+        """优化缓存保存"""
         try:
+            # 添加版本号和时间戳
             cache_data = {
+                "version": "1.0",
+                "timestamp": datetime.now().isoformat(),
                 "documents": self.documents,
-                "vectors": vectors
+                "vectors": vectors,
+                "metadata": {
+                    "vector_dim": self.vector_dim,
+                    "total_docs": len(self.documents),
+                    "model_name": self.embedding_model.__class__.__name__
+                }
             }
+            
+            # 使用压缩存储
             with open(self.cache_path, 'wb') as f:
-                pickle.dump(cache_data, f)
-            PrettyOutput.print(f"保存了 {len(self.documents)} 个文档片段", 
+                pickle.dump(cache_data, f, protocol=pickle.HIGHEST_PROTOCOL)
+            
+            # 创建备份
+            backup_path = f"{self.cache_path}.backup"
+            shutil.copy2(self.cache_path, backup_path)
+            
+            PrettyOutput.print(f"缓存已保存: {len(self.documents)} 个文档片段", 
                             output_type=OutputType.INFO)
         except Exception as e:
             PrettyOutput.print(f"保存缓存失败: {str(e)}", 
@@ -209,100 +226,74 @@ class RAGTool:
 
     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)
+        # 添加IVF索引以提高大规模检索性能
+        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)
         
-        # 添加向量到索引
         if vectors.shape[0] > 0:
-            self.index.add_with_ids(vectors, np.arange(vectors.shape[0]))
+            # 训练IVF索引
+            self.index.train(vectors)
+            self.index.add(vectors)
+            # 设置搜索时探测的聚类数
+            self.index.nprobe = min(nlist, 10)
         else:
             self.index = None
 
     def _split_text(self, text: str) -> List[str]:
-        """将文本分割成段落
+        """使用更智能的分块策略"""
+        # 添加重叠分块以保持上下文连贯性
+        overlap_size = min(200, self.max_paragraph_length // 4)
         
-        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)
+        current_chunk = []
+        current_length = 0
         
-        # 处理最后一个段落
-        if current_paragraph:
-            paragraph_text = ' '.join(current_paragraph)
-            if len(paragraph_text) >= self.min_paragraph_length:
-                paragraphs.append(paragraph_text)
+        # 首先按句子分割
+        sentences = []
+        current_sentence = []
+        sentence_ends = {'。', '！', '？', '…', '.', '!', '?'}
         
-        # 处理过长的段落
-        final_paragraphs = []
-        for paragraph in paragraphs:
-            if len(paragraph) <= self.max_paragraph_length:
-                final_paragraphs.append(paragraph)
-            else:
-                # 按句子分割过长的段落
-                sentences = []
+        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 = []
-                
-                # 中文句子结束标记
-                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_sentence:
+            sentence = ''.join(current_sentence)
+            if sentence.strip():
+                sentences.append(sentence)
+        
+        # 基于句子构建重叠块
+        for sentence in sentences:
+            if current_length + len(sentence) > self.max_paragraph_length:
                 if current_chunk:
-                    final_paragraphs.append(''.join(current_chunk))
+                    chunk_text = ' '.join(current_chunk)
+                    if len(chunk_text) >= self.min_paragraph_length:
+                        paragraphs.append(chunk_text)
+                        
+                    # 保留部分内容作为重叠
+                    overlap_text = ' '.join(current_chunk[-2:])  # 保留最后两句
+                    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)
         
-        # 过滤掉太短的段落
-        final_paragraphs = [p for p in final_paragraphs if len(p) >= self.min_paragraph_length]
+        # 处理最后一个chunk
+        if current_chunk:
+            chunk_text = ' '.join(current_chunk)
+            if len(chunk_text) >= self.min_paragraph_length:
+                paragraphs.append(chunk_text)
         
-        return final_paragraphs
+        return paragraphs
 
     def _get_embedding(self, text: str) -> np.ndarray:
         """获取文本的向量表示"""
@@ -410,82 +401,131 @@ class RAGTool:
                         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:
             PrettyOutput.print("索引未构建，正在构建...", output_type=OutputType.INFO)
             self.build_index(self.root_dir)
+        
+        # 实现MMR (Maximal Marginal Relevance) 来增加结果多样性
+        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)):
+                    # 计算与查询的相似度
+                    query_sim = float(np.dot(query_vec, doc_vec))
+                    
+                    # 计算与已选文档的最大相似度
+                    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
+        
+        # 获取查询向量
         query_vector = self._get_embedding(query)
         query_vector = query_vector.reshape(1, -1)
         
-        # 搜索最相似的向量
-        distances, indices = self.index.search(query_vector, top_k)
+        # 初始搜索更多结果用于MMR
+        initial_k = min(top_k * 2, len(self.documents))
+        distances, indices = self.index.search(query_vector, initial_k)
+        
+        # 获取有效结果
+        valid_indices = indices[0][indices[0] != -1]
+        valid_vectors = np.vstack([self._get_embedding(self.documents[idx].content) for idx in valid_indices])
         
-        # 返回结果
+        # 应用MMR
+        final_indices = mmr(query_vector[0], valid_vectors, valid_indices, n_docs=top_k)
+        
+        # 构建结果
         results = []
-        current_length = 0
+        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))
         
-        for idx, distance in zip(indices[0], distances[0]):
-            if idx == -1:  # FAISS返回-1表示无效结果
-                continue
+        return results
+
+    def _rerank_results(self, query: str, initial_results: List[Tuple[Document, float]]) -> List[Tuple[Document, float]]:
+        """使用 rerank 模型重新排序搜索结果"""
+        try:
+            import torch
+            model, tokenizer = load_rerank_model()
+            
+            # 准备数据
+            pairs = []
+            for doc, _ in initial_results:
+                # 组合文档信息
+                doc_content = f"""
+文件: {doc.metadata['file_path']}
+内容: {doc.content}
+"""
+                pairs.append([query, doc_content])
                 
-            doc = self.documents[idx]
-            similarity = 1.0 / (1.0 + float(distance))
+            # 对每个文档对进行打分
+            scores = []
+            batch_size = 8
             
-            # 获取同一文件中的所有文档片段
-            file_docs = [d for d in self.documents if d.metadata['file_path'] == doc.metadata['file_path']]
-            file_docs.sort(key=lambda x: x.metadata['chunk_index'])
+            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())
             
-            # 找到当前片段的索引
-            current_idx = file_docs.index(doc)
+            # 归一化分数到 0-1 范围
+            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]
             
-            # 尝试不同的上下文窗口大小，从最大到最小
-            added = False
-            for window_size in range(self.context_window, -1, -1):
-                start_idx = max(0, current_idx - window_size)
-                end_idx = min(len(file_docs), current_idx + window_size + 1)
-                
-                # 合并内容，包含上下文
-                content_parts = []
-                content_parts.extend(file_docs[i].content for i in range(start_idx, current_idx))
-                content_parts.append(doc.content)
-                content_parts.extend(file_docs[i].content for i in range(current_idx + 1, end_idx))
-                
-                merged_content = "\n".join(content_parts)
-                
-                # 创建文档对象
-                context_doc = Document(
-                    content=merged_content,
-                    metadata={
-                        **doc.metadata,
-                        "similarity": similarity
-                    }
-                )
-                
-                # 计算添加这个结果后的总长度
-                total_content_length = len(merged_content)
-                
-                # 检查是否在长度限制内
-                if current_length + total_content_length <= self.max_context_length:
-                    results.append((context_doc, similarity))
-                    current_length += total_content_length
-                    added = True
-                    break
+            # 将分数与文档组合并排序
+            scored_results = []
+            for (doc, _), score in zip(initial_results, scores):
+                if score >= 0.5:  # 只保留关联度大于 0.5 的结果
+                    scored_results.append((doc, float(score)))
+                    
+            # 按分数降序排序
+            scored_results.sort(key=lambda x: x[1], reverse=True)
             
-            # 如果即使没有上下文也无法添加，就停止添加更多结果
-            if not added:
-                break
-                
-        return results
+            return scored_results
+            
+        except Exception as e:
+            PrettyOutput.print(f"重排序失败，使用原始排序: {str(e)}", output_type=OutputType.WARNING)
+            return initial_results
 
     def is_index_built(self):
         """检查索引是否已构建"""
@@ -567,9 +607,12 @@ def main():
     args = parser.parse_args()
 
     try:
-        current_dir = find_git_root()
+        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"正在处理目录: {args.dir}", output_type=OutputType.INFO)
             rag.build_index(args.dir)

{jarvis_ai_assistant-0.1.86 → jarvis_ai_assistant-0.1.87}/src/jarvis/models/ai8.py

@@ -111,7 +111,7 @@ class AI8Model(BasePlatform):
             session_data = {
                 **self.conversation,
                 "model": self.model_name,
-                "contextCount": 1024,
+                "contextCount": 65536,
                 "prompt": self.system_message,
                 "plugins": [],
                 "localPlugins": None,

{jarvis_ai_assistant-0.1.86 → jarvis_ai_assistant-0.1.87}/src/jarvis/models/oyi.py

@@ -2,7 +2,7 @@ import mimetypes
 import os
 from typing import Dict, List
 from jarvis.models.base import BasePlatform
-from jarvis.utils import PrettyOutput, OutputType
+from jarvis.utils import PrettyOutput, OutputType, get_max_context_length
 import requests
 import json
 
@@ -72,10 +72,10 @@ class OyiModel(BasePlatform):
                     "is_webSearch": True,
                     "message": [],
                     "systemMessage": None,
-                    "requestMsgCount": 1000,
+                    "requestMsgCount": 65536,
                     "temperature": 0.8,
                     "speechVoice": "Alloy",
-                    "max_tokens": 8192,
+                    "max_tokens": get_max_context_length(),
                     "chatPluginIds": []
                 })
             }

{jarvis_ai_assistant-0.1.86 → jarvis_ai_assistant-0.1.87}/src/jarvis/tools/registry.py

@@ -19,7 +19,8 @@ class ToolRegistry:
         # 加载内置工具和外部工具
         self._load_builtin_tools()
         self._load_external_tools()
-        self.max_context_length = get_max_context_length() * 0.8
+        # 确保 max_context_length 是整数
+        self.max_context_length = int(get_max_context_length() * 0.8)
 
     @staticmethod
     def get_global_tool_registry():
@@ -181,10 +182,11 @@ class ToolRegistry:
                         PrettyOutput.print("输出较长，正在总结...", OutputType.PROGRESS)
                         model = PlatformRegistry.get_global_platform_registry().get_normal_platform()
                         
-                        # 如果输出超过30k，只取最后30k字符
-                        if len(output) > self.max_context_length:
-                            output_to_summarize = output[-self.max_context_length:]
-                            truncation_notice = f"\n(注意: 由于输出过长，仅总结最后{self.max_context_length}字符)"
+                        # 如果输出超过最大上下文长度，只取最后部分
+                        max_len = self.max_context_length
+                        if len(output) > max_len:
+                            output_to_summarize = output[-max_len:]
+                            truncation_notice = f"\n(注意: 由于输出过长，仅总结最后{max_len}字符)"
                         else:
                             output_to_summarize = output
                             truncation_notice = ""

{jarvis_ai_assistant-0.1.86 → jarvis_ai_assistant-0.1.87}/src/jarvis/utils.py

@@ -10,6 +10,8 @@ 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
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
+import torch
 
 # 初始化colorama
 colorama.init()
@@ -211,7 +213,7 @@ def find_git_root(dir="."):
 
 def load_embedding_model():
     os.environ["TOKENIZERS_PARALLELISM"] = "false"
-    model_name = os.environ.get("JARVIS_EMBEDDING_MODEL", "BAAI/bge-large-zh-v1.5")
+    model_name = "BAAI/bge-large-zh-v1.5"
     PrettyOutput.print(f"正在加载嵌入模型: {model_name}...", OutputType.INFO)
     try:
         # 首先尝试离线加载
@@ -234,6 +236,44 @@ def load_embedding_model():
     
     return embedding_model
 
+def load_rerank_model():
+    """加载重排序模型"""
+    model_name = "BAAI/bge-reranker-v2-m3"
+    PrettyOutput.print(f"正在加载重排序模型: {model_name}...", OutputType.INFO)
+    
+    try:
+        # 首先尝试离线加载
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_name,
+            local_files_only=True,
+            cache_dir=os.path.expanduser("~/.cache/huggingface/hub")
+        )
+        model = AutoModelForSequenceClassification.from_pretrained(
+            model_name,
+            local_files_only=True,
+            cache_dir=os.path.expanduser("~/.cache/huggingface/hub")
+        )
+        PrettyOutput.print("使用本地缓存加载模型成功", OutputType.SUCCESS)
+    except Exception as local_error:
+        PrettyOutput.print(f"本地加载失败，尝试在线下载: {str(local_error)}", OutputType.WARNING)
+        # 如果离线加载失败，尝试在线下载
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_name,
+            cache_dir=os.path.expanduser("~/.cache/huggingface/hub")
+        )
+        model = AutoModelForSequenceClassification.from_pretrained(
+            model_name,
+            cache_dir=os.path.expanduser("~/.cache/huggingface/hub")
+        )
+        PrettyOutput.print("模型下载并加载成功", OutputType.SUCCESS)
+    
+    # 如果有 GPU 就使用 GPU
+    if torch.cuda.is_available():
+        model = model.cuda()
+    model.eval()
+    
+    return model, tokenizer
+
 def get_max_context_length():
     return int(os.getenv('JARVIS_MAX_CONTEXT_LENGTH', '131072'))  # 默认128k
 

{jarvis_ai_assistant-0.1.86 → jarvis_ai_assistant-0.1.87/src/jarvis_ai_assistant.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: jarvis-ai-assistant
-Version: 0.1.86
+Version: 0.1.87
 Summary: Jarvis: An AI assistant that uses tools to interact with the system
 Home-page: https://github.com/skyfireitdiy/Jarvis
 Author: skyfire
@@ -134,7 +134,6 @@ Jarvis supports configuration through environment variables that can be set in t
 | 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 |
 | JARVIS_CHEAP_MODEL | Model name for cheap operations | Same as JARVIS_MODEL | No |
-| JARVIS_EMBEDDING_MODEL | Embedding model for code analysis | BAAI/bge-large-zh-v1.5 | No |
 | OPENAI_API_KEY | API key for OpenAI platform | - | Required for OpenAI |
 | OPENAI_API_BASE | Base URL for OpenAI API | https://api.deepseek.com | No |
 | OPENAI_MODEL_NAME | Model name for OpenAI | deepseek-chat | No |