from flask import Flask, request, Response
from waitress import serve
import requests
import os

app = Flask(__name__)

OLLAMA_API_URL = "http://localhost:11434"

def generate(data):
    with requests.post(f"{OLLAMA_API_URL}/api/chat", json=data, stream=True) as r:
        for line in r.iter_lines():
            if line:
                yield line.decode('utf-8') + '\n'

@app.route('/chat', methods=['POST'])
def chat():
    data = request.get_json(silent=True)
    return Response(generate(data), mimetype='text/plain')

if __name__ == '__main__':
    port = int(os.environ.get('PORT', 7860))
    serve(app, host='0.0.0.0', port=port)