Update app.py

app.py

@@ -1,332 +1,187 @@
 import gradio as gr
 import os
-import json
 import asyncio
-import httpx
-from typing import Optional, Dict, Any
-
-class MCPHTTPClient:
-    """MCP Client für Streamable HTTP Transport"""
-    
-    def __init__(self, server_url: str):
-        self.server_url = server_url.rstrip('/')
-        self.tools = []
-        self.client = httpx.AsyncClient(timeout=30.0)
-        
-    async def initialize(self):
-        """Initialisiere den Client und hole verfügbare Tools"""
-        try:
-            # MCP über HTTP: Sende initialize request
-            response = await self.client.post(
-                f"{self.server_url}/initialize",
-                json={
-                    "jsonrpc": "2.0",
-                    "id": 1,
-                    "method": "initialize",
-                    "params": {
-                        "protocolVersion": "2024-11-05",
-                        "capabilities": {},
-                        "clientInfo": {
-                            "name": "gradio-mcp-client",
-                            "version": "1.0.0"
-                        }
-                    }
-                }
-            )
-            
-            if response.status_code == 200:
-                print("MCP-Server initialisiert")
-                await self.list_tools()
-                return True
-            else:
-                print(f"Initialisierung fehlgeschlagen: {response.status_code}")
-                return False
-                
-        except Exception as e:
-            print(f"Verbindungsfehler bei Initialisierung: {e}")
-            return False
-    
-    async def list_tools(self):
-        """Rufe verfügbare Tools vom MCP-Server ab"""
-        try:
-            response = await self.client.post(
-                f"{self.server_url}/list_tools",
-                json={
-                    "jsonrpc": "2.0",
-                    "id": 2,
-                    "method": "tools/list",
-                    "params": {}
-                }
-            )
-            
-            if response.status_code == 200:
-                data = response.json()
-                if "result" in data and "tools" in data["result"]:
-                    self.tools = data["result"]["tools"]
-                    print(f"Verfügbare Tools: {[tool['name'] for tool in self.tools]}")
-                else:
-                    print("Keine Tools gefunden")
-            else:
-                print(f"Fehler beim Abrufen der Tools: {response.status_code}")
-                
-        except Exception as e:
-            print(f"Fehler beim Abrufen der Tools: {e}")
-    
-    async def call_tool(self, tool_name: str, arguments: Dict[str, Any]):
-        """Rufe ein Tool auf dem MCP-Server auf"""
-        try:
-            response = await self.client.post(
-                f"{self.server_url}/call_tool",
-                json={
-                    "jsonrpc": "2.0",
-                    "id": 3,
-                    "method": "tools/call",
-                    "params": {
-                        "name": tool_name,
-                        "arguments": arguments
-                    }
-                }
-            )
-            
-            if response.status_code == 200:
-                data = response.json()
-                if "result" in data:
-                    return data["result"]
-                else:
-                    print(f"Unerwartete Antwort: {data}")
-                    return None
-            else:
-                print(f"Tool-Aufruf fehlgeschlagen: {response.status_code}")
-                print(f"Response: {response.text}")
-                return None
-                
-        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:
-                return self._parse_result(result)
-            return None
-            
-        except Exception as e:
-            print(f"Fehler bei der Verbindungssuche: {e}")
-            return None
-    
-    def _parse_result(self, result):
-        """Parse die Antwort vom MCP-Server"""
-        try:
-            # MCP-Server gibt Ergebnisse im content-Array zurück
-            if isinstance(result, dict) and "content" in result:
-                content = result["content"]
-                if isinstance(content, list) and len(content) > 0:
-                    if isinstance(content[0], dict) and "text" in content[0]:
-                        return content[0]["text"]
-                    elif hasattr(content[0], 'text'):
-                        return content[0].text
-                return str(content)
-            
-            return str(result)
-            
-        except Exception as e:
-            print(f"Fehler beim Parsen: {e}")
-            return str(result)
+import json
+from google import genai
+from google.genai import types
+
+# MCP Imports
+from mcp import ClientSession, StdioServerParameters
+from mcp.client.sse import sse_client
+from mcp.types import CallToolResult
+
+# Konfiguration
+MCP_SERVER_URL = "https://mgokg-db-timetable-api.hf.space/gradio_api/mcp/"
+MODEL_ID = "gemini-2.0-flash-exp" # Oder "gemini-flash-latest" je nach Verfügbarkeit
+
+async def generate(input_text, history):
+    """
+    Hauptfunktion:
+    1. Verbindet sich via SSE mit dem MCP Server.
+    2. Holt Tool-Definitionen.
+    3. Sendet User-Input + Tools an Gemini.
+    4. Führt Tool-Calls aus (falls Gemini das will).
+    5. Gibt die Antwort zurück.
+    """
     
