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

tinyagent/code_agent/utils.py

@@ -2,9 +2,12 @@ import sys
 import cloudpickle
 import subprocess
 import os
-from typing import Dict, Any, List
+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]:
@@ -20,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.
@@ -69,8 +184,31 @@ def _run_shell(
         
         # 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
-            if command[0] == "bash" and len(command) >= 3 and command[1] in ["-c", "-lc"]:
+            # 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
@@ -80,33 +218,60 @@ def _run_shell(
                     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)
+            # 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(
-                    ["bash", "-c", shell_command],
-                    shell=False,  # Using explicit bash -c instead of shell=True
+                    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 {

tinyagent/prompts/truncation.yaml

@@ -0,0 +1,13 @@
+truncation_messages:
+  python_output:
+    message: |-
+      ---
+      **Output Truncated**: The original output was {original_size} {size_unit} ({original_lines} lines). Showing only the first {max_lines} lines.
+      To get more detailed output, please make your request more specific or adjust the output size.
+      ---
+  bash_output:
+    message: |-
+      ---
+      **Output Truncated**: The original output was {original_size} {size_unit} ({original_lines} lines). Showing only the first {max_lines} lines.
+      To get more detailed output, please use more specific commands or add filtering.
+      ---