Update app.py

app.py

@@ -1,84 +1,86 @@
-import base64
-import gradio as gr
-import os
+import requests
 import json
-from google import genai
-from google.genai import types
-from gradio_client import Client
+from datetime import datetime
+from mcp.server.fastmcp import FastMCP
 
-def clean_json_string(json_str):
-    """
-    Removes any comments or prefixes before the actual JSON content.
-    """
-    # Find the first occurrence of '{'
-    json_start = json_str.find('{')
-    if json_start == -1:
-        # If no '{' is found, try with '[' for arrays
-        json_start = json_str.find('[')
-        if json_start == -1:
-            return json_str  # Return original if no JSON markers found
-    
-    # Extract everything from the first JSON marker
-    cleaned_str = json_str[json_start:]
-    return cleaned_str
-    # Verify it's valid JSON
-    try:
-        json.loads(cleaned_str)
-        return cleaned_str
-    except json.JSONDecodeError:
-        return json_str  # Return original if cleaning results in invalid JSON
+# Initialize FastMCP Server
+mcp = FastMCP("DB-Train-Helper")
 
-def generate(input_text):
+BASE_URL = "https://v6.db.transport.rest"
+
+def get_station_id(query):
     try:
-        client = genai.Client(
-            api_key=os.environ.get("GEMINI_API_KEY"),
-        )
-    except Exception as e:
-        return f"Error initializing client: {e}.  Make sure GEMINI_API_KEY is set."
+        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
 
-    model = "gemini-flash-latest"
-    contents = [
-        types.Content(
-            role="user",
-            parts=[
-                types.Part.from_text(text=input_text),
-            ],
-        ),
-    ]
-    tools = [
-        types.Tool(google_search=types.GoogleSearch()),
-    ]
-    generate_content_config = types.GenerateContentConfig(
-        temperature=0.4,
-        thinking_config = types.ThinkingConfig(
-            thinking_budget=0,
-        ),
-        tools=tools,
-        response_mime_type="text/plain",
-    )
+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"
 
+@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)
+
+    if not start_id or not dest_id:
+        return f"Error: Could not find station IDs for {start_loc} or {dest_loc}."
 
-    response_text = ""
     try:
-      for chunk in client.models.generate_content_stream(
-          model=model,
-          contents=contents,
-          config=generate_content_config,
-      ):
-          response_text += chunk.text
-    except Exception as e:
-        return f"Error during generation: {e}"
-    data = clean_json_string(response_text)
-    data = data[:-1]
-    return response_text, ""
-   
+        response = requests.get(f"{BASE_URL}/journeys", params={
+            "from": start_id,
+            "to": dest_id,
+            "results": 3
+        })
+        journey_data = response.json()
+        
+        connections_list = []
+        for j in journey_data.get("journeys", []):
+            legs = j.get("legs", [])
+            if not legs: continue
+            
+            first, last = legs[0], legs[-1]
+            
+            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 __name__ == '__main__':
+        # 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>
+        """
+
+    except Exception as e:
+        return f"Error fetching journeys: {str(e)}"
 
-    with gr.Blocks() as demo:
-        title=gr.Markdown("# Gemini 2.0 Flash + Websearch")
-        output_textbox = gr.Markdown()
-        input_textbox = gr.Textbox(lines=3, label="", placeholder="Enter message here...")
-        submit_button = gr.Button("send")
-        submit_button.click(fn=generate,inputs=input_textbox,outputs=[output_textbox, input_textbox])
-    demo.launch(show_error=True)
+if __name__ == "__main__":
+    mcp.run()