tools/files.py

from langchain_community.agent_toolkits import FileManagementToolkit
from config import CACHE_DIR
from langchain_core.tools import tool
import pandas as pd

__all__ = [
    'save_file',
    'read_file',
    'analyze_csv',
    'extract_text_from_image'
]

@tool
def save_file(filename: str, content: str) -> bool:
    """
    Saves content to a file.

    Args:
        filename (str): The name of the file in the cache directory.
        content (str): The content to save.

    Returns:
        bool: True if the file was saved successfully, False otherwise.
    """
    try:
        file_path = CACHE_DIR / filename
        with open(file_path, "w") as f:
            f.write(content)
        return True
    except Exception as e:
        print(f"Error saving file {file_path}: {e}")
        return False
    
@tool
def read_file(filename: str) -> str:
    """
    Reads content from a file.

    Args:
        filename (str): The name of the file in the cache directory.

    Returns:
        str: The content of the file.
    """
    try:
        file_path = CACHE_DIR / filename
        with open(file_path, "r") as f:
            return f.read()
    except Exception as e:
        return (f"Error reading file {file_path}: {e}")
    
@tool
def analyze_csv(filename: str) -> str:
    """
    Analyzes a CSV file using pandas, returning statistics about the file.
    
    Args:
        filename (str): The name of the CSV file in the cache directory.
    
    Returns:
        str: The analysis of the file.
    """
    df = pd.read_csv(CACHE_DIR / filename)
    if df is not None:
        # Perform analysis on the DataFrame
        result = f"CSV file: {len(df)} rows; {len(df.columns)} columns.\n"
        result += f"Columns: {', '.join(df.columns)}\n\n"

        result += "Summary:\n"
        result += str(df.describe())

    return result

@tool
def extract_text_from_image(filename: str) -> str:
    """
    Extracts text from an image file using OCR.
    
    Args:
        filename (str): The name of the image file in the cache directory.
    
    Returns:
        str: The extracted text from the image.
    """
    from PIL import Image
    import pytesseract

    try:
        image_path = CACHE_DIR / filename
        image = Image.open(image_path)
        text = pytesseract.image_to_string(image)
        return text
    except Exception as e:
        return f"Error extracting text from image {filename}: {e}"