app.py

import gradio as gr
from openai import OpenAI
from ddgs import DDGS
import os
import base64
import io
from PIL import Image

# --- 1. CONFIGURATION ---
NEBIUS_API_KEY = os.environ.get("NEBIUS_API_KEY")
NEBIUS_BASE_URL = "https://api.studio.nebius.ai/v1/"
TEXT_MODEL = "meta-llama/Llama-3.3-70B-Instruct-fast"
IMAGE_MODEL = "black-forest-labs/flux-dev"

if not NEBIUS_API_KEY:
    print("WARNING: NEBIUS_API_KEY environment variable is not set.")

nebius_client = OpenAI(
    api_key=NEBIUS_API_KEY,
    base_url=NEBIUS_BASE_URL,
)

# Helper to convert bytes to PIL
def bytes_to_pil(image_bytes: bytes):
    return Image.open(io.BytesIO(image_bytes))

# --- 2. MCP TOOL FUNCTIONS ---

def tool_web_search(query: str):
    """Performs a market research search using DuckDuckGo."""
    print(f"🔍 Search Tool: {query}...")
    try:
        results = DDGS().text(query, max_results=4)
        if not results:
            return "No specific market data found."
        formatted_results = "\n".join([f"- {r.get('title')}: {r.get('body')}" for r in results])
        return formatted_results
    except Exception as e:
        return f"Search Error: {str(e)}"

def tool_generate_image(prompt: str):
    """Generates an image and returns the local file path."""
    print(f"🎨 Image Tool: {prompt[:40]}...")
    try:
        response = nebius_client.images.generate(
            model=IMAGE_MODEL,
            prompt=prompt,
            n=1,
            size="1024x1024",
            response_format="b64_json",
            style="flat",
        )
        b64_data = response.data[0].b64_json
        image_bytes = base64.b64decode(b64_data)
        img = Image.open(io.BytesIO(image_bytes))

        filename = f"gen_{abs(hash(prompt))}.png"
        img.save(filename)
        return filename
    except Exception as e:
        print(f"❌ Image Tool Error: {e}")
        return None

# --- Text Generation Helper ---
def tool_generate_text(prompt: str, system_message: str = "You are a helpful assistant.", max_tokens: int = 1024) -> str:
    try:
        response = nebius_client.chat.completions.create(
            messages=[
                {"role": "system", "content": system_message},
                {"role": "user", "content": prompt},
            ],
            model=TEXT_MODEL,
            max_tokens=max_tokens,
            temperature=0.7,
        )
        return response.choices[0].message.content
    except Exception as e:
        return f"Error: {e}"

# --- Market Insights Summarization Tool ---
def tool_summarize_market_insights(text: str):
    """
    Summarizes long-form market research into insights, trends, and opportunities.
    Exposed to MCP for agent workflows.
    """
    print("🧠 Summarization Tool Running...")
    try:
        prompt = f"""
        Summarize the following market research into:
        • Key trends
        • Industry opportunities
        • Consumer behavior signals
        • Competitive insights
        Make the summary short, actionable, and business-oriented.
        TEXT:
        {text}
        """
        # Increased max_tokens for long summaries
        summary = tool_generate_text(
            prompt,
            system_message="You are an expert market analyst. Provide concise insights only.",
            max_tokens=2000
        )
        return summary
    except Exception as e:
        return f"Summarization Error: {str(e)}"

# --- 3. AUTONOMOUS BRAND AGENT ---
def run_autonomous_agent(company_name, business_type, style, progress=gr.Progress()):
    """Autonomous branding agent logic with market insights summarization."""
    # PHASE 1: COLOR PALETTE
    progress(0.1, desc="Step 1: Planning & 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"✅ Brand Plan Started for {company_name}\n"
    log += f"🎨 Palette: {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."
    )
    logo_path = tool_generate_image(logo_prompt)
    log += "🖼️ Logo Generated.\n"

    # PHASE 3: COLOR PALETTE ASSET
    progress(0.5, desc="Step 3: Color Palette Asset...")
    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 += "🖼️ Palette Asset Generated.\n"

    # PHASE 4: RAW MARKET RESEARCH
    progress(0.65, desc="Step 4: Raw Market Research...")
    search_query = f"marketing trends for {business_type} 2025"
    raw_research = tool_web_search(search_query)
    log += f"🔍 Market Data Retrieved: {search_query}\n"

    # PHASE 4.5: SUMMARIZE MARKET INSIGHTS
    progress(0.8, desc="Step 4.5: Summarizing Insights...")
    summarized = tool_summarize_market_insights(raw_research)
    log += f"🧠 Market Insights Summary: {summarized[:150]}...\n"

    # PHASE 5: SOCIAL MEDIA COPY
    progress(0.95, desc="Step 5: Writing Social Content...")
    copy_prompt = f"""
    Write 3 engaging Twitter/X posts for a new {business_type} called '{company_name}'.
    Tone: {style}.
    Use these insights:
    {summarized}
    """
    social_copy = tool_generate_text(copy_prompt, "You are a professional social media strategist.", max_tokens=1500)
    log += "✍️ Social Copy Generated.\n"
    log += "✅ Autonomy Cycle Complete."

    return logo_path, asset_path, social_copy, log, summarized

# --- 4. BRAND CONSULTANT CHATBOT ---
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:
        response = nebius_client.chat.completions.create(
            messages=messages,
            model=TEXT_MODEL,
            max_tokens=300,
        )
        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:
        history.append({"role": "assistant", "content": f"Error: {e}"})
        return history, ""

# --- 5. GRADIO UI ---
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)",
)

custom_css = """#main-header {text-align: center; margin-bottom: 2rem;}"""

with gr.Blocks(theme=theme, css=custom_css, title="AutoBrand MCP Studio") as demo:
    # HEADER
    with gr.Column(elem_id="main-header"):
        gr.Markdown("# 🖼️ AutoBrand Studio")
        gr.Markdown(
            "Autonomous Agent powered by AI Models, MCP Tools & Web Search."
        )

    with gr.Row():
        # LEFT PANEL
        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 PANEL
        with gr.Column(scale=2):
            with gr.Tabs():
                # BRAND IDENTITY TAB
                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)

                # BRAND CONSULTANT CHATBOT TAB
                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")

                # MARKET INSIGHT SUMMARY TAB
                with gr.TabItem("📊 Market Insight Summary"):
                    summary_box = gr.TextArea(label="Summarized Market Insights", lines=20, interactive=False, show_copy_button=True)

    # HANDLERS
    def run_agent_and_extract_summary(company, btype, style, progress=gr.Progress()):
        logo, asset, copy, log, summarized = run_autonomous_agent(company, btype, style, progress)
        return logo, asset, copy, log, summarized

    generate_btn.click(
        run_agent_and_extract_summary,
        [company_input, type_input, style_input],
        [out_logo, out_asset, out_copy, out_log, summary_box],
    )

    chat_btn.click(chat_response, [chat_input, chatbot], [chatbot, chat_input])
    chat_input.submit(chat_response, [chat_input, chatbot], [chatbot, chat_input])

    # --- MCP TOOLS (Hidden Row) ---
    with gr.Row(visible=False):
        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")

        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")

        summarize_in = gr.Textbox()
        summarize_out = gr.Textbox()
        btn_summarize = gr.Button("Summarize Tool")
        btn_summarize.click(tool_summarize_market_insights, inputs=summarize_in, outputs=summarize_out, api_name="summarize_market_insights")

if __name__ == "__main__":
    demo.launch(ssr_mode=False, mcp_server=True)