Update app.py

app.py

@@ -8,40 +8,42 @@ API_KEY = os.getenv("GOOGLE_API_KEY")  # Ensure to set this in Hugging Face Secr
 client = genai.Client(api_key=API_KEY)
 MODEL_ID = "gemini-2.0-flash-exp"  # Replace with your desired model ID
 
-def google_search_query(question):
+def google_search_query(question, use_web_search):
     try:
-        # Define the Google Search Tool
-        google_search_tool = Tool(google_search=GoogleSearch())
-        
-        # Generate the response
-        response = client.models.generate_content(
-            model=MODEL_ID,
-            contents=question,
-            config=GenerateContentConfig(tools=[google_search_tool]),
-        )
-
-        # Extract AI response and search results
-        ai_response = response.text  # AI response as plain text
-        search_results = response.candidates[0].grounding_metadata.search_entry_point.rendered_content
-
-        # Parse and format search results to separate links
-        link_list = "\n".join([f"<a href='{entry.url}' target='_blank'>{entry.title}</a>" for entry in response.candidates[0].grounding_metadata.search_entry_point.entries])
-        return ai_response, link_list
+        if use_web_search:
+            google_search_tool = Tool(google_search=GoogleSearch())
+            response = client.models.generate_content(
+                model=MODEL_ID,
+                contents=question,
+                config=GenerateContentConfig(tools=[google_search_tool]),
+            )
+            ai_response = response.text  # AI response as plain text
+            link_list = "\n".join([f"<a href='{entry.url}' target='_blank'>{entry.title}</a>" for entry in response.candidates[0].grounding_metadata.search_entry_point.entries])
+            return ai_response, link_list
+        else:
+            # Return response without web search
+            response = client.models.generate_content(
+                model=MODEL_ID,
+                contents=question,
+            )
+            return response.text, "Web search not used."
     except Exception as e:
         return f"Error: {str(e)}", "No results found."
 
-# Gradio Interface
-app = gr.Interface(
-    fn=google_search_query,
-    inputs=gr.Textbox(lines=2, label="Ask a Question"),
-    outputs=[
-        gr.Textbox(label="AI Response"),
-        gr.HTML(label="Search Results"),
-    ],
-    title="Google Search with Gemini AI",
-    description="Ask a question, and the AI will use Google search tools to provide an answer along with contextual search results.",
-    theme="huggingface",
-)
+with gr.Blocks() as app:
+    with gr.Row():
+        question_input = gr.Textbox(label="Ask a Question", lines=2)
+        web_search_checkbox = gr.Checkbox(label="Enable Web Search", value=False)
+    with gr.Row():
+        submit_button = gr.Button("Submit")
+    with gr.Row():
+        ai_response_output = gr.Textbox(label="AI Response")
+        search_results_output = gr.HTML(label="Search Results")
+
+    submit_button.click(
+        google_search_query,
+        inputs=[question_input, web_search_checkbox],
+        outputs=[ai_response_output, search_results_output]
+    )
 
-if __name__ == "__main__":
-    app.launch(share=True)
+app.launch(share=True)