tools/ai_tools.py

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)}")

    @tool("Translation Tool")
    def translate_tool(text: str, target_lang: str) -> str:
        """Translate a given text into the target language.
        
        Args:
            text (str): The text to translate.
            target_lang (str): The target language (e.g., 'french', 'de', 'zh').
        
        Returns:
            str: Translated text.
        """
        try:
            client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
            
            prompt = f"Translate this text to {target_lang}: '{text}'"
            
            response = client.models.generate_content(
                model="gemini-2.5-flash-preview-04-17",
                contents=prompt
            )
            
            return response.text
        except Exception as e:
            raise RuntimeError(f"Translation failed: {str(e)}")
    
    @tool("Summarization Tool")
    def summarize_tool(text: str) -> str:
        """Summarize the given text input.
        
        Args:
            text (str): Long content to summarize.
        
        Returns:
            str: A concise summary.
        """
        try:
            client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
            
            prompt = f"Summarize this content: {text}"
            
            response = client.models.generate_content(
                model="gemini-2.5-flash-preview-04-17",
                contents=prompt
            )
            
            return response.text
        except Exception as e:
            raise RuntimeError(f"Summarization failed: {str(e)}")