Update app.py

app.py

@@ -20,9 +20,13 @@ nebius_client = OpenAI(
     base_url=NEBIUS_BASE_URL,
 )
 
-# --- 2. DEFINE TOOL FUNCTIONS (Standard Python) ---
+# Helper to convert bytes to PIL
+def bytes_to_pil(image_bytes: bytes):
+    return Image.open(io.BytesIO(image_bytes))
 
-def web_search(query: str):
+# --- 2. DEFINE TOOLS (Used by Agent AND exposed to MCP) ---
+
+def tool_web_search(query: str):
     """
     Performs a market research search using DuckDuckGo.
     """
@@ -36,7 +40,7 @@ def web_search(query: str):
     except Exception as e:
         return f"Search Error: {str(e)}"
 
-def generate_image(prompt: str):
+def tool_generate_image(prompt: str):
     """
     Generates an image and returns the local file path.
     """
@@ -48,6 +52,7 @@ def generate_image(prompt: str):
             n=1,
             size="1024x1024",
             response_format="b64_json",
+            style="flat",
         )
         b64_data = response.data[0].b64_json
         image_bytes = base64.b64decode(b64_data)
@@ -58,102 +63,192 @@ def generate_image(prompt: str):
         img.save(filename)
         return filename
     except Exception as e:
-        return None # Return None on error for image components
+        print(f"❌ Image Tool Error: {e}")
+        return None
+
+# Text generation helper (Internal only, usually not exposed as MCP tool directly)
+def tool_generate_text(prompt: str, system_message: str = "You are a helpful assistant.") -> str:
+    try:
+        response = nebius_client.chat.completions.create(
+            messages=[
+                {"role": "system", "content": system_message},
+                {"role": "user", "content": prompt},
+            ],
+            model=TEXT_MODEL,
+            max_tokens=1024,
+            temperature=0.7,
+        )
+        return response.choices[0].message.content
+    except Exception as e:
+        return f"Error: {e}"
 
-# --- 3. INTERNAL AGENT LOGIC ---
+# --- 3. AGENT LOGIC ---
 
 def run_autonomous_agent(company_name, business_type, style, progress=gr.Progress()):
-    log = f"🚀 Agent Started for '{company_name}'\n"
+    """Autonomous branding agent logic."""
     
-    # 1. Planning
-    progress(0.1, desc="Step 1: Planning...")
-    plan_prompt = f"Create a 5-color palette for '{company_name}' ({business_type}), style: {style}. Return hex codes."
+    # PHASE 1: PLANNING
+    progress(0.1, desc="Step 1: Planning & Color Palette...")
+    color_prompt = (
+        f"Create a distinct, 5-color palette for a {business_type} named '{company_name}' "
+        f"with a {style} style. Return ONLY a comma-separated list of hex codes."
+    )
+    colors_text = tool_generate_text(
+        color_prompt, "You are a design expert. Output only the requested list."
+    ).strip().strip('.')
+    
+    log = f"✅ Plan Started for {company_name}\n"
+    log += f"🎨 Palette Decided: {colors_text}\n"
+
+    # PHASE 2: LOGO
+    progress(0.3, desc="Step 2: Generating Logo...")
+    logo_prompt = (
+        f"A professional, minimalist logo for a {business_type} named '{company_name}'. "
+        f"Use ONLY these colors: {colors_text}. Vector flat style, centered, no text, no Transparent background."
+    )
+    logo_path = tool_generate_image(logo_prompt)
+    log += "🖼️ Logo Generated.\n"
+
+    # PHASE 3: ASSET
+    progress(0.5, desc="Step 3: Generating Brand Guide...")
+    asset_prompt = (
+        f"A professional color palette guide with 5 circular swatches, arranged in a row. "
+        f"Use these colors: {colors_text}. Minimalist white background."
+    )
+    asset_path = tool_generate_image(asset_prompt)
+    log += "🖼️ Brand Asset Generated.\n"
+
+    # PHASE 4: SEARCH
+    progress(0.7, desc="Step 4: Researching Market Trends...")
+    search_query = f"marketing trends for {business_type} 2025"
+    search_data = tool_web_search(search_query)
+    log += f"🔍 Search Query: {search_query}\n"
+    log += f"📊 Research: {search_data[:100]}...\n"
+
+    # PHASE 5: COPY
+    progress(0.9, desc="Step 5: Writing Content...")
+    copy_prompt = f"""
+    Write 3 engaging Twitter/X posts for a new {business_type} called '{company_name}'.
+    Tone: {style}.
+    Use these market insights if helpful:
+    {search_data}
+    """
+    social_copy = tool_generate_text(copy_prompt, "You are a professional social media manager.")
+    log += "✍️ Social Copy Written.\n✅ Execution Complete."
+
+    return logo_path, asset_path, social_copy, log
+
+# --- 4. CHATBOT LOGIC ---
+
+def chat_response(message, history):
+    system_message = "You are a creative brand consultant."
+    messages = [{"role": "system", "content": system_message}]
+    for msg in history:
+        messages.append({"role": msg["role"], "content": msg["content"]})
+    messages.append({"role": "user", "content": message})
+
     try:
