simple_ui.py

"""
Gradio-based user interface for AIris - Your AI Fashion Advisor.
"""

import os
from typing import List, Tuple
from pydantic import BaseModel, ConfigDict
import gradio as gr
from dotenv import load_dotenv
from src.agents.fashion_advisor import AIris
from src.data.fashion_knowledge import STYLE_AESTHETICS, OCCASIONS
from PIL import Image

# Load environment variables
load_dotenv()

# Pydantic configuration to allow arbitrary types
class Config:
    model_config = ConfigDict(arbitrary_types_allowed=True)

def initialize_agent(offline_mode: bool) -> AIris:
    """Initialize the AIris agent."""
    return AIris(offline_mode=offline_mode)

def process_request(
    message: str,
    image: Image.Image,
    weather: str,
    occasion: str,
    style: str,
    offline_mode: bool,
    history: List[Tuple[str, str]],
    agent: AIris
) -> tuple[List[Tuple[str, str]], List[Tuple[str, str]]]:
    """Process a fashion advice request."""
    if not message.strip():
        return history, history
        
    try:
        response = None
        
        # If asking for inspiration
        if any(word in message.lower() for word in ["inspire", "inspiration", "ideas", "trending"]):
            response = agent.get_inspiration(
                style=style if style != "Select style" else None,
                occasion=occasion if occasion != "Select occasion" else None,
                season=weather if weather != "Select weather" else None
            )
        
        # If an image is provided and the message asks about styling
        elif image is not None and any(word in message.lower() for word in ["style", "wear", "match", "pair", "combine"]):
            response = str(agent.analyze_image(image))
        
        # If it's a color combination request
        elif "color" in message.lower() and "combination" in message.lower():
            # Extract color from message
            colors = ["black", "white", "blue", "red", "green", "yellow", "purple", "pink", "brown", "gray"]
            for color in colors:
                if color in message.lower():
                    response = agent.suggest_color_combinations(color)
                    break
        
        # Default to chat
        if response is None:
            response = agent.chat(message)
        
        # Update history
        history = history or []
        history.append((message, str(response)))
        
        return history, history
        
    except Exception as e:
        error_msg = f"⚠️ Error: {str(e)}"
        history = history or []
        history.append((message, error_msg))
        return history, history

def create_ui() -> gr.Blocks:
    """Create and configure the Gradio interface."""
    # Initialize agent with offline mode
    agent = initialize_agent(offline_mode=True)  # Start in offline mode by default
    
    # Create the interface
    with gr.Blocks(title="AIris - Your AI Fashion Advisor") as interface:
        gr.Markdown(
            """
            # 👗 AIris - Your AI Fashion Advisor
            Your intelligent fashion companion that helps you create perfect outfits and analyze your style!
            
            ## What can I do?
            - Suggest outfits based on weather, occasion, and style
            - Analyze your fashion items and suggest styling options
            - Get inspiration for different styles and occasions
            - Recommend color combinations
            - Provide styling tips and tricks
            """
        )
        
        with gr.Row():
            with gr.Column(scale=1):
                image_input = gr.Image(
                    label="Upload a fashion item to get styling suggestions",
                    type="pil"
                )
                
                weather_dropdown = gr.Dropdown(
                    choices=["Select weather", "hot", "cold", "rainy", "windy"],
                    value="Select weather",
                    label="Weather"
                )
                occasion_dropdown = gr.Dropdown(
                    choices=["Select occasion"] + list(OCCASIONS.keys()),
                    value="Select occasion",
                    label="Occasion"
                )
            
            with gr.Column(scale=1):
                style_dropdown = gr.Dropdown(
                    choices=["Select style"] + list(STYLE_AESTHETICS.keys()),
                    value="Select style",
                    label="Style Preference"
                )
                
                offline_checkbox = gr.Checkbox(
                    label="Offline Mode",
                    value=True,  # Start in offline mode
                    info="When unchecked, AIris will use the Groq API for better responses"
                )
        
        # Create a state for the agent and chat history
        agent_state = gr.State(agent)
        chat_history = gr.State([])
        
        # Handle offline mode toggle
        def toggle_offline_mode(offline: bool, agent_state: AIris) -> AIris:
            """Toggle offline mode and reinitialize agent."""
            return initialize_agent(offline_mode=offline)
            
        offline_checkbox.change(
            fn=toggle_offline_mode,
            inputs=[offline_checkbox, agent_state],
            outputs=[agent_state]
        )
        
        # Chat interface
        with gr.Group():
            chatbot = gr.Chatbot(
                label="Chat with AIris",
                height=400,
                show_label=True,
                render_markdown=True,
                avatar_images=["👤", "👗"]
            )
            msg = gr.Textbox(
                label="Type your message here...",
                placeholder="Ask me for outfit suggestions, upload images for styling advice, or get fashion inspiration!",
                container=False
            )
            clear = gr.ClearButton([msg, chatbot])
            
            msg.submit(
                fn=process_request,
                inputs=[
                    msg,
                    image_input,
                    weather_dropdown,
                    occasion_dropdown,
                    style_dropdown,
                    offline_checkbox,
                    chat_history,
                    agent_state
                ],
                outputs=[chatbot, chat_history]
            )
            
        # Example prompts
        gr.Examples(
            examples=[
                "How do I style this uploaded item?",
                "Show me inspiration for casual summer outfits",
                "Suggest an outfit for a rainy day",
                "What colors go well with blue?",
                "How can I style a white button-up shirt?",
                "What's a good outfit for a business meeting?"
            ],
            inputs=msg
        )
    
    return interface

if __name__ == "__main__":
    # Create and launch the interface
    interface = create_ui()
    interface.launch(
        share=False,
        server_name="0.0.0.0",
        server_port=7860,
        show_api=False
    )