import interpreter
import gradio as gr
import os
import io
import contextlib

# Use environ variables for saving and using OpenAI API key as below or use from a textbox in gradio UI
# interpreter.api_key = os.environ.get('openai_api_key')

interpreter.auto_run = True

def chat_with_interpreter(message, history, openai_api_key):
    interpreter.api_key = openai_api_key
    if message == 'reset':
      interpreter.reset()
    # Redirect stdout to capture the streamed output
    new_stdout = io.StringIO()
    with contextlib.redirect_stdout(new_stdout):
        interpreter.chat(message)
    output = new_stdout.getvalue()

    # Return this output so Gradio's ChatInterface can display it
    return output

openai_api_key = gr.Textbox(label='OpenAI API Key', intercative=True)
additional_inputs = [openai_api_key]
examples=[["what is 2+2?"], 
          ["Can you solve for x: 10x -65=0?"], 
          ["What are top 10 headlines from BBC from last week?"]
         ],
                        
demo = gr.ChatInterface(fn=chat_with_interpreter, 
                        title="Open-Interpreter Gradio ChatInterface",
                        description="Open Interpreter lets LLMs run code (Python, Javascript, Shell, and more) locally",
                        clear_btn=None,
                        retry_btn=None,
                        undo_btn=None,
                        #examples=examples,
                        additional_inputs=additional_inputs,
                        additional_inputs_accordion_name = "OpenAI API Key",
                        ).queue()
demo.launch(debug=True)