Enhance file handling tools: save_file, read_file, analyze_csv, and extract_text_from_image functions

agent/agent.py

@@ -28,8 +28,8 @@ class OracleBot:
         """
         # Enhance question with file context if available
         if file_path and os.path.exists(file_path):
-            question = f"{question}\n\nNote: There is an associated file at: {file_path}\nYou can use the file management tools to read and analyze this file."
-        
+            question = f"{question}\n\nNote: There is an associated file named {os.path.basename(file_path)}\nYou can use the file management tools to read and analyze this file."
+
         messages = [HumanMessage(content=question)]
 
         for mode, chunk in self.graph.stream({"messages": messages}, config=self.config, stream_mode=["messages", "updates"]): # type: ignore

tools/files.py

@@ -1,4 +1,94 @@
 from langchain_community.agent_toolkits import FileManagementToolkit
 from config import CACHE_DIR
+from langchain_core.tools import tool
+import pandas as pd
 
-file_management_toolkit = FileManagementToolkit(root_dir=str(CACHE_DIR))
+__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}"