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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of jarvis-ai-assistant might be problematic. Click here for more details.

Files changed (47) hide show
  1. {jarvis_ai_assistant-0.1.85/src/jarvis_ai_assistant.egg-info → jarvis_ai_assistant-0.1.87}/PKG-INFO +2 -4
  2. {jarvis_ai_assistant-0.1.85 → jarvis_ai_assistant-0.1.87}/README.md +1 -3
  3. {jarvis_ai_assistant-0.1.85 → jarvis_ai_assistant-0.1.87}/pyproject.toml +1 -1
  4. {jarvis_ai_assistant-0.1.85 → jarvis_ai_assistant-0.1.87}/setup.py +1 -1
  5. {jarvis_ai_assistant-0.1.85 → jarvis_ai_assistant-0.1.87}/src/jarvis/__init__.py +1 -1
  6. {jarvis_ai_assistant-0.1.85 → jarvis_ai_assistant-0.1.87}/src/jarvis/agent.py +8 -7
  7. {jarvis_ai_assistant-0.1.85 → jarvis_ai_assistant-0.1.87}/src/jarvis/jarvis_codebase/main.py +115 -100
  8. {jarvis_ai_assistant-0.1.85 → jarvis_ai_assistant-0.1.87}/src/jarvis/jarvis_coder/main.py +104 -36
  9. {jarvis_ai_assistant-0.1.85 → jarvis_ai_assistant-0.1.87}/src/jarvis/jarvis_rag/main.py +191 -148
  10. {jarvis_ai_assistant-0.1.85 → jarvis_ai_assistant-0.1.87}/src/jarvis/models/ai8.py +1 -1
  11. jarvis_ai_assistant-0.1.87/src/jarvis/models/ollama.py +150 -0
  12. {jarvis_ai_assistant-0.1.85 → jarvis_ai_assistant-0.1.87}/src/jarvis/models/openai.py +2 -2
  13. {jarvis_ai_assistant-0.1.85 → jarvis_ai_assistant-0.1.87}/src/jarvis/models/oyi.py +3 -3
  14. {jarvis_ai_assistant-0.1.85 → jarvis_ai_assistant-0.1.87}/src/jarvis/tools/ask_user.py +13 -15
  15. {jarvis_ai_assistant-0.1.85 → jarvis_ai_assistant-0.1.87}/src/jarvis/tools/registry.py +10 -8
  16. {jarvis_ai_assistant-0.1.85 → jarvis_ai_assistant-0.1.87}/src/jarvis/utils.py +48 -2
  17. {jarvis_ai_assistant-0.1.85 → jarvis_ai_assistant-0.1.87/src/jarvis_ai_assistant.egg-info}/PKG-INFO +2 -4
  18. {jarvis_ai_assistant-0.1.85 → jarvis_ai_assistant-0.1.87}/src/jarvis_ai_assistant.egg-info/SOURCES.txt +1 -0
  19. {jarvis_ai_assistant-0.1.85 → jarvis_ai_assistant-0.1.87}/LICENSE +0 -0
  20. {jarvis_ai_assistant-0.1.85 → jarvis_ai_assistant-0.1.87}/MANIFEST.in +0 -0
  21. {jarvis_ai_assistant-0.1.85 → jarvis_ai_assistant-0.1.87}/setup.cfg +0 -0
  22. {jarvis_ai_assistant-0.1.85 → jarvis_ai_assistant-0.1.87}/src/jarvis/jarvis_codebase/__init__.py +0 -0
  23. {jarvis_ai_assistant-0.1.85 → jarvis_ai_assistant-0.1.87}/src/jarvis/jarvis_coder/__init__.py +0 -0
  24. {jarvis_ai_assistant-0.1.85 → jarvis_ai_assistant-0.1.87}/src/jarvis/jarvis_rag/__init__.py +0 -0
  25. {jarvis_ai_assistant-0.1.85 → jarvis_ai_assistant-0.1.87}/src/jarvis/jarvis_smart_shell/__init__.py +0 -0
  26. {jarvis_ai_assistant-0.1.85 → jarvis_ai_assistant-0.1.87}/src/jarvis/jarvis_smart_shell/main.py +0 -0
  27. {jarvis_ai_assistant-0.1.85 → jarvis_ai_assistant-0.1.87}/src/jarvis/main.py +0 -0
  28. {jarvis_ai_assistant-0.1.85 → jarvis_ai_assistant-0.1.87}/src/jarvis/models/__init__.py +0 -0
  29. {jarvis_ai_assistant-0.1.85 → jarvis_ai_assistant-0.1.87}/src/jarvis/models/base.py +0 -0
  30. {jarvis_ai_assistant-0.1.85 → jarvis_ai_assistant-0.1.87}/src/jarvis/models/kimi.py +0 -0
  31. {jarvis_ai_assistant-0.1.85 → jarvis_ai_assistant-0.1.87}/src/jarvis/models/registry.py +0 -0
  32. {jarvis_ai_assistant-0.1.85 → jarvis_ai_assistant-0.1.87}/src/jarvis/tools/__init__.py +0 -0
  33. {jarvis_ai_assistant-0.1.85 → jarvis_ai_assistant-0.1.87}/src/jarvis/tools/base.py +0 -0
  34. {jarvis_ai_assistant-0.1.85 → jarvis_ai_assistant-0.1.87}/src/jarvis/tools/chdir.py +0 -0
  35. {jarvis_ai_assistant-0.1.85 → jarvis_ai_assistant-0.1.87}/src/jarvis/tools/codebase_qa.py +0 -0
  36. {jarvis_ai_assistant-0.1.85 → jarvis_ai_assistant-0.1.87}/src/jarvis/tools/coder.py +0 -0
  37. {jarvis_ai_assistant-0.1.85 → jarvis_ai_assistant-0.1.87}/src/jarvis/tools/file_ops.py +0 -0
  38. {jarvis_ai_assistant-0.1.85 → jarvis_ai_assistant-0.1.87}/src/jarvis/tools/generator.py +0 -0
  39. {jarvis_ai_assistant-0.1.85 → jarvis_ai_assistant-0.1.87}/src/jarvis/tools/methodology.py +0 -0
  40. {jarvis_ai_assistant-0.1.85 → jarvis_ai_assistant-0.1.87}/src/jarvis/tools/search.py +0 -0
  41. {jarvis_ai_assistant-0.1.85 → jarvis_ai_assistant-0.1.87}/src/jarvis/tools/shell.py +0 -0
  42. {jarvis_ai_assistant-0.1.85 → jarvis_ai_assistant-0.1.87}/src/jarvis/tools/sub_agent.py +0 -0
  43. {jarvis_ai_assistant-0.1.85 → jarvis_ai_assistant-0.1.87}/src/jarvis/tools/webpage.py +0 -0
  44. {jarvis_ai_assistant-0.1.85 → jarvis_ai_assistant-0.1.87}/src/jarvis_ai_assistant.egg-info/dependency_links.txt +0 -0
  45. {jarvis_ai_assistant-0.1.85 → jarvis_ai_assistant-0.1.87}/src/jarvis_ai_assistant.egg-info/entry_points.txt +0 -0
  46. {jarvis_ai_assistant-0.1.85 → jarvis_ai_assistant-0.1.87}/src/jarvis_ai_assistant.egg-info/requires.txt +0 -0
  47. {jarvis_ai_assistant-0.1.85 → jarvis_ai_assistant-0.1.87}/src/jarvis_ai_assistant.egg-info/top_level.txt +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.2
