airtrain 0.1.58__py3-none-any.whl → 0.1.62__py3-none-any.whl
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.
- airtrain/__init__.py +72 -44
- airtrain/__pycache__/__init__.cpython-313.pyc +0 -0
- airtrain/core/__pycache__/__init__.cpython-313.pyc +0 -0
- airtrain/core/__pycache__/schemas.cpython-313.pyc +0 -0
- airtrain/core/__pycache__/skills.cpython-313.pyc +0 -0
- airtrain/core/credentials.py +59 -13
- airtrain/integrations/__init__.py +21 -2
- airtrain/integrations/combined/list_models_factory.py +80 -41
- airtrain/integrations/perplexity/__init__.py +49 -0
- airtrain/integrations/perplexity/credentials.py +43 -0
- airtrain/integrations/perplexity/list_models.py +112 -0
- airtrain/integrations/perplexity/models_config.py +128 -0
- airtrain/integrations/perplexity/skills.py +279 -0
- airtrain/integrations/search/__init__.py +21 -0
- airtrain/integrations/search/exa/__init__.py +23 -0
- airtrain/integrations/search/exa/credentials.py +30 -0
- airtrain/integrations/search/exa/schemas.py +114 -0
- airtrain/integrations/search/exa/skills.py +114 -0
- airtrain/tools/__init__.py +9 -5
- airtrain/tools/command.py +248 -61
- airtrain/tools/search.py +450 -0
- airtrain/tools/testing.py +135 -0
- {airtrain-0.1.58.dist-info → airtrain-0.1.62.dist-info}/METADATA +1 -1
- {airtrain-0.1.58.dist-info → airtrain-0.1.62.dist-info}/RECORD +27 -15
- {airtrain-0.1.58.dist-info → airtrain-0.1.62.dist-info}/WHEEL +1 -1
- {airtrain-0.1.58.dist-info → airtrain-0.1.62.dist-info}/entry_points.txt +0 -0
- {airtrain-0.1.58.dist-info → airtrain-0.1.62.dist-info}/top_level.txt +0 -0
airtrain/tools/command.py
CHANGED
@@ -8,52 +8,61 @@ import os
|
|
8
8
|
import subprocess
|
9
9
|
from typing import Dict, Any, List, Optional
|
10
10
|
|
11
|
-
from .registry import StatelessTool, register_tool
|
11
|
+
from .registry import StatelessTool, StatefulTool, register_tool
|
12
12
|
|
13
13
|
|
14
14
|
@register_tool("execute_command")
|
15
15
|
class ExecuteCommandTool(StatelessTool):
|
16
16
|
"""Tool for executing shell commands."""
|
17
|
-
|
17
|
+
|
18
18
|
def __init__(self):
|
19
19
|
self.name = "execute_command"
|
20
20
|
self.description = "Execute a shell command and return its output"
|
21
21
|
self.parameters = {
|
22
22
|
"type": "object",
|
23
23
|
"properties": {
|
24
|
-
"command": {
|
25
|
-
"type": "string",
|
26
|
-
"description": "The command to execute"
|
27
|
-
},
|
24
|
+
"command": {"type": "string", "description": "The command to execute"},
|
28
25
|
"working_dir": {
|
29
26
|
"type": "string",
|
30
|
-
"description": "Working directory for the command"
|
31
|
-
},
|
32
|
-
"timeout": {
|
33
|
-
"type": "number",
|
34
|
-
"description": "Timeout in seconds"
|
27
|
+
"description": "Working directory for the command",
|
35
28
|
},
|
29
|
+
"timeout": {"type": "number", "description": "Timeout in seconds"},
|
36
30
|
"env_vars": {
|
37
31
|
"type": "object",
|
38
|
-
"description": "Environment variables to set for the command"
|
39
|
-
}
|
32
|
+
"description": "Environment variables to set for the command",
|
33
|
+
},
|
40
34
|
},
|
41
|
-
"required": ["command"]
|
35
|
+
"required": ["command"],
|
42
36
|
}
|
43
|
-
|
37
|
+
|
44
38
|
# List of disallowed commands for security
|
45
39
|
self.disallowed_commands = [
|
46
|
-
"rm -rf",
|
47
|
-
"
|
48
|
-
"
|
40
|
+
"rm -rf",
|
41
|
+
"sudo",
|
42
|
+
"su",
|
43
|
+
"chown",
|
44
|
+
"chmod",
|
45
|
+
"mkfs",
|
46
|
+
"dd",
|
47
|
+
"shred",
|
48
|
+
">",
|
49
|
+
">>",
|
50
|
+
"|",
|
51
|
+
"perl -e",
|
52
|
+
"python -c",
|
53
|
+
"ruby -e",
|
54
|
+
":(){ :|:& };:",
|
55
|
+
"eval",
|
56
|
+
"exec",
|
57
|
+
"`",
|
49
58
|
]
|
50
|
-
|
59
|
+
|
51
60
|
def __call__(
|
52
|
-
self,
|
61
|
+
self,
|
53
62
|
command: str,
|
54
63
|
working_dir: Optional[str] = None,
|
55
64
|
timeout: Optional[float] = 30.0,
|
56
|
-
env_vars: Optional[Dict[str, str]] = None
|
65
|
+
env_vars: Optional[Dict[str, str]] = None,
|
57
66
|
) -> Dict[str, Any]:
|
58
67
|
"""Execute a shell command and return its output."""
|
59
68
|
try:
|
@@ -62,14 +71,14 @@ class ExecuteCommandTool(StatelessTool):
|
|
62
71
|
if disallowed in command:
|
63
72
|
return {
|
64
73
|
"success": False,
|
65
|
-
"error": f"Command contains disallowed pattern: {disallowed}"
|
74
|
+
"error": f"Command contains disallowed pattern: {disallowed}",
|
66
75
|
}
|
67
|
-
|
76
|
+
|
68
77
|
# Prepare environment
|
69
78
|
env = os.environ.copy()
|
70
79
|
if env_vars:
|
71
80
|
env.update(env_vars)
|
72
|
-
|
81
|
+
|
73
82
|
# Execute command
|
74
83
|
result = subprocess.run(
|
75
84
|
command,
|
@@ -78,26 +87,23 @@ class ExecuteCommandTool(StatelessTool):
|
|
78
87
|
text=True,
|
79
88
|
cwd=working_dir,
|
80
89
|
timeout=timeout,
|
81
|
-
env=env
|
90
|
+
env=env,
|
82
91
|
)
|
83
|
-
|
92
|
+
|
84
93
|
return {
|
85
94
|
"success": result.returncode == 0,
|
86
95
|
"return_code": result.returncode,
|
87
96
|
"stdout": result.stdout,
|
88
|
-
"stderr": result.stderr
|
97
|
+
"stderr": result.stderr,
|
89
98
|
}
|
90
99
|
except subprocess.TimeoutExpired:
|
91
100
|
return {
|
92
101
|
"success": False,
|
93
|
-
"error": f"Command timed out after {timeout} seconds"
|
102
|
+
"error": f"Command timed out after {timeout} seconds",
|
94
103
|
}
|
95
104
|
except Exception as e:
|
96
|
-
return {
|
97
|
-
|
98
|
-
"error": f"Error executing command: {str(e)}"
|
99
|
-
}
|
100
|
-
|
105
|
+
return {"success": False, "error": f"Error executing command: {str(e)}"}
|
106
|
+
|
101
107
|
def to_dict(self):
|
102
108
|
"""Convert tool to dictionary format for LLM function calling."""
|
103
109
|
return {
|
@@ -105,15 +111,15 @@ class ExecuteCommandTool(StatelessTool):
|
|
105
111
|
"function": {
|
106
112
|
"name": self.name,
|
107
113
|
"description": self.description,
|
108
|
-
"parameters": self.parameters
|
109
|
-
}
|
114
|
+
"parameters": self.parameters,
|
115
|
+
},
|
110
116
|
}
|
111
117
|
|
112
118
|
|
113
119
|
@register_tool("find_files")
|
114
120
|
class FindFilesTool(StatelessTool):
|
115
121
|
"""Tool for finding files matching patterns."""
|
116
|
-
|
122
|
+
|
117
123
|
def __init__(self):
|
118
124
|
self.name = "find_files"
|
119
125
|
self.description = "Find files matching the specified pattern"
|
@@ -122,83 +128,264 @@ class FindFilesTool(StatelessTool):
|
|
122
128
|
"properties": {
|
123
129
|
"directory": {
|
124
130
|
"type": "string",
|
125
|
-
"description": "Directory to search in"
|
131
|
+
"description": "Directory to search in",
|
126
132
|
},
|
127
133
|
"pattern": {
|
128
134
|
"type": "string",
|
129
|
-
"description": "Glob pattern to match (e.g., *.txt, **/*.py)"
|
135
|
+
"description": "Glob pattern to match (e.g., *.txt, **/*.py)",
|
130
136
|
},
|
131
137
|
"max_results": {
|
132
138
|
"type": "integer",
|
133
|
-
"description": "Maximum number of results to return"
|
139
|
+
"description": "Maximum number of results to return",
|
134
140
|
},
|
135
141
|
"show_hidden": {
|
136
142
|
"type": "boolean",
|
137
|
-
"description": "Whether to include hidden files (starting with .)"
|
138
|
-
}
|
143
|
+
"description": "Whether to include hidden files (starting with .)",
|
144
|
+
},
|
139
145
|
},
|
140
|
-
"required": ["directory", "pattern"]
|
146
|
+
"required": ["directory", "pattern"],
|
141
147
|
}
|
142
|
-
|
148
|
+
|
143
149
|
def __call__(
|
144
150
|
self,
|
145
151
|
directory: str,
|
146
152
|
pattern: str,
|
147
153
|
max_results: int = 100,
|
148
|
-
show_hidden: bool = False
|
154
|
+
show_hidden: bool = False,
|
149
155
|
) -> Dict[str, Any]:
|
150
156
|
"""Find files matching the specified pattern."""
|
151
157
|
try:
|
152
158
|
import glob
|
153
159
|
from pathlib import Path
|
154
|
-
|
160
|
+
|
155
161
|
directory = os.path.expanduser(directory)
|
156
162
|
if not os.path.exists(directory):
|
157
163
|
return {
|
158
164
|
"success": False,
|
159
|
-
"error": f"Directory '{directory}' does not exist"
|
165
|
+
"error": f"Directory '{directory}' does not exist",
|
160
166
|
}
|
161
|
-
|
167
|
+
|
162
168
|
if not os.path.isdir(directory):
|
163
169
|
return {
|
164
170
|
"success": False,
|
165
|
-
"error": f"Path '{directory}' is not a directory"
|
171
|
+
"error": f"Path '{directory}' is not a directory",
|
166
172
|
}
|
167
|
-
|
173
|
+
|
168
174
|
# Construct search path
|
169
175
|
search_path = os.path.join(directory, pattern)
|
170
|
-
|
176
|
+
|
171
177
|
# Find matching files
|
172
178
|
files = []
|
173
179
|
for file_path in glob.glob(search_path, recursive=True):
|
174
|
-
if not show_hidden and os.path.basename(file_path).startswith(
|
180
|
+
if not show_hidden and os.path.basename(file_path).startswith("."):
|
175
181
|
continue
|
176
|
-
|
182
|
+
|
177
183
|
file_info = {
|
178
184
|
"path": file_path,
|
179
185
|
"name": os.path.basename(file_path),
|
180
186
|
"type": "dir" if os.path.isdir(file_path) else "file",
|
181
|
-
"size":
|
187
|
+
"size": (
|
188
|
+
os.path.getsize(file_path)
|
189
|
+
if os.path.isfile(file_path)
|
190
|
+
else None
|
191
|
+
),
|
182
192
|
}
|
183
193
|
files.append(file_info)
|
184
|
-
|
194
|
+
|
185
195
|
if len(files) >= max_results:
|
186
196
|
break
|
187
|
-
|
197
|
+
|
188
198
|
return {
|
189
199
|
"success": True,
|
190
200
|
"directory": directory,
|
191
201
|
"pattern": pattern,
|
192
202
|
"files": files,
|
193
203
|
"count": len(files),
|
194
|
-
"truncated": len(files) >= max_results
|
204
|
+
"truncated": len(files) >= max_results,
|
195
205
|
}
|
206
|
+
except Exception as e:
|
207
|
+
return {"success": False, "error": f"Error finding files: {str(e)}"}
|
208
|
+
|
209
|
+
def to_dict(self):
|
210
|
+
"""Convert tool to dictionary format for LLM function calling."""
|
211
|
+
return {
|
212
|
+
"type": "function",
|
213
|
+
"function": {
|
214
|
+
"name": self.name,
|
215
|
+
"description": self.description,
|
216
|
+
"parameters": self.parameters,
|
217
|
+
},
|
218
|
+
}
|
219
|
+
|
220
|
+
|
221
|
+
@register_tool("terminal_navigation", tool_type="stateful")
|
222
|
+
class TerminalNavigationTool(StatefulTool):
|
223
|
+
"""Tool for navigating through the terminal with state memory."""
|
224
|
+
|
225
|
+
def __init__(self):
|
226
|
+
self.name = "terminal_navigation"
|
227
|
+
self.description = (
|
228
|
+
"Navigate through the terminal with persistent directory state"
|
229
|
+
)
|
230
|
+
self.parameters = {
|
231
|
+
"type": "object",
|
232
|
+
"properties": {
|
233
|
+
"action": {
|
234
|
+
"type": "string",
|
235
|
+
"enum": ["cd", "pwd", "pushd", "popd", "dirs"],
|
236
|
+
"description": "Navigation action to perform",
|
237
|
+
},
|
238
|
+
"directory": {
|
239
|
+
"type": "string",
|
240
|
+
"description": "Target directory for cd and pushd actions",
|
241
|
+
},
|
242
|
+
},
|
243
|
+
"required": ["action"],
|
244
|
+
}
|
245
|
+
self.reset()
|
246
|
+
|
247
|
+
@classmethod
|
248
|
+
def create_instance(cls):
|
249
|
+
"""Create a new instance with fresh state."""
|
250
|
+
return cls()
|
251
|
+
|
252
|
+
def reset(self):
|
253
|
+
"""Reset the terminal navigation state."""
|
254
|
+
self.current_dir = os.getcwd()
|
255
|
+
self.dir_stack = []
|
256
|
+
|
257
|
+
def __call__(self, action: str, directory: Optional[str] = None) -> Dict[str, Any]:
|
258
|
+
"""Execute the terminal navigation action."""
|
259
|
+
try:
|
260
|
+
# Handle the different navigation actions
|
261
|
+
if action == "cd" and directory:
|
262
|
+
# Expand user path if present
|
263
|
+
target_dir = os.path.expanduser(directory)
|
264
|
+
|
265
|
+
# Handle relative paths
|
266
|
+
if not os.path.isabs(target_dir):
|
267
|
+
target_dir = os.path.join(self.current_dir, target_dir)
|
268
|
+
|
269
|
+
# Normalize the path
|
270
|
+
target_dir = os.path.normpath(target_dir)
|
271
|
+
|
272
|
+
# Check if directory exists
|
273
|
+
if not os.path.exists(target_dir):
|
274
|
+
return {
|
275
|
+
"success": False,
|
276
|
+
"error": f"Directory does not exist: {target_dir}",
|
277
|
+
"current_dir": self.current_dir,
|
278
|
+
}
|
279
|
+
|
280
|
+
if not os.path.isdir(target_dir):
|
281
|
+
return {
|
282
|
+
"success": False,
|
283
|
+
"error": f"Path is not a directory: {target_dir}",
|
284
|
+
"current_dir": self.current_dir,
|
285
|
+
}
|
286
|
+
|
287
|
+
# Change directory
|
288
|
+
self.current_dir = target_dir
|
289
|
+
return {
|
290
|
+
"success": True,
|
291
|
+
"action": "cd",
|
292
|
+
"previous_dir": self.current_dir,
|
293
|
+
"current_dir": self.current_dir,
|
294
|
+
}
|
295
|
+
|
296
|
+
elif action == "pwd":
|
297
|
+
# Print working directory
|
298
|
+
return {
|
299
|
+
"success": True,
|
300
|
+
"action": "pwd",
|
301
|
+
"current_dir": self.current_dir,
|
302
|
+
}
|
303
|
+
|
304
|
+
elif action == "pushd" and directory:
|
305
|
+
# Expand user path if present
|
306
|
+
target_dir = os.path.expanduser(directory)
|
307
|
+
|
308
|
+
# Handle relative paths
|
309
|
+
if not os.path.isabs(target_dir):
|
310
|
+
target_dir = os.path.join(self.current_dir, target_dir)
|
311
|
+
|
312
|
+
# Normalize the path
|
313
|
+
target_dir = os.path.normpath(target_dir)
|
314
|
+
|
315
|
+
# Check if directory exists
|
316
|
+
if not os.path.exists(target_dir):
|
317
|
+
return {
|
318
|
+
"success": False,
|
319
|
+
"error": f"Directory does not exist: {target_dir}",
|
320
|
+
"current_dir": self.current_dir,
|
321
|
+
"dir_stack": self.dir_stack,
|
322
|
+
}
|
323
|
+
|
324
|
+
if not os.path.isdir(target_dir):
|
325
|
+
return {
|
326
|
+
"success": False,
|
327
|
+
"error": f"Path is not a directory: {target_dir}",
|
328
|
+
"current_dir": self.current_dir,
|
329
|
+
"dir_stack": self.dir_stack,
|
330
|
+
}
|
331
|
+
|
332
|
+
# Push current directory onto stack and change to new directory
|
333
|
+
self.dir_stack.append(self.current_dir)
|
334
|
+
self.current_dir = target_dir
|
335
|
+
|
336
|
+
return {
|
337
|
+
"success": True,
|
338
|
+
"action": "pushd",
|
339
|
+
"previous_dir": self.dir_stack[-1],
|
340
|
+
"current_dir": self.current_dir,
|
341
|
+
"dir_stack": self.dir_stack,
|
342
|
+
}
|
343
|
+
|
344
|
+
elif action == "popd":
|
345
|
+
# Check if directory stack is empty
|
346
|
+
if not self.dir_stack:
|
347
|
+
return {
|
348
|
+
"success": False,
|
349
|
+
"error": "Directory stack is empty",
|
350
|
+
"current_dir": self.current_dir,
|
351
|
+
"dir_stack": [],
|
352
|
+
}
|
353
|
+
|
354
|
+
# Pop directory from stack and change to it
|
355
|
+
previous_dir = self.current_dir
|
356
|
+
self.current_dir = self.dir_stack.pop()
|
357
|
+
|
358
|
+
return {
|
359
|
+
"success": True,
|
360
|
+
"action": "popd",
|
361
|
+
"previous_dir": previous_dir,
|
362
|
+
"current_dir": self.current_dir,
|
363
|
+
"dir_stack": self.dir_stack,
|
364
|
+
}
|
365
|
+
|
366
|
+
elif action == "dirs":
|
367
|
+
# Display directory stack
|
368
|
+
return {
|
369
|
+
"success": True,
|
370
|
+
"action": "dirs",
|
371
|
+
"current_dir": self.current_dir,
|
372
|
+
"dir_stack": self.dir_stack,
|
373
|
+
}
|
374
|
+
|
375
|
+
else:
|
376
|
+
return {
|
377
|
+
"success": False,
|
378
|
+
"error": f"Invalid action '{action}' or missing required parameters",
|
379
|
+
"current_dir": self.current_dir,
|
380
|
+
}
|
381
|
+
|
196
382
|
except Exception as e:
|
197
383
|
return {
|
198
384
|
"success": False,
|
199
|
-
"error": f"Error
|
385
|
+
"error": f"Error navigating terminal: {str(e)}",
|
386
|
+
"current_dir": self.current_dir,
|
200
387
|
}
|
201
|
-
|
388
|
+
|
202
389
|
def to_dict(self):
|
203
390
|
"""Convert tool to dictionary format for LLM function calling."""
|
204
391
|
return {
|
@@ -206,6 +393,6 @@ class FindFilesTool(StatelessTool):
|
|
206
393
|
"function": {
|
207
394
|
"name": self.name,
|
208
395
|
"description": self.description,
|
209
|
-
"parameters": self.parameters
|
210
|
-
}
|
211
|
-
}
|
396
|
+
"parameters": self.parameters,
|
397
|
+
},
|
398
|
+
}
|