jarvis-ai-assistant 0.1.130__py3-none-any.whl → 0.1.132__py3-none-any.whl

jarvis/jarvis_utils/embedding.py

@@ -3,9 +3,16 @@ import numpy as np
 import torch
 from sentence_transformers import SentenceTransformer
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
-from typing import List, Any, Tuple
+from typing import List, Any, Optional, Tuple
+import functools
+
+from yaspin.api import Yaspin
 from jarvis.jarvis_utils.output import PrettyOutput, OutputType
 
+# 全局缓存，避免重复加载模型
+_global_models = {}
+_global_tokenizers = {}
+
 def get_context_token_count(text: str) -> int:
     """使用分词器获取文本的token数量。
     
@@ -26,9 +33,10 @@ def get_context_token_count(text: str) -> int:
         # 回退到基于字符的粗略估计
         return len(text) // 4  # 每个token大约4个字符的粗略估计
 
+@functools.lru_cache(maxsize=1)
 def load_embedding_model() -> SentenceTransformer:
     """
-    加载句子嵌入模型。
+    加载句子嵌入模型，使用缓存避免重复加载。
     
     返回：
         SentenceTransformer: 加载的嵌入模型
@@ -36,6 +44,10 @@ def load_embedding_model() -> SentenceTransformer:
     model_name = "BAAI/bge-m3"
     cache_dir = os.path.expanduser("~/.cache/huggingface/hub")
     
+    # 检查全局缓存中是否已有模型
+    if model_name in _global_models:
+        return _global_models[model_name]
+    
     try:
         embedding_model = SentenceTransformer(
             model_name,
@@ -49,6 +61,13 @@ def load_embedding_model() -> SentenceTransformer:
             local_files_only=False
         )
     
+    # 如果可用，将模型移到GPU上
+    if torch.cuda.is_available():
+        embedding_model.to(torch.device("cuda"))
+    
+    # 保存到全局缓存
+    _global_models[model_name] = embedding_model
+    
     return embedding_model
 
 def get_embedding(embedding_model: Any, text: str) -> np.ndarray:
@@ -67,77 +86,249 @@ def get_embedding(embedding_model: Any, text: str) -> np.ndarray:
                                      show_progress_bar=False)
     return np.array(embedding, dtype=np.float32)
 
-def get_embedding_batch(embedding_model: Any, texts: List[str]) -> np.ndarray:
+def get_embedding_batch(embedding_model: Any, prefix: str, texts: List[str], spinner: Optional[Yaspin] = None, batch_size: int = 8) -> np.ndarray:
     """
-    为一批文本生成嵌入向量。
+    为一批文本生成嵌入向量，使用高效的批处理，针对RAG优化。
     
     参数：
         embedding_model: 使用的嵌入模型
+        prefix: 进度条前缀
         texts: 要嵌入的文本列表
+        spinner: 可选的进度指示器
+        batch_size: 批处理大小，更大的值可能更快但需要更多内存
         
     返回：
         np.ndarray: 堆叠的嵌入向量
     """
+    # 简单嵌入缓存，避免重复计算相同文本块
+    embedding_cache = {}
+    cache_hits = 0
+    
     try:
+        # 预处理：将所有文本分块
+        all_chunks = []
+        chunk_indices = []  # 跟踪每个原始文本对应的块索引
+        
+        for i, text in enumerate(texts):
+            if spinner:
+                spinner.text = f"{prefix} 预处理中 ({i+1}/{len(texts)}) ..."
+            
+            # 预处理文本：移除多余空白，规范化
+            text = ' '.join(text.split()) if text else ""
+            
+            # 使用更优化的分块函数
+            chunks = split_text_into_chunks(text, 512)
+            start_idx = len(all_chunks)
+            all_chunks.extend(chunks)
+            end_idx = len(all_chunks)
+            chunk_indices.append((start_idx, end_idx))
+        
+        if not all_chunks:
+            return np.zeros((0, embedding_model.get_sentence_embedding_dimension()), dtype=np.float32)
+        
+        # 批量处理所有块
         all_vectors = []
