Update app.py

app.py

@@ -1,84 +1,112 @@
-import base64
-import gradio as gr
 import os
 import json
+import asyncio
+import gradio as gr
 from google import genai
 from google.genai import types
-from gradio_client import Client
+from mcp import ClientSession
+from mcp.client.sse import sse_client
 
-def clean_json_string(json_str):
-    """
-    Removes any comments or prefixes before the actual JSON content.
-    """
-    # Find the first occurrence of '{'
-    json_start = json_str.find('{')
-    if json_start == -1:
-        # If no '{' is found, try with '[' for arrays
-        json_start = json_str.find('[')
-        if json_start == -1:
-            return json_str  # Return original if no JSON markers found
-    
-    # Extract everything from the first JSON marker
-    cleaned_str = json_str[json_start:]
-    return cleaned_str
-    # Verify it's valid JSON
-    try:
-        json.loads(cleaned_str)
-        return cleaned_str
-    except json.JSONDecodeError:
-        return json_str  # Return original if cleaning results in invalid JSON
+# --- CONFIGURATION ---
+MCP_SERVER_URL = "https://mgokg-db-api-mcp.hf.space/sse"
+GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")
 
-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."
+# 1. Function to connect to your MCP Server and run the tool
+async def call_mcp_tool(start_loc, dest_loc):
+    async with sse_client(MCP_SERVER_URL) as (read_stream, write_stream):
+        async with ClientSession(read_stream, write_stream) as session:
+            await session.initialize()
+            # Calling the tool defined in your MCP server
+            result = await session.call_tool("get_train_connections", arguments={
+                "start_loc": start_loc,
+                "dest_loc": dest_loc
+            })
+            # result.content[0].text contains the JSON train data
+            return result.content[0].text
 
-    model = "gemini-flash-latest"
-    contents = [
-        types.Content(
-            role="user",
-            parts=[
-                types.Part.from_text(text=input_text),
-            ],
-        ),
-    ]
-    tools = [
-        types.Tool(google_search=types.GoogleSearch()),
+# 2. Define the tool for Gemini
+# This describes the tool so Gemini knows when to use it
+train_tool = types.Tool(
+    function_declarations=[
+        types.FunctionDeclaration(
+            name="get_train_connections",
+            description="Finds live train connections between two German stations/cities.",
+            parameters=types.Schema(
+                type="OBJECT",
+                properties={
+                    "start_loc": types.Schema(type="STRING", description="Departure city/station"),
+                    "dest_loc": types.Schema(type="STRING", description="Destination city/station")
+                },
+                required=["start_loc", "dest_loc"]
+            )
+        )
     ]
-    generate_content_config = types.GenerateContentConfig(
-        temperature=0.4,
-        thinking_config = types.ThinkingConfig(
-            thinking_budget=0,
-        ),
-        tools=tools,
-        response_mime_type="text/plain",
-    )
+)
+
+def generate(input_text):
+    if not GEMINI_API_KEY:
+        return "Error: GEMINI_API_KEY is not set.", ""
+
+    client = genai.Client(api_key=GEMINI_API_KEY)
+    model_id = "gemini-2.0-flash-exp" # Recommended for tool use
 
+    # Create a chat session to handle the tool-use loop
+    chat = client.chats.create(
+        model=model_id,
+        config=types.GenerateContentConfig(
+            tools=[train_tool, types.Tool(google_search=types.GoogleSearch())],
+            temperature=0.3
+        )
+    )
 
-    response_text = ""
     try:
-      for chunk in client.models.generate_content_stream(
-          model=model,
-          contents=contents,
-          config=generate_content_config,
-      ):
-          response_text += chunk.text
+        # Step 1: Send user input to Gemini
+        response = chat.send_message(input_text)
+        
+        # Step 2: Tool Loop
+        # If Gemini emits a tool call, we execute it and send results back
+        while response.candidates[0].content.parts[0].tool_call:
+            tool_call = response.candidates[0].content.parts[0].tool_call
+            
+            if tool_call.name == "get_train_connections":
+                args = tool_call.args
+                # Execute the actual MCP request
+                train_data = asyncio.run(call_mcp_tool(args["start_loc"], args["dest_loc"]))
+                
+                # Feed the train data back to Gemini
+                response = chat.send_message(
+                    types.Part.from_function_response(
+                        name="get_train_connections",
+                        response={"result": train_data}
+                    )
+                )
+            else:
+                break # Exit if it's a tool we don't recognize here
+
+        return response.text, ""
+
     except Exception as e:
-        return f"Error during generation: {e}"
-    data = clean_json_string(response_text)
-    data = data[:-1]
-    return response_text, ""
-   
+        return f"Error: {str(e)}", ""
 
-if __name__ == '__main__':
+# --- GRADIO UI ---
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 🚄 Gemini 2.0 Flash + Live DB Trains")
+    gr.Markdown("I can search the web and check live train connections via an MCP Server.")
+    
+    with gr.Column():
+        output_textbox = gr.Markdown(label="Response")
+        input_textbox = gr.Textbox(
+            lines=3, 
+            label="Ask me anything", 
+            placeholder="e.g. 'Is there a train from Berlin to Munich tonight?'"
+        )
+        submit_button = gr.Button("Send", variant="primary")
+    
+    submit_button.click(
+        fn=generate,
+        inputs=input_textbox,
+        outputs=[output_textbox, input_textbox]
+    )
 
-    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__':
+    demo.launch()