State in Router

1from kern.agent import Agent
2from kern.models.openai import OpenAIResponses
3from kern.workflow.router import Router
4from kern.workflow.step import Step, StepInput, StepOutput
5from kern.workflow.workflow import Workflow
6from kern.run import RunContext
7
8
9def route_based_on_user_preference(step_input: StepInput, run_context: RunContext) -> Step:
10    """
11    Router selector that chooses an agent based on user preferences in session_state.
12
13    Args:
14        step_input: The input for this step (contains user query)
15        run_context: The run context object
16
17    Returns:
18        Step: The step to execute based on user preference
19    """
20    print("\n=== Routing Decision ===")
21    print(f"User ID: {run_context.session_state.get('current_user_id')}")
22    print(f"Session ID: {run_context.session_state.get('current_session_id')}")
23
24    # Get user preference from session state
25    user_preference = run_context.session_state.get("agent_preference", "general")
26    interaction_count = run_context.session_state.get("interaction_count", 0)
27
28    print(f"User Preference: {user_preference}")
29    print(f"Interaction Count: {interaction_count}")
30
31    # Update interaction count
32    run_context.session_state["interaction_count"] = interaction_count + 1
33
34    # Route based on preference
35    if user_preference == "technical":
36        print("→ Routing to Technical Expert")
37        return technical_step
38    elif user_preference == "friendly":
39        print("→ Routing to Friendly Assistant")
40        return friendly_step
41    else:
42        # For first interaction, route to onboarding
43        if interaction_count == 0:
44            print("→ Routing to Onboarding (first interaction)")
45            return onboarding_step
46        else:
47            print("→ Routing to General Assistant")
48            return general_step
49
50
51def set_user_preference(step_input: StepInput, run_context: RunContext) -> StepOutput:
52    """Custom function that sets user preference based on onboarding."""
53    print("\n=== Setting User Preference ===")
54
55    # In a real scenario, this would analyze the user's response
56    # For demo purposes, we'll set it based on interaction count
57    interaction_count = run_context.session_state.get("interaction_count", 0)
58
59    if interaction_count % 3 == 1:
60        run_context.session_state["agent_preference"] = "technical"
61        preference = "technical"
62    elif interaction_count % 3 == 2:
63        run_context.session_state["agent_preference"] = "friendly"
64        preference = "friendly"
65    else:
66        run_context.session_state["agent_preference"] = "general"
67        preference = "general"
68
69    print(f"Set preference to: {preference}")
70    return StepOutput(content=f"Preference set to: {preference}")
71
72
73# Create specialized agents
74onboarding_agent = Agent(
75    name="Onboarding Agent",
76    model=OpenAIResponses(id="gpt-5.2"),
77    instructions=(
78        "Welcome new users and ask about their preferences. "
79        "Determine if they prefer technical or friendly assistance."
80    ),
81    markdown=True,
82)
83
84technical_agent = Agent(
85    name="Technical Expert",
86    model=OpenAIResponses(id="gpt-5.2"),
87    instructions=(
88        "You are a technical expert. Provide detailed, technical answers with code examples and best practices."
89    ),
90    markdown=True,
91)
92
93friendly_agent = Agent(
94    name="Friendly Assistant",
95    model=OpenAIResponses(id="gpt-5.2"),
96    instructions=(
97        "You are a friendly, casual assistant. Use simple language, emojis, and make the conversation fun."
98    ),
99    markdown=True,
100)
101
102general_agent = Agent(
103    name="General Assistant",
104    model=OpenAIResponses(id="gpt-5.2"),
105    instructions=(
106        "You are a balanced assistant. Provide helpful answers that are neither too technical nor too casual."
107    ),
108    markdown=True,
109)
110
111# Create steps for routing
112onboarding_step = Step(
113    name="Onboard User",
114    description="Onboard new user and set preferences",
115    agent=onboarding_agent,
116)
117
118technical_step = Step(
119    name="Technical Response",
120    description="Provide technical assistance",
121    agent=technical_agent,
122)
123
124friendly_step = Step(
125    name="Friendly Response",
126    description="Provide friendly assistance",
127    agent=friendly_agent,
128)
129
130general_step = Step(
131    name="General Response",
132    description="Provide general assistance",
133    agent=general_agent,
134)
135
136# Create workflow with router
137workflow = Workflow(
138    name="Adaptive Assistant Workflow",
139    steps=[
140        # Router that selects agent based on session state
141        Router(
142            name="Route to Appropriate Agent",
143            description="Route to the appropriate agent based on user preferences",
144            selector=route_based_on_user_preference,
145            choices=[
146                onboarding_step,
147                technical_step,
148                friendly_step,
149                general_step,
150            ],
151        ),
152        # After first interaction, update preferences
153        Step(
154            name="Update Preferences",
155            description="Update user preferences based on interaction",
156            executor=set_user_preference,
157        ),
158    ],
159    session_state={
160        "agent_preference": "general",
161        "interaction_count": 0,
162    },
163)
164
165
166def run_example():
167    """Run the example workflow multiple times to see dynamic routing."""
168
169    queries = [
170        "Hello! I'm new here.",
171        "How do I implement a binary search tree in Python?",
172        "What's the best pizza topping?",
173        "Explain quantum computing",
174    ]
175
176    for i, query in enumerate(queries, 1):
177        print("\n" + "=" * 80)
178        print(f"Interaction {i}: {query}")
179        print("=" * 80)
180
181        workflow.print_response(
182            input=query,
183            session_id="user-456",
184            user_id="user-456",
185            stream=True,
186        )
187
188
189if __name__ == "__main__":
190    run_example()

1uv pip install -U kern-ai openai

1python access_session_state_in_router_selector_function.py