2
2
  Name: jarvis-ai-assistant
3
- Version: 0.1.85
3
+ Version: 0.1.87
4
4
  Summary: Jarvis: An AI assistant that uses tools to interact with the system
5
5
  Home-page: https://github.com/skyfireitdiy/Jarvis
6
6
  Author: skyfire
@@ -134,15 +134,13 @@ Jarvis supports configuration through environment variables that can be set in t
134
134
  | JARVIS_CODEGEN_MODEL | Model name for code generation | Same as JARVIS_MODEL | No |
135
135
  | JARVIS_CHEAP_PLATFORM | AI platform for cheap operations | Same as JARVIS_PLATFORM | No |
136
136
  | JARVIS_CHEAP_MODEL | Model name for cheap operations | Same as JARVIS_MODEL | No |
137
- | JARVIS_EMBEDDING_MODEL | Embedding model for code analysis | BAAI/bge-large-zh-v1.5 | No |
138
137
  | OPENAI_API_KEY | API key for OpenAI platform | - | Required for OpenAI |
139
138
  | OPENAI_API_BASE | Base URL for OpenAI API | https://api.deepseek.com | No |
140
139
  | OPENAI_MODEL_NAME | Model name for OpenAI | deepseek-chat | No |
141
140
  | AI8_API_KEY | API key for AI8 platform | - | Required for AI8 |
142
- | AI8_MODEL | Model name for AI8 platform | deepseek-chat | No |
143
141
  | KIMI_API_KEY | API key for Kimi platform | - | Required for Kimi |
144
142
  | OYI_API_KEY | API key for OYI platform | - | Required for OYI |
