from typing import List, Callable, Optional, Union, Dict
from pydantic import BaseModel

class Agent(BaseModel):
    name: str = "Agent"
    model: str = "llama3"
    instructions: Union[str, Callable] = "You are a helpful assistant."
    functions: List[Callable] = []
    context_variables: Dict = {}
    
    class Config:
        arbitrary_types_allowed = True

class Response(BaseModel):
    agent: Optional[Agent]
    messages: List[Dict]
    context_variables: Dict = {}
    trace: List[Dict] = []