Update app.py

app.py

@@ -1,3 +1,132 @@
+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"
+# Ensure your API key is set in your environment variables
+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
+
+async def generate(input_text):
+    """Async generator to handle tool loops and UI updates."""
+    if not GEMINI_API_KEY:
+        yield "### Error\nGEMINI_API_KEY is not set.", ""
+        return
+
+    client = genai.Client(api_key=GEMINI_API_KEY)
+    # Using Gemini 3 Flash as requested
+    model_id = "gemini-flash-latest" 
+
+    # 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.",
+                parameters={
+                    "type": "OBJECT",
+                    "properties": {
+                        "start_loc": {"type": "STRING", "description": "Departure city"},
+                        "dest_loc": {"type": "STRING", "description": "Destination city"}
+                    },
+                    "required": ["start_loc", "dest_loc"]
+                }
+            )
+        ]
+    )
+
+    try:
+        yield "🔍 Thinking...", ""
+        
+        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 = await chat.send_message_async(input_text)
+
+        # --- Manual Tool Loop ---
+        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":
+                    yield f"🚄 Fetching train data for **{call.args['start_loc']}** to **{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(
+                        types.Part.from_function_response(
+                            name=call.name,
+                            response={"result": train_data}
+                        )
+                    )
+            
+            if tool_responses:
+                yield "📝 Processing train data...", ""
+                response = await chat.send_message_async(tool_responses)
+            else:
+                break
+
+        # Final output
+        yield response.text, ""
+
+    except Exception as e:
+        yield f"### Logic Error\n{str(e)}", ""
+
+# --- 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.")
+    
+    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")
+    
+    # Gradio handles async generators automatically
+    submit_button.click(
+        fn=generate,
+        inputs=input_textbox,
+        outputs=[output_textbox, input_textbox]
+    )
+
+if __name__ == '__main__':
+    demo.launch()
+
+
+
+
+"""
 import base64
 import gradio as gr
 import os
@@ -62,4 +191,4 @@ if __name__ == '__main__':
         submit_button = gr.Button("send")
         submit_button.click(fn=generate,inputs=input_textbox,outputs=[output_textbox, input_textbox])
     demo.launch(show_error=True)
-
+"""