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.
@@ -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
- from tinyagent.hooks.jupyter_notebook_callback import JupyterNotebookCallback
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, Docker, local execution, etc.).
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 exceeds 30000 characters, output will be truncated before being returned to you.
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: 30).
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 = 30) -> str:
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.get_logger(__name__).warning(f"Unknown UI type: {ui_type}. No UI callback will be added.")
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