Update app.py

app.py

@@ -1,44 +1,58 @@
+import os
 import asyncio
-from mcp import ClientSession
-from mcp.client.sse import sse_client
-
-async def run_mcp_client():
-    # Die Basis-URL deines Hugging Face Spaces (mit /sse Endpunkt für das Python SDK)
-    server_url = "https://mgokg-db-timetable-api.hf.space"
+import gradio as gr
+from google import genai
+from google.genai import types
 
-    print(f"Verbinde zu MCP Server: {server_url}...")
+from mcp import ClientSession
+from mcp.client.sse import sse_client  # Spezifischer Transport für Gradio/HF
 
-    # 1. Transport-Layer aufbauen (SSE)
-    async with sse_client(url=server_url) as (read_stream, write_stream):
+async def generate(input_text):
+    # WICHTIG: Gradio MCP Server benötigen oft das Suffix /sse
+    mcp_url = "https://mgokg-db-timetable-api.hf.space"
+    
+    try:
+        client = genai.Client(api_key=os.environ.get("GEMINI_API_KEY"))
         
-        # 2. MCP-Session initialisieren
-        async with ClientSession(read_stream, write_stream) as session:
-            await session.initialize()
-            print("Verbindung erfolgreich hergestellt!\n")
-
-            # --- Beispiel 1: Verfügbare Tools auflisten ---
-            print("Verfügbare Tools:")
-            tools_result = await session.list_tools()
-            for tool in tools_result.tools:
-                print(f"- {tool.name}: {tool.description}")
+        # SSE Transport nutzen, um den 'text/html' Fehler zu vermeiden
+        async with sse_client(url=mcp_url) as (read_stream, write_stream):
+            async with ClientSession(read_stream, write_stream) as mcp_session:
+                await mcp_session.initialize()
 
-            # --- Beispiel 2: Ein Tool aufrufen ---
-            # Ersetze 'get_timetable' und die Argumente durch die echten Tool-Namen deiner API
-            if tools_result.tools:
-                tool_to_call = tools_result.tools[0].name
-                print(f"\nRufe Tool '{tool_to_call}' auf...")
+                model_id = "gemini-2.0-flash" 
                 
-                # Beispielhafter Aufruf (Passe die Argumente an deine API an!)
-                result = await session.call_tool(
-                    name=tool_to_call,
-                    arguments={"station": "Berlin Hbf"} 
+                generate_content_config = types.GenerateContentConfig(
+                    temperature=0.4,
+                    tools=[
+                        types.Tool(google_search=types.GoogleSearch()),
+                        mcp_session # Reicht die Tools des DB-Servers an Gemini durch
+                    ],
                 )
-                print(f"Ergebnis: {result.content}")
 
-if __name__ == "__main__":
-    try:
-        asyncio.run(run_mcp_client())
-    except KeyboardInterrupt:
-        pass
+                response_text = ""
+                async for chunk in client.aio.models.generate_content_stream(
+                    model=model_id,
+                    contents=input_text,
+                    config=generate_content_config,
+                ):
+                    if chunk.text:
+                        response_text += chunk.text
+                
+                return response_text, ""
+
     except Exception as e:
-        print(f"Fehler: {e}")
+        return f"Verbindung zum DB-Fahrplan fehlgeschlagen: {str(e)}", ""
+
+def gradio_wrapper(input_text):
+    return asyncio.run(generate(input_text))
+
+if __name__ == '__main__':
+    with gr.Blocks() as demo:
+        gr.Markdown("# Gemini Flash + DB Timetable")
+        input_tx = gr.Textbox(label="Anfrage", placeholder="Wann fährt der nächste Zug von Berlin nach Hamburg?")
+        btn = gr.Button("Senden")
+        output_md = gr.Markdown()
+        
+        btn.click(fn=gradio_wrapper, inputs=input_tx, outputs=[output_md, input_tx])
+        
+    demo.launch()