main.py

from flask import Flask, render_template, request, jsonify
import ollama

app = Flask(__name__)

DOCKERFILE_PROMPT = """
You must ONLY generate a Dockerfile for the {language} programming language.
DO NOT generate a Dockerfile for any other language under any circumstances.

Create a production-ready Dockerfile for {language} with:
1. Base image MUST be specific to {language}
2. All necessary {language} runtime/compiler installations
3. {language}-specific dependency management and build tools
4. Working directory setup
5. Source code handling
6. Default {language} application run command
7. Best practices for {language} containerization

Only output the Dockerfile content, no explanations.
"""

def sanitize_language(language):
    """
    Sanitize and validate the language input.
    Returns cleaned language name or raises ValueError if invalid.
    """
    if not language or not isinstance(language, str):
        raise ValueError("Invalid language input")
    
    # Clean the language input
    cleaned = ' '.join(language.strip().split())
    if not cleaned:
        raise ValueError("Empty language input")
    
    return cleaned

def generate_dockerfile(language):
    """Generate a Dockerfile for the specified language."""
    # Sanitize and validate the language input
    sanitized_language = sanitize_language(language)
        
    response = ollama.chat(
        model="llama3.2:1b", # Replace with the actual model name you are using
        messages=[
            {
                "role": "system", 
                "content": f"You are a Dockerfile expert specializing in {sanitized_language}. You must ONLY generate Dockerfiles for {sanitized_language}."
            },
            {
                "role": "user", 
                "content": DOCKERFILE_PROMPT.format(language=sanitized_language)
            }
        ]
    )
    
    dockerfile = response['message']['content'].strip()
    
    # Basic validation of the generated Dockerfile
    if "FROM" not in dockerfile:
        raise ValueError("Generated Dockerfile is invalid - missing FROM instruction")
    
    return dockerfile

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/generate', methods=['POST'])
def generate():
    try:
        language = request.json.get('language', '').strip()
        if not language:
            return jsonify({'error': 'Language not specified'}), 400
        
        dockerfile = generate_dockerfile(language)
        return jsonify({'dockerfile': dockerfile})
    except ValueError as e:
        return jsonify({'error': str(e)}), 400
    except Exception as e:
        return jsonify({'error': str(e)}), 500

if __name__ == "__main__":
    app.run(debug=True)