Update app.py

app.py

@@ -10,7 +10,7 @@ from google.genai import types
 from PIL import Image
 
 # -------------------------------------------------------------------
-# Config
+# App Metadata
 # -------------------------------------------------------------------
 
 APP_TITLE = "ZEN AI Co. Module 2 | Agent Assembler"
@@ -20,12 +20,11 @@ and create images using GPT-5, Gemini 2.5 Pro, Gemini 3 Pro, Nano Banana,
 Nano Banana Pro, and DALL·E 3.
 """
 
-# Reasonable defaults if user doesn't touch sliders
 DEFAULT_TEMPERATURE = 0.6
 DEFAULT_MAX_TOKENS = 1024
 
 # -------------------------------------------------------------------
-# Helpers: API clients
+# API Clients
 # -------------------------------------------------------------------
 
 def get_openai_client(key_override: Optional[str] = None) -> OpenAI:
@@ -33,14 +32,12 @@ def get_openai_client(key_override: Optional[str] = None) -> OpenAI:
     Returns an OpenAI client using either:
     1) key from the UI override, or
     2) OPENAI_API_KEY environment variable.
-
-    This satisfies the “two places for API keys” requirement.
     """
     api_key = (key_override or "").strip() or os.getenv("OPENAI_API_KEY", "").strip()
     if not api_key:
         raise ValueError(
             "OpenAI API key missing. "
-            "Either set OPENAI_API_KEY env var or paste it in the sidebar."
+            "Set OPENAI_API_KEY env var or paste it into the sidebar."
         )
     return OpenAI(api_key=api_key)
 
@@ -55,12 +52,12 @@ def get_google_client(key_override: Optional[str] = None) -> genai.Client:
     if not api_key:
         raise ValueError(
             "Google Gemini API key missing. "
-            "Either set GOOGLE_API_KEY env var or paste it in the sidebar."
+            "Set GOOGLE_API_KEY env var or paste it into the sidebar."
         )
     return genai.Client(api_key=api_key)
 
 # -------------------------------------------------------------------
