bstraehle commited on
Commit
eaa21a6
·
verified ·
1 Parent(s): ee32d98

Update crew.py

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