autogen/test/agentchat/contrib/vectordb/test_vectordb_utils.py

50 lines
1.5 KiB
Python
Raw Normal View History

#!/usr/bin/env python3 -m pytest
import os
import sys
import pytest
from autogen.agentchat.contrib.vectordb.utils import chroma_results_to_query_results, filter_results_by_distance
def test_retrieve_config():
results = [
[("id1", 1), ("id2", 2)],
[("id3", 2), ("id4", 3)],
]
print(filter_results_by_distance(results, 2.1))
filter_results = [
[("id1", 1), ("id2", 2)],
[("id3", 2)],
]
assert filter_results == filter_results_by_distance(results, 2.1)
def test_chroma_results_to_query_results():
data_dict = {
"key1s": [[1, 2, 3], [4, 5, 6], [7, 8, 9]],
"key2s": [["a", "b", "c"], ["c", "d", "e"], ["e", "f", "g"]],
"key3s": None,
"key4s": [["x", "y", "z"], ["1", "2", "3"], ["4", "5", "6"]],
"distances": [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6], [0.7, 0.8, 0.9]],
}
results = [
[
({"key1": 1, "key2": "a", "key4": "x"}, 0.1),
({"key1": 2, "key2": "b", "key4": "y"}, 0.2),
({"key1": 3, "key2": "c", "key4": "z"}, 0.3),
],
[
({"key1": 4, "key2": "c", "key4": "1"}, 0.4),
({"key1": 5, "key2": "d", "key4": "2"}, 0.5),
({"key1": 6, "key2": "e", "key4": "3"}, 0.6),
],
[
({"key1": 7, "key2": "e", "key4": "4"}, 0.7),
({"key1": 8, "key2": "f", "key4": "5"}, 0.8),
({"key1": 9, "key2": "g", "key4": "6"}, 0.9),
],
]
assert chroma_results_to_query_results(data_dict) == results