agent.py

from smolagents import CodeAgent, ToolCallingAgent, LiteLLMModel, tool, Tool, load_tool, WebSearchTool, DuckDuckGoSearchTool #, WikipediaSearchTool
import asyncio
import os
import re
import pandas as pd
from typing import Optional

from token_bucket import Limiter, MemoryStorage
import yaml
from PIL import Image, ImageOps
import requests
from io import BytesIO
from markdownify import markdownify
import whisper
import time
import shutil
import traceback
from langchain_community.document_loaders import ArxivLoader


@tool
def search_arxiv(query: str) -> str:
    """Search Arxiv for a query and return maximum 3 result.
    
    Args:
        query: The search query.
     Returns:
        str: Formatted search results
    """
    search_docs = ArxivLoader(query=query, load_max_docs=3).load()
    formatted_search_docs = "\n\n---\n\n".join(
        [
            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content[:1000]}\n</Document>'
            for doc in search_docs
        ])
    return {"arxiv_results": formatted_search_docs}

from tenacity import retry, stop_after_attempt, wait_exponential, retry_if_exception
import requests


def is_429_error(exception):
    return isinstance(exception, requests.exceptions.HTTPError) and exception.response.status_code == 429

class VisitWebpageTool(Tool):
    name = "visit_webpage"
    description = "Visits a webpage at the given url and reads its content as a markdown string. Use this to browse webpages."
    inputs = {'url': {'type': 'string', 'description': 'The url of the webpage to visit.'}}
    output_type = "string"

    @retry(
        stop=stop_after_attempt(3),
        wait=wait_exponential(multiplier=1, min=4, max=10),
        retry=retry_if_exception(is_429_error)
    )
    def forward(self, url: str) -> str:
        try:
            response = requests.get(url, timeout=50)
            response.raise_for_status()
            markdown_content = markdownify(response.text).strip()
            markdown_content = re.sub(r"\n{3,}", "\n\n", markdown_content)
            #from smolagents.utils import truncate_content
            #return truncate_content(markdown_content, 10000)
            return markdown_content
        except requests.exceptions.HTTPError as e:
            if e.response.status_code == 429:
                raise  # Retry on 429
            return f"Error fetching the webpage: {str(e)}"
        except requests.exceptions.Timeout:
            return "The request timed out. Please try again later or check the URL."
        except requests.exceptions.RequestException as e:
            return f"Error fetching the webpage: {str(e)}"
        except Exception as e:
            return f"An unexpected error occurred: {str(e)}"

    def __init__(self, *args, **kwargs):
        self.is_initialized = False

class SpeechToTextTool(Tool):
    name = "speech_to_text"
    description = (
        "Converts an audio file to text using OpenAI Whisper."
    )
    inputs = {
        "audio_path": {"type": "string", "description": "Path to audio file (.mp3, .wav)"},
    }
    output_type = "string"

    def __init__(self):
        super().__init__()
        try:
            print("Loading whisper model 'base'...")
            self.model = whisper.load_model("base")
            print("Whisper model loaded successfully.")
        except Exception as e:
            print(f"Error loading whisper model: {e}")
            self.model = None

    def forward(self, audio_path: str) -> str:
        if not os.path.exists(audio_path):
            return f"Error: File not found at {audio_path}"
        if self.model is None:
            return "Error: Whisper model failed to load during initialization."
        try:
            print(f"Starting transcription for {audio_path}...")
            result = self.model.transcribe(audio_path)
            print(f"Transcription completed for {audio_path}.")
            return result.get("text", "")
        except Exception as e:
            return f"Error processing audio file: {str(e)}"

class ExcelReaderTool(Tool):
    name = "excel_reader"
    description = """
    Reads an Excel file (.xlsx, .xls) and returns its content as a CSV string.
    Use pandas.read_csv(StringIO(data)) to parse the output into a DataFrame.
    """
    inputs = {
        "excel_path": {
            "type": "string",
            "description": "The path to the Excel file to read",
        },
        "sheet_name": {
            "type": "string",
            "description": "The name of the sheet to read (optional, defaults to first sheet)",
            "nullable": True
        }
    }
    output_type = "string"
    
    def forward(self, excel_path: str, sheet_name: str = None) -> str:
        try:
            if not os.path.exists(excel_path):
                return f"Error: Excel file not found at {excel_path}"
            if sheet_name:
                df = pd.read_excel(excel_path, sheet_name=sheet_name)
            else:
                df = pd.read_excel(excel_path)
            return df.to_csv(index=False)
        except Exception as e:
            return f"Error reading Excel file: {str(e)}"


class PythonCodeReaderTool(Tool):
    name = "read_python_code"
    description = "Reads a Python (.py) file and returns its content as a string."
    inputs = {
        "file_path": {"type": "string", "description": "The path to the Python file to read"
}
    }
    output_type = "string"

    def forward(self, file_path: str) -> str:
        try:
            if not os.path.exists(file_path):
                return f"Error: Python file not found at {file_path}"
            with open(file_path, "r", encoding="utf-8") as file:
                content = file.read()
            return content
        except Exception as e:
            return f"Error reading Python file: {str(e)}"

