Update app.py

app.py

@@ -8,8 +8,7 @@ from mcp import ClientSession
 from mcp.client.sse import sse_client
 
 # --- CONFIGURATION ---
-MCP_SERVER_URL = "https://mgokg-db-api-mcp.hf.space/sse"
-# Ensure your API key is set in your environment variables
+MCP_SERVER_URL = "https://mgokg-db-timetable-api.hf.space/gradio_api/mcp/"
 GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")
 
 async def call_mcp_tool(start_loc, dest_loc):
@@ -24,16 +23,15 @@ async def call_mcp_tool(start_loc, dest_loc):
             return result.content[0].text
 
 async def generate(input_text):
-    """Async generator to handle tool loops and UI updates."""
+    """Async generator handling the tool loop with the corrected aio client."""
     if not GEMINI_API_KEY:
         yield "### Error\nGEMINI_API_KEY is not set.", ""
         return
 
+    # Initialize the standard client
     client = genai.Client(api_key=GEMINI_API_KEY)
-    # Using Gemini 3 Flash as requested
-    model_id = "gemini-flash-latest" 
+    model_id = "gemini-flash-latest"  # Ensure this is a valid model ID for your tier
 
-    # Define the train tool schema
     train_tool = types.Tool(
         function_declarations=[
             types.FunctionDeclaration(
@@ -59,13 +57,16 @@ async def generate(input_text):
             temperature=0.3
         )
 
-        chat = client.chats.create(model=model_id, config=config)
-        response = await chat.send_message_async(input_text)
+        # FIX: Use client.aio to create an async chat session
+        chat = client.aio.chats.create(model=model_id, config=config)
+        
+        # FIX: Use await chat.send_message (no _async suffix needed)
+        response = await chat.send_message(input_text)
 
         # --- Manual Tool Loop ---
         max_turns = 5
         for _ in range(max_turns):
-            if not response.candidates[0].content.parts:
+            if not response.candidates or not response.candidates[0].content.parts:
                 break
             
             tool_calls = [p.tool_call for p in response.candidates[0].content.parts if p.tool_call]
@@ -75,9 +76,8 @@ async def generate(input_text):
             tool_responses = []
             for call in tool_calls:
                 if call.name == "get_train_connections":
-                    yield f"🚄 Fetching train data for **{call.args['start_loc']}** to **{call.args['dest_loc']}**...", ""
+                    yield f"🚄 Fetching train data: **{call.args['start_loc']}** → **{call.args['dest_loc']}**...", ""
                     
-                    # Execute the MCP call
                     train_data = await call_mcp_tool(call.args["start_loc"], call.args["dest_loc"])
                     
                     tool_responses.append(
@@ -88,12 +88,12 @@ async def generate(input_text):
                     )
             
             if tool_responses:
-                yield "📝 Processing train data...", ""
-                response = await chat.send_message_async(tool_responses)
+                yield "📝 Finalizing response...", ""
+                # Send the tool results back using the async chat
+                response = await chat.send_message(tool_responses)
             else:
                 break
 
-        # Final output
         yield response.text, ""
 
     except Exception as e:
@@ -101,19 +101,12 @@ async def generate(input_text):
 
 # --- GRADIO UI ---
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 🚄 Gemini 3 Flash + Live DB Trains")
-    gr.Markdown("I search the web and check live train connections via an MCP Server.")
+    gr.Markdown("# 🚄 Gemini MCP Train Assistant")
     
-    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")
+    output_textbox = gr.Markdown(label="Response")
+    input_textbox = gr.Textbox(label="Ask about German train connections", placeholder="Berlin to Munich?")
+    submit_button = gr.Button("Send", variant="primary")
     
-    # Gradio handles async generators automatically
     submit_button.click(
         fn=generate,
         inputs=input_textbox,