jarvis-ai-assistant 0.1.96__py3-none-any.whl → 0.1.98__py3-none-any.whl

jarvis/jarvis_codebase/main.py

@@ -3,6 +3,8 @@ import os
 import numpy as np
 import faiss
 from typing import List, Tuple, Optional, Dict
+
+import yaml
 from jarvis.models.registry import PlatformRegistry
 import concurrent.futures
 from threading import Lock
@@ -10,10 +12,10 @@ from concurrent.futures import ThreadPoolExecutor
 from jarvis.utils import OutputType, PrettyOutput, find_git_root, get_file_md5, 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
 import pickle
 import lzma  # 添加 lzma 导入
 from tqdm import tqdm
+import re
 
 class CodeBase:
     def __init__(self, root_dir: str):
@@ -91,19 +93,28 @@ class CodeBase:
 
     def make_description(self, file_path: str, content: str) -> str:
         model = PlatformRegistry.get_global_platform_registry().get_cheap_platform()
-        model.set_suppress_output(True)
-        prompt = f"""请分析以下代码文件，并生成一个详细的描述。描述应该包含以下要点：
-1. 整个文件的功能描述，不超过100个字
-2. 每个全局变量的函数、类型定义、类、方法等代码元素的一句话描述，不超过50字
-
-请用简洁专业的语言描述，突出代码的技术功能，以便后续进行关联代码检索。
-文件路径：{file_path}
-代码内容：
+        if self.thread_count > 1:
+            model.set_suppress_output(True)
+        prompt = f"""Please analyze the following code file and generate a detailed description. The description should include:
+1. Overall file functionality description, no more than 100 characters
+2. One-sentence description (max 50 characters) for each global variable, function, type definition, class, method, and other code elements
+3. 5 potential questions users might ask about this file
+
+Please use concise and professional language, emphasizing technical functionality to facilitate subsequent code retrieval.
+File path: {file_path}
+Code content:
 {content}
 """
-        response = model.chat(prompt)
+        response = model.chat_until_success(prompt)
         return response
 
+    def export(self):
+        """导出当前索引数据到标准输出"""
+        for file_path, data in self.vector_cache.items():
+            print(f"## {file_path}")
+            print(f"- path: {file_path}")
+            print(f"- description: {data['description']}")
+    
     def _save_cache(self):
         """保存缓存数据"""
         try:
@@ -187,14 +198,13 @@ class CodeBase:
                 return cached_vector
                 
             # 读取文件内容并组合信息
-            with open(file_path, "r", encoding="utf-8") as f:
-                content = f.read()[:self.max_context_length]  # 限制文件内容长度
+            content = open(file_path, "r", encoding="utf-8").read()[:self.max_context_length]  # 限制文件内容长度
             
             # 组合文件信息，包含文件内容
             combined_text = f"""
