Refactor agent.py to utilize LiteLLMModel for enhanced image processing capabilities. Implement improved file handling for various file types, including images and Excel files, allowing the agent to process content directly from provided files. Update task descriptions for clarity and error handling for file reading issues.

agent.py

@@ -1,20 +1,26 @@
-from smolagents import CodeAgent, InferenceClientModel, WebSearchTool
+import json
+from smolagents import CodeAgent, LiteLLMModel
+from tools import perplexity_search_tool
 import os
 from dotenv import load_dotenv
+from PIL import Image # Added for loading images
+import pandas as pd # Added for reading Excel
+import io # Added for StringIO
 
 load_dotenv()
 
-# Configure the underlying LLM model for the agent
-smol_model = InferenceClientModel(token=os.getenv("HF_TOKEN")) # You can choose to not pass any model_id to InferenceClientModel to use a default model
-# you can also specify a particular provider e.g. provider="together" or provider="sambanova"
+# Configure the underlying LLM model for the agent - using a vision-capable model
+smol_model = LiteLLMModel(model_id="gemini/gemini-2.5-flash-preview-04-17", api_key=os.getenv("GEMINI_API_KEY"))
 
-# Instantiate the CodeAgent with authorized imports for file operations
-# WebSearchTool is kept as an example, you might want to add/remove tools as needed.
+# Instantiate the CodeAgent with authorized imports
+# Removed base64, litellm. Kept PIL for now, though framework might handle image display.
 smol_code_agent_instance = CodeAgent(
-    tools=[WebSearchTool()],
+    tools=[perplexity_search_tool],
     model=smol_model,
-    add_base_tools=True,
-    additional_authorized_imports=["os", "io"] # Authorize os and io for file access
+    add_base_tools=False, 
+    # Authorized imports for the agent's generated code if it needs to manipulate images/files.
+    # PIL might be needed if agent code directly interacts with image objects.
+    additional_authorized_imports=["os", "io", "requests", "pandas", "PIL", "pydub", "subprocess"] 
 )
 
 # This is the agent class that app.py will instantiate and use.