-    async def close(self):
-        """Schließe die HTTP-Verbindung"""
-        await self.client.aclose()
-
-
-# Globaler MCP-Client
-mcp_client: Optional[MCPHTTPClient] = 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
+    # 1. Gemini Client initialisieren
     try:
-        mcp_client = MCPHTTPClient(MCP_SERVER_URL)
-        success = await mcp_client.initialize()
-        
-        if success:
-            return "✅ MCP-Client erfolgreich initialisiert"
-        else:
-            return "⚠️ MCP-Client initialisiert, aber keine Tools verfügbar"
-            
+        client = genai.Client(api_key=os.environ.get("GEMINI_API_KEY"))
     except Exception as e:
-        print(f"Fehler bei der Initialisierung: {e}")
-        return f"❌ Fehler bei der Initialisierung: {e}"
+        yield f"Error initializing Gemini client: {e}. Make sure GEMINI_API_KEY is set."
+        return
 
-
-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:
-        init_msg = await initialize_mcp_client()
-        if "❌" in init_msg:
-            return f"{init_msg}\n\nBitte versuchen Sie es später erneut."
-    
+    # 2. Verbindung zum MCP Server aufbauen (Context Manager)
     try:
-        # Extrahiere Start und Ziel aus der Benutzereingabe
-        user_input_lower = user_input.lower().strip()
-        
-        # Verschiedene Parsing-Strategien
-        origin, destination = None, None
-        
-        if " nach " in user_input_lower:
-            parts = user_input_lower.split(" nach ")
-            origin = parts[0].replace("von", "").strip()
-            destination = parts[1].strip()
-        elif " - " in user_input_lower:
-            parts = user_input_lower.split(" - ")
-            origin = parts[0].strip()
-            destination = parts[1].strip()
-        else:
-            # Versuche einfache Tokenisierung
-            words = user_input_lower.replace("von", "").replace(",", "").split()
-            if len(words) >= 2:
-                origin = words[0]
-                destination = words[-1]
-        
-        if not origin or not destination:
-            return """❓ Bitte geben Sie Start- und Zielort an.
-
-**Beispiele:**
-- Frankfurt nach Berlin
-- Von München nach Hamburg  
-- Köln - Stuttgart
-- Hamburg Berlin"""
-        
-        # Kapitalisiere die Namen
-        origin = origin.capitalize()
-        destination = destination.capitalize()
-        
-        # Zeige Suche an
-        print(f"Suche Verbindungen: {origin} → {destination}")
-        
-        # 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 f"""⚠️ Keine Verbindungen gefunden oder Fehler bei der Abfrage.
+        async with sse_client(url=MCP_SERVER_URL) as streams:
+            async with ClientSession(streams.read, streams.write) as session:
+                await session.initialize()
+                
+                # Verfügbare Tools vom MCP Server abrufen
+                mcp_list_tools_result = await session.list_tools()
+                mcp_tools = mcp_list_tools_result.tools
+                
+                # 3. Tools für Gemini konvertieren
+                gemini_tools_declarations = []
+                for tool in mcp_tools:
+                    # Fallback Description, falls leer (wie in der Warnung beschrieben)
+                    description = tool.description or "Tool to query DB (Deutsche Bahn) train timetables and connections."
+                    
+                    gemini_tools_declarations.append(
+                        types.FunctionDeclaration(
+                            name=tool.name,
+                            description=description,
+                            parameters=tool.inputSchema
+                        )
+                    )
+
+                # Zusammenbau der Tools (Websearch + MCP Functions)
+                tools_config = [
+                    types.Tool(google_search=types.GoogleSearch()),
+                    types.Tool(function_declarations=gemini_tools_declarations)
+                ]
+
+                # Konfiguration für den Chat
+                generate_content_config = types.GenerateContentConfig(
+                    temperature=0.4,
+                    tools=tools_config,
+                    system_instruction="You are a helpful assistant. You have access to Google Search and a tool called 'db_timetable_api_ui_wrapper' to check train connections (Deutsche Bahn). Use the train tool when the user asks for train times.",
+                    response_mime_type="text/plain",
+                )
 
