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/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", "sudo", "su", "chown", "chmod", "mkfs",
47
- "dd", "shred", ">", ">>", "|", "perl -e", "python -c",
48
- "ruby -e", ":(){ :|:& };:", "eval", "exec", "`"
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
- "success": False,
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": os.path.getsize(file_path) if os.path.isfile(file_path) else None
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 finding files: {str(e)}"
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
+ }