Update app.py

app.py

@@ -1,39 +1,12 @@
-import os
-import json
-import asyncio
-import gradio as gr
-from google import genai
-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):
-    """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()
-            result = await session.call_tool("get_train_connections", arguments={
-                "start_loc": start_loc,
-                "dest_loc": dest_loc
-            })
-            return result.content[0].text
-
-def generate(input_text):
+# Change the function to be async
+async def generate(input_text):
     if not GEMINI_API_KEY:
-        return "Error: GEMINI_API_KEY is not set.", ""
+        yield "Error: GEMINI_API_KEY is not set."
+        return
 
     client = genai.Client(api_key=GEMINI_API_KEY)
-    
-    # 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-3-flash-preview"  # Korrekte Modell-ID verwenden
-
+    model_id = "gemini-3-flash-preview" 
 
-    # Define the train tool schema
     train_tool = types.Tool(
         function_declarations=[
             types.FunctionDeclaration(
@@ -52,24 +25,20 @@ def generate(input_text):
     )
 
     try:
-        # 1. Setup config WITHOUT thinking_config
         config = types.GenerateContentConfig(
             tools=[train_tool, types.Tool(google_search=types.GoogleSearch())],
             temperature=0.3
         )
 
         chat = client.chats.create(model=model_id, config=config)
-        response = chat.send_message(input_text)
+        # Use await for the async call
+        response = await chat.send_message_async(input_text)
 
-        # 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
             
-            # 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
@@ -77,9 +46,8 @@ def generate(input_text):
             tool_responses = []
             for call in tool_calls:
                 if call.name == "get_train_connections":
-                    # 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"]))
+                    # 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(
@@ -88,36 +56,12 @@ def generate(input_text):
                         )
                     )
             
-            # If we have results, send them back to the model
             if tool_responses:
-                response = chat.send_message(tool_responses)
+                response = await chat.send_message_async(tool_responses)
             else:
                 break
 
         return response.text, ""
 
     except Exception as e:
-        return f"### Logic Error\n{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]
-    )
-
-if __name__ == '__main__':
-    demo.launch()
+        return f"### Logic Error\n{str(e)}", ""