jarvis-ai-assistant 0.1.76__tar.gz → 0.1.78__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.
Files changed (42) hide show
  1. {jarvis_ai_assistant-0.1.76/src/jarvis_ai_assistant.egg-info → jarvis_ai_assistant-0.1.78}/PKG-INFO +33 -16
  2. {jarvis_ai_assistant-0.1.76 → jarvis_ai_assistant-0.1.78}/README.md +27 -15
  3. {jarvis_ai_assistant-0.1.76 → jarvis_ai_assistant-0.1.78}/pyproject.toml +7 -1
  4. {jarvis_ai_assistant-0.1.76 → jarvis_ai_assistant-0.1.78}/setup.py +7 -2
  5. {jarvis_ai_assistant-0.1.76 → jarvis_ai_assistant-0.1.78}/src/jarvis/__init__.py +1 -1
  6. {jarvis_ai_assistant-0.1.76 → jarvis_ai_assistant-0.1.78}/src/jarvis/agent.py +9 -41
  7. {jarvis_ai_assistant-0.1.76 → jarvis_ai_assistant-0.1.78}/src/jarvis/jarvis_codebase/main.py +7 -29
  8. {jarvis_ai_assistant-0.1.76 → jarvis_ai_assistant-0.1.78}/src/jarvis/jarvis_coder/main.py +0 -11
  9. {jarvis_ai_assistant-0.1.76 → jarvis_ai_assistant-0.1.78}/src/jarvis/main.py +1 -16
  10. {jarvis_ai_assistant-0.1.76 → jarvis_ai_assistant-0.1.78}/src/jarvis/models/registry.py +19 -14
  11. jarvis_ai_assistant-0.1.78/src/jarvis/rag/__init__.py +0 -0
  12. jarvis_ai_assistant-0.1.78/src/jarvis/rag/main.py +566 -0
  13. {jarvis_ai_assistant-0.1.76 → jarvis_ai_assistant-0.1.78}/src/jarvis/tools/coder.py +4 -4
  14. {jarvis_ai_assistant-0.1.76 → jarvis_ai_assistant-0.1.78}/src/jarvis/tools/generator.py +1 -5
  15. {jarvis_ai_assistant-0.1.76 → jarvis_ai_assistant-0.1.78}/src/jarvis/tools/search.py +88 -21
  16. {jarvis_ai_assistant-0.1.76 → jarvis_ai_assistant-0.1.78}/src/jarvis/utils.py +27 -1
  17. {jarvis_ai_assistant-0.1.76 → jarvis_ai_assistant-0.1.78/src/jarvis_ai_assistant.egg-info}/PKG-INFO +33 -16
  18. {jarvis_ai_assistant-0.1.76 → jarvis_ai_assistant-0.1.78}/src/jarvis_ai_assistant.egg-info/SOURCES.txt +2 -0
  19. {jarvis_ai_assistant-0.1.76 → jarvis_ai_assistant-0.1.78}/src/jarvis_ai_assistant.egg-info/entry_points.txt +1 -0
  20. {jarvis_ai_assistant-0.1.76 → jarvis_ai_assistant-0.1.78}/src/jarvis_ai_assistant.egg-info/requires.txt +4 -0
  21. {jarvis_ai_assistant-0.1.76 → jarvis_ai_assistant-0.1.78}/LICENSE +0 -0
  22. {jarvis_ai_assistant-0.1.76 → jarvis_ai_assistant-0.1.78}/MANIFEST.in +0 -0
  23. {jarvis_ai_assistant-0.1.76 → jarvis_ai_assistant-0.1.78}/setup.cfg +0 -0
  24. {jarvis_ai_assistant-0.1.76 → jarvis_ai_assistant-0.1.78}/src/jarvis/jarvis_codebase/__init__.py +0 -0
  25. {jarvis_ai_assistant-0.1.76 → jarvis_ai_assistant-0.1.78}/src/jarvis/jarvis_coder/__init__.py +0 -0
  26. {jarvis_ai_assistant-0.1.76 → jarvis_ai_assistant-0.1.78}/src/jarvis/models/__init__.py +0 -0
  27. {jarvis_ai_assistant-0.1.76 → jarvis_ai_assistant-0.1.78}/src/jarvis/models/ai8.py +0 -0
  28. {jarvis_ai_assistant-0.1.76 → jarvis_ai_assistant-0.1.78}/src/jarvis/models/base.py +0 -0
  29. {jarvis_ai_assistant-0.1.76 → jarvis_ai_assistant-0.1.78}/src/jarvis/models/kimi.py +0 -0
  30. {jarvis_ai_assistant-0.1.76 → jarvis_ai_assistant-0.1.78}/src/jarvis/models/openai.py +0 -0
  31. {jarvis_ai_assistant-0.1.76 → jarvis_ai_assistant-0.1.78}/src/jarvis/models/oyi.py +0 -0
  32. {jarvis_ai_assistant-0.1.76 → jarvis_ai_assistant-0.1.78}/src/jarvis/tools/__init__.py +0 -0
  33. {jarvis_ai_assistant-0.1.76 → jarvis_ai_assistant-0.1.78}/src/jarvis/tools/base.py +0 -0
  34. {jarvis_ai_assistant-0.1.76 → jarvis_ai_assistant-0.1.78}/src/jarvis/tools/codebase_qa.py +0 -0
  35. {jarvis_ai_assistant-0.1.76 → jarvis_ai_assistant-0.1.78}/src/jarvis/tools/file_ops.py +0 -0
  36. {jarvis_ai_assistant-0.1.76 → jarvis_ai_assistant-0.1.78}/src/jarvis/tools/methodology.py +0 -0
  37. {jarvis_ai_assistant-0.1.76 → jarvis_ai_assistant-0.1.78}/src/jarvis/tools/registry.py +0 -0
  38. {jarvis_ai_assistant-0.1.76 → jarvis_ai_assistant-0.1.78}/src/jarvis/tools/shell.py +0 -0
  39. {jarvis_ai_assistant-0.1.76 → jarvis_ai_assistant-0.1.78}/src/jarvis/tools/sub_agent.py +0 -0
  40. {jarvis_ai_assistant-0.1.76 → jarvis_ai_assistant-0.1.78}/src/jarvis/tools/webpage.py +0 -0
  41. {jarvis_ai_assistant-0.1.76 → jarvis_ai_assistant-0.1.78}/src/jarvis_ai_assistant.egg-info/dependency_links.txt +0 -0
  42. {jarvis_ai_assistant-0.1.76 → jarvis_ai_assistant-0.1.78}/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.76
