vibe refactor

Author: nicucalcea

README.md

@@ -78,7 +78,6 @@ output_csv: path/to/processed_data.csv
 model:
   provider: openai
   name: gpt-4o-mini
-query_col: DonorName
 search:
   engine: brightdata_google
 prompt:
@@ -161,7 +160,6 @@ You will need to edit this file to suit your project. Let's break all this down:
 
 - `input_csv` and `output_csv` are the names of the data you want to process and where you want to save the results, respectively.
 - `model`: The LLM you want to use. You can find a list of supported models [here](https://ai.pydantic.dev/models/). Note that you need to provide both a `provider` and model `name` (ie. `anthropic` and `claude-3.5-sonnet`). You will also likely need to set up an API key (see [credentials below](#credentials)).
-- `query_col`: The name of the column in your input CSV that you want to use as the main search term (eg. company name).
 - `search`: The search engine you want to use. You can find a list of supported search engines [here](/docs/search.md). You will also likely need to set up an API key here (see [credentials](#credentials)).
 - `prompt`: LLMs take in a [system prompt](https://docs.anthropic.com/en/docs/build-with-claude/prompt-engineering/system-prompts) and a user prompt. Think of the system prompt as explaining to the LLM what its role is, and the user prompt as the instructions you want it to follow. You can use double curly braces (`{{ }}`) to refer to columns in your input CSV. Therea are some tips on writing good prompts [here](docs/prompt.md).
 - `structure`: The structure of the output data. You can think of this as the columns you want added to your original CSV.

augmenta/agent.py

@@ -39,12 +39,16 @@ def __init__(
             verbose: Whether to enable verbose logging with logfire
             system_prompt: Default system prompt for the agent
         """
-        # Create model settings with all available parameters
-        model_settings = {'temperature': temperature}
-        if rate_limit is not None:
-            model_settings['rate_limit'] = rate_limit
-        if max_tokens is not None:
-            model_settings['max_tokens'] = max_tokens
+        # Store parameters for reuse
+        self.model = model
+        self.temperature = temperature
+        self.rate_limit = rate_limit
+        self.max_tokens = max_tokens
+        self.verbose = verbose
+        self.system_prompt = system_prompt
+        
+        # Create model settings
+        model_settings = self._create_model_settings(temperature)
 
         # Load MCP servers from config
         try:
@@ -59,13 +63,23 @@ def __init__(
             tools=[search_web, visit_webpages],
             mcp_servers=mcp_servers
         )
-        self.model = model
-        self.temperature = temperature
-        self.rate_limit = rate_limit
-        self.max_tokens = max_tokens
-        self.verbose = verbose
-        self.system_prompt = system_prompt
 
+    def _create_model_settings(self, temperature: float) -> Dict[str, Any]:
+        """Create model settings dictionary with proper parameters.
+        
+        Args:
+            temperature: Temperature setting for the model
+            
+        Returns:
+            Dictionary with model settings
+        """
+        model_settings = {'temperature': temperature}
+        if self.rate_limit is not None:
+            model_settings['rate_limit'] = self.rate_limit
+        if self.max_tokens is not None:
+            model_settings['max_tokens'] = self.max_tokens
+        return model_settings
+    
     @staticmethod
     def create_structure_class(yaml_file_path: Union[str, Path]) -> Type[BaseModel]:
         """Creates a Pydantic model from YAML structure definition.
@@ -77,7 +91,6 @@ def create_structure_class(yaml_file_path: Union[str, Path]) -> Type[BaseModel]:
             A Pydantic model class based on the YAML structure
         """
         yaml_file_path = Path(yaml_file_path)
-        
         try:
             with open(yaml_file_path, 'r', encoding='utf-8') as f:
                 yaml_content = yaml.safe_load(f)
@@ -90,9 +103,12 @@ def create_structure_class(yaml_file_path: Union[str, Path]) -> Type[BaseModel]:
                 if not isinstance(field_info, dict):
                     raise ValueError(f"Invalid field definition for {field_name}")
 
-                field_type = (Literal[tuple(str(opt) for opt in field_info['options'])] 
-                           if 'options' in field_info 
-                           else AugmentaAgent.TYPE_MAPPING.get(field_info.get('type', 'str'), str))
+                # Determine field type based on options or type specification
+                if 'options' in field_info:
+                    field_type = Literal[tuple(str(opt) for opt in field_info['options'])]
+                else:
+                    type_str = field_info.get('type', 'str')
+                    field_type = AugmentaAgent.TYPE_MAPPING.get(type_str, str)
 
                 fields[field_name] = (
                     field_type,
@@ -101,8 +117,9 @@ def create_structure_class(yaml_file_path: Union[str, Path]) -> Type[BaseModel]:
 
             return create_model('Structure', **fields, __base__=BaseModel)
 
-        except (yaml.YAMLError, OSError) as e:            raise ValueError(f"Failed to parse YAML: {e}")
-            
+        except (yaml.YAMLError, OSError) as e:
+            raise ValueError(f"Failed to parse YAML: {e}")
+    
     async def run(
         self,
         prompt: str,
@@ -122,18 +139,14 @@ async def run(
             The agent's response after researching, either as string, dict or Pydantic model
         """
         try:
-            # Create model_settings for this specific request if temperature is provided
-            model_settings = None
-            if temperature is not None:
-                model_settings = {'temperature': temperature}
-                if self.rate_limit is not None:
-                    model_settings['rate_limit'] = self.rate_limit
-                if self.max_tokens is not None:
-                    model_settings['max_tokens'] = self.max_tokens
-                
             # Set the system prompt
             self.agent.system_prompt = system_prompt or self.system_prompt
 
+            # Prepare model settings only if temperature override is provided
+            model_settings = None
+            if temperature is not None and temperature != self.temperature:
+                model_settings = self._create_model_settings(temperature)
+            
             async with self.agent.run_mcp_servers():
                 result = await self.agent.run(
                     prompt,