Update app.py

app.py

@@ -1,65 +1,87 @@
-import os
-import asyncio
+import base64
 import gradio as gr
+import os
+import json
 from google import genai
 from google.genai import types
-from mcp import ClientSession
-from mcp.client.sse import sse_client
+from gradio_client import Client
 
-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]
-                    ],
-                )
 
-                response_text = ""
-                # Stream-Generierung der Antwort [4]
-                async for chunk in client.aio.models.generate_content_stream(
-                    model="gemini-flash-latest",
-                    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]
+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):
+    try:
+        client = genai.Client(
+            api_key=os.environ.get("GEMINI_API_KEY"),
+        )
     except Exception as e:
-        return f"Verbindung zum DB-Fahrplan fehlgeschlagen: {str(e)}", ""
+        return f"Error initializing client: {e}.  Make sure GEMINI_API_KEY is set."
 
-# Wrapper für die synchrone Gradio-Umgebung [2]
-def gradio_wrapper(input_text):
-    return asyncio.run(generate_response(input_text))
+    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",
+    )
 
-# 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]
 
-    # Event-Handling: Klick löst die Generierung aus
-    btn.click(fn=gradio_wrapper, inputs=input_tx, outputs=[output_md, input_tx])
+    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 = response_text  
+    #data = clean_json_string(response_text)
+    data = data[:-1]
+    return response_text, ""
+   
 
 if __name__ == '__main__':
-    demo.launch()
+
+    with gr.Blocks() as demo:
+        title=gr.Markdown("# Gemini 2.0 Flash + Websearch")
+        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)
+