Update app.py

app.py

@@ -1,66 +1,65 @@
 import os
 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
 
-# Konfiguration des MCP-Servers (Beispiel DB-Timetable auf Hugging Face)
-MCP_URL = "https://mgokg-db-timetable-api.hf.space/gradio_api/mcp/"
-
-async def fetch_train_connections(prompt: str):
-    client = genai.Client(api_key=os.environ.get("GEMINI_API_KEY"))
-    model_id = "gemini-flash-latest"
+async def generate_response(input_text):
+    mcp_url = "https://mgokg-db-timetable-api.hf.space/gradio_api/mcp/"
+    try:
+        client = genai.Client(api_key=os.environ.get("GEMINI_API_KEY"))
+        
+        # Verbindung zum MCP-Server via SSE (Server-Sent Events) [3]
+        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() [4]
+                
+                # Konfiguration der Tools (Google Search + MCP Tools) [4]
+                generate_content_config = types.GenerateContentConfig(
+                    temperature=0.4,
+                    tools=[
+                        types.Tool(google_search=types.GoogleSearch()),
+                        mcp_session # Reicht MCP-Tools an Gemini durch [4]
+                    ],
+                )
 
-    # 1. Verbindung zum MCP-Server über SSE herstellen [3]
-    async with sse_client(url=MCP_URL) as (read_stream, write_stream):
-        async with ClientSession(read_stream, write_stream) as session:
-            # Initialisierung der MCP-Sitzung [4], [5]
-            await session.initialize()
+                response_text = ""
+                # Stream-Generierung der Antwort [4]
+                async for chunk in client.aio.models.generate_content_stream(
+                    model="gemini-2.0-flash",
+                    contents=input_text,
+                    config=generate_content_config,
+                ):
+                    if chunk.text:
+                        response_text += chunk.text
+                
+                return response_text, "" # Rückgabe für Markdown und Textbox-Reset [2]
 
-            # 2. MCP-Tools abrufen und für Gemini konvertieren [4], [2]
-            mcp_tools = await session.list_tools()
-            tools = types.Tool(function_declarations=[
-                {
-                    "name": tool.name,
-                    "description": tool.description,
-                    "parameters": tool.inputSchema,
-                }
-                for tool in mcp_tools.tools
-            ])
+    except Exception as e:
+        return f"Verbindung zum DB-Fahrplan fehlgeschlagen: {str(e)}", ""
 
-            # 3. Anfrage an das Modell senden [6]
-            contents = [types.Content(role="user", parts=[types.Part(text=prompt)])]
-            response = await client.aio.models.generate_content(
-                model=model_id,
-                contents=contents,
-                config=types.GenerateContentConfig(
-                    tools=[tools],
-                    temperature=0.4
-                )
-            )
+# Wrapper für die synchrone Gradio-Umgebung [2]
+def gradio_wrapper(input_text):
+    return asyncio.run(generate_response(input_text))
 
-            # 4. Tool-Calling Loop: Falls das Modell eine Zugverbindung sucht [7], [8]
-            if response.candidates.content.parts.function_call:
-                fc = response.candidates.content.parts.function_call
-                
-                # Tool auf dem MCP-Server ausführen
-                tool_result = await session.call_tool(fc.name, fc.args)
-                
-                # Ergebnis an das Modell zurückgeben für die finale Antwort [8], [9]
-                tool_response_part = types.Part.from_function_response(
-                    name=fc.name, 
-                    response={"result": tool_result.content.text}
-                )
-                
-                contents.append(response.candidates.content)
-                contents.append(types.Content(role="user", parts=[tool_response_part]))
+# Aufbau des ursprünglichen Interfaces [2]
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 🚆 Gemini Flash + DB Timetable")
+    
+    with gr.Row():
+        input_tx = gr.Textbox(
+            label="Anfrage", 
+            placeholder="Wann fährt der nächste Zug von Berlin nach Hamburg?",
+            lines=3
+        )
+    
+    btn = gr.Button("Senden", variant="primary")
+    output_md = gr.Markdown() # Bereich für die Antwort [2]
 
-                final_response = await client.aio.models.generate_content(
-                    model=model_id,
-                    contents=contents,
-                    config=types.GenerateContentConfig(tools=[tools])
-                )
-                return final_response.text
+    # Event-Handling: Klick löst die Generierung aus
+    btn.click(fn=gradio_wrapper, inputs=input_tx, outputs=[output_md, input_tx])
 
-            return response.text
+if __name__ == '__main__':
+    demo.launch()