|
| 1 | +# Copyright (c) Microsoft. All rights reserved. |
| 2 | +import asyncio |
| 3 | +import os |
| 4 | +import uuid |
| 5 | + |
| 6 | +from agent_framework.azure import AzureAIClient |
| 7 | +from azure.ai.projects.aio import AIProjectClient |
| 8 | +from azure.ai.projects.models import MemoryStoreDefaultDefinition, MemoryStoreDefaultOptions |
| 9 | +from azure.identity.aio import AzureCliCredential |
| 10 | + |
| 11 | +""" |
| 12 | +Azure AI Agent with Memory Search Example |
| 13 | +
|
| 14 | +This sample demonstrates usage of AzureAIClient with memory search capabilities |
| 15 | +to retrieve relevant past user messages and maintain conversation context across sessions. |
| 16 | +It shows explicit memory store creation using Azure AI Projects client and agent creation |
| 17 | +using the Agent Framework. |
| 18 | +
|
| 19 | +Prerequisites: |
| 20 | +1. Set AZURE_AI_PROJECT_ENDPOINT and AZURE_AI_MODEL_DEPLOYMENT_NAME environment variables. |
| 21 | +2. Set AZURE_AI_CHAT_MODEL_DEPLOYMENT_NAME for the memory chat model. |
| 22 | +3. Set AZURE_AI_EMBEDDING_MODEL_DEPLOYMENT_NAME for the memory embedding model. |
| 23 | +4. Deploy both a chat model (e.g. gpt-4.1) and an embedding model (e.g. text-embedding-3-small). |
| 24 | +""" |
| 25 | + |
| 26 | + |
| 27 | +async def main() -> None: |
| 28 | + endpoint = os.environ["AZURE_AI_PROJECT_ENDPOINT"] |
| 29 | + # Generate a unique memory store name to avoid conflicts |
| 30 | + memory_store_name = f"agent_framework_memory_store_{uuid.uuid4().hex[:8]}" |
| 31 | + |
| 32 | + async with AzureCliCredential() as credential: |
| 33 | + # Create the memory store using Azure AI Projects client |
| 34 | + async with AIProjectClient(endpoint=endpoint, credential=credential) as project_client: |
| 35 | + # Create a memory store using proper model classes |
| 36 | + memory_store_definition = MemoryStoreDefaultDefinition( |
| 37 | + chat_model=os.environ["AZURE_AI_CHAT_MODEL_DEPLOYMENT_NAME"], |
| 38 | + embedding_model=os.environ["AZURE_AI_EMBEDDING_MODEL_DEPLOYMENT_NAME"], |
| 39 | + options=MemoryStoreDefaultOptions(user_profile_enabled=True, chat_summary_enabled=True), |
| 40 | + ) |
| 41 | + |
| 42 | + memory_store = await project_client.memory_stores.create( |
| 43 | + name=memory_store_name, |
| 44 | + description="Memory store for Agent Framework conversations", |
| 45 | + definition=memory_store_definition, |
| 46 | + ) |
| 47 | + print(f"Created memory store: {memory_store.name} ({memory_store.id}): {memory_store.description}") |
| 48 | + |
| 49 | + # Then, create the agent using Agent Framework |
| 50 | + async with AzureAIClient(async_credential=credential).create_agent( |
| 51 | + name="MyMemoryAgent", |
| 52 | + instructions="""You are a helpful assistant that remembers past conversations. |
| 53 | + Use the memory search tool to recall relevant information from previous interactions.""", |
| 54 | + tools={ |
| 55 | + "type": "memory_search", |
| 56 | + "memory_store_name": memory_store.name, |
| 57 | + "scope": "user_123", |
| 58 | + "update_delay": 1, # Wait 1 second before updating memories (use higher value in production) |
| 59 | + }, |
| 60 | + ) as agent: |
| 61 | + # First interaction - establish some preferences |
| 62 | + print("=== First conversation ===") |
| 63 | + query1 = "I prefer dark roast coffee" |
| 64 | + print(f"User: {query1}") |
| 65 | + result1 = await agent.run(query1) |
| 66 | + print(f"Agent: {result1}\n") |
| 67 | + |
| 68 | + # Wait for memories to be processed |
| 69 | + print("Waiting for memories to be stored...") |
| 70 | + await asyncio.sleep(5) # Reduced wait time for demo purposes |
| 71 | + |
| 72 | + # Second interaction - test memory recall |
| 73 | + print("=== Second conversation ===") |
| 74 | + query2 = "Please order my usual coffee" |
| 75 | + print(f"User: {query2}") |
| 76 | + result2 = await agent.run(query2) |
| 77 | + print(f"Agent: {result2}\n") |
| 78 | + |
| 79 | + # Clean up - delete the memory store |
| 80 | + async with AIProjectClient(endpoint=endpoint, credential=credential) as project_client: |
| 81 | + await project_client.memory_stores.delete(memory_store_name) |
| 82 | + print("Memory store deleted") |
| 83 | + |
| 84 | + |
| 85 | +if __name__ == "__main__": |
| 86 | + asyncio.run(main()) |
0 commit comments