AI Agent Design

{
  "type": "function",
  "function": {
    "name": "search_knowledge_base",
    "description": "Search the internal knowledge base for relevant documents",
    "parameters": {
      "type": "object",
      "properties": {
        "query": {
          "type": "string",
          "description": "The search query string"
        },
        "max_results": {
          "type": "integer",
          "description": "Maximum number of results to return (1-20)",
          "minimum": 1,
          "maximum": 20
        },
        "filter_by_date": {
          "type": "string",
          "description": "Optional date filter in ISO 8601 format"
        }
      },
      "required": ["query"]
    }
  }
}

from typing import Any
import json

class AgentTool:
    """Base class for agent tools."""
    
    def __init__(self, name: str, description: str):
        self.name = name
        self.description = description
    
    def get_schema(self) -> dict:
        """Return the function calling schema for this tool."""
        raise NotImplementedError
    
    async def execute(self, **kwargs) -> Any:
        """Execute the tool with validated parameters."""
        raise NotImplementedError

class SearchTool(AgentTool):
    def __init__(self):
        super().__init__(
            name="search",
            description="Search documents by query string"
        )
    
    def get_schema(self):
        return {
            "type": "function",
            "function": {
                "name": self.name,
                "description": self.description,
                "parameters": {
                    "type": "object",
                    "properties": {
                        "query": {"type": "string"},
                        "limit": {"type": "integer", "default": 5}
                    },
                    "required": ["query"]
                }
            }
        }
    
    async def execute(self, query: str, limit: int = 5):
        # Implementation
        results = await database.search(query, limit=limit)
        return json.dumps({"results": results, "count": len(results)})

class ShortTermMemory:
    """In-memory conversation buffer."""
    
    def __init__(self, max_tokens: int = 8000):
        self.messages = []
        self.max_tokens = max_tokens
    
    def add(self, role: str, content: str):
        self.messages.append({"role": role, "content": content})
        self._trim()
    
    def _trim(self):
        """Remove oldest messages when over capacity."""
        total = sum(len(m["content"]) for m in self.messages)
        while total > self.max_tokens and len(self.messages) > 1:
            removed = self.messages.pop(0)
            total -= len(removed["content"])
    
    def get_context(self) -> list:
        return self.messages

import chromadb

class LongTermMemory:
    """Persistent memory using vector storage."""
    
    def __init__(self, collection_name: str = "agent_memory"):
        self.client = chromadb.Client()
        self.collection = self.client.get_or_create_collection(collection_name)
    
    def store(self, content: str, metadata: dict = None):
        """Store a memory with embedding."""
        self.collection.add(
            documents=[content],
            metadatas=[metadata or {}],
            ids=[f"mem_{hash(content)}"]
        )
    
    def recall(self, query: str, n: int = 5) -> list:
        """Retrieve relevant memories."""
        results = self.collection.query(
            query_texts=[query],
            n_results=n
        )
        return [
            {"content": doc, "metadata": meta}
            for doc, meta in zip(results["documents"][0], results["metadatas"][0])
        ]

class EpisodicMemory:
    """Records agent actions and outcomes for learning."""
    
    def __init__(self):
        self.episodes = []
    
    def record(self, action: str, context: dict, outcome: str, success: bool):
        self.episodes.append({
            "timestamp": datetime.now().isoformat(),
            "action": action,
            "context": context,
            "outcome": outcome,
            "success": success
        })
    
    def get_similar_episodes(self, action: str, n: int = 3) -> list:
        """Find similar past episodes to inform decisions."""
        relevant = [e for e in self.episodes if e["action"] == action]
        return sorted(relevant, key=lambda x: x["timestamp"], reverse=True)[:n]

[User] → [LLM + Tools + Memory] → [Response]

                    [Supervisor Agent]
                    /        |        \
            [Research]  [Analysis]  [Writing]
            Agent        Agent        Agent

class SupervisorAgent:
    """Coordinates specialized worker agents."""
    
    def __init__(self):
        self.workers = {
            "researcher": ResearchAgent(),
            "analyst": AnalysisAgent(),
            "writer": WritingAgent()
        }
    
    async def process(self, task: str) -> str:
        # Step 1: Analyze the task
        plan = await self._create_plan(task)
        
        # Step 2: Delegate to workers
        results = {}
        for step in plan["steps"]:
            worker = self.workers[step["agent"]]
            results[step["id"]] = await worker.execute(step["instruction"])
        
        # Step 3: Synthesize
        return await self._synthesize(plan, results)

class Router:
    """Routes requests to the appropriate agent based on intent."""
    
    def __init__(self):
        self.routes = {
            "technical_support": TechnicalSupportAgent(),
            "billing": BillingAgent(),
            "general": GeneralAgent()
        }
    
    async def route(self, user_input: str):
        # Use LLM to classify intent
        intent = await self._classify_intent(user_input)
        
        # Route to the appropriate handler
        agent = self.routes.get(intent, self.routes["general"])
        return await agent.handle(user_input)