3
+ Version: 0.1.78
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
@@ -35,6 +35,7 @@ Classifier: Programming Language :: Python :: 3.8
35
35
  Classifier: Programming Language :: Python :: 3.9
36
36
  Classifier: Programming Language :: Python :: 3.10
37
37
  Classifier: Programming Language :: Python :: 3.11
38
+ Classifier: Operating System :: POSIX :: Linux
38
39
  Requires-Python: >=3.8
39
40
  Description-Content-Type: text/markdown
40
41
  License-File: LICENSE
@@ -48,6 +49,10 @@ Requires-Dist: numpy>=1.24.0
48
49
  Requires-Dist: faiss-cpu>=1.8.0
49
50
  Requires-Dist: sentence-transformers>=2.2.2
50
51
  Requires-Dist: bs4>=0.0.1
52
+ Requires-Dist: PyMuPDF>=1.21.0
53
+ Requires-Dist: python-docx>=0.8.11
54
+ Requires-Dist: tiktoken>=0.3.0
55
+ Requires-Dist: tqdm>=4.65.0
51
56
  Provides-Extra: dev
52
57
  Requires-Dist: pytest; extra == "dev"
53
58
  Requires-Dist: black; extra == "dev"
@@ -124,6 +129,7 @@ Jarvis supports configuration through environment variables that can be set in t
124
129
  |---------|------|--------|------|
125
130
  | JARVIS_PLATFORM | AI platform to use, supports kimi/openai/ai8 etc | kimi | Yes |
126
131
  | JARVIS_MODEL | Model name to use | - | No |
132
+ | JARVIS_THREAD_COUNT | Number of threads for parallel processing | 10 | No |
127
133
  | JARVIS_CODEGEN_PLATFORM | AI platform for code generation | Same as JARVIS_PLATFORM | No |
