phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-08-05 10:13:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

use maybe_register_memory

Browse source
This commit is contained in:
Dinesh Yeduguru 2025-01-06 12:50:36 -08:00
parent 9efe30c9d3
commit f3304abfba
1 changed files with 5 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

8
tests/client-sdk/agents/test_agents.py
View file

@ -9,7 +9,7 @@ from typing import Dict, List
from uuid import uuid4 from uuid import uuid4
import pytest 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.client_tool import ClientTool
from llama_stack_client.lib.agents.event_logger import EventLogger from llama_stack_client.lib.agents.event_logger import EventLogger
from llama_stack_client.types import ToolResponseMessage from llama_stack_client.types import ToolResponseMessage
@ -227,7 +227,8 @@ def test_code_execution(llama_stack_client):
"overlap_size_in_tokens": 64, "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) codex_agent = Agent(llama_stack_client, agent_config)
session_id = codex_agent.create_session("test-session") 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) 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) agent = Agent(llama_stack_client, agent_config)
llama_stack_client.memory.insert( llama_stack_client.memory.insert(
bank_id=memory_bank_id, bank_id=memory_bank_id,