app.py

import os
from dotenv import load_dotenv
from smolagents import CodeAgent, ToolCallingAgent, LiteLLMModel, MCPClient
from mcp import StdioServerParameters
import base64
from PIL import Image
import io

# --- 1. Environment and Model Setup ---
# Load environment variables from a .env file (for API keys)
load_dotenv()

# Initialize the language model that our agents will use.
# Ensure your GEMINI_API_KEY is set in your .env file.
model = LiteLLMModel(
    model_id="gemini/gemini-2.0-flash-exp",
    api_key=os.getenv("GEMINI_API_KEY")
)

# --- 2. MCP Server Configuration ---
# Define the connection parameters for your MCP servers.
# These commands will run in the background to connect to your deployed tools.
kgb_server_parameters = StdioServerParameters(
    command="npx",
    args=[        
        "mcp-remote",
        "https://agents-mcp-hackathon-kgb-mcp.hf.space/gradio_api/mcp/sse",
        "--transport",
        "sse-only"],
)

t2i_server_parameters = StdioServerParameters(
    command="npx",
    args=[        
        "mcp-remote",
        "https://agents-mcp-hackathon-t2i.hf.space/gradio_api/mcp/sse",
        "--transport",
        "sse-only"],
)

server_parameters = [kgb_server_parameters, t2i_server_parameters]

# --- 3. Main Execution Block ---
# We use the MCPClient as a context manager to handle the lifecycle of the servers.
with MCPClient(server_parameters) as mcp:
    print("Connecting to MCP servers and fetching tools...")
    all_tools = mcp.get_tools()
    print(f"Found {len(all_tools)} tools.")

    # --- 4. Tool Integration ---
    # Find our specific tools from the list provided by the MCP servers.
    # We will look for them by name.
    knowledge_tool = next((tool for tool in all_tools if "knowledge_graph" in tool.name.lower()), None)
    image_tool = next((tool for tool in all_tools if "text_to_image" in tool.name.lower()), None)

    if not knowledge_tool:
        print("Warning: Knowledge graph tool not found.")
    if not image_tool:
        print("Warning: Text-to-image tool not found.")
        
    writer_tools = [knowledge_tool] if knowledge_tool else []
    illustrator_tools = [image_tool] if image_tool else []

    # --- 5. Agent Definitions ---
    # We define our agent team, now equipped with the tools from your MCP servers.

    # The Writer Agent
    writer_agent = ToolCallingAgent(
        tools=writer_tools,
        model=model,
        name="writer",
        description="A creative agent that writes short stories. It can use a knowledge graph tool to research topics for inspiration."
    )

    # The Illustrator Agent
    illustrator_agent = ToolCallingAgent(
        tools=illustrator_tools,
        model=model,
        name="illustrator",
        description="An artist agent that creates illustrations based on a descriptive prompt using a text-to-image tool."
    )

    # The Director Agent
    director_agent = CodeAgent(
        tools=[],
        model=model,
        managed_agents=[writer_agent, illustrator_agent],
        system_prompt="""
        You are the Director of Agentic Storycrafter, a creative team. Your job is to manage the writer and illustrator agents to create a story with an illustration.

        Here is your workflow:
        1.  Receive a user's prompt for a story.
        2.  Call the `writer` agent to write a story based on the user's prompt.
        3.  After the story is written, create a short, descriptive prompt for an illustration that captures the essence of the story.
        4.  Call the `illustrator` agent with this new prompt to generate an image. The result will be a dictionary containing image data.
        5.  Return a dictionary containing both the final 'story' and the 'image_data' from the illustrator.
        """
    )
    
    # --- 6. The Creative Workflow ---
    if __name__ == "__main__":
        user_prompt = "a story about a wise old owl living in a library of forgotten books"

        print(f"\n--- Director's Task ---")
        print(f"Prompt: {user_prompt}\n")

        # The director now runs the full workflow.
        final_output = director_agent.run(f"Create a story and illustration for the following prompt: {user_prompt}")
        
        print("\n--- Agentic Storycrafter Result ---")
        
        # The output from the director is code that needs to be executed to get the result
        result_dict = eval(final_output)

        story = result_dict.get("story")
        image_data = result_dict.get("image_data")

        print("\n--- STORY ---")
        print(story)

        if image_data and 'b64_json' in image_data:
            print("\n--- ILLUSTRATION ---")
            print("Illustration created. Saving to 'story_illustration.png'")
            # Decode the base64 string and save it as an image file
            try:
                img_bytes = base64.b64decode(image_data['b64_json'])
                img = Image.open(io.BytesIO(img_bytes))
                img.save("story_illustration.png")
                print("Image saved successfully.")
            except Exception as e:
                print(f"Error saving image: {e}")
        else:
            print("\n--- ILLUSTRATION ---")
            print("No illustration was generated.")