-**Gesuchte Route:** {origin} → {destination}
+                # Chat History aufbauen (optional, hier vereinfacht nur aktueller Turn)
+                # Für echte History müsste man 'history' parsen.
+                contents = [
+                    types.Content(
+                        role="user",
+                        parts=[types.Part.from_text(text=input_text)],
+                    ),
+                ]
+
+                # --- Die Chat-Schleife (Model -> Tool -> Model) ---
+                while True:
+                    response_stream = await client.models.generate_content_stream(
+                        model=MODEL_ID,
+                        contents=contents,
+                        config=generate_content_config,
+                    )
+
+                    full_response_text = ""
+                    function_calls = []
+                    
+                    # Stream verarbeiten
+                    async for chunk in response_stream:
+                        # Text Teile sammeln
+                        if chunk.text:
+                            full_response_text += chunk.text
+                            yield full_response_text
+                        
+                        # Funktionsaufrufe sammeln (passieren meist am Ende oder in einem Block)
+                        if chunk.function_calls:
+                            for fc in chunk.function_calls:
+                                function_calls.append(fc)
+
+                    # Wenn keine Funktionsaufrufe -> Fertig
+                    if not function_calls:
+                        break
+                    
+                    # Wenn Funktionsaufrufe vorhanden sind -> Ausführen
+                    # Wir fügen die Antwort des Modells (die den Aufruf enthält) zur Historie hinzu
+                    contents.append(types.Content(
+                        role="model",
+                        parts=[types.Part.from_function_call(name=fc.name, args=fc.args) for fc in function_calls]
+                    ))
+
+                    # Jeden Call ausführen
+                    for fc in function_calls:
+                        yield f"\n\n*Calling MCP Tool: {fc.name}...*\n"
+                        
+                        try:
+                            # Aufruf an den MCP Server
+                            result: CallToolResult = await session.call_tool(
+                                name=fc.name,
+                                arguments=fc.args
+                            )
+                            
+                            # Ergebnis extrahieren (Text aus dem Content)
+                            tool_output_text = ""
+                            if result.content:
+                                for content_item in result.content:
+                                    if content_item.type == "text":
+                                        tool_output_text += content_item.text
+                            
+                            # Ergebnis zur Historie hinzufügen
+                            contents.append(types.Content(
+                                role="tool",
+                                parts=[types.Part.from_function_response(
+                                    name=fc.name,
+                                    response={"result": tool_output_text}
+                                )]
+                            ))
+                            
+                        except Exception as tool_error:
+                            error_msg = f"Error executing tool {fc.name}: {str(tool_error)}"
+                            contents.append(types.Content(
+                                role="tool",
+                                parts=[types.Part.from_function_response(
+                                    name=fc.name,
+                                    response={"error": error_msg}
+                                )]
+                            ))
+                    
+                    # Schleife läuft weiter -> Nächster Aufruf an Gemini mit den Tool-Ergebnissen
 