128
134
  | JARVIS_CODEGEN_MODEL | Model name for code generation | Same as JARVIS_MODEL | No |
129
135
  | JARVIS_CHEAP_PLATFORM | AI platform for cheap operations | Same as JARVIS_PLATFORM | No |
@@ -141,36 +147,47 @@ Jarvis supports configuration through environment variables that can be set in t
141
147
 
142
148
  ## 🎯 Usage
143
149
 
144
- ### Basic Usage
150
+ ### Main Assistant
145
151
  ```bash
146
152
  jarvis
147
153
  ```
148
154
 
149
-
150
- ### With Specific Model
155
+ ### Code Generation
151
156
  ```bash
152
- jarvis -p kimi # Use Kimi platform
153
- jarvis -p openai # Use OpenAI platform
154
- ```
155
-
156
- ### Code Modification
157
- ```bash
158
- jarvis-coder --feature "Add new feature" # Modify code to add new feature
157
+ jarvis-coder
159
158
  ```
160
159
 
161
160
  ### Codebase Search
162
161
  ```bash
163
- jarvis-codebase --search "database connection" # Search codebase
162
+ # Generate codebase index
163
+ jarvis-codebase --generate
164
+
165
+ # Search similar code
166
+ jarvis-codebase --search "your search query"
167
+
168
+ # Ask questions about codebase
169
+ jarvis-codebase --ask "your question"
164
170
  ```
165
171
 
166
- ### Codebase Question
172
+ ### Document Analysis (RAG)
167
173
  ```bash
168
- jarvis-codebase --ask "How to use the database?" # Ask about codebase
174
+ # Build document index
175
+ jarvis-rag --dir /path/to/documents --build
176
+
177
+ # Search documents
178
+ jarvis-rag --query "your search query"
169
179
  ```
170
180
 
171
- ### Keep Chat History
181
+ ### Search Tool
172
182
  ```bash
173
- jarvis --keep-history # Don't delete chat session after completion
183
+ # Basic search
184
+ jarvis-search "your query"
185
+
186
+ # Show only URLs
187
+ jarvis-search "your query" --url-only
188
+
189
+ # Limit results
190
+ jarvis-search "your query" --max 3
174
191
  ```
175
192
 
176
193
  ## 🛠️ Tools
@@ -65,6 +65,7 @@ Jarvis supports configuration through environment variables that can be set in t
65
65
  |---------|------|--------|------|
66
66
  | JARVIS_PLATFORM | AI platform to use, supports kimi/openai/ai8 etc | kimi | Yes |
67
67
  | JARVIS_MODEL | Model name to use | - | No |
68
+ | JARVIS_THREAD_COUNT | Number of threads for parallel processing | 10 | No |
68
69
  | JARVIS_CODEGEN_PLATFORM | AI platform for code generation | Same as JARVIS_PLATFORM | No |
69
70
  | JARVIS_CODEGEN_MODEL | Model name for code generation | Same as JARVIS_MODEL | No |
70
71
  | JARVIS_CHEAP_PLATFORM | AI platform for cheap operations | Same as JARVIS_PLATFORM | No |
@@ -82,36 +83,47 @@ Jarvis supports configuration through environment variables that can be set in t
82
83
 
83
84
  ## 🎯 Usage
84
85
 
85
- ### Basic Usage
86
+ ### Main Assistant
86
87
  ```bash
87
88
  jarvis
88
89
  ```
89
90
 
90
-
91
- ### With Specific Model
91
+ ### Code Generation
92
92
  ```bash
93
- jarvis -p kimi # Use Kimi platform
94
- jarvis -p openai # Use OpenAI platform
95
- ```
96
-
97
- ### Code Modification
98
- ```bash
99
- jarvis-coder --feature "Add new feature" # Modify code to add new feature
93
+ jarvis-coder
100
94
  ```
101
95
 
102
96
  ### Codebase Search
103
97
  ```bash
104
- jarvis-codebase --search "database connection" # Search codebase
98
+ # Generate codebase index
99
+ jarvis-codebase --generate
100
+
101
+ # Search similar code
102
+ jarvis-codebase --search "your search query"
103
+
104
+ # Ask questions about codebase
105
+ jarvis-codebase --ask "your question"
105
106
  ```
