Update app.py

app.py

@@ -1,305 +1,176 @@
-import base64
 import gradio as gr
 import os
-import json
-import asyncio
-import aiohttp
-from typing import Dict, Any, Optional, List
 from google import genai
 from google.genai import types
-from mcp import ClientSession, StdioServerParameters
-from mcp.client.stdio import stdio_client
-from mcp.client.streamable_http import streamablehttp_client
-import nest_asyncio
+from gradio_client import Client
 
-# Allow nested event loops for Jupyter notebooks
-nest_asyncio.apply()
+# --- 1. Robust Client Initialization for the External Tool ---
+db_client = None
 
-class MCPTrainClient:
-    """MCP Client for train timetable API using Streamable HTTP transport"""
-    
-    def __init__(self):
-        self.session = None
-        self.exit_stack = None
-        self.mcp_server_url = "https://mgokg-db-timetable-api.hf.space/gradio_api/mcp/"
-        
-    async def connect(self) -> str:
-        """Connect to the MCP server using Streamable HTTP transport"""
-        try:
-            if self.exit_stack:
-                await self.exit_stack.aclose()
-            
-            self.exit_stack = AsyncExitStack()
-            
-            # Connect using Streamable HTTP transport - FIXED VERSION
-            self.session = await self.exit_stack.enter_async_context(
-                streamablehttp_client(self.mcp_server_url)
-            )
-            
-            # Initialize the session
-            await self.session.initialize()
-            
-            # List available tools
-            response = await self.session.list_tools()
-            tool_names = [tool.name for tool in response.tools]
-            
-            return f"✅ Connected to MCP server. Available tools: {', '.join(tool_names)}"
-            
-        except Exception as e:
-            return f"❌ Failed to connect to MCP server: {str(e)}"
-    
-    async def get_train_connections(self, departure: str, destination: str) -> str:
-        """Get train connections using the MCP tool"""
-        if not self.session:
-            return "❌ Not connected to MCP server. Please connect first."
-        
-        try:
-            # Call the db_timetable_api_ui_wrapper tool
-            result = await self.session.call_tool(
-                "db_timetable_api_ui_wrapper",
-                arguments={"dep": departure, "dest": destination}
-            )
-            
-            if result.content and len(result.content) > 0:
-                return result.content[0].text
-            else:
-                return "❌ No results returned from train API"
-                
-        except Exception as e:
-            return f"❌ Error calling train API: {str(e)}"
-    
-    async def disconnect(self):
-        """Disconnect from the MCP server"""
-        if self.exit_stack:
-            await self.exit_stack.aclose()
-            self.exit_stack = None
-            self.session = None
+def init_db_client():
+    global db_client
+    try:
+        # Use the full URL to avoid DNS resolution issues
+        print("Connecting to DB Timetable API...")
+        db_client = Client("https://mgokg-db-timetable-api.hf.space/")
+        print("Successfully connected.")
+    except Exception as e:
+        print(f"Warning: Could not connect to DB Timetable API: {e}")
 
+# Attempt connection on startup
+init_db_client()
 
