import gradio as gr
import openai
from openai import OpenAI
import os
import base64

# Set API key and organization ID from environment variables
openai.api_key = os.getenv('OPENAI_API_KEY')
openai.organization = os.getenv('OPENAI_ORG_ID')
client = OpenAI(api_key= os.getenv('OPENAI_API_KEY'), organization=os.getenv('OPENAI_ORG_ID'))

# Define the model to be used
MODEL = "gpt-4o"

def process_text(text_input):
    if text_input:
        completion = client.chat.completions.create(
            model=MODEL,
            messages=[
                {"role": "system", "content": "You are a helpful assistant. Help me with my math homework!"},
                {"role": "user", "content": f"Hello! Could you solve {text_input}?"}
            ]
        )
        return "Assistant: " + completion.choices[0].message.content

def process_image(image_input):
    if image_input is not None:
        with open(image_input.name, "rb") as f:
            base64_image = base64.b64encode(f.read()).decode("utf-8")
        response = client.chat.completions.create(
            model=MODEL,
            messages=[
                {"role": "system", "content": "You are a helpful assistant that responds in Markdown."},
                {"role": "user", "content": [
                    {"type": "text", "text": "Help me understand what is it"},
                    {"type": "image_url", "image_url": {
                        "url": f"data:image/png;base64,{base64_image}"}
                    }
                ]}
            ],
            temperature=0.0,
        )
        return response.choices[0].message.content

def main(text_input="", image_input=None):
    if text_input and image_input is None:
        return process_text(text_input)
    elif image_input is not None:
        return process_image(image_input)

iface = gr.Interface(fn=main, inputs=["text", gr.inputs.Image()], outputs="text")
iface.launch()