jarvis-ai-assistant 0.1.5__tar.gz → 0.1.6__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 (63) hide show
  1. {jarvis_ai_assistant-0.1.5/src/jarvis_ai_assistant.egg-info → jarvis_ai_assistant-0.1.6}/PKG-INFO +1 -1
  2. {jarvis_ai_assistant-0.1.5 → jarvis_ai_assistant-0.1.6}/pyproject.toml +1 -1
  3. {jarvis_ai_assistant-0.1.5 → jarvis_ai_assistant-0.1.6}/setup.py +1 -1
  4. {jarvis_ai_assistant-0.1.5 → jarvis_ai_assistant-0.1.6}/src/jarvis/__init__.py +1 -1
  5. jarvis_ai_assistant-0.1.6/src/jarvis/__pycache__/__init__.cpython-313.pyc +0 -0
  6. jarvis_ai_assistant-0.1.6/src/jarvis/__pycache__/agent.cpython-313.pyc +0 -0
  7. jarvis_ai_assistant-0.1.6/src/jarvis/__pycache__/main.cpython-313.pyc +0 -0
  8. jarvis_ai_assistant-0.1.6/src/jarvis/__pycache__/models.cpython-313.pyc +0 -0
  9. jarvis_ai_assistant-0.1.6/src/jarvis/__pycache__/utils.cpython-313.pyc +0 -0
  10. jarvis_ai_assistant-0.1.6/src/jarvis/__pycache__/zte_llm.cpython-313.pyc +0 -0
  11. jarvis_ai_assistant-0.1.6/src/jarvis/agent.py +195 -0
  12. {jarvis_ai_assistant-0.1.5 → jarvis_ai_assistant-0.1.6}/src/jarvis/main.py +8 -8
  13. jarvis_ai_assistant-0.1.6/src/jarvis/models.py +130 -0
  14. {jarvis_ai_assistant-0.1.5 → jarvis_ai_assistant-0.1.6}/src/jarvis/tools/__init__.py +3 -3
  15. jarvis_ai_assistant-0.1.6/src/jarvis/tools/__pycache__/__init__.cpython-313.pyc +0 -0
  16. jarvis_ai_assistant-0.1.6/src/jarvis/tools/__pycache__/base.cpython-313.pyc +0 -0
  17. jarvis_ai_assistant-0.1.6/src/jarvis/tools/__pycache__/shell.cpython-313.pyc +0 -0
  18. jarvis_ai_assistant-0.1.6/src/jarvis/tools/__pycache__/sub_agent.cpython-313.pyc +0 -0
  19. jarvis_ai_assistant-0.1.6/src/jarvis/tools/__pycache__/user_input.cpython-313.pyc +0 -0
  20. jarvis_ai_assistant-0.1.6/src/jarvis/tools/base.py +130 -0
  21. jarvis_ai_assistant-0.1.6/src/jarvis/tools/shell.py +95 -0
  22. jarvis_ai_assistant-0.1.6/src/jarvis/tools/sub_agent.py +141 -0
  23. jarvis_ai_assistant-0.1.6/src/jarvis/tools/user_input.py +74 -0
  24. jarvis_ai_assistant-0.1.6/src/jarvis/utils.py +164 -0
  25. {jarvis_ai_assistant-0.1.5 → jarvis_ai_assistant-0.1.6}/src/jarvis/zte_llm.py +5 -1
  26. {jarvis_ai_assistant-0.1.5 → jarvis_ai_assistant-0.1.6/src/jarvis_ai_assistant.egg-info}/PKG-INFO +1 -1
  27. {jarvis_ai_assistant-0.1.5 → jarvis_ai_assistant-0.1.6}/src/jarvis_ai_assistant.egg-info/SOURCES.txt +2 -2
  28. jarvis_ai_assistant-0.1.5/src/jarvis/.jarvis +0 -1
  29. jarvis_ai_assistant-0.1.5/src/jarvis/__pycache__/__init__.cpython-313.pyc +0 -0
  30. jarvis_ai_assistant-0.1.5/src/jarvis/__pycache__/agent.cpython-313.pyc +0 -0
  31. jarvis_ai_assistant-0.1.5/src/jarvis/__pycache__/main.cpython-313.pyc +0 -0
  32. jarvis_ai_assistant-0.1.5/src/jarvis/__pycache__/models.cpython-313.pyc +0 -0
  33. jarvis_ai_assistant-0.1.5/src/jarvis/__pycache__/utils.cpython-313.pyc +0 -0
  34. jarvis_ai_assistant-0.1.5/src/jarvis/__pycache__/zte_llm.cpython-313.pyc +0 -0
  35. jarvis_ai_assistant-0.1.5/src/jarvis/agent.py +0 -133
  36. jarvis_ai_assistant-0.1.5/src/jarvis/models.py +0 -110
  37. jarvis_ai_assistant-0.1.5/src/jarvis/tools/__pycache__/__init__.cpython-313.pyc +0 -0
  38. jarvis_ai_assistant-0.1.5/src/jarvis/tools/__pycache__/base.cpython-313.pyc +0 -0
  39. jarvis_ai_assistant-0.1.5/src/jarvis/tools/__pycache__/shell.cpython-313.pyc +0 -0
  40. jarvis_ai_assistant-0.1.5/src/jarvis/tools/__pycache__/sub_agent.cpython-313.pyc +0 -0
  41. jarvis_ai_assistant-0.1.5/src/jarvis/tools/base.py +0 -198
  42. jarvis_ai_assistant-0.1.5/src/jarvis/tools/python_script.py +0 -150
  43. jarvis_ai_assistant-0.1.5/src/jarvis/tools/shell.py +0 -81
  44. jarvis_ai_assistant-0.1.5/src/jarvis/tools/sub_agent.py +0 -79
  45. jarvis_ai_assistant-0.1.5/src/jarvis/utils.py +0 -128
  46. {jarvis_ai_assistant-0.1.5 → jarvis_ai_assistant-0.1.6}/MANIFEST.in +0 -0
  47. {jarvis_ai_assistant-0.1.5 → jarvis_ai_assistant-0.1.6}/README.md +0 -0
  48. {jarvis_ai_assistant-0.1.5 → jarvis_ai_assistant-0.1.6}/setup.cfg +0 -0
  49. {jarvis_ai_assistant-0.1.5 → jarvis_ai_assistant-0.1.6}/src/jarvis/__pycache__/tools.cpython-313.pyc +0 -0
  50. {jarvis_ai_assistant-0.1.5 → jarvis_ai_assistant-0.1.6}/src/jarvis/tools/__pycache__/file_ops.cpython-313.pyc +0 -0
  51. {jarvis_ai_assistant-0.1.5 → jarvis_ai_assistant-0.1.6}/src/jarvis/tools/__pycache__/python_script.cpython-313.pyc +0 -0
  52. {jarvis_ai_assistant-0.1.5 → jarvis_ai_assistant-0.1.6}/src/jarvis/tools/__pycache__/rag.cpython-313.pyc +0 -0
  53. {jarvis_ai_assistant-0.1.5 → jarvis_ai_assistant-0.1.6}/src/jarvis/tools/__pycache__/search.cpython-313.pyc +0 -0
  54. {jarvis_ai_assistant-0.1.5 → jarvis_ai_assistant-0.1.6}/src/jarvis/tools/__pycache__/user_confirmation.cpython-313.pyc +0 -0
  55. {jarvis_ai_assistant-0.1.5 → jarvis_ai_assistant-0.1.6}/src/jarvis/tools/__pycache__/user_interaction.cpython-313.pyc +0 -0
  56. {jarvis_ai_assistant-0.1.5 → jarvis_ai_assistant-0.1.6}/src/jarvis/tools/__pycache__/webpage.cpython-313.pyc +0 -0
  57. {jarvis_ai_assistant-0.1.5 → jarvis_ai_assistant-0.1.6}/src/jarvis/tools/file_ops.py +0 -0
  58. {jarvis_ai_assistant-0.1.5 → jarvis_ai_assistant-0.1.6}/src/jarvis/tools/search.py +0 -0
  59. {jarvis_ai_assistant-0.1.5 → jarvis_ai_assistant-0.1.6}/src/jarvis/tools/webpage.py +0 -0
  60. {jarvis_ai_assistant-0.1.5 → jarvis_ai_assistant-0.1.6}/src/jarvis_ai_assistant.egg-info/dependency_links.txt +0 -0
  61. {jarvis_ai_assistant-0.1.5 → jarvis_ai_assistant-0.1.6}/src/jarvis_ai_assistant.egg-info/entry_points.txt +0 -0
  62. {jarvis_ai_assistant-0.1.5 → jarvis_ai_assistant-0.1.6}/src/jarvis_ai_assistant.egg-info/requires.txt +0 -0
  63. {jarvis_ai_assistant-0.1.5 → jarvis_ai_assistant-0.1.6}/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.5
