phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-07-29 07:14:20 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

improve memory test, but it fails on chromadb :/

Browse source
This commit is contained in:
Ashwin Bharambe 2024-10-07 23:09:31 -07:00 • committed by Ashwin Bharambe
parent 4ab6e1b81a
commit f21ad1173e
2 changed files with 56 additions and 13 deletions
Download patch file Download diff file Expand all files Collapse all files

2
llama_stack/providers/adapters/memory/chroma/chroma.py
View file

@ -99,7 +99,7 @@ class ChromaMemoryAdapter(Memory):
memory_bank.type == MemoryBankType.vector.value memory_bank.type == MemoryBankType.vector.value
), f"Only vector banks are supported {memory_bank.type}" ), f"Only vector banks are supported {memory_bank.type}"
collection = await self.client.create_collection( collection = await self.client.get_or_create_collection(
name=memory_bank.identifier, name=memory_bank.identifier,
) )
bank_index = BankWithIndex( bank_index = BankWithIndex(

67
llama_stack/providers/tests/memory/test_memory.py
View file

@ -22,13 +22,29 @@ async def memory_impl():
@pytest.fixture @pytest.fixture
def sample_document(): def sample_documents():
return MemoryBankDocument( return [
document_id="doc1", MemoryBankDocument(
content="This is a sample document for testing.", document_id="doc1",
mime_type="text/plain", content="Python is a high-level programming language.",
metadata={"author": "Test Author"}, metadata={"category": "programming", "difficulty": "beginner"},
) ),
MemoryBankDocument(
document_id="doc2",
content="Machine learning is a subset of artificial intelligence.",
metadata={"category": "AI", "difficulty": "advanced"},
),
MemoryBankDocument(
document_id="doc3",
content="Data structures are fundamental to computer science.",
metadata={"category": "computer science", "difficulty": "intermediate"},
),
MemoryBankDocument(
document_id="doc4",
content="Neural networks are inspired by biological neural networks.",
metadata={"category": "AI", "difficulty": "advanced"},
),
]
async def register_memory_bank(memory_impl: Memory): async def register_memory_bank(memory_impl: Memory):
@ -44,17 +60,44 @@ async def register_memory_bank(memory_impl: Memory):
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_query_documents(memory_impl, sample_document): async def test_query_documents(memory_impl, sample_documents):
with pytest.raises(ValueError): with pytest.raises(ValueError):
await memory_impl.insert_documents("test_bank", [sample_document]) await memory_impl.insert_documents("test_bank", sample_documents)
await register_memory_bank(memory_impl) await register_memory_bank(memory_impl)
await memory_impl.insert_documents("test_bank", [sample_document]) await memory_impl.insert_documents("test_bank", sample_documents)
query = ["sample ", "document"] query1 = "programming language"
response = await memory_impl.query_documents("test_bank", query) response1 = await memory_impl.query_documents("test_bank", query1)
assert_valid_response(response1)
assert any("Python" in chunk.content for chunk in response1.chunks)
# Test case 3: Query with semantic similarity
query3 = "AI and brain-inspired computing"
response3 = await memory_impl.query_documents("test_bank", query3)
assert_valid_response(response3)
assert any("neural networks" in chunk.content.lower() for chunk in response3.chunks)
# Test case 4: Query with limit on number of results
query4 = "computer"
params4 = {"max_chunks": 2}
response4 = await memory_impl.query_documents("test_bank", query4, params4)
assert_valid_response(response4)
assert len(response4.chunks) <= 2
# Test case 5: Query with threshold on similarity score
query5 = "quantum computing" # Not directly related to any document
params5 = {"score_threshold": 0.5}
response5 = await memory_impl.query_documents("test_bank", query5, params5)
assert_valid_response(response5)
assert all(score >= 0.5 for score in response5.scores)
def assert_valid_response(response: QueryDocumentsResponse):
assert isinstance(response, QueryDocumentsResponse) assert isinstance(response, QueryDocumentsResponse)
assert len(response.chunks) > 0 assert len(response.chunks) > 0
assert len(response.scores) > 0 assert len(response.scores) > 0
assert len(response.chunks) == len(response.scores) assert len(response.chunks) == len(response.scores)
for chunk in response.chunks:
assert isinstance(chunk.content, str)
assert chunk.document_id is not None