jarvis-ai-assistant 0.1.96__py3-none-any.whl → 0.1.97__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

@@ -16,13 +16,10 @@ def has_uncommitted_files() -> bool:
     
     return bool(unstaged or staged or untracked)
 
-def generate_commit_message(git_diff: str, feature: str) -> str:
+def generate_commit_message(git_diff: str) -> str:
     """根据git diff和功能描述生成commit信息"""
     prompt = f"""你是一个经验丰富的程序员，请根据以下代码变更和功能描述生成简洁明了的commit信息：
 
-功能描述：
-{feature}
-
 代码变更：
 Git Diff:
 {git_diff}
@@ -36,10 +33,10 @@ Git Diff:
 6. 仅生成commit信息的文本，不要输出任何其他内容
 """
     
-    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:
     """保存代码修改记录"""

jarvis/jarvis_coder/main.py

@@ -97,7 +97,7 @@ class JarvisCoder:
             PrettyOutput.print("提交未保存的更改", OutputType.INFO)
             os.system("git add .")
             git_diff = os.popen("git diff --cached").read()
-            commit_message = generate_commit_message(git_diff, "Pre-edit commit")
+            commit_message = generate_commit_message(git_diff)
             os.system(f"git commit -m '{commit_message}'")
 
     def _init_codebase(self):
@@ -129,10 +129,9 @@ class JarvisCoder:
             self._codebase.generate_codebase()
             
         related_files = self._codebase.search_similar(feature)
-        for file, score, _ in related_files:
+        for file, score in related_files:
             PrettyOutput.print(f"相关文件: {file} 相关度: {score:.3f}", OutputType.SUCCESS)
-            with open(file, "r", encoding="utf-8") as f:
-                content = f.read()
+            content = open(file, "r", encoding="utf-8").read()
             ret.append({"file_path": file, "file_content": content})
         return ret
 
@@ -354,7 +353,7 @@ class JarvisCoder:
         git_diff = os.popen("git diff --cached").read()
         
         # 自动生成commit信息，传入feature
-        commit_message = generate_commit_message(git_diff, feature)
+        commit_message = generate_commit_message(git_diff)
         
         # 显示并确认commit信息
         PrettyOutput.print(f"自动生成的commit信息: {commit_message}", OutputType.INFO)
@@ -381,28 +380,24 @@ class JarvisCoder:
             model.set_suppress_output(True)
             file_path = file_info["file_path"]
             content = file_info["file_content"]
-            
-            # 生成分析提示
-            system_message = f"""你是一个代码分析专家，可以从代码中提取出与需求相关的片段。
-请按以下格式返回：
+
+            try:
+                prompt = f"""You are a code analysis expert who can extract relevant snippets from code.
+Please return in the following format:
 <PART>
 content
 </PART>
 
-可返回多个片段。如果文件内容与需求无关，则返回空。
-"""
-            model.set_system_message(system_message)
-            
-            try:
+Multiple snippets can be returned. If the file content is not relevant to the requirement, return empty.
 
-                prompt = f"""需求：{feature}
-文件路径：{file_path}
-代码内容：
+Requirement: {feature}
+File path: {file_path}
+Code content:
 {content}
 """
 
                 # 调用大模型进行分析
-                response = while_success(lambda: model.chat(prompt))
+                response = model.chat_until_success(prompt)
 
                 parts = re.findall(r'<PART>\n(.*?)\n</PART>', response, re.DOTALL)
                 file_info["parts"] = parts
@@ -426,16 +421,16 @@ content
                     file["parts"] = [file["file_content"]]
             
             # 获取修改方案
-            modification_plan = PlanGenerator().generate_plan(feature, selected_files)
-            if not modification_plan:
+            raw_plan, structed_plan = PlanGenerator().generate_plan(feature, selected_files)
+            if not raw_plan or not structed_plan:
                 return {
                     "success": False,
                     "stdout": "",
-                    "stderr": "用户取消修改",
+                    "stderr": "修改计划生成失败，请修改需求后重试",
                 }
             
             # 执行修改
-            if PatchHandler().handle_patch_application(selected_files, feature, modification_plan):
+            if PatchHandler().handle_patch_application(feature ,raw_plan, structed_plan):
                 self._finalize_changes(feature)
                 return {
                     "success": True,