tinyagent-py 0.0.13__py3-none-any.whl → 0.0.16__py3-none-any.whl

tinyagent/code_agent/utils.py

@@ -1,7 +1,13 @@
 import sys
 import cloudpickle
-from typing import Dict, Any, List
+import subprocess
+import os
+from typing import Dict, Any, List, Tuple
 from .safety import validate_code_safety, function_safety_context
+import shlex
+import yaml
+from pathlib import Path
+import re
 
 
 def clean_response(resp: Dict[str, Any]) -> Dict[str, Any]:
@@ -17,6 +23,118 @@ def clean_response(resp: Dict[str, Any]) -> Dict[str, Any]:
     return {k: v for k, v in resp.items() if k in ['printed_output', 'return_value', 'stderr', 'error_traceback']}
 
 
+def truncate_output(output: str, max_tokens: int = 3000, max_lines: int = 250) -> Tuple[str, bool, int, int]:
+    """
+    Truncate output based on token count and line count.
+    
+    Args:
+        output: The output string to truncate
+        max_tokens: Maximum number of tokens to keep
+        max_lines: Maximum number of lines to keep
+        
+    Returns:
+        Tuple containing:
+        - Truncated output
+        - Boolean indicating if truncation occurred
+        - Original token count
+        - Original line count
+    """
+    # Count original lines
+    lines = output.splitlines()
+    original_line_count = len(lines)
+    
+    # Approximate token count (rough estimate: 4 chars ≈ 1 token)
+    original_token_count = len(output) // 4
+    
+    # Check if truncation is needed
+    if original_line_count <= max_lines and original_token_count <= max_tokens:
+        return output, False, original_token_count, original_line_count
+    
+    # Truncate by lines first
+    if original_line_count > max_lines:
+        lines = lines[:max_lines]  # Keep only the first max_lines
+    
+    # Join lines back together
+    truncated = '\n'.join(lines)
+    
+    # If still too many tokens, truncate further
+    if len(truncated) // 4 > max_tokens:
+        # Keep the first max_tokens*4 characters (approximate)
+        truncated = truncated[:max_tokens*4]
+        
+        # Try to start at a newline to avoid partial lines
+        newline_pos = truncated.find('\n')
+        if newline_pos > 0:
+            truncated = truncated[newline_pos+1:]
+    
+    return truncated, True, original_token_count, original_line_count
+
+
+def load_truncation_template(template_type: str = "python_output") -> str:
+    """
+    Load the truncation message template.
+    
+    Args:
+        template_type: Type of template to load ("python_output" or "bash_output")
+        
+    Returns:
+        Template string for the truncation message
+    """
+    template_path = Path(__file__).parent.parent / "prompts" / "truncation.yaml"
+    
+    try:
+        with open(template_path, 'r') as f:
+            templates = yaml.safe_load(f)
+        
+        return templates.get("truncation_messages", {}).get(template_type, {}).get("message", 
+            "--- Output truncated due to size limitations ---")
+    except Exception:
+        # Fallback template if file can't be loaded
+        return "--- Output truncated due to size limitations ---"
+
+
+def format_truncation_message(output: str, is_truncated: bool, original_tokens: int, 
+                             original_lines: int, max_lines: int, template_type: str = "python_output") -> str:
+    """
+    Format the truncated output with a truncation message if needed.
+    
+    Args:
+        output: The truncated output
+        is_truncated: Whether truncation occurred
+        original_tokens: Original token count
+        original_lines: Original line count
+        max_lines: Maximum line count used for truncation
+        template_type: Type of template to use
+        
+    Returns:
+        Formatted output with truncation message if needed
+    """
+    if not is_truncated:
+        return output
+    
+    # Load the appropriate template
+    template = load_truncation_template(template_type)
+    
+    # Determine size unit (tokens or KB)
+    if original_tokens > 1000:
+        size_value = original_tokens / 1000
+        size_unit = "K tokens"
+    else:
+        size_value = original_tokens
+        size_unit = "tokens"
+    
+    # Format the message
+    message = template.format(
+        original_size=round(size_value, 1),
+        size_unit=size_unit,
+        original_lines=original_lines,
+        max_lines=max_lines
+    )
+    
+    # Append the message to the output
+    return f"{output}\n\n{message}"
+
+
 def make_session_blob(ns: dict) -> bytes:
     """
     Create a serialized blob of the session namespace, excluding unserializable objects.
@@ -41,6 +159,146 @@ def make_session_blob(ns: dict) -> bytes:
     return cloudpickle.dumps(clean)
 
 
+def _run_shell(
+    command: List[str],
+    timeout: int = 10,
+    workdir: str = None
+) -> Dict[str, Any]:
+    """
+    Execute a shell command securely with proper timeout and error handling.
+    
+    Args:
+        command: List of command parts to execute
+        timeout: Maximum execution time in seconds
+        workdir: Working directory for command execution
+        
+    Returns:
+        Dictionary containing execution results with keys:
+        - stdout: stdout from the execution
+        - stderr: stderr from the execution
+        - exit_code: exit code from the command
+    """
+    try:
+        # Set working directory if provided
+        cwd = os.path.expanduser(workdir) if workdir else None
+        
+        # Check if this is a command that needs bash -c wrapping
+        if len(command) > 0:
+            # Special handling for bash login shells to avoid profile loading errors
+            if command[0] == "bash" and len(command) >= 3 and command[1] == "-lc":
+                # Create a clean environment that doesn't load user profile files
+                env = os.environ.copy()
+                env.update({
+                    "BASH_ENV": "/dev/null",
+                    "ENV": "/dev/null",
+                    "BASH_PROFILE": "/dev/null",
+                    "PROFILE": "/dev/null"
+                })
+                # Replace -lc with -c to avoid loading login profiles
+                modified_command = ["bash", "-c", command[2]]
+                process = subprocess.run(
+                    modified_command,
+                    shell=False,
+                    capture_output=True,
+                    text=True,
+                    timeout=timeout,
+                    cwd=cwd,
+                    check=False,
+                    env=env
+                )
+            # If the command already uses bash -c, use it directly
+            # This handles heredoc syntax and other complex shell constructs
+            elif command[0] == "bash" and len(command) >= 3 and command[1] == "-c":
+                process = subprocess.run(
+                    command,
+                    shell=False,  # No need for shell=True as we're explicitly using bash -c
+                    capture_output=True,
+                    text=True,
+                    timeout=timeout,
+                    cwd=cwd,
+                    check=False
+                )
+            # Special handling for interpreter commands with inline code execution flags
+            # This covers python -c, node -e, ruby -e, perl -e, etc.
+            elif len(command) >= 3 and command[0] in ["python", "node", "ruby", "perl", "php", "deno"] and command[1] in ["-c", "-e", "--eval", "--execute"]:
+                # Execute the interpreter command directly without shell wrapping
+                process = subprocess.run(
+                    command,
+                    shell=False,
+                    capture_output=True,
+                    text=True,
+                    timeout=timeout,
+                    cwd=cwd,
+                    check=False
+                )
+            else:
+                # Check if the command contains heredoc syntax
+                command_str = " ".join(command)
+                if "<<" in command_str and any(f"<<'{token}'" in command_str or f'<<"{token}"' in command_str or f"<<{token}" in command_str for token in ["EOF", "EOL", "END", "HEREDOC", "PY", "JS", "RUBY", "PHP"]):
+                    # For commands with heredoc, pass directly to bash -c without additional quoting
+                    process = subprocess.run(
+                        ["bash", "-c", command_str],
+                        shell=False,
+                        capture_output=True,
+                        text=True,
+                        timeout=timeout,
+                        cwd=cwd,
+                        check=False
+                    )
+                else:
+                    # For all other commands, wrap in bash -c to handle shell operators
+                    # and properly quote arguments that need quoting
+                    
+                    # Shell operators that should not be quoted
+                    shell_operators = ['|', '&&', '||', '>', '<', '>>', '<<', ';']
+                    
+                    # Quote each part that needs quoting
+                    quoted_parts = []
+                    for part in command:
+                        if part in shell_operators:
+                            # Don't quote shell operators
+                            quoted_parts.append(part)
+                        else:
+                            # Use shlex.quote to properly escape special characters
+                            quoted_parts.append(shlex.quote(part))
+                    
+                    shell_command = " ".join(quoted_parts)
+                    process = subprocess.run(
+                        ["bash", "-c", shell_command],
+                        shell=False,  # Using explicit bash -c instead of shell=True
+                        capture_output=True,
+                        text=True,
+                        timeout=timeout,
+                        cwd=cwd,
+                        check=False
+                    )
+        else:
+            # Empty command
+            return {
+                "stdout": "",
+                "stderr": "Empty command",
+                "exit_code": 1
+            }
+        
+        return {
+            "stdout": process.stdout,
+            "stderr": process.stderr,
+            "exit_code": process.returncode
+        }
+    except subprocess.TimeoutExpired:
+        return {
+            "stdout": "",
+            "stderr": f"Command timed out after {timeout} seconds",
+            "exit_code": 124  # Standard timeout exit code
+        }
+    except Exception as e:
+        return {
+            "stdout": "",
+            "stderr": f"Error executing command: {str(e)}",
+            "exit_code": 1
+        }
+
+
 def _run_python(
     code: str,
     globals_dict: Dict[str, Any] | None = None,
@@ -48,6 +306,7 @@ def _run_python(
     authorized_imports: List[str] | None = None,
     authorized_functions: List[str] | None = None,
     trusted_code: bool = False,
+    check_string_obfuscation: bool = True,
 ):
     """
     Execute Python code in a controlled environment with proper error handling.
