State Management

1
class DataCollectionAgent(OutputAgentNode):
2
    def __init__(self):
3
        super().__init__(name="data-agent")
4
        
5
        # Session state
6
        self.customer_name = None
7
        self.email = None
8
        self.phone = None
9
        self.order_id = None
10
11
    @function_tool()
12
    def save_customer_info(
13
        self,
14
        name: str = None,
15
        email: str = None,
16
        phone: str = None
17
    ) -> dict:
18
        """Save customer information."""
19
        if name:
20
            self.customer_name = name
21
        if email:
22
            self.email = email
23
        if phone:
24
            self.phone = phone
25
        
26
        return {"saved": True, "collected": self._get_collected()}
27
    
28
    def _get_collected(self) -> dict:
29
        return {
30
            "name": self.customer_name,
31
            "email": self.email,
32
            "phone": self.phone
33
        }

1
async def start(self, init_event, task_manager):
2
    await super().start(init_event, task_manager)
3
    
4
    ctx = init_event.session_context
5
    
6
    # Access initial variables
7
    self.customer_id = ctx.initial_variables.get("customer_id")
8
    self.account_tier = ctx.initial_variables.get("tier")
9
    self.language = ctx.initial_variables.get("language", "en")
10
    
11
    # Use them to customize behavior
12
    if self.account_tier == "premium":
13
        self.context.add_message({
14
            "role": "system",
15
            "content": "This is a premium customer. Prioritize their requests."
16
        })

1
class MyAgent(OutputAgentNode):
2
    async def generate_response(self):
3
        # Access full message history
4
        all_messages = self.context.messages
5
        
6
        # Get the last user message
7
        user_messages = [m for m in all_messages if m["role"] == "user"]
8
        last_user = user_messages[-1]["content"] if user_messages else ""
9
        
10
        # Count turns
11
        user_turns = sum(1 for m in all_messages if m["role"] == "user")

1
class IntakeAgent(OutputAgentNode):
2
    def __init__(self):
3
        super().__init__(name="intake-agent")
4
        
5
        self.collected = {}
6
        self.tool_registry = ToolRegistry()
7
        self.tool_registry.discover(self)
8
        self.tool_schemas = self.tool_registry.get_schemas()
9
10
    @function_tool()
11
    def record_information(
12
        self,
13
        field: str,
14
        value: str
15
    ) -> dict:
16
        """
17
        Record a piece of information from the conversation.
18
        
19
        Args:
20
            field: The type of information (name, email, phone, etc.)
21
            value: The value to record
22
        """
23
        self.collected[field] = value
24
        
25
        return {
26
            "recorded": f"{field}: {value}",
27
            "remaining": self._get_remaining_fields()
28
        }
29
    
30
    def _get_remaining_fields(self):
31
        required = ["name", "email", "reason"]
32
        return [f for f in required if f not in self.collected]

1
class ProgressAgent(OutputAgentNode):
2
    def __init__(self):
3
        super().__init__(name="progress-agent")
4
        self.turn_count = 0
5
        self.started_at = None
6
7
    async def start(self, init_event, task_manager):
8
        await super().start(init_event, task_manager)
9
        self.started_at = datetime.now()
10
11
    async def generate_response(self):
12
        self.turn_count += 1
13
        
14
        # Different behavior based on turn count
15
        if self.turn_count == 1:
16
            yield "Welcome! This is our first exchange."
17
        elif self.turn_count > 10:
18
            yield "We've been chatting for a while. "
19
            yield "Is there anything else I can help with?"

1
class RouterNode(Node):
2
    async def process_event(self, event):
3
        # Store routing decision in event metadata
4
        event.metadata["routed_to"] = "sales"
5
        await self.send_event(event)
6
7
8
class SalesAgent(OutputAgentNode):
9
    async def process_event(self, event):
10
        # Read state from event
11
        if event.metadata.get("routed_to") == "sales":
12
            await super().process_event(event)

1
import redis
2
3
4
class PersistentAgent(OutputAgentNode):
5
    def __init__(self):
6
        super().__init__(name="persistent-agent")
7
        self.redis = redis.Redis()
8
9
    async def start(self, init_event, task_manager):
10
        await super().start(init_event, task_manager)
11
        
12
        phone = init_event.session_context.initial_variables.get("phone")
13
        
14
        # Load previous state
15
        stored = self.redis.get(f"user:{phone}")
16
        if stored:
17
            self.user_data = json.loads(stored)
18
        else:
19
            self.user_data = {}
20
21
    async def stop(self):
22
        # Save state before session ends
23
        phone = self.init_event.session_context.initial_variables.get("phone")
24
        self.redis.set(f"user:{phone}", json.dumps(self.user_data))
25
        await super().stop()

1
class ValidationAgent(OutputAgentNode):
2
    def __init__(self):
3
        super().__init__(name="validation-agent")
4
        self.data = {}
5
6
    @function_tool()
7
    def submit_form(self) -> dict:
8
        """Submit the collected information."""
9
        errors = []
10
        
11
        if not self.data.get("name"):
12
            errors.append("Name is required")
13
        
14
        if not self.data.get("email"):
15
            errors.append("Email is required")
16
        elif "@" not in self.data["email"]:
17
            errors.append("Email format is invalid")
18
        
19
        if errors:
20
            return {"success": False, "errors": errors}
21
        
22
        return {"success": True, "data": self.data}