quantalogic 0.61.2__py3-none-any.whl → 0.80__py3-none-any.whl

quantalogic/flow/flow_manager_schema.py

@@ -223,16 +223,48 @@ class TransitionDefinition(BaseModel):
     )
 
 
+class LoopDefinition(BaseModel):
+    """Definition of a loop within the workflow."""
+    nodes: List[str] = Field(..., description="List of node names in the loop.")
+    condition: str = Field(..., description="Python expression using 'ctx' for the loop condition.")
+    exit_node: str = Field(..., description="Node to transition to when the loop ends.")
+
+
 class WorkflowStructure(BaseModel):
     """Structure defining the workflow's execution flow."""
     start: Optional[str] = Field(None, description="Name of the starting node.")
     transitions: List[TransitionDefinition] = Field(
         default_factory=list, description="List of transitions between nodes."
     )
+    loops: List[LoopDefinition] = Field(
+        default_factory=list, description="List of loop definitions (optional, for explicit loop support)."
+    )
     convergence_nodes: List[str] = Field(
         default_factory=list, description="List of nodes where branches converge."
     )
 
+    @model_validator(mode="before")
+    @classmethod
+    def check_loop_nodes(cls, data: Any) -> Any:
+        """Ensure all nodes in loops exist in the workflow.
+
+        Args:
+            data: Raw data to validate.
+
+        Returns:
+            Validated data.
+
+        Raises:
+            ValueError: If loop nodes are not defined.
+        """
+        loops = data.get("loops", [])
+        nodes = set(data.get("nodes", {}).keys())
+        for loop in loops:
+            for node in loop["nodes"] + [loop["exit_node"]]:
+                if node not in nodes:
+                    raise ValueError(f"Loop node '{node}' not defined in nodes")
+        return data
+
 
 class WorkflowDefinition(BaseModel):
     """Top-level definition of the workflow."""

quantalogic/tools/action_gen.py

@@ -8,8 +8,8 @@ from typing import Callable, Dict, List
 import litellm
 import typer
 from loguru import logger
+from quantalogic_pythonbox.python_interpreter import execute_async
 
-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

quantalogic/tools/action_gen_safe.py

