quantalogic 0.59.3__py3-none-any.whl → 0.61.0__py3-none-any.whl

quantalogic/tools/action_gen.py

@@ -0,0 +1,366 @@
+import ast
+import asyncio
+from asyncio import TimeoutError
+from contextlib import AsyncExitStack
+from functools import partial
+from typing import Callable, Dict, List
+
+import litellm
+import typer
+from loguru import logger
+
+from quantalogic.python_interpreter import execute_async
+from quantalogic.tools.tool import Tool, ToolArgument
+
+# Configure loguru to log to a file with rotation, matching original
+logger.add("action_gen.log", rotation="10 MB", level="DEBUG")
+
+# Initialize Typer app, unchanged
+app = typer.Typer()
+
+# Define tool classes with logging in async_execute, preserving original structure
+class AddTool(Tool):
+    def __init__(self):
+        super().__init__(
+            name="add_tool",
+            description="Adds two numbers and returns the sum.",
+            arguments=[
+                ToolArgument(name="a", arg_type="int", description="First number", required=True),
+                ToolArgument(name="b", arg_type="int", description="Second number", required=True)
+            ],
+            return_type="int"
+        )
+    
+    async def async_execute(self, **kwargs) -> str:
+        logger.info(f"Starting tool execution: {self.name}")
+        logger.info(f"Adding {kwargs['a']} and {kwargs['b']}")
+        result = str(int(kwargs["a"]) + int(kwargs["b"]))
+        logger.info(f"Finished tool execution: {self.name}")
+        return result
+
+class MultiplyTool(Tool):
+    def __init__(self):
+        super().__init__(
+            name="multiply_tool",
+            description="Multiplies two numbers and returns the product.",
+            arguments=[
+                ToolArgument(name="x", arg_type="int", description="First number", required=True),
+                ToolArgument(name="y", arg_type="int", description="Second number", required=True)
+            ],
+            return_type="int"
+        )
+    
+    async def async_execute(self, **kwargs) -> str:
+        logger.info(f"Starting tool execution: {self.name}")
+        logger.info(f"Multiplying {kwargs['x']} and {kwargs['y']}")
+        result = str(int(kwargs["x"]) * int(kwargs["y"]))
+        logger.info(f"Finished tool execution: {self.name}")
+        return result
+
+class ConcatTool(Tool):
+    def __init__(self):
+        super().__init__(
+            name="concat_tool",
+            description="Concatenates two strings.",
+            arguments=[
+                ToolArgument(name="s1", arg_type="string", description="First string", required=True),
+                ToolArgument(name="s2", arg_type="string", description="Second string", required=True)
+            ],
+            return_type="string"
+        )
+    
+    async def async_execute(self, **kwargs) -> str:
+        logger.info(f"Starting tool execution: {self.name}")
+        logger.info(f"Concatenating '{kwargs['s1']}' and '{kwargs['s2']}'")
+        result = kwargs["s1"] + kwargs["s2"]
+        logger.info(f"Finished tool execution: {self.name}")
+        return result
+
+class AgentTool(Tool):
+    def __init__(self, model: str = "gemini/gemini-2.0-flash"):
+        super().__init__(
+            name="agent_tool",
+            description="Generates text using a language model based on a system prompt and user prompt.",
+            arguments=[
+                ToolArgument(name="system_prompt", arg_type="string", description="System prompt to guide the model's behavior", required=True),
+                ToolArgument(name="prompt", arg_type="string", description="User prompt to generate a response for", required=True),
+                ToolArgument(name="temperature", arg_type="float", description="Temperature for generation (0 to 1)", required=True)
+            ],
+            return_type="string"
+        )
+        self.model = model
+    
+    async def async_execute(self, **kwargs) -> str:
+        logger.info(f"Starting tool execution: {self.name}")
+        system_prompt = kwargs["system_prompt"]
+        prompt = kwargs["prompt"]
+        temperature = float(kwargs["temperature"])
+        
+        # Validate temperature, unchanged
+        if not 0 <= temperature <= 1:
+            logger.error(f"Temperature {temperature} is out of range (0-1)")
+            raise ValueError("Temperature must be between 0 and 1")
+        
+        logger.info(f"Generating text with model {self.model}, temperature {temperature}")
+        try:
+            async with AsyncExitStack() as stack:
+                timeout_cm = asyncio.timeout(30)
+                await stack.enter_async_context(timeout_cm)
+                
+                logger.debug(f"Making API call to {self.model}")
+                response = await litellm.acompletion(
+                    model=self.model,
+                    messages=[
+                        {"role": "system", "content": system_prompt},
+                        {"role": "user", "content": prompt}
+                    ],
+                    temperature=temperature,
+                    max_tokens=1000  # Original default
+                )
+                generated_text = response.choices[0].message.content.strip()
+                logger.debug(f"Generated text: {generated_text}")
+                result = generated_text
+                logger.info(f"Finished tool execution: {self.name}")
+                return result
+        except TimeoutError as e:
+            error_msg = f"API call to {self.model} timed out after 30 seconds"
+            logger.error(error_msg)
+            raise TimeoutError(error_msg) from e
+        except Exception as e:
+            logger.error(f"Failed to generate text with {self.model}: {str(e)}")
+            raise RuntimeError(f"Text generation failed: {str(e)}")
+
+# Asynchronous function to generate the program, matching original behavior with updated prompt
+async def generate_program(task_description: str, tools: List[Tool], model: str, max_tokens: int) -> str:
+    """
+    Asynchronously generate a Python program that solves a given task using a list of tools.
+
+    Args:
+        task_description (str): A description of the task to be solved.
+        tools (List[Tool]): A list of Tool objects available for use.
+        model (str): The litellm model to use for code generation.
+        max_tokens (int): Maximum number of tokens for the generated response.
+
+    Returns:
+        str: A string containing a complete Python program.
+    """
+    logger.debug(f"Generating program for task: {task_description}")
+    tool_docstrings = "\n\n".join([tool.to_docstring() for tool in tools])
+
+    # Updated prompt with reinforced instruction to exclude __main__ block
+    prompt = f"""
+You are a Python code generator. Your task is to create a Python program that solves the following task:
+"{task_description}"
+
+You have access to the following pre-defined async tool functions, as defined with their signatures and descriptions:
+
+{tool_docstrings}
+
+Instructions:
+1. Generate a Python program as a single string.
+2. Include only the import for asyncio (import asyncio).
+3. Define an async function named main() that solves the task.
+4. Use the pre-defined tool functions (e.g., add_tool, multiply_tool, concat_tool) directly by calling them with await and the appropriate arguments as specified in their descriptions.
+5. Do not redefine the tool functions within the program; assume they are already available in the namespace.
+6. Return the program as markdown code block.
+7. Strictly exclude asyncio.run(main()) or any code outside the main() function definition, including any 'if __name__ == "__main__":' block, as the runtime will handle execution of main().
+8. Do not include explanatory text outside the program string.
+9. Express all string variables as multiline strings
+string, always start a string at the beginning of a line.
+10. Always print the result at the end of the program.
+
+Example task: "Add 5 and 7 and print the result"
+Example output:
+```python
+import asyncio
+
+async def main():
+    result = await add_tool(a=5, b=7)
+    print(result)
+```
+"""
+    
+    logger.debug(f"Prompt sent to litellm:\n{prompt}")
+    
+    try:
+        logger.debug(f"Calling litellm with model {model}")
+        response = await litellm.acompletion(
+            model=model,
+            messages=[
+                {"role": "system", "content": "You are a Python code generator."},
+                {"role": "user", "content": prompt}
+            ],
+            max_tokens=max_tokens,
+            temperature=0.3
+        )
+        generated_code = response.choices[0].message.content.strip()
+        logger.debug("Code generation successful")
+    except Exception as e:
+        logger.error(f"Failed to generate code: {str(e)}")
+        raise typer.BadParameter(f"Failed to generate code with model '{model}': {str(e)}")
+
+    # Clean up output, preserving original logic
+    if generated_code.startswith('"""') and generated_code.endswith('"""'):
+        generated_code = generated_code[3:-3]
+    elif generated_code.startswith("```python") and generated_code.endswith("```"):
+        generated_code = generated_code[9:-3].strip()
+    
+    # Post-processing to remove any __main__ block if generated despite instructions
+    if "if __name__ == \"__main__\":" in generated_code:
+        lines = generated_code.splitlines()
+        main_end_idx = next(
+            (i for i in range(len(lines)) if "if __name__" in lines[i]),
+            len(lines)
+        )
+        generated_code = "\n".join(lines[:main_end_idx]).strip()
+        logger.warning("Removed unexpected __main__ block from generated code")
+
+    return generated_code
+
+# Updated async core logic with improved interpreter usage
+async def generate_core(task: str, model: str, max_tokens: int) -> None:
+    """
+    Core logic to generate and execute a Python program based on a task description.
+
+    Args:
+        task (str): The task description to generate a program for.
+        model (str): The litellm model to use for generation.
+        max_tokens (int): Maximum number of tokens for the generated response.
+    """
+    logger.info(f"Starting generate command for task: {task}")
+    # Input validation, unchanged
+    if not task.strip():
+        logger.error("Task description is empty")
+        raise typer.BadParameter("Task description cannot be empty")
+    if max_tokens <= 0:
+        logger.error("max-tokens must be positive")
+        raise typer.BadParameter("max-tokens must be a positive integer")
+
+    # Initialize tools, unchanged
+    tools = [
+        AddTool(),
+        MultiplyTool(),
+        ConcatTool(),
+        AgentTool(model=model)
+    ]
+
+    # Generate the program
+    try:
+        program = await generate_program(task, tools, model, max_tokens)
+    except Exception as e:
+        logger.error(f"Failed to generate program: {str(e)}")
+        typer.echo(typer.style(f"Error: {str(e)}", fg=typer.colors.RED))
+        raise typer.Exit(code=1)
+
+    logger.debug(f"Generated program:\n{program}")
+    # Output the generated program with original style
+    typer.echo(typer.style("Generated Python Program:", fg=typer.colors.GREEN, bold=True))
+    typer.echo(program)
+    
+    # Validate program structure
+    try:
+        ast_tree = ast.parse(program)
+        has_async_main = any(
+            isinstance(node, ast.AsyncFunctionDef) and node.name == "main"
+            for node in ast.walk(ast_tree)
+        )
+        if not has_async_main:
+            logger.warning("Generated code lacks an async main() function")
+            typer.echo(typer.style("Warning: Generated code lacks an async main() function", fg=typer.colors.YELLOW))
+            return
+    except SyntaxError as e:
+        logger.error(f"Syntax error in generated code: {str(e)}")
+        typer.echo(typer.style(f"Syntax error in generated code: {str(e)}", fg=typer.colors.RED))
+        return
+
+    # Prepare namespace with tool instances
+    namespace: Dict[str, Callable] = {
+        "asyncio": asyncio,
+        "add_tool": partial(AddTool().async_execute),
+        "multiply_tool": partial(MultiplyTool().async_execute),
+        "concat_tool": partial(ConcatTool().async_execute),
+        "agent_tool": partial(AgentTool(model=model).async_execute),
+    }
+
+    # Check for namespace collisions
+    reserved_names = set(vars(__builtins__))
+    for name in namespace:
+        if name in reserved_names and name != "asyncio":
+            logger.warning(f"Namespace collision detected: '{name}' shadows a builtin")
+            typer.echo(typer.style(f"Warning: Tool name '{name}' shadows a builtin", fg=typer.colors.YELLOW))
+
+    # Execute the program
+    typer.echo("\n" + typer.style("Executing the program:", fg=typer.colors.GREEN, bold=True))
+    try:
+        logger.debug("Executing generated code with execute_async")
+        execution_result = await execute_async(
+            code=program,
+            timeout=30,
+            entry_point="main",
+            allowed_modules=["asyncio"],
+            namespace=namespace,
+        )
+        
+        # Detailed error handling
+        if execution_result.error:
+            if "SyntaxError" in execution_result.error:
+                logger.error(f"Syntax error: {execution_result.error}")
+                typer.echo(typer.style(f"Syntax error: {execution_result.error}", fg=typer.colors.RED))
+            elif "TimeoutError" in execution_result.error:
+                logger.error(f"Timeout: {execution_result.error}")
+                typer.echo(typer.style(f"Timeout: {execution_result.error}", fg=typer.colors.RED))
+            else:
+                logger.error(f"Runtime error: {execution_result.error}")
+                typer.echo(typer.style(f"Runtime error: {execution_result.error}", fg=typer.colors.RED))
+        else:
+            logger.info(f"Execution completed in {execution_result.execution_time:.2f} seconds")
+            typer.echo(typer.style(f"Execution completed in {execution_result.execution_time:.2f} seconds", fg=typer.colors.GREEN))
+            
+            # Display the result if it's not None
+            if execution_result.result is not None:
+                typer.echo("\n" + typer.style("Result:", fg=typer.colors.BLUE, bold=True))
+                typer.echo(str(execution_result.result))
+    except ValueError as e:
+        logger.error(f"Invalid code generated: {str(e)}")
+        typer.echo(typer.style(f"Invalid code: {str(e)}", fg=typer.colors.RED))
+    except Exception as e:
+        logger.error(f"Unexpected execution error: {str(e)}")
+        typer.echo(typer.style(f"Unexpected error during execution: {str(e)}", fg=typer.colors.RED))
+    else:
+        logger.info("Program executed successfully")
+
+@app.command()
+def generate(
+    task: str = typer.Argument(
+        ...,
+        help="The task description to generate a program for (e.g., 'Add 5 and 7 and print the result')"
+    ),
+    model: str = typer.Option(
+        "gemini/gemini-2.0-flash",
+        "--model",
+        "-m",
+        help="The litellm model to use for generation (e.g., 'gpt-3.5-turbo', 'gpt-4')"
+    ),
+    max_tokens: int = typer.Option(
+        4000,
+        "--max-tokens",
+        "-t",
+        help="Maximum number of tokens for the generated response (default: 4000)"
+    )
+) -> None:
+    """Generate and execute a Python program based on a task description"""
+    try:
+        # Run async core logic, preserving original execution style
+        asyncio.run(generate_core(task, model, max_tokens))
+    except Exception as e:
+        logger.error(f"Command failed: {str(e)}")
+        typer.echo(typer.style(f"Error: {str(e)}", fg=typer.colors.RED))
+        raise typer.Exit(code=1)
+
+# Entry point, unchanged
+def main() -> None:
+    logger.debug("Starting script execution")
+    app()
+
+if __name__ == "__main__":
+    main()

