Update app.py

app.py

@@ -1,67 +1,21 @@
-# Change the function to be async
 async def generate(input_text):
     if not GEMINI_API_KEY:
-        yield "Error: GEMINI_API_KEY is not set."
-        return
+        yield "Error: GEMINI_API_KEY is not set.", ""
+        return # Stops the generator
 
-    client = genai.Client(api_key=GEMINI_API_KEY)
-    model_id = "gemini-3-flash-preview" 
-
-    train_tool = types.Tool(
-        function_declarations=[
-            types.FunctionDeclaration(
-                name="get_train_connections",
-                description="Finds live train connections between two German stations.",
-                parameters={
-                    "type": "OBJECT",
-                    "properties": {
-                        "start_loc": {"type": "STRING", "description": "Departure city"},
-                        "dest_loc": {"type": "STRING", "description": "Destination city"}
-                    },
-                    "required": ["start_loc", "dest_loc"]
-                }
-            )
-        ]
-    )
+    # ... (setup code) ...
 
     try:
-        config = types.GenerateContentConfig(
-            tools=[train_tool, types.Tool(google_search=types.GoogleSearch())],
-            temperature=0.3
-        )
-
-        chat = client.chats.create(model=model_id, config=config)
-        # Use await for the async call
-        response = await chat.send_message_async(input_text)
-
-        max_turns = 5
-        for _ in range(max_turns):
-            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]
-            if not tool_calls:
-                break
-
-            tool_responses = []
-            for call in tool_calls:
-                if call.name == "get_train_connections":
-                    # Now we can simply await the tool call directly
-                    train_data = await call_mcp_tool(call.args["start_loc"], call.args["dest_loc"])
-                    
-                    tool_responses.append(
-                        types.Part.from_function_response(
-                            name=call.name,
-                            response={"result": train_data}
-                        )
-                    )
-            
-            if tool_responses:
-                response = await chat.send_message_async(tool_responses)
-            else:
-                break
+        yield "Thinking...", "" # UI update
+        
+        # ... (logic loop) ...
+        
+        if tool_calls:
+            yield "Fetching live train data from MCP server...", ""
+            # ... (execute tool) ...
 
-        return response.text, ""
+        # Final result
+        yield response.text, ""
 
     except Exception as e:
-        return f"### Logic Error\n{str(e)}", ""
+        yield f"### Logic Error\n{str(e)}", ""