-        for text in texts:
-            vectors = get_embedding_with_chunks(embedding_model, text)
-            all_vectors.extend(vectors)
-        return np.vstack(all_vectors)
+        for i in range(0, len(all_chunks), batch_size):
+            if spinner:
+                spinner.text = f"{prefix} 批量处理嵌入 ({i+1}/{len(all_chunks)}) ..."
+            
+            batch = all_chunks[i:i+batch_size]
+            batch_to_process = []
+            batch_indices = []
+            
+            # 检查缓存，避免重复计算
+            for j, chunk in enumerate(batch):
+                chunk_hash = hash(chunk)
+                if chunk_hash in embedding_cache:
+                    all_vectors.append(embedding_cache[chunk_hash])
+                    cache_hits += 1
+                else:
+                    batch_to_process.append(chunk)
+                    batch_indices.append(j)
+            
+            if batch_to_process:
+                # 对未缓存的块处理
+                batch_vectors = embedding_model.encode(
+                    batch_to_process, 
+                    normalize_embeddings=True,
+                    show_progress_bar=False,
+                    convert_to_numpy=True,
+                )
+                
+                # 处理结果并更新缓存
+                if len(batch_to_process) == 1:
+                    vec = batch_vectors
+                    chunk_hash = hash(batch_to_process[0])
+                    embedding_cache[chunk_hash] = vec
+                    all_vectors.append(vec)
+                else:
+                    for j, vec in enumerate(batch_vectors):
+                        chunk_hash = hash(batch_to_process[j])
+                        embedding_cache[chunk_hash] = vec
+                        all_vectors.append(vec)
+        
+        # 组织结果到原始文本顺序
+        result_vectors = []
+        for start_idx, end_idx in chunk_indices:
+            text_vectors = []
+            for j in range(start_idx, end_idx):
+                if j < len(all_vectors):
+                    text_vectors.append(all_vectors[j])
+            
+            if text_vectors:
+                # 当一个文本被分成多个块时，采用加权平均
+                if len(text_vectors) > 1:
+                    # 针对RAG优化：对多个块进行加权平均，前面的块权重略高
+                    weights = np.linspace(1.0, 0.8, len(text_vectors))
+                    weights = weights / weights.sum()  # 归一化权重
+                    
+                    # 应用权重并求和
+                    weighted_sum = np.zeros_like(text_vectors[0])
+                    for i, vec in enumerate(text_vectors):
+                        # 确保向量形状一致，处理可能的维度不匹配问题
+                        vec_array = np.asarray(vec).reshape(weighted_sum.shape)
+                        weighted_sum += vec_array * weights[i]
+                    
+                    # 归一化结果向量
+                    norm = np.linalg.norm(weighted_sum)
+                    if norm > 0:
+                        weighted_sum = weighted_sum / norm
+                    
+                    result_vectors.append(weighted_sum)
+                else:
+                    # 单块直接使用
+                    result_vectors.append(text_vectors[0])
+        
+        if spinner and cache_hits > 0:
+            spinner.text = f"{prefix} 缓存命中: {cache_hits}/{len(all_chunks)} 块"
+        
+        return np.vstack(result_vectors)
+    
     except Exception as e:
         PrettyOutput.print(f"批量嵌入失败: {str(e)}", OutputType.ERROR)
         return np.zeros((0, embedding_model.get_sentence_embedding_dimension()), dtype=np.float32)
     
-def split_text_into_chunks(text: str, max_length: int = 512) -> List[str]:
-    """将文本分割成带重叠窗口的块。
+def split_text_into_chunks(text: str, max_length: int = 512, min_length: int = 50) -> List[str]:
+    """将文本分割成带重叠窗口的块，优化RAG检索效果。
     
     参数：
         text: 要分割的输入文本
         max_length: 每个块的最大长度
+        min_length: 每个块的最小长度（除了最后一块可能较短）
         
     返回：
