From f3304abfba9ef251eca71ae177c76e8a52c89bf1 Mon Sep 17 00:00:00 2001 From: Dinesh Yeduguru Date: Mon, 6 Jan 2025 12:50:36 -0800 Subject: [PATCH] use maybe_register_memory --- tests/client-sdk/agents/test_agents.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/tests/client-sdk/agents/test_agents.py b/tests/client-sdk/agents/test_agents.py index a77bb6cab..522a8a4eb 100644 --- a/tests/client-sdk/agents/test_agents.py +++ b/tests/client-sdk/agents/test_agents.py @@ -9,7 +9,7 @@ from typing import Dict, List from uuid import uuid4 import pytest -from llama_stack_client.lib.agents.agent import Agent, AugmentConfigWithMemoryTool +from llama_stack_client.lib.agents.agent import Agent, maybe_register_memory_tool from llama_stack_client.lib.agents.client_tool import ClientTool from llama_stack_client.lib.agents.event_logger import EventLogger from llama_stack_client.types import ToolResponseMessage @@ -227,7 +227,8 @@ def test_code_execution(llama_stack_client): "overlap_size_in_tokens": 64, }, ) - AugmentConfigWithMemoryTool(agent_config, llama_stack_client) + tool_name, _ = maybe_register_memory_tool(llama_stack_client) + agent_config["tools"].append(tool_name) codex_agent = Agent(llama_stack_client, agent_config) session_id = codex_agent.create_session("test-session") @@ -324,7 +325,8 @@ def test_rag_agent(llama_stack_client, agent_config): for i, url in enumerate(urls) ] - memory_bank_id = AugmentConfigWithMemoryTool(agent_config, llama_stack_client) + tool_name, memory_bank_id = maybe_register_memory_tool(llama_stack_client) + agent_config["tools"].append(tool_name) agent = Agent(llama_stack_client, agent_config) llama_stack_client.memory.insert( bank_id=memory_bank_id,