3
+ Version: 0.1.6
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
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
4
4
 
5
5
  [project]
6
6
  name = "jarvis-ai-assistant"
7
- version = "0.1.5"
7
+ version = "0.1.6"
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",
5
+ version="0.1.6",
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.5"
3
+ __version__ = "0.1.6"
@@ -0,0 +1,195 @@
1
+ import json
2
+ import subprocess
3
+ from typing import Dict, Any, List, Optional
4
+ from .tools import ToolRegistry
5
+ from .utils import PrettyOutput, OutputType, get_multiline_input
6
+ from .models import BaseModel
7
+ import re
8
+ import os
9
+ from datetime import datetime
10
+
11
+ class Agent:
12
+ def __init__(self, model: BaseModel, tool_registry: ToolRegistry, name: str = "Jarvis"):
13
+ """Initialize Agent with a model, optional tool registry and name"""
14
+ self.model = model
15
+ self.tool_registry = tool_registry or ToolRegistry(model)
16
+ self.name = name
17
+
18
+ # 构建工具说明
19
+ tools_prompt = "Available Tools:\n"
20
+ for tool in self.tool_registry.get_all_tools():
21
+ tools_prompt += f"- Tool: {tool['function']['name']}\n"
22
+ tools_prompt += f" Description: {tool['function']['description']}\n"
23
+ tools_prompt += f" Arguments: {tool['function']['parameters']}\n"
24
+
25
+ self.messages = [
26
+ {
27
+ "role": "system",
28
+ "content": f"""You are {name}, an AI assistant that follows the ReAct (Reasoning + Acting) framework to solve tasks step by step.
29
+
30
+ FRAMEWORK:
31
+ 1. Thought: Analyze the current situation and plan the next step
32
+ 2. Action: Execute ONE specific tool call
33
+ 3. Observation: Review the result
34
+ 4. Next: Plan the next step or conclude
35
+
36
+ FORMAT:
37
+ Thought: I need to [reasoning about the current situation]...
38
+ Action: I will use [tool] to [purpose]...
39
+ <START_TOOL_CALL>
40
+ name: tool_name
41
+ arguments:
42
+ param1: value1
43
+ <END_TOOL_CALL>
44
+
45
+ After receiving result:
46
+ Observation: The tool returned [analyze result]...
47
+ Next: Based on this, I will [next step]...
48
+
49
+ CORE RULES:
50
+ 1. ONE Action Per Response
51
+ - Only ONE tool call per response
52
+ - Additional tool calls will be ignored
53
+ - Complete current step before next
54
+
55
+ 2. Clear Reasoning
56
+ - Explain your thought process
57
+ - Justify tool selection
58
+ - Analyze results thoroughly
59
+
60
+ Examples:
61
+ ✓ Good Response:
62
+ Thought: I need to check the content of utils.py first to understand its structure.
63
+ Action: I will read the file content.
64
+ <START_TOOL_CALL>
65
+ name: file_operation
66
+ arguments:
67
+ operation: read
68
+ filepath: src/utils.py
69
+ <END_TOOL_CALL>
70
+
71
+ ✗ Bad Response:
72
+ Thought: Let's analyze the code.
73
+ Action: I'll read and check everything.
74
+ [Multiple or vague tool calls...]
75
+
76
+ Remember:
77
+ - Always start with "Thought:"
78
+ - Use exactly ONE tool per response
79
+ - Wait for results before next step
80
+ - Clearly explain your reasoning
81
+
82
+ {tools_prompt}"""
83
+ }
84
+ ]
85
+
86
+ def _call_model(self, messages: List[Dict], use_tools: bool = True) -> Dict:
87
+ """调用模型获取响应"""
88
+ try:
89
+ return self.model.chat(
90
+ messages=messages,
91
+ tools=self.tool_registry.get_all_tools() if use_tools else []
92
+ )
93
+ except Exception as e:
94
+ raise Exception(f"{self.name}: 模型调用失败: {str(e)}")
95
+
96
+ def run(self, user_input: str) -> str:
97
+ """处理用户输入并返回响应,返回任务总结报告"""
98
+ self.clear_history()
99
+
100
+ # 显示任务开始
101
+ PrettyOutput.section(f"开始新任务: {self.name}", OutputType.PLANNING)
102
+
103
+ self.messages.append({
104
+ "role": "user",
105
+ "content": user_input
106
+ })
107
+
108
+ while True:
109
+ try:
110
+ # 显示思考状态
111
+ PrettyOutput.print("分析任务...", OutputType.PROGRESS)
112
+ response = self._call_model(self.messages)
113
+ current_response = response
114
+
115
+ # 流式输出已经在model中处理,这里添加换行
116
+ PrettyOutput.print_stream_end()
117
+
118
+ self.messages.append({
119
+ "role": "assistant",
120
+ "content": response["message"].get("content", ""),
121
+ "tool_calls": current_response["message"]["tool_calls"]
122
+ })
123
+
124
+ if len(current_response["message"]["tool_calls"]) > 0:
125
+ if current_response["message"].get("content"):
126
+ PrettyOutput.print(current_response["message"]["content"], OutputType.SYSTEM)
127
+
128
+ try:
129
+ # 显示工具调用
130
+ PrettyOutput.print("执行工具调用...", OutputType.PROGRESS)
131
+ tool_result = self.tool_registry.handle_tool_calls(current_response["message"]["tool_calls"])
132
+ PrettyOutput.print(tool_result, OutputType.RESULT)
133
+ except Exception as e:
134
+ PrettyOutput.print(str(e), OutputType.ERROR)
135
+ tool_result = f"Tool call failed: {str(e)}"
136
+
137
+ self.messages.append({
138
+ "role": "tool",
139
+ "content": tool_result
140
+ })
141
+ continue
142
+
143
+ # 获取用户输入
144
+ user_input = get_multiline_input(f"{self.name}: 您可以继续输入,或输入空行结束当前任务")
145
+ if not user_input:
146
+ PrettyOutput.print("生成任务总结...", OutputType.PROGRESS)
147
+
148
+ # 生成任务总结
149
+ summary_prompt = {
150
+ "role": "user",
151
+ "content": """The task has been completed. Based on the previous analysis and execution results, provide a task summary including:
152
+
153
+ 1. Key Information:
154
+ - Essential findings from analysis
155
+ - Important results from tool executions
156
+ - Critical data discovered
157
+
158
+ 2. Task Results:
159
+ - Final outcome
160
+ - Actual achievements
161
+ - Concrete results
162
+
163
+ Focus only on facts and actual results. Be direct and concise."""
164
+ }
165
+
166
+ while True:
167
+ try:
168
+ summary_response = self._call_model(self.messages + [summary_prompt], use_tools=False)
169
+ summary = summary_response["message"].get("content", "")
170
+
171
+ # 显示任务总结
172
+ PrettyOutput.section("任务总结", OutputType.SUCCESS)
173
+ PrettyOutput.print(summary, OutputType.SYSTEM)
174
+ PrettyOutput.section("任务完成", OutputType.SUCCESS)
175
+
176
+ return summary
177
+
178
+ except Exception as e:
179
+ PrettyOutput.print(str(e), OutputType.ERROR)
180
+
181
+ if user_input == "__interrupt__":
182
+ PrettyOutput.print("任务已取消", OutputType.WARNING)
183
+ return "Task cancelled by user"
184
+
185
+ self.messages.append({
186
+ "role": "user",
187
+ "content": user_input
188
+ })
189
+
190
+ except Exception as e:
191
+ PrettyOutput.print(str(e), OutputType.ERROR)
192
+
193
+ def clear_history(self):
194
+ """清除对话历史,只保留系统提示"""
195
+ self.messages = [self.messages[0]]
@@ -19,7 +19,7 @@ from jarvis.zte_llm import create_zte_llm
19
19
  # 定义支持的平台和模型
