mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-12-22 18:22:26 +00:00
99 lines
3.1 KiB
Python
99 lines
3.1 KiB
Python
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
|
# All rights reserved.
|
|
#
|
|
# This source code is licensed under the terms described in the LICENSE file in
|
|
# the root directory of this source tree.
|
|
|
|
import os
|
|
|
|
import pytest
|
|
|
|
from llama_stack.apis.memory import MemoryBankDocument
|
|
from llama_stack.apis.memory_banks import VectorMemoryBankParams
|
|
from llama_stack.apis.tools import ToolInvocationResult
|
|
from llama_stack.providers.datatypes import Api
|
|
|
|
|
|
@pytest.fixture
|
|
def sample_search_query():
|
|
return "What are the latest developments in quantum computing?"
|
|
|
|
|
|
@pytest.fixture
|
|
def sample_documents():
|
|
urls = [
|
|
"memory_optimizations.rst",
|
|
"chat.rst",
|
|
"llama3.rst",
|
|
"datasets.rst",
|
|
"qat_finetune.rst",
|
|
"lora_finetune.rst",
|
|
]
|
|
return [
|
|
MemoryBankDocument(
|
|
document_id=f"num-{i}",
|
|
content=f"https://raw.githubusercontent.com/pytorch/torchtune/main/docs/source/tutorials/{url}",
|
|
mime_type="text/plain",
|
|
metadata={},
|
|
)
|
|
for i, url in enumerate(urls)
|
|
]
|
|
|
|
|
|
class TestTools:
|
|
@pytest.mark.asyncio
|
|
async def test_brave_search_tool(self, tools_stack, sample_search_query):
|
|
"""Test the Brave search tool functionality."""
|
|
if "TAVILY_SEARCH_API_KEY" not in os.environ:
|
|
pytest.skip("TAVILY_SEARCH_API_KEY not set, skipping test")
|
|
|
|
tools_impl = tools_stack.impls[Api.tool_runtime]
|
|
|
|
# Execute the tool
|
|
response = await tools_impl.invoke_tool(
|
|
tool_name="brave_search", tool_args={"query": sample_search_query}
|
|
)
|
|
|
|
# Verify the response
|
|
assert isinstance(response, ToolInvocationResult)
|
|
assert response.content is not None
|
|
assert len(response.content) > 0
|
|
assert isinstance(response.content, str)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_memory_tool(self, tools_stack, sample_documents):
|
|
"""Test the memory tool functionality."""
|
|
memory_banks_impl = tools_stack.impls[Api.memory_banks]
|
|
memory_impl = tools_stack.impls[Api.memory]
|
|
tools_impl = tools_stack.impls[Api.tools]
|
|
|
|
# Register memory bank
|
|
await memory_banks_impl.register_memory_bank(
|
|
memory_bank_id="test_memory_bank",
|
|
params=VectorMemoryBankParams(
|
|
embedding_model="all-MiniLM-L6-v2",
|
|
chunk_size_in_tokens=512,
|
|
overlap_size_in_tokens=64,
|
|
),
|
|
provider_id="faiss",
|
|
)
|
|
|
|
# Insert documents into memory
|
|
memory_impl.insert_documents(
|
|
bank_id="test_memory_bank",
|
|
documents=sample_documents,
|
|
)
|
|
|
|
# Execute the memory tool
|
|
response = await tools_impl.invoke_tool(
|
|
tool_name="memory",
|
|
tool_args={
|
|
"query": "What are the main topics covered in the documentation?",
|
|
},
|
|
)
|
|
|
|
# Verify the response
|
|
assert isinstance(response, ToolInvocationResult)
|
|
assert response.content is not None
|
|
assert len(response.content) > 0
|
|
assert isinstance(response.content, str)
|