import openai
import os
from flask import Flask, request, jsonify
from flask_cors import CORS

# Set your OpenAI API key
openai.api_key = os.environ.get("OPENAI_API_KEY")

app = Flask(__name__)
CORS(app)

@app.route("/api/chat", methods=["POST"])
def chat():
    user_messages = request.json["messages"]
    
    # Add a system message to set the context for the assistant
    system_message = {
        'role': 'system',
        'content': 'You are ChatGPT, a large language model trained by OpenAI. Answer as concisely as possible. Knowledge cutoff: 2021-09 Current date: 2023-04-15'
    }
    user_messages.insert(0, system_message)

    # Create a prompt for GPT-3.5 Turbo
    prompt = ""
    for message in user_messages:
        prompt += f"{message['role']}: {message['content']}\n"
    prompt += "Assistant:"
    
    # Make a request to the OpenAI API
    response = openai.Completion.create(
        engine="gpt-3.5-turbo",
        prompt=prompt,
        max_tokens=150,
        n=1,
        stop=None,
        temperature=0.5,
    )

    # Extract the generated response from the API response
    assistant_message = response.choices[0].text.strip()
    total_tokens = response['usage']['total_tokens']

    return jsonify({"message": assistant_message, "total_tokens": total_tokens})

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