File size: 5,022 Bytes
6a0be2a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 | import gradio as gr
import os
import time
from openai import OpenAI
from dotenv import load_dotenv
# --- Load Environment Variables and Initialize Client ---
load_dotenv()
client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
# --- Configuration ---
ASSISTANT_INSTRUCTIONS = """
You are an expert data analyst. When the user uploads a file and asks a question,
your role is to use your Code Interpreter tool to write and run Python code to answer the question.
- Analyze the data thoroughly.
- If the user asks for a visualization, create it and display it as an image in your response.
- If the user doesn't asks for visualization, just make the analysis or whatever user ask.
- Provide a text explanation of your findings and describe any visualizations you create.
- Use Markdown formatting for all text in your output.This includes headings, bullet points, code blocks (if any), and emphasis for clarity.
"""
PLOT_FILENAME = "assistant_plot.png"
# --- Assistant Logic ---
def analyze_with_assistant(file_obj, user_prompt):
"""
Orchestrates the OpenAI Assistant to analyze a user-provided file.
"""
# 1. Input Validation
if file_obj is None:
return "Error: Please upload a file first.", None
if not user_prompt:
return "Error: Please enter a question.", None
# Cleanup old plot file
if os.path.exists(PLOT_FILENAME):
os.remove(PLOT_FILENAME)
try:
# 2. Upload the user's file to OpenAI
uploaded_file = client.files.create(
file=open(file_obj.name, "rb"),
purpose="assistants"
)
# 3. Create an Assistant
# This defines the agent's behavior and tools.
assistant = client.beta.assistants.create(
name="Data Analyst Assistant",
instructions=ASSISTANT_INSTRUCTIONS,
tools=[{"type": "code_interpreter"}],
model="gpt-4-turbo",
tool_resources={"code_interpreter": {"file_ids": [uploaded_file.id]}}
)
# 4. Create a Thread (a conversation session)
thread = client.beta.threads.create()
# 5. Add the User's Message to the Thread
client.beta.threads.messages.create(
thread_id=thread.id,
role="user",
content=user_prompt
)
# 6. Run the Assistant
run = client.beta.threads.runs.create(
thread_id=thread.id,
assistant_id=assistant.id,
)
# 7. Wait for the Run to complete
while run.status in ['queued', 'in_progress', 'cancelling']:
time.sleep(1)
run = client.beta.threads.runs.retrieve(thread_id=thread.id, run_id=run.id)
if run.status != 'completed':
return f"Error: The Assistant run failed with status {run.status}", None
# 8. Retrieve and Process Messages
messages = client.beta.threads.messages.list(thread_id=thread.id)
assistant_response = ""
image_file_id = None
# The latest message is from the assistant
for content_part in messages.data[0].content:
if content_part.type == "text":
assistant_response += content_part.text.value
elif content_part.type == "image_file":
image_file_id = content_part.image_file.file_id
# 9. Download the image if it exists
image_path = None
if image_file_id:
image_data = client.files.content(image_file_id)
image_data_bytes = image_data.read()
with open(PLOT_FILENAME, "wb") as f:
f.write(image_data_bytes)
image_path = PLOT_FILENAME
return assistant_response, image_path
except Exception as e:
error_message = f"An unexpected error occurred: {str(e)}"
print(error_message)
return error_message, None
# --- Gradio UI (Largely unchanged) ---
with gr.Blocks(theme=gr.themes.Soft(primary_hue="orange")) as demo:
gr.Markdown(
"""
# 🤖 Agentic Data Analysis with OpenAI Assistants
This version uses the official OpenAI Assistants API with the Code Interpreter tool.
"""
)
with gr.Row():
with gr.Column(scale=1):
file_input = gr.File(label="Upload your CSV", file_types=[".csv"])
text_input = gr.Textbox(
label="What would you like to know?",
placeholder="e.g., 'What is the correlation between column A and B?' or 'Create a bar chart of sales by category.'"
)
submit_button = gr.Button("🚀 Run Analysis", variant="primary")
with gr.Column(scale=2):
text_output = gr.Markdown(label="📝 Agent's Answer")
plot_output = gr.Image(label="📊 Generated Visualization", type="filepath")
submit_button.click(
fn=analyze_with_assistant,
inputs=[file_input, text_input],
outputs=[text_output, plot_output]
)
if __name__ == "__main__":
demo.launch()
|