from flask import Flask, jsonify, render_template, request, Response
from openai import OpenAI
from dotenv import load_dotenv
import os
import json

load_dotenv()

app = Flask(__name__,
            static_folder='static',
            template_folder='templates'
)
client = OpenAI(
    api_key=os.getenv('API_KEY'),
    base_url=os.getenv('BASE_URL')
)

class ChatBot:
    def __init__(self):
        self.messages = []

    def new_conversation(self):
        self.messages = []
        return {"status": "success", "message": "New conversation started"}

    def get_stream_response(self, user_input):
        self.messages.append({"role": "user", "content": user_input})

        try:
            response = client.chat.completions.create(
                model=os.getenv('MODEL'),
                messages=self.messages,
                stream=True
            )

            for chunk in response:
                if chunk.choices[0].delta.content:
                    yield f"data: {json.dumps({'content': chunk.choices[0].delta.content})}\n\n"

            # 保存完整对话到历史记录
            full_response = ''.join(chunk.choices[0].delta.content or '' 
                                  for chunk in response)
            self.messages.append({"role": "assistant", "content": full_response})

        except Exception as e:
            yield f"data: {json.dumps({'error': str(e)})}\n\n"

chatbot = ChatBot()

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

@app.route('/chat', methods=['POST'])
def chat():
    user_message = request.json.get('message', '')
    return Response(
        chatbot.get_stream_response(user_message),
        mimetype='text/event-stream'
    )
@app.route('/new_conversation', methods=['POST'])
def new_conversation():
    return jsonify(chatbot.new_conversation())

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=7860, debug=True)