Update app.py

app.py

@@ -1,21 +1,84 @@
-async def generate(input_text):
-    if not GEMINI_API_KEY:
-        yield "Error: GEMINI_API_KEY is not set.", ""
-        return # Stops the generator
+import base64
+import gradio as gr
+import os
+import json
+from google import genai
+from google.genai import types
+from gradio_client import Client
 
-    # ... (setup code) ...
+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):
     try:
-        yield "Thinking...", "" # UI update
-        
-        # ... (logic loop) ...
-        
-        if tool_calls:
-            yield "Fetching live train data from MCP server...", ""
-            # ... (execute tool) ...
+        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."
+
+    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",
+    )
 
-        # Final result
-        yield 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:
-        yield f"### Logic Error\n{str(e)}", ""
+        return f"Error during generation: {e}"
+    data = clean_json_string(response_text)
+    data = data[:-1]
+    return response_text, ""
+   
+
+if __name__ == '__main__':
+
+    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)