from tenacity import retry, stop_after_attempt, wait_exponential, retry_if_exception_type

class RetryDuckDuckGoSearchTool(DuckDuckGoSearchTool):
    @retry(
        stop=stop_after_attempt(3),
        wait=wait_exponential(multiplier=1, min=4, max=10),
        retry=retry_if_exception_type(Exception)
    )

    def forward(self, query: str) -> str:
        return super().forward(query
)



##############################
#   MAG Agent
##############################

class MagAgent:
    def __init__(self, rate_limiter: Optional[Limiter] = None):
        """Initialize the MagAgent with search tools."""
        self.rate_limiter = rate_limiter

        print("Initializing MagAgent with search tools...")
        model = LiteLLMModel(
            model_id="gemini/gemini-2.0-flash",
            api_key=os.environ.get("GEMINI_KEY"),
            max_tokens=8192
        )

        self.imports = [
            "pandas",
            "numpy",
            "os",
            "requests",
            "tempfile",
            "datetime",
            "json",
            "time",
            "re",
            "openpyxl",
            "pathlib",
            "sys",
            "bs4",
            "arxiv",
            "whisper",
            "io"
        ]

        self.tools = [
#            RetryDuckDuckGoSearchTool(),

#            WikipediaSearchTool(),
            SpeechToTextTool(),
            ExcelReaderTool(),
#            VisitWebpageTool(),
            PythonCodeReaderTool(),
            search_arxiv,
        ]

        self.prompt_template = (
            """
            You are an advanced AI assistant specialized in solving complex, real-world tasks, requiring multi-step reasoning, factual accuracy, and use of external tools.

            Follow these principles:
            - Reason step-by-step. Think through the solution logically and plan your actions carefully before answering.
            - Validate information. Always verify facts when possible instead of guessing.
            - When processing external data (e.g., YouTube transcripts, web searches), expect potential issues like missing punctuation, inconsistent formatting, or conversational text.
            - When asked to process Excel files, use the `excel_reader` tool, which returns a pandas DataFrame.
            - When calculating sales, make sure you multiply volume on price per each product or category.
            - When asked to transcript YouTube video, try searching it in www.youtubetotranscript.com.
            - If the input is ambiguous, prioritize extracting key information relevant to the question.
            - Use code if needed. For calculations, parsing, or transformations, generate Python code and execute it. Be cautious, as some questions contain time-consuming tasks, so analyze the question and choose the most efficient solution.
            - Be precise and concise. The final answer must strictly match the required format with no extra commentary.
            - Use tools intelligently. If a question involves external information, structured data, images, or audio, call the appropriate tool to retrieve or process it.
            - If the question includes direct speech or quoted text (e.g., "Isn't that hot?"), treat it as a precise query and preserve the quoted structure in your response, including quotation marks for direct quotes (e.g., final_answer('"Extremely."')).
            - If asked about the name of a place or city, use the full complete name without abbreviations (e.g., use Saint Petersburg instead of St.Petersburg).
            - If asked to look up page numbers, make sure you don't mix them with problem or excercise numbers.
            - If you cannot retrieve or process data (e.g., due to blocked requests), retry after 15 seconds delay, try another tool (try wikipedia_search, then web_search, then search_arxiv). Otherwise, return a clear error message: "Unable to retrieve data. Search has failed."
            - Use `final_answer` to give the final answer.
            
            QUESTION: {question}
            
            {file_section}
            
            ANSWER:
            """
        )

        web_agent = ToolCallingAgent(
            tools=[
#                RetryDuckDuckGoSearchTool(),
#                WikipediaSearchTool(),
#                SpeechToTextTool(),
                WebSearchTool(),
                VisitWebpageTool(),
#                ExcelReaderTool(),
#                PythonCodeReaderTool(),
                search_arxiv,
            ],
            model=model,
            max_steps=15,
            name="web_search_agent",
            description="Runs web searches for you.",
        )        
        
        self.agent = CodeAgent(
            model=model,
            managed_agents=[web_agent],
            tools=self.tools,
            add_base_tools=True,
            additional_authorized_imports=self.imports,
            verbosity_level=2
,
            max_steps=10
        )
        print("MagAgent initialized.")

    async def __call__(self, question: str, file_path: Optional[str] = None) -> str:
        """Process a question asynchronously using the MagAgent."""
        print(f"MagAgent received question (first 50 chars): {question[:50]}... File path: {file_path}")
        try:
            if self.rate_limiter:
                while not self.rate_limiter.consume(1):
                    print(f"Rate limit reached. Waiting...")
                    await asyncio.sleep(4)
            # Conditionally include FILE: section only if file_path is provided
            file_section = f"FILE: {file_path}" if file_path else ""
            task = self.prompt_template.format(
                question=question,
                file_section=file_section
            )
            print(f"Calling agent.run...")
            response = await asyncio.to_thread(self.agent.run, task=task)
            print(f"Agent.run completed.")
            response = str(response)
            if not response:
                print(f"No answer found.")
                response = "No answer found."
            print(f"MagAgent response: {response[:50]}...")
            return response
        except Exception as e:
            error_msg = f"Error processing question: {str(e)}. Check API key or network connectivity."
            print(error_msg)
            return error_msg