Update app.py

app.py

@@ -1,10 +1,16 @@
+
+
+
+
+
 import base64
 import gradio as gr
 import os
 import json
-import requests
 from google import genai
 from google.genai import types
+from gradio_client import Client
+
 
 route="""	
 how to handle special case "zugverbindung".
@@ -12,13 +18,18 @@ Wichtig: Dies Regeln gelten nur wenn eine zugverbindung angefragt wird, else ans
 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):
     """
     Removes any comments or prefixes before the actual JSON content.
@@ -33,125 +44,22 @@ def clean_json_string(json_str):
     
     # Extract everything from the first JSON marker
     cleaned_str = json_str[json_start:]
-    
-    # Remove trailing markdown code fence if present
-    if cleaned_str.endswith('```'):
-        cleaned_str = cleaned_str[:-3].rstrip()
-    
     return cleaned_str
-
-def is_train_connection_query(text):
-    """
-    Checks if the query is about train connections.
-    """
-    keywords = ['zugverbindung', 'zug', 'bahn', 'von', 'nach', 'fahrt', 'reise', 'verbindung']
-    text_lower = text.lower()
-    return any(keyword in text_lower for keyword in keywords) and ('von' in text_lower or 'nach' in text_lower)
-
-def get_train_connections(start_loc, dest_loc):
-    """
-    Calls the MCP server via HTTP to get train connections.
-    """
+    # Verify it's valid JSON
     try:
-        mcp_url = "https://mgokg-db-timetable-api.hf.space/gradio_api/mcp/"
-        
-        # MCP protocol request structure
-        payload = {
-            "jsonrpc": "2.0",
-            "id": 1,
-            "method": "tools/call",
-            "params": {
-                "name": "db_timetable_api_ui_wrapper",
-                "arguments": {
-                    "query": f"Verbindung von {start_loc} nach {dest_loc}"
-                }
-            }
-        }
-        
-        headers = {
-            "Content-Type": "application/json"
-        }
-        
-        response = requests.post(mcp_url, json=payload, headers=headers, timeout=30)
-        response.raise_for_status()
-        
-        result = response.json()
-        
-        # Extract the result from MCP response
-        if "result" in result:
-            content = result["result"].get("content", [])
-            if content and len(content) > 0:
-                return content[0].get("text", "Keine Verbindungen gefunden")
-        
-        return f"Unerwartete Antwortstruktur: {result}"
-        
-    except requests.exceptions.Timeout:
-        return "Die Anfrage hat zu lange gedauert. Bitte versuche es erneut."
-    except requests.exceptions.RequestException as e:
-        return f"Fehler beim Abrufen der Zugverbindungen: {e}"
-    except Exception as e:
-        return f"Unerwarteter Fehler: {e}"
+        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):
-    # Check if this is a train connection query
-    if is_train_connection_query(input_text):
-        try:
-            # Use Gemini to extract start and destination
-            client = genai.Client(
-                api_key=os.environ.get("GEMINI_API_KEY"),
-            )
-            
-            model = "gemini-flash-latest"
-            contents = [
-                types.Content(
-                    role="user",
-                    parts=[
-                        types.Part.from_text(text=f"{route}\n\nUser query: {input_text}"),
-                    ],
-                ),
-            ]
-            
-            generate_content_config = types.GenerateContentConfig(
-                temperature=0.1,
-                response_mime_type="application/json",
-            )
-            
-            response = client.models.generate_content(
-                model=model,
-                contents=contents,
-                config=generate_content_config,
-            )
-            
-            # Parse the JSON response
-            json_str = clean_json_string(response.text)
-            location_data = json.loads(json_str)
-            
-            start_loc = location_data.get("start_loc", "")
-            dest_loc = location_data.get("dest_loc", "")
-            
-            if start_loc and dest_loc:
-                # Call MCP server for train connections
-                train_data = get_train_connections(start_loc, dest_loc)
-                
-                # Format the response nicely
-                formatted_response = f"## Zugverbindung von {start_loc} nach {dest_loc}\n\n{train_data}"
-                return formatted_response, ""
-            else:
-                return "Konnte Start- oder Zielort nicht identifizieren. Bitte gib beide Orte an.", ""
-                
-        except json.JSONDecodeError as e:
-            return f"Error parsing location data: {e}\nResponse was: {response.text}", ""
-        except Exception as e:
-            return f"Error processing train connection request: {e}", ""
-    
-    # If not a train query, use regular Gemini with web search
     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"
     contents = [
         types.Content(
@@ -161,39 +69,41 @@ def generate(input_text):
             ],
         ),
     ]
-    
     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,
-        ):
-            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 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 + Websearch + Deutsche Bahn Verbindungen")
-        gr.Markdown("Frage nach Zugverbindungen (z.B. 'Zugverbindung von Berlin nach München') oder stelle allgemeine Fragen.")
+        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)
+        submit_button.click(fn=generate,inputs=input_textbox,outputs=[output_textbox, input_textbox])
+    demo.launch(show_error=True)
+    """"""