-        resp = nebius_client.chat.completions.create(
-            messages=[{"role":"user", "content":plan_prompt}], model=TEXT_MODEL
+        response = nebius_client.chat.completions.create(
+            messages=messages,
+            model=TEXT_MODEL,
+            max_tokens=300,
         )
-        palette = resp.choices[0].message.content
-        log += f"🎨 Palette: {palette}\n"
+        bot_reply = response.choices[0].message.content
+        history.append({"role": "user", "content": message})
+        history.append({"role": "assistant", "content": bot_reply})
+        return history, ""
     except Exception as e:
-        palette = "Blue, White"
-        log += f"⚠️ Planning Error: {e}\n"
-
-    # 2. Generate Logo (Direct Call)
-    progress(0.3, desc="Step 2: Logo...")
-    logo_path = generate_image(f"Minimalist logo for {company_name}, colors: {palette}. White background.")
-    log += "🖼️ Logo Created.\n"
-
-    # 3. Research (Direct Call)
-    progress(0.6, desc="Step 3: Research...")
-    research = web_search(f"marketing trends {business_type} 2025")
-    log += f"📊 Research: {research[:50]}...\n"
-
-    # 4. Copywriting
-    progress(0.8, desc="Step 4: Copywriting...")
-    copy_prompt = f"Write a tweet for {company_name} based on: {research}"
-    try:
-        resp = nebius_client.chat.completions.create(
-            messages=[{"role":"user", "content":copy_prompt}], model=TEXT_MODEL
-        )
-        copy = resp.choices[0].message.content
-    except:
-        copy = "Error generating copy."
-    
-    log += "✅ Done."
-    return logo_path, copy, log
+        history.append({"role": "assistant", "content": f"Error: {e}"})
+        return history, ""
+
+# --- 5. UI CONSTRUCTION ---
+
+theme = gr.themes.Soft(
+    primary_hue="indigo",
+    secondary_hue="blue",
+    neutral_hue="slate",
+).set(
+    button_primary_background_fill="linear-gradient(90deg, #6366f1 0%, #4338ca 100%)",
+    button_primary_background_fill_hover="linear-gradient(90deg, #4f46e5 0%, #3730a3 100%)",
+    button_primary_text_color="white",
+    block_title_text_weight="600",
+    block_shadow="0 4px 6px rgba(0,0,0,0.1)",
+)
 
-# --- 4. UI & MCP REGISTRATION ---
+custom_css = """#main-header {text-align: center; margin-bottom: 2rem;}"""
 
-with gr.Blocks(title="AutoBrand MCP") as demo:
-    gr.Markdown("# 🚀 AutoBrand Agent + MCP Server")
+with gr.Blocks(theme=theme, css=custom_css, title="AutoBrand MCP Studio") as demo:
     
-    # --- VISIBLE UI ---
+    with gr.Column(elem_id="main-header"):
+        gr.Markdown("# 🖼️ AutoBrand Studio (MCP Enabled)")
+        gr.Markdown("Autonomous Agent powered by AI Models & Web Search. Connect Claude Desktop to use these tools!")
+
     with gr.Row():
-        with gr.Column():
-            c_input = gr.Textbox(label="Company Name", value="Lumina")
-            t_input = gr.Textbox(label="Type", value="Candles")
-            s_input = gr.Textbox(label="Style", value="Minimal")
-            btn = gr.Button("Run Agent", variant="primary")
-            log_box = gr.TextArea(label="Logs")
-        
-        with gr.Column():
-            out_img = gr.Image(label="Logo", type="filepath")
-            out_txt = gr.TextArea(label="Social Copy")
-
-    # Connect the UI Agent
-    btn.click(
-        run_autonomous_agent, 
-        [c_input, t_input, s_input], 
-        [out_img, out_txt, log_box]
+        # LEFT PANE
+        with gr.Column(scale=1):
+            gr.Markdown("### 🚀 Brand Configuration")
+            company_input = gr.Textbox(label="Company Name", value="Lumina")
+            type_input = gr.Textbox(label="Business Type", value="Organic Candle Shop")
+            style_input = gr.Textbox(label="Style / Vibe", value="Minimalist, calming")
+            
+            generate_btn = gr.Button("✨ Generate Brand Kit", variant="primary")
+            
+            gr.Markdown("### ⚙️ System Logs")
+            out_log = gr.TextArea(label="Agent Log", lines=15, show_copy_button=True)
+
+        # RIGHT PANE
+        with gr.Column(scale=2):
+            with gr.Tabs():
+                with gr.TabItem("🎨 Brand Identity"):
+                    with gr.Row():
+                        out_logo = gr.Image(label="Logo Concept", height=350, type="filepath")
+                        out_asset = gr.Image(label="Color Palette", height=350, type="filepath")
+                    
+                    out_copy = gr.TextArea(
+                        label="Generated Social Content",
+                        lines=15,
+                        show_copy_button=True,
+                        interactive=False,
+                    )
+                
+                with gr.TabItem("💬 Brand Consultant"):
+                    chatbot = gr.Chatbot(label="AI Consultant", height=650, type="messages")
+                    chat_input = gr.Textbox(show_label=False, placeholder="Ask something...")
+                    chat_btn = gr.Button("Send")
+
+    # --- AGENT & CHATBOT HANDLERS ---
+    generate_btn.click(
+        run_autonomous_agent,
+        [company_input, type_input, style_input],
+        [out_logo, out_asset, out_copy, out_log]
     )
+    chat_btn.click(chat_response, [chat_input, chatbot], [chatbot, chat_input])
+    chat_input.submit(chat_response, [chat_input, chatbot], [chatbot, chat_input])
 
     # --- HIDDEN MCP TOOL REGISTRATION ---
-    # To expose functions to MCP, we attach them to hidden triggers with `api_name`.
-    # This tells Gradio: "These are API endpoints" -> which enables MCP.
-    
+    # This section registers the functions as MCP tools without showing them in the UI.
     with gr.Row(visible=False):
         # Tool 1: Web Search
-        search_input = gr.Textbox()
-        search_output = gr.Textbox()
-        search_btn = gr.Button("Search Tool")
-        search_btn.click(
-            web_search, 
-            inputs=search_input, 
-            outputs=search_output, 
-            api_name="web_search"  # <--- THIS MAKES IT AN MCP TOOL
+        search_in = gr.Textbox()
+        search_out = gr.Textbox()
+        btn_search = gr.Button("Search Tool")
+        btn_search.click(
+            tool_web_search, 
+            inputs=search_in, 
+            outputs=search_out, 
+            api_name="web_search"  # Exposes tool to MCP
         )
 
         # Tool 2: Image Gen
-        img_input = gr.Textbox()
-        img_output = gr.Image(type="filepath")
-        img_btn = gr.Button("Image Tool")
-        img_btn.click(
-            generate_image, 
-            inputs=img_input, 
-            outputs=img_output, 
-            api_name="generate_image" # <--- THIS MAKES IT AN MCP TOOL
+        img_in = gr.Textbox()
+        img_out = gr.Image(type="filepath")
+        btn_img = gr.Button("Image Tool")
+        btn_img.click(
+            tool_generate_image, 
+            inputs=img_in, 
+            outputs=img_out, 
+            api_name="generate_image" # Exposes tool to MCP
         )
 
 if __name__ == "__main__":
-    # Launch with mcp_server=True
     demo.launch(ssr_mode=False, mcp_server=True)