minor fixes

This commit is contained in:
Dinesh Yeduguru 2025-01-07 16:37:39 -08:00
parent 6632d7e410
commit c3865faf37
6 changed files with 19 additions and 9 deletions

View file

@ -9,7 +9,7 @@ from typing import Dict, List
from uuid import uuid4
import pytest
from llama_stack_client.lib.agents.agent import Agent, maybe_register_memory_tool
from llama_stack_client.lib.agents.agent import Agent
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
@ -293,9 +293,18 @@ def test_rag_agent(llama_stack_client, agent_config):
for i, url in enumerate(urls)
]
tool_name, memory_bank_id = maybe_register_memory_tool(llama_stack_client)
agent_config["tools"].append(tool_name)
agent_config["tools"].append("builtin::memory")
agent = Agent(llama_stack_client, agent_config)
memory_bank_id = "test-memory-bank"
llama_stack_client.memory_banks.register(
memory_bank_id=memory_bank_id,
params={
"memory_bank_type": "vector",
"embedding_model": "all-MiniLM-L6-v2",
"chunk_size_in_tokens": 512,
"overlap_size_in_tokens": 64,
},
)
llama_stack_client.memory.insert(
bank_id=memory_bank_id,
documents=documents,