Update app.py

app.py

@@ -5,17 +5,73 @@ import json
 from google import genai
 from google.genai import types
 from gradio_client import Client
+import requests
+from typing import Optional
 
 
-def generate(input_text):
+def get_train_connections(departure: str, destination: str) -> str:
+    """
+    Get train connections using the db_timetable_api MCP tool.
+    
+    Args:
+        departure: Departure station (e.g., "Schweinfurt HBF")
+        destination: Destination station (e.g., "Oerlenbach")
+    
+    Returns:
+        Formatted train connection information
+    """
+    try:
+        # Use the db_timetable_api tool directly
+        result = db_timetable_api_ui_wrapper(dep=departure, dest=destination)
+        return result
+    except Exception as e:
+        return f"Error getting train connections: {str(e)}"
+
+
+def generate_with_tools(input_text: str, use_websearch: bool = True, use_train_api: bool = False) -> str:
+    """
+    Generate response using Gemini with optional web search and train API integration.
+    
+    Args:
+        input_text: User input text
+        use_websearch: Whether to use web search functionality
+        use_train_api: Whether to use train API for connections
+    
+    Returns:
+        Generated response text
+    """
     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."
+        return f"Error initializing client: {e}. Make sure GEMINI_API_KEY is set."
 
     model = "gemini-flash-latest"
+    
+    # Check if user is asking for train connections
+    train_keywords = ["zug", "bahn", "train", "connection", "fahrplan", "timetable", "abfahrt", "ankunft"]
+    is_train_query = any(keyword in input_text.lower() for keyword in train_keywords)
+    
+    # If it's a train query, try to extract stations and use train API
+    if is_train_query and use_train_api:
+        # Simple extraction - in production you might want more sophisticated NLP
+        words = input_text.split()
+        potential_stations = []
+        
+        # Common German station patterns
+        for i, word in enumerate(words):
+            if word.upper() == "HBF" or "bahnhof" in word.lower():
+                # Look for station name before HBF
+                if i > 0:
+                    station = words[i-1] + " " + word
+                    potential_stations.append(station)
+        
+        if len(potential_stations) >= 2:
+            train_result = get_train_connections(potential_stations[0], potential_stations[1])
+            return f"**Train Connections Found:**\n\n{train_result}\n\n---\n\n*Powered by Deutsche Bahn API*"
+    
+    # Original websearch functionality
     contents = [
         types.Content(
             role="user",
@@ -24,42 +80,112 @@ def generate(input_text):
             ],
         ),
     ]
-    tools = [
-        types.Tool(google_search=types.GoogleSearch()),
-    ]
+    
+    tools = []
+    if use_websearch:
+        tools.append(types.Tool(google_search=types.GoogleSearch()))
+    
     generate_content_config = types.GenerateContentConfig(
         temperature=0.4,
-        thinking_config = types.ThinkingConfig(
+        thinking_config=types.ThinkingConfig(
             thinking_budget=0,
         ),
         tools=tools,
         response_mime_type="text/plain",
     )
 
-
     response_text = ""
     try:
-      for chunk in client.models.generate_content_stream(
-          model=model,
-          contents=contents,
-          config=generate_content_config,
-      ):
-          response_text += chunk.text
+        for chunk in client.models.generate_content_stream(
+            model=model,
+            contents=contents,
+            config=generate_content_config,
+        ):
+            if hasattr(chunk, 'text'):
+                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, ""
-   
+    
+    return response_text
 
-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)
+def process_input(input_text: str, mode: str = "Auto-detect") -> str:
+    """
+    Process user input based on selected mode.
+    
+    Args:
+        input_text: User input text
+        mode: Processing mode ("Auto-detect", "Web Search", "Train Connections", "Both")
+    
+    Returns:
+        Processed response
+    """
+    if not input_text.strip():
+        return "Please enter a message."
+    
+    if mode == "Web Search":
+        return generate_with_tools(input_text, use_websearch=True, use_train_api=False)
+    elif mode == "Train Connections":
+        return generate_with_tools(input_text, use_websearch=False, use_train_api=True)
+    elif mode == "Both":
+        return generate_with_tools(input_text, use_websearch=True, use_train_api=True)
+    else:  # Auto-detect
+        return generate_with_tools(input_text, use_websearch=True, use_train_api=True)
+
 
+if __name__ == '__main__':
+    with gr.Blocks(title="Gemini 2.0 Flash + Websearch + Train Connections") as demo:
+        gr.Markdown("# 🤖 Gemini 2.0 Flash + Websearch + Train Connections")
+        gr.Markdown("Ask me anything! I can search the web and find train connections for you.")
+        
+        with gr.Row():
+            with gr.Column(scale=3):
+                input_textbox = gr.Textbox(
+                    lines=3, 
+                    label="Your Message", 
+                    placeholder="Enter your message here...\nExamples:\n- 'What's the weather in Berlin?'\n- 'Train from Schweinfurt HBF to Oerlenbach'\n- 'Latest news about AI'"
+                )
+            with gr.Column(scale=1):
+                mode_dropdown = gr.Dropdown(
+                    choices=["Auto-detect", "Web Search", "Train Connections", "Both"],
+                    value="Auto-detect",
+                    label="Mode",
+                    info="Choose how to process your query"
+                )
+        
+        submit_button = gr.Button("Send", variant="primary")
+        
+        output_textbox = gr.Markdown(
+            label="Response",
+            show_copy_button=True
+        )
+        
+        # Examples
+        gr.Markdown("### 🎯 Quick Examples")
+        with gr.Row():
+            example1 = gr.Button("Weather in Munich")
+            example2 = gr.Button("Train Schweinfurt HBF to Oerlenbach")
+            example3 = gr.Button("Latest AI news")
+        
+        def set_example(example_text):
+            return example_text
+        
+        example1.click(fn=lambda: set_example("What's the weather in Munich today?"), outputs=input_textbox)
+        example2.click(fn=lambda: set_example("Train connections from Schweinfurt HBF to Oerlenbach"), outputs=input_textbox)
+        example3.click(fn=lambda: set_example("Latest news about artificial intelligence developments"), outputs=input_textbox)
+        
+        # Main functionality
+        submit_button.click(
+            fn=process_input,
+            inputs=[input_textbox, mode_dropdown],
+            outputs=output_textbox
+        )
+        
+        # Allow Enter key submission
+        input_textbox.submit(
+            fn=process_input,
+            inputs=[input_textbox, mode_dropdown],
+            outputs=output_textbox
+        )
+    
+    demo.launch(show_error=True)