Update app.py

app.py

@@ -2,64 +2,311 @@ 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 gradio_client import Client
+from contextlib import AsyncExitStack
+from mcp import ClientSession
+from mcp.client.streamable_http import streamablehttp_client
+import nest_asyncio
 
+# Allow nested event loops for Jupyter notebooks
+nest_asyncio.apply()
 
-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."
+class MCPTrainClient:
+    """MCP Client for train timetable API using Streamable HTTP transport"""
+    
+    def __init__(self):
+        self.session = None
+        self.exit_stack = None
+        self.tools = []
+        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
+            transport = await self.exit_stack.enter_async_context(
+                streamablehttp_client(self.mcp_server_url)
+            )
+            
+            self.read_stream, self.write_stream = transport
+            self.session = await self.exit_stack.enter_async_context(
+                ClientSession(self.read_stream, self.write_stream)
+            )
+            
+            await self.session.initialize()
+            
+            # List available tools
+            response = await self.session.list_tools()
+            self.tools = response.tools
+            
+            tool_names = [tool.name for tool in self.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
+            self.tools = []
 
-    model = "gemini-flash-latest"
-    contents = [
-        types.Content(
-            role="user",
-            parts=[
-                types.Part.from_text(text=f"{input_text}"),
-            ],
-        ),
-    ]
-    tools = [
-        types.Tool(google_search=types.GoogleSearch()),
-    ]
-    generate_content_config = types.GenerateContentConfig(
-        temperature=0.4,
-        thinking_config = types.ThinkingConfig(
-            thinking_budget=0,
-        ),
-        tools=tools,
-        response_mime_type="text/plain",
-    )
 
+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)}"
 
-    response_text = ""
-    try:
-      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 = response_text  
-    #data = clean_json_string(response_text)
-    data = data[:-1]
-    return response_text, ""
-   
 
-if __name__ == '__main__':
+# Global instance
+mcp_integration = GeminiMCPIntegration()
+
+
+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)
+
+
+def connect_to_train_api() -> str:
+    """Connect to train API MCP server"""
+    return mcp_integration.connect_to_mcp()
 
-    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 (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
+        )
+        
+        status_textbox = gr.Textbox(
+            label="Connection Status",
+            interactive=False,
+            max_lines=2
+        )
+        
+        # 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)
+        
+        # Main functionality
+        submit_button.click(
+            fn=process_user_input,
+            inputs=[input_textbox, mode_dropdown],
+            outputs=output_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)