-Bitte prüfen Sie:
-- Sind die Ortsnamen korrekt geschrieben?
-- Existieren die Bahnhöfe?
-- Versuchen Sie es mit vollständigen Stadtnamen"""
-            
     except Exception as e:
-        error_msg = str(e)
-        print(f"Fehler bei der Suche: {error_msg}")
-        return f"""❌ Fehler bei der Suche: {error_msg}
-
-Bitte versuchen Sie es erneut oder verwenden Sie ein anderes Format.
+        yield f"An error occurred: {e}"
 
-**Beispiele:**
-- Frankfurt nach Berlin
-- Von München nach Hamburg"""
 
+if __name__ == '__main__':
+    with gr.Blocks() as demo:
+        gr.Markdown("# Gemini 2.0 Flash + Websearch + MCP (DB Timetable)")
+        
+        chatbot = gr.Chatbot(height=600, type="messages")
+        msg = gr.Textbox(label="Nachricht eingeben (z.B. 'Zug von Berlin nach München morgen früh')", lines=2)
+        clear = gr.Button("Clear")
 
-def search_train_connections(user_input: str):
-    """Synchrone Wrapper-Funktion für Gradio"""
-    if not user_input or not user_input.strip():
-        return "Bitte geben Sie eine Anfrage ein.", ""
-    
-    try:
-        # 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"❌ Fehler: {str(e)}", ""
+        async def user_message(user_input, history):
+            # Zeige User Nachricht sofort an
+            return "", history + [{"role": "user", "content": user_input}]
 
+        async def bot_response(history):
+            # Letzte User Nachricht holen
+            user_input = history[-1]["content"]
+            
+            # Generator starten
+            # Wir bauen die Antwort Stück für Stück auf
+            history.append({"role": "assistant", "content": ""})
+            
+            async for chunk in generate(user_input, history[:-1]):
+                history[-1]["content"] = chunk
+                yield history
 
-if __name__ == '__main__':
-    # Initialisiere MCP-Client beim Start
-    print("=" * 60)
-    print("🚂 Deutsche Bahn MCP-Client wird gestartet...")
-    print("=" * 60)
-    
-    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.")
-    
-    print("=" * 60)
-    
-    # Gradio Interface
-    with gr.Blocks(theme=gr.themes.Soft()) as demo:
-        title = gr.Markdown("# 🚂 Deutsche Bahn Zugverbindungen")
-        gr.Markdown("### MCP-Client für DB Timetable API")
-        
-        description = gr.Markdown("""
-        Geben Sie Start- und Zielort ein, um aktuelle Zugverbindungen zu finden.
-        
-        **Beispiele für Eingaben:**
-        - `Frankfurt nach Berlin`
-        - `Von München nach Hamburg`
-        - `Köln - Stuttgart`
-        - `Hamburg Berlin`
-        """)
-        
-        with gr.Row():
-            with gr.Column(scale=4):
-                input_textbox = gr.Textbox(
-                    lines=2, 
-                    label="Ihre Anfrage", 
-                    placeholder="z.B. 'Frankfurt nach Berlin' oder 'Von München nach Hamburg'",
-                    autofocus=True
-                )
-            with gr.Column(scale=1):
-                submit_button = gr.Button("🔍 Suchen", variant="primary", size="lg")
-        
-        output_textbox = gr.Markdown(label="Ergebnisse")
-        
-        gr.Markdown("""
-        ---
-        💡 **Hinweis:** Dieser Client nutzt die Deutsche Bahn Timetable API über einen MCP-Server.
-        """)
-        
-        # Event Handler
-        submit_button.click(
-            fn=search_train_connections,
-            inputs=input_textbox,
-            outputs=[output_textbox, input_textbox]
+        msg.submit(user_message, [msg, chatbot], [msg, chatbot], queue=False).then(
+            bot_response, [chatbot], [chatbot]
         )
-        
-        # Enter-Taste funktioniert auch
-        input_textbox.submit(
-            fn=search_train_connections,
-            inputs=input_textbox,
-            outputs=[output_textbox, input_textbox]
-        )
-    
+        clear.click(lambda: None, None, chatbot, queue=False)
+
     demo.launch(show_error=True)