application.py

from flask import Flask, request, jsonify, Response
from flask_cors import CORS
import asyncio
import logging
import queue
import threading
from open_ai_modules import ChromeExtensionAssistant

app = Flask(__name__)
CORS(app)  # Enable CORS for all routes

# Configure logging
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")

# Create a global instance of the assistant
assistant = ChromeExtensionAssistant()

@app.route('/create_thread', methods=['POST'])
def create_thread():
    """Create or update a thread for a website"""
    try:
        data = request.json
        website_url = data.get('website_url')
        existing_thread_id = data.get('thread_id')  # Check if client sent a thread_id
        
        if not website_url:
            return jsonify({"error": "Missing website_url parameter"}), 400
        
        # Run the async function in the Flask context
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        
        if existing_thread_id:
            logging.info(f"Updating existing thread {existing_thread_id} with new URL {website_url}")
            assistant.active_thread_id = existing_thread_id
            result = loop.run_until_complete(assistant.update_website_context_async(existing_thread_id, website_url))
        else:
            logging.info(f"Creating new thread for URL {website_url}")
            result = loop.run_until_complete(assistant.create_thread(website_url))
        
        loop.close()
        
        if "error" in result:
            return jsonify({"error": result["error"]}), 500
        
        return jsonify({"thread_id": result["thread_id"]})
    
    except Exception as e:
        logging.error(f"Error in create_thread: {str(e)}")
        return jsonify({"error": str(e)}), 500

@app.route('/chat', methods=['POST'])
def chat():
    """Send a message to the AI assistant (non-streaming)"""
    try:
        data = request.json
        thread_id = data.get('thread_id')
        message = data.get('message')
        
        if not thread_id:
            return jsonify({"error": "Missing thread_id parameter"}), 400
        if not message:
            return jsonify({"error": "Missing message parameter"}), 400
        
        if thread_id not in assistant.threads:
            return jsonify({"error": f"Thread {thread_id} not found"}), 404
        
        response = assistant.chat(thread_id, message)
        
        return jsonify({"response": response})
    
    except Exception as e:
        logging.error(f"Error in chat: {str(e)}")
        return jsonify({"error": str(e)}), 500

@app.route('/chat_stream', methods=['POST'])
def chat_stream():
    """Stream responses from the AI assistant"""
    try:
        data = request.json
        thread_id = data.get('thread_id')
        message = data.get('message')
        
        if not thread_id:
            return jsonify({"error": "Missing thread_id parameter"}), 400
        if not message:
            return jsonify({"error": "Missing message parameter"}), 400
        
        if thread_id not in assistant.threads:
            return jsonify({"error": f"Thread {thread_id} not found"}), 404
        
        token_queue = queue.Queue()
        
        def token_collector(token):
            token_queue.put(token)
        
        def run_chat_stream():
            try:
                assistant.chat_stream(thread_id, message, token_collector)
                token_queue.put(None)
            except Exception as e:
                logging.error(f"Error in streaming thread: {str(e)}")
                token_queue.put(None)
        
        threading.Thread(target=run_chat_stream, daemon=True).start()
        
        def generate():
            while True:
                token = token_queue.get()
                if token is None:
                    break
                yield token
        
        return Response(generate(), mimetype='text/plain')
    
    except Exception as e:
        logging.error(f"Error in chat_stream: {str(e)}")
        return jsonify({"error": str(e)}), 500

@app.route('/get_active_thread', methods=['GET'])
def get_active_thread():
    """Get the currently active thread ID"""
    try:
        active_thread_id = assistant.get_active_threads()
        return jsonify({"thread_ids": active_thread_id})
    except Exception as e:
        logging.error(f"Error getting active thread: {str(e)}")
        return jsonify({"error": str(e)}), 500

if __name__ == '__main__':
    app.run(debug=True, port=5000)