import gradio as gr
from swarm import Swarm, Agent

# Initialize Swarm client
client = Swarm()

# Define the agent functions
def transfer_to_agent_b():
    return agent_b

def transfer_to_agent_c():
    return agent_c

def transfer_to_spanish_agent():
    """Transfer Spanish speaking users immediately."""
    return spanish_agent

# Define the agents
agent_a = Agent(
    name="Agent A",
    instructions="You are a helpful agent.",
    functions=[transfer_to_agent_b],
)

agent_b = Agent(
    name="Agent B",
    instructions="Only speak in Haikus.",
    functions=[transfer_to_spanish_agent],
)

agent_c = Agent(
    name="Agent C",
    instructions="Welcome to the matrix.",
)

spanish_agent = Agent(
    name="Spanish Agent",
    instructions="You only speak Spanish.",
)

english_agent = Agent(
    name="English Agent",
    instructions="You only speak English.",
)

# Append function to transfer to Spanish agent for English agent
english_agent.functions.append(transfer_to_spanish_agent)

# Create the Gradio interface function
def chat_with_agent(user_input):
    messages = [{"role": "user", "content": user_input}]
    response = client.run(agent=english_agent, messages=messages)
    return response.messages[-1]["content"]

# Create Gradio interface
iface = gr.Interface(fn=chat_with_agent, inputs="text", outputs="text", live=True)

# Launch the interface
iface.launch()