Update app.py

app.py

@@ -1,87 +1,204 @@
-import base64
 import gradio as gr
 import os
 import json
-from google import genai
-from google.genai import types
-from gradio_client import Client
+import asyncio
+from typing import Optional
+from mcp import ClientSession, StdioServerParameters
+from mcp.client.stdio import stdio_client
+import httpx
 
-
-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
+class MCPTrainClient:
+    """MCP Client für Zugverbindungsabfragen"""
     
-    # 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
+    def __init__(self, server_url: str):
+        self.server_url = server_url
+        self.session: Optional[ClientSession] = None
+        self.tools = []
+        
+    async def connect(self):
+        """Stelle Verbindung zum MCP-Server her"""
+        try:
+            # Für Streamable HTTP Transport verwenden wir httpx
+            # Die Session wird über die MCP-URL initialisiert
+            self.session = ClientSession()
+            
+            # Tools vom Server abrufen
+            await self.list_tools()
+            return True
+        except Exception as e:
+            print(f"Verbindungsfehler: {e}")
+            return False
+    
+    async def list_tools(self):
+        """Rufe verfügbare Tools vom MCP-Server ab"""
+        try:
+            if self.session:
+                response = await self.session.list_tools()
+                self.tools = response.tools
+                print(f"Verfügbare Tools: {[tool.name for tool in self.tools]}")
+        except Exception as e:
+            print(f"Fehler beim Abrufen der Tools: {e}")
+    
+    async def call_tool(self, tool_name: str, arguments: dict):
+        """Rufe ein spezifisches Tool auf dem MCP-Server auf"""
+        try:
+            if not self.session:
+                await self.connect()
+            
+            result = await self.session.call_tool(tool_name, arguments)
+            return result
+        except Exception as e:
+            print(f"Fehler beim Tool-Aufruf: {e}")
+            return None
+    
+    async def search_connections(self, origin: str, destination: str):
+        """Suche Zugverbindungen über den MCP-Server"""
+        try:
+            # Verwende das db_timetable_api_ui_wrapper Tool
+            result = await self.call_tool(
+                "db_timetable_api_ui_wrapper",
+                {
+                    "origin": origin,
+                    "destination": destination
+                }
+            )
+            
+            if result and hasattr(result, 'content'):
+                return self._parse_result(result.content)
+            return None
+        except Exception as e:
+            print(f"Fehler bei der Verbindungssuche: {e}")
+            return None
+    
+    def _parse_result(self, content):
+        """Parse die Antwort vom MCP-Server"""
+        try:
+            if isinstance(content, list) and len(content) > 0:
+                # Extrahiere den Text aus der Antwort
+                if hasattr(content[0], 'text'):
+                    return content[0].text
+                elif isinstance(content[0], dict) and 'text' in content[0]:
+                    return content[0]['text']
+            return str(content)
+        except Exception as e:
+            print(f"Fehler beim Parsen: {e}")
+            return None
+
 
-def generate(input_text):
+# Globaler MCP-Client
+mcp_client = None
+MCP_SERVER_URL = "https://mgokg-db-timetable-api.hf.space/gradio_api/mcp/"
+
+
+async def initialize_mcp_client():
+    """Initialisiere den MCP-Client"""
+    global mcp_client
     try:
-        client = genai.Client(
-            api_key=os.environ.get("GEMINI_API_KEY"),
-        )
+        mcp_client = MCPTrainClient(MCP_SERVER_URL)
+        await mcp_client.connect()
+        return "MCP-Client erfolgreich initialisiert"
     except Exception as e:
-        return f"Error initializing client: {e}.  Make sure GEMINI_API_KEY is set."
+        return f"Fehler bei der Initialisierung: {e}"
 
