Update app.py

app.py

@@ -1,19 +1,41 @@
-import os
 import gradio as gr
+import os
 from google import genai
 from google.genai import types
 from gradio_client import Client
 
-# 1. Initialize the client for the external DB Timetable App
-# We use the Hugging Face Space ID provided in your documentation
-db_client = Client("mgokg/db-timetable-api")
+# --- 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}")
+
+# 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
+    # 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: db_timetable_api_ui_wrapper
+        # Calling the specific endpoint mentioned in the MCP docs
         result = db_client.predict(
             dep=dep,
             dest=dest,
@@ -23,8 +45,7 @@ def get_train_connection(dep: str, dest: str):
     except Exception as e:
         return f"Error fetching timetable: {str(e)}"
 
-# 2. Define the tool for Gemini
-# This tells the model how to use the Python function above
+# 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).",
@@ -38,130 +59,117 @@ train_tool = types.FunctionDeclaration(
     )
 )
 
-# Map the string name to the actual python function
+# Map string name to the actual function
 tools_map = {
     "get_train_connection": get_train_connection
 }
 
-def generate(input_text, history):
-    # Initialize Gemini Client
+# --- 3. Generation Logic ---
+
+def generate(input_text):
+    if not input_text:
+        yield "", ""
+        return
+
     try:
         client = genai.Client(
             api_key=os.environ.get("GEMINI_API_KEY"),
         )
     except Exception as e:
-        yield f"Error initializing client: {e}. Make sure GEMINI_API_KEY is set.", history
+        yield f"Error initializing client: {e}. Make sure GEMINI_API_KEY is set.", input_text
         return
 
-    model = "gemini-2.0-flash-exp" # Or "gemini-2.0-flash" depending on availability
+    model = "gemini-2.0-flash-exp" # Ensure you use a model version that supports tools
 
-    # Prepare the conversation history for context
-    # (Optional: You can add previous history here if you want multi-turn chat)
-    contents = [
-        types.Content(
-            role="user",
-            parts=[types.Part.from_text(text=input_text)],
-        ),
-    ]
-
-    # 3. Configure tools (Google Search + Our Custom DB Tool)
+    # Configure tools (Google Search + Our Custom DB Tool)
     tools = [
-        types.Tool(google_search=types.GoogleSearch()),
-        types.Tool(function_declarations=[train_tool]),
+        types.Tool(
+            google_search=types.GoogleSearch(),
+            function_declarations=[train_tool]
+        )
     ]
 
     generate_content_config = types.GenerateContentConfig(
         temperature=0.4,
         tools=tools,
-        # Automatic function calling allows the SDK to handle the loop, 
-        # but for granular control in Gradio, we often handle it manually below 
-        # or rely on the model to return a function call part.
     )
 
+    contents = [
+        types.Content(
+            role="user",
+            parts=[types.Part.from_text(text=input_text)],
+        ),
+    ]
+
     response_text = ""
     
-    # First API Call: Ask the model what to do
     try:
+        # First API Call: Ask Gemini what to do
         response = client.models.generate_content(
             model=model,
             contents=contents,
             config=generate_content_config,
         )
-    except Exception as e:
-        yield f"Error during generation: {e}", history
-        return
 
-    # 4. Check if the model wants to call a function
-    # We look at the first candidate's first part
-    if response.candidates and response.candidates[0].content.parts:
-        first_part = response.candidates[0].content.parts[0]
-        
-        # If it's a function call
-        if first_part.function_call:
-            fn_name = first_part.function_call.name
-            fn_args = first_part.function_call.args
+        # Check if Gemini wants to call a function
+        if response.candidates and response.candidates[0].content.parts:
+            part = response.candidates[0].content.parts[0]
             
-            # Execute the tool
-            if fn_name in tools_map:
-                status_msg = f"🔄 Checking trains from {fn_args.get('dep')} to {fn_args.get('dest')}..."
-                yield status_msg, history
-                
-                api_result = tools_map[fn_name](**fn_args)
+            # If it's a function call
+            if part.function_call:
+                fn_name = part.function_call.name
+                fn_args = part.function_call.args
                 
-                # Send the result back to Gemini
-                # We append the model's function call and our function response to 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}
-                            )
-                        ]
+                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}
+                                )
+                            ]
+                        )
                     )
-                )
-                
-                # Second API Call: Get the final natural language answer
-                stream = client.models.generate_content_stream(
-                    model=model,
-                    contents=contents,
-                    config=generate_content_config # Keep tools enabled just in case
-                )
-                
-                final_text = ""
-                for chunk in stream:
-                    if chunk.text:
-                        final_text += chunk.text
-                        yield final_text, history
-                return
+                    
+                    # 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, ""
+            
+    except Exception as e:
+        yield f"Error during generation: {e}", input_text
 
-    # If no function call, just return the text (e.g., normal chat or Google Search result)
-    if response.text:
-        yield response.text, history
+# --- 4. UI Setup ---
 
 if __name__ == '__main__':
     with gr.Blocks() as demo:
-        gr.Markdown("# Gemini 2.0 Flash + DB Timetable Tool")
+        title = gr.Markdown("# Gemini 2.0 Flash + DB Timetable Tool")
+        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")
         
-        chatbot = gr.Chatbot(label="Conversation", height=400)
-        msg = gr.Textbox(lines=1, label="Ask about trains (e.g., 'Train from Berlin to Munich')", placeholder="Enter message here...")
-        clear = gr.Button("Clear")
-
-        def user(user_message, history):
-            return "", history + [[user_message, None]]
-
-        def bot(history):
-            user_message = history[-1][0]
-            # Call generate and update the last message in history
-            for partial_response, _ in generate(user_message, history):
-                history[-1][1] = partial_response
-                yield history
-
-        msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
-            bot, chatbot, chatbot
+        submit_button.click(
+            fn=generate,
+            inputs=input_textbox,
+            outputs=[output_textbox, input_textbox]
         )
-        clear.click(lambda: None, None, chatbot, queue=False)
-
+        
     demo.launch(show_error=True)