Update app.py

app.py

@@ -6,25 +6,19 @@ from google import genai
 from google.genai import types
 from gradio_client import Client
 
-
 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):
     """
     Removes any comments or prefixes before the actual JSON content.
@@ -39,22 +33,95 @@ 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
-    # Verify it's valid JSON
+
+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 to get train connections.
+    """
     try:
-        json.loads(cleaned_str)
-        return cleaned_str
-    except json.JSONDecodeError:
-        return json_str  # Return original if cleaning results in invalid JSON
+        client = Client("mgokg/db-timetable-api")
+        result = client.predict(
+            query=f"Verbindung von {start_loc} nach {dest_loc}",
+            api_name="/db_timetable_api_ui_wrapper"
+        )
+        return result
+    except Exception as e:
+        return f"Error calling MCP server: {e}"
 
 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}", ""
+        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(
@@ -64,41 +131,39 @@ 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}"
-    data = clean_json_string(response_text)
-    data = data[:-1]
+        return f"Error during generation: {e}", ""
+    
     return response_text, ""
-   
 
 if __name__ == '__main__':
-
     with gr.Blocks() as demo:
-        title=gr.Markdown("# Gemini 2.0 Flash + Websearch")
+        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.")
         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])
+        submit_button.click(fn=generate, inputs=input_textbox, outputs=[output_textbox, input_textbox])
     demo.launch(show_error=True)
-