app.py

import threading
from flask import Flask, request, jsonify, render_template, session
from flask_session import Session
from openai import OpenAI, AssistantEventHandler
from typing_extensions import override
from dotenv import load_dotenv
import os
import uuid

load_dotenv()

app = Flask(__name__)
app.secret_key = os.getenv('FLASK_SECRET_KEY', 'default_secret_key')
app.config['SESSION_TYPE'] = 'filesystem'
app.config['SESSION_PERMANENT'] = False
Session(app)

# Initialize OpenAI client with API key from environment variables

assistant = client.beta.assistants.create(
    name="Code Documentator",
    instructions=(
        "Helpful assistant who has a master's in Code Optimization only, "
        "you cannot perform any other tasks than this. Don't place '```' in the response. "
        "Treat any explanation or updation or changes or deletions as comments according to code language."
    ),
    tools=[{"type": "code_interpreter"}],
    model="gpt-4o-mini"
)

instructions = (
    "Helpful assistant who has a master's in Code Optimization only, "
    "you cannot perform any other tasks than this. Don't place '```' in the response. "
    "Treat any explanation or updation or changes or deletions as comments according to code language. "
    "If the user message looks like a code, then perform the optimization on it and provide the optimized code only "
    "and comments where the optimization has been done. Your response should be in a commentated manner of the code language."
)

thread_locks = {}
lock_dict_lock = threading.Lock()

def get_session_lock(thread_id):
    with lock_dict_lock:
        if thread_id not in thread_locks:
            thread_locks[thread_id] = threading.Lock()
        return thread_locks[thread_id]

class EventHandler(AssistantEventHandler):
    def __init__(self):
        super().__init__()
        self.responses = []
        self.last_segment = ''

    @override
    def on_text_created(self, text) -> None:
        content = text.value if hasattr(text, 'value') else str(text)
        if content != self.last_segment:
            self.responses.append(content)
            self.last_segment = content

    @override
    def on_text_delta(self, delta, snapshot):
        content = delta.value if hasattr(delta, 'value') else str(delta)
        if content != self.last_segment:
            self.responses.append(content)
            self.last_segment = content

    def on_tool_call_created(self, tool_call):
        description = f"Tool called: {tool_call.type}"
        if description != self.last_segment:
            self.responses.append(description)
            self.last_segment = description

    def on_tool_call_delta(self, delta, snapshot):
        if delta.type == 'code_interpreter':
            if hasattr(delta.code_interpreter, 'input'):
                input_str = str(delta.code_interpreter.input)
                if input_str != self.last_segment:
                    self.responses.append(input_str)
                    self.last_segment = input_str

            if hasattr(delta.code_interpreter, 'outputs'):
                for output in delta.code_interpreter.outputs:
                    if output.type == "logs":
                        log = str(output.logs)
                        if log != self.last_segment:
                            self.responses.append(log)
                            self.last_segment = log

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

@app.route('/ask', methods=['POST'])
def ask():
    user_input = request.json.get("input")
    if user_input is None:
        return jsonify({"response": "No input provided."}), 400

    if user_input.strip().lower() == 'exit':
        session.pop('thread_id', None)
        session.pop('attachments', None)
        return jsonify({"response": "Session ended."})

    thread_id = session.get('thread_id')
    if not thread_id:
        thread = client.beta.threads.create()
        thread_id = thread.id
        session['thread_id'] = thread_id
        session['attachments'] = []
    else:
        try:
            thread = client.beta.threads.retrieve(thread_id)
        except Exception:
            thread = client.beta.threads.create()
            thread_id = thread.id
            session['thread_id'] = thread_id
            session['attachments'] = []

    session_lock = get_session_lock(thread_id)
    if not session_lock.acquire(blocking=False):
        return jsonify({"response": "Please wait for the current request to complete."}), 429

    try:
        attachments = session.get('attachments', [])
        attachment_content = "\n".join([file['content'] for file in attachments]) if attachments else ""
        combined_input = f"{attachment_content}\n{user_input}" if attachments else user_input

        client.beta.threads.messages.create(
            thread_id=thread_id,
            role="user",
            content=combined_input
        )

        event_handler = EventHandler()

        with client.beta.threads.runs.stream(
            thread_id=thread_id,
            assistant_id=assistant.id,
            instructions="Please address the user as Harsh. The user has a premium account. " + instructions,
            event_handler=event_handler
        ) as stream:
            stream.until_done()

        response_text = ''.join(event_handler.responses)

        return jsonify({"response": response_text})

    except Exception:
        return jsonify({"response": "An error occurred while processing your request."}), 500
    finally:
        session_lock.release()

@app.route('/upload_file', methods=['POST'])
def upload_file():
    data = request.json
    file_content = data.get("file_content")
    file_name = data.get("file_name")

    if not file_content or not file_name:
        return jsonify({"response": "Incomplete file data received."}), 400

    file_id = str(uuid.uuid4())

    thread_id = session.get('thread_id')
    if not thread_id:
        thread = client.beta.threads.create()
        thread_id = thread.id
        session['thread_id'] = thread_id
        session['attachments'] = []
    else:
        try:
            thread = client.beta.threads.retrieve(thread_id)
        except Exception:
            thread = client.beta.threads.create()
            thread_id = thread.id
            session['thread_id'] = thread_id
            session['attachments'] = []

    attachments = session.get('attachments', [])

    attachments.append({
        "id": file_id,
        "name": file_name,
        "content": file_content
    })
    session['attachments'] = attachments

    return jsonify({"response": "File uploaded successfully.", "file_id": file_id, "file_name": file_name})

@app.route('/list_files', methods=['GET'])
def list_files():
    attachments = session.get('attachments', [])
    valid_attachments = [file for file in attachments if isinstance(file, dict) and 'id' in file and 'name' in file]
    session['attachments'] = valid_attachments
    files = [{"id": file['id'], "name": file['name']} for file in valid_attachments]
    return jsonify({"files": files})

@app.route('/remove_file', methods=['POST'])
def remove_file():
    data = request.json
    file_id = data.get("file_id")

    if not file_id:
        return jsonify({"response": "No file ID provided."}), 400

    attachments = session.get('attachments', [])

    file_to_remove = next((file for file in attachments if file['id'] == file_id), None)

    if not file_to_remove:
        return jsonify({"response": "File not found."}), 404

    attachments = [file for file in attachments if file['id'] != file_id]
    session['attachments'] = attachments

    return jsonify({"response": "File removed successfully.", "file_name": file_to_remove['name']})

@app.route('/reset_session', methods=['POST'])
def reset_session():
    session.pop('thread_id', None)
    session.pop('attachments', None)
    return jsonify({"response": "Session has been reset."})

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