Update app.py

app.py

@@ -1,104 +1,69 @@
-import base64
-import gradio as gr
 import os
-import json
+import asyncio
+import gradio as gr
 from google import genai
 from google.genai import types
-from gradio_client import Client
-
-
-route="""	
-how to handle special case "zugverbindung".
-Wichtig: Dies Regeln gelten nur wenn eine zugverbindung angefragt wird, else answer prompt
-Regeln:
-Wenn eine Zugverbindung von {Startort} nach {Zielort} angefragt wird, return json object with Startort and Zielort.
-always follow json scheme below.
-
-Wichtig: Gib absolut keinen Text vor oder nach dem JSON aus (keine Erklärungen, kein "Hier ist das Ergebnis").
-
-{
-  "start_loc": "fill in Startort here",
-  "dest_loc": "fill in Zielort here"
-}
-
-"""
-
-
+# Import MCP client for Python (standard in 2026)
+from mcp import Client as MCPClient
 
-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
-    
-    # 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 generate(input_text):
+async def generate(input_text):
     try:
+        # Initialize Gemini Client
         client = genai.Client(
             api_key=os.environ.get("GEMINI_API_KEY"),
         )
-    except Exception as e:
-        return f"Error initializing client: {e}.  Make sure GEMINI_API_KEY is set."
-
-    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",
-    )
+        
+        # 1. Connect to the Streamable HTTP MCP server on Hugging Face
+        # 2026 SDKs support direct URL initialization
+        mcp_url = "https://mgokg-db-timetable-api.hf.space/gradio_api/mcp/"
+        
+        async with MCPClient(url=mcp_url) as mcp_session:
+            
+            model = "gemini-2.5-flash" # Version 2026
+            
+            # 2. Define tools: Google Search + MCP Session
+            tools = [
+                types.Tool(google_search=types.GoogleSearch()),
+                mcp_session # Pass the entire MCP server as a toolset here
+            ]
+            
+            generate_content_config = types.GenerateContentConfig(
+                temperature=0.4,
+                tools=tools,
+                response_mime_type="text/plain",
+            )
 
+            response_text = ""
+            # 3. Stream call with MCP tools
+            async for chunk in client.aio.models.generate_content_stream(
+                model=model,
+                contents=input_text,
+                config=generate_content_config,
+            ):
+                if chunk.text:
+                    response_text += chunk.text
+            
+            return response_text, ""
 
-    response_text = ""
-    try:
-      for chunk in client.models.generate_content_stream(
-          model=model,
-          contents=contents,
-          config=generate_content_config,
-      ):
-          response_text += chunk.text
     except Exception as e:
-        return f"Error during generation: {e}"
-    data = clean_json_string(response_text)
-    data = data[:-1]
-    return response_text, ""
-   
+        return f"Error during processing: {str(e)}", ""
 
-if __name__ == '__main__':
+# Helper function for asynchronous Gradio calls
+def gradio_wrapper(input_text):
+    return asyncio.run(generate(input_text))
 
+if __name__ == '__main__':
     with gr.Blocks() as demo:
-        title=gr.Markdown("# Gemini 2.0 Flash + Websearch")
+        gr.Markdown("# Gemini Flash + DB Timetable (MCP)")
         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=3, label="Request",
+                                  placeholder="When is the next train from Berlin to Hamburg?")
+        submit_button = gr.Button("Send Request")
+        
+        submit_button.click(
+            fn=gradio_wrapper,
+            inputs=input_textbox,
+            outputs=[output_textbox, input_textbox]
+        )
+        
+    demo.launch(show_error=True)