bstraehle commited on
Commit
234b682
·
verified ·
1 Parent(s): b0cfa9f

Update crew.py

Browse files
Files changed (1) hide show
  1. crew.py +47 -6
crew.py CHANGED
@@ -25,6 +25,7 @@ IMAGE_ANALYSIS_MODEL = "gemini-2.5-flash-preview-04-17"
25
  AUDIO_ANALYSIS_MODEL = "gemini-2.5-flash-preview-04-17"
26
  VIDEO_ANALYSIS_MODEL = "gemini-2.5-flash-preview-04-17"
27
  YOUTUBE_ANALYSIS_MODEL = "gemini-2.5-flash-preview-04-17"
 
28
  CODE_GENERATION_MODEL = "gemini-2.5-flash-preview-04-17"
29
  CODE_EXECUTION_MODEL = "gemini-2.5-flash-preview-04-17"
30
 
@@ -152,6 +153,33 @@ def run_crew(question, file_path):
152
  return response.text
153
  except Exception as e:
154
  raise RuntimeError(f"Processing failed: {str(e)}")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
155
 
156
  @tool("YouTube Analysis Tool")
157
  def youtube_analysis_tool(question: str, url: str) -> str:
@@ -178,7 +206,7 @@ def run_crew(question, file_path):
178
  )
179
  except Exception as e:
180
  raise RuntimeError(f"Processing failed: {str(e)}")
181
-
182
  @tool("Code Generation Tool")
183
  def code_generation_tool(question: str, json_data: str) -> str:
184
  """Given a question and JSON data, generate and execute code to answer the question.
@@ -298,6 +326,17 @@ def run_crew(question, file_path):
298
  verbose=True
299
  )
300
 
 
 
 
 
 
 
 
 
 
 
 
301
  code_generation_agent = Agent(
302
  role="Code Generation Agent",
303
  goal="Given a question and JSON data, generate and execute code to answer the question: {question}",
@@ -336,12 +375,13 @@ def run_crew(question, file_path):
336
  agent=manager_agent,
337
  description="Answer the following question. If needed, delegate to one of your coworkers: "
338
  "- Web Search Agent requires a question only. "
339
- "- Image Analysis Agent requires a question and **image file**. "
340
- "- Audio Analysis Agent requires a question and **audio file**. "
341
- "- Video Analysis Agent requires a question and **video file**. "
 
342
  "- YouTube Analysis Agent requires a question and **YouTube URL**. "
343
  "- Code Generation Agent requires a question and **JSON data**. "
344
- "- Code Execution Agent requires a question and **Python file**. "
345
  "Question: {question}",
346
  expected_output="The answer to the question."
347
  )
@@ -354,7 +394,8 @@ def run_crew(question, file_path):
354
  audio_analysis_agent,
355
  video_analysis_agent,
356
  youtube_analysis_agent,
357
- code_generation_agent,
 
358
  code_execution_agent],
359
  manager_agent=manager_agent,
360
  tasks=[manager_task],
 
25
  AUDIO_ANALYSIS_MODEL = "gemini-2.5-flash-preview-04-17"
26
  VIDEO_ANALYSIS_MODEL = "gemini-2.5-flash-preview-04-17"
27
  YOUTUBE_ANALYSIS_MODEL = "gemini-2.5-flash-preview-04-17"
28
+ DOCUMENT_ANALYSIS_MODEL = "gemini-2.5-flash-preview-04-17"
29
  CODE_GENERATION_MODEL = "gemini-2.5-flash-preview-04-17"
30
  CODE_EXECUTION_MODEL = "gemini-2.5-flash-preview-04-17"
31
 
 
153
  return response.text
154
  except Exception as e:
155
  raise RuntimeError(f"Processing failed: {str(e)}")
156
+
157
+ @tool("Document Analysis Tool")
158
+ def document_analysis_tool(question: str, file_path: str) -> str:
159
+ """Given a question and document file, analyze the document to answer the question.
160
+
161
+ Args:
162
+ question (str): Question about a document file
163
+ file_path (str): The document file path
164
+
165
+ Returns:
166
+ str: Answer to the question about the document file
167
+
168
+ Raises:
169
+ RuntimeError: If processing fails"""
170
+ try:
171
+ client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
172
+
173
+ file = client.files.upload(file=file_path)
174
+
175
+ response = client.models.generate_content(
176
+ model=DOCUMENT_ANALYSIS_MODEL,
177
+ contents=[file, question]
178
+ )
179
+
180
+ return response.text
181
+ except Exception as e:
182
+ raise RuntimeError(f"Processing failed: {str(e)}")
183
 
184
  @tool("YouTube Analysis Tool")
185
  def youtube_analysis_tool(question: str, url: str) -> str:
 
206
  )
207
  except Exception as e:
208
  raise RuntimeError(f"Processing failed: {str(e)}")
209
+
210
  @tool("Code Generation Tool")
211
  def code_generation_tool(question: str, json_data: str) -> str:
212
  """Given a question and JSON data, generate and execute code to answer the question.
 
326
  verbose=True
327
  )
328
 
329
+ document_analysis_agent = Agent(
330
+ role="Document Analysis Agent",
331
+ goal="Given a question and document file, analyze the document and answer the question: {question}",
332
+ backstory="As an expert document analysis assistant, you analyze the document to answer the question.",
333
+ allow_delegation=False,
334
+ llm=AGENT_MODEL,
335
+ max_iter=2,
336
+ tools=[document_analysis_tool],
337
+ verbose=True
338
+ )
339
+
340
  code_generation_agent = Agent(
341
  role="Code Generation Agent",
342
  goal="Given a question and JSON data, generate and execute code to answer the question: {question}",
 
375
  agent=manager_agent,
376
  description="Answer the following question. If needed, delegate to one of your coworkers: "
377
  "- Web Search Agent requires a question only. "
378
+ "- Image Analysis Agent requires a question and **.png, .jpeg, .webp, .heic, or .heif image file**. "
379
+ "- Audio Analysis Agent requires a question and **.wav, .mp3, .aiff, .aac, .ogg, or .flac audio file**. "
380
+ "- Video Analysis Agent requires a question and **.mp4, .mpeg, .mov, .avi, .x-flv, .mpg, .webm, .wmv, or .3gpp video file**. "
381
+ "- Document Analysis Agent requires a question and **.pdf, .txt, .html, css, .js, .md, .xml, or .rtf document file**. "
382
  "- YouTube Analysis Agent requires a question and **YouTube URL**. "
383
  "- Code Generation Agent requires a question and **JSON data**. "
384
+ "- Code Execution Agent requires a question and **.py Python file**. "
385
  "Question: {question}",
386
  expected_output="The answer to the question."
387
  )
 
394
  audio_analysis_agent,
395
  video_analysis_agent,
396
  youtube_analysis_agent,
397
+ document_analysis_agent,
398
+ code_generation_agent,
399
  code_execution_agent],
400
  manager_agent=manager_agent,
401
  tasks=[manager_task],