-文件路径: {file_path}
-文件描述: {description}
-文件内容: {content}
+{file_path}
+{description}
+{content}
 """
             vector = self.get_embedding(combined_text)
             
@@ -449,7 +459,7 @@ class CodeBase:
         score = len(matched_keywords) / len(keywords)
         return score
 
-    def rerank_results(self, query: str, initial_results: List[Tuple[str, float, str]]) -> List[Tuple[str, float, str]]:
+    def rerank_results(self, query: str, initial_results: List[Tuple[str, float, str]]) -> List[Tuple[str, float]]:
         """使用多种策略对搜索结果重新排序"""
         if not initial_results:
             return []
@@ -465,16 +475,15 @@ class CodeBase:
             
             for path, _, desc in initial_results:
                 try:
-                    with open(path, "r", encoding="utf-8") as f:
-                        content = f.read()[:512]  # 限制内容长度
+                    content = open(path, "r", encoding="utf-8").read()[:512]  # 限制内容长度
                     
                     # 组合文件信息
-                    doc_content = f"文件: {path}\n描述: {desc}\n内容: {content}"
+                    doc_content = f"File path: {path}\nDescription: {desc}\nContent: {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}"
+                    doc_content = f"File path: {path}\nDescription: {desc}"
                     pairs.append([query, doc_content])
             
             # 使用更大的batch size提高处理速度
@@ -507,11 +516,9 @@ class CodeBase:
             
             # 将重排序分数与原始分数结合
             scored_results = []
-            for (path, orig_score, desc), rerank_score in zip(initial_results, batch_scores):
-                # 综合分数 = 0.3 * 原始分数 + 0.7 * 重排序分数
-                combined_score = 0.3 * float(orig_score) + 0.7 * float(rerank_score)
-                if combined_score >= 0.5:  # 只保留相关度较高的结果
-                    scored_results.append((path, combined_score, desc))
+            for (path,_, desc), rerank_score in zip(initial_results, batch_scores):
+                if rerank_score >= 0.5:  # 只保留相关度较高的结果
+                    scored_results.append((path, rerank_score))
                     
             # 按综合分数降序排序
             scored_results.sort(key=lambda x: x[1], reverse=True)
@@ -521,7 +528,7 @@ class CodeBase:
         except Exception as e:
             PrettyOutput.print(f"重排序失败: {str(e)}", 
                             output_type=OutputType.ERROR)
-            return initial_results  # 发生错误时返回原始结果
+            return [(path, score) for path, score, _ in initial_results]  # 发生错误时返回原始结果
 
     def _generate_query_variants(self, query: str) -> List[str]:
         """生成查询的不同表述变体
@@ -533,12 +540,12 @@ class CodeBase:
             List[str]: 查询变体列表
         """
         model = PlatformRegistry.get_global_platform_registry().get_normal_platform()
-        prompt = f"""请根据以下查询，生成3个不同的表述，每个表述都要完整表达原始查询的意思。这些表述将用于代码搜索，要保持专业性和准确性。
-原始查询: {query}
+        prompt = f"""Please generate 3 different expressions based on the following query, each expression should fully convey the meaning of the original query. These expressions will be used for code search, maintain professionalism and accuracy.
+Original query: {query}
 
-请直接输出3个表述，用换行分隔，不要有编号或其他标记。
+Please output 3 expressions directly, separated by two line breaks, without numbering or other markers.
 """
-        variants = model.chat(prompt).strip().split('\n')
+        variants = model.chat_until_success(prompt).strip().split('\n\n')
         variants.append(query)  # 添加原始查询
         return variants
 
@@ -574,7 +581,7 @@ class CodeBase:
         return results
 
 
-    def search_similar(self, query: str, top_k: int = 30) -> List[Tuple[str, float, str]]:
+    def search_similar(self, query: str, top_k: int = 30) -> List[Tuple[str, float]]:
         """搜索关联文件"""
         try:
             if self.index is None:
@@ -594,6 +601,9 @@ class CodeBase:
             # 如果没有找到结果，直接返回
             if not initial_results:
                 return []
+            
+            # 过滤低分结果
+            initial_results = [(path, score, desc) for path, score, desc in initial_results if score >= 0.5]
                 
             # 对初步结果进行重排序
             return self.rerank_results(query, initial_results)
@@ -610,21 +620,21 @@ class CodeBase:
             return ""
         
         PrettyOutput.print(f"找到的关联文件: ", output_type=OutputType.SUCCESS)
-        for path, score, _ in results:
+        for path, score in results:
             PrettyOutput.print(f"文件: {path} 关联度: {score:.3f}", 
                              output_type=OutputType.INFO)
         
         prompt = f"""你是一个代码专家，请根据以下文件信息回答用户的问题：
 """
-        for path, _, _ in results:
+        for path, _ in results:
             try:
                 if len(prompt) > self.max_context_length:
                     PrettyOutput.print(f"避免上下文超限，丢弃低相关度文件：{path}", OutputType.WARNING)
                     continue
                 content = open(path, "r", encoding="utf-8").read()
                 prompt += f"""
-文件路径: {path}prompt
-文件内容:
+File path: {path}prompt
+File content:
 {content}
 ========================================
 """