@@ -0,0 +1,340 @@
+import ast
+import asyncio
+from functools import partial
+from typing import Callable, Dict, List
+
+import litellm
+import typer
+from loguru import logger
+
+from quantalogic.tools.tool import Tool, ToolArgument
+from quantalogic.utils.python_interpreter import ASTInterpreter
+
+# Configure loguru to log to a file with rotation
+logger.add("action_gen_safe.log", rotation="10 MB", level="DEBUG")
+
+# Initialize Typer app
+app = typer.Typer()
+
+# Define tool classes with logging in async_execute
+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"Adding {kwargs['a']} and {kwargs['b']}")
+        return str(int(kwargs["a"]) + int(kwargs["b"]))
+
+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"Multiplying {kwargs['x']} and {kwargs['y']}")
+        return str(int(kwargs["x"]) * int(kwargs["y"]))
+
+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"Concatenating '{kwargs['s1']}' and '{kwargs['s2']}'")
+        return kwargs["s1"] + kwargs["s2"]
+
+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:
+        system_prompt = kwargs["system_prompt"]
+        prompt = kwargs["prompt"]
+        temperature = float(kwargs["temperature"])
+        
+        # Validate temperature
+        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:
+            response = await litellm.acompletion(
+                model=self.model,
+                messages=[
+                    {"role": "system", "content": system_prompt},
+                    {"role": "user", "content": prompt}
+                ],
+                temperature=temperature,
+                max_tokens=1000  # Reasonable default for text generation
+            )
+            generated_text = response.choices[0].message.content.strip()
+            logger.debug(f"Generated text: {generated_text}")
+            return generated_text
+        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
+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}")
+    # Collect tool docstrings
+    tool_docstrings = "\n\n".join([tool.to_docstring() for tool in tools])
+
+    # Construct the prompt for litellm
+    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 a string enclosed in triple quotes (\"\"\"program\"\"\")).
+7. Do not include asyncio.run(main()) or any code outside the main() function definition.
+8. Do not include explanatory text outside the program string.
+9. Express all string variables as multiline strings (\"\"\"\nstring\n\"\"\"), always start a string at the beginning of a line.
+10. Always print the result at the end of the program.
+11. Never call the main() function.
+12. Never use not defined variables, modules, or functions.
+
+Example task: "Add 5 and 7 and print the result"
+Example output:
+\"\"\"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:
+        # Call litellm asynchronously to generate the program
+        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 the output
+    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()
+    
+    return generated_code
+
+# Async core logic for generate
+async def generate_core(task: str, model: str, max_tokens: int):
+    logger.info(f"Starting generate command for task: {task}")
+    # Input validation
+    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
+    tools = [
+        AddTool(),
+        MultiplyTool(),
+        ConcatTool(),
+        AgentTool(model=model)  # Include AgentTool with the specified 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
+    typer.echo(typer.style("Generated Python Program:", fg=typer.colors.GREEN, bold=True))
+    typer.echo(program)
+    
+    # Attempt to execute the program using the safe interpreter
+    typer.echo("\n" + typer.style("Executing the program safely:", fg=typer.colors.GREEN, bold=True))
+    try:
+        # Create instances of tools
+        add_tool_instance = AddTool()
+        multiply_tool_instance = MultiplyTool()
+        concat_tool_instance = ConcatTool()
+        agent_tool_instance = AgentTool(model=model)
+
+        # Define the allowed modules
+        allowed_modules = ["asyncio"]
+        
+        # Create a namespace with our tools
+        namespace = {
+            "add_tool": add_tool_instance.async_execute,
+            "multiply_tool": multiply_tool_instance.async_execute,
+            "concat_tool": concat_tool_instance.async_execute,
+            "agent_tool": agent_tool_instance.async_execute,
+        }
+        
+        # Parse the program to AST
+        program_ast = ast.parse(program)
+        
+        # Extract imports first to ensure they're allowed
+        for node in program_ast.body:
+            if isinstance(node, ast.Import):
+                for name in node.names:
+                    if name.name not in allowed_modules:
+                        raise ValueError(f"Import not allowed: {name.name}")
+            elif isinstance(node, ast.ImportFrom):
+                if node.module not in allowed_modules:
+                    raise ValueError(f"Import from not allowed: {node.module}")
+        
+        # Find the main function in the AST
+        main_func_node = None
+        for node in program_ast.body:
+            if isinstance(node, ast.AsyncFunctionDef) and node.name == "main":
+                main_func_node = node
+                break
+        
+        if not main_func_node:
+            logger.warning("No async main() function found in generated code")
+            typer.echo(typer.style("Warning: No async main() function found", fg=typer.colors.YELLOW))
+            return
+            
+        # Execute the code safely to define functions in the namespace
+        # Use a custom ASTInterpreter directly instead of calling interpret_code
+        interpreter = ASTInterpreter(allowed_modules=allowed_modules, source=program)
+        
+        # Add tools to the interpreter's environment
+        interpreter.env_stack[0].update({
+            "add_tool": add_tool_instance.async_execute,
+            "multiply_tool": multiply_tool_instance.async_execute,
+            "concat_tool": concat_tool_instance.async_execute,
+            "agent_tool": agent_tool_instance.async_execute,
+        })
+        
+        # First, execute all non-main code (imports and other top-level code)
+        for node in program_ast.body:
+            if not isinstance(node, ast.AsyncFunctionDef) or node.name != "main":
+                await interpreter.visit(node, wrap_exceptions=True)
+        
+        # Then execute the main function definition
+        for node in program_ast.body:
+            if isinstance(node, ast.AsyncFunctionDef) and node.name == "main":
+                await interpreter.visit(node, wrap_exceptions=True)
+        
+        # Get the main function from the environment
+        main_func = interpreter.get_variable("main")
+        
+        # Execute the main function in the current event loop
+        logger.debug("Executing main() function")
+        await main_func()
+        
+        logger.info("Program executed successfully")
+        typer.echo(typer.style("Execution completed successfully", fg=typer.colors.GREEN))
+    
+    except SyntaxError as e:
+        logger.error(f"Syntax error in generated code: {e}")
+        typer.echo(typer.style(f"Syntax error: {e}", fg=typer.colors.RED))
+    except Exception as e:
+        logger.error(f"Execution error: {e}")
+        typer.echo(typer.style(f"Execution failed: {e}", fg=typer.colors.RED))
+
+# Synchronous callback to invoke async generate_core
+@app.callback(invoke_without_command=True)
+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)"
+    )
+):
+    """
+    Asynchronously generate a Python program based on a task description using specified tools and model.
+    Executes the program in a controlled, safe environment.
+
+    Examples:
+        $ python action_gen_safe.py "Add 5 and 7 and print the result"
+        $ python action_gen_safe.py "Concatenate 'Hello' and 'World' and print it" --model gpt-4 --max-tokens 5000
+    """
+    asyncio.run(generate_core(task, model, max_tokens))
+
+# Entry point to start the app
+def main():
+    logger.debug("Starting script execution")
+    app()
+
+if __name__ == "__main__":
+    main()