# Agents

## Schema Hierarchy

The Atomic Agents framework uses Pydantic for schema validation and serialization. All input and output schemas follow this inheritance pattern:

```
pydantic.BaseModel
    └── BaseIOSchema
        ├── BaseAgentInputSchema
        └── BaseAgentOutputSchema
```

### BaseIOSchema

The base schema class that all agent input/output schemas inherit from.

```{eval-rst}
.. py:class:: BaseIOSchema

    Base schema class for all agent input/output schemas. Inherits from :class:`pydantic.BaseModel`.

    All agent schemas must inherit from this class to ensure proper serialization and validation.

    **Inheritance:**
        - :class:`pydantic.BaseModel`
```

### BaseAgentInputSchema

The default input schema for agents.

```{eval-rst}
.. py:class:: BaseAgentInputSchema

    Default input schema for agent interactions.

    **Inheritance:**
        - :class:`BaseIOSchema` → :class:`pydantic.BaseModel`

    .. py:attribute:: chat_message
        :type: str

        The message to send to the agent.

    Example:
        >>> input_schema = BaseAgentInputSchema(chat_message="Hello, agent!")
        >>> agent.run(input_schema)
```

### BaseAgentOutputSchema

The default output schema for agents.

```{eval-rst}
.. py:class:: BaseAgentOutputSchema

    Default output schema for agent responses.

    **Inheritance:**
        - :class:`BaseIOSchema` → :class:`pydantic.BaseModel`

    .. py:attribute:: chat_message
        :type: str

        The response message from the agent.

    Example:
        >>> response = agent.run(input_schema)
        >>> print(response.chat_message)
```

### Creating Custom Schemas

You can create custom input/output schemas by inheriting from `BaseIOSchema`:

```python
from pydantic import Field
from typing import List
from atomic_agents.lib.base.base_io_schema import BaseIOSchema

class CustomInputSchema(BaseIOSchema):
    chat_message: str = Field(..., description="User's message")
    context: str = Field(None, description="Optional context for the agent")

class CustomOutputSchema(BaseIOSchema):
    chat_message: str = Field(..., description="Agent's response")
    follow_up_questions: List[str] = Field(
        default_factory=list,
        description="Suggested follow-up questions"
    )
    confidence: float = Field(
        ...,
        description="Confidence score for the response",
        ge=0.0,
        le=1.0
    )
```

## Base Agent

The `BaseAgent` class is the foundation for building AI agents in the Atomic Agents framework. It handles chat interactions, memory management, system prompts, and responses from language models.

```python
from atomic_agents.agents.base_agent import BaseAgent, BaseAgentConfig
from atomic_agents.lib.components.agent_memory import AgentMemory
from atomic_agents.lib.components.system_prompt_generator import SystemPromptGenerator

# Create agent with basic configuration
agent = BaseAgent(
    config=BaseAgentConfig(
        client=instructor.from_openai(OpenAI()),
        model="gpt-4-turbo-preview",
        memory=AgentMemory(),
        system_prompt_generator=SystemPromptGenerator()
    )
)

# Run the agent
response = agent.run(user_input)

# Stream responses
async for partial_response in agent.run_async(user_input):
    print(partial_response)
```

### Configuration

The `BaseAgentConfig` class provides configuration options:

```python
class BaseAgentConfig:
    client: instructor.Instructor  # Client for interacting with the language model
    model: str = "gpt-4-turbo-preview"  # Model to use
    memory: Optional[AgentMemory] = None  # Memory component
    system_prompt_generator: Optional[SystemPromptGenerator] = None  # Prompt generator
    input_schema: Optional[Type[BaseModel]] = None  # Custom input schema
    output_schema: Optional[Type[BaseModel]] = None  # Custom output schema
    model_api_parameters: Optional[dict] = None  # Additional API parameters
```

### Input/Output Schemas

Default schemas for basic chat interactions:

```python
class BaseAgentInputSchema(BaseIOSchema):
    """Input from the user to the AI agent."""
    chat_message: str = Field(
        ...,
        description="The chat message sent by the user."
    )

class BaseAgentOutputSchema(BaseIOSchema):
    """Response generated by the chat agent."""
    chat_message: str = Field(
        ...,
        description="The markdown-enabled response generated by the chat agent."
    )
```

### Key Methods

- `run(user_input: Optional[BaseIOSchema] = None) -> BaseIOSchema`: Process user input and get response
- `run_async(user_input: Optional[BaseIOSchema] = None)`: Stream responses asynchronously
- `get_response(response_model=None) -> Type[BaseModel]`: Get direct model response
- `reset_memory()`: Reset memory to initial state
- `get_context_provider(provider_name: str)`: Get a registered context provider
- `register_context_provider(provider_name: str, provider: SystemPromptContextProviderBase)`: Register a new context provider
- `unregister_context_provider(provider_name: str)`: Remove a context provider

### Context Providers

Context providers can be used to inject dynamic information into the system prompt:

```python
from atomic_agents.lib.components.system_prompt_generator import SystemPromptContextProviderBase

class SearchResultsProvider(SystemPromptContextProviderBase):
    def __init__(self, title: str):
        super().__init__(title=title)
        self.results = []

    def get_info(self) -> str:
        return "\n\n".join([
            f"Result {idx}:\n{result}"
            for idx, result in enumerate(self.results, 1)
        ])

# Register with agent
agent.register_context_provider(
    "search_results",
    SearchResultsProvider("Search Results")
)
```

### Streaming Support

The agent supports streaming responses for more interactive experiences:

```python
async def chat():
    async for partial_response in agent.run_async(user_input):
        # Handle each chunk of the response
        print(partial_response.chat_message)
```

### Memory Management

The agent automatically manages conversation history through the `AgentMemory` component:

```python
# Access memory
history = agent.memory.get_history()

# Reset to initial state
agent.reset_memory()

# Save/load memory state
serialized = agent.memory.dump()
agent.memory.load(serialized)
```

### Custom Schemas

You can use custom input/output schemas for structured interactions:

```python
from pydantic import BaseModel, Field
from typing import List

class CustomInput(BaseIOSchema):
    """Custom input with specific fields"""
    question: str = Field(..., description="User's question")
    context: str = Field(..., description="Additional context")

class CustomOutput(BaseIOSchema):
    """Custom output with structured data"""
    answer: str = Field(..., description="Answer to the question")
    sources: List[str] = Field(..., description="Source references")

# Create agent with custom schemas
agent = BaseAgent[CustomInput, CustomOutput](
    config=BaseAgentConfig(
        client=client,
        model=model,
        input_schema=CustomInput,
        output_schema=CustomOutput
    )
)
```

For full API details:
```{eval-rst}
.. automodule:: atomic_agents.agents.base_agent
   :members:
   :undoc-members:
   :show-inheritance:
```