Update app.py

app.py

@@ -19,7 +19,6 @@ model = LiteLLMModel(
 
 # --- 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=[        
@@ -41,22 +40,21 @@ t2i_server_parameters = StdioServerParameters(
 server_parameters = [kgb_server_parameters, t2i_server_parameters]
 
 # --- 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:
+    # Instantiate the MCPClient *before* the 'with' block, as per the working example.
+    mcp = MCPClient(server_parameters)
+    
+    # Use the created MCPClient instance as a context manager.
+    with 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,
@@ -65,7 +63,6 @@ def run_storycrafter():
         )
 
         # 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,
@@ -75,7 +72,7 @@ def run_storycrafter():
 
         # The Director Agent orchestrates the other two agents.
         director_agent = CodeAgent(
-            tools=[], # The director itself doesn't call external tools, it directs other agents.
+            tools=[],
             model=model,
             managed_agents=[writer_agent, illustrator_agent],
             system_prompt="""
@@ -96,12 +93,10 @@ def run_storycrafter():
         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")
@@ -113,7 +108,6 @@ def run_storycrafter():
         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))
@@ -125,7 +119,7 @@ def run_storycrafter():
             print("\n--- ILLUSTRATION ---")
             print("No illustration was generated.")
 
-
 # --- 6. Execution Start ---
 if __name__ == "__main__":
-    run_storycrafter()
+    run_storycrafter()
+