Code Executor

import asyncio
import sys
from io import StringIO
from ai_query import generate_text
from ai_query.providers import anthropic, tool, Field, step_count_is

@tool(description="Execute Python code and return the output")
def run_python(
    code: str = Field(description="Python code to execute")
) -> str:
    """Execute Python code in a restricted environment."""
    # Capture stdout
    old_stdout = sys.stdout
    sys.stdout = StringIO()

    try:
        # Create a restricted namespace
        namespace = {
            "__builtins__": {
                "print": print,
                "range": range,
                "len": len,
                "sum": sum,
                "min": min,
                "max": max,
                "abs": abs,
                "round": round,
                "sorted": sorted,
                "list": list,
                "dict": dict,
                "set": set,
                "str": str,
                "int": int,
                "float": float,
                "bool": bool,
                "enumerate": enumerate,
                "zip": zip,
                "map": map,
                "filter": filter,
            }
        }

        exec(code, namespace)
        output = sys.stdout.getvalue()

        if not output:
            # Check if there's a result variable
            if "result" in namespace:
                output = str(namespace["result"])
            else:
                output = "Code executed successfully (no output)"

        return output.strip()

    except Exception as e:
        return f"Error: {type(e).__name__}: {e}"

    finally:
        sys.stdout = old_stdout

async def main():
    print("Code Executor Agent")
    print("=" * 40)

    problems = [
        "Calculate the factorial of 10",
        "Find all prime numbers up to 50",
        "Generate the first 15 Fibonacci numbers",
    ]

    for problem in problems:
        print(f"\n📝 Problem: {problem}")
        print("-" * 40)

        result = await generate_text(
            model=anthropic("claude-sonnet-4-20250514"),
            system="""You are a Python coding assistant. When given a problem:
            1. Write clean, efficient Python code to solve it
            2. Use the run_python tool to execute your code
            3. The code should print the result

            Keep your code simple and focused.""",
            prompt=problem,
            tools={"run_python": run_python},
            stop_when=step_count_is(3)  # Limit retries
        )

        print(f"Result: {result.text}")

if __name__ == "__main__":
    asyncio.run(main())

Code Executor Agent
========================================

📝 Problem: Calculate the factorial of 10
----------------------------------------
Result: The factorial of 10 is 3,628,800.

I calculated this by multiplying all integers from 1 to 10:
10! = 10 × 9 × 8 × 7 × 6 × 5 × 4 × 3 × 2 × 1 = 3,628,800

📝 Problem: Find all prime numbers up to 50
----------------------------------------
Result: The prime numbers up to 50 are:
2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47

There are 15 prime numbers in this range.

📝 Problem: Generate the first 15 Fibonacci numbers
----------------------------------------
Result: The first 15 Fibonacci numbers are:
1, 1, 2, 3, 5, 8, 13, 21, 34, 55, 89, 144, 233, 377, 610

Each number is the sum of the two preceding numbers.

namespace = {
    "__builtins__": {
        "print": print,
        "range": range,
        "len": len,
        # ... only safe functions
    }
}