tinyagent-py 0.0.15__py3-none-any.whl → 0.0.16rc0__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.
- tinyagent/code_agent/providers/__init__.py +14 -1
- tinyagent/code_agent/providers/base.py +29 -1
- tinyagent/code_agent/providers/modal_provider.py +9 -0
- tinyagent/code_agent/providers/seatbelt_provider.py +1065 -0
- tinyagent/code_agent/tiny_code_agent.py +692 -8
- tinyagent/code_agent/utils.py +187 -22
- tinyagent/prompts/truncation.yaml +13 -0
- tinyagent/tiny_agent.py +402 -49
- {tinyagent_py-0.0.15.dist-info → tinyagent_py-0.0.16rc0.dist-info}/METADATA +25 -1
- {tinyagent_py-0.0.15.dist-info → tinyagent_py-0.0.16rc0.dist-info}/RECORD +13 -11
- {tinyagent_py-0.0.15.dist-info → tinyagent_py-0.0.16rc0.dist-info}/WHEEL +0 -0
- {tinyagent_py-0.0.15.dist-info → tinyagent_py-0.0.16rc0.dist-info}/licenses/LICENSE +0 -0
- {tinyagent_py-0.0.15.dist-info → tinyagent_py-0.0.16rc0.dist-info}/top_level.txt +0 -0
@@ -7,10 +7,20 @@ from pathlib import Path
|
|
7
7
|
from tinyagent import TinyAgent, tool
|
8
8
|
from tinyagent.hooks.logging_manager import LoggingManager
|
9
9
|
from tinyagent.hooks.rich_code_ui_callback import RichCodeUICallback
|
10
|
-
|
10
|
+
# Conditional import for Jupyter callback - only import when needed
|
11
|
+
try:
|
12
|
+
from tinyagent.hooks.jupyter_notebook_callback import JupyterNotebookCallback, OptimizedJupyterNotebookCallback
|
13
|
+
JUPYTER_CALLBACKS_AVAILABLE = True
|
14
|
+
except ImportError:
|
15
|
+
JUPYTER_CALLBACKS_AVAILABLE = False
|
16
|
+
JupyterNotebookCallback = None
|
17
|
+
OptimizedJupyterNotebookCallback = None
|
11
18
|
from .providers.base import CodeExecutionProvider
|
12
19
|
from .providers.modal_provider import ModalProvider
|
20
|
+
from .providers.seatbelt_provider import SeatbeltProvider
|
13
21
|
from .helper import translate_tool_for_code_agent, load_template, render_system_prompt, prompt_code_example, prompt_qwen_helper
|
22
|
+
from .utils import truncate_output, format_truncation_message
|
23
|
+
import datetime
|
14
24
|
|
15
25
|
|
16
26
|
DEFAULT_SUMMARY_SYSTEM_PROMPT = (
|
@@ -26,7 +36,16 @@ class TinyCodeAgent:
|
|
26
36
|
A TinyAgent specialized for code execution tasks.
|
27
37
|
|
28
38
|
This class provides a high-level interface for creating agents that can execute
|
29
|
-
Python code using various providers (Modal,
|
39
|
+
Python code using various providers (Modal, SeatbeltProvider for macOS sandboxing, etc.).
|
40
|
+
|
41
|
+
Features include:
|
42
|
+
- Code execution in sandboxed environments
|
43
|
+
- Shell command execution with safety checks
|
44
|
+
- Environment variable management (SeatbeltProvider)
|
45
|
+
- File system access controls
|
46
|
+
- Memory management and conversation summarization
|
47
|
+
- Git checkpoint automation
|
48
|
+
- Output truncation controls
|
30
49
|
"""
|
31
50
|
|
32
51
|
def __init__(
|
@@ -48,6 +67,8 @@ class TinyCodeAgent:
|
|
48
67
|
default_workdir: Optional[str] = None,
|
49
68
|
summary_config: Optional[Dict[str, Any]] = None,
|
50
69
|
ui: Optional[str] = None,
|
70
|
+
truncation_config: Optional[Dict[str, Any]] = None,
|
71
|
+
auto_git_checkpoint: bool = False,
|
51
72
|
**agent_kwargs
|
52
73
|
):
|
53
74
|
"""
|
@@ -72,7 +93,33 @@ class TinyCodeAgent:
|
|
72
93
|
default_workdir: Default working directory for shell commands. If None, the current working directory is used.
|
73
94
|
summary_config: Optional configuration for generating conversation summaries
|
74
95
|
ui: The user interface callback to use ('rich', 'jupyter', or None).
|
96
|
+
truncation_config: Configuration for output truncation (max_tokens, max_lines)
|
97
|
+
auto_git_checkpoint: If True, automatically create git checkpoints after each successful shell command
|
75
98
|
**agent_kwargs: Additional arguments passed to TinyAgent
|
99
|
+
|
100
|
+
Provider Config Options:
|
101
|
+
For SeatbeltProvider:
|
102
|
+
- seatbelt_profile: String containing seatbelt profile rules
|
103
|
+
- seatbelt_profile_path: Path to a file containing seatbelt profile rules
|
104
|
+
- python_env_path: Path to the Python environment to use
|
105
|
+
- bypass_shell_safety: If True, bypass shell command safety checks (default: True for seatbelt)
|
106
|
+
- additional_safe_shell_commands: Additional shell commands to consider safe
|
107
|
+
- additional_safe_control_operators: Additional shell control operators to consider safe
|
108
|
+
- additional_read_dirs: List of additional directories to allow read access to
|
109
|
+
- additional_write_dirs: List of additional directories to allow write access to
|
110
|
+
- environment_variables: Dictionary of environment variables to make available in the sandbox
|
111
|
+
|
112
|
+
For ModalProvider:
|
113
|
+
- pip_packages: List of additional Python packages to install
|
114
|
+
- authorized_imports: List of authorized Python imports
|
115
|
+
- bypass_shell_safety: If True, bypass shell command safety checks (default: False for modal)
|
116
|
+
- additional_safe_shell_commands: Additional shell commands to consider safe
|
117
|
+
- additional_safe_control_operators: Additional shell control operators to consider safe
|
118
|
+
|
119
|
+
Truncation Config Options:
|
120
|
+
- max_tokens: Maximum number of tokens to keep in output (default: 3000)
|
121
|
+
- max_lines: Maximum number of lines to keep in output (default: 250)
|
122
|
+
- enabled: Whether truncation is enabled (default: True)
|
76
123
|
"""
|
77
124
|
self.model = model
|
78
125
|
self.api_key = api_key
|
@@ -87,6 +134,15 @@ class TinyCodeAgent:
|
|
87
134
|
self.provider = provider # Store provider type for reuse
|
88
135
|
self.check_string_obfuscation = check_string_obfuscation
|
89
136
|
self.default_workdir = default_workdir or os.getcwd() # Default to current working directory if not specified
|
137
|
+
self.auto_git_checkpoint = auto_git_checkpoint # Enable/disable automatic git checkpoints
|
138
|
+
|
139
|
+
# Set up truncation configuration with defaults
|
140
|
+
default_truncation = {
|
141
|
+
"max_tokens": 3000,
|
142
|
+
"max_lines": 250,
|
143
|
+
"enabled": True
|
144
|
+
}
|
145
|
+
self.truncation_config = {**default_truncation, **(truncation_config or {})}
|
90
146
|
|
91
147
|
# Create the code execution provider
|
92
148
|
self.code_provider = self._create_provider(provider, self.provider_config)
|
@@ -139,12 +195,70 @@ class TinyCodeAgent:
|
|
139
195
|
final_config["authorized_imports"] = final_authorized_imports
|
140
196
|
final_config["check_string_obfuscation"] = self.check_string_obfuscation
|
141
197
|
|
198
|
+
# Shell safety configuration (default to False for Modal)
|
199
|
+
bypass_shell_safety = config.get("bypass_shell_safety", False)
|
200
|
+
additional_safe_shell_commands = config.get("additional_safe_shell_commands", None)
|
201
|
+
additional_safe_control_operators = config.get("additional_safe_control_operators", None)
|
202
|
+
|
142
203
|
return ModalProvider(
|
143
204
|
log_manager=self.log_manager,
|
144
205
|
code_tools=self.code_tools,
|
145
206
|
local_execution=self.local_execution,
|
207
|
+
bypass_shell_safety=bypass_shell_safety,
|
208
|
+
additional_safe_shell_commands=additional_safe_shell_commands,
|
209
|
+
additional_safe_control_operators=additional_safe_control_operators,
|
146
210
|
**final_config
|
147
211
|
)
|
212
|
+
elif provider_type.lower() == "seatbelt":
|
213
|
+
# Check if seatbelt is supported on this system
|
214
|
+
if not SeatbeltProvider.is_supported():
|
215
|
+
raise ValueError("Seatbelt provider is not supported on this system. It requires macOS with sandbox-exec.")
|
216
|
+
|
217
|
+
# Seatbelt only works with local execution
|
218
|
+
if not self.local_execution:
|
219
|
+
raise ValueError("Seatbelt provider requires local execution mode. Please set local_execution=True.")
|
220
|
+
|
221
|
+
# Create a copy of the config without the parameters we'll pass directly
|
222
|
+
filtered_config = config.copy()
|
223
|
+
for key in ['seatbelt_profile', 'seatbelt_profile_path', 'python_env_path',
|
224
|
+
'bypass_shell_safety', 'additional_safe_shell_commands',
|
225
|
+
'additional_safe_control_operators', 'additional_read_dirs',
|
226
|
+
'additional_write_dirs', 'environment_variables']:
|
227
|
+
if key in filtered_config:
|
228
|
+
filtered_config.pop(key)
|
229
|
+
|
230
|
+
# Get seatbelt profile configuration
|
231
|
+
seatbelt_profile = config.get("seatbelt_profile", None)
|
232
|
+
seatbelt_profile_path = config.get("seatbelt_profile_path", None)
|
233
|
+
python_env_path = config.get("python_env_path", None)
|
234
|
+
|
235
|
+
# Shell safety configuration (default to True for Seatbelt)
|
236
|
+
bypass_shell_safety = config.get("bypass_shell_safety", True)
|
237
|
+
additional_safe_shell_commands = config.get("additional_safe_shell_commands", None)
|
238
|
+
additional_safe_control_operators = config.get("additional_safe_control_operators", None)
|
239
|
+
|
240
|
+
# Additional directory access configuration
|
241
|
+
additional_read_dirs = config.get("additional_read_dirs", None)
|
242
|
+
additional_write_dirs = config.get("additional_write_dirs", None)
|
243
|
+
|
244
|
+
# Environment variables to make available in the sandbox
|
245
|
+
environment_variables = config.get("environment_variables", {})
|
246
|
+
|
247
|
+
# Create the seatbelt provider
|
248
|
+
return SeatbeltProvider(
|
249
|
+
log_manager=self.log_manager,
|
250
|
+
code_tools=self.code_tools,
|
251
|
+
seatbelt_profile=seatbelt_profile,
|
252
|
+
seatbelt_profile_path=seatbelt_profile_path,
|
253
|
+
python_env_path=python_env_path,
|
254
|
+
bypass_shell_safety=bypass_shell_safety,
|
255
|
+
additional_safe_shell_commands=additional_safe_shell_commands,
|
256
|
+
additional_safe_control_operators=additional_safe_control_operators,
|
257
|
+
additional_read_dirs=additional_read_dirs,
|
258
|
+
additional_write_dirs=additional_write_dirs,
|
259
|
+
environment_variables=environment_variables,
|
260
|
+
**filtered_config
|
261
|
+
)
|
148
262
|
else:
|
149
263
|
raise ValueError(f"Unsupported provider type: {provider_type}")
|
150
264
|
|
@@ -307,6 +421,24 @@ class TinyCodeAgent:
|
|
307
421
|
# This ensures they stay in sync
|
308
422
|
self.user_variables = self.code_provider.get_user_variables()
|
309
423
|
|
424
|
+
# Apply truncation if enabled
|
425
|
+
if self.truncation_config["enabled"] and "printed_output" in result:
|
426
|
+
truncated_output, is_truncated, original_tokens, original_lines = truncate_output(
|
427
|
+
result["printed_output"],
|
428
|
+
max_tokens=self.truncation_config["max_tokens"],
|
429
|
+
max_lines=self.truncation_config["max_lines"]
|
430
|
+
)
|
431
|
+
|
432
|
+
if is_truncated:
|
433
|
+
result["printed_output"] = format_truncation_message(
|
434
|
+
truncated_output,
|
435
|
+
is_truncated,
|
436
|
+
original_tokens,
|
437
|
+
original_lines,
|
438
|
+
self.truncation_config["max_lines"],
|
439
|
+
"python_output"
|
440
|
+
)
|
441
|
+
|
310
442
|
return json.dumps(result)
|
311
443
|
except Exception as e:
|
312
444
|
print("!"*100)
|
@@ -348,7 +480,7 @@ class TinyCodeAgent:
|
|
348
480
|
- You have the capability to call multiple tools in a single response. When multiple independent pieces of information are requested, batch your tool calls together for optimal performance.
|
349
481
|
- You can specify an optional timeout in milliseconds (up to 600000ms / 10 minutes). If not specified, commands will timeout after 120000ms (2 minutes).
|
350
482
|
- It is very helpful if you write a clear, concise description of what this command does in 5-10 words.
|
351
|
-
- If the output
|
483
|
+
- If the output is too large, it will be truncated before being returned to you.
|
352
484
|
|
353
485
|
- If you _still_ need to run `grep`, STOP. ALWAYS USE ripgrep at `rg` first, which all ${PRODUCT_NAME} users have pre-installed.
|
354
486
|
" - When issuing multiple commands, use the ; or && operator to separate them. DO NOT use newlines (newlines are ok in quoted strings).\n" +
|
@@ -359,19 +491,66 @@ class TinyCodeAgent:
|
|
359
491
|
<bad-example>
|
360
492
|
cd /foo/bar && pytest tests
|
361
493
|
</bad-example>
|
494
|
+
|
495
|
+
## IMPORTANT: Bash Tool Usage
|
496
|
+
|
497
|
+
When using the bash tool, you MUST provide all required parameters:
|
498
|
+
|
499
|
+
**Correct Usage:**
|
500
|
+
```
|
501
|
+
bash(
|
502
|
+
command=["ls", "-la"],
|
503
|
+
absolute_workdir="/path/to/directory",
|
504
|
+
description="List files in directory"
|
505
|
+
)
|
506
|
+
```
|
507
|
+
|
508
|
+
**For creating files with content, use these safe patterns:**
|
509
|
+
|
510
|
+
1. **Simple file creation:**
|
511
|
+
```
|
512
|
+
bash(
|
513
|
+
command=["touch", "filename.txt"],
|
514
|
+
absolute_workdir="/working/directory",
|
515
|
+
description="Create empty file"
|
516
|
+
)
|
517
|
+
```
|
518
|
+
|
519
|
+
2. **Write content using cat and heredoc:**
|
520
|
+
```
|
521
|
+
bash(
|
522
|
+
command=["sh", "-c", "cat > filename.txt << 'EOF'\nYour content here\nEOF"],
|
523
|
+
absolute_workdir="/working/directory",
|
524
|
+
description="Create file with content"
|
525
|
+
)
|
526
|
+
```
|
527
|
+
|
528
|
+
3. **Write content using echo:**
|
529
|
+
```
|
530
|
+
bash(
|
531
|
+
command=["sh", "-c", "echo 'Your content' > filename.txt"],
|
532
|
+
absolute_workdir="/working/directory",
|
533
|
+
description="Write content to file"
|
534
|
+
)
|
535
|
+
```
|
536
|
+
|
537
|
+
**Never:**
|
538
|
+
- Call bash() without all required parameters
|
539
|
+
- Use complex nested quotes without testing
|
540
|
+
- Try to create large files in a single command (break into parts)
|
362
541
|
|
363
542
|
Args:
|
364
543
|
command: list[str]: The shell command to execute as a list of strings. Example: ["ls", "-la"] or ["cat", "file.txt"]
|
365
544
|
|
366
545
|
absolute_workdir: str: could be presented workdir in the system prompt or one of the subdirectories of the workdir. This is the only allowed path, and accessing else will result in an error.
|
367
546
|
description: str: A clear, concise description of what this command does in 5-10 words.
|
368
|
-
timeout: int: Maximum execution time in seconds (default:
|
547
|
+
timeout: int: Maximum execution time in seconds (default: 60).
|
369
548
|
Returns:
|
370
549
|
Dictionary with stdout, stderr, and exit_code from the command execution.
|
371
550
|
If the command is rejected for security reasons, stderr will contain the reason.
|
372
551
|
The stdout will include information about which working directory was used.
|
373
552
|
"""))
|
374
|
-
async def run_shell(command: List[str], absolute_workdir: str, description: str, timeout: int =
|
553
|
+
async def run_shell(command: List[str], absolute_workdir: str, description: str, timeout: int = 60) -> str:
|
375
554
|
"""Execute shell commands securely using the configured provider."""
|
376
555
|
try:
|
377
556
|
# Use the default working directory if none is specified
|
@@ -393,6 +572,30 @@ class TinyCodeAgent:
|
|
393
572
|
})
|
394
573
|
|
395
574
|
result = await self.code_provider.execute_shell(command, timeout, effective_workdir)
|
575
|
+
|
576
|
+
# Apply truncation if enabled
|
577
|
+
if self.truncation_config["enabled"] and "stdout" in result and result["stdout"]:
|
578
|
+
truncated_output, is_truncated, original_tokens, original_lines = truncate_output(
|
579
|
+
result["stdout"],
|
580
|
+
max_tokens=self.truncation_config["max_tokens"],
|
581
|
+
max_lines=self.truncation_config["max_lines"]
|
582
|
+
)
|
583
|
+
|
584
|
+
if is_truncated:
|
585
|
+
result["stdout"] = format_truncation_message(
|
586
|
+
truncated_output,
|
587
|
+
is_truncated,
|
588
|
+
original_tokens,
|
589
|
+
original_lines,
|
590
|
+
self.truncation_config["max_lines"],
|
591
|
+
"bash_output"
|
592
|
+
)
|
593
|
+
|
594
|
+
# Create a git checkpoint if auto_git_checkpoint is enabled
|
595
|
+
if self.auto_git_checkpoint and result.get("exit_code", 1) == 0:
|
596
|
+
checkpoint_result = await self._create_git_checkpoint(command, description, effective_workdir)
|
597
|
+
self.log_manager.get_logger(__name__).info(f"Git checkpoint {effective_workdir} result: {checkpoint_result}")
|
598
|
+
|
396
599
|
return json.dumps(result)
|
397
600
|
except Exception as e:
|
398
601
|
COLOR = {
|
@@ -407,6 +610,64 @@ class TinyCodeAgent:
|
|
407
610
|
self.agent.add_tool(run_python)
|
408
611
|
self.agent.add_tool(run_shell)
|
409
612
|
|
613
|
+
async def _create_git_checkpoint(self, command: List[str], description: str, workdir: str) -> Dict[str, Any]:
|
614
|
+
"""
|
615
|
+
Create a git checkpoint after command execution.
|
616
|
+
|
617
|
+
Args:
|
618
|
+
command: The command that was executed
|
619
|
+
description: Description of the command
|
620
|
+
workdir: Working directory where the command was executed
|
621
|
+
|
622
|
+
Returns:
|
623
|
+
Dictionary with stdout and stderr from the git operations
|
624
|
+
"""
|
625
|
+
try:
|
626
|
+
# Format the command for the commit message
|
627
|
+
cmd_str = " ".join(command)
|
628
|
+
|
629
|
+
# Check if there are changes to commit
|
630
|
+
git_check_cmd = ["bash", "-c", "if ! git diff-index --quiet HEAD --; then echo 'changes_exist'; else echo 'no_changes'; fi"]
|
631
|
+
check_result = await self.code_provider.execute_shell(git_check_cmd, 10, workdir)
|
632
|
+
|
633
|
+
# If no changes or check failed, return early
|
634
|
+
if check_result.get("exit_code", 1) != 0 or "no_changes" in check_result.get("stdout", ""):
|
635
|
+
return {"stdout": "No changes detected, skipping git checkpoint", "stderr": ""}
|
636
|
+
|
637
|
+
# Stage all changes
|
638
|
+
git_add_cmd = ["git", "add", "-A"]
|
639
|
+
add_result = await self.code_provider.execute_shell(git_add_cmd, 30, workdir)
|
640
|
+
|
641
|
+
if add_result.get("exit_code", 1) != 0:
|
642
|
+
return {
|
643
|
+
"stdout": "",
|
644
|
+
"stderr": f"Failed to stage changes: {add_result.get('stderr', '')}"
|
645
|
+
}
|
646
|
+
|
647
|
+
# Create commit with command description and timestamp
|
648
|
+
timestamp = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
|
649
|
+
commit_msg = f"Checkpoint: {description} @ {timestamp}\n\nCommand: {cmd_str}"
|
650
|
+
git_commit_cmd = ["git", "commit", "-m", commit_msg, "--no-gpg-sign"]
|
651
|
+
commit_result = await self.code_provider.execute_shell(git_commit_cmd, 30, workdir)
|
652
|
+
|
653
|
+
if commit_result.get("exit_code", 1) != 0:
|
654
|
+
return {
|
655
|
+
"stdout": "",
|
656
|
+
"stderr": f"Failed to create commit: {commit_result.get('stderr', '')}"
|
657
|
+
}
|
658
|
+
|
659
|
+
# Get the first line of the commit message without using split with \n in f-string
|
660
|
+
first_line = commit_msg.split("\n")[0]
|
661
|
+
return {
|
662
|
+
"stdout": f"✓ Git checkpoint created: {first_line}",
|
663
|
+
"stderr": ""
|
664
|
+
}
|
665
|
+
except Exception as e:
|
666
|
+
return {
|
667
|
+
"stdout": "",
|
668
|
+
"stderr": f"Error creating git checkpoint: {str(e)}"
|
669
|
+
}
|
670
|
+
|
410
671
|
def set_default_workdir(self, workdir: str, create_if_not_exists: bool = False):
|
411
672
|
"""
|
412
673
|
Set the default working directory for shell commands.
|
@@ -676,6 +937,17 @@ class TinyCodeAgent:
|
|
676
937
|
"""
|
677
938
|
return self.authorized_imports.copy()
|
678
939
|
|
940
|
+
@classmethod
|
941
|
+
def is_seatbelt_supported(cls) -> bool:
|
942
|
+
"""
|
943
|
+
Check if the seatbelt provider is supported on this system.
|
944
|
+
|
945
|
+
Returns:
|
946
|
+
True if seatbelt is supported (macOS with sandbox-exec), False otherwise
|
947
|
+
"""
|
948
|
+
from .providers.seatbelt_provider import SeatbeltProvider
|
949
|
+
return SeatbeltProvider.is_supported()
|
950
|
+
|
679
951
|
def remove_authorized_import(self, import_name: str):
|
680
952
|
"""
|
681
953
|
Remove an authorized import.
|
@@ -777,8 +1049,13 @@ class TinyCodeAgent:
|
|
777
1049
|
)
|
778
1050
|
self.add_callback(ui_callback)
|
779
1051
|
elif ui_type == 'jupyter':
|
1052
|
+
if not JUPYTER_CALLBACKS_AVAILABLE:
|
1053
|
+
raise ImportError(
|
1054
|
+
"Jupyter notebook callbacks are not available. "
|
1055
|
+
"Install the required dependencies with: pip install ipython ipywidgets"
|
1056
|
+
)
|
1057
|
+
|
780
1058
|
if optimized:
|
781
|
-
from tinyagent.hooks.jupyter_notebook_callback import OptimizedJupyterNotebookCallback
|
782
1059
|
ui_callback = OptimizedJupyterNotebookCallback(
|
783
1060
|
logger=self.log_manager.get_logger('tinyagent.hooks.jupyter_notebook_callback') if self.log_manager else None,
|
784
1061
|
max_visible_turns=20, # Limit visible turns for performance
|
@@ -792,7 +1069,123 @@ class TinyCodeAgent:
|
|
792
1069
|
)
|
793
1070
|
self.add_callback(ui_callback)
|
794
1071
|
else:
|
795
|
-
self.log_manager
|
1072
|
+
if self.log_manager:
|
1073
|
+
self.log_manager.get_logger(__name__).warning(f"Unknown UI type: {ui_type}. No UI callback will be added.")
|
1074
|
+
else:
|
1075
|
+
print(f"Warning: Unknown UI type: {ui_type}. No UI callback will be added.")
|
1076
|
+
|
1077
|
+
def set_truncation_config(self, config: Dict[str, Any]):
|
1078
|
+
"""
|
1079
|
+
Set the truncation configuration.
|
1080
|
+
|
1081
|
+
Args:
|
1082
|
+
config: Dictionary containing truncation configuration options:
|
1083
|
+
- max_tokens: Maximum number of tokens to keep in output
|
1084
|
+
- max_lines: Maximum number of lines to keep in output
|
1085
|
+
- enabled: Whether truncation is enabled
|
1086
|
+
"""
|
1087
|
+
self.truncation_config.update(config)
|
1088
|
+
|
1089
|
+
def get_truncation_config(self) -> Dict[str, Any]:
|
1090
|
+
"""
|
1091
|
+
Get the current truncation configuration.
|
1092
|
+
|
1093
|
+
Returns:
|
1094
|
+
Dictionary containing truncation configuration
|
1095
|
+
"""
|
1096
|
+
return self.truncation_config.copy()
|
1097
|
+
|
1098
|
+
def enable_truncation(self, enabled: bool = True):
|
1099
|
+
"""
|
1100
|
+
Enable or disable output truncation.
|
1101
|
+
|
1102
|
+
Args:
|
1103
|
+
enabled: Whether to enable truncation
|
1104
|
+
"""
|
1105
|
+
self.truncation_config["enabled"] = enabled
|
1106
|
+
|
1107
|
+
def enable_auto_git_checkpoint(self, enabled: bool = True):
|
1108
|
+
"""
|
1109
|
+
Enable or disable automatic git checkpoint creation after successful shell commands.
|
1110
|
+
|
1111
|
+
Args:
|
1112
|
+
enabled: If True, automatically create git checkpoints. If False, do not create them.
|
1113
|
+
"""
|
1114
|
+
self.auto_git_checkpoint = enabled
|
1115
|
+
|
1116
|
+
def get_auto_git_checkpoint_status(self) -> bool:
|
1117
|
+
"""
|
1118
|
+
Get the current status of auto_git_checkpoint.
|
1119
|
+
|
1120
|
+
Returns:
|
1121
|
+
True if auto_git_checkpoint is enabled, False otherwise.
|
1122
|
+
"""
|
1123
|
+
return self.auto_git_checkpoint
|
1124
|
+
|
1125
|
+
def set_environment_variables(self, env_vars: Dict[str, str]):
|
1126
|
+
"""
|
1127
|
+
Set environment variables for the code execution provider.
|
1128
|
+
Currently only supported for SeatbeltProvider.
|
1129
|
+
|
1130
|
+
Args:
|
1131
|
+
env_vars: Dictionary of environment variable name -> value pairs
|
1132
|
+
|
1133
|
+
Raises:
|
1134
|
+
AttributeError: If the provider doesn't support environment variables
|
1135
|
+
"""
|
1136
|
+
if hasattr(self.code_provider, 'set_environment_variables'):
|
1137
|
+
self.code_provider.set_environment_variables(env_vars)
|
1138
|
+
else:
|
1139
|
+
raise AttributeError(f"Provider {self.provider} does not support environment variables")
|
1140
|
+
|
1141
|
+
def add_environment_variable(self, name: str, value: str):
|
1142
|
+
"""
|
1143
|
+
Add a single environment variable for the code execution provider.
|
1144
|
+
Currently only supported for SeatbeltProvider.
|
1145
|
+
|
1146
|
+
Args:
|
1147
|
+
name: Environment variable name
|
1148
|
+
value: Environment variable value
|
1149
|
+
|
1150
|
+
Raises:
|
1151
|
+
AttributeError: If the provider doesn't support environment variables
|
1152
|
+
"""
|
1153
|
+
if hasattr(self.code_provider, 'add_environment_variable'):
|
1154
|
+
self.code_provider.add_environment_variable(name, value)
|
1155
|
+
else:
|
1156
|
+
raise AttributeError(f"Provider {self.provider} does not support environment variables")
|
1157
|
+
|
1158
|
+
def remove_environment_variable(self, name: str):
|
1159
|
+
"""
|
1160
|
+
Remove an environment variable from the code execution provider.
|
1161
|
+
Currently only supported for SeatbeltProvider.
|
1162
|
+
|
1163
|
+
Args:
|
1164
|
+
name: Environment variable name to remove
|
1165
|
+
|
1166
|
+
Raises:
|
1167
|
+
AttributeError: If the provider doesn't support environment variables
|
1168
|
+
"""
|
1169
|
+
if hasattr(self.code_provider, 'remove_environment_variable'):
|
1170
|
+
self.code_provider.remove_environment_variable(name)
|
1171
|
+
else:
|
1172
|
+
raise AttributeError(f"Provider {self.provider} does not support environment variables")
|
1173
|
+
|
1174
|
+
def get_environment_variables(self) -> Dict[str, str]:
|
1175
|
+
"""
|
1176
|
+
Get a copy of current environment variables from the code execution provider.
|
1177
|
+
Currently only supported for SeatbeltProvider.
|
1178
|
+
|
1179
|
+
Returns:
|
1180
|
+
Dictionary of current environment variables
|
1181
|
+
|
1182
|
+
Raises:
|
1183
|
+
AttributeError: If the provider doesn't support environment variables
|
1184
|
+
"""
|
1185
|
+
if hasattr(self.code_provider, 'get_environment_variables'):
|
1186
|
+
return self.code_provider.get_environment_variables()
|
1187
|
+
else:
|
1188
|
+
raise AttributeError(f"Provider {self.provider} does not support environment variables")
|
796
1189
|
|
797
1190
|
|
798
1191
|
# Example usage demonstrating both LLM tools and code tools
|
@@ -836,7 +1229,12 @@ async def run_example():
|
|
836
1229
|
authorized_imports=["tinyagent", "gradio", "requests", "numpy", "pandas"], # Explicitly specify authorized imports
|
837
1230
|
local_execution=False, # Remote execution via Modal (default)
|
838
1231
|
check_string_obfuscation=True,
|
839
|
-
default_workdir=os.path.join(os.getcwd(), "examples") # Set a default working directory for shell commands
|
1232
|
+
default_workdir=os.path.join(os.getcwd(), "examples"), # Set a default working directory for shell commands
|
1233
|
+
truncation_config={
|
1234
|
+
"max_tokens": 3000,
|
1235
|
+
"max_lines": 250,
|
1236
|
+
"enabled": True
|
1237
|
+
}
|
840
1238
|
)
|
841
1239
|
|
842
1240
|
# Connect to MCP servers
|
@@ -961,6 +1359,292 @@ async def run_example():
|
|
961
1359
|
print("Shell Execution with Explicit Working Directory:")
|
962
1360
|
print(response_shell_explicit)
|
963
1361
|
|
1362
|
+
# Test truncation functionality
|
1363
|
+
print("\n" + "="*80)
|
1364
|
+
print("✂️ Testing output truncation")
|
1365
|
+
|
1366
|
+
# Configure truncation with smaller limits for testing
|
1367
|
+
agent_remote.set_truncation_config({
|
1368
|
+
"max_tokens": 100, # Very small limit for testing
|
1369
|
+
"max_lines": 5 # Very small limit for testing
|
1370
|
+
})
|
1371
|
+
|
1372
|
+
# Generate a large output to test truncation
|
1373
|
+
large_output_prompt = """
|
1374
|
+
Generate a large output by printing a lot of text. Create a Python script that:
|
1375
|
+
1. Prints numbers from 1 to 1000
|
1376
|
+
2. For each number, also print its square and cube
|
1377
|
+
3. Add random text for each line to make it longer
|
1378
|
+
"""
|
1379
|
+
|
1380
|
+
response_truncated = await agent_remote.run(large_output_prompt)
|
1381
|
+
print("Truncated Output Response:")
|
1382
|
+
print(response_truncated)
|
1383
|
+
|
1384
|
+
# Test disabling truncation
|
1385
|
+
print("\n" + "="*80)
|
1386
|
+
print("🔄 Testing with truncation disabled")
|
1387
|
+
|
1388
|
+
agent_remote.enable_truncation(False)
|
1389
|
+
response_untruncated = await agent_remote.run("Run the same script again but limit to 20 numbers")
|
1390
|
+
print("Untruncated Output Response:")
|
1391
|
+
print(response_untruncated)
|
1392
|
+
|
1393
|
+
# Test git checkpoint functionality
|
1394
|
+
print("\n" + "="*80)
|
1395
|
+
print("🔄 Testing git checkpoint functionality")
|
1396
|
+
|
1397
|
+
# Enable git checkpoints
|
1398
|
+
agent_remote.enable_auto_git_checkpoint(True)
|
1399
|
+
print(f"Auto Git Checkpoint enabled: {agent_remote.get_auto_git_checkpoint_status()}")
|
1400
|
+
|
1401
|
+
# Create a test file to demonstrate git checkpoint
|
1402
|
+
git_test_prompt = """
|
1403
|
+
Create a new file called test_file.txt with some content, then modify it, and observe
|
1404
|
+
that git checkpoints are created automatically after each change.
|
1405
|
+
"""
|
1406
|
+
|
1407
|
+
git_response = await agent_remote.run(git_test_prompt)
|
1408
|
+
print("Git Checkpoint Response:")
|
1409
|
+
print(git_response)
|
1410
|
+
|
1411
|
+
# Disable git checkpoints
|
1412
|
+
agent_remote.enable_auto_git_checkpoint(False)
|
1413
|
+
print(f"Auto Git Checkpoint disabled: {agent_remote.get_auto_git_checkpoint_status()}")
|
1414
|
+
|
1415
|
+
# Test seatbelt provider if supported
|
1416
|
+
if TinyCodeAgent.is_seatbelt_supported():
|
1417
|
+
print("\n" + "="*80)
|
1418
|
+
print("🔒 Testing TinyCodeAgent with SEATBELT provider (sandboxed execution)")
|
1419
|
+
|
1420
|
+
# Create a test directory for read/write access
|
1421
|
+
test_read_dir = os.path.join(os.getcwd(), "test_read_dir")
|
1422
|
+
test_write_dir = os.path.join(os.getcwd(), "test_write_dir")
|
1423
|
+
|
1424
|
+
# Create directories if they don't exist
|
1425
|
+
os.makedirs(test_read_dir, exist_ok=True)
|
1426
|
+
os.makedirs(test_write_dir, exist_ok=True)
|
1427
|
+
|
1428
|
+
# Create a test file in the read directory
|
1429
|
+
with open(os.path.join(test_read_dir, "test.txt"), "w") as f:
|
1430
|
+
f.write("This is a test file for reading")
|
1431
|
+
|
1432
|
+
# Create a simple seatbelt profile
|
1433
|
+
seatbelt_profile = """(version 1)
|
1434
|
+
|
1435
|
+
; Default to deny everything
|
1436
|
+
(deny default)
|
1437
|
+
|
1438
|
+
; Allow network connections with proper DNS resolution
|
1439
|
+
(allow network*)
|
1440
|
+
(allow network-outbound)
|
1441
|
+
(allow mach-lookup)
|
1442
|
+
|
1443
|
+
; Allow process execution
|
1444
|
+
(allow process-exec)
|
1445
|
+
(allow process-fork)
|
1446
|
+
(allow signal (target self))
|
1447
|
+
|
1448
|
+
; Restrict file read to current path and system files
|
1449
|
+
(deny file-read* (subpath "/Users"))
|
1450
|
+
(allow file-read*
|
1451
|
+
(subpath "{os.getcwd()}")
|
1452
|
+
(subpath "/usr")
|
1453
|
+
(subpath "/System")
|
1454
|
+
(subpath "/Library")
|
1455
|
+
(subpath "/bin")
|
1456
|
+
(subpath "/sbin")
|
1457
|
+
(subpath "/opt")
|
1458
|
+
(subpath "/private/tmp")
|
1459
|
+
(subpath "/private/var/tmp")
|
1460
|
+
(subpath "/dev")
|
1461
|
+
(subpath "/etc")
|
1462
|
+
(literal "/")
|
1463
|
+
(literal "/."))
|
1464
|
+
|
1465
|
+
; Allow write access to specified folder and temp directories
|
1466
|
+
(deny file-write* (subpath "/"))
|
1467
|
+
(allow file-write*
|
1468
|
+
(subpath "{os.getcwd()}")
|
1469
|
+
(subpath "/private/tmp")
|
1470
|
+
(subpath "/private/var/tmp")
|
1471
|
+
(subpath "/dev"))
|
1472
|
+
|
1473
|
+
; Allow standard device operations
|
1474
|
+
(allow file-write-data
|
1475
|
+
(literal "/dev/null")
|
1476
|
+
(literal "/dev/dtracehelper")
|
1477
|
+
(literal "/dev/tty")
|
1478
|
+
(literal "/dev/stdout")
|
1479
|
+
(literal "/dev/stderr"))
|
1480
|
+
|
1481
|
+
; Allow iokit operations needed for system functions
|
1482
|
+
(allow iokit-open)
|
1483
|
+
|
1484
|
+
; Allow shared memory operations
|
1485
|
+
(allow ipc-posix-shm)
|
1486
|
+
|
1487
|
+
; Allow basic system operations
|
1488
|
+
(allow file-read-metadata)
|
1489
|
+
(allow process-info-pidinfo)
|
1490
|
+
(allow process-info-setcontrol)
|
1491
|
+
"""
|
1492
|
+
|
1493
|
+
# Create TinyCodeAgent with seatbelt provider
|
1494
|
+
agent_seatbelt = TinyCodeAgent(
|
1495
|
+
model="gpt-4.1-mini",
|
1496
|
+
tools=[search_web], # LLM tools
|
1497
|
+
code_tools=[data_processor], # Code tools
|
1498
|
+
user_variables={
|
1499
|
+
"sample_data": [1, 2, 3, 4, 5, 10, 15, 20]
|
1500
|
+
},
|
1501
|
+
provider="seatbelt", # Use seatbelt provider
|
1502
|
+
provider_config={
|
1503
|
+
"seatbelt_profile": seatbelt_profile,
|
1504
|
+
# Alternatively, you can specify a path to a seatbelt profile file:
|
1505
|
+
# "seatbelt_profile_path": "/path/to/seatbelt.sb",
|
1506
|
+
# "python_env_path": "/path/to/python/env", # Optional path to Python environment
|
1507
|
+
|
1508
|
+
# Specify additional directories for read/write access
|
1509
|
+
"additional_read_dirs": [test_read_dir],
|
1510
|
+
"additional_write_dirs": [test_write_dir],
|
1511
|
+
|
1512
|
+
# Allow git commands
|
1513
|
+
"bypass_shell_safety": True,
|
1514
|
+
"additional_safe_shell_commands": ["git"],
|
1515
|
+
|
1516
|
+
# Environment variables to make available in the sandbox
|
1517
|
+
"environment_variables": {
|
1518
|
+
"TEST_READ_DIR": test_read_dir,
|
1519
|
+
"TEST_WRITE_DIR": test_write_dir,
|
1520
|
+
"PROJECT_NAME": "TinyAgent Seatbelt Demo",
|
1521
|
+
"BUILD_VERSION": "1.0.0"
|
1522
|
+
}
|
1523
|
+
},
|
1524
|
+
local_execution=True, # Required for seatbelt
|
1525
|
+
check_string_obfuscation=True,
|
1526
|
+
truncation_config={
|
1527
|
+
"max_tokens": 500,
|
1528
|
+
"max_lines": 20,
|
1529
|
+
"enabled": True
|
1530
|
+
}
|
1531
|
+
)
|
1532
|
+
|
1533
|
+
# Connect to MCP servers
|
1534
|
+
await agent_seatbelt.connect_to_server("npx", ["-y", "@openbnb/mcp-server-airbnb", "--ignore-robots-txt"])
|
1535
|
+
await agent_seatbelt.connect_to_server("npx", ["-y", "@modelcontextprotocol/server-sequential-thinking"])
|
1536
|
+
|
1537
|
+
# Test the seatbelt agent
|
1538
|
+
response_seatbelt = await agent_seatbelt.run("""
|
1539
|
+
I have some sample data. Please use the data_processor tool in Python to analyze my sample_data
|
1540
|
+
and show me the results.
|
1541
|
+
""")
|
1542
|
+
|
1543
|
+
print("Seatbelt Agent Response:")
|
1544
|
+
print(response_seatbelt)
|
1545
|
+
|
1546
|
+
# Test shell execution in sandbox
|
1547
|
+
shell_prompt_sandbox = "Run 'ls -la' to list files in the current directory."
|
1548
|
+
|
1549
|
+
response_shell_sandbox = await agent_seatbelt.run(shell_prompt_sandbox)
|
1550
|
+
print("Shell Execution in Sandbox:")
|
1551
|
+
print(response_shell_sandbox)
|
1552
|
+
|
1553
|
+
# Test reading from the additional read directory
|
1554
|
+
read_prompt = f"Read the contents of the file in the test_read_dir directory."
|
1555
|
+
|
1556
|
+
response_read = await agent_seatbelt.run(read_prompt)
|
1557
|
+
print("Reading from Additional Read Directory:")
|
1558
|
+
print(response_read)
|
1559
|
+
|
1560
|
+
# Test writing to the additional write directory
|
1561
|
+
write_prompt = f"Write a file called 'output.txt' with the text 'Hello from sandbox!' in the test_write_dir directory."
|
1562
|
+
|
1563
|
+
response_write = await agent_seatbelt.run(write_prompt)
|
1564
|
+
print("Writing to Additional Write Directory:")
|
1565
|
+
print(response_write)
|
1566
|
+
|
1567
|
+
# Test environment variables
|
1568
|
+
print("\n" + "="*80)
|
1569
|
+
print("🔧 Testing environment variables functionality")
|
1570
|
+
|
1571
|
+
# Add additional environment variables dynamically
|
1572
|
+
agent_seatbelt.add_environment_variable("CUSTOM_VAR", "custom_value")
|
1573
|
+
agent_seatbelt.add_environment_variable("DEBUG_MODE", "true")
|
1574
|
+
|
1575
|
+
# Get and display current environment variables
|
1576
|
+
current_env_vars = agent_seatbelt.get_environment_variables()
|
1577
|
+
print(f"Current environment variables: {list(current_env_vars.keys())}")
|
1578
|
+
|
1579
|
+
# Test accessing environment variables in Python and shell
|
1580
|
+
env_test_prompt = """
|
1581
|
+
Test the environment variables we set:
|
1582
|
+
1. In Python, use os.environ to check for CUSTOM_VAR and DEBUG_MODE
|
1583
|
+
2. In a shell command, use 'echo $CUSTOM_VAR' and 'echo $DEBUG_MODE'
|
1584
|
+
3. Also check the TEST_READ_DIR and TEST_WRITE_DIR variables that were set during initialization
|
1585
|
+
4. Show all environment variables that start with 'TEST_' or 'CUSTOM_' or 'DEBUG_'
|
1586
|
+
"""
|
1587
|
+
|
1588
|
+
response_env_test = await agent_seatbelt.run(env_test_prompt)
|
1589
|
+
print("Environment Variables Test:")
|
1590
|
+
print(response_env_test)
|
1591
|
+
|
1592
|
+
# Update environment variables
|
1593
|
+
agent_seatbelt.set_environment_variables({
|
1594
|
+
"CUSTOM_VAR": "updated_value",
|
1595
|
+
"NEW_VAR": "new_value",
|
1596
|
+
"API_KEY": "test_api_key_123"
|
1597
|
+
})
|
1598
|
+
|
1599
|
+
# Test updated environment variables
|
1600
|
+
updated_env_test_prompt = """
|
1601
|
+
Test the updated environment variables:
|
1602
|
+
1. Check that CUSTOM_VAR now has the value 'updated_value'
|
1603
|
+
2. Check that NEW_VAR is available with value 'new_value'
|
1604
|
+
3. Check that API_KEY is available with value 'test_api_key_123'
|
1605
|
+
4. Verify that DEBUG_MODE is no longer available (should have been removed by set operation)
|
1606
|
+
"""
|
1607
|
+
|
1608
|
+
response_updated_env = await agent_seatbelt.run(updated_env_test_prompt)
|
1609
|
+
print("Updated Environment Variables Test:")
|
1610
|
+
print(response_updated_env)
|
1611
|
+
|
1612
|
+
# Remove a specific environment variable
|
1613
|
+
agent_seatbelt.remove_environment_variable("API_KEY")
|
1614
|
+
|
1615
|
+
# Test that the removed variable is no longer available
|
1616
|
+
removed_env_test_prompt = """
|
1617
|
+
Test that API_KEY environment variable has been removed:
|
1618
|
+
1. Try to access API_KEY in Python - it should not be available
|
1619
|
+
2. Use shell command 'echo $API_KEY' - it should be empty
|
1620
|
+
3. List all current environment variables that start with 'CUSTOM_' or 'NEW_'
|
1621
|
+
"""
|
1622
|
+
|
1623
|
+
response_removed_env = await agent_seatbelt.run(removed_env_test_prompt)
|
1624
|
+
print("Removed Environment Variable Test:")
|
1625
|
+
print(response_removed_env)
|
1626
|
+
|
1627
|
+
# Test git commands with the custom configuration
|
1628
|
+
git_prompt = "Run 'git status' to show the current git status."
|
1629
|
+
|
1630
|
+
response_git = await agent_seatbelt.run(git_prompt)
|
1631
|
+
print("Git Command Execution:")
|
1632
|
+
print(response_git)
|
1633
|
+
|
1634
|
+
# Clean up test directories
|
1635
|
+
import shutil
|
1636
|
+
try:
|
1637
|
+
shutil.rmtree(test_read_dir)
|
1638
|
+
shutil.rmtree(test_write_dir)
|
1639
|
+
print("Cleaned up test directories")
|
1640
|
+
except Exception as e:
|
1641
|
+
print(f"Error cleaning up test directories: {str(e)}")
|
1642
|
+
|
1643
|
+
await agent_seatbelt.close()
|
1644
|
+
else:
|
1645
|
+
print("\n" + "="*80)
|
1646
|
+
print("⚠️ Seatbelt provider is not supported on this system. Skipping seatbelt tests.")
|
1647
|
+
|
964
1648
|
await agent_remote.close()
|
965
1649
|
await agent_local.close()
|
966
1650
|
|