jarvis-ai-assistant 0.1.57__tar.gz → 0.1.58__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.
- {jarvis_ai_assistant-0.1.57/src/jarvis_ai_assistant.egg-info → jarvis_ai_assistant-0.1.58}/PKG-INFO +1 -1
- {jarvis_ai_assistant-0.1.57 → jarvis_ai_assistant-0.1.58}/pyproject.toml +1 -1
- {jarvis_ai_assistant-0.1.57 → jarvis_ai_assistant-0.1.58}/setup.py +1 -1
- {jarvis_ai_assistant-0.1.57 → jarvis_ai_assistant-0.1.58}/src/jarvis/__init__.py +1 -1
- {jarvis_ai_assistant-0.1.57 → jarvis_ai_assistant-0.1.58}/src/jarvis/jarvis_coder/main.py +291 -242
- {jarvis_ai_assistant-0.1.57 → jarvis_ai_assistant-0.1.58}/src/jarvis/models/ai8.py +9 -10
- {jarvis_ai_assistant-0.1.57 → jarvis_ai_assistant-0.1.58}/src/jarvis/models/base.py +9 -0
- {jarvis_ai_assistant-0.1.57 → jarvis_ai_assistant-0.1.58}/src/jarvis/models/kimi.py +16 -8
- {jarvis_ai_assistant-0.1.57 → jarvis_ai_assistant-0.1.58}/src/jarvis/models/openai.py +9 -4
- {jarvis_ai_assistant-0.1.57 → jarvis_ai_assistant-0.1.58}/src/jarvis/models/oyi.py +7 -1
- {jarvis_ai_assistant-0.1.57 → jarvis_ai_assistant-0.1.58/src/jarvis_ai_assistant.egg-info}/PKG-INFO +1 -1
- {jarvis_ai_assistant-0.1.57 → jarvis_ai_assistant-0.1.58}/LICENSE +0 -0
- {jarvis_ai_assistant-0.1.57 → jarvis_ai_assistant-0.1.58}/MANIFEST.in +0 -0
- {jarvis_ai_assistant-0.1.57 → jarvis_ai_assistant-0.1.58}/README.md +0 -0
- {jarvis_ai_assistant-0.1.57 → jarvis_ai_assistant-0.1.58}/setup.cfg +0 -0
- {jarvis_ai_assistant-0.1.57 → jarvis_ai_assistant-0.1.58}/src/jarvis/agent.py +0 -0
- {jarvis_ai_assistant-0.1.57 → jarvis_ai_assistant-0.1.58}/src/jarvis/main.py +0 -0
- {jarvis_ai_assistant-0.1.57 → jarvis_ai_assistant-0.1.58}/src/jarvis/models/__init__.py +0 -0
- {jarvis_ai_assistant-0.1.57 → jarvis_ai_assistant-0.1.58}/src/jarvis/models/registry.py +0 -0
- {jarvis_ai_assistant-0.1.57 → jarvis_ai_assistant-0.1.58}/src/jarvis/tools/__init__.py +0 -0
- {jarvis_ai_assistant-0.1.57 → jarvis_ai_assistant-0.1.58}/src/jarvis/tools/base.py +0 -0
- {jarvis_ai_assistant-0.1.57 → jarvis_ai_assistant-0.1.58}/src/jarvis/tools/bing_search.py +0 -0
- {jarvis_ai_assistant-0.1.57 → jarvis_ai_assistant-0.1.58}/src/jarvis/tools/file_ops.py +0 -0
- {jarvis_ai_assistant-0.1.57 → jarvis_ai_assistant-0.1.58}/src/jarvis/tools/generator.py +0 -0
- {jarvis_ai_assistant-0.1.57 → jarvis_ai_assistant-0.1.58}/src/jarvis/tools/methodology.py +0 -0
- {jarvis_ai_assistant-0.1.57 → jarvis_ai_assistant-0.1.58}/src/jarvis/tools/registry.py +0 -0
- {jarvis_ai_assistant-0.1.57 → jarvis_ai_assistant-0.1.58}/src/jarvis/tools/search.py +0 -0
- {jarvis_ai_assistant-0.1.57 → jarvis_ai_assistant-0.1.58}/src/jarvis/tools/shell.py +0 -0
- {jarvis_ai_assistant-0.1.57 → jarvis_ai_assistant-0.1.58}/src/jarvis/tools/sub_agent.py +0 -0
- {jarvis_ai_assistant-0.1.57 → jarvis_ai_assistant-0.1.58}/src/jarvis/tools/webpage.py +0 -0
- {jarvis_ai_assistant-0.1.57 → jarvis_ai_assistant-0.1.58}/src/jarvis/utils.py +0 -0
- {jarvis_ai_assistant-0.1.57 → jarvis_ai_assistant-0.1.58}/src/jarvis_ai_assistant.egg-info/SOURCES.txt +0 -0
- {jarvis_ai_assistant-0.1.57 → jarvis_ai_assistant-0.1.58}/src/jarvis_ai_assistant.egg-info/dependency_links.txt +0 -0
- {jarvis_ai_assistant-0.1.57 → jarvis_ai_assistant-0.1.58}/src/jarvis_ai_assistant.egg-info/entry_points.txt +0 -0
- {jarvis_ai_assistant-0.1.57 → jarvis_ai_assistant-0.1.58}/src/jarvis_ai_assistant.egg-info/requires.txt +0 -0
- {jarvis_ai_assistant-0.1.57 → jarvis_ai_assistant-0.1.58}/src/jarvis_ai_assistant.egg-info/top_level.txt +0 -0
|
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
|
|
|
4
4
|
|
|
5
5
|
[project]
|
|
6
6
|
name = "jarvis-ai-assistant"
|
|
7
|
-
version = "0.1.
|
|
7
|
+
version = "0.1.58"
|
|
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.
|
|
5
|
+
version="0.1.58",
|
|
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,7 +1,9 @@
|
|
|
1
|
+
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
1
2
|
import hashlib
|
|
2
3
|
import os
|
|
3
4
|
import re
|
|
4
5
|
import sqlite3
|
|
6
|
+
import threading
|
|
5
7
|
import time
|
|
6
8
|
from typing import Dict, Any, List, Optional, Tuple
|
|
7
9
|
|
|
@@ -10,15 +12,16 @@ from jarvis.models.base import BasePlatform
|
|
|
10
12
|
from jarvis.utils import OutputType, PrettyOutput, get_multiline_input, load_env_from_file
|
|
11
13
|
from jarvis.models.registry import PlatformRegistry
|
|
12
14
|
|
|
15
|
+
# 全局锁对象
|
|
16
|
+
index_lock = threading.Lock()
|
|
17
|
+
|
|
13
18
|
class JarvisCoder:
|
|
14
19
|
def __init__(self, root_dir: str, language: str):
|
|
15
20
|
"""初始化代码修改工具"""
|
|
16
|
-
|
|
17
|
-
self.db_path = ""
|
|
21
|
+
|
|
18
22
|
self.root_dir = root_dir
|
|
19
23
|
self.platform = os.environ.get("JARVIS_CODEGEN_PLATFORM")
|
|
20
24
|
self.model = os.environ.get("JARVIS_CODEGEN_MODEL")
|
|
21
|
-
self.language = language
|
|
22
25
|
|
|
23
26
|
self.root_dir = self._find_git_root_dir(self.root_dir)
|
|
24
27
|
if not self.root_dir:
|
|
@@ -117,6 +120,7 @@ class JarvisCoder:
|
|
|
117
120
|
Optional[Dict[str, Any]]: 文件信息,包含文件描述
|
|
118
121
|
"""
|
|
119
122
|
model = self._new_model() # 创建新的模型实例
|
|
123
|
+
model.set_suppress_output(True)
|
|
120
124
|
|
|
121
125
|
prompt = f"""你是一个资深程序员,请根据文件内容,生成文件的关键信息,要求如下,除了代码,不要输出任何内容:
|
|
122
126
|
|
|
@@ -125,20 +129,10 @@ class JarvisCoder:
|
|
|
125
129
|
<CONTENT_START>
|
|
126
130
|
{content}
|
|
127
131
|
<CONTENT_END>
|
|
128
|
-
3. 关键信息:
|
|
129
|
-
<FILE_INFO_START>
|
|
130
|
-
file_description: 这个文件的主要功能和作用描述
|
|
131
|
-
<FILE_INFO_END>
|
|
132
|
+
3. 关键信息: 请生成这个文件的主要功能和作用描述,包含的特征符号(函数和类、变量等),不超过100字
|
|
132
133
|
"""
|
|
133
134
|
try:
|
|
134
|
-
|
|
135
|
-
model.delete_chat() # 删除会话历史
|
|
136
|
-
old_response = response
|
|
137
|
-
response = response.replace("<FILE_INFO_START>", "").replace("<FILE_INFO_END>", "")
|
|
138
|
-
if old_response != response:
|
|
139
|
-
return yaml.safe_load(response)
|
|
140
|
-
else:
|
|
141
|
-
return None
|
|
135
|
+
return model.chat(prompt)
|
|
142
136
|
except Exception as e:
|
|
143
137
|
PrettyOutput.print(f"解析文件信息失败: {str(e)}", OutputType.ERROR)
|
|
144
138
|
return None
|
|
@@ -155,49 +149,57 @@ file_description: 这个文件的主要功能和作用描述
|
|
|
155
149
|
"""获取文件MD5"""
|
|
156
150
|
return hashlib.md5(open(file_path, "rb").read()).hexdigest()
|
|
157
151
|
|
|
152
|
+
|
|
158
153
|
def _create_index_db(self):
|
|
159
154
|
"""创建索引数据库"""
|
|
160
|
-
|
|
161
|
-
|
|
162
|
-
|
|
163
|
-
|
|
164
|
-
|
|
165
|
-
|
|
166
|
-
|
|
167
|
-
|
|
168
|
-
|
|
169
|
-
|
|
170
|
-
|
|
171
|
-
|
|
155
|
+
with index_lock:
|
|
156
|
+
if not os.path.exists(self.index_db_path):
|
|
157
|
+
PrettyOutput.print("Index database does not exist, creating...", OutputType.INFO)
|
|
158
|
+
index_db = sqlite3.connect(self.index_db_path)
|
|
159
|
+
index_db.execute(
|
|
160
|
+
"CREATE TABLE files (file_path TEXT PRIMARY KEY, file_md5 TEXT, file_description TEXT)")
|
|
161
|
+
index_db.commit()
|
|
162
|
+
index_db.close()
|
|
163
|
+
PrettyOutput.print("Index database created", OutputType.SUCCESS)
|
|
164
|
+
# commit
|
|
165
|
+
os.chdir(self.root_dir)
|
|
166
|
+
os.system(f"git add .gitignore -f")
|
|
167
|
+
os.system(f"git commit -m 'add index database'")
|
|
172
168
|
|
|
169
|
+
|
|
173
170
|
def _find_file_by_md5(self, file_md5: str) -> Optional[str]:
|
|
174
171
|
"""根据文件MD5查找文件路径"""
|
|
175
|
-
|
|
176
|
-
|
|
177
|
-
|
|
178
|
-
|
|
179
|
-
|
|
180
|
-
|
|
181
|
-
|
|
172
|
+
with index_lock:
|
|
173
|
+
index_db = sqlite3.connect(self.index_db_path)
|
|
174
|
+
cursor = index_db.cursor()
|
|
175
|
+
cursor.execute(
|
|
176
|
+
"SELECT file_path FROM files WHERE file_md5 = ?", (file_md5,))
|
|
177
|
+
result = cursor.fetchone()
|
|
178
|
+
index_db.close()
|
|
179
|
+
return result[0] if result else None
|
|
182
180
|
|
|
181
|
+
|
|
183
182
|
def _update_file_path(self, file_path: str, file_md5: str):
|
|
184
183
|
"""更新文件路径"""
|
|
185
|
-
|
|
186
|
-
|
|
187
|
-
|
|
188
|
-
|
|
189
|
-
|
|
190
|
-
|
|
184
|
+
with index_lock:
|
|
185
|
+
index_db = sqlite3.connect(self.index_db_path)
|
|
186
|
+
cursor = index_db.cursor()
|
|
187
|
+
cursor.execute(
|
|
188
|
+
"UPDATE files SET file_path = ? WHERE file_md5 = ?", (file_path, file_md5))
|
|
189
|
+
index_db.commit()
|
|
190
|
+
index_db.close()
|
|
191
191
|
|
|
192
|
+
|
|
192
193
|
def _insert_info(self, file_path: str, file_md5: str, file_description: str):
|
|
193
194
|
"""插入文件信息"""
|
|
194
|
-
|
|
195
|
-
|
|
196
|
-
|
|
197
|
-
|
|
198
|
-
|
|
199
|
-
|
|
200
|
-
|
|
195
|
+
with index_lock:
|
|
196
|
+
index_db = sqlite3.connect(self.index_db_path)
|
|
197
|
+
cursor = index_db.cursor()
|
|
198
|
+
cursor.execute("DELETE FROM files WHERE file_path = ?", (file_path,))
|
|
199
|
+
cursor.execute("INSERT INTO files (file_path, file_md5, file_description) VALUES (?, ?, ?)",
|
|
200
|
+
(file_path, file_md5, file_description))
|
|
201
|
+
index_db.commit()
|
|
202
|
+
index_db.close()
|
|
201
203
|
|
|
202
204
|
def _is_text_file(self, file_path: str) -> bool:
|
|
203
205
|
"""判断文件是否是文本文件"""
|
|
@@ -216,6 +218,9 @@ file_description: 这个文件的主要功能和作用描述
|
|
|
216
218
|
|
|
217
219
|
def _index_project(self):
|
|
218
220
|
"""建立代码库索引"""
|
|
221
|
+
import threading
|
|
222
|
+
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
223
|
+
|
|
219
224
|
git_files = os.popen("git ls-files").read().splitlines()
|
|
220
225
|
|
|
221
226
|
index_db = sqlite3.connect(self.index_db_path)
|
|
@@ -229,59 +234,113 @@ file_description: 这个文件的主要功能和作用描述
|
|
|
229
234
|
index_db.commit()
|
|
230
235
|
index_db.close()
|
|
231
236
|
|
|
232
|
-
|
|
233
|
-
|
|
234
|
-
if self._is_text_file(file_path):
|
|
235
|
-
|
|
236
|
-
|
|
237
|
+
def process_file(file_path: str):
|
|
238
|
+
"""处理单个文件的索引任务"""
|
|
239
|
+
if not self._is_text_file(file_path):
|
|
240
|
+
return
|
|
241
|
+
|
|
242
|
+
# 计算文件MD5
|
|
243
|
+
file_md5 = self._get_file_md5(file_path)
|
|
244
|
+
|
|
245
|
+
# 查找文件
|
|
246
|
+
file_path_in_db = self._find_file_by_md5(file_md5)
|
|
247
|
+
if file_path_in_db:
|
|
248
|
+
PrettyOutput.print(
|
|
249
|
+
f"文件 {file_path} 重复,跳过", OutputType.INFO)
|
|
250
|
+
if file_path_in_db != file_path:
|
|
251
|
+
self._update_file_path(file_path, file_md5)
|
|
252
|
+
PrettyOutput.print(
|
|
253
|
+
f"文件 {file_path} 重复,更新路径为 {file_path}", OutputType.INFO)
|
|
254
|
+
return
|
|
237
255
|
|
|
238
|
-
|
|
239
|
-
|
|
240
|
-
|
|
256
|
+
with open(file_path, "r", encoding="utf-8") as f:
|
|
257
|
+
file_content = f.read()
|
|
258
|
+
key_info = self._get_key_info(file_path, file_content)
|
|
259
|
+
if not key_info:
|
|
241
260
|
PrettyOutput.print(
|
|
242
|
-
f"文件 {file_path}
|
|
243
|
-
|
|
244
|
-
|
|
245
|
-
|
|
246
|
-
|
|
247
|
-
|
|
261
|
+
f"文件 {file_path} 索引失败", OutputType.INFO)
|
|
262
|
+
return
|
|
263
|
+
|
|
264
|
+
self._insert_info(file_path, file_md5, key_info)
|
|
265
|
+
PrettyOutput.print(
|
|
266
|
+
f"文件 {file_path} 已建立索引", OutputType.INFO)
|
|
267
|
+
|
|
268
|
+
|
|
269
|
+
# 使用线程池处理文件索引
|
|
270
|
+
with ThreadPoolExecutor(max_workers=10) as executor:
|
|
271
|
+
futures = [executor.submit(process_file, file_path) for file_path in git_files]
|
|
272
|
+
for future in as_completed(futures):
|
|
273
|
+
try:
|
|
274
|
+
future.result()
|
|
275
|
+
except Exception as e:
|
|
276
|
+
PrettyOutput.print(f"处理文件时发生错误: {str(e)}", OutputType.ERROR)
|
|
248
277
|
|
|
249
|
-
with open(file_path, "r", encoding="utf-8") as f:
|
|
250
|
-
file_content = f.read()
|
|
251
|
-
key_info = self._get_key_info(file_path, file_content)
|
|
252
|
-
if not key_info:
|
|
253
|
-
PrettyOutput.print(
|
|
254
|
-
f"文件 {file_path} 索引失败", OutputType.INFO)
|
|
255
|
-
continue
|
|
256
|
-
if "file_description" in key_info:
|
|
257
|
-
self._insert_info(file_path, file_md5, key_info["file_description"])
|
|
258
|
-
PrettyOutput.print(
|
|
259
|
-
f"文件 {file_path} 已建立索引", OutputType.INFO)
|
|
260
|
-
else:
|
|
261
|
-
PrettyOutput.print(
|
|
262
|
-
f"文件 {file_path} 不是代码文件,跳过", OutputType.INFO)
|
|
263
278
|
PrettyOutput.print("项目索引完成", OutputType.INFO)
|
|
264
279
|
|
|
265
|
-
def
|
|
266
|
-
"""
|
|
280
|
+
def _get_files_from_db(self) -> List[Tuple[str, str]]:
|
|
281
|
+
"""从数据库获取所有文件信息
|
|
282
|
+
|
|
283
|
+
Returns:
|
|
284
|
+
List[Tuple[str, str]]: [(file_path, file_description), ...]
|
|
285
|
+
"""
|
|
267
286
|
try:
|
|
268
|
-
# Get all files from database
|
|
269
287
|
index_db = sqlite3.connect(self.index_db_path)
|
|
270
288
|
cursor = index_db.cursor()
|
|
271
289
|
cursor.execute("SELECT file_path, file_description FROM files")
|
|
272
290
|
all_files = cursor.fetchall()
|
|
273
291
|
index_db.close()
|
|
292
|
+
return all_files
|
|
274
293
|
except sqlite3.Error as e:
|
|
275
294
|
PrettyOutput.print(f"数据库操作失败: {str(e)}", OutputType.ERROR)
|
|
276
295
|
return []
|
|
277
296
|
|
|
278
|
-
|
|
279
|
-
|
|
297
|
+
def _analyze_files_in_batches(self, all_files: List[Tuple[str, str]], feature: str, batch_size: int = 100) -> List[Dict]:
|
|
298
|
+
"""批量分析文件相关性
|
|
299
|
+
|
|
300
|
+
Args:
|
|
301
|
+
all_files: 所有文件列表
|
|
302
|
+
feature: 需求描述
|
|
303
|
+
batch_size: 批处理大小
|
|
304
|
+
|
|
305
|
+
Returns:
|
|
306
|
+
List[Dict]: 带评分的文件列表
|
|
307
|
+
"""
|
|
308
|
+
batch_results = []
|
|
309
|
+
|
|
310
|
+
with ThreadPoolExecutor(max_workers=10) as executor:
|
|
311
|
+
futures = []
|
|
312
|
+
for i in range(0, len(all_files), batch_size):
|
|
313
|
+
batch_files = all_files[i:i + batch_size]
|
|
314
|
+
prompt = self._create_batch_analysis_prompt(batch_files, feature)
|
|
315
|
+
model = self._new_model()
|
|
316
|
+
model.set_suppress_output(True)
|
|
317
|
+
futures.append(executor.submit(self._call_model_with_retry, model, prompt))
|
|
318
|
+
|
|
319
|
+
for future in as_completed(futures):
|
|
320
|
+
success, response = future.result()
|
|
321
|
+
if not success:
|
|
322
|
+
continue
|
|
323
|
+
|
|
324
|
+
batch_start = futures.index(future) * batch_size
|
|
325
|
+
batch_end = min(batch_start + batch_size, len(all_files))
|
|
326
|
+
current_batch = all_files[batch_start:batch_end]
|
|
327
|
+
|
|
328
|
+
results = self._process_batch_response(response, current_batch)
|
|
329
|
+
batch_results.extend(results)
|
|
330
|
+
|
|
331
|
+
return batch_results
|
|
332
|
+
|
|
333
|
+
def _create_batch_analysis_prompt(self, batch_files: List[Tuple[str, str]], feature: str) -> str:
|
|
334
|
+
"""创建批量分析的提示词
|
|
280
335
|
|
|
281
|
-
|
|
282
|
-
batch_files
|
|
336
|
+
Args:
|
|
337
|
+
batch_files: 批次文件列表
|
|
338
|
+
feature: 需求描述
|
|
283
339
|
|
|
284
|
-
|
|
340
|
+
Returns:
|
|
341
|
+
str: 提示词
|
|
342
|
+
"""
|
|
343
|
+
prompt = """你是资深程序员,请根据需求描述,从以下文件路径中选出最相关的文件,按相关度从高到低排序。
|
|
285
344
|
|
|
286
345
|
相关度打分标准(0-9分):
|
|
287
346
|
- 9分:文件名直接包含需求中的关键词,且文件功能与需求完全匹配
|
|
@@ -299,102 +358,92 @@ file2.py: 7
|
|
|
299
358
|
|
|
300
359
|
文件列表:
|
|
301
360
|
"""
|
|
302
|
-
|
|
303
|
-
|
|
304
|
-
|
|
305
|
-
|
|
306
|
-
|
|
307
|
-
success, response = self._call_model_with_retry(self._new_model(), prompt)
|
|
308
|
-
if not success:
|
|
309
|
-
continue
|
|
310
|
-
|
|
311
|
-
try:
|
|
312
|
-
response = response.replace("<RELEVANT_FILES_START>", "").replace("<RELEVANT_FILES_END>", "")
|
|
313
|
-
result = yaml.safe_load(response)
|
|
314
|
-
|
|
315
|
-
# Convert results to file objects with scores
|
|
316
|
-
batch_files_dict = {f[0]: f[1] for f in batch_files}
|
|
317
|
-
for file_path, score in result.items():
|
|
318
|
-
if isinstance(file_path, str) and isinstance(score, int):
|
|
319
|
-
score = max(0, min(9, score)) # Ensure score is between 0-9
|
|
320
|
-
if file_path in batch_files_dict:
|
|
321
|
-
batch_results.append({
|
|
322
|
-
"file_path": file_path,
|
|
323
|
-
"file_description": batch_files_dict[file_path],
|
|
324
|
-
"score": score
|
|
325
|
-
})
|
|
326
|
-
|
|
327
|
-
except Exception as e:
|
|
328
|
-
PrettyOutput.print(f"处理批次文件失败: {str(e)}", OutputType.ERROR)
|
|
329
|
-
continue
|
|
330
|
-
|
|
331
|
-
# Sort all results by score
|
|
332
|
-
batch_results.sort(key=lambda x: x["score"], reverse=True)
|
|
333
|
-
top_files = batch_results[:5]
|
|
361
|
+
for file_path, _ in batch_files:
|
|
362
|
+
prompt += f"- {file_path}\n"
|
|
363
|
+
prompt += f"\n需求描述: {feature}\n"
|
|
364
|
+
prompt += "\n注意:\n1. 只输出最相关的文件,不超过5个\n2. 根据上述打分标准判断相关性\n3. 相关度必须是0-9的整数"
|
|
334
365
|
|
|
335
|
-
|
|
336
|
-
|
|
337
|
-
|
|
338
|
-
|
|
339
|
-
"file_path": f[0],
|
|
340
|
-
"file_description": f[1],
|
|
341
|
-
"score": 0
|
|
342
|
-
} for f in remaining_files[:5-len(top_files)]])
|
|
343
|
-
|
|
344
|
-
# Now do content relevance analysis on these files
|
|
345
|
-
score = [[], [], [], [], [], [], [], [], [], []]
|
|
366
|
+
return prompt
|
|
367
|
+
|
|
368
|
+
def _process_batch_response(self, response: str, batch_files: List[Tuple[str, str]]) -> List[Dict]:
|
|
369
|
+
"""处理批量分析的响应
|
|
346
370
|
|
|
347
|
-
|
|
371
|
+
Args:
|
|
372
|
+
response: 模型响应
|
|
373
|
+
batch_files: 批次文件列表
|
|
374
|
+
|
|
375
|
+
Returns:
|
|
376
|
+
List[Dict]: 处理后的文件列表
|
|
377
|
+
"""
|
|
378
|
+
try:
|
|
379
|
+
response = response.replace("<RELEVANT_FILES_START>", "").replace("<RELEVANT_FILES_END>", "")
|
|
380
|
+
result = yaml.safe_load(response)
|
|
381
|
+
|
|
382
|
+
batch_files_dict = {f[0]: f[1] for f in batch_files}
|
|
383
|
+
results = []
|
|
384
|
+
for file_path, score in result.items():
|
|
385
|
+
if isinstance(file_path, str) and isinstance(score, int):
|
|
386
|
+
score = max(0, min(9, score)) # Ensure score is between 0-9
|
|
387
|
+
if file_path in batch_files_dict:
|
|
388
|
+
results.append({
|
|
389
|
+
"file_path": file_path,
|
|
390
|
+
"file_description": batch_files_dict[file_path],
|
|
391
|
+
"score": score
|
|
392
|
+
})
|
|
393
|
+
return results
|
|
394
|
+
except Exception as e:
|
|
395
|
+
PrettyOutput.print(f"处理批次文件失败: {str(e)}", OutputType.ERROR)
|
|
396
|
+
return []
|
|
348
397
|
|
|
349
|
-
相关度打分标准(0-9分):
|
|
350
|
-
- 9分:文件内容与需求完全匹配,是实现需求的核心文件
|
|
351
|
-
- 7-8分:文件内容与需求高度相关,需要较大改动
|
|
352
|
-
- 5-6分:文件内容与需求部分相关,需要中等改动
|
|
353
|
-
- 3-4分:文件内容与需求相关性较低,但需要配合修改
|
|
354
|
-
- 1-2分:文件内容与需求关系较远,只需极少改动
|
|
355
|
-
- 0分:文件内容与需求完全无关
|
|
356
398
|
|
|
357
|
-
|
|
358
|
-
|
|
359
|
-
"""
|
|
360
|
-
for i, file in enumerate(top_files):
|
|
361
|
-
prompt += f"""{i}. {file["file_path"]} : {file["file_description"]}\n"""
|
|
362
|
-
prompt += f"""需求描述: {feature}\n"""
|
|
363
|
-
prompt += "<FILE_LIST_END>\n"
|
|
364
|
-
prompt += """请根据需求描述和文件描述,分析文件的相关性,输出每个编号的相关性[0~9],仅输出以下格式内容(key为文件编号,value为相关性):
|
|
365
|
-
<FILE_RELATION_START>
|
|
366
|
-
"0": 5
|
|
367
|
-
"1": 3
|
|
368
|
-
<FILE_RELATION_END>"""
|
|
369
|
-
|
|
370
|
-
success, response = self._call_model_with_retry(self._new_model(), prompt)
|
|
371
|
-
if not success:
|
|
372
|
-
return top_files[:5] # Return top 5 files from filename matching if model fails
|
|
373
|
-
|
|
399
|
+
def _process_content_response(self, response: str, top_files: List[Dict]) -> List[Dict]:
|
|
400
|
+
"""处理内容分析的响应"""
|
|
374
401
|
try:
|
|
375
402
|
response = response.replace("<FILE_RELATION_START>", "").replace("<FILE_RELATION_END>", "")
|
|
376
403
|
file_relation = yaml.safe_load(response)
|
|
377
404
|
if not file_relation:
|
|
378
405
|
return top_files[:5]
|
|
379
406
|
|
|
407
|
+
score = [[] for _ in range(10)] # 创建10个空列表,对应0-9分
|
|
380
408
|
for file_id, relation in file_relation.items():
|
|
381
409
|
id = int(file_id)
|
|
382
410
|
relation = max(0, min(9, relation)) # 确保范围在0-9之间
|
|
383
411
|
score[relation].append(top_files[id])
|
|
384
412
|
|
|
413
|
+
files = []
|
|
414
|
+
for scores in reversed(score): # 从高分到低分遍历
|
|
415
|
+
files.extend(scores)
|
|
416
|
+
if len(files) >= 5: # 直接取相关性最高的5个文件
|
|
417
|
+
break
|
|
418
|
+
|
|
419
|
+
return files[:5]
|
|
385
420
|
except Exception as e:
|
|
386
421
|
PrettyOutput.print(f"处理文件关系失败: {str(e)}", OutputType.ERROR)
|
|
387
422
|
return top_files[:5]
|
|
423
|
+
|
|
424
|
+
def _find_related_files(self, feature: str) -> List[Dict]:
|
|
425
|
+
"""根据需求描述,查找相关文件
|
|
388
426
|
|
|
389
|
-
|
|
390
|
-
|
|
391
|
-
|
|
392
|
-
|
|
393
|
-
|
|
394
|
-
|
|
427
|
+
Args:
|
|
428
|
+
feature: 需求描述
|
|
429
|
+
|
|
430
|
+
Returns:
|
|
431
|
+
List[Dict]: 相关文件列表
|
|
432
|
+
"""
|
|
433
|
+
# 1. 从数据库获取所有文件
|
|
434
|
+
all_files = self._get_files_from_db()
|
|
435
|
+
if not all_files:
|
|
436
|
+
return []
|
|
395
437
|
|
|
396
|
-
|
|
397
|
-
|
|
438
|
+
# 2. 批量分析文件相关性
|
|
439
|
+
batch_results = self._analyze_files_in_batches(all_files, feature)
|
|
440
|
+
|
|
441
|
+
# 3. 排序并获取前5个文件
|
|
442
|
+
batch_results.sort(key=lambda x: x["score"], reverse=True)
|
|
443
|
+
return batch_results[:5]
|
|
444
|
+
|
|
445
|
+
|
|
446
|
+
|
|
398
447
|
def _remake_patch(self, prompt: str) -> List[str]:
|
|
399
448
|
success, response = self._call_model_with_retry(self.main_model, prompt, max_retries=5) # 增加重试次数
|
|
400
449
|
if not success:
|
|
@@ -417,7 +466,7 @@ file2.py: 7
|
|
|
417
466
|
<PATCH_START>
|
|
418
467
|
>>>>>> path/to/file
|
|
419
468
|
要替换的内容
|
|
420
|
-
|
|
469
|
+
=======
|
|
421
470
|
新的内容
|
|
422
471
|
<<<<<<
|
|
423
472
|
<PATCH_END>
|
|
@@ -425,7 +474,7 @@ file2.py: 7
|
|
|
425
474
|
2. 如果是新文件,格式如下:
|
|
426
475
|
<PATCH_START>
|
|
427
476
|
>>>>>> path/to/new/file
|
|
428
|
-
|
|
477
|
+
=======
|
|
429
478
|
新文件的完整内容
|
|
430
479
|
<<<<<<
|
|
431
480
|
<PATCH_END>
|
|
@@ -488,14 +537,14 @@ file2.py: 7
|
|
|
488
537
|
|
|
489
538
|
# 解析补丁内容
|
|
490
539
|
patch_content = "\n".join(lines[1:])
|
|
491
|
-
parts = patch_content.split("
|
|
540
|
+
parts = patch_content.split("=======")
|
|
492
541
|
|
|
493
542
|
if len(parts) != 2:
|
|
494
543
|
error_info.append(f"补丁格式错误: {file_path}")
|
|
495
544
|
return False, "\n".join(error_info)
|
|
496
545
|
|
|
497
|
-
old_content = parts[0]
|
|
498
|
-
new_content = parts[1].split("<<<<<<")[0]
|
|
546
|
+
old_content = parts[0]
|
|
547
|
+
new_content = parts[1].split("<<<<<<")[0]
|
|
499
548
|
|
|
500
549
|
# 处理新文件
|
|
501
550
|
if not old_content:
|
|
@@ -582,14 +631,85 @@ file2.py: 7
|
|
|
582
631
|
return None
|
|
583
632
|
return root_dir
|
|
584
633
|
|
|
634
|
+
|
|
635
|
+
def _prepare_execution(self) -> None:
|
|
636
|
+
"""准备执行环境"""
|
|
637
|
+
self.main_model = self._new_model()
|
|
638
|
+
self._index_project()
|
|
639
|
+
|
|
640
|
+
def _load_related_files(self, feature: str) -> List[Dict]:
|
|
641
|
+
"""加载相关文件内容"""
|
|
642
|
+
related_files = self._find_related_files(feature)
|
|
643
|
+
for file in related_files:
|
|
644
|
+
PrettyOutput.print(f"Related file: {file['file_path']}", OutputType.INFO)
|
|
645
|
+
with open(file["file_path"], "r", encoding="utf-8") as f:
|
|
646
|
+
file["file_content"] = f.read()
|
|
647
|
+
return related_files
|
|
648
|
+
|
|
649
|
+
def _handle_patch_application(self, related_files: List[Dict], patches: List[str], feature: str) -> Dict[str, Any]:
|
|
650
|
+
"""处理补丁应用流程"""
|
|
651
|
+
while True:
|
|
652
|
+
PrettyOutput.print(f"生成{len(patches)}个补丁", OutputType.INFO)
|
|
653
|
+
|
|
654
|
+
if not patches:
|
|
655
|
+
retry_prompt = f"""未生成补丁,请重新生成补丁"""
|
|
656
|
+
patches = self._remake_patch(retry_prompt)
|
|
657
|
+
continue
|
|
658
|
+
|
|
659
|
+
success, error_info = self._apply_patch(related_files, patches)
|
|
660
|
+
|
|
661
|
+
if success:
|
|
662
|
+
user_confirm = input("是否确认修改?(y/n)")
|
|
663
|
+
if user_confirm.lower() == "y":
|
|
664
|
+
self._finalize_changes(feature, patches)
|
|
665
|
+
return {
|
|
666
|
+
"success": True,
|
|
667
|
+
"stdout": f"已完成功能开发{feature}",
|
|
668
|
+
"stderr": "",
|
|
669
|
+
"error": None
|
|
670
|
+
}
|
|
671
|
+
else:
|
|
672
|
+
self._revert_changes()
|
|
673
|
+
return {
|
|
674
|
+
"success": False,
|
|
675
|
+
"stdout": "",
|
|
676
|
+
"stderr": "修改被用户取消,文件未发生任何变化",
|
|
677
|
+
"error": UserWarning("用户取消修改")
|
|
678
|
+
}
|
|
679
|
+
else:
|
|
680
|
+
PrettyOutput.print(f"补丁应用失败,请求重新生成: {error_info}", OutputType.WARNING)
|
|
681
|
+
retry_prompt = f"""补丁应用失败,请根据以下错误信息重新生成补丁:
|
|
682
|
+
|
|
683
|
+
错误信息:
|
|
684
|
+
{error_info}
|
|
685
|
+
|
|
686
|
+
请确保:
|
|
687
|
+
1. 准确定位要修改的代码位置
|
|
688
|
+
2. 正确处理代码缩进
|
|
689
|
+
3. 考虑代码上下文
|
|
690
|
+
4. 对新文件不要包含原始内容
|
|
691
|
+
"""
|
|
692
|
+
patches = self._remake_patch(retry_prompt)
|
|
693
|
+
|
|
694
|
+
def _finalize_changes(self, feature: str, patches: List[str]) -> None:
|
|
695
|
+
"""完成修改并提交"""
|
|
696
|
+
PrettyOutput.print("修改确认成功,提交修改", OutputType.INFO)
|
|
697
|
+
os.system(f"git add .")
|
|
698
|
+
os.system(f"git commit -m '{feature}'")
|
|
699
|
+
self._save_edit_record(feature, patches)
|
|
700
|
+
self._index_project()
|
|
701
|
+
|
|
702
|
+
def _revert_changes(self) -> None:
|
|
703
|
+
"""回退所有修改"""
|
|
704
|
+
PrettyOutput.print("修改已取消,回退更改", OutputType.INFO)
|
|
705
|
+
os.system(f"git reset --hard")
|
|
706
|
+
os.system(f"git clean -df")
|
|
707
|
+
|
|
585
708
|
def execute(self, feature: str) -> Dict[str, Any]:
|
|
586
709
|
"""执行代码修改
|
|
587
710
|
|
|
588
711
|
Args:
|
|
589
|
-
|
|
590
|
-
- feature: 要实现的功能描述
|
|
591
|
-
- root_dir: 代码库根目录
|
|
592
|
-
- language: 编程语言
|
|
712
|
+
feature: 要实现的功能描述
|
|
593
713
|
|
|
594
714
|
Returns:
|
|
595
715
|
Dict[str, Any]: 包含执行结果的字典
|
|
@@ -599,81 +719,10 @@ file2.py: 7
|
|
|
599
719
|
- error: 错误对象(如果有)
|
|
600
720
|
"""
|
|
601
721
|
try:
|
|
602
|
-
self.
|
|
603
|
-
|
|
604
|
-
|
|
605
|
-
# 4. 开始建立代码库索引
|
|
606
|
-
self._index_project()
|
|
607
|
-
|
|
608
|
-
# 5. 根据索引和需求,查找相关文件
|
|
609
|
-
related_files = self._find_related_files(feature)
|
|
610
|
-
for file in related_files:
|
|
611
|
-
PrettyOutput.print(f"Related file: {file['file_path']}", OutputType.INFO)
|
|
612
|
-
for file in related_files:
|
|
613
|
-
with open(file["file_path"], "r", encoding="utf-8") as f:
|
|
614
|
-
file_content = f.read()
|
|
615
|
-
file["file_content"] = file_content
|
|
722
|
+
self._prepare_execution()
|
|
723
|
+
related_files = self._load_related_files(feature)
|
|
616
724
|
patches = self._make_patch(related_files, feature)
|
|
617
|
-
|
|
618
|
-
# 生成修改方案
|
|
619
|
-
PrettyOutput.print(f"生成{len(patches)}个补丁", OutputType.INFO)
|
|
620
|
-
|
|
621
|
-
if not patches:
|
|
622
|
-
retry_prompt = f"""未生成补丁,请重新生成补丁"""
|
|
623
|
-
patches = self._remake_patch(retry_prompt)
|
|
624
|
-
continue
|
|
625
|
-
|
|
626
|
-
# 尝试应用补丁
|
|
627
|
-
success, error_info = self._apply_patch(related_files, patches)
|
|
628
|
-
|
|
629
|
-
if success:
|
|
630
|
-
# 用户确认修改
|
|
631
|
-
user_confirm = input("是否确认修改?(y/n)")
|
|
632
|
-
if user_confirm.lower() == "y":
|
|
633
|
-
PrettyOutput.print("修改确认成功,提交修改", OutputType.INFO)
|
|
634
|
-
|
|
635
|
-
os.system(f"git add .")
|
|
636
|
-
os.system(f"git commit -m '{feature}'")
|
|
637
|
-
|
|
638
|
-
# 保存修改记录
|
|
639
|
-
self._save_edit_record(feature, patches)
|
|
640
|
-
# 重新建立代码库索引
|
|
641
|
-
self._index_project()
|
|
642
|
-
|
|
643
|
-
return {
|
|
644
|
-
"success": True,
|
|
645
|
-
"stdout": f"已完成功能开发{feature}",
|
|
646
|
-
"stderr": "",
|
|
647
|
-
"error": None
|
|
648
|
-
}
|
|
649
|
-
else:
|
|
650
|
-
PrettyOutput.print("修改已取消,回退更改", OutputType.INFO)
|
|
651
|
-
|
|
652
|
-
os.system(f"git reset --hard") # 回退已修改的文件
|
|
653
|
-
os.system(f"git clean -df") # 删除新创建的文件和目录
|
|
654
|
-
|
|
655
|
-
return {
|
|
656
|
-
"success": False,
|
|
657
|
-
"stdout": "",
|
|
658
|
-
"stderr": "修改被用户取消,文件未发生任何变化",
|
|
659
|
-
"error": UserWarning("用户取消修改")
|
|
660
|
-
}
|
|
661
|
-
else:
|
|
662
|
-
# 补丁应用失败,让模型重新生成
|
|
663
|
-
PrettyOutput.print(f"补丁应用失败,请求重新生成: {error_info}", OutputType.WARNING)
|
|
664
|
-
retry_prompt = f"""补丁应用失败,请根据以下错误信息重新生成补丁:
|
|
665
|
-
|
|
666
|
-
错误信息:
|
|
667
|
-
{error_info}
|
|
668
|
-
|
|
669
|
-
请确保:
|
|
670
|
-
1. 准确定位要修改的代码位置
|
|
671
|
-
2. 正确处理代码缩进
|
|
672
|
-
3. 考虑代码上下文
|
|
673
|
-
4. 对新文件不要包含原始内容
|
|
674
|
-
"""
|
|
675
|
-
patches = self._remake_patch(retry_prompt)
|
|
676
|
-
continue
|
|
725
|
+
return self._handle_patch_application(related_files, patches, feature)
|
|
677
726
|
|
|
678
727
|
except Exception as e:
|
|
679
728
|
return {
|
|
@@ -14,6 +14,7 @@ class AI8Model(BasePlatform):
|
|
|
14
14
|
|
|
15
15
|
def __init__(self):
|
|
16
16
|
"""Initialize model"""
|
|
17
|
+
super().__init__()
|
|
17
18
|
self.system_message = ""
|
|
18
19
|
self.conversation = None
|
|
19
20
|
self.files = []
|
|
@@ -156,15 +157,11 @@ class AI8Model(BasePlatform):
|
|
|
156
157
|
self.system_message = message
|
|
157
158
|
|
|
158
159
|
def chat(self, message: str) -> str:
|
|
159
|
-
"""
|
|
160
|
-
|
|
161
|
-
Args:
|
|
162
|
-
message: User input message
|
|
163
|
-
|
|
164
|
-
Returns:
|
|
165
|
-
str: Model response
|
|
166
|
-
"""
|
|
160
|
+
"""执行对话"""
|
|
167
161
|
try:
|
|
162
|
+
if not self.suppress_output:
|
|
163
|
+
PrettyOutput.print("发送请求...", OutputType.PROGRESS)
|
|
164
|
+
|
|
168
165
|
# 确保有会话ID
|
|
169
166
|
if not self.conversation:
|
|
170
167
|
if not self.create_conversation():
|
|
@@ -219,12 +216,14 @@ class AI8Model(BasePlatform):
|
|
|
219
216
|
chunk = data.get('data', '')
|
|
220
217
|
if chunk:
|
|
221
218
|
full_response += chunk
|
|
222
|
-
|
|
219
|
+
if not self.suppress_output:
|
|
220
|
+
PrettyOutput.print_stream(chunk)
|
|
223
221
|
|
|
224
222
|
except json.JSONDecodeError:
|
|
225
223
|
continue
|
|
226
224
|
|
|
227
|
-
|
|
225
|
+
if not self.suppress_output:
|
|
226
|
+
PrettyOutput.print_stream_end()
|
|
228
227
|
|
|
229
228
|
return full_response
|
|
230
229
|
|
|
@@ -7,6 +7,7 @@ class BasePlatform(ABC):
|
|
|
7
7
|
|
|
8
8
|
def __init__(self):
|
|
9
9
|
"""初始化模型"""
|
|
10
|
+
self.suppress_output = False # 添加输出控制标志
|
|
10
11
|
pass
|
|
11
12
|
|
|
12
13
|
def set_model_name(self, model_name: str):
|
|
@@ -39,3 +40,11 @@ class BasePlatform(ABC):
|
|
|
39
40
|
def set_system_message(self, message: str):
|
|
40
41
|
"""设置系统消息"""
|
|
41
42
|
raise NotImplementedError("set_system_message is not implemented")
|
|
43
|
+
|
|
44
|
+
def set_suppress_output(self, suppress: bool):
|
|
45
|
+
"""设置是否屏蔽输出
|
|
46
|
+
|
|
47
|
+
Args:
|
|
48
|
+
suppress: 是否屏蔽输出
|
|
49
|
+
"""
|
|
50
|
+
self.suppress_output = suppress
|
|
@@ -17,6 +17,7 @@ class KimiModel(BasePlatform):
|
|
|
17
17
|
"""
|
|
18
18
|
初始化Kimi模型
|
|
19
19
|
"""
|
|
20
|
+
super().__init__()
|
|
20
21
|
self.api_key = os.getenv("KIMI_API_KEY")
|
|
21
22
|
if not self.api_key:
|
|
22
23
|
PrettyOutput.print("\n需要设置 KIMI_API_KEY 才能使用 Jarvis。请按以下步骤操作:", OutputType.INFO)
|
|
@@ -217,7 +218,8 @@ class KimiModel(BasePlatform):
|
|
|
217
218
|
def chat(self, message: str) -> str:
|
|
218
219
|
"""发送消息并获取响应"""
|
|
219
220
|
if not self.chat_id:
|
|
220
|
-
|
|
221
|
+
if not self.suppress_output:
|
|
222
|
+
PrettyOutput.print("创建新的对话会话...", OutputType.PROGRESS)
|
|
221
223
|
if not self._create_chat():
|
|
222
224
|
raise Exception("Failed to create chat session")
|
|
223
225
|
|
|
@@ -228,13 +230,15 @@ class KimiModel(BasePlatform):
|
|
|
228
230
|
refs_file = []
|
|
229
231
|
if self.first_chat:
|
|
230
232
|
if self.uploaded_files:
|
|
231
|
-
|
|
233
|
+
if not self.suppress_output:
|
|
234
|
+
PrettyOutput.print(f"首次对话,引用 {len(self.uploaded_files)} 个文件...", OutputType.PROGRESS)
|
|
232
235
|
refs = [f["id"] for f in self.uploaded_files]
|
|
233
236
|
refs_file = self.uploaded_files
|
|
234
237
|
message = self.system_message + "\n" + message
|
|
235
238
|
self.first_chat = False
|
|
236
239
|
|
|
237
|
-
|
|
240
|
+
if not self.suppress_output:
|
|
241
|
+
PrettyOutput.print("发送请求...", OutputType.PROGRESS)
|
|
238
242
|
payload = {
|
|
239
243
|
"messages": [{"role": "user", "content": message}],
|
|
240
244
|
"use_search": True,
|
|
@@ -259,7 +263,9 @@ class KimiModel(BasePlatform):
|
|
|
259
263
|
search_results = []
|
|
260
264
|
ref_sources = []
|
|
261
265
|
|
|
262
|
-
|
|
266
|
+
if not self.suppress_output:
|
|
267
|
+
PrettyOutput.print("接收响应...", OutputType.PROGRESS)
|
|
268
|
+
|
|
263
269
|
for line in response.iter_lines():
|
|
264
270
|
if not line:
|
|
265
271
|
continue
|
|
@@ -276,7 +282,8 @@ class KimiModel(BasePlatform):
|
|
|
276
282
|
# 处理补全文本
|
|
277
283
|
text = data.get("text", "")
|
|
278
284
|
if text:
|
|
279
|
-
|
|
285
|
+
if not self.suppress_output:
|
|
286
|
+
PrettyOutput.print_stream(text)
|
|
280
287
|
full_response += text
|
|
281
288
|
|
|
282
289
|
elif event == "search_plus":
|
|
@@ -311,11 +318,12 @@ class KimiModel(BasePlatform):
|
|
|
311
318
|
except json.JSONDecodeError:
|
|
312
319
|
continue
|
|
313
320
|
|
|
314
|
-
|
|
321
|
+
if not self.suppress_output:
|
|
322
|
+
PrettyOutput.print_stream_end()
|
|
315
323
|
|
|
316
324
|
|
|
317
325
|
# 显示搜索结果摘要
|
|
318
|
-
if search_results:
|
|
326
|
+
if search_results and not self.suppress_output:
|
|
319
327
|
PrettyOutput.print("\n搜索结果:", OutputType.PROGRESS)
|
|
320
328
|
for result in search_results:
|
|
321
329
|
PrettyOutput.print(f"- {result['title']}", OutputType.PROGRESS)
|
|
@@ -328,7 +336,7 @@ class KimiModel(BasePlatform):
|
|
|
328
336
|
PrettyOutput.print("", OutputType.PROGRESS)
|
|
329
337
|
|
|
330
338
|
# 显示引用来源
|
|
331
|
-
if ref_sources:
|
|
339
|
+
if ref_sources and not self.suppress_output:
|
|
332
340
|
PrettyOutput.print("\n引用来源:", OutputType.PROGRESS)
|
|
333
341
|
for source in ref_sources:
|
|
334
342
|
PrettyOutput.print(f"- [{source['ref_id']}] {source['title']} ({source['source']})", OutputType.PROGRESS)
|
|
@@ -13,6 +13,7 @@ class OpenAIModel(BasePlatform):
|
|
|
13
13
|
"""
|
|
14
14
|
初始化DeepSeek模型
|
|
15
15
|
"""
|
|
16
|
+
super().__init__()
|
|
16
17
|
self.api_key = os.getenv("OPENAI_API_KEY")
|
|
17
18
|
if not self.api_key:
|
|
18
19
|
PrettyOutput.print("\n需要设置以下环境变量才能使用 OpenAI 模型:", OutputType.INFO)
|
|
@@ -54,7 +55,8 @@ class OpenAIModel(BasePlatform):
|
|
|
54
55
|
def chat(self, message: str) -> str:
|
|
55
56
|
"""执行对话"""
|
|
56
57
|
try:
|
|
57
|
-
|
|
58
|
+
if not self.suppress_output:
|
|
59
|
+
PrettyOutput.print("发送请求...", OutputType.PROGRESS)
|
|
58
60
|
|
|
59
61
|
# 添加用户消息到历史记录
|
|
60
62
|
self.messages.append({"role": "user", "content": message})
|
|
@@ -65,16 +67,19 @@ class OpenAIModel(BasePlatform):
|
|
|
65
67
|
stream=True
|
|
66
68
|
)
|
|
67
69
|
|
|
68
|
-
|
|
70
|
+
if not self.suppress_output:
|
|
71
|
+
PrettyOutput.print("接收响应...", OutputType.PROGRESS)
|
|
69
72
|
full_response = ""
|
|
70
73
|
|
|
71
74
|
for chunk in response:
|
|
72
75
|
if chunk.choices[0].delta.content:
|
|
73
76
|
text = chunk.choices[0].delta.content
|
|
74
|
-
|
|
77
|
+
if not self.suppress_output:
|
|
78
|
+
PrettyOutput.print_stream(text)
|
|
75
79
|
full_response += text
|
|
76
80
|
|
|
77
|
-
|
|
81
|
+
if not self.suppress_output:
|
|
82
|
+
PrettyOutput.print_stream_end()
|
|
78
83
|
|
|
79
84
|
# 添加助手回复到历史记录
|
|
80
85
|
self.messages.append({"role": "assistant", "content": full_response})
|
|
@@ -14,6 +14,7 @@ class OyiModel(BasePlatform):
|
|
|
14
14
|
|
|
15
15
|
def __init__(self):
|
|
16
16
|
"""Initialize model"""
|
|
17
|
+
super().__init__()
|
|
17
18
|
PrettyOutput.section("支持的模型", OutputType.SUCCESS)
|
|
18
19
|
|
|
19
20
|
# 获取可用模型列表
|
|
@@ -114,6 +115,9 @@ class OyiModel(BasePlatform):
|
|
|
114
115
|
str: Model response
|
|
115
116
|
"""
|
|
116
117
|
try:
|
|
118
|
+
if not self.suppress_output:
|
|
119
|
+
PrettyOutput.print("发送请求...", OutputType.PROGRESS)
|
|
120
|
+
|
|
117
121
|
# 确保有会话ID
|
|
118
122
|
if not self.conversation:
|
|
119
123
|
if not self.create_conversation():
|
|
@@ -182,7 +186,9 @@ class OyiModel(BasePlatform):
|
|
|
182
186
|
)
|
|
183
187
|
|
|
184
188
|
if response.status_code == 200:
|
|
185
|
-
|
|
189
|
+
if not self.suppress_output:
|
|
190
|
+
PrettyOutput.print("接收响应...", OutputType.PROGRESS)
|
|
191
|
+
PrettyOutput.print(response.text, OutputType.SYSTEM)
|
|
186
192
|
self.messages.append({"role": "assistant", "content": response.text})
|
|
187
193
|
return response.text
|
|
188
194
|
else:
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|