|
|
import os |
|
|
|
|
|
from crewai.tools import tool |
|
|
from crewai_tools import StagehandTool |
|
|
from google import genai |
|
|
from google.genai import types |
|
|
from utils import ( |
|
|
read_file_json, |
|
|
read_docx_text, |
|
|
read_pptx_text, |
|
|
is_ext |
|
|
) |
|
|
|
|
|
WEB_SEARCH_MODEL = "gemini-2.5-flash-preview-04-17" |
|
|
IMAGE_ANALYSIS_MODEL = "gemini-2.5-flash-preview-04-17" |
|
|
AUDIO_ANALYSIS_MODEL = "gemini-2.5-flash-preview-04-17" |
|
|
VIDEO_ANALYSIS_MODEL = "gemini-2.5-flash-preview-04-17" |
|
|
YOUTUBE_ANALYSIS_MODEL = "gemini-2.5-flash-preview-04-17" |
|
|
DOCUMENT_ANALYSIS_MODEL = "gemini-2.5-flash-preview-04-17" |
|
|
ARITHMETIC_MODEL = "gemini-2.5-flash-preview-04-17" |
|
|
CODE_GENERATION_MODEL = "gemini-2.5-flash-preview-04-17" |
|
|
CODE_EXECUTION_MODEL = "gemini-2.5-flash-preview-04-17" |
|
|
|
|
|
WEB_BROWSER_MODEL = "claude-3-7-sonnet-latest" |
|
|
|
|
|
class AITools(): |
|
|
@tool("Web Search Tool") |
|
|
def web_search_tool(question: str) -> str: |
|
|
"""Given a question only, search the web to answer the question. |
|
|
|
|
|
Args: |
|
|
question (str): Question to answer |
|
|
|
|
|
Returns: |
|
|
str: Answer to the question |
|
|
|
|
|
Raises: |
|
|
RuntimeError: If processing fails""" |
|
|
try: |
|
|
client = genai.Client(api_key=os.environ["GEMINI_API_KEY"]) |
|
|
|
|
|
response = client.models.generate_content( |
|
|
model=WEB_SEARCH_MODEL, |
|
|
contents=question, |
|
|
config=types.GenerateContentConfig( |
|
|
tools=[types.Tool(google_search=types.GoogleSearchRetrieval())] |
|
|
) |
|
|
) |
|
|
|
|
|
return response.text |
|
|
except Exception as e: |
|
|
raise RuntimeError(f"Processing failed: {str(e)}") |
|
|
|
|
|
@tool("Web Browser Tool") |
|
|
def web_browser_tool(question: str, url: str) -> str: |
|
|
"""Given a question and URL, load the URL and act, extract, or observe to answer the question. |
|
|
|
|
|
Args: |
|
|
question (str): Question about a URL |
|
|
url (str): The URL |
|
|
|
|
|
Returns: |
|
|
str: Answer to the question |
|
|
|
|
|
Raises: |
|
|
RuntimeError: If processing fails""" |
|
|
try: |
|
|
stagehand_tool = StagehandTool( |
|
|
api_key=os.environ["BROWSERBASE_API_KEY"], |
|
|
project_id=os.environ["BROWSERBASE_PROJECT_ID"], |
|
|
model_api_key=os.environ["MODEL_API_KEY"], |
|
|
model_name=WEB_BROWSER_MODEL, |
|
|
dom_settle_timeout_ms=5000, |
|
|
headless=True, |
|
|
self_heal=True, |
|
|
wait_for_captcha_solves=True, |
|
|
verbose=3 |
|
|
) |
|
|
|
|
|
return stagehand_tool.run( |
|
|
instruction=question, |
|
|
url=url |
|
|
) |
|
|
except Exception as e: |
|
|
raise RuntimeError(f"Processing failed: {str(e)}") |
|
|
finally: |
|
|
stagehand_tool.close() |
|
|
|
|
|
@tool("Image Analysis Tool") |
|
|
def image_analysis_tool(question: str, file_path: str) -> str: |
|
|
"""Given a question and image file, analyze the image to answer the question. |
|
|
|
|
|
Args: |
|
|
question (str): Question about an image file |
|
|
file_path (str): The image file path |
|
|
|
|
|
Returns: |
|
|
str: Answer to the question about the image file |
|
|
|
|
|
Raises: |
|
|
RuntimeError: If processing fails""" |
|
|
try: |
|
|
client = genai.Client(api_key=os.environ["GEMINI_API_KEY"]) |
|
|
|
|
|
file = client.files.upload(file=file_path) |
|
|
|
|
|
response = client.models.generate_content( |
|
|
model=IMAGE_ANALYSIS_MODEL, |
|
|
contents=[file, question] |
|
|
) |
|
|
|
|
|
return response.text |
|
|
except Exception as e: |
|
|
raise RuntimeError(f"Processing failed: {str(e)}") |
|
|
|
|
|
@tool("Audio Analysis Tool") |
|
|
def audio_analysis_tool(question: str, file_path: str) -> str: |
|
|
"""Given a question and audio file, analyze the audio to answer the question. |
|
|
|
|
|
Args: |
|
|
question (str): Question about an audio file |
|
|
file_path (str): The audio file path |
|
|
|
|
|
Returns: |
|
|
str: Answer to the question about the audio file |
|
|
|
|
|
Raises: |
|
|
RuntimeError: If processing fails""" |
|
|
try: |
|
|
client = genai.Client(api_key=os.environ["GEMINI_API_KEY"]) |
|
|
|
|
|
file = client.files.upload(file=file_path) |
|
|
|
|
|
response = client.models.generate_content( |
|
|
model=AUDIO_ANALYSIS_MODEL, |
|
|
contents=[file, question] |
|
|
) |
|
|
|
|
|
return response.text |
|
|
except Exception as e: |
|
|
raise RuntimeError(f"Processing failed: {str(e)}") |
|
|
|
|
|
@tool("Video Analysis Tool") |
|
|
def video_analysis_tool(question: str, file_path: str) -> str: |
|
|
"""Given a question and video file, analyze the video to answer the question. |
|
|
|
|
|
Args: |
|
|
question (str): Question about a video file |
|
|
file_path (str): The video file path |
|
|
|
|
|
Returns: |
|
|
str: Answer to the question about the video file |
|
|
|
|
|
Raises: |
|
|
RuntimeError: If processing fails""" |
|
|
try: |
|
|
client = genai.Client(api_key=os.environ["GEMINI_API_KEY"]) |
|
|
|
|
|
file = client.files.upload(file=file_path) |
|
|
|
|
|
response = client.models.generate_content( |
|
|
model=VIDEO_ANALYSIS_MODEL, |
|
|
contents=[file, question] |
|
|
) |
|
|
|
|
|
return response.text |
|
|
except Exception as e: |
|
|
raise RuntimeError(f"Processing failed: {str(e)}") |
|
|
|
|
|
@tool("YouTube Analysis Tool") |
|
|
def youtube_analysis_tool(question: str, url: str) -> str: |
|
|
"""Given a question and YouTube URL, analyze the video to answer the question. |
|
|
|
|
|
Args: |
|
|
question (str): Question about a YouTube video |
|
|
url (str): The YouTube URL |
|
|
|
|
|
Returns: |
|
|
str: Answer to the question about the YouTube video |
|
|
|
|
|
Raises: |
|
|
RuntimeError: If processing fails""" |
|
|
try: |
|
|
client = genai.Client(api_key=os.environ["GEMINI_API_KEY"]) |
|
|
|
|
|
return client.models.generate_content( |
|
|
model=YOUTUBE_ANALYSIS_MODEL, |
|
|
contents=types.Content( |
|
|
parts=[types.Part(file_data=types.FileData(file_uri=url)), |
|
|
types.Part(text=question)] |
|
|
) |
|
|
) |
|
|
except Exception as e: |
|
|
raise RuntimeError(f"Processing failed: {str(e)}") |
|
|
|
|
|
@tool("Document Analysis Tool") |
|
|
def document_analysis_tool(question: str, file_path: str) -> str: |
|
|
"""Given a question and document file, analyze the document to answer the question. |
|
|
|
|
|
Args: |
|
|
question (str): Question about a document file |
|
|
file_path (str): The document file path |
|
|
|
|
|
Returns: |
|
|
str: Answer to the question about the document file |
|
|
|
|
|
Raises: |
|
|
RuntimeError: If processing fails""" |
|
|
try: |
|
|
client = genai.Client(api_key=os.environ["GEMINI_API_KEY"]) |
|
|
|
|
|
contents = [] |
|
|
|
|
|
if is_ext(file_path, ".docx"): |
|
|
text_data = read_docx_text(file_path) |
|
|
contents = [f"{question}\n{text_data}"] |
|
|
print(f"=> Text data:\n{text_data}") |
|
|
elif is_ext(file_path, ".pptx"): |
|
|
text_data = read_pptx_text(file_path) |
|
|
contents = [f"{question}\n{text_data}"] |
|
|
print(f"=> Text data:\n{text_data}") |
|
|
else: |
|
|
file = client.files.upload(file=file_path) |
|
|
contents = [file, question] |
|
|
|
|
|
response = client.models.generate_content( |
|
|
model=DOCUMENT_ANALYSIS_MODEL, |
|
|
contents=contents |
|
|
) |
|
|
|
|
|
return response.text |
|
|
except Exception as e: |
|
|
raise RuntimeError(f"Processing failed: {str(e)}") |
|
|
|
|
|
@tool("Arithmetic Tool") |
|
|
def arithmetic_tool(question: str, a: float, b: float) -> float: |
|
|
"""Given a question and two numbers, perform the calculation to answer the question. |
|
|
|
|
|
Args: |
|
|
question (str): Question to answer |
|
|
a (float): First number |
|
|
b (float): Second number |
|
|
|
|
|
Returns: |
|
|
float: Result number |
|
|
|
|
|
Raises: |
|
|
RuntimeError: If processing fails""" |
|
|
try: |
|
|
client = genai.Client(api_key=os.environ["GEMINI_API_KEY"]) |
|
|
|
|
|
response = client.models.generate_content( |
|
|
model=ARITHMETIC_MODEL, |
|
|
contents=question, |
|
|
config=types.GenerateContentConfig( |
|
|
tools=[add, subtract, multiply, divide, modulus] |
|
|
) |
|
|
) |
|
|
|
|
|
return response.text |
|
|
except Exception as e: |
|
|
raise RuntimeError(f"Processing failed: {str(e)}") |
|
|
|
|
|
@tool("Code Generation Tool") |
|
|
def code_generation_tool(question: str, json_data: str) -> str: |
|
|
"""Given a question and JSON data, generate and execute code to answer the question. |
|
|
Args: |
|
|
question (str): Question to answer |
|
|
file_path (str): The JSON data |
|
|
Returns: |
|
|
str: Answer to the question |
|
|
|
|
|
Raises: |
|
|
RuntimeError: If processing fails""" |
|
|
try: |
|
|
client = genai.Client(api_key=os.environ["GEMINI_API_KEY"]) |
|
|
|
|
|
response = client.models.generate_content( |
|
|
model=CODE_GENERATION_MODEL, |
|
|
contents=[f"{question}\n{json_data}"], |
|
|
config=types.GenerateContentConfig( |
|
|
tools=[types.Tool(code_execution=types.ToolCodeExecution)] |
|
|
), |
|
|
) |
|
|
|
|
|
for part in response.candidates[0].content.parts: |
|
|
if part.code_execution_result is not None: |
|
|
return part.code_execution_result.output |
|
|
except Exception as e: |
|
|
raise RuntimeError(f"Processing failed: {str(e)}") |
|
|
|
|
|
@tool("Code Execution Tool") |
|
|
def code_execution_tool(question: str, file_path: str) -> str: |
|
|
"""Given a question and Python file, execute the file to answer the question. |
|
|
|
|
|
Args: |
|
|
question (str): Question to answer |
|
|
file_path (str): The Python file path |
|
|
|
|
|
Returns: |
|
|
str: Answer to the question |
|
|
|
|
|
Raises: |
|
|
RuntimeError: If processing fails""" |
|
|
try: |
|
|
client = genai.Client(api_key=os.environ["GEMINI_API_KEY"]) |
|
|
|
|
|
file = client.files.upload(file=file_path) |
|
|
|
|
|
response = client.models.generate_content( |
|
|
model=CODE_EXECUTION_MODEL, |
|
|
contents=[file, question], |
|
|
config=types.GenerateContentConfig( |
|
|
tools=[types.Tool(code_execution=types.ToolCodeExecution)] |
|
|
), |
|
|
) |
|
|
|
|
|
for part in response.candidates[0].content.parts: |
|
|
if part.code_execution_result is not None: |
|
|
return part.code_execution_result.output |
|
|
except Exception as e: |
|
|
raise RuntimeError(f"Processing failed: {str(e)}") |
