remove fixed agent

Author: nicucalcea

augmenta/__init__.py

@@ -0,0 +1,5 @@
+"""Augmenta - Web research and data augmentation package."""
+
+from .agent import AugmentaAgent
+
+__all__ = ['AugmentaAgent']

augmenta/agents/base_agent.py augmenta/agent.pyaugmenta/agents/base_agent.py renamed to augmenta/agent.py

@@ -1,17 +1,19 @@
+"""Unified agent module for Augmenta."""
+
 from typing import Type, Optional, Union, Any, Dict, ClassVar, Literal, List
 from pathlib import Path
 import yaml
 from pydantic import BaseModel, Field, create_model
 from pydantic_ai import Agent
-# from pydantic_ai.usage import UsageLimits
-# from pydantic_ai.mcp import MCPServerStdio
 import logfire
-from ..tools.mcp import load_mcp_servers
+from .tools.mcp import load_mcp_servers
+from .tools.search_web import search_web
+from .tools.visit_webpages import visit_webpages
 
 Agent.instrument_all()
 
-class BaseAgent:
-    """Base class providing core LLM functionality"""
+class AugmentaAgent:
+    """Agent that provides LLM functionality with web research capabilities."""
 
     TYPE_MAPPING: ClassVar[Dict[str, type]] = {
         'str': str, 'bool': bool, 'int': int,
@@ -25,19 +27,17 @@ def __init__(
         rate_limit: Optional[float] = None,
         max_tokens: Optional[int] = None,
         verbose: bool = False,
-        tools: Optional[list] = None,
-        mcp_servers: Optional[list] = None
+        system_prompt: str = "You are a web research assistant. Use the provided tools to search for information and analyse web pages."
     ):
-        """Initialize the base agent.
+        """Initialize the agent.
         
         Args:
             model: The LLM model identifier
             temperature: Temperature setting for the model
             rate_limit: Optional rate limit between requests
             max_tokens: Optional maximum tokens for response
             verbose: Whether to enable verbose logging with logfire
-            tools: Optional list of tool functions to register
-            mcp_servers: Optional list of pre-configured MCP servers (overrides config)
+            system_prompt: Default system prompt for the agent
         """
         # Create model settings with all available parameters
         model_settings = {'temperature': temperature}
@@ -46,25 +46,25 @@ def __init__(
         if max_tokens is not None:
             model_settings['max_tokens'] = max_tokens
 
-        # Load MCP servers from config if not provided explicitly
-        if mcp_servers is None:
-            try:
-                mcp_servers = load_mcp_servers()
-            except RuntimeError:
-                # Config not loaded yet, proceed without MCP servers
-                mcp_servers = []
+        # Load MCP servers from config
+        try:
+            mcp_servers = load_mcp_servers()
+        except RuntimeError:
+            # Config not loaded yet, proceed without MCP servers
+            mcp_servers = []
 
         self.agent = Agent(
             model,
             model_settings=model_settings,
-            tools=tools or [],
+            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
 
     @staticmethod
     def create_structure_class(yaml_file_path: Union[str, Path]) -> Type[BaseModel]:
@@ -92,7 +92,7 @@ def create_structure_class(yaml_file_path: Union[str, Path]) -> Type[BaseModel]:
 
                 field_type = (Literal[tuple(str(opt) for opt in field_info['options'])] 
                            if 'options' in field_info 
-                           else BaseAgent.TYPE_MAPPING.get(field_info.get('type', 'str'), str))
+                           else AugmentaAgent.TYPE_MAPPING.get(field_info.get('type', 'str'), str))
 
                 fields[field_name] = (
                     field_type,
@@ -135,23 +135,31 @@ async def complete(
             # Set the system prompt
             self.agent.system_prompt = prompt_system
 
-            # Run the agent with the appropriate parameters
-            # result = await self.agent.run(
-            #     prompt_user,
-            #     output_type=response_format,
-            #     model_settings=model_settings
-            #     # usage_limits=UsageLimits(request_limit=5)
-            # )
-
             async with self.agent.run_mcp_servers():
                 result = await self.agent.run(
                     prompt_user,
                     output_type=response_format,
                     model_settings=model_settings
-            )
+                )
 
             # Return the appropriate result format
             return result.data.model_dump() if response_format else result.data
         except Exception as e:
             logfire.error(f"LLM request failed: {e}")
             raise RuntimeError(f"LLM request failed: {e}")
+            
+    async def run(self, prompt: str, response_format: Optional[Type[BaseModel]] = None) -> Union[str, dict[str, Any], BaseModel]:
+        """Run the agent to perform web research.
+        
+        Args:
+            prompt: The research query or task
+            response_format: Optional Pydantic model for structured output
+            
+        Returns:
+            The agent's response after researching, either as string, dict or Pydantic model
+        """
+        return await self.complete(
+            prompt_system=self.system_prompt,
+            prompt_user=prompt,
+            response_format=response_format
+        )

augmenta/agents/__init__.py

@@ -8,9 +8,7 @@
 from dataclasses import dataclass
 
 from augmenta.utils.prompt_formatter import format_examples
-from augmenta.agents.base_agent import BaseAgent
-from augmenta.agents.autonomous_agent import AutonomousAgent
-from augmenta.agents.fixed_agent import FixedAgent
+from augmenta.agent import AugmentaAgent
 from augmenta.cache import CacheManager
 from augmenta.cache.process import handle_process_resumption, setup_caching, apply_cached_results
 from augmenta.config.read_config import load_config, get_config_values
@@ -45,6 +43,7 @@ async def process_augmenta(
 
     # Initialize agent once for all processing
     config_values = get_config_values(config_data)
+    
     agent_settings = {
         "model": config_values["model_id"],
         "temperature": config_values["temperature"],
@@ -53,12 +52,11 @@ async def process_augmenta(
         "system_prompt": config_data["prompt"]["system"]
     }
 
-    agent_mode = config_data.get("agent", "autonomous")
-    agent = (AutonomousAgent(**agent_settings) if agent_mode == "autonomous"
-            else FixedAgent(**agent_settings))
+    # Create agent instance
+    agent = AugmentaAgent(**agent_settings)
 
     # Create structure class once
-    response_format = BaseAgent.create_structure_class(config_data["config_path"])
+    response_format = AugmentaAgent.create_structure_class(config_data["config_path"])
 
     # Load and validate input data
     df = pd.read_csv(config_data["input_csv"])
@@ -139,7 +137,7 @@ async def process_with_limit(row):
 async def process_row(
     row_data: Dict[str, Any],
     config: Dict[str, Any],
-    agent: BaseAgent,
+    agent: AugmentaAgent,
     response_format: Type,
     cache_manager: Optional[CacheManager] = None,
     process_id: Optional[str] = None,
@@ -161,11 +159,8 @@ async def process_row(
             if examples_text := format_examples(examples_yaml):
                 prompt_user = f'{prompt_user}\n\n{examples_text}'
 
-        # Run prompt using the provided agent
-        if isinstance(agent, FixedAgent):
-            response = await agent.run(prompt_user, query=query, response_format=response_format)
-        else:
-            response = await agent.run(prompt_user, response_format=response_format)
+        # Run prompt using the agent
+        response = await agent.run(prompt_user, response_format=response_format)
 
         # Handle caching if enabled
         if cache_manager and process_id:

augmenta/agents/autonomous_agent.py

@@ -1,11 +1,37 @@
-# Fixed agent
+# Augmenta Agent
 
-By default, Augmenta runs in an "autonomous" mode, where it will autonomatically make decisions about what to search for, which resuults to read, and when it has enough information to make a decision.
+Augmenta uses a unified agent architecture that handles web research through a streamlined interface.
 
-You may want to switch to a "pre-determined" mode, where you can specify the search terms and the number of results to scrape. This is useful if you want to have more control over the search process, or if you want to use Augmenta in a more deterministic way.
+The `AugmentaAgent` class provides built-in capabilities for web search, content extraction, and LLM-based analysis.
 
-To start, add this to your `config.yml`:
+## Features
+
+- Web search using various search providers
+- Content extraction from web pages
+- Integration with Model Context Protocol (MCP) servers
+- Structured output formatting via Pydantic models
+- Customizable system prompts and model parameters
+
+## Usage
+
+The agent is automatically configured when you run Augmenta with your configuration:
 
 ```yaml
-agent: fixed
+# Basic configuration example
+model: "openai/gpt-4o"
+temperature: 0.2
+max_tokens: 2048
+rate_limit: 10  # Optional: Requests per minute
 ```
+
+## Advanced Configuration
+
+You can customize the agent's behavior through your configuration file:
+
+```yaml
+prompt:
+  system: "You are a research assistant specialized in analyzing company donations data. Your task is to categorize each donation and provide reasoning."
+  user: "Analyze this donation: {{amount}} from {{donor}} to {{recipient}}. Provide category and reasoning."
+```
+
+The agent integrates seamlessly with Augmenta's search and extraction tools to provide comprehensive research capabilities.