eaglelandsonce commited on
Commit
aa910e9
1 Parent(s): fb84bde

Update crewai/tools/gemini_tools.py

Browse files
Files changed (1) hide show
  1. crewai/tools/gemini_tools.py +18 -8
crewai/tools/gemini_tools.py CHANGED
@@ -4,6 +4,7 @@ import json
4
  import os
5
 
6
  import google.generativeai as genai
 
7
 
8
  # Retrieve API Key from Environment Variable
9
  GOOGLE_AI_STUDIO = os.environ.get('GOOGLE_AI_STUDIO2')
@@ -22,14 +23,23 @@ model = genai.GenerativeModel('gemini-pro')
22
  class GeminiSearchTools():
23
  @tool("Gemini search the internet")
24
  def gemini_search(query):
25
- """Useful to search the internet
26
- about a a given topic and return relevant results"""
27
- response = model.generate_content(query)
28
- return response.text
 
 
 
 
 
29
 
30
  @tool("Gemini search news on the internet")
31
  def gemini_search_news(query):
32
- """Useful to search news about a company, stock or any other
33
- topic and return relevant results"""""
34
- response = model.generate_content(query)
35
- return response.text
 
 
 
 
 
4
  import os
5
 
6
  import google.generativeai as genai
7
+ from google.api_core import exceptions
8
 
9
  # Retrieve API Key from Environment Variable
10
  GOOGLE_AI_STUDIO = os.environ.get('GOOGLE_AI_STUDIO2')
 
23
  class GeminiSearchTools():
24
  @tool("Gemini search the internet")
25
  def gemini_search(query):
26
+ try:
27
+ response = model.generate_content(query)
28
+ return response.text
29
+ except exceptions.DeadlineExceeded as e:
30
+ # Handle the DeadlineExceeded exception here
31
+ print("Error: Deadline Exceeded -", str(e))
32
+ # You can return a custom message or take other appropriate actions
33
+ return "Error: The request timed out. Please try again later."
34
+
35
 
36
  @tool("Gemini search news on the internet")
37
  def gemini_search_news(query):
38
+ try:
39
+ response = model.generate_content(query)
40
+ return response.text
41
+ except exceptions.DeadlineExceeded as e:
42
+ # Handle the DeadlineExceeded exception here
43
+ print("Error: Deadline Exceeded -", str(e))
44
+ # You can return a custom message or take other appropriate actions
45
+ return "Error: The request timed out. Please try again later."