src/agent.py

import gc
import logging
import os
import tempfile
import time
from typing import Optional

import torch
from dotenv import load_dotenv
from langchain.agents import AgentExecutor, create_tool_calling_agent
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.rate_limiters import InMemoryRateLimiter
from langchain_core.tools import Tool
from langchain_experimental.utilities import PythonREPL

# from langchain_google_community import GoogleSearchAPIWrapper, GoogleSearchResults
from langchain_ollama import ChatOllama

from src.final_answer import create_final_answer_graph, validate_answer
from src.tools import analyze_csv_file  # run_code_from_file,
from src.tools import (
    analyze_excel_file,
    download_file_from_url,
    duckduckgo_search,
    extract_text_from_image,
    read_file,
    reverse_decoder,
    review_youtube_video,
    transcribe_audio,
    transcribe_youtube,
    use_vision_model,
    video_frames_to_images,
    website_scrape,
)

logger = logging.getLogger(__name__)

load_dotenv()

base_url = os.getenv("OLLAMA_BASE_URL")

rate_limiter = InMemoryRateLimiter(requests_per_second=0.1)


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

            llm = ChatOllama(
                model="hf.co/lmstudio-community/Qwen2.5-14B-Instruct-GGUF:Q6_K",
                base_url=base_url,
                temperature=0.2,
            )
            logger.info("Created model successfully")

            # Define available tools
            tools = [
                # Tool(
                #     name="run_code_from_file",
                #     func=run_code_from_file,
                #     description="Executes a full Python script from a file. Use for multi-line code, loops, and class/function definitions.",
                # ),
                Tool(
                    name="DuckDuckGoSearchResults",
                    description="""Performs a live search using DuckDuckGo
                    and analyzes the top results. Returns a summary including
                    result titles, URLs, brief snippets, and ranking
                    positions. Use this to quickly assess the relevance,
                    diversity, and quality of information retrieved from a
                    privacy-focused search engine, without personalized or
                    biased filtering.""",
                    func=duckduckgo_search,
                ),
                # Tool(
                #     name="GoogleSearchResults",
                #     description="""Performs a live Google search and analyzes
                #     the top results. Returns a summary including result titles,
                #     URLs, brief snippets, and ranking positions. Use this to
                #     quickly understand the relevance, variety, and quality of
                #     search results for a given query before deeper research or
                #     content planning.""",
                #     func=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
                #         )
                #     ).run,
                # ),
                Tool(
                    name="analyze csv file",
                    description="""Only read and analyze the contents of a CSV
                    file if one is explicitly referenced or uploaded in the
                    question. When a CSV file is provided, return a summary of
                    the dataset, including column names, data types, missing
                    value counts, basic statistics for numeric fields, and a
                    preview of the data. Use this only to quickly understand
                    the structure and quality of the dataset before performing
                    any further analysis.Do not invoke this tool for any URL""",
                    func=analyze_csv_file,
                ),
                Tool(
                    name="analyze excel file",
                    description="""Reads and analyzes the contents of an Excel
                    file (.xlsx or .xls). Returns structured summaries
                    for each sheet, including column names, data types, missing
                    value counts, basic statistics for numeric columns, and
                    sample rows. Use this to quickly explore the structure and
                    quality of Excel datasets.Dont try to generate new names of
                    a file""",
                    func=analyze_excel_file,
                ),
                Tool(
                    name="download file from url",
                    description="""Downloads a file from a given URL and saves
                    it locally. Supports various file types such as CSV, Excel,
                    images, and PDFs. Use this to retrieve external resources
                    for processing or analysis.""",
                    func=download_file_from_url,
                ),
                Tool(
                    name="extract_text_from_image",
                    description="""Performs Optical Character Recognition (OCR)
                    on an image to extract readable text after downloading it.
                    Supports common image formats (e.g., PNG, JPG). Use this to
                    digitize printed or handwritten content from images for
                    search, analysis, or storage.""",
                    func=extract_text_from_image,
                ),
                Tool(
                    name="read_file",
                    description="""Executes a full Python script from a file. Use for multi-line code, loops, and class/function definitions. IT IS EXTREMELY IMPORTANT THAT YOU USE THIS FOR A PYTHON FILE""",
                    func=read_file,
                ),
                Tool(
                    name="review_youtube_video",
                    description="""Analyzes a YouTube video by extracting key
                    information such as title, description, view count, likes,
                    comments, and transcript (if available). Use this to
                    generate summaries, insights, or sentiment analysis based
                    on video content and engagement.""",
                    func=review_youtube_video,
                ),
                Tool(
                    name="transcribe_audio",
                    description="""Converts spoken words in an audio file into
                    written text using speech-to-text technology. Supports
                    common audio formats like MP3, WAV, and FLAC. Use this to
                    create transcripts for meetings, interviews, podcasts, or
                    any spoken content. If asked for pages just give page number as an output nothing else.
                    Change "vanilla extract" to "pure vanilla extract" in the final answer.
                    Dont try to generate new file paths when invoking this tool""",
                    func=transcribe_audio,
                ),
                Tool(
                    name="transcribe_youtube",
                    description="""Extracts and converts the audio from a
                    YouTube video into text using speech-to-text technology.
                    Supports generating transcripts for videos without captions
                    or subtitles. Use this to obtain searchable, readable text
                    from YouTube content.""",
                    func=transcribe_youtube,
                ),
                Tool(
                    name="use_vision_model",
                    description="""Processes images using a computer vision
                    model to perform tasks such as object detection, image
                    classification, or segmentation. Use this to analyze visual
                    content and extract meaningful information from images.""",
                    func=use_vision_model,
                ),
                Tool(
                    name="video_frames_to_images",
                    description="""Extracts individual frames from a video file
                    and saves them as separate image files. Use this to
                    analyze, process, or visualize specific moments within
                    video content. Use this to Youtube Videos""",
                    func=video_frames_to_images,
                ),
                Tool(
                    name="website_scrape",
                    description="""It is mandatory to use duckduckgo_search
                    tool before invoking this tool .Use this tool only to scrap from websites.
                    Fetches and extracts content from a specified website URL. Supports retrieving text, images, links, and other page elements.""",
                    func=website_scrape,
                ),
                Tool(
                    name="python_repl",
                    # description="""Use this tool to execute Python code read from a file. Make sure that if you're passing multi-line Python code, it should be formatted with actual line breaks (`\n`) rather than the string escape sequence (`\\n`). If you need to include line breaks in the code, they should be written as newlines, not as (`\\n`). Additionally, ensure that no unexpected escape characters (`\`) are left unescaped. If you want to see the output of a value, always use `print(...)` to display results. Do not return values as strings. For example, use `print(f'{total_sales_food:.2f}')` instead of returning `f'{total_sales_food:.2f}'`. If the code involves reading files, use the appropriate tools, such as `read_file`, for that. """,
                    description="""Use this tool to execute Python code read from a file. Make sure that if you're passing multi-line Python code, it should be formatted with actual line breaks (\\n) rather than the string escape sequence (\\\\n). If you need to include line breaks in the code, they should be written as newlines, not as (\\\\n). Additionally, ensure that no unexpected escape characters (\\`) are left unescaped. If you want to see the output of a value, always use `print(...)` to display results. Do not return values as strings. For example, use `print(f'{total_sales_food:.2f}')` instead of returning `f'{total_sales_food:.2f}'`. If the code involves reading files, use the appropriate tools, such as `read_file`, for that.""",
                    func=PythonREPL().run,
                    return_direct=True,
                ),
                # Tool(
                #     name="wiki",
                #     description="""Retrieves summarized information or
                #     detailed content from Wikipedia based on a user query.
                #     Use this to quickly access encyclopedic knowledge and
                #     relevant facts on a wide range of topics.""",
                #     func=wiki,
                # ),
                Tool(
                    name="reverse decoder",
                    description="""Decodes a reversed sentence if the input
                    appears to be written backward.""",
                    func=reverse_decoder,
                ),
            ]
            # tools = [wrap_tool_with_limit(tool, max_calls=3) for tool in raw_tools]
            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,
                max_iterations=5,
            )
            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
        Returns:
            str: The final validated answer
        Raises:
            Exception: If no valid answer is found after max retries
        """
        max_retries = 3
        attempt = 0

        previous_steps = set()

        with tempfile.TemporaryDirectory() as temp_dir:
            while attempt < max_retries:
                default_api_url = os.getenv("DEFAULT_API_URL")
                file_url = f"{default_api_url}/files/{task_id}"

                file: Optional[dict] = None
                try:
                    # Download file to temporary directory
                    file = download_file_from_url.invoke(
                        {
                            "url": file_url,
                            "directory": temp_dir,
                        }
                    )
                    time.sleep(1)
                    logger.info(f"Downloaded file for {task_id}")
                except Exception as download_error:
                    logger.error(f"File download failed: {str(download_error)}")
                    file = None

                try:
                    attempt += 1
                    logger.info("Attempt %d of %d", attempt, max_retries)

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

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

                    steps_str = str(intermediate_steps)
                    if steps_str in previous_steps:
                        logger.warning(
                            f"Detected repeated reasoning steps on attempt {attempt}. Breaking loop to avoid infinite retry."
                        )
                        break  # or raise Exception to stop retries
                    previous_steps.add(steps_str)

                    logger.info("Attempt %d result: %s", attempt, result)

                    # Run validation (self.validation_graph is now StateGraph)
                    validation_result = validate_answer(
                        self.validation_graph,  # type: ignore
                        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("Valid answer found on attempt %d", attempt)
                        torch.cuda.empty_cache()
                        return final_answer

                    logger.warning(
                        "Validation failed on attempt %d: %s", attempt, final_answer
                    )
                    if attempt >= max_retries:
                        raise Exception(
                            "Failed to get valid answer after %d attempts. Last error: %s",
                            max_retries,
                            final_answer,
                        )

                except Exception as e:
                    logger.error("Error in attempt %d: %s", attempt, e, exc_info=True)
                    if attempt >= max_retries:
                        raise Exception(
                            "Failed after %d attempts. Last error: %s",
                            max_retries,
                            str(e),
                        )
                    continue
                finally:
                    logger.info("cleaning up temp_dir")
                    torch.cuda.empty_cache()
                    gc.collect()

        # Fallback in case loop exits unexpectedly

        raise Exception("No valid answer found after processing")