106
107
 
107
- ### Codebase Question
108
+ ### Document Analysis (RAG)
108
109
  ```bash
109
- jarvis-codebase --ask "How to use the database?" # Ask about codebase
110
+ # Build document index
111
+ jarvis-rag --dir /path/to/documents --build
112
+
113
+ # Search documents
114
+ jarvis-rag --query "your search query"
110
115
  ```
111
116
 
112
- ### Keep Chat History
117
+ ### Search Tool
113
118
  ```bash
114
- jarvis --keep-history # Don't delete chat session after completion
119
+ # Basic search
120
+ jarvis-search "your query"
121
+
122
+ # Show only URLs
123
+ jarvis-search "your query" --url-only
124
+
125
+ # Limit results
126
+ jarvis-search "your query" --max 3
115
127
  ```
116
128
 
117
129
  ## 🛠️ Tools
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
4
4
 
5
5
  [project]
6
6
  name = "jarvis-ai-assistant"
7
- version = "0.1.76"
7
+ version = "0.1.78"
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" }]
@@ -17,6 +17,7 @@ classifiers = [
17
17
  "Programming Language :: Python :: 3.9",
18
18
  "Programming Language :: Python :: 3.10",
19
19
  "Programming Language :: Python :: 3.11",
20
+ "Operating System :: POSIX :: Linux",
20
21
  ]
21
22
  keywords = ["jarvis", "ai", "assistant", "tools", "automation"]