@@ -634,12 +644,12 @@ class CodeBase:
                 continue
                 
         prompt += f"""
-用户问题: {query}
+User question: {query}
 
-请用专业的语言回答用户的问题，如果给出的文件内容不足以回答用户的问题，请告诉用户，绝对不要胡编乱造。
+Please answer the user's question in Chinese using professional language. If the provided file content is insufficient to answer the user's question, please inform the user. Never make up information.
 """
         model = PlatformRegistry.get_global_platform_registry().get_codegen_platform()
-        response = model.chat(prompt)
+        response = model.chat_until_success(prompt)
         return response
 
     def is_index_generated(self) -> bool:
@@ -669,46 +679,69 @@ class CodeBase:
 
 
 
+
+
 def main():
+
     parser = argparse.ArgumentParser(description='Codebase management and search tool')
-    parser.add_argument('--search', type=str, help='Search query to find similar code files')
-    parser.add_argument('--top-k', type=int, default=20, help='Number of results to return (default: 20)')
-    parser.add_argument('--ask', type=str, help='Ask a question about the codebase')
-    parser.add_argument('--generate', action='store_true', help='Generate codebase index')
+    subparsers = parser.add_subparsers(dest='command', help='Available commands')
+
+    # Generate command
+    generate_parser = subparsers.add_parser('generate', help='Generate codebase index')
+    generate_parser.add_argument('--force', action='store_true', help='Force rebuild index')
+
+    # Search command
+    search_parser = subparsers.add_parser('search', help='Search similar code files')
+    search_parser.add_argument('query', type=str, help='Search query')
+    search_parser.add_argument('--top-k', type=int, default=20, help='Number of results to return (default: 20)')
+
+    # Ask command
+    ask_parser = subparsers.add_parser('ask', help='Ask a question about the codebase')
+    ask_parser.add_argument('question', type=str, help='Question to ask')
+    ask_parser.add_argument('--top-k', type=int, default=20, help='Number of results to use (default: 20)')
+
+    export_parser = subparsers.add_parser('export', help='Export current index data')
     args = parser.parse_args()
     
     current_dir = find_git_root()
     codebase = CodeBase(current_dir)
 
-    # 如果没有生成索引，且不是生成命令，提示用户先生成索引
-    if not codebase.is_index_generated() and not args.generate:
-        PrettyOutput.print("索引尚未生成，请先运行 --generate 生成索引", output_type=OutputType.WARNING)
+    if args.command == 'export':
+        codebase.export()
         return
 
+    # 如果没有生成索引，且不是生成命令，提示用户先生成索引
+    if not codebase.is_index_generated() and args.command != 'generate':
+        PrettyOutput.print("索引尚未生成，请先运行 'generate' 命令生成索引", output_type=OutputType.WARNING)
+        return
 
-    if args.generate:
+    if args.command == 'generate':
         try:
-            codebase.generate_codebase(force=True)
+            codebase.generate_codebase(force=args.force)
             PrettyOutput.print("\nCodebase generation completed", output_type=OutputType.SUCCESS)
         except Exception as e:
             PrettyOutput.print(f"Error during codebase generation: {str(e)}", output_type=OutputType.ERROR)
     
-    if args.search:
-        results = codebase.search_similar(args.search, args.top_k)
+    elif args.command == 'search':
+        results = codebase.search_similar(args.query, args.top_k)
         if not results:
             PrettyOutput.print("No similar files found", output_type=OutputType.WARNING)
             return
             
         PrettyOutput.print("\nSearch Results:", output_type=OutputType.INFO)
-        for path, score, desc in results:
+        for path, score in results:
             PrettyOutput.print("\n" + "="*50, output_type=OutputType.INFO)
             PrettyOutput.print(f"File: {path}", output_type=OutputType.INFO)
             PrettyOutput.print(f"Similarity: {score:.3f}", output_type=OutputType.INFO)
-            PrettyOutput.print(f"Description: {desc[100:]}", output_type=OutputType.INFO)
 