-class GeminiMCPIntegration:
-    """Integration class that combines Gemini with MCP tools"""
-    
-    def __init__(self):
-        self.mcp_client = MCPTrainClient()
-        self.gemini_client = None
-        self.loop = asyncio.new_event_loop()
-        asyncio.set_event_loop(self.loop)
-        
-    def initialize_gemini(self) -> Optional[str]:
-        """Initialize Gemini client"""
-        try:
-            api_key = os.environ.get("GEMINI_API_KEY")
-            if not api_key:
-                return "GEMINI_API_KEY not found in environment variables"
-            
-            self.gemini_client = genai.Client(api_key=api_key)
-            return None
-        except Exception as e:
-            return f"Error initializing Gemini client: {str(e)}"
-    
-    def connect_to_mcp(self) -> str:
-        """Connect to MCP server (synchronous wrapper)"""
-        return self.loop.run_until_complete(self.mcp_client.connect())
-    
-    def extract_train_stations(self, text: str) -> Optional[Dict[str, str]]:
-        """Extract departure and destination stations from text"""
-        text_lower = text.lower()
-        
-        # Common patterns for train queries
-        patterns = [
-            r'(?:from|von)\s+(.+?)\s+(?:to|nach)\s+(.+)',
-            r'(.+?)\s+(?:to|nach)\s+(.+)',
-            r'(?:train|zug|bahn)(?:\s+from|\s+von)?\s+(.+?)\s+(?:to|nach)\s+(.+)',
-        ]
-        
-        import re
-        for pattern in patterns:
-            match = re.search(pattern, text_lower)
-            if match:
-                dep = match.group(1).strip()
-                dest = match.group(2).strip()
-                
-                # Clean up common suffixes
-                for suffix in ['?', '.', '!', 'please', 'bitte']:
-                    if dep.endswith(suffix):
-                        dep = dep[:-len(suffix)].strip()
-                    if dest.endswith(suffix):
-                        dest = dest[:-len(suffix)].strip()
-                
-                return {"departure": dep, "destination": dest}
-        
-        return None
-    
-    def process_query(self, query: str, mode: str = "auto") -> str:
-        """Process user query with appropriate tools"""
-        
-        # Initialize Gemini if not already done
-        if not self.gemini_client:
-            error = self.initialize_gemini()
-            if error:
-                return f"❌ {error}"
-        
-        # Check if this is a train query
-        station_info = self.extract_train_stations(query)
-        
-        if station_info and (mode == "auto" or mode == "train"):
-            # Use MCP for train connections
-            if not self.mcp_client.session:
-                connect_result = self.connect_to_mcp()
-                if "❌" in connect_result:
-                    return connect_result
-            
-            train_result = self.loop.run_until_complete(
-                self.mcp_client.get_train_connections(
-                    station_info["departure"], 
-                    station_info["destination"]
-                )
-            )
-            
-            # Also get Gemini's interpretation
-            gemini_prompt = f"User asked: '{query}'. I found these train connections: {train_result}. Please provide a helpful summary and any additional context."
-            
-            try:
-                contents = [types.Content(role="user", parts=[types.Part.from_text(text=gemini_prompt)])]
-                config = types.GenerateContentConfig(
-                    temperature=0.4,
-                    response_mime_type="text/plain"
-                )
-                
-                gemini_response = ""
-                for chunk in self.gemini_client.models.generate_content_stream(
-                    model="gemini-flash-latest",
-                    contents=contents,
-                    config=config
-                ):
-                    if hasattr(chunk, 'text'):
-                        gemini_response += chunk.text
-                
-                return f"🚂 **Train Connections:**\n\n{train_result}\n\n🤖 **AI Assistant:**\n\n{gemini_response}"
-                
-            except Exception as e:
-                return f"🚂 **Train Connections:**\n\n{train_result}\n\n⚠️ Note: Could not generate additional AI insights: {str(e)}"
-        
-        elif mode == "train" and not station_info:
-            return "❌ Could not extract station names from your query. Please use format like 'Train from [station] to [station]'"
-        
-        else:
-            # Use original Gemini + websearch functionality
-            return self.generate_with_websearch(query)
-    
-    def generate_with_websearch(self, query: str) -> str:
-        """Original Gemini + websearch functionality"""
-        try:
-            contents = [types.Content(role="user", parts=[types.Part.from_text(text=query)])]
-            tools = [types.Tool(google_search=types.GoogleSearch())]
-            
-            config = types.GenerateContentConfig(
-                temperature=0.4,
-                thinking_config=types.ThinkingConfig(thinking_budget=0),
-                tools=tools,
-                response_mime_type="text/plain"
-            )
-            
-            response_text = ""
-            for chunk in self.gemini_client.models.generate_content_stream(
-                model="gemini-flash-latest",
-                contents=contents,
-                config=config
-            ):
-                if hasattr(chunk, 'text'):
-                    response_text += chunk.text
-            
-            return f"🔍 **Web Search Results:**\n\n{response_text}"
-            
-        except Exception as e:
-            return f"❌ Error during web search: {str(e)}"
+# --- 2. Tool Definition ---
 
+def get_train_connection(dep: str, dest: str):
+    """
+    Fetches the train timetable between two cities using the external API.
+    """
+    global db_client
+    # Retry connection if it failed initially
+    if db_client is None:
+        init_db_client()
+        if db_client is None:
+            return "Error: The train database is currently unreachable."
 
-# Global instance
-mcp_integration = GeminiMCPIntegration()
+    try:
+        # Calling the specific endpoint mentioned in the MCP docs
+        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)}"
 
+# Define the tool schema for Gemini
+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"]
+    )
+)
 
-def process_user_input(input_text: str, mode: str = "auto") -> str:
-    """Process user input"""
-    if not input_text.strip():
-        return "Please enter a message."
-    
-    return mcp_integration.process_query(input_text, mode)
+# Map string name to the actual function
+tools_map = {
+    "get_train_connection": get_train_connection
+}
 
+# --- 3. Generation Logic (v1alpha) ---
 
-def connect_to_train_api() -> str:
-    """Connect to train API MCP server"""
-    return mcp_integration.connect_to_mcp()
+def generate(input_text):
+    if not input_text:
+        yield "", ""
+        return
 
+    try:
+        # Initialize Client with v1alpha
+        client = genai.Client(
+            api_key=os.environ.get("GEMINI_API_KEY"),
+            #http_options={'api_version': 'v1alpha'}  # <--- ENABLE v1alpha HERE
+        )
+    except Exception as e:
+        yield f"Error initializing client: {e}. Make sure GEMINI_API_KEY is set.", input_text
+        return
 
