app.py

import os
from http import HTTPStatus
from flask import Flask, request
from flask_caching import Cache
from flask_cors import CORS
from rubber_duck import codegen


codgen_model = os.environ.get('MODEL', 'Salesforce/codegen-350M-mono')
print(f'Attempting to setup codegen with the model {codgen_model}')
codegen.setup(
    os.environ.get('MODEL', 'Salesforce/codegen-350M-mono'),
    setup_torch=os.environ.get("SETUP_TORCH") == 'true',
)


app = Flask(__name__)
CORS(app)
cache = Cache(app, config={'CACHE_TYPE': 'simple'})


@app.route('/generate', methods=['GET'])
@cache.cached(timeout=1440)  # Cache for 24 hours
def generate():
    user_input = request.args.get('input', '', type=str)
    if not user_input:
        return 'Missing input', HTTPStatus.BAD_REQUEST

    return {
        'result': codegen.generate(user_input)
    }, HTTPStatus.OK


@app.route('/status', methods=['GET'])
def status():
    return {'status': 'OK'}