-        List[str]: 文本块列表
+        List[str]: 文本块列表，每个块的长度尽可能接近但不超过max_length
     """
+    if not text:
+        return []
+    
+    # 如果文本长度小于最大长度，直接返回整个文本
+    if len(text) <= max_length:
+        return [text]
+    
+    # 预处理：规范化文本，移除多余空白字符
+    text = ' '.join(text.split())
+    
+    # 中英文标点符号集合，优化RAG召回的句子边界
+    primary_punctuation = {'.', '!', '?', '\n', '。', '！', '？'}  # 主要句末标点
+    secondary_punctuation = {'；', '：', '…', ';', ':'}  # 次级分隔符
+    tertiary_punctuation = {',', '，', '、', ')', '）', ']', '】', '}', '》', '"', "'"}  # 最低优先级
+    
     chunks = []
     start = 0
-    while start < len(text):
-        end = start + max_length
-        # 找到最近的句子边界
-        if end < len(text):
-            while end > start and text[end] not in {'.', '!', '?', '\n'}:
-                end -= 1
-            if end == start:  # 未找到标点，强制分割
-                end = start + max_length
-        chunk = text[start:end]
-        chunks.append(chunk)
-        # 重叠20%的窗口
-        start = end - int(max_length * 0.2)
-    return chunks
-
-def get_embedding_with_chunks(embedding_model: Any, text: str) -> List[np.ndarray]:
-    """
-    为文本块生成嵌入向量。
     
-    参数：
-        embedding_model: 使用的嵌入模型
-        text: 要处理的输入文本
+    while start < len(text):
+        # 初始化结束位置为最大可能长度
+        end = min(start + max_length, len(text))
         
-    返回：
-        List[np.ndarray]: 每个块的嵌入向量列表
-    """
-    chunks = split_text_into_chunks(text, 512)
-    if not chunks:
-        return []
+        # 只有当不是最后一块且结束位置等于最大长度时，才尝试寻找句子边界
+        if end < len(text) and end == start + max_length:
+            # 优先查找段落边界，这对RAG特别重要
+            paragraph_boundary = text.rfind('\n\n', start, end)
+            if paragraph_boundary > start and paragraph_boundary < end - min_length:  # 确保不会切得太短
+                end = paragraph_boundary + 2
+            else:
+                # 寻找句子边界，从end-1位置开始
+                found_boundary = False
+                best_boundary = -1
+                
+                # 扩大搜索范围以找到更好的语义边界
+                search_range = min(120, end - start - min_length)  # 扩大搜索范围，但确保新块不小于min_length
+                
+                # 先尝试找主要标点（句号等）
+                for i in range(end-1, max(start, end-search_range), -1):
+                    if text[i] in primary_punctuation:
+                        best_boundary = i
+                        found_boundary = True
+                        break
+                
+                # 如果没找到主要标点，再找次要标点（分号、冒号等）
+                if not found_boundary:
+                    for i in range(end-1, max(start, end-search_range), -1):
+                        if text[i] in secondary_punctuation:
+                            best_boundary = i
+                            found_boundary = True
+                            break
+                
+                # 最后考虑逗号和其他可能的边界
+                if not found_boundary:
+                    for i in range(end-1, max(start, end-search_range), -1):
+                        if text[i] in tertiary_punctuation:
+                            best_boundary = i
+                            found_boundary = True
+                            break
+                
+                # 如果找到了合适的边界且不会导致太短的块，使用它
+                if found_boundary and (best_boundary - start) >= min_length:
+                    end = best_boundary + 1
+        
+        # 添加当前块，并确保删除开头和结尾的空白字符
+        chunk = text[start:end].strip()
+        if chunk and len(chunk) >= min_length:  # 只添加符合最小长度的非空块
+            chunks.append(chunk)
+        elif chunk and not chunks:  # 如果是第一个块且小于最小长度，也添加它
+            chunks.append(chunk)
+        elif chunk:  # 如果块太小，尝试与前一个块合并
+            if chunks:
+                if len(chunks[-1]) + len(chunk) <= max_length * 1.1:  # 允许略微超过最大长度
+                    chunks[-1] = chunks[-1] + " " + chunk
+                else:
+                    # 如果合并会导致太长，添加这个小块（特殊情况）
+                    chunks.append(chunk)
+        
+        # 计算下一块的开始位置，调整重叠窗口大小以提高RAG检索质量
+        next_start = end - int(max_length * 0.2)  # 20%的重叠窗口大小
+        
+        # 确保总是有前进，避免无限循环
+        if next_start <= start:
+            next_start = start + max(1, min_length // 2)
+            
+        start = next_start
+    
+    # 最后检查是否有太短的块，尝试合并相邻的短块
+    if len(chunks) > 1:
+        merged_chunks = []
+        i = 0
+        while i < len(chunks):
+            current = chunks[i]
+            # 如果当前块太短且不是最后一个块，尝试与下一个合并
+            if len(current) < min_length and i < len(chunks) - 1:
+                next_chunk = chunks[i + 1]
+                if len(current) + len(next_chunk) <= max_length * 1.1:
+                    merged_chunks.append(current + " " + next_chunk)
+                    i += 2  # 跳过下一个块
+                    continue
+            merged_chunks.append(current)
+            i += 1
+        chunks = merged_chunks
     
-    vectors = []
-    for chunk in chunks:
-        vector = get_embedding(embedding_model, chunk)
-        vectors.append(vector)
-    return vectors
+    return chunks
+
 
+@functools.lru_cache(maxsize=1)
 def load_tokenizer() -> AutoTokenizer:
     """