-    if args.ask:            
-        codebase.ask_codebase(args.ask, args.top_k)
+    elif args.command == 'ask':            
+        response = codebase.ask_codebase(args.question, args.top_k)
+        PrettyOutput.print("\nAnswer:", output_type=OutputType.INFO)
+        PrettyOutput.print(response, output_type=OutputType.INFO)
+
+    else:
+        parser.print_help()
 
 
 if __name__ == "__main__":
-    exit(main())
+    exit(main())

jarvis/jarvis_coder/git_utils.py

@@ -6,51 +6,48 @@ from jarvis.utils import OutputType, PrettyOutput, while_success
 from jarvis.models.registry import PlatformRegistry
 
 def has_uncommitted_files() -> bool:
-    """判断代码库是否有未提交的文件"""
-    # 获取未暂存的修改
+    """Check if there are uncommitted files in the repository"""
+    # Get unstaged modifications
     unstaged = os.popen("git diff --name-only").read()
-    # 获取已暂存但未提交的修改
+    # Get staged but uncommitted modifications
     staged = os.popen("git diff --cached --name-only").read()
-    # 获取未跟踪的文件
+    # Get untracked files
     untracked = os.popen("git ls-files --others --exclude-standard").read()
     
     return bool(unstaged or staged or untracked)
 
-def generate_commit_message(git_diff: str, feature: str) -> str:
-    """根据git diff和功能描述生成commit信息"""
-    prompt = f"""你是一个经验丰富的程序员，请根据以下代码变更和功能描述生成简洁明了的commit信息：
+def generate_commit_message(git_diff: str) -> str:
+    """Generate commit message based on git diff and feature description"""
+    prompt = f"""You are an experienced programmer, please generate a concise and clear commit message based on the following code changes and feature description:
 
-功能描述：
-{feature}
-
-代码变更：
+Code changes:
 Git Diff:
 {git_diff}
 
-请遵循以下规则：
-1. 使用英文编写
-2. 采用常规的commit message格式：<type>(<scope>): <subject>
-3. 保持简洁，不超过50个字符
-4. 准确描述代码变更的主要内容
-5. 优先考虑功能描述和git diff中的变更内容
-6. 仅生成commit信息的文本，不要输出任何其他内容
+Please follow these rules:
+1. Write in English
+2. Use conventional commit message format: <type>(<scope>): <subject>
+3. Keep it concise, no more than 50 characters
+4. Accurately describe the main content of code changes
+5. Prioritize feature description and changes in git diff
+6. Only generate the commit message text, do not output anything else
 """
     
-    model = PlatformRegistry().get_global_platform_registry().get_codegen_platform()
-    response = while_success(lambda: model.chat(prompt), 5)
+    model = PlatformRegistry().get_global_platform_registry().get_normal_platform()
+    response = model.chat_until_success(prompt)
         
-    return response.strip().split("\n")[0]
+    return ';'.join(response.strip().split("\n"))
 
 def save_edit_record(record_dir: str, commit_message: str, git_diff: str) -> None:
-    """保存代码修改记录"""
-    # 获取下一个序号
+    """Save code modification record"""
+    # Get next sequence number
     existing_records = [f for f in os.listdir(record_dir) if f.endswith('.yaml')]
     next_num = 1
     if existing_records:
         last_num = max(int(f[:4]) for f in existing_records)
         next_num = last_num + 1
     
-    # 创建记录文件
+    # Create record file
     record = {
         "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
         "commit_message": commit_message,
@@ -61,4 +58,4 @@ def save_edit_record(record_dir: str, commit_message: str, git_diff: str) -> Non
     with open(record_path, "w", encoding="utf-8") as f:
         yaml.safe_dump(record, f, allow_unicode=True)
     
-    PrettyOutput.print(f"已保存修改记录: {record_path}", OutputType.SUCCESS) 
+    PrettyOutput.print(f"Modification record saved: {record_path}", OutputType.SUCCESS)