145
- | OYI_MODEL | Model name for OYI platform | deepseek-chat | No |
143
+ | OLLAMA_API_BASE | Base URL for Ollama API | http://localhost:11434 | No |
146
144
 
147
145
 
148
146
  ## 🎯 Usage
@@ -70,15 +70,13 @@ Jarvis supports configuration through environment variables that can be set in t
70
70
  | JARVIS_CODEGEN_MODEL | Model name for code generation | Same as JARVIS_MODEL | No |
71
71
  | JARVIS_CHEAP_PLATFORM | AI platform for cheap operations | Same as JARVIS_PLATFORM | No |
72
72
  | JARVIS_CHEAP_MODEL | Model name for cheap operations | Same as JARVIS_MODEL | No |
73
- | JARVIS_EMBEDDING_MODEL | Embedding model for code analysis | BAAI/bge-large-zh-v1.5 | No |
74
73
  | OPENAI_API_KEY | API key for OpenAI platform | - | Required for OpenAI |
75
74
  | OPENAI_API_BASE | Base URL for OpenAI API | https://api.deepseek.com | No |
76
75
  | OPENAI_MODEL_NAME | Model name for OpenAI | deepseek-chat | No |
77
76
  | AI8_API_KEY | API key for AI8 platform | - | Required for AI8 |
78
- | AI8_MODEL | Model name for AI8 platform | deepseek-chat | No |
79
77
  | KIMI_API_KEY | API key for Kimi platform | - | Required for Kimi |
80
78
  | OYI_API_KEY | API key for OYI platform | - | Required for OYI |
81
- | OYI_MODEL | Model name for OYI platform | deepseek-chat | No |
79
+ | OLLAMA_API_BASE | Base URL for Ollama API | http://localhost:11434 | No |
82
80
 
83
81
 
84
82
  ## 🎯 Usage
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
4
4
 
5
5
  [project]
6
6
  name = "jarvis-ai-assistant"
7
- version = "0.1.85"
7
+ version = "0.1.87"
8
8
  description = "Jarvis: An AI assistant that uses tools to interact with the system"
9
9
  readme = "README.md"
10
10
  authors = [{ name = "Your Name", email = "your.email@example.com" }]
@@ -2,7 +2,7 @@ from setuptools import setup, find_packages
2
2
 
