Update app.py

app.py

@@ -1,213 +1,219 @@
 import gradio as gr
-import httpx  # Included in requirements
-from openai import OpenAI # We use the OpenAI client to talk to Nebius!
+from openai import OpenAI
+from duckduckgo_search import DDGS
 import os
 
-# --- 1. CONFIGURE YOUR CLIENTS ---
+# --- 1. CONFIGURATION ---
 
-# Get your secret key from Hugging Face Space Secrets
 NEBIUS_API_KEY = os.environ.get("NEBIUS_API_KEY")
+# Ensure the URL is correct for the OpenAI-compatible endpoint
+NEBIUS_BASE_URL = "https://api.studio.nebius.ai/v1/" 
+TEXT_MODEL = "meta-llama/Llama-3.3-70B-Instruct" 
+IMAGE_MODEL = "black-forest-labs/flux-schnell"
 
-# --- UPDATE THESE VALUES ---
-# Make sure to include "https://" and the "/v1/" path
-NEBIUS_BASE_URL = "https://api.tokenfactory.nebius.com/v1/" # Use the URL from your Nebius dashboard
-TEXT_MODEL = "meta-llama/Llama-3.3-70B-Instruct"           # Use a text model from your dashboard
-IMAGE_MODEL = "black-forest-labs/flux-dev"             # Use an image model from your dashboard
-
-# Check if the API key is set
 if not NEBIUS_API_KEY:
-    # This will show a nice error in the Gradio app if the secret isn't set
-    raise gr.Error("NEBIUS_API_KEY is not set. Please add it to your Hugging Face Space secrets.")
+    # Check for the key but don't crash immediately to allow UI to load with error message
+    print("WARNING: API Key is not set.")
 
-# Set up the single Nebius client (it's OpenAI-compatible)
+# Initialize the client
 nebius_client = OpenAI(
     api_key=NEBIUS_API_KEY,
     base_url=NEBIUS_BASE_URL,
 )
 
-# --- 2. DEFINE YOUR AGENT'S TOOLS ---
+# --- 2. DEFINE LOCAL TOOLS ---
 
-def call_text_gen_tool(prompt: str, system_message: str) -> str:
+def tool_generate_text(prompt: str, system_message: str = "You are a helpful assistant.") -> str:
     """
-    Tool 1: Calls the Nebius LLM for text generation.
-    Takes a customizable system message.
+    Tool: Generates text using the LLM.
     """
-    print(f"Calling Text Tool: {prompt[:30]}...")
+    print(f"📝 Text Tool calling API: {prompt[:40]}...")
     try:
-        chat_completion = nebius_client.chat.completions.create(
+        response = nebius_client.chat.completions.create(
             messages=[
                 {"role": "system", "content": system_message},
                 {"role": "user", "content": prompt},
             ],
             model=TEXT_MODEL,
-            max_tokens=256,
+            max_tokens=512,
+            temperature=0.7,
         )
-        return chat_completion.choices[0].message.content
+        return response.choices[0].message.content
     except Exception as e:
-        print(f"Error calling Nebius LLM: {e}")
-        raise gr.Error(f"Nebius Text API Error: {e}")
+        print(f"❌ Text Tool Error: {e}")
+        raise gr.Error(f"Text Generation Failed: {e}")
 
-def call_image_gen_tool(prompt: str) -> str:
-    """Tool 2: Calls the Nebius Image API for image generation."""
-    print(f"Calling Image Tool: {prompt[:30]}...")
+def tool_generate_image(prompt: str) -> str:
+    """
+    Tool: Generates an image using the Image API.
+    Returns a URL to the image.
+    """
+    print(f"🎨 Image Tool calling API: {prompt[:40]}...")
     try:
         response = nebius_client.images.generate(
             model=IMAGE_MODEL,
             prompt=prompt,
             n=1,
-            size="1024x1024",
+            size="1024x1024", # Standard size for Flux
         )
-        # The API returns a URL to the generated image
+        # Return the URL
         return response.data[0].url
     except Exception as e:
-        print(f"Error calling Nebius Image API: {e}")
-        raise gr.Error(f"Nebius Image API Error: {e}")
+        print(f"❌ Image Tool Error: {e}")
+        raise gr.Error(f"Image Generation Failed: {e}")
 
+def tool_web_search(query: str) -> str:
+    """
+    Tool: Searches the web using DuckDuckGo.
+    Returns a summarized string of results.
+    """
+    print(f"🔍 Search Tool searching for: {query}...")
+    try:
+        results = DDGS().text(query, max_results=3)
+        if not results:
+            return "No relevant search results found."
+        
+        formatted_results = ""
+        for result in results:
+            formatted_results += f"- {result['title']}: {result['body']}\n"
+            
+        return formatted_results
+    except Exception as e:
+        print(f"❌ Search Tool Error: {e}")
+        return "Error performing web search."
 
-# --- 3. DEFINE THE AUTONOMOUS AGENT LOGIC ---
+# --- 3. AGENT LOGIC (THE BRAIN) ---
 
-def run_brand_agent_plan(company_name, business_type, style, progress=gr.Progress()):
+def run_autonomous_agent(company_name, business_type, style, progress=gr.Progress()):
+    """
+    This is the 'Agent' that plans and executes the workflow.
+    """
     
-    # 1. PLANNING
-    progress(0, desc="Agent is creating a plan...")
+    # --- PHASE 1: PLANNING & COLOR ---
+    progress(0.1, desc="Step 1: Planning & Color Palette...")
     
-    plan = [
-        {"step": 1, "action": "generate_color_text", "prompt": f"Generate a 5-color palette for a {business_type} with a {style} vibe. List the colors as a simple, comma-separated string, e.g., 'Deep espresso brown, warm cream, rustic bronze, slate grey, soft olive green'."},
-        {"step": 2, "action": "generate_logo_from_colors"},
-        {"step": 3, "action": "generate_palette_from_colors"},
-        {"step": 4, "action": "generate_social_copy", "prompt": f"Write 3 short, catchy social media posts (for X/Twitter) announcing a new {business_type} called '{company_name}'. The vibe is {style}."}
-    ]
+    # Prompt to get a structured color palette from the text model
+    color_prompt = f"Create a distinct, 5-color palette for a {business_type} named '{company_name}' with a {style} style. Return ONLY a comma-separated list of color names/hex codes. No intro/outro."
+    colors_text = tool_generate_text(color_prompt, "You are a design expert. Output only the requested list.")
     
-    log_message = "Agent Plan Created:\n"
-    log_message += f"- Step 1: Generate color palette (as text)\n"
-    log_message += f"- Step 2: Generate logo from colors\n"
-    log_message += f"- Step 3: Generate palette image from colors\n"
-    log_message += f"- Step 4: Generate social media copy\n"
+    # Clean up the response
+    colors_text = colors_text.strip().strip('.')
     
-    results = {}
+    log = f"✅ Plan Started for {company_name}\n"
+    log += f"🎨 Palette Decided: {colors_text}\n"
     
-    # 2. EXECUTION
+    # --- PHASE 2: LOGO GENERATION ---
+    progress(0.3, desc="Step 2: Generating Logo...")
     
-    # Step 1: Generate Color Palette (as Text)
-    progress(0.2, desc="Executing Step 1: Calling Nebius Text Tool (for Colors)...")
-    system_prompt_colors = "You are a concise branding expert. Only output a single, comma-separated list of 5 colors. Do not add any other text, labels, or explanations."
-    color_palette_text = call_text_gen_tool(plan[0]["prompt"], system_prompt_colors)
-    results["color_text"] = color_palette_text.strip().strip('.') # Clean up the output
-    log_message += f"Step 1 Complete. Colors decided: {results['color_text']}\n"
+    logo_prompt = f"A {style} logo for a {business_type} named '{company_name}'. Use these colors: {colors_text}. Minimalist, vector style, white background. NO TEXT."
+    logo_url = tool_generate_image(logo_prompt)
     
-    # Step 2: Generate Logo (using colors from Step 1)
-    progress(0.4, desc="Executing Step 2: Calling Nebius Image Tool (Logo)...")
-    logo_prompt = f"A {style} logo for a {business_type} named '{company_name}'. Use *only* the following colors: {results['color_text']}. The logo should be an abstract icon or geometric mark. NO TEXT."
-    results["logo_url"] = call_image_gen_tool(logo_prompt)
-    log_message += f"Step 2 Complete. Logo URL acquired.\n"
+    log += f"🖼️ Logo Generated.\n"
 
-    # Step 3: Generate Color Palette Image (using colors from Step 1)
-    progress(0.7, desc="Executing Step 3: Calling Nebius Image Tool (Color Palette)...")
-    palette_prompt = f"A refined color palette presentation. Show 5 distinct, harmonious colors in a modern, clean layout. The colors are: {results['color_text']}. No text."
-    results["palette_url"] = call_image_gen_tool(palette_prompt)
-    log_message += f"Step 3 Complete. Color Palette URL acquired.\n"
+    # --- PHASE 3: BRAND ASSET GENERATION ---
+    progress(0.5, desc="Step 3: Generating Brand Asset...")
+    
+    asset_prompt = f"A sophisticated color palette display for a brand. Show swatches of these colors: {colors_text}. Clean, modern layout, high quality, {style} aesthetic. No text."
+    asset_url = tool_generate_image(asset_prompt)
+    
+    log += f"🖼️ Brand Asset Generated.\n"
+    
+    # --- PHASE 4: MARKET RESEARCH ---
+    progress(0.7, desc="Step 4: Researching Market Trends...")
+    
+    search_query = f"current social media marketing trends for {style} {business_type} 2024 2025"
+    search_data = tool_web_search(search_query)
     
-    # Step 4: Generate Social Copy
-    progress(0.9, desc="Executing Step 4: Calling Nebius Text Tool (Copy)...")
-    system_prompt_kit = "You are a helpful branding assistant. Be creative and concise."
-    results["copy_text"] = call_text_gen_tool(plan[3]["prompt"], system_prompt_kit)
-    log_message += "Step 4 Complete. Copy generated.\n"
+    log += f"🔍 Market Trends Found: {search_data[:100]}...\n"
     
-    progress(1.0, desc="Agent Execution Complete!")
-    log_message += "Plan executed successfully."
+    # --- PHASE 5: COPYWRITING ---
+    progress(0.9, desc="Step 5: Writing Content...")
+    
+    copy_system_msg = "You are a professional social media manager."
+    copy_prompt = f"""
+    Write 3 engaging social media posts (Twitter/X style) for a new {business_type} called '{company_name}'.
+    
+    Brand Vibe: {style}
+    
+    Integrate these real-time market insights:
+    {search_data}
+    """
     
-    # 3. RESULT
-    return results["logo_url"], results["palette_url"], results["copy_text"], log_message
+    social_copy = tool_generate_text(copy_prompt, copy_system_msg)
+    
+    log += f"✍️ Copy Written.\n✅ Agent Execution Complete."
+    
+    return logo_url, asset_url, social_copy, log
+
+# --- 4. CHATBOT LOGIC ---
 
-# --- 4. CHATBOT FUNCTION ---
-def chat_with_brand_agent(message, history):
-    system_prompt_chat = "You are a friendly and creative branding expert. A user is asking for ideas for their new brand. Give them suggestions, ask clarifying questions, and help them brainstorm."
+def chat_response(message, history):
+    """
+    A simple chat interface that remembers context.
+    """
+    system_message = "You are a creative brand consultant. Help the user brainstorm ideas."
     
-    # Build the message list from history
-    messages = [{"role": "system", "content": system_prompt_chat}]
-    for user_msg, assistant_msg in history:
+    messages = [{"role": "system", "content": system_message}]
+    
+    # Add history
+    for user_msg, bot_msg in history:
         messages.append({"role": "user", "content": user_msg})
-        messages.append({"role": "assistant", "content": assistant_msg})
+        messages.append({"role": "assistant", "content": bot_msg})
+    
+    # Add current message
     messages.append({"role": "user", "content": message})
-
+    
     try:
         response = nebius_client.chat.completions.create(
             messages=messages,
             model=TEXT_MODEL,
             max_tokens=300,
         )
-        bot_message = response.choices[0].message.content
-        
-        history.append((message, bot_message))
-        
+        bot_reply = response.choices[0].message.content
+        history.append((message, bot_reply))
         return history, ""
-    
     except Exception as e:
-        print(f"Error in chat: {e}")
-        # Return an error message in the chat
-        history.append((message, f"Sorry, I ran into an error: {e}"))
+        history.append((message, f"Error: {e}"))
         return history, ""
 
+# --- 5. UI CONSTRUCTION ---
 
-# --- 5. CREATE THE GRADIO UI ---
-
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown(
-        """
-        # 🤖 AutoBrand Studio
-        Powered by **Hugging Face** and **Nebius Token Factory**.
-        This agent autonomously plans and calls Nebius APIs for text and image generation to build a brand kit.
-        """
-    )
+with gr.Blocks(theme=gr.themes.Soft(), title="AutoBrand Studio") as demo:
+    gr.Markdown("# 🤖 AutoBrand Studio")
+    gr.Markdown("An Autonomous Agent powered by **AI Models** & **DuckDuckGo**. Generates a full brand kit with real-time market research.")
     
     with gr.Row():
-        company_input = gr.Textbox(label="Company Name", value="Aura")
-        business_type_input = gr.Textbox(label="Business Type / Industry", value="Skincare Brand")
-        style_input = gr.Textbox(label="Style / Vibe", value="modern, minimalist, clean")
-    
-    submit_btn = gr.Button("🚀 Generate Brand Kit")
-    
-    with gr.Tabs():
-        # --- TAB 1: BRAND KIT ---
-        with gr.TabItem("Brand Kit Results"):
-            with gr.Row():
-                logo_output = gr.Image(label="Generated Logo (from Nebius)", height=400)
-                color_palette_output = gr.Image(label="Generated Color Palette (from Nebius)", height=400)
-            copy_output = gr.Textbox(label="Generated Social Media Copy (from Nebius)", lines=8, interactive=True)
+        with gr.Column(scale=1):
+            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, sage green tones")
+            generate_btn = gr.Button("✨ Create Brand Kit", variant="primary")
         
-        # --- TAB 2: AGENT LOG ---
-        with gr.TabItem("Agent Log"):
-            log_output = gr.Textbox(label="Agent Plan & Execution Log", lines=10, interactive=False)
-
-        # --- TAB 3: CHATBOT ---
-        with gr.TabItem("Brand Chatbot"):
-            chatbot = gr.Chatbot(label="Brand Brainstormer")
-            chat_msg = gr.Textbox(label="Ask for ideas...", placeholder="e.g., 'What are some good names for my new skincare brand?'")
-            chat_btn = gr.Button("Send")
-
-    # --- Wire up the UI ---
-    
-    submit_btn.click(
-        fn=run_brand_agent_plan,
-        inputs=[company_input, business_type_input, style_input],
-        outputs=[logo_output, color_palette_output, copy_output, log_output]
+        with gr.Column(scale=2):
+            with gr.Tabs():
+                with gr.Tab("🎨 Brand Assets"):
+                    with gr.Row():
+                        out_logo = gr.Image(label="Logo Concept")
+                        out_asset = gr.Image(label="Color Palette / Mood")
+                    out_copy = gr.Textbox(label="Social Media Strategy (Data-Driven)", lines=10)
+                
+                with gr.Tab("⚙️ Agent Logs"):
+                    out_log = gr.Textbox(label="Execution Log", lines=15)
+                
+                with gr.Tab("💬 Consultant"):
+                    chatbot = gr.Chatbot(label="Brand Consultant", height=400)
+                    chat_input = gr.Textbox(label="Ask about your brand...")
+                    chat_btn = gr.Button("Send")
+
+    # Event Handlers
+    generate_btn.click(
+        fn=run_autonomous_agent,
+        inputs=[company_input, type_input, style_input],
+        outputs=[out_logo, out_asset, out_copy, out_log]
     )
     
-    # *** UPDATED: Simplified the click/submit handlers ***
-    
-    # Click handler for the new chatbot
-    chat_btn.click(
-        fn=chat_with_brand_agent,
-        inputs=[chat_msg, chatbot],
-        outputs=[chatbot, chat_msg] # Clears chat_msg
-    )
-    
-    # Submit handler for the new chatbot (pressing Enter)
-    chat_msg.submit(
-        fn=chat_with_brand_agent,
-        inputs=[chat_msg, chatbot],
-        outputs=[chatbot, chat_msg] # Clears chat_msg
-    )
+    chat_btn.click(fn=chat_response, inputs=[chat_input, chatbot], outputs=[chatbot, chat_input])
+    chat_input.submit(fn=chat_response, inputs=[chat_input, chatbot], outputs=[chatbot, chat_input])
 
 if __name__ == "__main__":
     demo.launch()