mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Python: add RedisContextProvider (#716)
* Setting up * Readme * Add redis tests path to all-tests * First pass integration * Keep provider convention * First pass integration * add redis integration tests * update README.md * Add basic sample for redis integration * Add partitioning, add partition-aware tests, improve sample script * Fix code quality check * Try to resolve pytest check * Try to identify if pytest is the cause of failed checks * Re-enable tests * Rename redis test file * Removing some tests to narrow down issue * Revert, no difference * Delete temp files * Starting refactor of RedisProvider * Build dynamic schema builder, still need to do dynamic embedding model config * Add scope control * Complete first pass functionality with OpenAI + HF vectors -> Tests, Samples, Demo to follow * Fix code quality * attempt to identify rootcause of failed test * attempt to identify rootcause of failed test * Attempt to resolve code quality fail * Update pyproject.toml for foundry to pin azure-ai-projects == 1.1.0b3,azure-ai-agents == 1.2.0b3 * Add tests for redisprovider * Remove invalid tests * Add API key handling for openai vectorizer * Update uv.locl * Use master uv.lock * Begin sample file, add lazy index creation, fix faulty override * Index drop and reinit depends on drop_redis_index not overwrite * Add samples, threading included, escaped queries, verify threading works, sample README.md * Refactor filters * Opinionated vars * Allow filter expression combination * Try inline stubs for mypy * Address mypy errors * Better docstrings, tweaks for feedback * Tweak example 3 in redis_threads.py sample * adjust confusing name * Enrich docstrings * Add descriptions and comments to samples, externalize vectorizer choice, remove nltk and sentencetransformers dependnecy * Add descriptions and comments to samples, externalize vectorizer choice, remove nltk and sentencetransformers dependnecy * Incorporate initial feedback from dmytrostruk * Fix uv.lock * Attempt to resolve conflict * Use remote .tomls * Sanity check * fix tests * Remove hardcoded API key from samples * Fix incorrect env var * Make add and redis_search private * Fix tests relying on private funcs * Expand tests * Explainer comments to each test * Add a 'get_conversation_history' function to RedisProvider - This just returns messages in sequential order. Added 'created_at_*' timestamps to facilitate sequential recovery. function has to be manually invoked by user * Add agent-framework-redis to python/pyproject.toml * Remove get_conversation_history * improve redis context provider with pydantic techniques and safe index handling patterns * add RedisChatMessageStore * remove integration test :( * fix mypy error * Remove unused params * Redo schema validation to be order-invariant, handle attrs (previously throwing errors due to strict ==) * Expand explanation * Add ChatMessageStore example * Fix comments in redis_conversation.py * Resolving uv.lock conflict, update to match main * Fix test in redis provider * Apply suggestion from @ekzhu * Update python/packages/main/pyproject.toml --------- Co-authored-by: Tyler Hutcherson <tyler.hutcherson@redis.com> Co-authored-by: Eric Zhu <ekzhu@users.noreply.github.com>
This commit is contained in:
Han
GitHub
@@ -8,3 +8,4 @@ This folder contains examples demonstrating different ways to manage conversatio
|
||||
|------|-------------|
|
||||
| [`custom_chat_message_store_thread.py`](custom_chat_message_store_thread.py) | Demonstrates how to implement a custom `ChatMessageStore` for persisting conversation history. Shows how to create a custom store with serialization/deserialization capabilities and integrate it with agents for thread management across multiple sessions. |
|
||||
| [`suspend_resume_thread.py`](suspend_resume_thread.py) | Shows how to suspend and resume conversation threads, allowing you to save the state of a conversation and continue it later. This is useful for long-running conversations or when you need to persist conversation state across application restarts. |
|
||||
| [`redis_chat_message_store_thread.py`](redis_chat_message_store_thread.py) | Comprehensive examples of using the Redis-backed `RedisChatMessageStore` for persistent conversation storage. Covers basic usage, user session management, conversation persistence across app restarts, thread serialization, and automatic message trimming. Requires Redis server and demonstrates production-ready patterns for scalable chat applications. |
|
||||
|
||||
@@ -0,0 +1,318 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
from uuid import uuid4
|
||||
|
||||
from agent_framework import AgentThread
|
||||
from agent_framework._threads import deserialize_thread_state
|
||||
from agent_framework.openai import OpenAIChatClient
|
||||
from agent_framework.redis import RedisChatMessageStore
|
||||
|
||||
|
||||
async def example_basic_redis_store() -> None:
|
||||
"""Basic example of using Redis chat message store."""
|
||||
print("=== Basic Redis Chat Message Store Example ===")
|
||||
|
||||
# Create Redis store with auto-generated thread ID
|
||||
redis_store = RedisChatMessageStore(
|
||||
redis_url="redis://localhost:6379",
|
||||
# thread_id will be auto-generated if not provided
|
||||
)
|
||||
|
||||
print(f"Created store with thread ID: {redis_store.thread_id}")
|
||||
|
||||
# Create thread with Redis store
|
||||
thread = AgentThread(message_store=redis_store)
|
||||
|
||||
# Create agent
|
||||
agent = OpenAIChatClient().create_agent(
|
||||
name="RedisBot",
|
||||
instructions="You are a helpful assistant that remembers our conversation using Redis.",
|
||||
)
|
||||
|
||||
# Have a conversation
|
||||
print("\n--- Starting conversation ---")
|
||||
query1 = "Hello! My name is Alice and I love pizza."
|
||||
print(f"User: {query1}")
|
||||
response1 = await agent.run(query1, thread=thread)
|
||||
print(f"Agent: {response1.text}")
|
||||
|
||||
query2 = "What do you remember about me?"
|
||||
print(f"User: {query2}")
|
||||
response2 = await agent.run(query2, thread=thread)
|
||||
print(f"Agent: {response2.text}")
|
||||
|
||||
# Show messages are stored in Redis
|
||||
messages = await redis_store.list_messages()
|
||||
print(f"\nTotal messages in Redis: {len(messages)}")
|
||||
|
||||
# Cleanup
|
||||
await redis_store.clear()
|
||||
await redis_store.aclose()
|
||||
print("Cleaned up Redis data\n")
|
||||
|
||||
|
||||
async def example_user_session_management() -> None:
|
||||
"""Example of managing user sessions with Redis."""
|
||||
print("=== User Session Management Example ===")
|
||||
|
||||
user_id = "alice_123"
|
||||
session_id = f"session_{uuid4()}"
|
||||
|
||||
# Create Redis store for specific user session
|
||||
def create_user_session_store():
|
||||
return RedisChatMessageStore(
|
||||
redis_url="redis://localhost:6379",
|
||||
thread_id=f"user_{user_id}_{session_id}",
|
||||
max_messages=10 # Keep only last 10 messages
|
||||
)
|
||||
|
||||
# Create agent with factory pattern
|
||||
agent = OpenAIChatClient().create_agent(
|
||||
name="SessionBot",
|
||||
instructions="You are a helpful assistant. Keep track of user preferences.",
|
||||
chat_message_store_factory=create_user_session_store,
|
||||
)
|
||||
|
||||
# Start conversation
|
||||
thread = agent.get_new_thread()
|
||||
|
||||
print(f"Started session for user {user_id}")
|
||||
if hasattr(thread.message_store, 'thread_id'):
|
||||
print(f"Thread ID: {thread.message_store.thread_id}") # type: ignore[union-attr]
|
||||
|
||||
# Simulate conversation
|
||||
queries = [
|
||||
"Hi, I'm Alice and I prefer vegetarian food.",
|
||||
"What restaurants would you recommend?",
|
||||
"I also love Italian cuisine.",
|
||||
"Can you remember my food preferences?"
|
||||
]
|
||||
|
||||
for i, query in enumerate(queries, 1):
|
||||
print(f"\n--- Message {i} ---")
|
||||
print(f"User: {query}")
|
||||
response = await agent.run(query, thread=thread)
|
||||
print(f"Agent: {response.text}")
|
||||
|
||||
# Show persistent storage
|
||||
if thread.message_store:
|
||||
messages = await thread.message_store.list_messages() # type: ignore[union-attr]
|
||||
print(f"\nMessages stored for user {user_id}: {len(messages)}")
|
||||
|
||||
# Cleanup
|
||||
if thread.message_store:
|
||||
await thread.message_store.clear() # type: ignore[union-attr]
|
||||
await thread.message_store.aclose() # type: ignore[union-attr]
|
||||
print("Cleaned up session data\n")
|
||||
|
||||
|
||||
async def example_conversation_persistence() -> None:
|
||||
"""Example of conversation persistence across application restarts."""
|
||||
print("=== Conversation Persistence Example ===")
|
||||
|
||||
conversation_id = "persistent_chat_001"
|
||||
|
||||
# Phase 1: Start conversation
|
||||
print("--- Phase 1: Starting conversation ---")
|
||||
store1 = RedisChatMessageStore(
|
||||
redis_url="redis://localhost:6379",
|
||||
thread_id=conversation_id,
|
||||
)
|
||||
|
||||
thread1 = AgentThread(message_store=store1)
|
||||
agent = OpenAIChatClient().create_agent(
|
||||
name="PersistentBot",
|
||||
instructions="You are a helpful assistant. Remember our conversation history.",
|
||||
)
|
||||
|
||||
# Start conversation
|
||||
query1 = "Hello! I'm working on a Python project about machine learning."
|
||||
print(f"User: {query1}")
|
||||
response1 = await agent.run(query1, thread=thread1)
|
||||
print(f"Agent: {response1.text}")
|
||||
|
||||
query2 = "I'm specifically interested in neural networks."
|
||||
print(f"User: {query2}")
|
||||
response2 = await agent.run(query2, thread=thread1)
|
||||
print(f"Agent: {response2.text}")
|
||||
|
||||
print(f"Stored {len(await store1.list_messages())} messages in Redis")
|
||||
await store1.aclose()
|
||||
|
||||
# Phase 2: Resume conversation (simulating app restart)
|
||||
print("\n--- Phase 2: Resuming conversation (after 'restart') ---")
|
||||
store2 = RedisChatMessageStore(
|
||||
redis_url="redis://localhost:6379",
|
||||
thread_id=conversation_id, # Same thread ID
|
||||
)
|
||||
|
||||
thread2 = AgentThread(message_store=store2)
|
||||
|
||||
# Continue conversation - agent should remember context
|
||||
query3 = "What was I working on before?"
|
||||
print(f"User: {query3}")
|
||||
response3 = await agent.run(query3, thread=thread2)
|
||||
print(f"Agent: {response3.text}")
|
||||
|
||||
query4 = "Can you suggest some Python libraries for neural networks?"
|
||||
print(f"User: {query4}")
|
||||
response4 = await agent.run(query4, thread=thread2)
|
||||
print(f"Agent: {response4.text}")
|
||||
|
||||
print(f"Total messages after resuming: {len(await store2.list_messages())}")
|
||||
|
||||
# Cleanup
|
||||
await store2.clear()
|
||||
await store2.aclose()
|
||||
print("Cleaned up persistent data\n")
|
||||
|
||||
|
||||
async def example_thread_serialization() -> None:
|
||||
"""Example of thread state serialization and deserialization."""
|
||||
print("=== Thread Serialization Example ===")
|
||||
|
||||
# Create initial thread with Redis store
|
||||
original_store = RedisChatMessageStore(
|
||||
redis_url="redis://localhost:6379",
|
||||
thread_id="serialization_test",
|
||||
max_messages=50,
|
||||
)
|
||||
|
||||
original_thread = AgentThread(message_store=original_store)
|
||||
|
||||
agent = OpenAIChatClient().create_agent(
|
||||
name="SerializationBot",
|
||||
instructions="You are a helpful assistant.",
|
||||
)
|
||||
|
||||
# Have initial conversation
|
||||
print("--- Initial conversation ---")
|
||||
query1 = "Hello! I'm testing serialization."
|
||||
print(f"User: {query1}")
|
||||
response1 = await agent.run(query1, thread=original_thread)
|
||||
print(f"Agent: {response1.text}")
|
||||
|
||||
# Serialize thread state
|
||||
serialized_thread = await original_thread.serialize()
|
||||
print(f"\nSerialized thread state: {serialized_thread}")
|
||||
|
||||
# Close original connection
|
||||
await original_store.aclose()
|
||||
|
||||
# Deserialize thread state (simulating loading from database/file)
|
||||
print("\n--- Deserializing thread state ---")
|
||||
|
||||
# Create a new thread with the same Redis store type
|
||||
# This ensures the correct store type is used for deserialization
|
||||
restored_store = RedisChatMessageStore(redis_url="redis://localhost:6379")
|
||||
restored_thread = AgentThread(message_store=restored_store)
|
||||
|
||||
# Deserialize the thread state into the properly typed thread
|
||||
await deserialize_thread_state(restored_thread, serialized_thread)
|
||||
|
||||
# Continue conversation with restored thread
|
||||
query2 = "Do you remember what I said about testing?"
|
||||
print(f"User: {query2}")
|
||||
response2 = await agent.run(query2, thread=restored_thread)
|
||||
print(f"Agent: {response2.text}")
|
||||
|
||||
# Cleanup
|
||||
if restored_thread.message_store:
|
||||
await restored_thread.message_store.clear() # type: ignore[union-attr]
|
||||
await restored_thread.message_store.aclose() # type: ignore[union-attr]
|
||||
print("Cleaned up serialization test data\n")
|
||||
|
||||
|
||||
async def example_message_limits() -> None:
|
||||
"""Example of automatic message trimming with limits."""
|
||||
print("=== Message Limits Example ===")
|
||||
|
||||
# Create store with small message limit
|
||||
store = RedisChatMessageStore(
|
||||
redis_url="redis://localhost:6379",
|
||||
thread_id="limits_test",
|
||||
max_messages=3, # Keep only 3 most recent messages
|
||||
)
|
||||
|
||||
thread = AgentThread(message_store=store)
|
||||
agent = OpenAIChatClient().create_agent(
|
||||
name="LimitBot",
|
||||
instructions="You are a helpful assistant with limited memory.",
|
||||
)
|
||||
|
||||
# Send multiple messages to test trimming
|
||||
messages = [
|
||||
"Message 1: Hello!",
|
||||
"Message 2: How are you?",
|
||||
"Message 3: What's the weather?",
|
||||
"Message 4: Tell me a joke.",
|
||||
"Message 5: This should trigger trimming.",
|
||||
]
|
||||
|
||||
for i, query in enumerate(messages, 1):
|
||||
print(f"\n--- Sending message {i} ---")
|
||||
print(f"User: {query}")
|
||||
response = await agent.run(query, thread=thread)
|
||||
print(f"Agent: {response.text}")
|
||||
|
||||
stored_messages = await store.list_messages()
|
||||
print(f"Messages in store: {len(stored_messages)}")
|
||||
if len(stored_messages) > 0:
|
||||
print(f"Oldest message: {stored_messages[0].text[:30]}...")
|
||||
|
||||
# Final check
|
||||
final_messages = await store.list_messages()
|
||||
print(f"\nFinal message count: {len(final_messages)} (should be <= 6: 3 messages × 2 per exchange)")
|
||||
|
||||
# Cleanup
|
||||
await store.clear()
|
||||
await store.aclose()
|
||||
print("Cleaned up limits test data\n")
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
"""Run all Redis chat message store examples."""
|
||||
print("Redis Chat Message Store Examples")
|
||||
print("=" * 50)
|
||||
print("Prerequisites:")
|
||||
print("- Redis server running on localhost:6379")
|
||||
print("- OPENAI_API_KEY environment variable set")
|
||||
print("=" * 50)
|
||||
|
||||
# Check prerequisites
|
||||
if not os.getenv("OPENAI_API_KEY"):
|
||||
print("ERROR: OPENAI_API_KEY environment variable not set")
|
||||
return
|
||||
|
||||
try:
|
||||
# Test Redis connection
|
||||
test_store = RedisChatMessageStore(redis_url="redis://localhost:6379")
|
||||
connection_ok = await test_store.ping()
|
||||
await test_store.aclose()
|
||||
if not connection_ok:
|
||||
raise Exception("Redis ping failed")
|
||||
print("✓ Redis connection successful\n")
|
||||
except Exception as e:
|
||||
print(f"ERROR: Cannot connect to Redis: {e}")
|
||||
print("Please ensure Redis is running on localhost:6379")
|
||||
return
|
||||
|
||||
try:
|
||||
# Run all examples
|
||||
await example_basic_redis_store()
|
||||
await example_user_session_management()
|
||||
await example_conversation_persistence()
|
||||
await example_thread_serialization()
|
||||
await example_message_limits()
|
||||
|
||||
print("All examples completed successfully!")
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error running examples: {e}")
|
||||
raise
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
Reference in New Issue
Block a user