3
3
  setup(
4
4
  name="jarvis-ai-assistant",
5
- version="0.1.85",
5
+ version="0.1.87",
6
6
  author="skyfire",
7
7
  author_email="skyfireitdiy@hotmail.com",
8
8
  description="An AI assistant that uses various tools to interact with the system",
@@ -1,3 +1,3 @@
1
1
  """Jarvis AI Assistant"""
2
2
 
3
- __version__ = "0.1.85"
3
+ __version__ = "0.1.87"
@@ -8,7 +8,7 @@ import json
8
8
 
9
9
  from .models.registry import PlatformRegistry
10
10
  from .tools import ToolRegistry
11
- from .utils import PrettyOutput, OutputType, get_multiline_input, load_embedding_model, while_success
11
+ from .utils import PrettyOutput, OutputType, get_max_context_length, get_multiline_input, load_embedding_model, while_success
12
12
  import os
13
13
  from datetime import datetime
14
14
  from prompt_toolkit import prompt
@@ -33,7 +33,7 @@ class Agent:
33
33
 
34
34
  # 从环境变量加载配置
35
35
  self.embedding_dimension = 1536 # Default for many embedding models
36
- self.max_context_length = int(os.getenv('JARVIS_MAX_CONTEXT_LENGTH', '65536')) # 默认30k
36
+ self.max_context_length = get_max_context_length()
37
37
 
38
38
  # 初始化嵌入模型
39
39
  try:
@@ -72,10 +72,10 @@ class Agent:
72
72
 
73
73
  # 逐行处理
74
74
  for line in lines:
75
- if '<START_TOOL_CALL>' in line:
75
+ if '<TOOL_CALL>' in line:
76
76
  in_tool_call = True
77
77
  continue
78
- elif '<END_TOOL_CALL>' in line:
78
+ elif '</TOOL_CALL>' in line:
79
79
  if in_tool_call and tool_call_lines:
80
80
  try:
81
81
  # 直接解析YAML
@@ -344,7 +344,7 @@ class Agent:
344
344
  4. 生成解决方案:生成一个或者多个具备可操作性的解决方案
345
345
  5. 评估解决方案:从众多解决方案中选择一种最优的方案
346
346
  6. 制定行动计划:根据目前可以使用的工具制定行动计划,使用 PlantUML 格式输出明确的执行流程
347
- 7. 执行行动计划:每步执行一个步骤,最多使用一个工具(工具执行完成后,等待工具结果再执行下一步)
347
+ 7. 执行行动计划:每步执行一个步骤,**最多使用一个工具**(工具执行完成后,等待工具结果再执行下一步)
348
348
  8. 监控与调整:如果执行结果与预期不符,则反思并调整行动计划,迭代之前的步骤
349
349
  9. 方法论:如果当前任务具有通用性且执行过程中遇到了值得记录的经验,使用方法论工具记录方法论,以提升后期处理类似问题的能力
350
350
  10. 任务结束:如果任务已经完成,使用任务结束指令结束任务
@@ -362,12 +362,12 @@ class Agent:
362
362
 
363
363
  工具使用格式:
364
364
 
365
- <START_TOOL_CALL>
365
+ <TOOL_CALL>
366
366
  name: tool_name
367
367
  arguments:
368
368
  param1: value1
369
369
  param2: value2
370
- <END_TOOL_CALL>
370
+ </TOOL_CALL>
371
371
 
372
372
  -------------------------------------------------------------
373
373
 
@@ -381,6 +381,7 @@ arguments:
381
381
  7. 在执行一些可能对系统或者用户代码库造成破坏的工具时,请先询问用户
382
382
  8. 在多次迭代却没有任何进展时,可请求用户指导
383
383
  9. 如果返回的yaml字符串中包含冒号,请将整个字符串用引号包裹,避免yaml解析错误
384
+ 10. yaml中包含多行字符串时,请使用 | 语法
384
385
 
385
386
  {methodology_prompt}
386
387
 
@@ -7,7 +7,7 @@ from jarvis.models.registry import PlatformRegistry
7
7
  import concurrent.futures
8
8
  from threading import Lock
9
9
  from concurrent.futures import ThreadPoolExecutor
10
- from jarvis.utils import OutputType, PrettyOutput, find_git_root, load_embedding_model
10
+ from jarvis.utils import OutputType, PrettyOutput, find_git_root, get_max_context_length, get_thread_count, load_embedding_model, load_rerank_model
11
11
  from jarvis.utils import load_env_from_file
12
12
  import argparse
13
13
  from sentence_transformers import SentenceTransformer
@@ -18,8 +18,8 @@ class CodeBase:
18
18
  load_env_from_file()
19
19
  self.root_dir = root_dir
20
20
  os.chdir(self.root_dir)
21
- self.thread_count = int(os.environ.get("JARVIS_THREAD_COUNT") or 1)
22
- self.max_context_length = int(os.getenv("JARVIS_MAX_CONTEXT_LENGTH", 65536))
21
+ self.thread_count = get_thread_count()
22
+ self.max_context_length = get_max_context_length()
23
23
 
24
24
  # 初始化数据目录
25
25
  self.data_dir = os.path.join(self.root_dir, ".jarvis-codebase")
@@ -98,7 +98,7 @@ class CodeBase:
98
98
  5. 关键业务逻辑和处理流程
99
99
  6. 特殊功能点和亮点特性
100
100
 
101
- 请用简洁专业的语言描述,突出代码的技术特征和功能特点,以便后续进行相似代码检索。
101
+ 请用简洁专业的语言描述,突出代码的技术特征和功能特点,以便后续进行关联代码检索。
102
102
 
103
103
  文件路径:{file_path}
104
104
  代码内容:
@@ -196,10 +196,15 @@ class CodeBase:
196
196
  if cached_vector is not None:
197
197
  return cached_vector
198
198
 
199
- # 组合文件信息
199
+ # 读取文件内容并组合信息
200
+ with open(file_path, "r", encoding="utf-8") as f:
201
+ content = f.read()[:self.max_context_length] # 限制文件内容长度
202
+
203
+ # 组合文件信息,包含文件内容
200
204
  combined_text = f"""
201
205
  文件路径: {file_path}
202
206
  文件描述: {description}
207
+ 文件内容: {content}
203
208
  """
204
209
  vector = self.get_embedding(combined_text)
205
210
 
@@ -398,116 +403,126 @@ class CodeBase:
398
403
  PrettyOutput.print("没有检测到文件变更,无需重建索引", output_type=OutputType.INFO)
399
404
 
400
405
  def rerank_results(self, query: str, initial_results: List[Tuple[str, float, str]]) -> List[Tuple[str, float, str]]:
401
- """使用大模型对搜索结果重新排序"""
406
+ """使用 BAAI/bge-reranker-v2-m3 对搜索结果重新排序"""
402
407
  if not initial_results:
403
408
  return []
404
-
405
- model = PlatformRegistry.get_global_platform_registry().get_codegen_platform()
406
- # model.set_suppress_output(True)
407
-
408
- # 构建重排序的prompt
409
- prompt = f"""请根据用户的查询,对以下代码文件进行相关性排序。对每个文件给出0-100的相关性分数,分数越高表示越相关。
410
- 只需要输出每个文件的分数,格式为:
411
- <RERANK_START>
412
- 文件路径: 分数
413
- 文件路径: 分数
414
- <RERANK_END>
415
-
416
- 用户查询: {query}
417
-
418
- 待评估文件:
419
- """
420
- for path, _, desc in initial_results:
421
- prompt += f"""
422
- 文件: {path}
423
- 描述: {desc}
424
- ---
425
- """
426
-
427
- response = model.chat(prompt)
428
-
429
- # 提取<RERANK_START>和<RERANK_END>之间的内容
430
- start_tag = "<RERANK_START>"
431
- end_tag = "<RERANK_END>"
432
- if start_tag in response and end_tag in response:
433
- response = response[response.find(start_tag) + len(start_tag):response.find(end_tag)]
434
-
435
- # 解析响应,提取文件路径和分数
436
- scored_results = []
437
- for line in response.split('\n'):
438
- if ':' not in line:
439
- continue
440
- try:
441
- file_path, score_str = line.split(':', 1)
442
- file_path = file_path.strip()
443
- score = float(score_str.strip()) / 100.0 # 转换为0-1范围
444
- # 只保留相关度大于等于0.7的结果
445
- if score >= 0.7:
446
- # 找到对应的原始描述
447
- desc = next((desc for p, _, desc in initial_results if p == file_path), "")
448
- scored_results.append((file_path, score, desc))
449
- except:
450
- continue
451
-
452
- # 按分数降序排序
453
- return sorted(scored_results, key=lambda x: x[1], reverse=True)
454
-
409
+
410
+ try:
411
+ import torch
412
+
413
+ # 加载模型和分词器
414
+ model, tokenizer = load_rerank_model()
415
+
416
+ # 准备数据 - 加入文件内容进行更准确的重排序
417
+ pairs = []
418
+ for path, _, desc in initial_results:
419
+ try:
420
+ with open(path, "r", encoding="utf-8") as f:
421
+ content = f.read()[:512] # 限制内容长度
422
+ # 组合文件路径、描述和内容
423
+ doc_content = f"文件: {path}\n描述: {desc}\n内容: {content}"
424
+ pairs.append([query, doc_content])
425
+ except Exception as e:
426
+ PrettyOutput.print(f"读取文件失败 {path}: {str(e)}",
427
+ output_type=OutputType.ERROR)
428
+ doc_content = f"文件: {path}\n描述: {desc}"
429
+ pairs.append([query, doc_content])
430
+
431
+ # 使用更大的batch size提高处理速度
432
+ batch_size = 16 # 根据GPU显存调整
433
+
434
+ with torch.no_grad():
435
+ for i in range(0, len(pairs), batch_size):
436
+ batch_pairs = pairs[i:i + batch_size]
437
+ encoded = tokenizer(
438
+ batch_pairs,
439
+ padding=True,
440
+ truncation=True,
441
+ max_length=512,
442
+ return_tensors='pt'
443
+ )
444
+
445
+ if torch.cuda.is_available():
446
+ encoded = {k: v.cuda() for k, v in encoded.items()}
447
+
448
+ outputs = model(**encoded)
449
+ # 修改这里:直接使用 outputs.logits 作为分数
450
+ batch_scores = outputs.logits.squeeze(-1).cpu().numpy()
451
+
452
+ # 归一化分数到 0-1 范围
453
+ if batch_scores:
454
+ min_score = min(batch_scores)
455
+ max_score = max(batch_scores)
456
+ if max_score > min_score:
457
+ batch_scores = [(s - min_score) / (max_score - min_score) for s in batch_scores]
458
+
459
+ # 将分数与原始结果组合并排序
460
+ scored_results = []
461
+ for (path, _, desc), score in zip(initial_results, batch_scores):
462
+ if score >= 0.5: # 只保留相关度大于 0.5 的结果
463
+ scored_results.append((path, float(score), desc))
464
+
465
+ # 按分数降序排序
466
+ scored_results.sort(key=lambda x: x[1], reverse=True)
467
+
468
+ return scored_results
469
+
470
+ except Exception as e:
471
+ PrettyOutput.print(f"重排序失败,使用原始排序: {str(e)}", output_type=OutputType.WARNING)
472
+ return initial_results
455
473
 
456
474
  def search_similar(self, query: str, top_k: int = 30) -> List[Tuple[str, float, str]]:
457
- """搜索相似文件"""
458
- model = PlatformRegistry.get_global_platform_registry().get_normal_platform()
459
- model.set_suppress_output(True)
460
-
461
-
462
- prompt = f"""请根据以下查询,生成意思完全相同的另一个表述。这个表述将用于代码搜索,所以要保持专业性和准确性。
475
+ """搜索关联文件"""
476
+ try:
477
+ # 生成多个查询变体以提高召回率
478
+ model = PlatformRegistry.get_global_platform_registry().get_normal_platform()
479
+ model.set_suppress_output(True)
480
+
481
+ prompt = f"""请根据以下查询,生成3个不同的表述,每个表述都要完整表达原始查询的意思。这些表述将用于代码搜索,要保持专业性和准确性。
463
482
  原始查询: {query}
464
483
 
465
- 请直接输出新表述,不要有编号或其他标记。
484
+ 请直接输出3个表述,用换行分隔,不要有编号或其他标记。
466
485
  """
467
-
468
- query = model.chat(prompt)
486
+ query_variants = model.chat(prompt).strip().split('\n')
487
+ query_variants.append(query) # 添加原始查询
469
488
 
470
-
471
-
472
- PrettyOutput.print(f"查询: {query}", output_type=OutputType.INFO)
473
-
474
- # 为每个查询获取相似文件
475
- q_vector = self.get_embedding(query)
476
- q_vector = q_vector.reshape(1, -1)
477
-
478
- distances, indices = self.index.search(q_vector, top_k)
479
-
480
- PrettyOutput.print(f"查询 {query} 的结果: ", output_type=OutputType.INFO)
481
-
482
- initial_results = []
483
-
484
- for i, distance in zip(indices[0], distances[0]):
485
- if i == -1: # faiss返回-1表示无效结果
486
- continue
489
+ # 对每个查询变体进行搜索
490
+ all_results = {}
491
+ for q in query_variants:
492
+ q_vector = self.get_embedding(q)
493
+ q_vector = q_vector.reshape(1, -1)
487
494
 
488
- similarity = 1.0 / (1.0 + float(distance))
489
- # 只保留相似度大于等于0.5的结果
490
- if similarity >= 0.5:
491
- PrettyOutput.print(f" {self.file_paths[i]} : 距离 {distance:.3f}, 相似度 {similarity:.3f}",
492
- output_type=OutputType.INFO)
495
+ distances, indices = self.index.search(q_vector, top_k)
493
496
 
494
- file_path = self.file_paths[i]
495
- data = self.vector_cache[file_path]
496
- initial_results.append((file_path, similarity, data["description"]))
497
-
498
- if not initial_results:
499
- PrettyOutput.print("没有找到相似度大于0.5的文件", output_type=OutputType.WARNING)
497
+ for i, distance in zip(indices[0], distances[0]):
498
+ if i == -1:
499
+ continue
500
+
501
+ similarity = 1.0 / (1.0 + float(distance))
502
+ if similarity >= 0.5:
503
+ file_path = self.file_paths[i]
504
+ # 使用最高的相似度分数
505
+ if file_path not in all_results or similarity > all_results[file_path][1]:
506
+ data = self.vector_cache[file_path]
507
+ all_results[file_path] = (file_path, similarity, data["description"])
508
+
509
+ # 转换为列表并排序
510
+ results = list(all_results.values())
511
+ results.sort(key=lambda x: x[1], reverse=True)
512
+
513
+ return results[:top_k]
514
+
515
+ except Exception as e:
516
+ PrettyOutput.print(f"搜索失败: {str(e)}", output_type=OutputType.ERROR)
500
517
  return []
501
518
 
502
- # 使用大模型重新排序
503
- PrettyOutput.print("使用大模型重新排序...", output_type=OutputType.INFO)
504
- reranked_results = self.rerank_results(query, initial_results)
505
-
506
- return reranked_results
507
-
508
519
  def ask_codebase(self, query: str, top_k: int=20) -> str:
509
520
  """查询代码库"""
510
521
  results = self.search_similar(query, top_k)
522
+ if not results:
523
+ PrettyOutput.print("没有找到关联的文件", output_type=OutputType.WARNING)
524
+ return ""
525
+
511
526
  PrettyOutput.print(f"找到的关联文件: ", output_type=OutputType.SUCCESS)
512
527
  for path, score, _ in results:
513
528
  PrettyOutput.print(f"文件: {path} 关联度: {score:.3f}",
@@ -6,7 +6,7 @@ from typing import Dict, Any, List, Tuple
6
6
 
7
7
  import yaml
8
8
  from jarvis.models.base import BasePlatform
9
- from jarvis.utils import OutputType, PrettyOutput, find_git_root, get_multiline_input, load_env_from_file
9
+ from jarvis.utils import OutputType, PrettyOutput, find_git_root, get_max_context_length, get_multiline_input, load_env_from_file
10
10
  from jarvis.models.registry import PlatformRegistry
11
11
  from jarvis.jarvis_codebase.main import CodeBase
12
12
  from prompt_toolkit import PromptSession
@@ -22,7 +22,7 @@ class JarvisCoder:
22
22
  def __init__(self, root_dir: str, language: str):
23
23
  """初始化代码修改工具"""
24
24
 
25
- self.max_context_length = int(os.getenv("JARVIS_MAX_CONTEXT_LENGTH", 65536))
25
+ self.max_context_length = get_max_context_length()
26
26
 
27
27
 
28
28
  self.root_dir = find_git_root(root_dir)
@@ -113,8 +113,8 @@ class JarvisCoder:
113
113
  return []
114
114
 
115
115
  try:
116
- patches = re.findall(r'<PATCH_START>.*?<PATCH_END>', response, re.DOTALL)
117
- return [patch.replace('<PATCH_START>', '').replace('<PATCH_END>', '').strip()
116
+ patches = re.findall(r'<PATCH>.*?</PATCH>', response, re.DOTALL)
117
+ return [patch.replace('<PATCH>', '').replace('</PATCH>', '').strip()
118
118
  for patch in patches if patch.strip()]
119
119
  except Exception as e:
120
120
  PrettyOutput.print(f"解析patch失败: {str(e)}", OutputType.WARNING)
@@ -126,29 +126,57 @@ class JarvisCoder:
126
126
 
127
127
  修改格式说明:
128
128
  1. 每个修改块格式如下:
129
- <PATCH_START>
130
- >>>>>> path/to/file
131
- 要替换的内容
129
+ <PATCH>
130
+ > path/to/file
131
+ def old_function():
132
+ print("old code")
133
+ return False
132
134
  =======
133
- 新的内容
134
- >>>>>>
135
- <PATCH_END>
135
+ def old_function():
136
+ print("new code")
137
+ return True
138
+ </PATCH>
136
139
 
137
140
  2. 如果是新文件或者替换整个文件内容,格式如下:
138
- <PATCH_START>
139
- >>>>>> path/to/new/file
141
+ <PATCH>
142
+ > src/new_module.py
140
143
  =======
141
- 新文件的完整内容
142
- >>>>>>
143
- <PATCH_END>
144
+ from typing import List
145
+
146
+ def new_function():
147
+ return "This is a new file"
148
+ </PATCH>
144
149
 
145
150
  3. 如果要删除文件中的某一段,格式如下:
146
- <PATCH_START>
147
- >>>>>> path/to/file
148
- 要删除的内容
151
+ <PATCH>
152
+ > path/to/file
153
+ # 这是要删除的注释
154
+ deprecated_code = True
155
+ if deprecated_code:
156
+ print("old feature")
149
157
  =======
150
- >>>>>>
151
- <PATCH_END>
158
+ </PATCH>
159
+
160
+ 4. 如果要修改导入语句,格式如下:
161
+ <PATCH>
162
+ > src/main.py
163
+ from old_module import old_class
164
+ =======
165
+ from new_module import new_class
166
+ </PATCH>
167
+
168
+ 5. 如果要修改类定义,格式如下:
169
+ <PATCH>
170
+ > src/models.py
171
+ class OldModel:
172
+ def __init__(self):
173
+ self.value = 0
174
+ =======
175
+ class OldModel:
176
+ def __init__(self):
177
+ self.value = 1
178
+ self.name = "new"
179
+ </PATCH>
152
180
 
153
181
  文件列表如下:
154
182
  """
@@ -158,27 +186,27 @@ class JarvisCoder:
158
186
  continue
159
187
  prompt += f"""{i}. {file["file_path"]}\n"""
160
188
  prompt += f"""文件内容:\n"""
161
- prompt += f"<FILE_CONTENT_START>\n"
189
+ prompt += f"<FILE_CONTENT>\n"
162
190
  prompt += f'{file["file_content"]}\n'
163
- prompt += f"<FILE_CONTENT_END>\n"
191
+ prompt += f"</FILE_CONTENT>\n"
164
192
 
165
193
  prompt += f"\n需求描述: {feature}\n"
166
194
  prompt += """
167
195
  注意事项:
168
- 1、仅输出补丁内容,不要输出任何其他内容,每个补丁必须用<PATCH_START>和<PATCH_END>标记
196
+ 1、仅输出补丁内容,不要输出任何其他内容,每个补丁必须用<PATCH>和</PATCH>标记
169
197
  2、如果在大段代码中有零星修改,生成多个补丁
170
198
  3、要替换的内容,一定要与文件内容完全一致,不要有任何多余或者缺失的内容
171
199
  4、每个patch不超过20行,超出20行,请生成多个patch
172
200
  """
173
-
201
+
174
202
  success, response = self._call_model_with_retry(self.main_model, prompt)
175
203
  if not success:
176
204
  return []
177
205
 
178
206
  try:
179
207
  # 使用正则表达式匹配每个patch块
180
- patches = re.findall(r'<PATCH_START>.*?<PATCH_END>', response, re.DOTALL)
181
- return [patch.replace('<PATCH_START>', '').replace('<PATCH_END>', '').strip()
208
+ patches = re.findall(r'<PATCH>.*?</PATCH>', response, re.DOTALL)
209
+ return [patch.replace('<PATCH>', '').replace('</PATCH>', '').strip()
182
210
  for patch in patches if patch.strip()]
183
211
  except Exception as e:
184
212
  PrettyOutput.print(f"解析patch失败: {str(e)}", OutputType.WARNING)
@@ -204,7 +232,7 @@ class JarvisCoder:
204
232
  continue
205
233
 
206
234
  # 获取文件路径
207
- file_path_match = re.search(r'>>>>>> (.*)', lines[0])
235
+ file_path_match = re.search(r'> (.*)', lines[0])
208
236
  if not file_path_match:
209
237
  error_info.append(f"无法解析文件路径: {lines[0]}")
210
238
  return False, "\n".join(error_info)
@@ -220,7 +248,7 @@ class JarvisCoder:
220
248
  return False, "\n".join(error_info)
221
249
 
222
250
  old_content = parts[0]
223
- new_content = parts[1].split(">>>>>>")[0]
251
+ new_content = parts[1].split("</PATCH>")[0]
224
252
 
225
253
  # 处理新文件
226
254
  if not old_content:
@@ -348,19 +376,59 @@ class JarvisCoder:
348
376
  }
349
377
  else:
350
378
  self._revert_changes()
351
- return {
352
- "success": False,
353
- "stdout": "",
354
- "stderr": "修改被用户取消,文件未发生任何变化",
355
- "error": UserWarning("用户取消修改")
356
- }
379
+
380
+ # 让用户输入调整意见
381
+ user_feedback = get_multiline_input("""
382
+ 请提供修改建议,帮助生成更好的补丁:
383
+ 1. 修改的位置是否正确?
384
+ 2. 修改的内容是否合适?
385
+ 3. 是否有遗漏的修改?
386
+ 4. 其他调整建议?
387
+
388
+ 请输入调整意见(直接回车跳过):""")
389
+
390
+ if not user_feedback:
391
+ return {
392
+ "success": False,
393
+ "stdout": "",
394
+ "stderr": "修改被用户取消,文件未发生任何变化",
395
+ "error": UserWarning("用户取消修改")
396
+ }
397
+
398
+ retry_prompt = f"""补丁被用户拒绝,请根据用户意见重新生成补丁:
399
+
400
+ 用户意见:
401
+ {user_feedback}
402
+
403
+ 请重新生成补丁,确保:
404
+ 1. 按照用户意见调整修改内容
405
+ 2. 准确定位要修改的代码位置
406
+ 3. 正确处理代码缩进
407
+ 4. 考虑代码上下文
408
+ """
409
+ patches = self._remake_patch(retry_prompt)
410
+ continue
357
411
  else:
358
- PrettyOutput.print(f"补丁应用失败,请求重新生成: {error_info}", OutputType.WARNING)
359
- retry_prompt = f"""补丁应用失败,请根据以下错误信息重新生成补丁:
412
+ PrettyOutput.print(f"补丁应用失败: {error_info}", OutputType.WARNING)
413
+
414
+ # 让用户输入补充信息
415
+ user_info = get_multiline_input("""
416
+ 补丁应用失败。请提供更多信息来帮助修复问题:
417
+ 1. 是否需要调整代码位置?
418
+ 2. 是否有特殊的格式要求?
419
+ 3. 是否需要考虑其他文件的依赖?
420
+ 4. 其他补充说明?
421
+
422
+ 请输入补充信息(直接回车跳过):""")
423
+
424
+ retry_prompt = f"""补丁应用失败,请根据以下信息重新生成补丁:
360
425
 
361
426
  错误信息:
362
427
  {error_info}
363
428
 
429
+ 用户补充信息:
430
+ {user_info if user_info else "用户未提供补充信息"}
431
+
364
432
  请确保:
365
433
  1. 准确定位要修改的代码位置
366
434
  2. 正确处理代码缩进