Update app.py

app.py

@@ -1,86 +1,108 @@
-import requests
+import os
 import json
-from datetime import datetime
-from mcp.server.fastmcp import FastMCP
+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
 
-# Initialize FastMCP Server
-mcp = FastMCP("DB-Train-Helper")
+# --- CONFIGURATION ---
+MCP_SERVER_URL = "https://mgokg/DB_API_MCP.hf.space/sse"
+GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")
 
-BASE_URL = "https://v6.db.transport.rest"
+# 1. Helper to talk to your MCP Server
+async def call_mcp_train_tool(start_loc, dest_loc):
+    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
+            })
+            # result.content[0].text contains the JSON list from your MCP server
+            return result.content[0].text
 
-def get_station_id(query):
-    try:
-        response = requests.get(f"{BASE_URL}/locations", params={
-            "poi": "false",
-            "addresses": "false",
-            "query": query
-        })
-        data = response.json()
-        if data and len(data) > 0:
-            return data[0]["id"]
-        return None
-    except:
-        return None
+# 2. Tool Definition for Gemini
+# This tells Gemini what the tool does and what arguments it needs
+train_tool_declaration = types.Tool(
+    function_declarations=[
+        types.FunctionDeclaration(
+            name="get_train_connections",
+            description="Get live train connections between two German cities/stations.",
+            parameters=types.Schema(
+                type="OBJECT",
+                properties={
+                    "start_loc": types.Schema(type="STRING", description="The departure station"),
+                    "dest_loc": types.Schema(type="STRING", description="The destination station")
+                },
+                required=["start_loc", "dest_loc"]
+            )
+        )
+    ]
+)
 
-def format_duration(departure_str, arrival_str):
-    fmt = "%Y-%m-%dT%H:%M:%S%z"
-    dep = datetime.strptime(departure_str, fmt)
-    arr = datetime.strptime(arrival_str, fmt)
-    diff = arr - dep
-    hours, remainder = divmod(diff.seconds, 3600)
-    minutes = remainder // 60
-    return f"{hours}h {minutes}min" if hours > 0 else f"{minutes}min"
+def generate(input_text):
+    if not GEMINI_API_KEY:
+        return "Error: GEMINI_API_KEY is not set.", ""
 
-@mcp.tool()
-def get_train_connections(start_loc: str, dest_loc: str) -> str:
-    """
-    Fetches train connections between two German cities and returns 
-    a beautifully formatted HTML snippet with departure times and platforms.
-    """
-    start_id = get_station_id(start_loc)
-    dest_id = get_station_id(dest_loc)
+    client = genai.Client(api_key=GEMINI_API_KEY)
+    model_id = "gemini-2.0-flash-exp" # Or gemini-flash-latest
 
-    if not start_id or not dest_id:
-        return f"Error: Could not find station IDs for {start_loc} or {dest_loc}."
+    # We use a chat session to handle the multi-turn tool loop
+    chat = client.chats.create(
+        model=model_id,
+        config=types.GenerateContentConfig(
+            tools=[train_tool_declaration, types.Tool(google_search=types.GoogleSearch())],
+            temperature=0.4
+        )
+    )
 
     try:
-        response = requests.get(f"{BASE_URL}/journeys", params={
-            "from": start_id,
-            "to": dest_id,
-            "results": 3
-        })
-        journey_data = response.json()
+        # Step 1: Send initial request to Gemini
+        response = chat.send_message(input_text)
         
-        connections_list = []
-        for j in journey_data.get("journeys", []):
-            legs = j.get("legs", [])
-            if not legs: continue
-            
-            first, last = legs[0], legs[-1]
+        # Step 2: Check if Gemini wants to use a tool
+        # We loop in case Gemini needs multiple tool calls
+        while response.candidates[0].content.parts[0].tool_call:
+            tool_call = response.candidates[0].content.parts[0].tool_call
             
-            conn_obj = {
-                "departure": datetime.strptime(first["departure"], "%Y-%m-%dT%H:%M:%S%z").strftime("%H:%M"),
-                "arrival": datetime.strptime(last["arrival"], "%Y-%m-%dT%H:%M:%S%z").strftime("%H:%M"),
-                "startLocation": first["origin"]["name"],
-                "destination": last["destination"]["name"],
-                "duration": format_duration(first["departure"], last["arrival"]),
-                "platform": f"Gl. {first.get('departurePlatform', '-')}"
-            }
-            connections_list.append(conn_obj)
+            if tool_call.name == "get_train_connections":
+                # Extract arguments Gemini provided
+                args = tool_call.args
+                
+                # Execute the MCP call (running async code in sync Gradio)
+                train_data = asyncio.run(call_mcp_train_tool(args["start_loc"], args["dest_loc"]))
+                
+                # Send the tool result back to Gemini
+                response = chat.send_message(
+                    types.Part.from_function_response(
+                        name="get_train_connections",
+                        response={"result": train_data}
+                    )
+                )
+            else:
+                break # Handle other tools or exit
 
-        # We return the HTML just as your original code did
-        connections_json = json.dumps(connections_list)
-        
-        return f"""
-        <div style="font-family: sans-serif; background: #1a1a2e; padding: 20px; border-radius: 10px;">
-            <h2 style="color: white;">Connections: {start_loc} to {dest_loc}</h2>
-            {connections_json}
-            <p style="color: #a0a0a0;">(Code generated for HTML rendering)</p>
-        </div>
-        """
+        return response.text, ""
 
     except Exception as e:
-        return f"Error fetching journeys: {str(e)}"
+        return f"Error during generation: {str(e)}", ""
+
+# --- GRADIO UI ---
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 🚄 Gemini 2.0 + DB Train MCP")
+    gr.Markdown("Ask about train connections (e.g., 'When is the next train from Munich to Berlin?')")
+    
+    output_textbox = gr.Markdown(label="Response")
+    input_textbox = gr.Textbox(lines=2, label="Your Question")
+    submit_button = gr.Button("Send", variant="primary")
+    
+    submit_button.click(
+        fn=generate,
+        inputs=input_textbox,
+        outputs=[output_textbox, input_textbox]
+    )
 
-if __name__ == "__main__":
-    mcp.run()
+if __name__ == '__main__':
+    demo.launch()