Self-Correcting Code Generation Loop with DSPy

An intelligent code generation system that uses DSPy and Ollama to automatically generate, test, and refine Python code based on natural language descriptions.

🚀 Features

• Intelligent Code Generation: Converts natural language task descriptions into working Python code
• Automatic Test Generation: Creates comprehensive unit tests for generated code
• Self-Correcting Loop: Automatically refines code based on test failures until tests pass
• DSPy Integration: Leverages DSPy's optimization capabilities for improved code generation
• Local LLM Support: Uses Ollama with Llama3 for privacy-focused, offline code generation

🔧 Prerequisites

Before running this project, ensure you have the following installed:

• Python 3.8+
• Ollama with Llama3 model
• pytest for test execution

Installing Ollama and Llama3

1. Install Ollama from https://ollama.ai
2. Pull the Llama3 model:

   ollama pull llama3

3. Verify Ollama is running:

   ollama serve

📦 Installation

1. Clone the repository:

   git clone https://github.com/evalops/dspy-code-generator.git
   cd dspy-code-generator

2. Create and activate a virtual environment:

   python -m venv venv
   source venv/bin/activate  # On Windows: venv\Scripts\activate

3. Install dependencies:

   pip install -r requirements.txt

🚦 Quick Start

1. Ensure Ollama is running with Llama3 model available
2. Run the code generator:

   python main.py

The default example will generate a function to filter even numbers from a list.

💡 How It Works

The self-correcting loop follows these steps:

1. Initial Generation: Takes a natural language task description and generates Python code
2. Test Generation: Creates comprehensive unit tests for the generated code
3. Test Execution: Runs the tests using pytest in a temporary environment
4. Refinement Loop: If tests fail, uses the error output to refine the code (up to 3 iterations)
5. Success: Returns the final working code and tests when all tests pass

Architecture Components

• CodeGeneratorSignature: Converts task descriptions to Python code
• TestGeneratorSignature: Generates pytest unit tests for given code
• CodeRefinerSignature: Refines failing code based on test output
• CodeGenProgram: Orchestrates the entire self-correcting pipeline

🔧 Configuration

Customizing the LLM Model

Edit the model configuration in main.py:

llm = dspy.LM(
    "ollama_chat/llama3",  # Change model here
)

Adjusting Refinement Iterations

Modify the max_refinements parameter:

prediction = compiled_program(task_description=new_task, max_refinements=5)  # Default: 3

📝 Usage Examples

Custom Task Example

from main import CodeGenProgram

# Initialize the program
code_gen_program = CodeGenProgram()

# Define your task
task = "Write a Python function that calculates the factorial of a number recursively"

# Generate code
result = code_gen_program(task_description=task, max_refinements=3)

print("Generated Code:")
print(result.python_code)
print("\nGenerated Tests:")
print(result.unit_test_code)
print(f"\nTests Passed: {result.test_passed}")

Training with Custom Examples

# Add your own training examples
custom_training_data = [
    dspy.Example(
        task_description="Your task description here",
        python_code="def your_function():\n    pass",
        unit_test_code="def test_your_function():\n    assert True"
    ).with_inputs("task_description")
]

# Compile with custom data
teleprompter = dspy.teleprompt.BootstrapFewShot(metric=test_execution_metric)
compiled_program = teleprompter.compile(code_gen_program, trainset=custom_training_data)

🐛 Troubleshooting

Common Issues

Ollama Connection Error:

• Ensure Ollama service is running: ollama serve
• Verify Llama3 model is available: ollama list

Pytest Not Found:

• Install pytest: pip install pytest
• Ensure pytest is in your PATH

Code Generation Fails:

• Check if the task description is clear and specific
• Try increasing max_refinements for complex tasks
• Verify internet connection for Ollama model downloads

Tests Never Pass:

• The task might be too complex or ambiguous
• Try breaking down complex tasks into smaller subtasks
• Check the generated test logic for correctness

🔬 Development

Project Structure

dspy-code-generator/
├── main.py              # Main application with DSPy pipeline
├── requirements.txt     # Python dependencies
├── README.md           # This file
└── venv/               # Virtual environment (created after setup)

Running Tests

The system generates its own tests, but you can also test the framework:

# Test with a simple task
python -c "
from main import CodeGenProgram
program = CodeGenProgram()
result = program('Write a function that adds two numbers')
print('Success!' if result.test_passed else 'Failed')
"

🤝 Contributing

1. Fork the repository
2. Create a feature branch: git checkout -b feature-name
3. Make your changes and test thoroughly
4. Commit with clear messages: git commit -m "Add feature description"
5. Push to your fork: git push origin feature-name
6. Submit a pull request

📄 License

This project is open source. Please check the license file for details.

🔗 Dependencies

• dspy-ai: DSPy framework for language model programming
• pytest: Testing framework for generated code
• ollama: Local LLM inference engine

📚 Additional Resources

• DSPy Documentation
• Ollama Documentation
• pytest Documentation