Update app.py

app.py

@@ -1,187 +1,3 @@
-import os
-import asyncio
-import json
-import gradio as gr
-
-from google import genai
-from google.genai import types
-
-from mcp import ClientSession
-from mcp.client.sse import sse_client
-
-# --- CONFIGURATION ---
-MCP_SERVER_URL = "https://mgokg-db-timetable-api.hf.space/gradio_api/mcp/"
-GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")
-
-
-async def call_mcp_tool(start_loc: str, dest_loc: str) -> str:
-    """Connect to the MCP server to fetch train data."""
-    async with sse_client(MCP_SERVER_URL) as (read_stream, write_stream):
-        async with ClientSession(read_stream, write_stream) as session:
-            await session.initialize()
-            result = await session.call_tool(
-                "get_train_connections",
-                arguments={"start_loc": start_loc, "dest_loc": dest_loc}
-            )
-
-            # Defensive handling of potential result formats
-            try:
-                if hasattr(result, "content") and result.content:
-                    first = result.content[0]
-                    if hasattr(first, "text") and first.text:
-                        return first.text
-                    # Fallback: serialize unknown structures
-                    return json.dumps(
-                        [getattr(c, "__dict__", str(c)) for c in result.content],
-                        ensure_ascii=False
-                    )
-            except Exception:
-                pass
-
-            return "No train data returned from MCP."
-
-
-async def generate(input_text: str):
-    """
-    Async generator handling the tool loop with the corrected aio client.
-    Yields (markdown_message, input_box_echo) tuples for Gradio.
-    """
-    if not GEMINI_API_KEY:
-        yield "### Error\n`GEMINI_API_KEY` is not set.", ""
-        return
-
-    client = genai.Client(api_key=GEMINI_API_KEY)
-    model_id = "gemini-2.5-flash"  # Ensure this is enabled for your API key
-
-    # Declare the function/tool to the model
-    train_tool = types.Tool(
-        function_declarations=[
-            types.FunctionDeclaration(
-                name="get_train_connections",
-                description="Finds live train connections between two German stations.",
-                parameters={
-                    "type": "OBJECT",
-                    "properties": {
-                        "start_loc": {"type": "STRING", "description": "Departure city"},
-                        "dest_loc": {"type": "STRING", "description": "Destination city"},
-                    },
-                    "required": ["start_loc", "dest_loc"],
-                },
-            )
-        ]
-    )
-
-    try:
-        yield "🔍 Thinking...", ""
-
-        config = types.GenerateContentConfig(
-            tools=[train_tool, types.Tool(google_search=types.GoogleSearch())],
-            temperature=0.3,
-        )
-
-        # Create async chat session (DO NOT await here)
-        chat = client.aio.chats.create(model=model_id, config=config)
-
-        # Send the user input to the model (await is correct here)
-        response =  chat.send_message(input_text)
-
-        # --- Manual Tool Loop ---
-        max_turns = 5
-        for _ in range(max_turns):
-            # Guards
-            if not getattr(response, "candidates", None):
-                break
-
-            candidate = response.candidates[0]
-            parts = getattr(candidate.content, "parts", []) or []
-            if not parts:
-                break
-
-            # Collect function calls from parts (correct attribute is function_call)
-            tool_calls = []
-            for p in parts:
-                call = getattr(p, "function_call", None)
-                if call and hasattr(call, "name") and hasattr(call, "args"):
-                    tool_calls.append(call)
-
-            # If no tools were requested, the model is done
-            if not tool_calls:
-                break
-
-            tool_responses = []
-
-            for call in tool_calls:
-                if call.name == "get_train_connections":
-                    start_loc = call.args.get("start_loc", "").strip()
-                    dest_loc = call.args.get("dest_loc", "").strip()
-                    yield f"🚄 Fetching train data: **{start_loc}** → **{dest_loc}**...", ""
-
-                    # Call the MCP tool
-                    train_data = await call_mcp_tool(start_loc, dest_loc)
-
-                    # Return the tool function result to the model
-                    tool_responses.append(
-                        types.Part.from_function_response(
-                            name=call.name,
-                            response={"result": train_data},
-                        )
-                    )
-                else:
-                    # Unknown tool – inform the model
-                    tool_responses.append(
-                        types.Part.from_function_response(
-                            name=call.name,
-                            response={
-                                "error": f"Unknown tool '{call.name}'. Only 'get_train_connections' is supported."
-                            },
-                        )
-                    )
-
-            if tool_responses:
-                yield "📝 Finalizing response...", ""
-                # Provide tool results back to the model so it can compose the final answer
-                response =  chat.send_message(tool_responses)
-            else:
-                break
-
-        # Prefer response.text; fallback to concatenating text parts
-        final_text = getattr(response, "text", "") or ""
-        if not final_text and getattr(response, "candidates", None):
-            parts = getattr(response.candidates[0].content, "parts", [])
-            texts = []
-            for p in parts:
-                if hasattr(p, "text") and p.text:
-                    texts.append(p.text)
-            final_text = "\n".join(texts)
-
-        if not final_text:
-            final_text = "I couldn't generate a response. Please try rephrasing your question."
-
-        yield final_text, ""
-
-    except Exception as e:
-        # Show the actual logic/exception error back in the UI
-        yield f"### Logic Error\n{str(e)}", ""
-
-
-# --- GRADIO UI ---
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 🚄 Gemini MCP Train Assistant")
-
-    output_textbox = gr.Markdown(label="Response")
-    input_textbox = gr.Textbox(label="Ask about German train connections", placeholder="Berlin to Munich?")
-    submit_button = gr.Button("Send", variant="primary")
-
-    submit_button.click(
-        fn=generate,
-        inputs=input_textbox,
-        outputs=[output_textbox, input_textbox]
-    )
-
-if __name__ == "__main__":
-    demo.launch()
-
-"""
 import base64
 import gradio as gr
 import os