20
20
  SUPPORTED_PLATFORMS = {
21
21
  "ollama": {
22
- "models": ["llama3.2", "qwen2.5:14b"],
22
+ "models": ["qwen2.5:14b", "qwq"],
23
23
  "default": "qwen2.5:14b"
24
24
  },
25
25
  "ddgs": {
@@ -88,13 +88,16 @@ def select_task(tasks: dict) -> str:
88
88
 
89
89
  def main():
90
90
  """Main entry point for Jarvis."""
91
+
92
+ load_env_from_file()
93
+
91
94
  parser = argparse.ArgumentParser(description="Jarvis AI Assistant")
92
95
 
93
96
  # 添加平台选择参数
94
97
  parser.add_argument(
95
98
  "--platform",
96
99
  choices=list(SUPPORTED_PLATFORMS.keys()),
97
- default="ddgs",
100
+ default=os.getenv("JARVIS_PLATFORM") or "ddgs",
98
101
  help="选择运行平台 (默认: ollama)"
99
102
  )
100
103
 
@@ -107,13 +110,13 @@ def main():
107
110
  # 添加API基础URL参数
108
111
  parser.add_argument(
109
112
  "--api-base",
110
- default="http://localhost:11434",
113
+ default=os.getenv("JARVIS_OLLAMA_API_BASE") or "http://localhost:11434",
111
114
  help="Ollama API基础URL (仅用于Ollama平台, 默认: http://localhost:11434)"
112
115
  )
113
116
 
114
117
  args = parser.parse_args()
115
118
 
116
- load_env_from_file()
119
+ args.model = args.model or os.getenv("JARVIS_MODEL")
117
120
 
118
121
  # 验证并设置默认模型
119
122
  if args.model:
@@ -162,12 +165,9 @@ def main():
162
165
  while True:
163
166
  try:
164
167
  user_input = get_multiline_input("请输入您的任务(输入空行退出):")
165
- if not user_input:
168
+ if not user_input or user_input == "__interrupt__":
166
169
  break
167
170
  agent.run(user_input)
168
- except KeyboardInterrupt:
169
- print("\n正在退出...")
170
- break
171
171
  except Exception as e:
172
172
  PrettyOutput.print(f"错误: {str(e)}", OutputType.ERROR)
173
173
 
@@ -0,0 +1,130 @@
1
+ import re
2
+ import time
3
+ from typing import Dict, List, Optional
4
+ from duckduckgo_search import DDGS
5
+ import ollama
6
+ from abc import ABC, abstractmethod
7
+ import yaml
8
+
9
+ from .utils import OutputType, PrettyOutput
10
+
11
+ class BaseModel(ABC):
12
+ """大语言模型基类"""
13
+
14
+ @abstractmethod
15
+ def chat(self, messages: List[Dict], tools: Optional[List[Dict]] = None) -> Dict:
16
+ """执行对话"""
17
+ pass
18
+
19
+ @staticmethod
20
+ def extract_tool_calls(content: str) -> List[Dict]:
21
+ """从内容中提取工具调用,只返回第一个有效的工具调用"""
22
+ # 分割内容为行
23
+ lines = content.split('\n')
24
+ tool_call_lines = []
25
+ in_tool_call = False
26
+
27
+ # 逐行处理
28
+ for line in lines:
29
+ if not line:
30
+ continue
31
+
32
+ if line == '<START_TOOL_CALL>':
33
+ tool_call_lines = []
34
+ in_tool_call = True
35
+ continue
36
+ elif line == '<END_TOOL_CALL>':
37
+ if in_tool_call and tool_call_lines:
38
+ try:
39
+ # 解析工具调用内容
40
+ tool_call_text = '\n'.join(tool_call_lines)
41
+ tool_call_data = yaml.safe_load(tool_call_text)
42
+
43
+ # 验证必要的字段
44
+ if "name" in tool_call_data and "arguments" in tool_call_data:
45
+ # 只返回第一个有效的工具调用
46
+ return [{
47
+ "function": {
48
+ "name": tool_call_data["name"],
49
+ "arguments": tool_call_data["arguments"]
50
+ }
51
+ }]
52
+ except yaml.YAMLError:
53
+ pass # 跳过无效的YAML
54
+ except Exception:
55
+ pass # 跳过其他错误
56
+ in_tool_call = False
57
+ elif in_tool_call:
58
+ tool_call_lines.append(line)
59
+
60
+ return [] # 如果没有找到有效的工具调用,返回空列表
61
+
62
+
63
+ class DDGSModel(BaseModel):
64
+ def __init__(self, model_name: str = "gpt-4o-mini"):
65
+ """
66
+ [1]: gpt-4o-mini
67
+ [2]: claude-3-haiku
68
+ [3]: llama-3.1-70b
69
+ [4]: mixtral-8x7b
70
+ """
71
+ self.model_name = model_name
72
+
73
+ def __make_prompt(self, messages: List[Dict], tools: Optional[List[Dict]] = None) -> str:
74
+ prompt = ""
75
+ for message in messages:
76
+ prompt += f"[{message['role']}]: {message['content']}\n"
77
+ return prompt
78
+
79
+ def chat(self, messages: List[Dict], tools: Optional[List[Dict]] = None) -> Dict:
80
+ ddgs = DDGS()
81
+ prompt = self.__make_prompt(messages, tools)
82
+ content = ddgs.chat(prompt)
83
+ PrettyOutput.print_stream(content, OutputType.SYSTEM)
84
+ tool_calls = BaseModel.extract_tool_calls(content)
85
+ return {
86
+ "message": {
87
+ "content": content,
88
+ "tool_calls": tool_calls
89
+ }
90
+ }
91
+
92
+
93
+ class OllamaModel(BaseModel):
94
+ """Ollama模型实现"""
95
+
96
+ def __init__(self, model_name: str = "qwen2.5:14b", api_base: str = "http://localhost:11434"):
97
+ self.model_name = model_name
98
+ self.api_base = api_base
99
+ self.client = ollama.Client(host=api_base)
100
+
101
+ def chat(self, messages: List[Dict], tools: Optional[List[Dict]] = None) -> Dict:
102
+ """调用Ollama API获取响应"""
103
+ try:
104
+ # 使用流式调用
105
+ stream = self.client.chat(
106
+ model=self.model_name,
107
+ messages=messages,
108
+ stream=True
109
+ )
110
+
111
+ # 收集完整响应
112
+ content_parts = []
113
+ for chunk in stream:
114
+ if chunk.message.content:
115
+ content_parts.append(chunk.message.content)
116
+ # 实时打印内容
117
+ PrettyOutput.print_stream(chunk.message.content, OutputType.SYSTEM)
118
+
119
+ # 合并完整内容
120
+ content = "".join(content_parts)
121
+ tool_calls = BaseModel.extract_tool_calls(content)
122
+
123
+ return {
124
+ "message": {
125
+ "content": content,
126
+ "tool_calls": tool_calls
127
+ }
128
+ }
129
+ except Exception as e:
130
+ raise Exception(f"Ollama API调用失败: {str(e)}")
@@ -1,16 +1,16 @@
1
1
  from .base import Tool, ToolRegistry
2
- from .python_script import PythonScript
3
2
  from .file_ops import FileOperationTool
4
3
  from .search import SearchTool
5
4
  from .shell import ShellTool
6
5
  from .webpage import WebpageTool
6
+ from .user_input import UserInputTool
7
7
 
8
8
  __all__ = [
9
9
  'Tool',
10
10
  'ToolRegistry',
11
- 'PythonScript',
12
11
  'FileOperationTool',
13
12
  'SearchTool',
14
13
  'ShellTool',
15
14
  'WebpageTool',
16
- ]
15
+ 'UserInputTool',
16
+ ]
@@ -0,0 +1,130 @@
1
+ from typing import Dict, Any, List, Optional, Callable
2
+ import json
3
+ from ..utils import PrettyOutput, OutputType
4
+ from ..models import BaseModel
5
+
6
+ class Tool:
7
+ def __init__(self, name: str, description: str, parameters: Dict, func: Callable):
8
+ self.name = name
9
+ self.description = description
10
+ self.parameters = parameters
11
+ self.func = func
12
+
13
+ def to_dict(self) -> Dict:
14
+ """转换为Ollama工具格式"""
15
+ return {
16
+ "type": "function",
17
+ "function": {
18
+ "name": self.name,
19
+ "description": self.description,
20
+ "parameters": self.parameters
21
+ }
22
+ }
23
+
24
+ def execute(self, arguments: Dict) -> Dict[str, Any]:
25
+ """执行工具函数"""
26
+ return self.func(arguments)
27
+
28
+ class ToolRegistry:
29
+ def __init__(self, model: BaseModel):
30
+ self.tools: Dict[str, Tool] = {}
31
+ self.model = model
32
+ self._register_default_tools()
33
+
34
+ def _register_default_tools(self):
35
+ """注册所有默认工具"""
36
+ from .search import SearchTool
37
+ from .shell import ShellTool
38
+ from .file_ops import FileOperationTool
39
+ from .webpage import WebpageTool
40
+ from .sub_agent import SubAgentTool
41
+ from .user_input import UserInputTool
42
+
43
+ tools = [
44
+ SearchTool(),
45
+ ShellTool(),
46
+ FileOperationTool(),
47
+ WebpageTool(),
48
+ UserInputTool(),
49
+ ]
50
+
51
+ for tool in tools:
52
+ self.register_tool(
53
+ name=tool.name,
54
+ description=tool.description,
55
+ parameters=tool.parameters,
56
+ func=tool.execute
57
+ )
58
+
59
+ sub_agent_tool = SubAgentTool(self.model)
60
+ self.register_tool(
61
+ name=sub_agent_tool.name,
62
+ description=sub_agent_tool.description,
63
+ parameters=sub_agent_tool.parameters,
64
+ func=sub_agent_tool.execute
65
+ )
66
+
67
+ def register_tool(self, name: str, description: str, parameters: Dict, func: Callable):
68
+ """注册新工具"""
69
+ self.tools[name] = Tool(name, description, parameters, func)
70
+
71
+ def get_tool(self, name: str) -> Optional[Tool]:
72
+ """获取工具"""
73
+ return self.tools.get(name)
74
+
75
+ def get_all_tools(self) -> List[Dict]:
76
+ """获取所有工具的Ollama格式定义"""
77
+ return [tool.to_dict() for tool in self.tools.values()]
78
+
79
+ def execute_tool(self, name: str, arguments: Dict) -> Dict[str, Any]:
80
+ """执行指定工具"""
81
+ tool = self.get_tool(name)
82
+ if tool is None:
83
+ return {"success": False, "error": f"Tool {name} does not exist"}
84
+ return tool.execute(arguments)
85
+
86
+ def handle_tool_calls(self, tool_calls: List[Dict]) -> str:
87
+ """处理工具调用,只处理第一个工具"""
88
+ if not tool_calls:
89
+ return ""
90
+
91
+ # 只处理第一个工具调用
92
+ tool_call = tool_calls[0]
93
+ name = tool_call["function"]["name"]
94
+ args = tool_call["function"]["arguments"]
95
+
96
+ if isinstance(args, str):
97
+ try:
98
+ args = json.loads(args)
99
+ except json.JSONDecodeError:
100
+ PrettyOutput.print(f"工具参数格式无效: {name}", OutputType.ERROR)
101
+ return ""
102
+
103
+ # 显示工具调用信息
104
+ PrettyOutput.section(f"执行工具: {name}", OutputType.TOOL)
105
+ if isinstance(args, dict):
106
+ for key, value in args.items():
107
+ PrettyOutput.print(f"参数: {key} = {value}", OutputType.DEBUG)
108
+ else:
109
+ PrettyOutput.print(f"参数: {args}", OutputType.DEBUG)
110
+
111
+ # 执行工具调用
112
+ result = self.execute_tool(name, args)
113
+
114
+ # 处理结果
115
+ if result["success"]:
116
+ stdout = result["stdout"]
117
+ stderr = result.get("stderr", "")
118
+ output_parts = []
119
+ if stdout:
120
+ output_parts.append(f"输出:\n{stdout}")
121
+ if stderr:
122
+ output_parts.append(f"错误:\n{stderr}")
123
+ output = "\n\n".join(output_parts)
124
+ PrettyOutput.section("执行成功", OutputType.SUCCESS)
125
+ else:
126
+ error_msg = result["error"]
127
+ output = f"执行失败: {error_msg}"
128
+ PrettyOutput.section("执行失败", OutputType.ERROR)
129
+
130
+ return output