import gradio as gr
from smolagents import ToolCallingAgent, tool, LiteLLMModel
import os
from dotenv import load_dotenv
load_dotenv()

CO_API_KEY = os.environ.get("CO_API_KEY")
model = "command-r-plus-08-2024"

preamble = """You are Alfred, a helpful butler assistant. 
When using a tool and getting a result, provide a clear response based on that result.
After calling a tool once, suggest a menu based on the result. 
Include a creative description of ingredients for each item.
Do not add any additional information or context such as recommendations or explanations.
Use gender-neutral language and avoid gendered pronouns."""

model = LiteLLMModel(
    model_id=f"cohere/{model}",
    api_key=CO_API_KEY,
    temperature=0.2,
    max_tokens=1000
)

@tool
def suggest_menu(occasion: str) -> str:
    """
    Suggests a menu based on the occasion.
    Args:
        occasion (str): The type of occasion for the party. Allowed values are:
                        - "casual": Menu for casual party.
                        - "formal": Menu for formal party.
                        - "superhero": Menu for superhero party.
                        - "custom": Custom menu.
    """
    if occasion == "casual":
        return "Pizza, and some snacks and drinks." 
    elif occasion == "formal":
        return "three-course dinner with wine and dessert."
    elif occasion == "superhero":
        return "Buffet with high-energy and healthy food."
    else:
        return "Custom menu for the butler."

agent = ToolCallingAgent(tools=[suggest_menu], model=model)

def plan_menu(user_request):
    try:
        result = agent.run(preamble + f"\n{user_request}", max_steps=1)
        return result
    except Exception as e:
        return f"Error: {str(e)}"

# Create Gradio interface
with gr.Blocks(title="Alfred's Menu Planner") as demo:
    gr.Markdown("# 🍽️ Alfred's Menu Planner")
    gr.Markdown("Ask Alfred to help plan a menu for your party!")
    
    with gr.Row():
        with gr.Column(scale=1):
            user_input = gr.Textbox(
                label="What kind of menu do you need?",
                placeholder="Enter a query, select an option below, or click the button for a surprise menu!",
                lines=3
            )

            submit_btn = gr.Button("Ask Alfred", variant="primary")

            gr.Examples(
            examples=[
                ["Plan a casual menu for my birthday party"],
                ["I need a formal dinner menu for 8 people"],
                ["What would be good for a superhero themed party in my backyard?"],
                ["Suggest a menu for a garden party. Aim to impress!"]
            ],
            inputs=user_input
    )
        
        with gr.Column(scale=2):
            output = gr.Markdown(
                label="Alfred's Response",
                value="Ask Alfred for a menu recommendation!"
            )
    
    submit_btn.click(
        fn=lambda: "Ask Alfred for a menu recommendation!",
        outputs=output,
        show_progress=True
    ).then(
        fn=plan_menu,
        inputs=user_input,
        outputs=output,
        show_progress=True
    )
    
    

if __name__ == "__main__":
    demo.launch()