Update app.py

app.py

@@ -1,109 +1,167 @@
-
-
-
-
-
-import base64
-import gradio as gr
 import os
-import json
+import gradio as gr
 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")
 
-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"
-}
-
-"""
-
-
-
-def clean_json_string(json_str):
+def get_train_connection(dep: str, dest: str):
     """
-    Removes any comments or prefixes before the actual JSON content.
+    Fetches the train timetable between two cities using the external API.
     """
-    # 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
+        # Calling the specific endpoint mentioned in the MCP docs: db_timetable_api_ui_wrapper
+        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)}"
+
+# 2. Define the tool for Gemini
+# This tells the model how to use the Python function above
+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 the string name to the actual python function
+tools_map = {
+    "get_train_connection": get_train_connection
+}
 
-def generate(input_text):
+def generate(input_text, history):
+    # Initialize Gemini Client
     try:
         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."
+        yield f"Error initializing client: {e}. Make sure GEMINI_API_KEY is set.", history
+        return
+
+    model = "gemini-2.0-flash-exp" # Or "gemini-2.0-flash" depending on availability
 
-    model = "gemini-flash-latest"
+    # 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=f"{input_text}"),
-            ],
+            parts=[types.Part.from_text(text=input_text)],
         ),
     ]
+
+    # 3. Configure tools (Google Search + Our Custom DB Tool)
     tools = [
         types.Tool(google_search=types.GoogleSearch()),
+        types.Tool(function_declarations=[train_tool]),
     ]
+
     generate_content_config = types.GenerateContentConfig(
         temperature=0.4,
-        thinking_config = types.ThinkingConfig(
-            thinking_budget=0,
-        ),
         tools=tools,
-        response_mime_type="text/plain",
+        # 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.
     )
 
-
     response_text = ""
+    
+    # First API Call: Ask the model what to do
     try:
-      for chunk in client.models.generate_content_stream(
-          model=model,
-          contents=contents,
-          config=generate_content_config,
-      ):
-          response_text += chunk.text
+        response = client.models.generate_content(
+            model=model,
+            contents=contents,
+            config=generate_content_config,
+        )
     except Exception as e:
-        return f"Error during generation: {e}"
-    data = clean_json_string(response_text)
-    data = data[:-1]
-    return response_text, ""
-   
+        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
+            
+            # 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)
+                
+                # 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}
+                            )
+                        ]
+                    )
+                )
+                
+                # 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
+
+    # If no function call, just return the text (e.g., normal chat or Google Search result)
+    if response.text:
+        yield response.text, history
 
 if __name__ == '__main__':
-
     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)
-    """"""
+        gr.Markdown("# Gemini 2.0 Flash + DB Timetable Tool")
+        
+        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
+        )
+        clear.click(lambda: None, None, chatbot, queue=False)
+
+    demo.launch(show_error=True)