@@ -22,55 +28,102 @@ smol_code_agent_instance = CodeAgent(
 class BasicAgent:
     def __init__(self):
         self.smol_agent = smol_code_agent_instance
-        print("BasicAgent (using smolagents.CodeAgent) initialized.")
+        print("BasicAgent (using smolagents.CodeAgent with direct image passing and content pre-loading) initialized.")
 
     def __call__(self, question: str, file_path: str | None = None) -> str:
         print(f"BasicAgent received question (first 100 chars): {question[:100]}...")
         
-        task_description = question
-        run_additional_args = {}
+        task_description = question # Default to original question
+        pil_images = []
+        run_kwargs = {} # Will hold {'images': pil_images} or {'additional_args': ...} or be empty
 
-        if file_path:
+        if file_path and os.path.exists(file_path):
             print(f"BasicAgent received file_path: {file_path}")
-            # Augment the task description to inform the agent about the file
-            # and how to access it via additional_args.
-            task_description += (
-                f"\n\nA file relevant to this task has been provided. "
-                f"You can access this file's content using standard Python file operations (e.g., open(), read()). "
-                f"The path to this file is available in the 'additional_args' dictionary, passed to the run method, under the key 'task_file'. "
-                f"The value is: '{file_path}'."
-            )
-            run_additional_args["additional_args"] = {"task_file": file_path}
-            print(f"Passing to smolagent: task='{task_description[:150]}...', additional_args={run_additional_args['additional_args']}")
+            file_extension = os.path.splitext(file_path)[1].lower()
+            try:
+                # Attempt to load as image first
+                img = Image.open(file_path).convert("RGB")
+                pil_images.append(img)
+                run_kwargs["images"] = pil_images
+                
+                task_description = f"""Please answer the following question: "{question}"
+
+An image has been provided to help answer this question. The image is directly available to you for analysis. Focus on interpreting the visual content of the image in relation to the question."""
+                print(f"Passing PIL image to smolagent. Task: '{task_description[:150]}...'")
+
+            except IOError: # Not a standard image or unreadable by PIL
+                print(f"File {file_path} is not a PIL-compatible image. Attempting other file type processing.")
+                try:
+                    if file_extension == ".xlsx":
+                        df = pd.read_excel(file_path)
+                        # Convert dataframe to CSV string to pass as content
+                        excel_content_str = df.to_csv(index=False)
+                        run_kwargs["additional_args"] = {"task_file_content": excel_content_str, "file_type": "excel_csv_string"}
+                        task_description = f"""Please answer the following question: "{question}"
+
+An Excel file's content (converted to CSV format) has been provided in 'additional_args' under the key 'task_file_content'.
+You can process this CSV string, for example, by using pandas: pandas.read_csv(io.StringIO(task_file_content)).
+The original file name was: {os.path.basename(file_path)}"""
+                        print(f"Passing Excel content (as CSV string) via additional_args. Task: '{task_description[:200]}...'")
+                    elif file_extension == ".mp3": # Example for mp3, passing as bytes
+                        with open(file_path, "rb") as f:
+                            file_bytes = f.read()
+                        run_kwargs["additional_args"] = {"task_file_content": file_bytes, "file_type": "mp3_bytes"}
+                        task_description = f"""Please answer the following question: "{question}"
+
+An MP3 file's content (as bytes) has been provided in 'additional_args' under the key 'task_file_content'.
+The original file name was: {os.path.basename(file_path)}.
+Note: You will need appropriate tools/capabilities to process raw audio bytes for tasks like transcription."""
+                        print(f"Passing MP3 file content (as bytes) via additional_args. Task: '{task_description[:200]}...'")
+                    else: # Generic binary/text file
+                        with open(file_path, "rb") as f: # Read as binary by default
+                            file_bytes = f.read()
+                        run_kwargs["additional_args"] = {"task_file_content": file_bytes, "file_type": "generic_bytes"}
+                        # Try to decode as text for the task description if it's small, otherwise indicate binary
+                        try:
+                            preview_text = file_bytes.decode('utf-8', errors='ignore')[:200]
+                            file_content_description = f"The file content (as bytes, starting with: '{preview_text}...') has been provided"
+                        except:
+                            file_content_description = "The file content (as binary bytes) has been provided"
+
+                        task_description = f"""Please answer the following question: "{question}"
+
+A file's content has been provided in 'additional_args' under the key 'task_file_content' (type: {run_kwargs['additional_args']['file_type']}).
+{file_content_description}.
+The original file name was: {os.path.basename(file_path)}."""
+                        print(f"Passing generic file content (as bytes) via additional_args. Task: '{task_description[:200]}...'")
+
+                except Exception as e_read:
+                    print(f"Error reading content from {file_path} after PIL failure: {e_read}")
+                    task_description = f"""Regarding the question: "{question}"
+A file was provided at '{file_path}', but an error occurred while trying to read its content: {e_read}.
+Please attempt to answer based on the text query alone. If the file was essential, state it could not be processed."""
+                    # run_kwargs remains as is, or cleared if preferred. Let's clear additional_args if read fails.
+                    if "additional_args" in run_kwargs: del run_kwargs["additional_args"]
+
+            except Exception as e_pil: # Catch other PIL errors if any beyond IOError
+                print(f"An unexpected PIL error occurred while processing file {file_path}: {e_pil}")
+                task_description = f"""Regarding the question: "{question}"
+A file was initially provided at '{file_path}', but an error occurred during its image processing: {e_pil}.
+Please attempt to answer the question based on the text query alone if possible. If the file was essential, please state that the file could not be processed."""
         else:
-            print("BasicAgent received no file_path.")
-            print(f"Passing to smolagent: task='{task_description[:150]}...'")
+            if file_path: # File path provided but does not exist
+                print(f"BasicAgent received file_path: {file_path}, but it does not exist. Proceeding without file content.")
+                task_description = f"""Regarding the question: "{question}"
+A file was expected at '{file_path}', but it was not found.
+Please answer the question based on the text query alone if possible. If the file was essential, please state that the file was missing."""
+            else: # No file path provided at all
+                print("BasicAgent received no file_path. Task is based on question only.")
+        
+        print(f"Final task description for smolagent (first 300 chars): {task_description[:300]}...")
 
         try:
-            # The CodeAgent's .run() method returns the final answer.
-            # Pass additional_args if they exist.
-            if run_additional_args:
-                answer = self.smol_agent.run(task_description, **run_additional_args)
-            else:
-                answer = self.smol_agent.run(task_description)
-                
-            print(f"smolagents.CodeAgent returned: {str(answer)[:200]}...") # Log a snippet of the answer
-            return str(answer)  # Ensure the output is a string
+            answer = self.smol_agent.run(task_description, **run_kwargs)
+            print(f"smolagents.CodeAgent returned: {str(answer)[:200]}...")
+            return str(answer)
         except Exception as e:
             print(f"Error during smolagents.CodeAgent run: {e}")
-            # It's often good to return a more informative error message
-            return f"AGENT_ERROR: An error occurred while processing with CodeAgent: {str(e)}"
-
-# Example test run (optional, can be commented out for deployment)
-# if __name__ == '__main__':
-#     print("Testing CodeAgent with a simple task...")
-#     test_agent = BasicAgent()
-#     # Test without file
-#     # response = test_agent("What is the capital of France?")
-#     # print(f"Test response (no file): {response}")
-#     # To test with a file, you'd create a dummy file and pass its path:
-#     # with open("dummy_test_file.txt", "w") as f:
-#     #     f.write("This is a test file for the agent.")
-#     # response_with_file = test_agent("Summarize the content of the provided file.", file_path="dummy_test_file.txt")
-#     # print(f"Test response (with file): {response_with_file}")
-#     # os.remove("dummy_test_file.txt")
+            # Be more specific about the error type if possible
+            return f"AGENT_ERROR: An error occurred in CodeAgent ({type(e).__name__}): {str(e)}"
+        
+