final_assignment_agentic_course

Sleeping

final_assignment_agentic_course / agent_gemini.py

Update model

715a773 over 1 year ago

8.01 kB

import logging
import os

import tempfile

from langchain.agents import AgentExecutor, create_tool_calling_agent
from langchain_google_community import GoogleSearchResults
from langchain_google_community import GoogleSearchAPIWrapper
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.tools import Tool
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_experimental.utilities import PythonREPL

from tools_audio import transcribe_audio

from tools_doc import (
    analyze_csv_file,
    analyze_excel_file,
    download_file_from_url,
    extract_text_from_image,
    read_file,
)
from tools_video import (
    review_youtube_video,
    use_vision_model,
    transcribe_youtube,
    video_frames_to_images,
)

from tools_browser import website_scrape, web_search

from answers import create_final_answer_graph, validate_answer

logger = logging.getLogger(__name__)
DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"


class BasicAgent:
    def __init__(self):
        try:
            logger.info("Initializing BasicAgent")

            # Create the prompt template
            prompt = ChatPromptTemplate.from_messages(
                [
                    (
                        "system",
                        """You are a general AI assistant. I will ask you a question. Report your thoughts, and finish your answer with the following template: FINAL ANSWER: [YOUR FINAL ANSWER]. YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise. If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string.
                """,
                    ),
                    ("placeholder", "{chat_history}"),
                    ("human", "{input}"),
                    ("placeholder", "{agent_scratchpad}"),
                ]
            )
            logger.info("Created prompt template")

            # Initialize Gemini model
            logger.info("Creating Gemini model...")
            llm = ChatGoogleGenerativeAI(
                # model="models/gemini-2.5-flash-preview-04-17",
                model="gemini-2.5-flash-preview-04-17",
                google_api_key=os.getenv("GEMINI_KEY"),
                temperature=0.2,
            )
            logger.info("Created Gemini model successfully")

            # Define available tools
            tools = [
                # GoogleSearchResults(
                #     api_wrapper=GoogleSearchAPIWrapper(
                #         google_api_key=os.getenv("GOOGLE_SEARCH_API_KEY"),
                #         google_cse_id=os.getenv("GOOGLE_CSE_ID"),
                #         k=5,  # Number of results to return
                #     )
                # ),
                web_search,
                analyze_csv_file,
                analyze_excel_file,
                download_file_from_url,
                extract_text_from_image,
                read_file,
                review_youtube_video,
                transcribe_audio,
                transcribe_youtube,
                use_vision_model,
                video_frames_to_images,
                website_scrape,
                Tool(
                    name="python_repl",
                    description="A Python shell. Use this to execute python commands. Input # should be a valid python command. If you want to see the output of a value, # you should print it out with `print(...)`.",
                    func=PythonREPL().run,
                ),
            ]
            logger.info("Tools: %s", tools)

            # Create the agent
            agent = create_tool_calling_agent(llm, tools, prompt)
            logger.info("Created tool calling agent")

            # Create the agent executor
            self.agent_executor = AgentExecutor(
                agent=agent,
                tools=tools,
                return_intermediate_steps=True,
                verbose=True,
            )
            logger.info("Created agent executor")

            # Create the graph
            self.validation_graph = create_final_answer_graph()

        except Exception as e:
            logger.error("Error initializing agent: %s", e, exc_info=True)
            raise

    def __call__(self, question: str, task_id: str) -> str:
        """Execute the agent with the given question and optional file.
        Args:
            question (str): The question to answer
            task_id (str): The task ID to fetch the file
        """
        max_retries = 3
        attempt = 0

        # Create a temporary directory that will be automatically cleaned up
        print("HELLO")
        with tempfile.TemporaryDirectory() as temp_dir:
            while attempt < max_retries:
                # default_api_url = os.getenv("DEFAULT_API_URL")
                default_api_url = DEFAULT_API_URL
                file_url = f"{default_api_url}/files/{task_id}"

                try:
                    print("HELLO-A")
                    # Download file to temporary directory
                    file = download_file_from_url.invoke(
                        {
                            "url": file_url,
                            "directory": temp_dir,
                        }
                    )
                except Exception as e:
                    logger.error(f"Error downloading file: {e}")
                    file = None

                try:
                    print("HELLO-B")
                    attempt += 1
                    logger.info(f"Attempt {attempt} of {max_retries}")

                    # Prepare input with file information
                    if file and file.get("type") != "error":
                        input_data = {
                            "input": question
                            + f" [File: type={file.get('type', 'None')}, path={file.get('path', 'None')}]",
                        }
                    else:
                        input_data = {
                            "input": question,
                        }

                    # Run the agent to get the answer
                    result = self.agent_executor.invoke(input_data)
                    answer = result.get("output", "")

                    logger.info(f"Attempt {attempt} result: {result}")

                    # Run validation
                    validation_result = validate_answer(
                        self.validation_graph,
                        answer,
                        [result.get("intermediate_steps", [])],
                    )

                    valid_answer = validation_result.get("valid_answer", False)
                    final_answer = validation_result.get("final_answer", "")

                    if valid_answer:
                        logger.info(f"Valid answer found on attempt {attempt}")
                        return final_answer

                    logger.warning(
                        f"Validation failed on attempt {attempt}: {final_answer}"
                    )
                    if attempt >= max_retries:
                        raise Exception(
                            f"Failed to get valid answer after {max_retries} attempts. Last error: {final_answer}"
                        )

                except Exception as e:
                    logger.error(
                        f"Error in attempt {attempt}: {e}", exc_info=True
                    )
                    if attempt >= max_retries:
                        raise Exception(
                            f"Failed after {max_retries} attempts. Last error: {str(e)}"
                        )
                    continue

	import logging
	import os

	import tempfile

	from langchain.agents import AgentExecutor, create_tool_calling_agent
	from langchain_google_community import GoogleSearchResults
	from langchain_google_community import GoogleSearchAPIWrapper
	from langchain_core.prompts import ChatPromptTemplate
	from langchain_core.tools import Tool
	from langchain_google_genai import ChatGoogleGenerativeAI
	from langchain_experimental.utilities import PythonREPL

	from tools_audio import transcribe_audio

	from tools_doc import (
	analyze_csv_file,
	analyze_excel_file,
	download_file_from_url,
	extract_text_from_image,
	read_file,
	)
	from tools_video import (
	review_youtube_video,
	use_vision_model,
	transcribe_youtube,
	video_frames_to_images,
	)

	from tools_browser import website_scrape, web_search

	from answers import create_final_answer_graph, validate_answer

	logger = logging.getLogger(__name__)
	DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"


	class BasicAgent:
	def __init__(self):
	try:
	logger.info("Initializing BasicAgent")

	# Create the prompt template
	prompt = ChatPromptTemplate.from_messages(
	[
	(
	"system",
	"""You are a general AI assistant. I will ask you a question. Report your thoughts, and finish your answer with the following template: FINAL ANSWER: [YOUR FINAL ANSWER]. YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise. If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string.
	""",
	),
	("placeholder", "{chat_history}"),
	("human", "{input}"),
	("placeholder", "{agent_scratchpad}"),
	]
	)
	logger.info("Created prompt template")

	# Initialize Gemini model
	logger.info("Creating Gemini model...")
	llm = ChatGoogleGenerativeAI(
	# model="models/gemini-2.5-flash-preview-04-17",
	model="gemini-2.5-flash-preview-04-17",
	google_api_key=os.getenv("GEMINI_KEY"),
	temperature=0.2,
	)
	logger.info("Created Gemini model successfully")

	# Define available tools
	tools = [
	# GoogleSearchResults(
	# api_wrapper=GoogleSearchAPIWrapper(
	# google_api_key=os.getenv("GOOGLE_SEARCH_API_KEY"),
	# google_cse_id=os.getenv("GOOGLE_CSE_ID"),
	# k=5, # Number of results to return
	# )
	# ),
	web_search,
	analyze_csv_file,
	analyze_excel_file,
	download_file_from_url,
	extract_text_from_image,
	read_file,
	review_youtube_video,
	transcribe_audio,
	transcribe_youtube,
	use_vision_model,
	video_frames_to_images,
	website_scrape,
	Tool(
	name="python_repl",
	description="A Python shell. Use this to execute python commands. Input # should be a valid python command. If you want to see the output of a value, # you should print it out with `print(...)`.",
	func=PythonREPL().run,
	),
	]
	logger.info("Tools: %s", tools)

	# Create the agent
	agent = create_tool_calling_agent(llm, tools, prompt)
	logger.info("Created tool calling agent")

	# Create the agent executor
	self.agent_executor = AgentExecutor(
	agent=agent,
	tools=tools,
	return_intermediate_steps=True,
	verbose=True,
	)
	logger.info("Created agent executor")

	# Create the graph
	self.validation_graph = create_final_answer_graph()

	except Exception as e:
	logger.error("Error initializing agent: %s", e, exc_info=True)
	raise

	def __call__(self, question: str, task_id: str) -> str:
	"""Execute the agent with the given question and optional file.
	Args:
	question (str): The question to answer
	task_id (str): The task ID to fetch the file
	"""
	max_retries = 3
	attempt = 0

	# Create a temporary directory that will be automatically cleaned up
	print("HELLO")
	with tempfile.TemporaryDirectory() as temp_dir:
	while attempt < max_retries:
	# default_api_url = os.getenv("DEFAULT_API_URL")
	default_api_url = DEFAULT_API_URL
	file_url = f"{default_api_url}/files/{task_id}"

	try:
	print("HELLO-A")
	# Download file to temporary directory
	file = download_file_from_url.invoke(
	{
	"url": file_url,
	"directory": temp_dir,
	}
	)
	except Exception as e:
	logger.error(f"Error downloading file: {e}")
	file = None

	try:
	print("HELLO-B")
	attempt += 1
	logger.info(f"Attempt {attempt} of {max_retries}")

	# Prepare input with file information
	if file and file.get("type") != "error":
	input_data = {
	"input": question
	+ f" [File: type={file.get('type', 'None')}, path={file.get('path', 'None')}]",
	}
	else:
	input_data = {
	"input": question,
	}

	# Run the agent to get the answer
	result = self.agent_executor.invoke(input_data)
	answer = result.get("output", "")

	logger.info(f"Attempt {attempt} result: {result}")

	# Run validation
	validation_result = validate_answer(
	self.validation_graph,
	answer,
	[result.get("intermediate_steps", [])],
	)

	valid_answer = validation_result.get("valid_answer", False)
	final_answer = validation_result.get("final_answer", "")

	if valid_answer:
	logger.info(f"Valid answer found on attempt {attempt}")
	return final_answer

	logger.warning(
	f"Validation failed on attempt {attempt}: {final_answer}"
	)
	if attempt >= max_retries:
	raise Exception(
	f"Failed to get valid answer after {max_retries} attempts. Last error: {final_answer}"
	)

	except Exception as e:
	logger.error(
	f"Error in attempt {attempt}: {e}", exc_info=True
	)
	if attempt >= max_retries:
	raise Exception(
	f"Failed after {max_retries} attempts. Last error: {str(e)}"
	)
	continue