retriever.py

import requests
from langchain_core.tools import Tool, tool
from typing import Optional
from PIL import Image
from io import BytesIO
import base64
import mimetypes
import urllib.error
import urllib.request
from urllib.parse import urlparse
import tempfile
import os
import pytesseract
import uuid
import pandas as pd

local_filename = "downloaded_data.xlsx"


# GAIA File Retriver for task_id
def get_file(task_id: str) -> str:
    """fetches the file from GAIA API for given task_id"""
    response = requests.get(
        f"https://agents-course-unit4-scoring.hf.space/files/{task_id}", timeout=120
    )
    return response.content


file_retriever_tool = Tool(
    name="get_file",
    description="fetches the file from Agent API for given task_id",
    func=get_file,
)


@tool
def download_file_from_url(url: str, filename: Optional[str] = None) -> str:
    """
    Download a file from a URL and save it to a temporary location.
    Args:
        url (str): the URL of the file to download.
        filename (str, optional): the name of the file. If not provided, a random name file will be created.
    """
    try:
        # Parse URL to get filename if not provided
        if not filename:
            path = urlparse(url).path
            filename = os.path.basename(path)
            if not filename:
                filename = f"downloaded_{uuid.uuid4().hex[:8]}"

        # Create temporary file
        temp_dir = tempfile.gettempdir()
        filepath = os.path.join(temp_dir, filename)

        # Download the file
        response = requests.get(url, stream=True, timeout=120)
        response.raise_for_status()

        # Save the file
        with open(filepath, "wb") as f:
            for chunk in response.iter_content(chunk_size=8192):
                f.write(chunk)

        return f"File downloaded to {filepath}. You can read this file to process its contents."
    except Exception as e:
        return f"Error downloading file: {str(e)}"


# text file retriever
@tool
def fetch_text_from_url(url: str) -> str:
    """Fetch the document from a URL"""
    req = urllib.request.Request(
        url,
        headers={"User-Agent": "Mozilla/5.0 (compatible; quickstart-research/1.0)"},
    )
    try:
        with urllib.request.urlopen(req, timeout=120) as resp:
            raw = resp.read()
    except urllib.error.URLError as e:
        return f"Fetch failed: {e}"
    text = raw.decode("utf-8", errors="replace")
    return text


# YT Video Frame Retriever


# Image Retriever
@tool
def image_decoder(img_url: str) -> str:
    """downaload image from url and generate base64 decoded image url to read in local
    Args:
        img_url: url of image to download
    """
    try:
        buffered = BytesIO()
        headers = {
            "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/114.0.0.0 Safari/537.36"
        }
        response = requests.get(img_url, headers, timeout=120)
        mime, _ = mimetypes.guess_type(img_url)
        image = Image.open(BytesIO(response.content))
        img_mime = (
            mime if mime else f"image/{image.format}" if image.format else "image/png"
        )
        img_format = image.format if image.format else "PNG"
        image.save(buffered, format=img_format)
        img_base64 = base64.b64encode(buffered.getvalue()).decode("utf-8")
        buffered.close()
        return f"data:{img_mime};base64,{img_base64}"
    except Exception as e:
        raise Exception(f"Error fetching image details from {img_url}: {str(e)}")


# @tool
# def extract_image_info(image_url: str) -> dict:
#     """
#     Extract text from an image file using a multimodal model.

#     Args:
#         image_url: The URL of the image to analyze

#     Returns:
#         A dictionary containing the description, table, and table content
#     """
#     try:
#         img_url = image_decoder(image_url)
#         all_text = ""
#         message = [
#             HumanMessage(
#                 content=[
#                     {
#                         "type": "text",
#                         "text": (
#                             "Extract all the text from this image. "
#                             "Return only the extracted text, no explanations."
#                         ),
#                     },
#                     {
#                         "type": "image_url",
#                         "image_url": {"url": img_url},
#                     },
#                 ]
#             )
#         ]

#         response = vision_llm.invoke(message)
#         all_text += response.content + "\n\n"
#         return all_text.strip()
#     except Exception as e:
#         return f"Error extracting image details from {image_url}: {str(e)}"


@tool
def extract_image_info(image_url: str) -> str:
    """
    Extract text from an image file using a OCR library pytesseract (if available).

    Args:
        image_url: The URL of the image to analyze

    Returns:
        A text containing the description of image
    """
    # 1. Fetch the image from the URL
    response = requests.get(image_url, timeout=120)

    # 2. Open the image from the downloaded bytes
    img = Image.open(BytesIO(response.content))

    # 3. Perform OCR
    text = pytesseract.image_to_string(img)
    return text


# Excel File Reader
@tool
def excel_data_retriever(file_path: str):
    """Download and read the excel file using given excel file url.

    Args:
        file_path: path of excel file to process
    """
    try:
        df = pd.read_excel(file_path)
        df_json = df.to_json()
        result = (
            f"Excel file loaded with {len(df)} rows and {len(df.columns)} columns.\n"
        )
        result += f"Detailed Info of records as given in JSON format: \n Here the JSON is created by columns, so each key in JOSN represent column names with its values\n {df_json}"
        print(df.to_json())
        return result
    except Exception as e:
        return f"Error reading excel file {file_path}: {str(e)}"


# CSV File Reader
@tool
def csv_data_retriever(file_path: str):
    """Download and read the csv file using given csv file url.

    Args:
        file_path: path of csv file to process
    """
    try:
        df = pd.read_csv(file_path)
        df_json = df.to_json()
        result = f"CSV file loaded with {len(df)} rows and {len(df.columns)} columns.\n"
        result += f"Detailed Info of records as given in JSON format: \n Here the JSON is created by columns, so each key in JOSN represent column names with its values\n {df_json}"
        print(df.to_json())
        return result
    except Exception as e:
        return f"Error reading CSV file {file_path}: {str(e)}"