handle file input

Author: nicucalcea

augmenta/agent.py

@@ -4,7 +4,7 @@
 from pathlib import Path
 import yaml
 from pydantic import BaseModel, Field, create_model
-from pydantic_ai import Agent
+from pydantic_ai import Agent, BinaryContent
 import logfire
 from .tools.mcp import load_mcp_servers
 from .tools.search_web import search_web
@@ -119,18 +119,18 @@ def create_structure_class(yaml_file_path: Union[str, Path]) -> Type[BaseModel]:
 
         except (yaml.YAMLError, OSError) as e:
             raise ValueError(f"Failed to parse YAML: {e}")
-    
+      
     async def run(
         self,
-        prompt: str,
+        prompt: Union[str, List[Union[str, BinaryContent]]],
         response_format: Optional[Type[BaseModel]] = None,
         temperature: Optional[float] = None,
         system_prompt: Optional[str] = None
     ) -> Union[str, dict[str, Any], BaseModel]:
         """Run the agent to perform web research.
         
         Args:
-            prompt: The research query or task
+            prompt: The research query/task as a string or a list containing text and binary content
             response_format: Optional Pydantic model for structured output
             temperature: Optional override for model temperature
             system_prompt: Optional override for system prompt (defaults to self.system_prompt)

augmenta/augmenta.py

@@ -4,14 +4,15 @@
 import asyncio
 import pandas as pd
 from pathlib import Path
-from typing import Optional, Tuple, Dict, Any, Callable, Type, Union
+from typing import Optional, Tuple, Dict, Any, Callable, Type, Union, List
 from dataclasses import dataclass
 
 from augmenta.utils.prompt_formatter import format_examples, substitute_template_variables, build_complete_prompt
 from augmenta.agent import AugmentaAgent
 from augmenta.cache import CacheManager
 from augmenta.cache.process import setup_cache_handling, apply_cached_results
 from augmenta.config.read_config import load_config, get_config_values
+from augmenta.tools.file import load_file
 import logfire
 
 @dataclass
@@ -206,12 +207,34 @@ async def process_row(
     """
     try:
         index = row_data['index']
-        row = row_data['data']
-          # Build complete prompt with data from row
+        row = row_data['data']        # Build complete prompt with data from row
         prompt_user = build_complete_prompt(config, row)
+          # Get the file column name from config (if available)
+        file_col = config.get("file_col")
 
-        # Run prompt using the agent
-        response = await agent.run(prompt_user, response_format=response_format)
+        # Check if a file column is specified and the row contains a file path
+        file_path = None
+        if file_col and file_col in row:
+            file_path = row.get(file_col)
+            logfire.debug(f"Using file from column '{file_col}': {file_path}")
+        elif file_col:
+            logfire.debug(f"File column '{file_col}' specified in config but not found in row data")
+        else:
+            logfire.debug("No file column specified in config")
+            
+        try:
+            binary_content = load_file(file_path) if file_path is not None else None
+            if binary_content:
+                # If file exists, create a message list with prompt and binary content
+                message_contents = [prompt_user, binary_content]
+                response = await agent.run(message_contents, response_format=response_format)
+            else:
+                # If file doesn't exist or couldn't be loaded, just use the text prompt
+                response = await agent.run(prompt_user, response_format=response_format)
+        except Exception as e:
+            logfire.warning(f"Error loading file at row {index}: {str(e)}. Proceeding with text prompt only.")
+            # Fallback to text-only prompt if file handling fails
+            response = await agent.run(prompt_user, response_format=response_format)
 
         # Handle caching and progress tracking
         handle_result_tracking(

augmenta/config/read_config.py

@@ -92,15 +92,15 @@ def get_config_values(config: Dict[str, Any]) -> Dict[str, Any]:
 
     # Construct model ID with provider
     model_id = f"{model_config['provider']}:{model_config['name']}"
-    
-    # Extract commonly used values with defaults
+      # Extract commonly used values with defaults
     return {
         "model_id": model_id,
         "temperature": model_config.get("temperature", 0.0),
         "max_tokens": model_config.get("max_tokens"),
         "rate_limit": model_config.get("rate_limit"),
         "search_engine": search_config.get("engine"),
-        "search_results": search_config.get("results", 3)
+        "search_results": search_config.get("results", 3),
+        "file_col": config.get("file_col")
     }
 
 def load_config(config_path: Union[str, Path]) -> Dict[str, Any]:

augmenta/tools/__init__.py

@@ -1,9 +1,11 @@
-"""Tools and utilities for web interaction."""
+"""Tools and utilities for web interaction and file handling."""
 
 from .search_web import search_web
 from .visit_webpages import visit_webpages
+from .file import load_file
 
 __all__ = [
     'search_web',
-    'visit_webpages'
+    'visit_webpages',
+    'load_file'
 ]

augmenta/tools/file.py

@@ -0,0 +1,57 @@
+"""
+File handling utilities for loading binary content.
+"""
+import mimetypes
+from pathlib import Path
+from typing import Optional, Union, Any
+
+from pydantic_ai import BinaryContent
+
+
+def load_file(file_path: Union[str, Path, Any]) -> Optional[BinaryContent]:
+    """
+    Load a file as binary content with appropriate MIME type detection.
+    
+    Args:
+        file_path: Path to the file to load
+        
+    Returns:
+        BinaryContent object with binary data and media type if file exists,
+        None otherwise
+    """
+    # Skip if file path is None, empty, or 'NA' or not a string-like object
+    if file_path is None:
+        return None
+    
+    # Handle non-string types by converting to string first
+    try:
+        file_path_str = str(file_path).strip()
+        if not file_path_str or file_path_str.upper() == 'NA':
+            return None
+    except:
+        # If we can't convert to string, it's not a valid path
+        return None
+    
+    try:
+        path = Path(file_path_str)
+        
+        # Check if file exists
+        if not path.exists():
+            return None
+        
+        # Read binary content
+        file_binary = path.read_bytes()
+        
+        # Determine MIME type
+        media_type = mimetypes.guess_type(str(path))[0]
+        
+        # Default to application/octet-stream if type cannot be determined
+        if not media_type:
+            media_type = "application/octet-stream"
+        
+        # Create and return binary content
+        return BinaryContent(data=file_binary, media_type=media_type)
+    
+    except Exception as e:
+        print(f"Error loading file {file_path}: {str(e)}")
+        return None