-if __name__ == '__main__':
-    with gr.Blocks(title="Gemini 2.0 Flash + Websearch + Train Connections (MCP)") as demo:
-        gr.Markdown("# 🤖 Gemini 2.0 Flash + Websearch + Train Connections (MCP)")
-        gr.Markdown("Advanced AI assistant with web search and real-time train connections via MCP.")
-        
-        with gr.Row():
-            with gr.Column(scale=3):
-                input_textbox = gr.Textbox(
-                    lines=3, 
-                    label="Your Message", 
-                    placeholder="Enter your message here...\nExamples:\n- 'Train from Schweinfurt HBF to Oerlenbach'\n- 'What's the weather in Berlin?'\n- 'Latest news about AI'"
-                )
-            with gr.Column(scale=1):
-                mode_dropdown = gr.Dropdown(
-                    choices=["auto", "websearch", "train", "both"],
-                    value="auto",
-                    label="Mode",
-                    info="Choose processing mode"
-                )
-        
-        with gr.Row():
-            submit_button = gr.Button("Send", variant="primary", scale=3)
-            connect_button = gr.Button("Connect to Train API", scale=1)
-        
-        output_textbox = gr.Markdown(
-            label="Response",
-            show_copy_button=True
+    model = "gemini-flash-latest" 
+
+    # Configure tools
+    tools = [
+        types.Tool(
+            google_search=types.GoogleSearch(),
+            function_declarations=[train_tool]
         )
-        
-        status_textbox = gr.Textbox(
-            label="Connection Status",
-            interactive=False,
-            max_lines=2
+    ]
+
+    generate_content_config = types.GenerateContentConfig(
+        temperature=0.4,
+        tools=tools,
+    )
+
+    contents = [
+        types.Content(
+            role="user",
+            parts=[types.Part.from_text(text=input_text)],
+        ),
+    ]
+
+    response_text = ""
+    
+    try:
+        # First API Call
+        response = client.models.generate_content(
+            model=model,
+            contents=contents,
+            config=generate_content_config,
         )
+
+        # Handle Function Calling
+        if response.candidates and response.candidates[0].content.parts:
+            part = response.candidates[0].content.parts[0]
+            
+            if part.function_call:
+                fn_name = part.function_call.name
+                fn_args = part.function_call.args
+                
+                if fn_name in tools_map:
+                    # Execute tool
+                    api_result = tools_map[fn_name](**fn_args)
+                    
+                    # Append 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 (Streamed)
+                    stream = client.models.generate_content_stream(
+                        model=model,
+                        contents=contents,
+                        config=generate_content_config
+                    )
+                    
+                    for chunk in stream:
+                        response_text += chunk.text
+                        yield response_text, ""
+                    return
+
+        # Handle standard response
+        if response.text:
+            yield response.text, ""
+            
+    except Exception as e:
+        yield f"Error during generation: {e}", input_text
+
+# --- 4. UI Setup ---
+
+if __name__ == '__main__':
+    with gr.Blocks() as demo:
+        gr.Markdown("# Gemini 2.0 Flash (v1alpha) + DB Timetable Tool")
         
-        # Examples
-        gr.Markdown("### 🎯 Quick Examples")
-        with gr.Row():
-            example1 = gr.Button("🚂 Train: Schweinfurt → Oerlenbach")
-            example2 = gr.Button("🌤️ Weather in Munich")
-            example3 = gr.Button("📰 Latest AI news")
-            example4 = gr.Button("🏙️ Restaurants in Berlin")
-        
-        def set_example(example_text):
-            return example_text
-        
-        example1.click(fn=lambda: set_example("Train connections from Schweinfurt HBF to Oerlenbach"), outputs=input_textbox)
-        example2.click(fn=lambda: set_example("What's the weather in Munich today?"), outputs=input_textbox)
-        example3.click(fn=lambda: set_example("Latest news about artificial intelligence developments"), outputs=input_textbox)
-        example4.click(fn=lambda: set_example("Find good restaurants in Berlin city center"), outputs=input_textbox)
+        output_textbox = gr.Markdown()
+        input_textbox = gr.Textbox(lines=3, label="", placeholder="Ask for a train connection (e.g., 'Train from Berlin to Frankfurt')...")
+        submit_button = gr.Button("Send")
         
-        # Main functionality
         submit_button.click(
-            fn=process_user_input,
-            inputs=[input_textbox, mode_dropdown],
-            outputs=output_textbox
+            fn=generate,
+            inputs=input_textbox,
+            outputs=[output_textbox, input_textbox]
         )
         
-        connect_button.click(
-            fn=connect_to_train_api,
-            outputs=status_textbox
-        )
-        
-        # Allow Enter key submission
-        input_textbox.submit(
-            fn=process_user_input,
-            inputs=[input_textbox, mode_dropdown],
-            outputs=output_textbox
-        )
-    
     demo.launch(show_error=True)