Custom Agents

Learn how to create custom agents that extend the MARSYS framework with specialized capabilities and behaviors.

Overview

Custom agents enable you to:

Extend BaseAgent: Create agents with specialized behavior
Add Domain Knowledge: Incorporate specific expertise
Implement State Machines: Complex multi-step workflows
Custom Response Processing: Specialized output formats

Creating Custom Agents

Extending BaseAgent

The fundamental pattern for custom agents:

from marsys.agents import BaseAgent
from marsys.agents.memory import Message
from typing import Dict, Any, Optional

class CustomAgent(BaseAgent):
    """Custom agent with specialized behavior."""

    def __init__(
        self,
        model,
        specialized_knowledge: Optional[Dict] = None,
        **kwargs
    ):
        super().__init__(
            model=model,
            name=kwargs.pop("name", "CustomAgent"),
            goal="Custom specialized agent",
            instruction="A custom agent with specialized capabilities.",
            **kwargs
        )
        self.specialized_knowledge = specialized_knowledge or {}

    async def _run(
        self,
        prompt: Any,
        context: Dict[str, Any],
        **kwargs
    ) -> Message:
        """
        Pure execution logic - NO side effects!
        This method must be stateless and pure.
        """

        # Prepare messages with base functionality
        messages = self._prepare_messages(prompt)

        # Add custom logic
        if self._should_use_knowledge(prompt):
            messages.append(Message(
                role="system",
                content=f"Use this knowledge: {self.specialized_knowledge}"
            ))

        # Call model
        response = await self.model.run(messages)

        # Return pure Message object
        return Message(
            role="assistant",
            content=response.content,
            tool_calls=getattr(response, 'tool_calls', None)
        )

    def _should_use_knowledge(self, prompt: str) -> bool:
        """Determine if specialized knowledge is relevant."""
        return True

Pure Execution Rule

The _run() method must be pure - no memory manipulation, no logging to files, no global state changes. All side effects are handled by the coordination layer.

Domain-Specific Agents

Create agents with domain expertise:

class FinancialAnalyst(BaseAgent):
    """Financial analysis specialist."""

    def __init__(self, model, market_data_api=None, **kwargs):
        super().__init__(
            model=model,
            name=kwargs.pop("name", "FinancialAnalyst"),
            goal="Expert financial analyst with real-time market access",
            instruction="You are a senior financial analyst.",
            tools={
                "get_stock_price": self._get_stock_price,
                "calculate_metrics": self._calculate_metrics
            },
            **kwargs
        )
        self.market_data_api = market_data_api

    async def _run(self, prompt, context, **kwargs):
        """Analyze financial data with expertise."""

        # Add market context for financial queries
        if self._is_financial_query(prompt):
            market_context = await self._get_market_context()
            enriched_prompt = f"{prompt}\n\nMarket: {market_context}"
        else:
            enriched_prompt = prompt

        messages = self._prepare_messages(enriched_prompt)
        response = await self.model.run(messages)

        return Message(role="assistant", content=response.content)

    def _get_stock_price(self, symbol: str) -> Dict:
        """Get real-time stock price."""
        if self.market_data_api:
            return self.market_data_api.get_price(symbol)
        return {"error": "No market data API configured"}

Advanced Patterns

State Machine Agent

Implement complex workflows with state transitions:

from enum import Enum

class ProcessingState(Enum):
    INITIAL = "initial"
    ANALYZING = "analyzing"
    VALIDATING = "validating"
    COMPLETE = "complete"

class StateMachineAgent(BaseAgent):
    """Agent with state machine behavior."""

    def __init__(self, model, **kwargs):
        super().__init__(model=model, **kwargs)
        self.state = ProcessingState.INITIAL

    async def _run(self, prompt, context, **kwargs):
        # Handle based on current state
        if self.state == ProcessingState.INITIAL:
            return await self._handle_initial(prompt)
        elif self.state == ProcessingState.ANALYZING:
            return await self._handle_analyzing(prompt)
        # ... etc

    async def _handle_initial(self, prompt):
        # Initial processing logic
        pass

Agent with Custom Validation

from pydantic import BaseModel

class AnalysisInput(BaseModel):
    query: str
    data_source: str
    output_format: str = "json"

class AnalysisOutput(BaseModel):
    result: Dict[str, Any]
    confidence: float
    methodology: str

class ValidatedAgent(BaseAgent):
    """Agent with input/output validation."""

    def __init__(self, model, **kwargs):
        super().__init__(
            model=model,
            input_schema=AnalysisInput,
            output_schema=AnalysisOutput,
            **kwargs
        )

Best Practices

Keep _run() pure: No side effects in the execution method
Use descriptive names: Agent names should reflect their purpose
Document tools: Clear docstrings help the model use tools correctly
Validate inputs: Use schemas for type safety
Handle edge cases: Consider empty inputs and error scenarios

Reusing Components

Consider creating reusable base classes for common patterns in your domain. For example, a DataAnalysisAgent base that multiple specialized analyzers can extend.

Navigation