Update app.py

app.py

@@ -7,106 +7,96 @@ 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")
 
-# 1. Function to connect to your MCP Server and run the tool
 async def call_mcp_tool(start_loc, dest_loc):
+    """Executes the train search via your MCP server."""
     async with sse_client(MCP_SERVER_URL) as (read_stream, write_stream):
         async with ClientSession(read_stream, write_stream) as session:
             await session.initialize()
-            # Calling the tool defined in your MCP server
             result = await session.call_tool("get_train_connections", arguments={
                 "start_loc": start_loc,
                 "dest_loc": dest_loc
             })
-            # result.content[0].text contains the JSON train data
             return result.content[0].text
 
-# 2. Define the tool for Gemini
-# This describes the tool so Gemini knows when to use it
-train_tool = types.Tool(
-    function_declarations=[
-        types.FunctionDeclaration(
-            name="get_train_connections",
-            description="Finds live train connections between two German stations/cities.",
-            parameters=types.Schema(
-                type="OBJECT",
-                properties={
-                    "start_loc": types.Schema(type="STRING", description="Departure city/station"),
-                    "dest_loc": types.Schema(type="STRING", description="Destination city/station")
-                },
-                required=["start_loc", "dest_loc"]
-            )
-        )
-    ]
-)
-
 def generate(input_text):
     if not GEMINI_API_KEY:
         return "Error: GEMINI_API_KEY is not set.", ""
 
     client = genai.Client(api_key=GEMINI_API_KEY)
-    model_id = "gemini-flash-latest" # Recommended for tool use
+    
+    # Use 2.0 Flash (Tools work here). 
+    # Do NOT use 'gemini-2.0-flash-thinking-exp' as it doesn't support tools yet.
+    model_id = "gemini-2.0-flash-exp"
 
-    # Create a chat session to handle the tool-use loop
-    chat = client.chats.create(
-        model=model_id,
-        config=types.GenerateContentConfig(
-            tools=[train_tool, types.Tool(google_search=types.GoogleSearch())],
-            temperature=0.3
-        )
+    # Define the tool explicitly
+    train_tool = types.Tool(
+        function_declarations=[
+            types.FunctionDeclaration(
+                name="get_train_connections",
+                description="Finds live train connections between two German stations/cities.",
+                parameters={
+                    "type": "OBJECT",
+                    "properties": {
+                        "start_loc": {"type": "STRING", "description": "Departure city"},
+                        "dest_loc": {"type": "STRING", "description": "Destination city"}
+                    },
+                    "required": ["start_loc", "dest_loc"]
+                }
+            )
+        ]
     )
 
     try:
-        # Step 1: Send user input to Gemini
+        # 1. Create the 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
+        )
+
+        # 2. Start the chat session
+        chat = client.chats.create(model=model_id, config=config)
         response = chat.send_message(input_text)
-        
-        # Step 2: Tool Loop
-        # If Gemini emits a tool call, we execute it and send results back
-        while response.candidates[0].content.parts[0].tool_call:
-            tool_call = response.candidates[0].content.parts[0].tool_call
-            
-            if tool_call.name == "get_train_connections":
-                args = tool_call.args
-                # Execute the actual MCP request
-                train_data = asyncio.run(call_mcp_tool(args["start_loc"], args["dest_loc"]))
+
+        # 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
+            if not response.candidates[0].content.parts:
+                break
                 
-                # Feed the train data back to Gemini
-                response = chat.send_message(
-                    types.Part.from_function_response(
-                        name="get_train_connections",
-                        response={"result": train_data}
+            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"]))
+                    tool_responses.append(
+                        types.Part.from_function_response(
+                            name=call.name,
+                            response={"result": result_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 tool_responses:
+                response = chat.send_message(tool_responses)
             else:
-                break # Exit if it's a tool we don't recognize here
+                break
 
         return response.text, ""
 
     except Exception as e:
-        return f"Error: {str(e)}", ""
-
-# --- GRADIO UI ---
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 🚄 Gemini 2.0 Flash + Live DB Trains")
-    gr.Markdown("I can search the web and check live train connections via an MCP Server.")
-    
-    with gr.Column():
-        output_textbox = gr.Markdown(label="Response")
-        input_textbox = gr.Textbox(
-            lines=3, 
-            label="Ask me anything", 
-            placeholder="e.g. 'Is there a train from Berlin to Munich tonight?'"
-        )
-        submit_button = gr.Button("Send", variant="primary")
-    
-    submit_button.click(
-        fn=generate,
-        inputs=input_textbox,
-        outputs=[output_textbox, input_textbox]
-    )
+        return f"### Error encountered\n{str(e)}", ""
 
-if __name__ == '__main__':
-    demo.launch()
+# (Keep your existing Gradio Blocks code below)