-    model = "gemini-flash-latest"
-    contents = [
-        types.Content(
-            role="user",
-            parts=[
-                types.Part.from_text(text=f"{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",
-    )
+
+async def search_train_connections_async(user_input: str):
+    """Asynchrone Funktion zur Suche von Zugverbindungen"""
+    global mcp_client
+    
+    # Initialisiere Client falls noch nicht geschehen
+    if mcp_client is None:
+        await initialize_mcp_client()
+    
+    try:
+        # Extrahiere Start und Ziel aus der Benutzereingabe
+        # Einfache Parsing-Logik (kann erweitert werden)
+        parts = user_input.lower().replace("von", "").replace("nach", "").strip().split()
+        
+        if len(parts) < 2:
+            return "Bitte geben Sie Start- und Zielort an (z.B. 'Frankfurt nach Berlin')"
+        
+        # Versuche, Start und Ziel zu identifizieren
+        if "nach" in user_input.lower():
+            split_parts = user_input.lower().split("nach")
+            origin = split_parts[0].replace("von", "").strip()
+            destination = split_parts[1].strip()
+        else:
+            origin = parts[0].capitalize()
+            destination = parts[-1].capitalize()
+        
+        # Rufe MCP-Server auf
+        result = await mcp_client.search_connections(origin, destination)
+        
+        if result:
+            # Formatiere Ergebnis als Markdown
+            response = f"## Zugverbindungen von {origin} nach {destination}\n\n"
+            response += result
+            return response
+        else:
+            return "Keine Verbindungen gefunden oder Fehler bei der Abfrage."
+            
+    except Exception as e:
+        return f"Fehler bei der Suche: {str(e)}\n\nBitte versuchen Sie es erneut oder prüfen Sie Ihre Eingabe."
 
 
-    response_text = ""
+def search_train_connections(user_input: str):
+    """Synchrone Wrapper-Funktion für Gradio"""
     try:
-      for chunk in client.models.generate_content_stream(
-          model=model,
-          contents=contents,
-          config=generate_content_config,
-      ):
-          response_text += chunk.text
+        # Erstelle neuen Event Loop für die asynchrone Funktion
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+        result = loop.run_until_complete(search_train_connections_async(user_input))
+        loop.close()
+        return result, ""
     except Exception as e:
-        return f"Error during generation: {e}"
-    data = response_text  
-    #data = clean_json_string(response_text)
-    data = data[:-1]
-    return response_text, ""
-   
+        return f"Fehler: {str(e)}", ""
 
-if __name__ == '__main__':
 
+if __name__ == '__main__':
+    # Initialisiere MCP-Client beim Start
+    print("Initialisiere MCP-Client...")
+    try:
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+        init_result = loop.run_until_complete(initialize_mcp_client())
+        print(init_result)
+        loop.close()
+    except Exception as e:
+        print(f"Warnung bei Initialisierung: {e}")
+        print("Client wird bei erster Anfrage initialisiert.")
+    
+    # Gradio Interface
     with gr.Blocks() as demo:
-        title=gr.Markdown("# Gemini 2.0 Flash + Websearch")
+        title = gr.Markdown("# 🚂 Deutsche Bahn Zugverbindungen (MCP-Client)")
+        description = gr.Markdown("""
+        Geben Sie Start- und Zielort ein, um Zugverbindungen zu finden.
+        
+        **Beispiele:**
+        - Frankfurt nach Berlin
+        - Von München nach Hamburg
+        - Köln Stuttgart
+        """)
+        
         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)
-
+        input_textbox = gr.Textbox(
+            lines=2, 
+            label="Ihre Anfrage", 
+            placeholder="z.B. 'Frankfurt nach Berlin' oder 'Von München nach Hamburg'"
+        )
+        submit_button = gr.Button("🔍 Verbindungen suchen")
+        
+        # Event Handler
+        submit_button.click(
+            fn=search_train_connections,
+            inputs=input_textbox,
+            outputs=[output_textbox, input_textbox]
+        )
+        
+        # Enter-Taste funktioniert auch
+        input_textbox.submit(
+            fn=search_train_connections,
+            inputs=input_textbox,
+            outputs=[output_textbox, input_textbox]
+        )
+    
+    demo.launch(show_error=True)