from dotenv import load_dotenv
load_dotenv()
import os

import gradio as gr
import time

from openai import OpenAI

client = OpenAI()


directory = './documents'
file_ids = []
for filename in os.listdir(directory):
    if filename.endswith(".pdf"):  # Assuming you're only interested in PDF files
        file_path = os.path.join(directory, filename)
        with open(file_path, 'rb') as file:
            uploaded_file = client.files.create(file=file, purpose='assistants')
            file_ids.append(uploaded_file.id)

assistant = client.beta.assistants.create(
    instructions="You are an expert consultant for SBIR grant proposals. Your retrievable files include SBIR proposal guidelines as well as several real SBIR grant applications you can use as examples to guide the user.",
    model="gpt-4-1106-preview",
    tools=[{"type": "retrieval"}],
    file_ids=file_ids
)

thread = client.beta.threads.create()

def chat(prompt):
    message = client.beta.threads.messages.create(
      thread_id=thread.id,
      role="user",
      content=prompt
    )
    run = client.beta.threads.runs.create(
      thread_id=thread.id,
      assistant_id=assistant.id,
    )
    while run.status != "completed":
      time.sleep(0.1)
      run = client.beta.threads.runs.retrieve(
        thread_id=thread.id,
        run_id=run.id
      )
    
    messages = client.beta.threads.messages.list(
      thread_id=thread.id
    )
    return messages.data[0].content[0].text.value

# Define Gradio interface
iface = gr.Interface(
    fn=chat,
    inputs="text",
    outputs=gr.Markdown(),
    title="💬 SBIR Consultant",
    description="🚀 An SBIR Grant Application Consultant"
)

iface.launch()