Update app.py

app.py

@@ -1,134 +1,196 @@
-import os, json
+import os
+import io
+import json
+import base64
+from typing import Any, Dict, List, Tuple
+
 import gradio as gr
-from openai import OpenAI
 import plotly.graph_objects as go
+from PIL import Image
+from openai import OpenAI
 
-client = OpenAI()
 
 SYSTEM_PROMPT = """
 You are OmniGen AI.
 
-You must choose ONLY ONE output mode:
+Decide the single best response mode for the user's request.
+
+Return one of these formats only:
 
-1. IMAGE → if visual needed
-Return JSON:
-{"type":"image","prompt":"detailed prompt for image generation"}
+For text:
+{"type":"text","content":"<helpful answer>"}
 
-2. CHART → if data visualization needed
-Return JSON:
-{"type":"chart","title":"...","data":[{"x":[...],"y":[...],"label":"..."}]}
+For image:
+{"type":"image","prompt":"<clear image generation prompt>"}
 
-3. TEXT → otherwise
-Return plain text only (no JSON)
+For chart:
+{"type":"chart","title":"<chart title>","data":[{"x":[...],"y":[...],"label":"<series label>"}]}
 
-Do NOT explain your choice.
+Rules:
+- Return valid JSON only.
+- Choose exactly one type.
+- Use chart only when a chart is genuinely appropriate.
+- Use image only when a visual is genuinely appropriate.
+- Otherwise use text.
+- Do not wrap output in markdown fences.
 """
 
-def build_messages(history, user_msg):
+
+def build_messages(history: List[Dict[str, Any]], user_msg: str) -> List[Dict[str, str]]:
     messages = [{"role": "system", "content": SYSTEM_PROMPT}]
-    for u, a in history:
-        messages.append({"role": "user", "content": u})
-        messages.append({"role": "assistant", "content": a})
+    for item in history:
+        role = item.get("role")
+        content = item.get("content")
+        if role in {"user", "assistant"} and isinstance(content, str):
+            messages.append({"role": role, "content": content})
     messages.append({"role": "user", "content": user_msg})
     return messages
 
 
-def run_agent(api_key, user_msg, history, mode):
-    if not api_key:
+def make_client(api_key: str) -> OpenAI:
+    if not api_key or not api_key.strip():
         raise gr.Error("Please enter your OpenAI API key.")
+    return OpenAI(api_key=api_key.strip())
 
-    client.api_key = api_key
-    history = history or []
-
-    messages = build_messages(history, user_msg)
 
+def generate_router_response(client: OpenAI, messages: List[Dict[str, str]]) -> Dict[str, Any]:
     response = client.chat.completions.create(
         model="gpt-4.1",
         messages=messages,
-        temperature=0.7,
+        temperature=0.5,
+        response_format={"type": "json_object"},
+    )
+    content = response.choices[0].message.content or ""
+    try:
+        parsed = json.loads(content)
+    except json.JSONDecodeError as exc:
+        raise gr.Error(f"Model returned invalid JSON: {exc}") from exc
+
+    if "type" not in parsed:
+        raise gr.Error("Model response did not include a type field.")
+    return parsed
+
+
+def generate_image(client: OpenAI, prompt: str) -> Image.Image:
+    img_response = client.images.generate(
+        model="gpt-image-1",
+        prompt=prompt,
+        size="1024x1024",
     )
 
-    content = response.choices[0].message.content.strip()
+    image_base64 = img_response.data[0].b64_json
+    if not image_base64:
+        raise gr.Error("Image generation returned no image data.")
 
-    img, fig = None, None
+    image_bytes = base64.b64decode(image_base64)
+    return Image.open(io.BytesIO(image_bytes))
 
-    try:
-        parsed = json.loads(content)
 
