Update app.py

app.py

@@ -3,13 +3,18 @@ import os
 import google.generativeai as genai
 import logging
 import time
-#import backoff
+import backoff
+import google.ai.generativelanguage as glm
 
 # Configure Logging
 logging.basicConfig(level=logging.ERROR, format='%(asctime)s - %(levelname)s - %(message)s')
 
 # Load environment variables
-genai.configure(api_key=os.environ["geminiapikey"])
+try:
+    genai.configure(api_key=os.environ["geminiapikey"])
+except KeyError:
+    logging.error("Error: 'geminiapikey' environment variable not found.")
+    exit(1)
 
 read_key = os.environ.get('HF_TOKEN', None)
 
@@ -24,11 +29,13 @@ custom_css = """
 }
 """
 
-
+@backoff.on_exception(backoff.expo,
+                      (genai.APIError),
+                      max_tries=3)
 def predict(prompt):
     # Create the model
     generation_config = {
-        "temperature": 0.4,
+        "temperature": 0.7,
         "top_p": 0.95,
         "top_k": 40,
         "max_output_tokens": 2048,
@@ -40,10 +47,30 @@ def predict(prompt):
         generation_config=generation_config,
     )
 
-    #contents_to_send = [genai.Content(parts=[prompt])]
-    contents_to_send = [prompt]
-    response = model.generate_content(contents=contents_to_send, tools='google_search_retrieval')
-    return response.text
+    try:
+      # Create the tools configuration
+
+      tools_config = glm.ToolConfig(
+          function_declarations=[],
+          search_queries=[prompt],
+      )
+
+      response = model.generate_content(
+          contents=[prompt],  # Directly pass the prompt
+          tools=[tools_config]
+      )
+
+      if response and response.text:
+          return response.text
+      else:
+          logging.error(f"Unexpected response: {response}")
+          return "Error: Could not extract text from the response."
+    except genai.APIError as e:
+       logging.error(f"API error occurred: {e}")
+       raise
+    except Exception as e:
+        logging.error(f"An error occurred: {e}")
+        return f"An error occurred: {e}"
 
 # Create the Gradio interface
 with gr.Blocks(css=custom_css) as demo: