Pacama95/chatbot_agent

1from pydantic import BaseModel, Field2from langchain.tools import BaseTool3 4from typing import Type, Any5 6import pandas as pd7import requests8import io9import PyPDF210 11class DocumentReaderInput(BaseModel):12    file_url: str = Field(description='The URL to the file.')13    file_format: str = Field(description="The file format (one of: xlsx, csv, txt, pdf).")14 15class DocumentReader(BaseTool):16 17    name: str = "document_reader_tool"18    description: str = (19        "Given the URL to a file and its format, return the file content."20        "Use this tool when you need to read the content of a file in a document format like PDF, xlsx, csv, txt..."21        "This is not an OCR tool, so scanned documents cannot be read using this tool."22    )23    args_schema: Type[BaseModel] = DocumentReaderInput24 25    def __init__(self, **kwargs: Any) -> None:26        super().__init__()27 28    def _run(self, file_url: str, file_format: str) -> str:29        """30            Given the URL to a file and its format, return the file content.31            Args:32                file_url (str): The URL to the file.33                file_format (str): The file format (one of: xlsx, csv, txt, pdf).34 35            Returns:36                str: The file content37        """38        file_format = file_format.lower()39 40        if file_format == 'xlsx':41            return str(read_excel_file(file_url=file_url))42        elif file_format == 'csv':43            return str(read_csv_file(file_url=file_url))44        elif file_format == 'txt':45            return read_txt_file(file_url=file_url)46        elif file_format == 'pdf':47            return read_pdf_file(file_url=file_url)48        else:49            raise ValueError(f"Unsupported file format: {file_format}")50    51 52 53def read_txt_file(file_url: str):54    try:55        response = requests.get(file_url)56        response.raise_for_status()  # Raises HTTPError for bad responses (4xx or 5xx)57    58        59        return response.text60 61    except requests.exceptions.RequestException as e:62        raise RuntimeError(f"Error fetching the file: {e}")63    64def read_excel_file(file_url: str):65    try:66        response = requests.get(file_url)67        response.raise_for_status()68 69        excel_data = pd.read_excel(io.BytesIO(response.content), sheet_name=None)70        return {sheet: df.head().to_string() for sheet, df in excel_data.items()}71 72    except requests.exceptions.RequestException as e:73        raise RuntimeError(f"Error fetching the Excel file: {e}")74 75def read_csv_file(file_url: str, max_rows: int = 100):76    try:77        response = requests.get(file_url)78        response.raise_for_status()79 80        df = pd.read_csv(io.StringIO(response.text))81 82        if max_rows:83            df = df.head(max_rows)84 85        return df.to_string()86 87    except requests.exceptions.RequestException as e:88        raise RuntimeError(f"Error fetching the CSV file: {e}")89 90def read_pdf_file(file_url: str):91    try:92        response = requests.get(file_url)93        response.raise_for_status()94 95        with io.BytesIO(response.content) as file_stream:96            reader = PyPDF2.PdfReader(file_stream)97            text = ''98            for page in reader.pages:99                text += page.extract_text() or ''100 101            return text.strip()102 103    except requests.exceptions.RequestException as e:104        raise RuntimeError(f"Error fetching the PDF file: {e}")105 106