bstraehle commited on
Commit
54e1f0f
·
verified ·
1 Parent(s): 2aa67da

Update tools.py

Browse files
Files changed (1) hide show
  1. tools.py +268 -1
tools.py CHANGED
@@ -1,8 +1,271 @@
1
  # Automatic function calling
2
  # https://ai.google.dev/gemini-api/docs/function-calling
3
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
4
  # Arithmetic functions
5
 
 
6
  def add(a: float, b: float) -> float:
7
  """Add two numbers.
8
 
@@ -15,6 +278,7 @@ def add(a: float, b: float) -> float:
15
  """
16
  return a + b
17
 
 
18
  def subtract(a: float, b: float) -> float:
19
  """Subtract two numbers.
20
 
@@ -26,7 +290,8 @@ def subtract(a: float, b: float) -> float:
26
  Result number
27
  """
28
  return a - b
29
-
 
30
  def multiply(a: float, b: float) -> float:
31
  """Multiply two numbers.
32
  Args:
@@ -38,6 +303,7 @@ def multiply(a: float, b: float) -> float:
38
  """
39
  return a * b
40
 
 
41
  def divide(a: float, b: float) -> float:
42
  """Divide two numbers.
43
 
@@ -52,6 +318,7 @@ def divide(a: float, b: float) -> float:
52
  raise ValueError("Cannot divide by zero.")
53
  return a / b
54
 
 
55
  def modulus(a: float, b: float) -> float:
56
  """Get the modulus of two numbers.
57
 
 
1
  # Automatic function calling
2
  # https://ai.google.dev/gemini-api/docs/function-calling
3
 