-# Helpers: Prompt & style shaping
+# Prompt & Style Helpers
 # -------------------------------------------------------------------
 
 def build_system_instructions(
@@ -70,21 +67,27 @@ def build_system_instructions(
     tone: str,
 ) -> str:
     """
-    Builds a strong system prompt that shapes behavior according to theme,
-    output mode, and tone.
+    Builds a system prompt that reflects theme, output mode, and tone.
     """
     theme_map = {
         "ZEN Dark": "Use a sleek, modern, slightly futuristic tone. Be concise but high signal.",
         "ZEN Light": "Use a clear, friendly, educational tone suitable for learners of all ages.",
         "Research / Technical": "Write like a senior research engineer: rigorous, structured, and explicit.",
-        "Youth AI Pioneer": "Explain things in simple, motivating language suitable for ages 11–18, "
-                            "but never dumb it down.",
+        "Youth AI Pioneer": (
+            "Explain things in simple, motivating language for ages 11–18, "
+            "but keep the concepts accurate and serious."
+        ),
     }
 
     output_map = {
         "Standard Chat": "Respond like a normal assistant, but keep paragraphs tight and skimmable.",
-        "Executive Report": "Respond as a structured executive brief with headings, bullets, and 1–2 sentence insights.",
-        "Infographic Outline": "Respond as a bullet-point infographic blueprint with short, punchy lines and clear sections.",
+        "Executive Report": (
+            "Respond as a structured executive brief with headings, bullets, "
+            "and 1–2 sentence insights per section."
+        ),
+        "Infographic Outline": (
+            "Respond as a bullet-point infographic blueprint with short, punchy lines and clear sections."
+        ),
         "Bullet Summary": "Respond as a compact bullet summary with 5–10 bullets max.",
     }
 
@@ -112,7 +115,7 @@ def history_to_messages(
     system_instructions: str,
 ) -> List[dict]:
     """
-    Converts Gradio Chatbot history into OpenAI-style messages.
+    Convert chat history into OpenAI chat messages.
     """
     messages: List[dict] = []
     if system_instructions:
@@ -134,7 +137,7 @@ def history_to_gemini_prompt(
     system_instructions: str,
 ) -> str:
     """
-    Flattens history into a single text prompt for Gemini.
+    Flatten history into a single text prompt for Gemini.
     """
     lines = []
     if system_instructions:
@@ -151,7 +154,7 @@ def history_to_gemini_prompt(
     return "\n\n".join(lines)
 
 # -------------------------------------------------------------------
-# Helpers: Model calls (text)
+# Text Model Calls
 # -------------------------------------------------------------------
 
 def call_openai_text(
@@ -162,7 +165,7 @@ def call_openai_text(
 ) -> str:
     client = get_openai_client(openai_key)
     completion = client.chat.completions.create(
-        model="gpt-5",  # You can change to gpt-5.1 or whatever is available in your project
+        model="gpt-5",  # adjust to the exact GPT-5 variant available to you
         messages=messages,
         temperature=temperature,
         max_tokens=max_tokens,
@@ -199,7 +202,7 @@ def call_hybrid_text(
     max_tokens: int,
 ) -> str:
     """
-    Calls GPT-5 and Gemini (2.5 Pro or 3 Pro) and fuses their answers.
+    Call GPT-5 and a Gemini model, then fuse the outputs.
     """
     try:
         gpt_answer = call_openai_text(openai_key, messages, temperature, max_tokens)
@@ -222,7 +225,7 @@ def call_hybrid_text(
     return fused
 
 # -------------------------------------------------------------------
-# Helpers: Image generation
+# Image Generation Calls
 # -------------------------------------------------------------------
 
 def call_openai_dalle(
@@ -231,7 +234,7 @@ def call_openai_dalle(
     size: str = "1024x1024",
 ) -> Optional[Image.Image]:
     """
-    Uses DALL·E 3 via OpenAI Images API to generate a PIL image.
+    Use DALL·E 3 to generate a PIL image.
     """
     client = get_openai_client(openai_key)
     response = client.images.generate(
@@ -243,7 +246,6 @@ def call_openai_dalle(
     if not response.data:
         return None
 
-    # DALL·E responses can be URL or base64; here we handle base64
     img_data = response.data[0].b64_json
     img_bytes = base64.b64decode(img_data)
     return Image.open(BytesIO(img_bytes))
@@ -255,7 +257,7 @@ def call_gemini_image(
     prompt: str,
 ) -> Optional[Image.Image]:
     """
-    Uses Nano Banana (gemini-2.5-flash-image) or Nano Banana Pro 
+    Use Nano Banana (gemini-2.5-flash-image) or Nano Banana Pro 
     (gemini-3-pro-image-preview) via Google GenAI SDK.
     """
     client = get_google_client(google_key)
@@ -264,7 +266,7 @@ def call_gemini_image(
         contents=[prompt],
     )
 
-    # Follow pattern from official docs: walk parts for inline image data
+    # Walk candidate parts for inline image data
     for candidate in response.candidates:
         for part in candidate.content.parts:
             inline = getattr(part, "inline_data", None)
@@ -275,7 +277,7 @@ def call_gemini_image(
     return None
 
 # -------------------------------------------------------------------
-# Core chat function used by Gradio
+# Core Chat Logic
 # -------------------------------------------------------------------
 
 def agent_assembler_chat(
@@ -294,13 +296,13 @@ def agent_assembler_chat(
     image_backend: str,
 ) -> Tuple[List[Tuple[str, str]], Optional[Image.Image]]:
     """
-    Main callback for the app. Returns updated chat history & optional image.
+    Main callback: text + optional image generation.
     """
     if not user_message.strip():
         return chat_history, None
 
     base_system = (
-        "You are ZEN AI Co.'s **Agent Assembler**, a multi-model orchestrator. "
+        "You are ZEN AI Co.'s Agent Assembler, a multi-model orchestrator. "
         "You can:\n"
         "- Hold deep, contextual conversations about AI literacy, automation, and education.\n"
         "- Generate executive reports and structured briefs.\n"
@@ -320,11 +322,11 @@ def agent_assembler_chat(
         tone=tone,
     )
 
-    # Prepare conversations for both stacks
+    # Prepare inputs for OpenAI and Gemini
     messages = history_to_messages(chat_history, user_message, system_instructions)
     gemini_prompt = history_to_gemini_prompt(chat_history, user_message, system_instructions)
 
-    # Decide which text model(s) to call
+    # Route text generation
     if model_family == "OpenAI: GPT-5":
         ai_reply = call_openai_text(
             openai_key=openai_key_ui,
@@ -345,7 +347,7 @@ def agent_assembler_chat(
             temperature=temperature,
             max_tokens=max_tokens,
         )
-    else:  # Hybrid mode
+    else:  # Hybrid: GPT-5 + Gemini
         if gemini_model_choice == "Gemini 2.5 Pro":
             model_id = "gemini-2.5-pro"
         else:
@@ -361,18 +363,17 @@ def agent_assembler_chat(
             max_tokens=max_tokens,
         )
 
-    # Update chat history
+    # Update history with text
     chat_history = chat_history + [(user_message, ai_reply)]
 
     # Optional image generation
     generated_image: Optional[Image.Image] = None
     if generate_image:
-        # Build an image-oriented prompt from the last user query + output mode
         image_prompt = (
             f"{user_message.strip()}\n\n"
             f"Image intent: {output_mode}. "
-            "Render clean, readable text if any labels are required. "
-            "Use a style that would fit the ZEN AI Co. brand."
+            "Render clean, readable text if labels are required. "
+            "Use a visual style aligned with the ZEN AI Co. brand."
         )
 
         try:
@@ -393,7 +394,7 @@ def agent_assembler_chat(
                     prompt=image_prompt,
                 )
         except Exception as e:
-            # Append a note to the assistant message if image fails
+            # Attach error note to latest assistant message
             chat_history[-1] = (
                 chat_history[-1][0],
                 chat_history[-1][1]
@@ -403,7 +404,7 @@ def agent_assembler_chat(
     return chat_history, generated_image
 
 
-def clear_chat():
+def clear_chat() -> Tuple[List[Tuple[str, str]], Optional[Image.Image]]:
     return [], None
 
 # -------------------------------------------------------------------
@@ -416,11 +417,10 @@ def build_interface() -> gr.Blocks:
         gr.Markdown(APP_DESCRIPTION)
 
         with gr.Row():
-            # Left: Chat + image output
+            # Left side: chat + image output
             with gr.Column(scale=3):
                 chatbot = gr.Chatbot(
                     label="Agent Assembler Chat",
-                    type="messages",
                     height=520,
                 )
                 image_out = gr.Image(
@@ -430,7 +430,9 @@ def build_interface() -> gr.Blocks:
                 )
                 user_input = gr.Textbox(
                     label="Your message",
-                    placeholder="Ask for a chat, a report, an infographic outline, or an image...",
+                    placeholder=(
+                        "Ask for a chat, a report, an infographic outline, or an image..."
+                    ),
                     lines=3,
                 )
 
@@ -438,15 +440,15 @@ def build_interface() -> gr.Blocks:
                     send_btn = gr.Button("Send", variant="primary")
                     clear_btn = gr.Button("Clear")
 
-            # Right: Control panel
+            # Right side: control panel
             with gr.Column(scale=2):
                 gr.Markdown("## API Keys")
                 openai_key_ui = gr.Textbox(
-                    label="OpenAI API Key (optional, otherwise uses OPENAI_API_KEY env var)",
+                    label="OpenAI API Key (optional; otherwise uses OPENAI_API_KEY env var)",
                     type="password",
                 )
                 google_key_ui = gr.Textbox(
-                    label="Google Gemini API Key (optional, otherwise uses GOOGLE_API_KEY env var)",
+                    label="Google Gemini API Key (optional; otherwise uses GOOGLE_API_KEY env var)",
                     type="password",
                 )
 
@@ -496,7 +498,7 @@ def build_interface() -> gr.Blocks:
                     value="Neutral",
                 )
 
-                gr.Markdown("## Generation Controls")
+                gr.Markdown("## Text Generation Controls")
 
                 temperature = gr.Slider(
                     label="Temperature (creativity)",
@@ -531,10 +533,10 @@ def build_interface() -> gr.Blocks:
                     value="Nano Banana Pro (Gemini 3 Pro Image Preview)",
                 )
 
-        # State for chat history
+        # Shared chat state
         chat_state = gr.State([])
 
-        # Wire up events
+        # Send button wiring
         send_btn.click(
             fn=agent_assembler_chat,
             inputs=[
@@ -554,11 +556,12 @@ def build_interface() -> gr.Blocks:
             ],
             outputs=[chatbot, image_out],
         ).then(
-            fn=lambda h: (h, ""),  # update state + clear box
+            fn=lambda h: (h, ""),  # sync state, clear input
             inputs=chatbot,
             outputs=[chat_state, user_input],
         )
 
+        # Submit on Enter
         user_input.submit(
             fn=agent_assembler_chat,
             inputs=[
@@ -578,11 +581,12 @@ def build_interface() -> gr.Blocks:
             ],
             outputs=[chatbot, image_out],
         ).then(
-            fn=lambda h: (h, ""),  # update state + clear box
+            fn=lambda h: (h, ""),
             inputs=chatbot,
             outputs=[chat_state, user_input],
         )
 
+        # Clear button wiring
         clear_btn.click(
             fn=clear_chat,
             inputs=None,