Update app.py

app.py

@@ -40,65 +40,57 @@ t2i_server_parameters = StdioServerParameters(
 
 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.")
+# --- 3. Main Application Logic ---
+# We define the main function to encapsulate the agent setup and workflow.
+def run_storycrafter():
+    # 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...")
+        # Get all available tools from all connected MCP servers.
+        all_tools = mcp.get_tools()
+        print(f"Found {len(all_tools)} tools.")
+        if not all_tools:
+            print("Warning: No tools were loaded from the MCP servers. Agents will have limited capabilities.")
+            
+        # --- 4. Agent Definitions ---
+        # We define our agent team inside the 'with' block to give them access to the loaded tools.
+
+        # The Writer Agent is given all available tools. 
+        # It will intelligently select the correct tool (e.g., knowledge_graph) when needed.
+        writer_agent = ToolCallingAgent(
+            tools=all_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 is also given all available tools.
+        # It will select the 'text_to_image' tool when its services are required.
+        illustrator_agent = ToolCallingAgent(
+            tools=all_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 orchestrates the other two agents.
+        director_agent = CodeAgent(
+            tools=[], # The director itself doesn't call external tools, it directs other agents.
+            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.
+            """
+        )
         
-    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__":
+        # --- 5. The Creative Workflow ---
         user_prompt = "a story about a wise old owl living in a library of forgotten books"
 
         print(f"\n--- Director's Task ---")
@@ -133,3 +125,7 @@ with MCPClient(server_parameters) as mcp:
             print("\n--- ILLUSTRATION ---")
             print("No illustration was generated.")
 
+
+# --- 6. Execution Start ---
+if __name__ == "__main__":
+    run_storycrafter()