22
23
  dependencies = [
@@ -30,6 +31,10 @@ dependencies = [
30
31
  "faiss-cpu>=1.8.0",
31
32
  "sentence-transformers>=2.2.2",
32
33
  "bs4>=0.0.1",
34
+ "PyMuPDF>=1.21.0",
35
+ "python-docx>=0.8.11",
36
+ "tiktoken>=0.3.0",
37
+ "tqdm>=4.65.0",
33
38
  ]
34
39
  requires-python = ">=3.8"
35
40
 
@@ -43,3 +48,4 @@ Homepage = "https://github.com/skyfireitdiy/Jarvis"
43
48
  jarvis = "jarvis.main:main"
44
49
  jarvis-coder = "jarvis.jarvis_coder.main:main"
45
50
  jarvis-codebase = "jarvis.jarvis_codebase.main:main"
51
+ jarvis-rag = "jarvis.rag.main:main"
@@ -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.76",
5
+ version="0.1.78",
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",
@@ -23,12 +23,17 @@ setup(
23
23
  "faiss-cpu>=1.8.0",
24
24
  "sentence-transformers>=2.2.2",
25
25
  "bs4>=0.0.1",
26
+ "PyMuPDF>=1.21.0",
27
+ "python-docx>=0.8.11",
28
+ "tiktoken>=0.3.0",
29
+ "tqdm>=4.65.0",
26
30
  ],
27
31
  entry_points={
28
32
  "console_scripts": [
29
33
  "jarvis=jarvis.main:main",
30
34
  "jarvis-coder=jarvis.jarvis_coder.main:main",
31
35
  "jarvis-codebase=jarvis.jarvis_codebase.main:main",
36
+ "jarvis-rag=jarvis.rag.main:main",
32
37
  ],
33
38
  },
34
39
  python_requires=">=3.8",
@@ -36,7 +41,7 @@ setup(
36
41
  "Development Status :: 3 - Alpha",
37
42
  "Intended Audience :: Developers",
38
43
  "License :: OSI Approved :: MIT License",
39
- "Operating System :: OS Independent",
44
+ "Operating System :: POSIX :: Linux",
40
45
  "Programming Language :: Python :: 3",
41
46
  "Programming Language :: Python :: 3.8",
42
47
  "Programming Language :: Python :: 3.9",
@@ -1,3 +1,3 @@
1
1
  """Jarvis AI Assistant"""
2
2
 
3
- __version__ = "0.1.76"
3
+ __version__ = "0.1.78"
@@ -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, while_success
11
+ from .utils import PrettyOutput, OutputType, 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
@@ -24,7 +24,7 @@ class Agent:
24
24
  name: Agent名称,默认为"Jarvis"
25
25
  is_sub_agent: 是否为子Agent,默认为False
26
26
  """
27
- self.model = PlatformRegistry.get_global_platform()
27
+ self.model = PlatformRegistry.get_global_platform_registry().get_normal_platform()
28
28
  self.tool_registry = ToolRegistry.get_global_tool_registry()
29
29
  self.name = name
30
30
  self.is_sub_agent = is_sub_agent
@@ -32,20 +32,19 @@ class Agent:
32
32
  self.conversation_turns = 0
33
33
 
34
34
  # 从环境变量加载嵌入模型配置
35
- self.embedding_model_name = os.environ.get("JARVIS_EMBEDDING_MODEL", "BAAI/bge-large-zh-v1.5")
36
35
  self.embedding_dimension = 1536 # Default for many embedding models
37
36
 
38
37
  # 初始化嵌入模型
39
38
  try:
40
- os.environ["TOKENIZERS_PARALLELISM"] = "false"
41
- PrettyOutput.print(f"正在加载嵌入模型: {self.embedding_model_name}...", OutputType.INFO)
42
- self.embedding_model = SentenceTransformer(self.embedding_model_name)
39
+ self.embedding_model = load_embedding_model()
43
40
 
44
41
  # 预热模型并获取正确的维度
45
42
  test_text = "这是一段测试文本,用于确保模型完全加载。"
46
- test_embedding = self.embedding_model.encode(test_text,
47
- convert_to_tensor=True,
48
- normalize_embeddings=True)
43
+ test_embedding = self.embedding_model.encode(
44
+ test_text,
45
+ convert_to_tensor=True,
46
+ normalize_embeddings=True
47
+ )
49
48
  self.embedding_dimension = len(test_embedding)
50
49
  PrettyOutput.print("嵌入模型加载完成", OutputType.SUCCESS)
51
50
 
@@ -291,37 +290,6 @@ class Agent:
291
290
  self.prompt = summary_prompt
292
291
  return self._call_model(self.prompt)
293
292
 
294
- def choose_tools(self, user_input: str) -> List[Dict]:
295
- """根据用户输入选择工具"""
296
- PrettyOutput.print("选择工具...", OutputType.PLANNING)
297
- tools = self.tool_registry.get_all_tools()
298
- prompt = f"""你是一个工具选择专家,请根据用户输入选择合适的工具,返回可能使用到的工具的名称。以下是可用工具:
299
- """
300
- for tool in tools:
301
- prompt += f"- {tool['name']}: {tool['description']}\n"
302
- prompt += f"用户输入: {user_input}\n"
303
- prompt += f"请返回可能使用到的工具的名称,如果无法确定,请返回空列表。"
304
- prompt += f"返回的格式为:\n"
305
- prompt += f"<TOOL_CHOICE_START>\n"
306
- prompt += f"tool_name1\n"
307
- prompt += f"tool_name2\n"
308
- prompt += f"<TOOL_CHOICE_END>\n"
309
- model = PlatformRegistry.get_global_platform()
310
- model.set_suppress_output(True)
311
- try:
312
- response = model.chat(prompt)
313
- response = response.replace("<TOOL_CHOICE_START>", "").replace("<TOOL_CHOICE_END>", "")
314
- tools_name = response.split("\n")
315
- choosed_tools = []
316
- for tool_name in tools_name:
317
- for tool in tools:
318
- if tool['name'] == tool_name:
319
- choosed_tools.append(tool)
320
- break
321
- return choosed_tools
322
- except Exception as e:
323
- PrettyOutput.print(f"工具选择失败: {str(e)}", OutputType.ERROR)
324
- return []
325
293
 
326
294
  def run(self, user_input: str, file_list: Optional[List[str]] = None, keep_history: bool = False) -> str:
327
295
  """处理用户输入并返回响应,返回任务总结报告
@@ -350,7 +318,7 @@ class Agent:
350
318
  tools_prompt = ""
351
319
 
352
320
  # 选择工具
353
- tools = self.choose_tools(user_input)
321
+ tools = self.tool_registry.get_all_tools()
354
322
  if tools:
355
323
  tools_prompt += "可用工具:\n"
356
324
  for tool in tools:
@@ -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
10
+ from jarvis.utils import OutputType, PrettyOutput, find_git_root, load_embedding_model
11
11
  from jarvis.utils import load_env_from_file
12
12
  import argparse
13
13
  from sentence_transformers import SentenceTransformer
@@ -19,22 +19,6 @@ class CodeBase:
19
19
  self.root_dir = root_dir
20
20
  os.chdir(self.root_dir)
21
21
  self.thread_count = int(os.environ.get("JARVIS_THREAD_COUNT") or 10)
22
- self.cheap_platform = os.environ.get("JARVIS_CHEAP_PLATFORM") or os.environ.get("JARVIS_PLATFORM") or "kimi"
23
- self.cheap_model = os.environ.get("JARVIS_CHEAP_MODEL") or os.environ.get("JARVIS_MODEL") or "kimi"
24
- self.normal_platform = os.environ.get("JARVIS_PLATFORM") or "kimi"
25
- self.codegen_platform = os.environ.get("JARVIS_CODEGEN_PLATFORM") or os.environ.get("JARVIS_PLATFORM") or "kimi"
26
- self.codegen_model = os.environ.get("JARVIS_CODEGEN_MODEL") or os.environ.get("JARVIS_MODEL") or "kimi"
27
- self.normal_model = os.environ.get("JARVIS_MODEL") or "kimi"
28
- self.embedding_model_name = os.environ.get("JARVIS_EMBEDDING_MODEL") or "BAAI/bge-large-zh-v1.5"
29
- if not self.cheap_platform or not self.cheap_model or not self.codegen_platform or not self.codegen_model or not self.embedding_model_name or not self.normal_platform or not self.normal_model:
30
- raise ValueError("JARVIS_CHEAP_PLATFORM or JARVIS_CHEAP_MODEL or JARVIS_CODEGEN_PLATFORM or JARVIS_CODEGEN_MODEL or JARVIS_EMBEDDING_MODEL or JARVIS_PLATFORM or JARVIS_MODEL is not set")
31
-
32
- PrettyOutput.print(f"廉价模型使用平台: {self.cheap_platform} 模型: {self.cheap_model}", output_type=OutputType.INFO)
33
- PrettyOutput.print(f"代码生成模型使用平台: {self.codegen_platform} 模型: {self.codegen_model}", output_type=OutputType.INFO)
34
- PrettyOutput.print(f"分析模型使用平台: {self.normal_platform} 模型: {self.normal_model}", output_type=OutputType.INFO)
35
- PrettyOutput.print(f"嵌入模型: {self.embedding_model_name}", output_type=OutputType.INFO)
36
- PrettyOutput.print(f"索引建立线程数: {self.thread_count}", output_type=OutputType.INFO)
37
- PrettyOutput.print(f"检索算法:分层导航小世界算法", output_type=OutputType.INFO)
38
22
 
39
23
  # 初始化数据目录
40
24
  self.data_dir = os.path.join(self.root_dir, ".jarvis-codebase")
@@ -43,9 +27,7 @@ class CodeBase:
43
27
 
44
28
  # 初始化嵌入模型,使用系统默认缓存目录
45
29
  try:
46
- os.environ["TOKENIZERS_PARALLELISM"] = "false"
47
- PrettyOutput.print("正在加载/下载模型,请稍候...", output_type=OutputType.INFO)
48
- self.embedding_model = SentenceTransformer(self.embedding_model_name)
30
+ self.embedding_model = load_embedding_model()
49
31
 
50
32
  # 强制完全加载所有模型组件
51
33
  test_text = """
@@ -64,7 +46,7 @@ class CodeBase:
64
46
  self.vector_dim = self.embedding_model.get_sentence_embedding_dimension()
65
47
 
66
48
  self.git_file_list = self.get_git_file_list()
67
- self.platform_registry = PlatformRegistry().get_global_platform_registry()
49
+ self.platform_registry = PlatformRegistry.get_global_platform_registry()
68
50
 
69
51
  # 初始化缓存和索引
70
52
  self.cache_path = os.path.join(self.data_dir, "cache.pkl")
@@ -104,8 +86,7 @@ class CodeBase:
104
86
  return False
105
87
 
106
88
  def make_description(self, file_path: str) -> str:
107
- model = self.platform_registry.create_platform(self.cheap_platform)
108
- model.set_model_name(self.cheap_model)
89
+ model = PlatformRegistry.get_global_platform_registry().get_codegen_platform()
109
90
  model.set_suppress_output(True)
110
91
  content = open(file_path, "r", encoding="utf-8").read()
111
92
  prompt = f"""请分析以下代码文件,并生成一个详细的描述。描述应该包含以下要点:
@@ -413,8 +394,7 @@ class CodeBase:
413
394
  if not initial_results:
414
395
  return []
415
396
 
416
- model = self.platform_registry.create_platform(self.normal_platform)
417
- model.set_model_name(self.normal_model)
397
+ model = PlatformRegistry.get_global_platform_registry().get_normal_platform()
418
398
  model.set_suppress_output(True)
419
399
 
420
400
  try:
@@ -472,8 +452,7 @@ class CodeBase:
472
452
 
473
453
  def search_similar(self, query: str, top_k: int = 20) -> List[Tuple[str, float, str]]:
474
454
  """搜索相似文件"""
475
- model = self.platform_registry.create_platform(self.normal_platform)
476
- model.set_model_name(self.normal_model)
455
+ model = PlatformRegistry.get_global_platform_registry().get_normal_platform()
477
456
  model.set_suppress_output(True)
478
457
 
479
458
  try:
@@ -556,8 +535,7 @@ class CodeBase:
556
535
 
557
536
  请用专业的语言回答用户的问题,如果给出的文件内容不足以回答用户的问题,请告诉用户,绝对不要胡编乱造。
558
537
  """
559
- model = self.platform_registry.create_platform(self.codegen_platform)
560
- model.set_model_name(self.codegen_model)
538
+ model = PlatformRegistry.get_global_platform_registry().get_codegen_platform()
561
539
  try:
562
540
  response = model.chat(prompt)
563
541
  return response
@@ -21,13 +21,7 @@ index_lock = threading.Lock()
21
21
  class JarvisCoder:
22
22
  def __init__(self, root_dir: str, language: str):
23
23
  """初始化代码修改工具"""
24
-
25
- self.platform = os.environ.get("JARVIS_CODEGEN_PLATFORM") or os.environ.get("JARVIS_PLATFORM")
26
- self.model = os.environ.get("JARVIS_CODEGEN_MODEL") or os.environ.get("JARVIS_MODEL")
27
-
28
24
 
29
- if not self.platform or not self.model:
30
- raise ValueError("JARVIS_CODEGEN_PLATFORM or JARVIS_CODEGEN_MODEL is not set")
31
25
 
32
26
  self.root_dir = find_git_root(root_dir)
33
27
  if not self.root_dir:
@@ -486,8 +480,6 @@ def main():
486
480
  load_env_from_file()
487
481
 
488
482
  parser = argparse.ArgumentParser(description='代码修改工具')
489
- parser.add_argument('-p', '--platform', help='AI平台名称', default=os.environ.get('JARVIS_CODEGEN_PLATFORM'))
490
- parser.add_argument('-m', '--model', help='模型名称', default=os.environ.get('JARVIS_CODEGEN_MODEL'))
491
483
  parser.add_argument('-d', '--dir', help='项目根目录', default=os.getcwd())
492
484
  parser.add_argument('-l', '--language', help='编程语言', default="python")
493
485
  args = parser.parse_args()
@@ -495,9 +487,6 @@ def main():
495
487
  # 设置平台
496
488
  if not args.platform:
497
489
  print("错误: 未指定AI平台,请使用 -p 参数")
498
- # 设置模型
499
- if args.model:
500
- os.environ['JARVIS_CODEGEN_MODEL'] = args.model
501
490
 
502
491
  tool = JarvisCoder(args.dir, args.language)
503
492
 
@@ -104,31 +104,16 @@ def main():
104
104
  parser = argparse.ArgumentParser(description='Jarvis AI 助手')
105
105
  parser.add_argument('-f', '--files', nargs='*', help='要处理的文件列表')
106
106
  parser.add_argument('--keep-history', action='store_true', help='保持聊天历史(不删除会话)')
107
- parser.add_argument('-p', '--platform', default=os.getenv('JARVIS_PLATFORM') or 'kimi', help='选择AI平台')
108
- parser.add_argument('-m', '--model', default=os.getenv('JARVIS_MODEL') or '', help='模型') # 用于指定使用的模型名称,默认使用环境变量或平台默认模型
109
107
  args = parser.parse_args()
110
108
 
111
- platform = args.platform if args.platform else os.getenv('JARVIS_PLATFORM')
112
-
113
- if not platform:
114
- PrettyOutput.print("未指定AI平台,请使用 -p 参数或者设置 JARVIS_PLATFORM 环境变量", OutputType.ERROR)
115
- return 1
116
-
117
- PlatformRegistry.get_global_platform_registry().set_global_platform_name(platform)
118
-
119
- if args.model:
120
- os.environ["JARVIS_MODEL"] = args.model
121
-
122
109
  try:
123
110
  # 获取全局模型实例
124
111
  agent = Agent()
125
112
 
126
113
  # 如果用户传入了模型参数,则更换当前模型为用户指定的模型
127
- if args.model:
128
- agent.model.set_model_name(args.model)
129
114
 
130
115
  # 欢迎信息
131
- PrettyOutput.print(f"Jarvis 已初始化 - With {platform} 平台,模型: {agent.model.name()}", OutputType.SYSTEM)
116
+ PrettyOutput.print(f"Jarvis 已初始化 - With {agent.model.name()}", OutputType.SYSTEM)
132
117
  if args.keep_history:
133
118
  PrettyOutput.print("已启用历史保留模式", OutputType.INFO)
134
119
 
@@ -152,14 +152,26 @@ class PlatformRegistry:
152
152
  """
153
153
  self.platforms: Dict[str, Type[BasePlatform]] = {}
154
154
 
155
- @staticmethod
156
- def get_global_platform() -> BasePlatform:
157
- """获取全局平台实例"""
158
- platform = PlatformRegistry.get_global_platform_registry().create_platform(PlatformRegistry.global_platform_name)
159
- if not platform:
160
- raise Exception(f"Failed to create platform: {PlatformRegistry.global_platform_name}")
155
+ def get_normal_platform(self) -> BasePlatform:
156
+ platform_name = os.environ.get("JARVIS_PLATFORM", "kimi")
157
+ model_name = os.environ.get("JARVIS_MODEL", "kimi")
158
+ platform = self.create_platform(platform_name)
159
+ platform.set_model_name(model_name)
160
+ return platform
161
+
162
+ def get_codegen_platform(self) -> BasePlatform:
163
+ platform_name = os.environ.get("JARVIS_CODEGEN_PLATFORM", os.environ.get("JARVIS_PLATFORM", "kimi"))
164
+ model_name = os.environ.get("JARVIS_CODEGEN_MODEL", os.environ.get("JARVIS_MODEL", "kimi"))
165
+ platform = self.create_platform(platform_name)
166
+ platform.set_model_name(model_name)
167
+ return platform
168
+
169
+ def get_cheap_platform(self) -> BasePlatform:
170
+ platform_name = os.environ.get("JARVIS_CHEAP_PLATFORM", os.environ.get("JARVIS_PLATFORM", "kimi"))
171
+ model_name = os.environ.get("JARVIS_CHEAP_MODEL", os.environ.get("JARVIS_MODEL", "kimi"))
172
+ platform = self.create_platform(platform_name)
173
+ platform.set_model_name(model_name)
161
174
  return platform
162
-
163
175
 
164
176
  def register_platform(self, name: str, platform_class: Type[BasePlatform]):
165
177
  """注册平台类
@@ -195,10 +207,3 @@ class PlatformRegistry:
195
207
  """获取可用平台列表"""
196
208
  return list(self.platforms.keys())
197
209
 
198
- def set_global_platform_name(self, platform_name: str):
199
- """设置全局平台"""
200
- PlatformRegistry.global_platform_name = platform_name
201
-
202
- def get_global_platform_name(self) -> str:
203
- """获取全局平台名称"""
204
- return PlatformRegistry.global_platform_name
File without changes