57 lines
1.5 KiB
Python
57 lines
1.5 KiB
Python
|
"""Tests for vector memory with cosine similarity."""
|
||
|
|
|
||
|
|
from fusionagi.memory.service import VectorMemory
|
||
|
|
|
||
|
|
|
||
|
|
def test_add_and_search():
|
||
|
|
vm = VectorMemory()
|
||
|
|
vm.add("doc1", [1.0, 0.0, 0.0], {"text": "hello"})
|
||
|
|
vm.add("doc2", [0.0, 1.0, 0.0], {"text": "world"})
|
||
|
|
results = vm.search([1.0, 0.0, 0.0], top_k=1)
|
||
|
|
assert len(results) == 1
|
||
|
|
assert results[0]["id"] == "doc1"
|
||
|
|
assert results[0]["score"] > 0.99
|
||
|
|
|
||
|
|
|
||
|
|
def test_cosine_similarity():
|
||
|
|
assert abs(VectorMemory._cosine_similarity([1, 0], [1, 0]) - 1.0) < 0.001
|
||
|
|
assert abs(VectorMemory._cosine_similarity([1, 0], [0, 1])) < 0.001
|
||
|
|
assert abs(VectorMemory._cosine_similarity([1, 1], [1, 1]) - 1.0) < 0.001
|
||
|
|
|
||
|
|
|
||
|
|
def test_zero_vector():
|
||
|
|
assert VectorMemory._cosine_similarity([0, 0], [1, 0]) == 0.0
|
||
|
|
|
||
|
|
|
||
|
|
def test_delete():
|
||
|
|
vm = VectorMemory()
|
||
|
|
vm.add("doc1", [1.0, 0.0])
|
||
|
|
assert vm.count() == 1
|
||
|
|
assert vm.delete("doc1") is True
|
||
|
|
assert vm.count() == 0
|
||
|
|
|
||
|
|
|
||
|
|
def test_max_entries():
|
||
|
|
vm = VectorMemory(max_entries=2)
|
||
|
|
vm.add("a", [1.0])
|
||
|
|
vm.add("b", [2.0])
|
||
|
|
vm.add("c", [3.0])
|
||
|
|
assert vm.count() == 2
|
||
|
|
|
||
|
|
|
||
|
|
def test_search_top_k():
|
||
|
|
vm = VectorMemory()
|
||
|
|
vm.add("a", [1.0, 0.0])
|
||
|
|
vm.add("b", [0.9, 0.1])
|
||
|
|
vm.add("c", [0.0, 1.0])
|
||
|
|
results = vm.search([1.0, 0.0], top_k=2)
|
||
|
|
assert len(results) == 2
|
||
|
|
assert results[0]["id"] == "a"
|
||
|
|
|
||
|
|
|
||
|
|
def test_search_with_metadata():
|
||
|
|
vm = VectorMemory()
|
||
|
|
vm.add("doc", [1.0], {"key": "value"})
|
||
|
|
results = vm.search([1.0])
|
||
|
|
assert results[0]["metadata"]["key"] == "value"
|