agent_mistral.py

import logging
import os

import tempfile

from langchain.agents import AgentExecutor, create_tool_calling_agent
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.tools import Tool
from langchain_mistralai import ChatMistralAI
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 = ChatMistralAI(
                # model="models/gemini-2.5-flash-preview-04-17",
                model="mistral-small-latest",
                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