Update app.py

app.py

@@ -10,20 +10,15 @@ from mcp import ClientSession, StdioServerParameters
 from mcp.client.sse import sse_client
 from mcp.types import CallToolResult
 
-# Konfiguration
+# --- 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
+# Stelle sicher, dass du einen validen Model-Namen hast (z.B. gemini-2.0-flash-exp oder gemini-1.5-flash)
+MODEL_ID = "gemini-2.0-flash-exp" 
 
-async def generate(input_text, history):
+async def generate_response(user_input, 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.
+    Core Logic: Verbindet MCP, Google Search und Gemini.
     """
-    
     # 1. Gemini Client initialisieren
     try:
         client = genai.Client(api_key=os.environ.get("GEMINI_API_KEY"))
@@ -31,54 +26,52 @@ async def generate(input_text, history):
         yield f"Error initializing Gemini client: {e}. Make sure GEMINI_API_KEY is set."
         return
 
-    # 2. Verbindung zum MCP Server aufbauen (Context Manager)
+    # 2. Verbindung zum MCP Server aufbauen (SSE)
     try:
         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
+                # Tools vom MCP Server abrufen
+                mcp_list = await session.list_tools()
                 
-                # 3. Tools für Gemini konvertieren
+                # 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."
-                    
+                for tool in mcp_list.tools:
+                    desc = tool.description or "Tool to query DB (Deutsche Bahn) train timetables."
                     gemini_tools_declarations.append(
                         types.FunctionDeclaration(
                             name=tool.name,
-                            description=description,
+                            description=desc,
                             parameters=tool.inputSchema
                         )
                     )
 
-                # Zusammenbau der Tools (Websearch + MCP Functions)
+                # Tools zusammenstellen (Google Search + MCP)
                 tools_config = [
                     types.Tool(google_search=types.GoogleSearch()),
                     types.Tool(function_declarations=gemini_tools_declarations)
                 ]
 
-                # Konfiguration für den Chat
+                # Config
                 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.",
+                    system_instruction="You are a helpful assistant. Use 'db_timetable_api_ui_wrapper' for German train connections. Use Google Search for other current info.",
                     response_mime_type="text/plain",
                 )
 
-                # Chat History aufbauen (optional, hier vereinfacht nur aktueller Turn)
-                # Für echte History müsste man 'history' parsen.
+                # History für Gemini aufbereiten
+                # Wir nehmen hier vereinfacht nur den aktuellen User-Input, 
+                # für echte Multi-Turn-Gespräche müsste man 'history' parsen.
                 contents = [
                     types.Content(
                         role="user",
-                        parts=[types.Part.from_text(text=input_text)],
+                        parts=[types.Part.from_text(text=user_input)],
                     ),
                 ]
 
-                # --- Die Chat-Schleife (Model -> Tool -> Model) ---
+                # --- Loop für Tool Calls ---
                 while True:
                     response_stream = await client.models.generate_content_stream(
                         model=MODEL_ID,
@@ -86,102 +79,101 @@ async def generate(input_text, history):
                         config=generate_content_config,
                     )
 
-                    full_response_text = ""
+                    full_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
+                            full_text += chunk.text
+                            yield full_text # Live-Update des Textes
                         
-                        # 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
+                        break # Keine Tools mehr, wir sind fertig
                     
-                    # Wenn Funktionsaufrufe vorhanden sind -> Ausführen
-                    # Wir fügen die Antwort des Modells (die den Aufruf enthält) zur Historie hinzu
+                    # Tool Calls verarbeiten
+                    # 1. Modell-Antwort zur History hinzufügen (damit es weiß, dass es gefragt hat)
                     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
+                    # 2. Tools ausführen
                     for fc in function_calls:
-                        yield f"\n\n*Calling MCP Tool: {fc.name}...*\n"
+                        yield f"\n\n*Rufe MCP Tool auf: {fc.name}...*\n"
                         
                         try:
-                            # Aufruf an den MCP Server
-                            result: CallToolResult = await session.call_tool(
-                                name=fc.name,
-                                arguments=fc.args
-                            )
+                            result = await session.call_tool(name=fc.name, arguments=fc.args)
                             
-                            # Ergebnis extrahieren (Text aus dem Content)
-                            tool_output_text = ""
+                            # Ergebnis extrahieren
+                            tool_output = ""
                             if result.content:
-                                for content_item in result.content:
-                                    if content_item.type == "text":
-                                        tool_output_text += content_item.text
+                                for item in result.content:
+                                    if item.type == "text":
+                                        tool_output += item.text
                             
-                            # Ergebnis zur Historie hinzufügen
+                            # Ergebnis zur History
                             contents.append(types.Content(
                                 role="tool",
                                 parts=[types.Part.from_function_response(
                                     name=fc.name,
-                                    response={"result": tool_output_text}
+                                    response={"result": tool_output}
                                 )]
                             ))
-                            
-                        except Exception as tool_error:
-                            error_msg = f"Error executing tool {fc.name}: {str(tool_error)}"
+                        except Exception as e:
                             contents.append(types.Content(
                                 role="tool",
                                 parts=[types.Part.from_function_response(
                                     name=fc.name,
-                                    response={"error": error_msg}
+                                    response={"error": str(e)}
                                 )]
                             ))
                     
-                    # Schleife läuft weiter -> Nächster Aufruf an Gemini mit den Tool-Ergebnissen
+                    # Schleife läuft weiter -> Gemini bekommt Tool-Ergebnisse und antwortet neu
 
     except Exception as e:
-        yield f"An error occurred: {e}"
+        yield f"Ein Fehler ist aufgetreten: {e}"
 
 
 if __name__ == '__main__':
     with gr.Blocks() as demo:
-        gr.Markdown("# Gemini 2.0 Flash + Websearch + MCP (DB Timetable)")
+        gr.Markdown("# Gemini 2.0 Flash + Websearch + MCP (DB)")
         
-        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")
-
-        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
+        # Chat-Fenster
+        chatbot = gr.Chatbot(height=500, type="messages")
+        
+        # Eingabebereich
+        with gr.Row():
+            msg = gr.Textbox(
+                scale=4, 
+                show_label=False, 
+                placeholder="Nachricht eingeben (z.B. Zug von Berlin nach München)...",
+                container=False
+            )
+            submit_btn = gr.Button("Send", scale=1, variant="primary")
+
+        # Hilfsfunktionen für Gradio Event-Handling
+        async def user_turn(user_message, history):
+            return "", history + [{"role": "user", "content": user_message}]
+
+        async def bot_turn(history):
+            user_message = history[-1]["content"]
             history.append({"role": "assistant", "content": ""})
             
-            async for chunk in generate(user_input, history[:-1]):
+            async for chunk in generate_response(user_message, history):
                 history[-1]["content"] = chunk
                 yield history
 
-        msg.submit(user_message, [msg, chatbot], [msg, chatbot], queue=False).then(
-            bot_response, [chatbot], [chatbot]
+        # Events: Sowohl Enter (msg.submit) als auch Klick (submit_btn.click) lösen das Gleiche aus
+        msg.submit(user_turn, [msg, chatbot], [msg, chatbot], queue=False).then(
+            bot_turn, [chatbot], [chatbot]
+        )
+        
+        submit_btn.click(user_turn, [msg, chatbot], [msg, chatbot], queue=False).then(
+            bot_turn, [chatbot], [chatbot]
         )
-        clear.click(lambda: None, None, chatbot, queue=False)
 
     demo.launch(show_error=True)