Update app.py

app.py

@@ -7,11 +7,12 @@ from google.genai import types
 from mcp import ClientSession
 from mcp.client.sse import sse_client
 
+# --- CONFIGURATION ---
 MCP_SERVER_URL = "https://mgokg-db-api-mcp.hf.space/sse"
 GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")
 
 async def call_mcp_tool(start_loc, dest_loc):
-    """Executes the train search via your MCP server."""
+    """Connects to the MCP server to fetch train data."""
     async with sse_client(MCP_SERVER_URL) as (read_stream, write_stream):
         async with ClientSession(read_stream, write_stream) as session:
             await session.initialize()
@@ -27,16 +28,16 @@ def generate(input_text):
 
     client = genai.Client(api_key=GEMINI_API_KEY)
     
-    # Use 2.0 Flash (Tools work here). 
-    # Do NOT use 'gemini-2.0-flash-thinking-exp' as it doesn't support tools yet.
+    # CRITICAL: Use gemini-2.0-flash-exp. 
+    # Do NOT use gemini-2.0-flash-thinking-exp as it doesn't support tools.
     model_id = "gemini-2.0-flash-exp"
 
-    # Define the tool explicitly
+    # Define the train tool schema
     train_tool = types.Tool(
         function_declarations=[
             types.FunctionDeclaration(
                 name="get_train_connections",
-                description="Finds live train connections between two German stations/cities.",
+                description="Finds live train connections between two German stations.",
                 parameters={
                     "type": "OBJECT",
                     "properties": {
@@ -50,45 +51,43 @@ def generate(input_text):
     )
 
     try:
-        # 1. Create the config WITHOUT thinking_config
+        # 1. Setup config WITHOUT thinking_config
         config = types.GenerateContentConfig(
             tools=[train_tool, types.Tool(google_search=types.GoogleSearch())],
-            temperature=0.3,
-            # thinking_config REMOVED - this is the key fix
+            temperature=0.3
         )
 
-        # 2. Start the chat session
         chat = client.chats.create(model=model_id, config=config)
         response = chat.send_message(input_text)
 
-        # 3. Handle the Tool-Use Loop manually
-        # This handles both your MCP tool and Google Search
-        max_iterations = 5
-        for _ in range(max_iterations):
-            # Check if there is a tool call in the first part of the message
+        # 2. Manual Tool Loop
+        # We must manually process tool calls because we are calling an external MCP server
+        max_turns = 5
+        for _ in range(max_turns):
+            # Exit if there are no parts or no tool calls
             if not response.candidates[0].content.parts:
                 break
-                
-            tool_calls = [p.tool_call for p in response.candidates[0].content.parts if p.tool_call]
             
+            # Find all tool calls in the message parts
+            tool_calls = [p.tool_call for p in response.candidates[0].content.parts if p.tool_call]
             if not tool_calls:
                 break
 
             tool_responses = []
             for call in tool_calls:
                 if call.name == "get_train_connections":
-                    # Run the MCP call
-                    result_data = asyncio.run(call_mcp_tool(call.args["start_loc"], call.args["dest_loc"]))
+                    # Execute the MCP call
+                    # We use asyncio.run because Gradio's click handler is synchronous
+                    train_data = asyncio.run(call_mcp_tool(call.args["start_loc"], call.args["dest_loc"]))
+                    
                     tool_responses.append(
                         types.Part.from_function_response(
                             name=call.name,
-                            response={"result": result_data}
+                            response={"result": train_data}
                         )
                     )
-                # Google Search is handled automatically by the model if configured, 
-                # but if it returns a call, we let the loop handle the logic.
-
-            # Send all tool results back to the model
+            
+            # If we have results, send them back to the model
             if tool_responses:
                 response = chat.send_message(tool_responses)
             else:
@@ -97,6 +96,6 @@ def generate(input_text):
         return response.text, ""
 
     except Exception as e:
-        return f"### Error encountered\n{str(e)}", ""
+        return f"### Logic Error\n{str(e)}", ""
 
-# (Keep your existing Gradio Blocks code below)
+# (The rest of your Gradio code remains the same)