app.py

import os
import gradio as gr
from gradio.utils import get_space
from huggingface_hub import InferenceClient
from e2b_code_interpreter import Sandbox
from pathlib import Path
from transformers import AutoTokenizer
import json
from openai import OpenAI
from huggingface_hub import HfApi, HfFolder
from jupyter_handler import JupyterNotebook

if not get_space():
    try:
        from dotenv import load_dotenv

        load_dotenv()
    except (ImportError, ModuleNotFoundError):
        pass


from utils import (
    run_interactive_notebook,
)

E2B_API_KEY = os.environ["E2B_API_KEY"]
HF_TOKEN = os.environ["HF_TOKEN"] #HfFolder.get_token() #
DEFAULT_MAX_TOKENS = 512
SANDBOXES = {}
SANDBOX_TIMEOUT = 5
TMP_DIR = './tmp/'
if not os.path.exists(TMP_DIR):
    os.makedirs(TMP_DIR)

with open(TMP_DIR+"jupyter-agent.ipynb", 'w', encoding='utf-8') as f:
    json.dump(JupyterNotebook().data, f, indent=2)

with open("ds-system-prompt.txt", "r") as f:
    DEFAULT_SYSTEM_PROMPT = f.read()
DEFAULT_SYSTEM_PROMPT = None

def execute_jupyter_agent(
    sytem_prompt, user_input, max_new_tokens, model, files, message_history, request: gr.Request
):
    if request.session_hash not in SANDBOXES:
        SANDBOXES[request.session_hash] = Sandbox(api_key=E2B_API_KEY, timeout=SANDBOX_TIMEOUT)
    sbx = SANDBOXES[request.session_hash]

    save_dir = os.path.join(TMP_DIR, request.session_hash)
    os.makedirs(save_dir, exist_ok=True)
    save_dir = os.path.join(save_dir, 'jupyter-agent.ipynb')

    client = OpenAI(
        base_url="https://router.huggingface.co/v1",
        api_key=HF_TOKEN,
    )

    tokenizer = AutoTokenizer.from_pretrained(model)
    # model = "meta-llama/Llama-3.1-8B-Instruct"

    filenames = []
    if files is not None:
        for filepath in files:
            filpath = Path(filepath)
            with open(filepath, "rb") as file:
                print(f"uploading {filepath}...")
                sbx.files.write(filpath.name, file)
                filenames.append(filpath.name)

    # Initialize message_history if it doesn't exist
    if len(message_history) == 0:
        message_history.append(
            {
                "role": "system",
                "content": sytem_prompt.format("- " + "\n- ".join(filenames)),
            }
        )
    message_history.append({"role": "user", "content": user_input})

    print("history:", message_history)

    for notebook_html, notebook_data, messages in run_interactive_notebook(
        client, model, tokenizer, message_history, sbx, max_new_tokens=max_new_tokens
    ):
        message_history = messages
        
        yield notebook_html, message_history, TMP_DIR+"jupyter-agent.ipynb"
    
    with open(save_dir, 'w', encoding='utf-8') as f:
        json.dump(notebook_data, f, indent=2)
    yield notebook_html, message_history, save_dir

def clear(msg_state):
    msg_state = []
    return JupyterNotebook().render(), msg_state


css = """
#component-0 {
    height: 100vh;
    overflow-y: auto;
    padding: 20px;
}

.gradio-container {
    height: 100vh !important;
}

.contain {
    height: 100vh !important;
}
"""


# Create the interface
with gr.Blocks() as demo:
    msg_state = gr.State(value=[])

    html_output = gr.HTML(value=JupyterNotebook().render())
    
    user_input = gr.Textbox(
        value="Write code to multiply three numbers: 10048, 32, 19", lines=3, label="User input"
        #value="Solve the Lotka-Volterra equation and plot the results. Do it step by step and explain what you are doing and in the end make a super nice and clean plot.", lines=3, label="User input"
    )

    with gr.Row():
        generate_btn = gr.Button("Let's go!")
        clear_btn = gr.Button("Clear")
    
    file = gr.File(TMP_DIR+"jupyter-agent.ipynb", label="Download Jupyter Notebook")
    
    with gr.Accordion("Upload files", open=False):
        files = gr.File(label="Upload files to use", file_count="multiple")

    with gr.Accordion("Advanced Settings", open=False):
        system_input = gr.Textbox(
            label="System Prompt",
            value=DEFAULT_SYSTEM_PROMPT,
            elem_classes="input-box",
            lines=8,
        )
        with gr.Row():
            max_tokens = gr.Number(
                label="Max New Tokens",
                value=DEFAULT_MAX_TOKENS,
                minimum=128,
                maximum=2048,
                step=8,
                interactive=True,
            )

            model = gr.Dropdown(
                value="Qwen/Qwen3-Coder-480B-A35B-Instruct",
                choices=[
                    "Qwen/Qwen3-Coder-30B-A3B-Instruct",
                    "Qwen/Qwen3-Coder-480B-A35B-Instruct",
                ],
                label="Models"
            )

    generate_btn.click(
        fn=execute_jupyter_agent,
        inputs=[system_input, user_input, max_tokens, model, files, msg_state],
        outputs=[html_output, msg_state, file],
    )

    clear_btn.click(fn=clear, inputs=[msg_state], outputs=[html_output, msg_state])

    demo.load(
        fn=None,
        inputs=None,
        outputs=None,
        js=""" () => {
    if (document.querySelectorAll('.dark').length) {
        document.querySelectorAll('.dark').forEach(el => el.classList.remove('dark'));
    }
}
"""
    )

demo.launch(ssr_mode=False)