Update app.py

app.py

@@ -1,177 +1,65 @@
+import base64
 import gradio as gr
 import os
+import json
 from google import genai
 from google.genai import types
 from gradio_client import Client
 
-# --- 1. Robust Client Initialization for the External Tool ---
-db_client = None
-
-def init_db_client():
-    global db_client
-    try:
-        # Use the full URL to avoid DNS resolution issues
-        print("Connecting to DB Timetable API...")
-        db_client = Client("https://mgokg-db-timetable-api.hf.space/")
-        print("Successfully connected.")
-    except Exception as e:
-        print(f"Warning: Could not connect to DB Timetable API: {e}")
-
-# Attempt connection on startup
-init_db_client()
-
-# --- 2. Tool Definition ---
-
-def get_train_connection(dep: str, dest: str):
-    """
-    Fetches the train timetable between two cities using the external API.
-    """
-    global db_client
-    # Retry connection if it failed initially
-    if db_client is None:
-        init_db_client()
-        if db_client is None:
-            return "Error: The train database is currently unreachable."
-
-    try:
-        # Calling the specific endpoint mentioned in the MCP docs
-        result = db_client.predict(
-            dep=dep,
-            dest=dest,
-            api_name="/db_timetable_api_ui_wrapper"
-        )
-        return result
-    except Exception as e:
-        return f"Error fetching timetable: {str(e)}"
-
-# Define the tool schema for Gemini
-train_tool = types.FunctionDeclaration(
-    name="get_train_connection",
-    description="Find train connections and timetables between a start location (dep) and a destination (dest).",
-    parameters=types.Schema(
-        type=types.Type.OBJECT,
-        properties={
-            "dep": types.Schema(type=types.Type.STRING, description="Departure city or station"),
-            "dest": types.Schema(type=types.Type.STRING, description="Destination city or station"),
-        },
-        required=["dep", "dest"]
-    )
-)
-
-# Map string name to the actual function
-tools_map = {
-    "get_train_connection": get_train_connection
-}
-
-# --- 3. Generation Logic (v1alpha) ---
 
 def generate(input_text):
-    if not input_text:
-        yield "", ""
-        return
-
     try:
-        # Initialize Client with v1alpha
         client = genai.Client(
             api_key=os.environ.get("GEMINI_API_KEY"),
-            #http_options={'api_version': 'v1alpha'}  # <--- ENABLE v1alpha HERE
         )
     except Exception as e:
-        yield f"Error initializing client: {e}. Make sure GEMINI_API_KEY is set.", input_text
-        return
+        return f"Error initializing client: {e}.  Make sure GEMINI_API_KEY is set."
 
-    #model = "gemini-flash-latest"
     model = "gemini-flash-latest"
-
-    # Configure tools
+    contents = [
+        types.Content(
+            role="user",
+            parts=[
+                types.Part.from_text(text=f"{input_text}"),
+            ],
+        ),
+    ]
     tools = [
-        types.Tool(
-            google_search=types.GoogleSearch(),
-            function_declarations=[train_tool]
-        )
+        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",
     )
 
-    contents = [
-        types.Content(
-            role="user",
-            parts=[types.Part.from_text(text=input_text)],
-        ),
-    ]
 
     response_text = ""
-    
     try:
-        # First API Call
-        response = client.models.generate_content(
-            model=model,
-            contents=contents,
-            config=generate_content_config,
-        )
-
-        # Handle Function Calling
-        if response.candidates and response.candidates[0].content.parts:
-            part = response.candidates[0].content.parts[0]
-            
-            if part.function_call:
-                fn_name = part.function_call.name
-                fn_args = part.function_call.args
-                
-                if fn_name in tools_map:
-                    # Execute tool
-                    api_result = tools_map[fn_name](**fn_args)
-                    
-                    # Append history
-                    contents.append(response.candidates[0].content)
-                    contents.append(
-                        types.Content(
-                            role="tool",
-                            parts=[
-                                types.Part.from_function_response(
-                                    name=fn_name,
-                                    response={"result": api_result}
-                                )
-                            ]
-                        )
-                    )
-                    
-                    # Second API Call (Streamed)
-                    stream = client.models.generate_content_stream(
-                        model=model,
-                        contents=contents,
-                        config=generate_content_config
-                    )
-                    
-                    for chunk in stream:
-                        response_text += chunk.text
-                        yield response_text, ""
-                    return
-
-        # Handle standard response
-        if response.text:
-            yield response.text, ""
-            
+      for chunk in client.models.generate_content_stream(
+          model=model,
+          contents=contents,
+          config=generate_content_config,
+      ):
+          response_text += chunk.text
     except Exception as e:
-        yield f"Error during generation: {e}", input_text
-
-# --- 4. UI Setup ---
+        return f"Error during generation: {e}"
+    data = response_text  
+    #data = clean_json_string(response_text)
+    data = data[:-1]
+    return response_text, ""
+   
 
 if __name__ == '__main__':
+
     with gr.Blocks() as demo:
-        gr.Markdown("# Gemini 2.0 Flash  + DB train connections")
-        
+        title=gr.Markdown("# Gemini 2.0 Flash + Websearch")
         output_textbox = gr.Markdown()
-        input_textbox = gr.Textbox(lines=3, label="", placeholder="Ask for a train connection (e.g., 'Train from Berlin to Frankfurt')...")
-        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="", 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)
+