Thought: I need to find the latest sales data. Let me check the database.
Action: query_database({"query": "SELECT * FROM sales ORDER BY date DESC LIMIT 10"})
Observation: [{"date": "2024-01-15", "revenue": 45000}, ...]
Thought: I have the data. Now I need to identify trends.
Action: analyze_data({"data": [...], "analysis_type": "trend"})
Observation: Revenue has increased 12% month-over-month.
Thought: I can now answer the user's question about sales performance.
Answer: Sales revenue has grown 12% month-over-month, reaching $45,000 in January.

class ReActAgent:
    def __init__(self, llm, tools):
        self.llm = llm
        self.tools = {t.name: t for t in tools}
    
    async def run(self, task: str, max_steps: int = 10):
        messages = [{"role": "user", "content": task}]
        
        for step in range(max_steps):
            response = await self.llm.generate(messages)
            action = self._parse_action(response)
            
            if not action:
                return response  # Final answer
            
            tool = self.tools.get(action["name"])
            if not tool:
                return f"Error: Unknown tool {action['name']}"
            
            result = await tool.execute(**action["parameters"])
            messages.append({"role": "assistant", "content": response})
            messages.append({"role": "tool", "content": result})
        
        return "Reached maximum steps without resolution."

Plan:
1. Query database for Q4 sales data
2. Calculate year-over-year growth
3. Identify top-performing regions
4. Generate summary report
5. Schedule email to stakeholders

Executing step 1...
Executing step 2...
...

class ErrorRecovery:
    """Handles agent errors with graceful recovery strategies."""
    
    MAX_RETRIES = 3
    
    async def execute_with_recovery(self, tool, params):
        for attempt in range(self.MAX_RETRIES):
            try:
                return await tool.execute(**params)
            except ValidationError as e:
                # Parameter issue — fix and retry
                params = self._fix_params(e, params)
                continue
            except TimeoutError:
                # Transient failure — wait and retry
                await asyncio.sleep(2 ** attempt)  # Exponential backoff
                continue
            except PermissionError:
                # Can't recover — escalate
                return {"error": "permission_denied", "action": "escalate"}
        
        return {"error": "max_retries_exceeded", "action": "ask_human"}

class EscalationHandler:
    """Escalates to humans when the agent can't proceed."""
    
    async def should_escalate(self, error: dict, confidence: float) -> bool:
        """Determine if we need human intervention."""
        return (
            error.get("action") == "escalate" or
            confidence < 0.3 or
            error.get("type") in ["security_violation", "permission_denied"]
        )
    
    async def escalate(self, context: dict, error: dict):
        """Send to human operator with full context."""
        ticket = {
            "agent_id": context["agent_id"],
            "task": context["task"],
            "error": error,
            "conversation_history": context["history"][-10:],
            "timestamp": datetime.now().isoformat()
        }
        await notification_service.send_to_human(ticket)
        return "Escalated to human operator. They will review shortly."

class InputGuardrail:
    """Validates and sanitizes user input before it reaches the agent."""
    
    BLOCKED_PATTERNS = [
        r"ignore all previous instructions",
        r"you are now .*",
        r"system prompt",
        r"jailbreak",
    ]
    
    def check(self, user_input: str) -> tuple[bool, str]:
        """Returns (allowed, reason)."""
        for pattern in self.BLOCKED_PATTERNS:
            if re.search(pattern, user_input, re.IGNORECASE):
                return False, f"Input blocked: pattern '{pattern}' detected"
        return True, ""

class OutputGuardrail:
    """Validates agent output before sending to user."""
    
    def check(self, output: str) -> tuple[bool, str]:
        # Never output system prompts or internal instructions
        if "system:" in output.lower() and "instruction" in output.lower():
            return False, "Output contains internal instructions"
        
        # Never output harmful content
        if contains_harmful_content(output):
            return False, "Output flagged by safety classifier"
        
        return True, ""

class ActionGuardrail:
    """Requires confirmation for high-risk actions."""
    
    HIGH_RISK_ACTIONS = {
        "send_email", "delete_record", "modify_permissions", 
        "transfer_funds", "execute_code"
    }
    
    async def authorize(self, action: str, params: dict) -> bool:
        if action in self.HIGH_RISK_ACTIONS:
            return await self._ask_human(f"Confirm {action} with params: {params}")
        return True  # Low-risk actions auto-approved

[User Input] → [Input Guardrail] → [Agent] → [Output Guardrail] → [Action Guardrail] → [Execute]
                      |                  |            |                    |
                  Blocked/Pass       +Tool       Blocked/Pass       Confirm/Deny
                                   Results

# Bad
"name": "search"

# Good  
"name": "search_knowledge_base",
"description": "Search internal documentation for product information. Use this when users ask about product features, specifications, or troubleshooting."

# Dangerous — letting tool result influence behavior
system_prompt += f"\nHere's additional context: {tool_result}"

# Safe — tool result is data, not instructions
context = f"According to the database: {tool_result}"