quantalogic/tools/duckduckgo_search_tool.py

@@ -33,6 +33,7 @@ class DuckDuckGoSearchTool(Tool):
     """
 
     name: str = "duckduckgo_tool"
+    need_post_process: bool = False
     description: str = "Retrieves search results from DuckDuckGo. " "Provides structured output of search results."
     arguments: list = [
         ToolArgument(

quantalogic/tools/execute_bash_command_tool.py

@@ -6,6 +6,10 @@ import signal
 import subprocess
 import sys
 from typing import Dict, Optional, Union
+from pathlib import Path
+from datetime import datetime
+import shutil
+import re
 
 from loguru import logger
 
@@ -24,7 +28,7 @@ class ExecuteBashCommandTool(Tool):
     """Tool for executing bash commands with real-time I/O handling."""
 
     name: str = "execute_bash_tool"
-    description: str = "Executes a bash command and returns its output."
+    description: str = "Executes a bash command and returns its output. All commands are executed in /tmp for security."
     need_validation: bool = True
     arguments: list = [
         ToolArgument(
@@ -34,14 +38,6 @@ class ExecuteBashCommandTool(Tool):
             required=True,
             example="ls -la",
         ),
-        ToolArgument(
-            name="working_dir",
-            arg_type="string",
-            description="The working directory where the command will be executed. Defaults to the current directory.",
-            required=False,
-            example="/path/to/directory",
-            default=os.getcwd(),
-        ),
         ToolArgument(
             name="timeout",
             arg_type="int",
@@ -52,15 +48,67 @@ class ExecuteBashCommandTool(Tool):
         ),
     ]
 
+    def _validate_command(self, command: str) -> None:
+        """Validate the command for potential security risks."""
+        forbidden_commands = ["rm -rf /", "mkfs", "dd", ":(){ :|:& };:"]
+        for cmd in forbidden_commands:
+            if cmd in command.lower():
+                raise ValueError(f"Command '{command}' contains forbidden operation")
+
+    def _handle_mkdir_command(self, command: str) -> str:
+        """Handle mkdir command with automatic cleanup and date suffix.
+        
+        Args:
+            command: The original mkdir command
+            
+        Returns:
+            Modified command that handles existing directories and adds date suffix
+        """
+        try:
+            # Extract directory name from mkdir command
+            mkdir_pattern = r'mkdir\s+(?:-p\s+)?["\']?([^"\'>\n]+)["\']?'
+            match = re.search(mkdir_pattern, command)
+            if not match:
+                return command
+                
+            dir_name = match.group(1)
+            # Add timestamp suffix
+            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+            new_dir_name = f"{dir_name}_{timestamp}"
+            
+            # Create cleanup and mkdir commands
+            cleanup_cmd = f'rm -rf "{dir_name}" 2>/dev/null; '
+            mkdir_cmd = command.replace(dir_name, new_dir_name)
+            
+            return cleanup_cmd + mkdir_cmd
+            
+        except Exception as e:
+            logger.warning(f"Error processing mkdir command: {e}")
+            return command
+
+    def _format_output(self, stdout: str, return_code: int, error: Optional[str] = None) -> str:
+        """Format command output with stdout, return code, and optional error."""
+        formatted_result = "<command_output>\n"
+        formatted_result += f" <stdout>{stdout.strip()}</stdout>\n"
+        formatted_result += f" <returncode>{return_code}</returncode>\n"
+        if error:
+            formatted_result += f" <error>{error}</error>\n"
+        formatted_result += "</command_output>"
+        return formatted_result
+
     def _execute_windows(
         self,
         command: str,
-        cwd: str,
         timeout_seconds: int,
         env_vars: Dict[str, str],
+        cwd: Optional[str] = None,
     ) -> str:
         """Execute command on Windows platform."""
         try:
+            # Handle mkdir commands
+            if "mkdir" in command:
+                command = self._handle_mkdir_command(command)
+
             # On Windows, use subprocess with pipes
             process = subprocess.Popen(
                 command,
@@ -68,7 +116,7 @@ class ExecuteBashCommandTool(Tool):
                 stdin=subprocess.PIPE,
                 stdout=subprocess.PIPE,
                 stderr=subprocess.PIPE,
-                cwd=cwd,
+                cwd="/tmp",  # cwd, Force /tmp directory
                 env=env_vars,
                 text=True,
                 encoding="utf-8",
@@ -77,34 +125,27 @@ class ExecuteBashCommandTool(Tool):
             try:
                 stdout, stderr = process.communicate(timeout=timeout_seconds)
                 return_code = process.returncode
-
-                if return_code != 0 and stderr:
-                    logger.warning(f"Command failed with error: {stderr}")
-
-                formatted_result = (
-                    "<command_output>"
-                    f" <stdout>{stdout.strip()}</stdout>"
-                    f" <returncode>{return_code}</returncode>"
-                    f"</command_output>"
-                )
-                return formatted_result
+                return self._format_output(stdout, return_code, stderr if stderr else None)
 
             except subprocess.TimeoutExpired:
                 process.kill()
-                return f"Command timed out after {timeout_seconds} seconds."
+                return self._format_output("", 1, f"Command timed out after {timeout_seconds} seconds")
 
         except Exception as e:
-            return f"Unexpected error executing command: {str(e)}"
+            return self._format_output("", 1, f"Unexpected error executing command: {str(e)}")
 
     def _execute_unix(
         self,
         command: str,
-        cwd: str,
         timeout_seconds: int,
         env_vars: Dict[str, str],
     ) -> str:
         """Execute command on Unix platform."""
         try:
+            # Handle mkdir commands
+            if "mkdir" in command:
+                command = self._handle_mkdir_command(command)
+
             master, slave = pty.openpty()
             proc = subprocess.Popen(
                 command,
@@ -112,7 +153,7 @@ class ExecuteBashCommandTool(Tool):
                 stdin=slave,
                 stdout=slave,
                 stderr=subprocess.STDOUT,
-                cwd=cwd,
+                cwd="/tmp",                # cwd=cwd, # Force /tmp directory
                 env=env_vars,
                 preexec_fn=os.setsid,
                 close_fds=True,
@@ -127,74 +168,90 @@ class ExecuteBashCommandTool(Tool):
                     rlist, _, _ = select.select([master, sys.stdin], [], [], timeout_seconds)
                     if not rlist:
                         if proc.poll() is not None:
-                            break  # Process completed but select timed out
+                            break
                         raise subprocess.TimeoutExpired(command, timeout_seconds)
 
                     for fd in rlist:
                         if fd == master:
-                            data = os.read(master, 1024).decode()
-                            if not data:
+                            try:
+                                data = os.read(master, 1024).decode()
+                                if not data:
+                                    break_loop = True
+                                    break
+                                stdout_buffer.append(data)
+                                sys.stdout.write(data)
+                                sys.stdout.flush()
+                            except (OSError, UnicodeDecodeError) as e:
+                                logger.warning(f"Error reading output: {e}")
                                 break_loop = True
                                 break
-                            stdout_buffer.append(data)
-                            sys.stdout.write(data)
-                            sys.stdout.flush()
                         elif fd == sys.stdin:
-                            user_input = os.read(sys.stdin.fileno(), 1024)
-                            os.write(master, user_input)
+                            try:
+                                user_input = os.read(sys.stdin.fileno(), 1024)
+                                os.write(master, user_input)
+                            except OSError as e:
+                                logger.warning(f"Error handling input: {e}")
 
                     if break_loop or proc.poll() is not None:
-                        while True:
-                            data = os.read(master, 1024).decode()
-                            if not data:
-                                break
-                            stdout_buffer.append(data)
-                            sys.stdout.write(data)
-                            sys.stdout.flush()
+                        try:
+                            while True:
+                                data = os.read(master, 1024).decode()
+                                if not data:
+                                    break
+                                stdout_buffer.append(data)
+                                sys.stdout.write(data)
+                                sys.stdout.flush()
+                        except (OSError, UnicodeDecodeError):
+                            pass
                         break
 
             except subprocess.TimeoutExpired:
                 os.killpg(os.getpgid(proc.pid), signal.SIGTERM)
-                return f"Command timed out after {timeout_seconds} seconds."
+                return self._format_output("", 1, f"Command timed out after {timeout_seconds} seconds")
             except EOFError:
-                pass  # Process exited normally
+                pass
             finally:
                 os.close(master)
                 proc.wait()
 
             stdout_content = "".join(stdout_buffer)
             return_code = proc.returncode
-            formatted_result = (
-                "<command_output>"
-                f" <stdout>{stdout_content.strip()}</stdout>"
-                f" <returncode>{return_code}</returncode>"
-                f"</command_output>"
-            )
-            return formatted_result
+            return self._format_output(stdout_content, return_code)
+
         except Exception as e:
-            return f"Unexpected error executing command: {str(e)}"
+            return self._format_output("", 1, f"Unexpected error executing command: {str(e)}")
 
     def execute(
         self,
         command: str,
-        working_dir: Optional[str] = None,
+        working_dir: Optional[str] = None,  # Kept for backward compatibility but ignored
         timeout: Union[int, str, None] = 60,
         env: Optional[Dict[str, str]] = None,
     ) -> str:
-        """Executes a bash command with interactive input handling."""
+        """Executes a bash command with interactive input handling in /tmp directory."""
+        # Ensure /tmp exists and is writable
+        tmp_dir = Path("/tmp")
+        if not (tmp_dir.exists() and os.access(tmp_dir, os.W_OK)):
+            return self._format_output("", 1, "Error: /tmp directory is not accessible")
+
+        # Validate command
+        try:
+            self._validate_command(command)
+        except ValueError as e:
+            return self._format_output("", 1, str(e))
+
         timeout_seconds = int(timeout) if timeout else 60
-        cwd = working_dir or os.getcwd()
         env_vars = os.environ.copy()
         if env:
             env_vars.update(env)
 
         if sys.platform == "win32":
-            return self._execute_windows(command, cwd, timeout_seconds, env_vars)
+            return self._execute_windows(command, timeout_seconds, env_vars)
         else:
             if not pty:
                 logger.warning("PTY module not available, falling back to Windows-style execution")
-                return self._execute_windows(command, cwd, timeout_seconds, env_vars)
-            return self._execute_unix(command, cwd, timeout_seconds, env_vars)
+                return self._execute_windows(command, timeout_seconds, env_vars)
+            return self._execute_unix(command, timeout_seconds, env_vars)
 
 
 if __name__ == "__main__":