mellea/docs/examples/sampling/python_code_generation_preset.py at be3c214ab67304c164ea4992959d3929089497d4 · generative-computing/mellea · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
# pytest: e2e, qualitative, ollama
"""Python code generation sampling preset with automatic repair.

Demonstrates python_code_generation_sampling(), which bundles extraction, syntax,
and execution requirements with repair feedback into a reusable configuration.

Usage:
    uv run python docs/examples/sampling/python_code_generation_preset.py
"""

import mellea
from mellea.stdlib.sampling import python_code_generation_sampling


def example_basic_code_generation():
    """Generate Python code with automatic repair feedback.

    Demonstrates the basic usage of python_code_generation_sampling() with
    default parameters. The model generates code to compute the sum of 1 to 100.
    """
    session = mellea.start_session()

    preset = python_code_generation_sampling()

    prompt = (
        "Write Python code to compute the sum of integers from 1 to 100. "
        "Return the result using a variable named 'total'."
    )

    result = session.instruct(
        prompt, requirements=preset.requirements, strategy=preset.strategy
    )

    code = str(result)
    print("Generated code:")
    print(code)
    print()
    assert "range" in code or "101" in code, "Result should reference the computation"


def example_with_import_restrictions():
    """Generate code with strict import allowlist.

    Demonstrates restricting imports to safe modules (numpy, math).
    If the model tries to import subprocess or other restricted modules,
    the preset provides model-friendly feedback about forbidden imports.
    """
    session = mellea.start_session()

    preset = python_code_generation_sampling(
        allowed_imports=["numpy", "math"],
        loop_budget=3,  # Allow more repair attempts for restricted imports
    )

    prompt = (
        "Write Python code using numpy to create an array [1, 2, 3, 4, 5] "
        "and compute its mean. Only use numpy and math modules."
    )

    result = session.instruct(
        prompt, requirements=preset.requirements, strategy=preset.strategy
    )

    # If code generation fails due to import violations, the feedback will guide
    # the model: "Your code imports forbidden modules: [X]. These are not available.
    # Try: Remove these imports and use only standard library or whitelisted modules."

    code = str(result)
    print("Generated code (with import restrictions):")
    print(code)
    print()
    assert "numpy" in code, "Code should use numpy"
    # Verify no forbidden imports are present
    assert "subprocess" not in code

    # Use ' os' so that 'purposes' and other words don't trigger the assert.
    assert " os" not in code


def example_with_tighter_constraints():
    """Generate code with stricter resource constraints.

    Demonstrates configuring the preset with:
    - Smaller output limit (5KB instead of 10KB)
    - Shorter timeout (3 seconds instead of 5)
    - Only 2 repair attempts instead of the default

    Useful for resource-constrained environments or performance-critical tasks.
    """
    session = mellea.start_session()

    preset = python_code_generation_sampling(
        output_limit_chars=5000, timeout_seconds=3, loop_budget=2
    )

    prompt = (
        "Write a short Python function that returns the first 10 Fibonacci numbers "
        "as a list. Keep it concise."
    )

    result = session.instruct(
        prompt, requirements=preset.requirements, strategy=preset.strategy
    )

    # Verify output stays within the configured limit
    code = str(result)
    print("Generated code (with tight constraints):")
    print(code)
    print()
    assert len(code) < 5000, "Output should be within limit"


def main():
    """Run all examples."""
    print("Running example_basic_code_generation...")
    example_basic_code_generation()
    print("✓ Basic code generation example passed\n")

    print("Running example_with_import_restrictions...")
    example_with_import_restrictions()
    print("✓ Import restrictions example completed\n")

    print("Running example_with_tighter_constraints...")
    example_with_tighter_constraints()
    print("✓ Tighter constraints example completed\n")

    print("All examples completed!")


if __name__ == "__main__":
    main()