-        if parsed["type"] == "image":
-            img_resp = client.images.generate(
-                model="gpt-image-1",
-                prompt=parsed["prompt"],
-                size="1024x1024"
+def generate_chart(chart_payload: Dict[str, Any]) -> go.Figure:
+    fig = go.Figure()
+    series_list = chart_payload.get("data", [])
+
+    if not isinstance(series_list, list) or not series_list:
+        raise gr.Error("Chart response did not include usable data.")
+
+    for series in series_list:
+        x = series.get("x", [])
+        y = series.get("y", [])
+        label = series.get("label", "Series")
+        fig.add_trace(
+            go.Scatter(
+                x=x,
+                y=y,
+                mode="lines+markers",
+                name=label,
             )
-            img = img_resp.data[0].url
-            history.append([user_msg, f"![image]({img})"])
+        )
+
+    fig.update_layout(
+        title=chart_payload.get("title", "Generated Chart"),
+        template="plotly_white",
+        margin=dict(l=40, r=40, t=60, b=40),
+    )
+    return fig
+
+
+def respond(api_key: str, user_msg: str, history: List[Dict[str, Any]]):
+    if not user_msg or not user_msg.strip():
+        raise gr.Error("Please enter a prompt.")
+
+    history = history or []
+    client = make_client(api_key)
+    messages = build_messages(history, user_msg.strip())
+    routed = generate_router_response(client, messages)
 
-        elif parsed["type"] == "chart":
-            fig = go.Figure()
-            for s in parsed["data"]:
-                fig.add_trace(go.Scatter(
-                    x=s["x"],
-                    y=s["y"],
-                    mode="lines+markers",
-                    name=s.get("label", "")
-                ))
-            fig.update_layout(title=parsed.get("title", "Chart"))
-            history.append([user_msg, parsed.get("title", "Chart generated")])
+    result_type = routed.get("type")
+    image_output = None
+    chart_output = None
 
-        else:
-            history.append([user_msg, content])
+    history.append({"role": "user", "content": user_msg.strip()})
 
-    except:
-        history.append([user_msg, content])
+    if result_type == "text":
+        text = routed.get("content", "").strip() or "No text was generated."
+        history.append({"role": "assistant", "content": text})
 
-    return history, img, fig
+    elif result_type == "image":
+        prompt = routed.get("prompt", "").strip()
+        if not prompt:
+            raise gr.Error("Image mode selected, but no prompt was returned.")
+        image_output = generate_image(client, prompt)
+        history.append({"role": "assistant", "content": f"Generated an image for: {prompt}"})
 
+    elif result_type == "chart":
+        chart_output = generate_chart(routed)
+        title = routed.get("title", "Generated Chart")
+        history.append({"role": "assistant", "content": f"Generated chart: {title}"})
 
-with gr.Blocks(css="style.css") as demo:
+    else:
+        raise gr.Error(f"Unsupported response type: {result_type}")
 
+    return history, history, image_output, chart_output, ""
+
+
+with gr.Blocks(css_paths="style.css") as demo:
     gr.Markdown(
-        "# 🧠 OmniGen AI Studio\n"
-        "Generate text, images, or charts from one prompt.\n"
-        "Powered by GPT-4.1"
+        """
+# 🧠 OmniGen AI Studio
+Generate text, images, or charts from a single prompt using GPT-4.1.
+"""
     )
 
-    with gr.Row():
-        api_key = gr.Textbox(label="OpenAI API Key", type="password")
+    api_key = gr.Textbox(
+        label="OpenAI API Key",
+        type="password",
+        placeholder="sk-...",
+    )
 
-    chatbot = gr.Chatbot(height=400)
+    chatbot = gr.Chatbot(label="Conversation", type="messages", height=420)
 
     with gr.Row():
-        user_msg = gr.Textbox(label="Your Prompt", scale=4)
-        mode = gr.Radio(
-            ["Auto", "Force Text", "Force Image", "Force Chart"],
-            value="Auto",
-            label="Mode"
+        user_msg = gr.Textbox(
+            label="Your Prompt",
+            placeholder="Ask for a blog paragraph, a logo concept, or a chart of quarterly revenue...",
+            scale=5,
         )
-        send = gr.Button("Generate", variant="primary")
-
-    img_out = gr.Image(label="Generated Image")
-    chart_out = gr.Plot(label="Generated Chart")
+        send_btn = gr.Button("Generate", variant="primary")
 
-    def respond(api_key, user_msg, history, mode):
-        history, img, fig = run_agent(api_key, user_msg, history, mode)
+    with gr.Row():
+        image_out = gr.Image(label="Generated Image", type="pil")
+        chart_out = gr.Plot(label="Generated Chart")
 
-        return (
-            history,
-            gr.update(value=img) if img else gr.update(value=None),
-            gr.update(value=fig) if fig else gr.update(value=None),
-        )
+    state = gr.State([])
 
-    send.click(
-        respond,
-        inputs=[api_key, user_msg, chatbot, mode],
-        outputs=[chatbot, img_out, chart_out]
+    send_btn.click(
+        fn=respond,
+        inputs=[api_key, user_msg, state],
+        outputs=[chatbot, state, image_out, chart_out, user_msg],
     )
 
     user_msg.submit(
-        respond,
-        inputs=[api_key, user_msg, chatbot, mode],
-        outputs=[chatbot, img_out, chart_out]
+        fn=respond,
+        inputs=[api_key, user_msg, state],
+        outputs=[chatbot, state, image_out, chart_out, user_msg],
     )