4
+ @tool("Web Search Tool")
5
+ def web_search_tool(question: str) -> str:
6
+ """Given a question only, search the web to answer the question.
7
+
8
+ Args:
9
+ question (str): Question to answer
10
+
11
+ Returns:
12
+ str: Answer to the question
13
+
14
+ Raises:
15
+ RuntimeError: If processing fails"""
16
+ try:
17
+ client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
18
+
19
+ response = client.models.generate_content(
20
+ model=WEB_SEARCH_MODEL,
21
+ contents=question,
22
+ config=types.GenerateContentConfig(
23
+ tools=[types.Tool(google_search=types.GoogleSearchRetrieval())]
24
+ )
25
+ )
26
+
27
+ return response.text
28
+ except Exception as e:
29
+ raise RuntimeError(f"Processing failed: {str(e)}")
30
+
31
+ @tool("Image Analysis Tool")
32
+ def image_analysis_tool(question: str, file_path: str) -> str:
33
+ """Given a question and image file, analyze the image to answer the question.
34
+
35
+ Args:
36
+ question (str): Question about an image file
37
+ file_path (str): The image file path
38
+
39
+ Returns:
40
+ str: Answer to the question about the image file
41
+
42
+ Raises:
43
+ RuntimeError: If processing fails"""
44
+ try:
45
+ client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
46
+
47
+ file = client.files.upload(file=file_path)
48
+
49
+ response = client.models.generate_content(
50
+ model=IMAGE_ANALYSIS_MODEL,
51
+ contents=[file, question]
52
+ )
53
+
54
+ return response.text
55
+ except Exception as e:
56
+ raise RuntimeError(f"Processing failed: {str(e)}")
57
+
58
+ @tool("Audio Analysis Tool")
59
+ def audio_analysis_tool(question: str, file_path: str) -> str:
60
+ """Given a question and audio file, analyze the audio to answer the question.
61
+
62
+ Args:
63
+ question (str): Question about an audio file
64
+ file_path (str): The audio file path
65
+
66
+ Returns:
67
+ str: Answer to the question about the audio file
68
+
69
+ Raises:
70
+ RuntimeError: If processing fails"""
71
+ try:
72
+ client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
73
+
74
+ file = client.files.upload(file=file_path)
75
+
76
+ response = client.models.generate_content(
77
+ model=AUDIO_ANALYSIS_MODEL,
78
+ contents=[file, question]
79
+ )
80
+
81
+ return response.text
82
+ except Exception as e:
83
+ raise RuntimeError(f"Processing failed: {str(e)}")
84
+
85
+ @tool("Video Analysis Tool")
86
+ def video_analysis_tool(question: str, file_path: str) -> str:
87
+ """Given a question and video file, analyze the video to answer the question.
88
+
89
+ Args:
90
+ question (str): Question about a video file
91
+ file_path (str): The video file path
92
+
93
+ Returns:
94
+ str: Answer to the question about the video file
95
+
96
+ Raises:
97
+ RuntimeError: If processing fails"""
98
+ try:
99
+ client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
100
+
101
+ file = client.files.upload(file=file_path)
102
+
103
+ response = client.models.generate_content(
104
+ model=VIDEO_ANALYSIS_MODEL,
105
+ contents=[file, question]
106
+ )
107
+
108
+ return response.text
109
+ except Exception as e:
110
+ raise RuntimeError(f"Processing failed: {str(e)}")
111
+
112
+ @tool("YouTube Analysis Tool")
113
+ def youtube_analysis_tool(question: str, url: str) -> str:
114
+ """Given a question and YouTube URL, analyze the video to answer the question.
115
+
116
+ Args:
117
+ question (str): Question about a YouTube video
118
+ url (str): The YouTube URL
119
+
120
+ Returns:
121
+ str: Answer to the question about the YouTube video
122
+
123
+ Raises:
124
+ RuntimeError: If processing fails"""
125
+ try:
126
+ client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
127
+
128
+ return client.models.generate_content(
129
+ model=YOUTUBE_ANALYSIS_MODEL,
130
+ contents=types.Content(
131
+ parts=[types.Part(file_data=types.FileData(file_uri=url)),
132
+ types.Part(text=question)]
133
+ )
134
+ )
135
+ except Exception as e:
136
+ raise RuntimeError(f"Processing failed: {str(e)}")
137
+
138
+ @tool("Document Analysis Tool")
139
+ def document_analysis_tool(question: str, file_path: str) -> str:
140
+ """Given a question and document file, analyze the document to answer the question.
141
+
142
+ Args:
143
+ question (str): Question about a document file
144
+ file_path (str): The document file path
145
+
146
+ Returns:
147
+ str: Answer to the question about the document file
148
+
149
+ Raises:
150
+ RuntimeError: If processing fails"""
151
+ try:
152
+ client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
153
+
154
+ contents = []
155
+
156
+ if is_ext(file_path, ".docx"):
157
+ text_data = read_docx_text(file_path)
158
+ contents = [f"{question}\n{text_data}"]
159
+ print(f"=> Text data:\n{text_data}")
160
+ elif is_ext(file_path, ".pptx"):
161
+ text_data = read_pptx_text(file_path)
162
+ contents = [f"{question}\n{text_data}"]
163
+ print(f"=> Text data:\n{text_data}")
164
+ else:
165
+ file = client.files.upload(file=file_path)
166
+ contents = [file, question]
167
+
168
+ response = client.models.generate_content(
169
+ model=DOCUMENT_ANALYSIS_MODEL,
170
+ contents=contents
171
+ )
172
+
173
+ return response.text
174
+ except Exception as e:
175
+ raise RuntimeError(f"Processing failed: {str(e)}")
176
+
177
+ @tool("Arithmetic Tool")
178
+ def arithmetic_tool(question: str, a: float, b: float) -> float:
179
+ """Given a question and two numbers, perform the calculation to answer the question.
180
+
181
+ Args:
182
+ question (str): Question to answer
183
+ a (float): First number
184
+ b (float): Second number
185
+
186
+ Returns:
187
+ float: Result number
188
+
189
+ Raises:
190
+ RuntimeError: If processing fails"""
191
+ try:
192
+ client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
193
+
194
+ response = client.models.generate_content(
195
+ model=ARITHMETIC_MODEL,
196
+ contents=question,
197
+ config=types.GenerateContentConfig(
198
+ tools=[add, subtract, multiply, divide, modulus]
199
+ )
200
+ )
201
+
202
+ return response.text
203
+ except Exception as e:
204
+ raise RuntimeError(f"Processing failed: {str(e)}")
205
+
206
+ @tool("Code Execution Tool")
207
+ def code_execution_tool(question: str, file_path: str) -> str:
208
+ """Given a question and Python file, execute the file to answer the question.
209
+
210
+ Args:
211
+ question (str): Question to answer
212
+ file_path (str): The Python file path
213
+
214
+ Returns:
215
+ str: Answer to the question
216
+
217
+ Raises:
218
+ RuntimeError: If processing fails"""
219
+ try:
220
+ client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
221
+
222
+ file = client.files.upload(file=file_path)
223
+
224
+ response = client.models.generate_content(
225
+ model=CODE_EXECUTION_MODEL,
226
+ contents=[file, question],
227
+ config=types.GenerateContentConfig(
228
+ tools=[types.Tool(code_execution=types.ToolCodeExecution)]
229
+ ),
230
+ )
231
+
232
+ for part in response.candidates[0].content.parts:
233
+ if part.code_execution_result is not None:
234
+ return part.code_execution_result.output
235
+ except Exception as e:
236
+ raise RuntimeError(f"Processing failed: {str(e)}")
237
+
238
+ @tool("Code Generation Tool")
239
+ def code_generation_tool(question: str, json_data: str) -> str:
240
+ """Given a question and JSON data, generate and execute code to answer the question.
241
+ Args:
242
+ question (str): Question to answer
243
+ file_path (str): The JSON data
244
+ Returns:
245
+ str: Answer to the question
246
+
247
+ Raises:
248
+ RuntimeError: If processing fails"""
249
+ try:
250
+ client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
251
+
252
+ response = client.models.generate_content(
253
+ model=CODE_GENERATION_MODEL,
254
+ contents=[f"{question}\n{json_data}"],
255
+ config=types.GenerateContentConfig(
256
+ tools=[types.Tool(code_execution=types.ToolCodeExecution)]
257
+ ),
258
+ )
259
+
260
+ for part in response.candidates[0].content.parts:
261
+ if part.code_execution_result is not None:
262
+ return part.code_execution_result.output
263
+ except Exception as e:
264
+ raise RuntimeError(f"Processing failed: {str(e)}")
265
+
266
  # Arithmetic functions
267
 
268
+ @tool("Addition Tool")
269
  def add(a: float, b: float) -> float:
270
  """Add two numbers.
271
 
 
278
  """
279
  return a + b
280
 
281
+ @tool("Subtraction Tool")
282
  def subtract(a: float, b: float) -> float:
283
  """Subtract two numbers.
284
 
 
290
  Result number
291
  """
292
  return a - b
293
+
294
+ @tool("Multiplication Tool")
295
  def multiply(a: float, b: float) -> float:
296
  """Multiply two numbers.
297
  Args:
 
303
  """
304
  return a * b
305
 
306
+ @tool("Division Tool")
307
  def divide(a: float, b: float) -> float:
308
  """Divide two numbers.
309
 
 
318
  raise ValueError("Cannot divide by zero.")
319
  return a / b
320
 
321
+ @tool("Modulus Tool")
322
  def modulus(a: float, b: float) -> float:
323
  """Get the modulus of two numbers.
324