@@ -59,6 +318,7 @@ def _run_python(
         authorized_imports: List of authorized imports that user code may access. Wildcards (e.g. "numpy.*") are supported. A value of None disables the allow-list and only blocks dangerous modules.
         authorized_functions: List of authorized dangerous functions that user code may access. A value of None disables the allow-list and blocks all dangerous functions.
         trusted_code: If True, skip security checks. Should only be used for framework code, tools, or default executed code.
+        check_string_obfuscation: If True (default), check for string obfuscation techniques. Set to False to allow legitimate use of base64 encoding and other string manipulations.
         
     Returns:
         Dictionary containing execution results
@@ -74,7 +334,8 @@ def _run_python(
     # 1. Static safety analysis – refuse code containing dangerous imports or functions
     # ------------------------------------------------------------------
     validate_code_safety(code, authorized_imports=authorized_imports, 
-                        authorized_functions=authorized_functions, trusted_code=trusted_code)
+                        authorized_functions=authorized_functions, trusted_code=trusted_code,
+                        check_string_obfuscation=check_string_obfuscation)
 
     # Make copies to avoid mutating the original parameters
     globals_dict = globals_dict or {}

tinyagent/hooks/__init__.py

@@ -2,4 +2,6 @@
 from .rich_ui_callback import RichUICallback
 from .rich_code_ui_callback import RichCodeUICallback
 from .logging_manager import LoggingManager
-__all__ = ["RichUICallback", "RichCodeUICallback", "LoggingManager"]
+from .token_tracker import TokenTracker, UsageStats, create_token_tracker
+
+__all__ = ["RichUICallback", "RichCodeUICallback", "LoggingManager", "TokenTracker", "UsageStats", "create_token_tracker"]