Update app.py

app.py

@@ -1,67 +1,102 @@
+import base64, io
+from typing import Optional
+
 import gradio as gr
-import openai, os, asyncio
+from PIL import Image
+
+# OpenAI SDK v1 style
+from openai import OpenAI
 
-TITLE = "ZEN AI Playground — Ask, Learn, Create"
-TAGLINE = "A simple way to chat with GPT-5 and get an instant image suggestion."
+TITLE = "ZEN AI Playground — Beginner Edition"
+TAGLINE = "Paste one key. Pick a model. Get an answer (and an optional image)."
 
-def check_key(key):
-    if not key or not key.startswith("sk-"):
-        raise gr.Error("Please paste a valid OpenAI API key starting with 'sk-'.")
-    return key.strip()
+def ensure_key(key: str) -> str:
+    key = (key or "").strip()
+    if not key.startswith("sk-") or len(key) < 20:
+        raise gr.Error("Please paste a valid OpenAI API key (starts with 'sk-').")
+    return key
 
-async def ask_ai(api_key, question):
-    check_key(api_key)
-    openai.api_key = api_key
+def try_text(client: OpenAI, model: str, prompt: str) -> str:
+    """
+    Try a chat completion on the given model.
+    If that fails (e.g., model not available), fallback to gpt-4o.
+    """
     try:
-        # --- Text reply ---
-        text_resp = await asyncio.to_thread(
-            openai.ChatCompletion.create,
-            model="gpt-5",  # works if your account has access; fallback below otherwise
-            messages=[{"role": "user", "content": question}],
+        r = client.chat.completions.create(
+            model=model,
+            messages=[{"role": "user", "content": prompt}],
             temperature=0.7,
         )
-        answer = text_resp["choices"][0]["message"]["content"]
-
+        return r.choices[0].message.content
     except Exception as e:
-        # Fallback to gpt-4o if gpt-5 unavailable
-        try:
-            text_resp = await asyncio.to_thread(
-                openai.ChatCompletion.create,
+        # fallback if the selected model isn't available for the account
+        if model != "gpt-4o":
+            r = client.chat.completions.create(
                 model="gpt-4o",
-                messages=[{"role": "user", "content": question}],
+                messages=[{"role": "user", "content": prompt}],
                 temperature=0.7,
             )
-            answer = text_resp["choices"][0]["message"]["content"]
-        except Exception as ee:
-            raise gr.Error(f"OpenAI error: {ee}")
+            return r.choices[0].message.content
+        raise gr.Error(f"OpenAI text error: {e}")
 
-    # --- Optional Image generation ---
+def try_image(client: OpenAI, prompt: str) -> Optional[Image.Image]:
+    """
+    Generate an image with gpt-image-1. If not permitted, return None.
+    """
     try:
-        img_resp = await asyncio.to_thread(
-            openai.Image.create,
-            model="dall-e-3",
-            prompt=question,
-            n=1,
+        img = client.images.generate(
+            model="gpt-image-1",
+            prompt=prompt,
             size="512x512",
+            n=1,
         )
-        image_url = img_resp["data"][0]["url"]
+        b64 = img.data[0].b64_json
+        raw = base64.b64decode(b64)
+        return Image.open(io.BytesIO(raw)).convert("RGB")
     except Exception:
-        image_url = None
+        return None
+
+def run(api_key: str, model_choice: str, prompt: str):
+    api_key = ensure_key(api_key)
+    prompt = (prompt or "").strip()
+    if not prompt:
+        raise gr.Error("Please enter a prompt/question.")
 
-    return answer, image_url
+    client = OpenAI(api_key=api_key)
+
+    # 1) Text
+    answer = try_text(client, model_choice, prompt)
+
+    # 2) Image (best-effort)
+    image = try_image(client, prompt)
+
+    return answer, image
 
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown(f"## {TITLE}\n{TAGLINE}")
 
-    api_key = gr.Textbox(label="OpenAI API Key", type="password", placeholder="sk-...")
-    question = gr.Textbox(label="Ask anything about AI, science, or creativity", lines=3)
-    run_btn = gr.Button("Generate Response", variant="primary")
+    with gr.Row():
+        with gr.Column(scale=1):
+            key = gr.Textbox(label="OpenAI API Key", type="password", placeholder="sk-...")
+            model = gr.Dropdown(
+                label="Model",
+                choices=["gpt-4o", "gpt-4o-mini", "gpt-5"],
+                value="gpt-4o",
+                info="Use gpt-4o for widest access. Pick gpt-5 only if your account has it."
+            )
+        with gr.Column(scale=2):
+            prompt = gr.Textbox(
+                label="Ask anything",
+                placeholder="Explain how transformers work like I’m new to AI…",
+                lines=4
+            )
+            go = gr.Button("Generate", variant="primary")
 
     with gr.Row():
         text_out = gr.Markdown(label="AI Response")
-        img_out = gr.Image(label="Suggested Image", show_label=True)
+        image_out = gr.Image(label="Generated Image (optional)", type="pil")
 
-    run_btn.click(fn=ask_ai, inputs=[api_key, question], outputs=[text_out, img_out])
+    go.click(run, inputs=[key, model, prompt], outputs=[text_out, image_out])
 
 if __name__ == "__main__":
     demo.queue().launch()