Skip to content

Commit f93ceae

Browse files
Simplify memory sample to use session state (#4085)
- Rename UserNameProvider → UserMemoryProvider - Use session state (state dict) instead of instance variables - Use context.extend_instructions() instead of context.instructions.append() - Use DEFAULT_SOURCE_ID class attribute - Fix imports to use public agent_framework API - Add session state inspection at end of sample Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
1 parent 93bcc4a commit f93ceae

1 file changed

Lines changed: 29 additions & 25 deletions

File tree

python/samples/01-get-started/04_memory.py

Lines changed: 29 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
import os
55
from typing import Any
66

7-
from agent_framework._sessions import AgentSession, BaseContextProvider, SessionContext
7+
from agent_framework import AgentSession, BaseContextProvider, SessionContext
88
from agent_framework.azure import AzureOpenAIResponsesClient
99
from azure.identity import AzureCliCredential
1010
from dotenv import load_dotenv
@@ -13,11 +13,11 @@
1313
load_dotenv()
1414

1515
"""
16-
Agent Memory with Context Providers
16+
Agent Memory with Context Providers and Session State
1717
18-
Context providers let you inject dynamic instructions and context into each
19-
agent invocation. This sample defines a simple provider that tracks the user's
20-
name and enriches every request with personalization instructions.
18+
Context providers inject dynamic context into each agent call. This sample
19+
shows a provider that stores the user's name in session state and personalizes
20+
responses — the name persists across turns via the session.
2121
2222
Environment variables:
2323
AZURE_AI_PROJECT_ENDPOINT — Your Azure AI Foundry project endpoint
@@ -26,41 +26,45 @@
2626

2727

2828
# <context_provider>
29-
class UserNameProvider(BaseContextProvider):
30-
"""A simple context provider that remembers the user's name."""
29+
class UserMemoryProvider(BaseContextProvider):
30+
"""A context provider that remembers user info in session state."""
3131

32-
def __init__(self) -> None:
33-
super().__init__(source_id="user-name-provider")
34-
self.user_name: str | None = None
32+
DEFAULT_SOURCE_ID = "user_memory"
3533

3634
async def before_run(
3735
self,
3836
*,
3937
agent: Any,
40-
session: AgentSession,
38+
session: AgentSession | None,
4139
context: SessionContext,
4240
state: dict[str, Any],
4341
) -> None:
44-
"""Called before each agent invocation — add extra instructions."""
45-
if self.user_name:
46-
context.instructions.append(f"The user's name is {self.user_name}. Always address them by name.")
42+
"""Inject personalization instructions based on stored user info."""
43+
user_name = state.get("user_name")
44+
if user_name:
45+
context.extend_instructions(
46+
self.source_id,
47+
f"The user's name is {user_name}. Always address them by name.",
48+
)
4749
else:
48-
context.instructions.append("You don't know the user's name yet. Ask for it politely.")
50+
context.extend_instructions(
51+
self.source_id,
52+
"You don't know the user's name yet. Ask for it politely.",
53+
)
4954

5055
async def after_run(
5156
self,
5257
*,
5358
agent: Any,
54-
session: AgentSession,
59+
session: AgentSession | None,
5560
context: SessionContext,
5661
state: dict[str, Any],
5762
) -> None:
58-
"""Called after each agent invocation — extract information."""
63+
"""Extract and store user info in session state after each call."""
5964
for msg in context.input_messages:
6065
text = msg.text if hasattr(msg, "text") else ""
6166
if isinstance(text, str) and "my name is" in text.lower():
62-
# Simple extraction — production code should use structured extraction
63-
self.user_name = text.lower().split("my name is")[-1].strip().split()[0].capitalize()
67+
state["user_name"] = text.lower().split("my name is")[-1].strip().split()[0].capitalize()
6468
# </context_provider>
6569

6670

@@ -73,12 +77,10 @@ async def main() -> None:
7377
credential=credential,
7478
)
7579

76-
memory = UserNameProvider()
77-
7880
agent = client.as_agent(
7981
name="MemoryAgent",
8082
instructions="You are a friendly assistant.",
81-
context_providers=[memory],
83+
context_providers=[UserMemoryProvider()],
8284
)
8385
# </create_agent>
8486

@@ -89,15 +91,17 @@ async def main() -> None:
8991
result = await agent.run("Hello! What's the square root of 9?", session=session)
9092
print(f"Agent: {result}\n")
9193

92-
# Now provide the name — the provider extracts and stores it
94+
# Now provide the name — the provider stores it in session state
9395
result = await agent.run("My name is Alice", session=session)
9496
print(f"Agent: {result}\n")
9597

96-
# Subsequent calls are personalized
98+
# Subsequent calls are personalized — name persists via session state
9799
result = await agent.run("What is 2 + 2?", session=session)
98100
print(f"Agent: {result}\n")
99101

100-
print(f"[Memory] Stored user name: {memory.user_name}")
102+
# Inspect session state to see what the provider stored
103+
provider_state = session.state.get("user_memory", {})
104+
print(f"[Session State] Stored user name: {provider_state.get('user_name')}")
101105
# </run_with_memory>
102106

103107

0 commit comments

Comments
 (0)