Update app.py

app.py

@@ -1,108 +1,84 @@
+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 mcp import ClientSession
-from mcp.client.sse import sse_client
-
-# --- CONFIGURATION ---
-MCP_SERVER_URL = "https://mgokg/DB_API_MCP.hf.space/sse"
-GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")
+from gradio_client import Client
 
-# 1. Helper to talk to your MCP Server
-async def call_mcp_train_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()
-            result = await session.call_tool("get_train_connections", arguments={
-                "start_loc": start_loc,
-                "dest_loc": dest_loc
-            })
-            # result.content[0].text contains the JSON list from your MCP server
-            return result.content[0].text
-
-# 2. Tool Definition for Gemini
-# This tells Gemini what the tool does and what arguments it needs
-train_tool_declaration = types.Tool(
-    function_declarations=[
-        types.FunctionDeclaration(
-            name="get_train_connections",
-            description="Get live train connections between two German cities/stations.",
-            parameters=types.Schema(
-                type="OBJECT",
-                properties={
-                    "start_loc": types.Schema(type="STRING", description="The departure station"),
-                    "dest_loc": types.Schema(type="STRING", description="The destination station")
-                },
-                required=["start_loc", "dest_loc"]
-            )
-        )
-    ]
-)
+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
 
 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-flash-latest" # Or gemini-flash-latest
-
-    # We use a chat session to handle the multi-turn tool loop
-    chat = client.chats.create(
-        model=model_id,
-        config=types.GenerateContentConfig(
-            tools=[train_tool_declaration, types.Tool(google_search=types.GoogleSearch())],
-            temperature=0.4
+    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."
 
-    try:
-        # Step 1: Send initial request to Gemini
-        response = chat.send_message(input_text)
-        
-        # Step 2: Check if Gemini wants to use a tool
-        # We loop in case Gemini needs multiple tool calls
-        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":
-                # Extract arguments Gemini provided
-                args = tool_call.args
-                
-                # Execute the MCP call (running async code in sync Gradio)
-                train_data = asyncio.run(call_mcp_train_tool(args["start_loc"], args["dest_loc"]))
-                
-                # Send the tool result back to Gemini
-                response = chat.send_message(
-                    types.Part.from_function_response(
-                        name="get_train_connections",
-                        response={"result": train_data}
-                    )
-                )
-            else:
-                break # Handle other tools or exit
+    model = "gemini-flash-latest"
+    contents = [
+        types.Content(
+            role="user",
+            parts=[
+                types.Part.from_text(text=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",
+    )
 
-        return response.text, ""
 
+    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: {str(e)}", ""
-
-# --- GRADIO UI ---
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 🚄 Gemini 2.0 + DB Train MCP")
-    gr.Markdown("Ask about train connections (e.g., 'When is the next train from Munich to Berlin?')")
-    
-    output_textbox = gr.Markdown(label="Response")
-    input_textbox = gr.Textbox(lines=2, label="Your Question")
-    submit_button = gr.Button("Send", variant="primary")
-    
-    submit_button.click(
-        fn=generate,
-        inputs=input_textbox,
-        outputs=[output_textbox, input_textbox]
-    )
+        return f"Error during generation: {e}"
+    data = clean_json_string(response_text)
+    data = data[:-1]
+    return response_text, ""
+   
 
 if __name__ == '__main__':
-    demo.launch()
+
+    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)