-    加载用于文本处理的分词器。
+    加载用于文本处理的分词器，使用缓存避免重复加载。
     
     返回：
         AutoTokenizer: 加载的分词器
@@ -145,6 +336,10 @@ def load_tokenizer() -> AutoTokenizer:
     model_name = "gpt2"
     cache_dir = os.path.expanduser("~/.cache/huggingface/hub")
     
+    # 检查全局缓存
+    if model_name in _global_tokenizers:
+        return _global_tokenizers[model_name]
+    
     try:
         tokenizer = AutoTokenizer.from_pretrained(
             model_name,
@@ -158,11 +353,15 @@ def load_tokenizer() -> AutoTokenizer:
             local_files_only=False
         )
     
+    # 保存到全局缓存
+    _global_tokenizers[model_name] = tokenizer
+    
     return tokenizer # type: ignore
 
+@functools.lru_cache(maxsize=1)
 def load_rerank_model() -> Tuple[AutoModelForSequenceClassification, AutoTokenizer]:
     """
-    加载重排序模型和分词器。
+    加载重排序模型和分词器，使用缓存避免重复加载。
     
     返回：
         Tuple[AutoModelForSequenceClassification, AutoTokenizer]: 加载的模型和分词器
@@ -170,7 +369,10 @@ def load_rerank_model() -> Tuple[AutoModelForSequenceClassification, AutoTokeniz
     model_name = "BAAI/bge-reranker-v2-m3"
     cache_dir = os.path.expanduser("~/.cache/huggingface/hub")
     
-    PrettyOutput.print(f"加载重排序模型: {model_name}...", OutputType.INFO)
+    # 检查全局缓存
+    key = f"rerank_{model_name}"
+    if key in _global_models and f"{key}_tokenizer" in _global_tokenizers:
+        return _global_models[key], _global_tokenizers[f"{key}_tokenizer"]
     
     try:
         tokenizer = AutoTokenizer.from_pretrained(
@@ -199,4 +401,101 @@ def load_rerank_model() -> Tuple[AutoModelForSequenceClassification, AutoTokeniz
         model = model.cuda()
     model.eval()
     
+    # 保存到全局缓存
+    _global_models[key] = model
+    _global_tokenizers[f"{key}_tokenizer"] = tokenizer
+    
     return model, tokenizer # type: ignore
+
+def rerank_results(query: str, documents: List[str], initial_scores: Optional[List[float]] = None, 
+                  batch_size: int = 8, spinner: Optional[Yaspin] = None) -> List[float]:
+    """
+    使用交叉编码器重排序检索结果，提高RAG精度。
+    
+    参数：
+        query: 查询文本
+        documents: 要重排序的文档内容列表
+        initial_scores: 初始检索分数，可选。如果提供，将与重排序分数融合
+        batch_size: 批处理大小
+        spinner: 可选的进度指示器
+        
+    返回：
+        List[float]: 重排序后的分数列表，与输入文档对应
+    """
+    try:
+        if not documents:
+            return []
+            
+        # 加载重排序模型
+        if spinner:
+            spinner.text = "加载重排序模型..."
+        model, tokenizer = load_rerank_model()
+        
+        # 准备评分
+        all_scores = []
+        
+        # 批量处理
+        for i in range(0, len(documents), batch_size):
+            if spinner:
+                spinner.text = f"重排序进度: {i}/{len(documents)}..."
+                
+            # 准备当前批次
+            batch_docs = documents[i:i+batch_size]
+            pairs = [(query, doc) for doc in batch_docs]
+            
+            # 编码输入
+            with torch.no_grad():
+                # 使用类型忽略以避免mypy错误
+                inputs = tokenizer(  # type: ignore
+                    pairs,
+                    padding=True,
+                    truncation=True,
+                    return_tensors="pt",
+                    max_length=512
+                )
+                
+                # 使用GPU加速（如果可用）
+                if torch.cuda.is_available():
+                    inputs = {k: v.cuda() for k, v in inputs.items()}
+                
+                # 获取分数
+                outputs = model(**inputs)  # type: ignore
+                scores = outputs.logits.squeeze(-1).cpu().tolist()
+                
+                # 如果只有一个文档，确保返回列表
+                if len(batch_docs) == 1:
+                    all_scores.append(float(scores))
+                else:
+                    all_scores.extend(scores)
+        
+        # 归一化分数到0-1范围
+        if all_scores:
+            min_score = min(all_scores)
+            max_score = max(all_scores)
+            if max_score > min_score:
+                normalized_scores = [(score - min_score) / (max_score - min_score) for score in all_scores]
+            else:
+                normalized_scores = [0.5] * len(all_scores)
+                
+            # 融合初始分数（如果提供）
+            if initial_scores and len(initial_scores) == len(normalized_scores):
+                # 使用加权平均融合分数：初始分数权重0.3，重排序分数权重0.7
+                final_scores = [0.3 * init_score + 0.7 * rerank_score 
+                               for init_score, rerank_score in zip(initial_scores, normalized_scores)]
+                return final_scores
+                
+            return normalized_scores
+            
+        if spinner:
+            spinner.text = "重排序完成"
+            
+        # 如果重排序失败，返回初始分数或默认分数
+        return initial_scores if initial_scores else [0.5] * len(documents)
+        
+    except Exception as e:
+        PrettyOutput.print(f"重排序失败: {str(e)}", OutputType.ERROR)
+        if spinner:
+            spinner.text = f"重排序失败: {str(e)}"
+            
+        # 发生错误时回退到初始分数
+        return initial_scores if initial_scores else [0.5] * len(documents)

jarvis/jarvis_utils/git_utils.py

@@ -14,9 +14,25 @@ import subprocess
 from typing import List, Tuple, Dict
 from jarvis.jarvis_utils.output import PrettyOutput, OutputType
 def find_git_root(start_dir="."):
-    """切换到给定路径的Git根目录"""
+    """
+    切换到给定路径的Git根目录，如果不是Git仓库则初始化。
+    
+    参数:
+        start_dir (str): 起始查找目录，默认为当前目录。
+    
+    返回:
+        str: Git仓库根目录路径。如果目录不是Git仓库，则会初始化一个新的Git仓库。
+    """
     os.chdir(start_dir)
-    git_root = os.popen("git rev-parse --show-toplevel").read().strip()
+    try:
+        git_root = os.popen("git rev-parse --show-toplevel").read().strip()
+        if not git_root:
+            subprocess.run(["git", "init"], check=True)
+            git_root = os.path.abspath(".")
+    except subprocess.CalledProcessError:
+        # 如果不是Git仓库，初始化一个新的
+        subprocess.run(["git", "init"], check=True)
+        git_root = os.path.abspath(".")
     os.chdir(git_root)
     return git_root
 def has_uncommitted_changes():

jarvis/jarvis_utils/input.py

@@ -15,7 +15,7 @@ from prompt_toolkit.document import Document
 from prompt_toolkit.key_binding import KeyBindings
 from fuzzywuzzy import process
 from colorama import Fore, Style as ColoramaStyle
-from ..jarvis_utils.output import PrettyOutput, OutputType
+from jarvis.jarvis_utils.output import PrettyOutput, OutputType
 def get_single_line_input(tip: str) -> str:
     """
     获取支持历史记录的单行输入。
@@ -74,10 +74,13 @@ class FileCompleter(Completer):
         # 添加默认建议
         if not text_after_at.strip():
             # 默认建议列表
+            from jarvis.jarvis_utils.utils import ot
             default_suggestions = [
-                ('<CodeBase>', '查询代码库'),
-                ('<Web>', '网页搜索'),
-                ('<RAG>', '知识库检索')
+                (ot("CodeBase"), '查询代码库'),
+                (ot("Web"), '网页搜索'),
+                (ot("RAG"), '知识库检索'),
+                (ot("Summary"), '总结'),
+                (ot("Clear"), '清除历史'),
             ]
             for name, desc in default_suggestions:
                 yield Completion(