Update app.py

app.py

@@ -5,73 +5,17 @@ import json
 from google import genai
 from google.genai import types
 from gradio_client import Client
-import requests
-from typing import Optional
 
 
-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
-    """
+def generate(input_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",
@@ -80,112 +24,42 @@ def generate_with_tools(input_text: str, use_websearch: bool = True, use_train_a
             ],
         ),
     ]
-    
-    tools = []
-    if use_websearch:
-        tools.append(types.Tool(google_search=types.GoogleSearch()))
-    
+    tools = [
+        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,
-        ):
-            if hasattr(chunk, 'text'):
-                response_text += chunk.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:
         return f"Error during generation: {e}"
-    
-    return response_text
+    data = response_text  
+    #data = clean_json_string(response_text)
+    data = data[:-1]
+    return response_text, ""
+   
 
+if __name__ == '__main__':
 
-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)
-
+    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)
 
-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)