Update app.py

app.py

@@ -1,177 +1,104 @@
+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 ---
-# We initialize this globally but handle errors so the app doesn't crash on startup.
-db_client = None
 
-def init_db_client():
-    global db_client
-    try:
-        # Use the DIRECT URL to avoid DNS resolution issues with the Hub API
-        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}")
+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.
 
-# Attempt connection on startup
-init_db_client()
+Wichtig: Gib absolut keinen Text vor oder nach dem JSON aus (keine Erklärungen, kein "Hier ist das Ergebnis").
 
-# --- 2. Tool Definition ---
+{
+  "start_loc": "fill in Startort here",
+  "dest_loc": "fill in Zielort here"
+}
 
-def get_train_connection(dep: str, dest: str):
-    """
-    Fetches the train timetable between two cities using the external API.
-    """
-    global db_client
-    # If client failed to load initially, try one more time
-    if db_client is None:
-        init_db_client()
-        if db_client is None:
-            return "Error: The train database is currently unreachable. Please check your network connection."
+"""
 
-    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 ---
+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):
-    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
-
-    model = "gemini-2.0-flash-latest" # Ensure you use a model version that supports tools
+        return f"Error initializing client: {e}.  Make sure GEMINI_API_KEY is set."
 
-    # Configure tools (Google Search + Our Custom DB Tool)
+    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(),
-            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: Ask Gemini what to do
-        response = client.models.generate_content(
-            model=model,
-            contents=contents,
-            config=generate_content_config,
-        )
-
-        # Check if Gemini wants to call a function
-        if response.candidates and response.candidates[0].content.parts:
-            part = response.candidates[0].content.parts[0]
-            
-            # If it's a function call
-            if part.function_call:
-                fn_name = part.function_call.name
-                fn_args = part.function_call.args
-                
-                if fn_name in tools_map:
-                    # 1. Execute the Python function (calls the external Gradio app)
-                    api_result = tools_map[fn_name](**fn_args)
-                    
-                    # 2. Feed the result back to Gemini
-                    contents.append(response.candidates[0].content) # Add the model's call to history
-                    contents.append(
-                        types.Content(
-                            role="tool",
-                            parts=[
-                                types.Part.from_function_response(
-                                    name=fn_name,
-                                    response={"result": api_result}
-                                )
-                            ]
-                        )
-                    )
-                    
-                    # 3. Get the final natural language answer
-                    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
-
-        # If no function call, just return the text (e.g. normal chat or Google Search)
-        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 = clean_json_string(response_text)
+    data = data[:-1]
+    return response_text, ""
+   
 
 if __name__ == '__main__':
+
     with gr.Blocks() as demo:
-        title = gr.Markdown("# Gemini 2.0 Flash + DB Timetable Tool")
+        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)
+    """"""