1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
|
'''
Copyright (C) 2024-2025 Mo Zhou <lumin@debian.org>
This program is free software: you can redistribute it and/or modify
it under the terms of the GNU Lesser General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.
This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
GNU Lesser General Public License for more details.
You should have received a copy of the GNU Lesser General Public License
along with this program. If not, see <https://www.gnu.org/licenses/>.
'''
from types import SimpleNamespace
import sys
import os
import numpy as np
import pytest
from debgpt import defaults
from debgpt import retrieval
from debgpt import reader
def test_vectorretriever_add(tmpdir):
conf = defaults.Config()
conf.db = os.path.join(tmpdir, 'test.db')
embedding_frontend = conf.embedding_frontend
api_key = conf[f'{embedding_frontend}_api_key']
if api_key.startswith('your-') and api_key.endswith('-key'):
pytest.skip(f'API Key for {embedding_frontend} not configured')
retriever = retrieval.VectorRetriever(conf)
# add some documents
for i in range(2):
retriever.add(f'temp{i}', f'fruit{i}')
def test_vectorretriever_retrieve_onfly(tmpdir):
conf = defaults.Config()
conf.db = os.path.join(tmpdir, 'test.db')
embedding_frontend = conf.embedding_frontend
api_key = conf[f'{embedding_frontend}_api_key']
if api_key.startswith('your-') and api_key.endswith('-key'):
pytest.skip(f'API Key for {embedding_frontend} not configured')
retriever = retrieval.VectorRetriever(conf)
# on-the-fly retrieval
query = 'fruit'
documents = ['fruit', 'sky', 'orange', 'dog', 'cat', 'apple', 'banana']
results = retriever.retrieve_onfly(query, documents, topk=3)
assert len(results) == 3
for i, result in enumerate(results):
score, source, text = result
assert text in documents
assert score >= 0.0 - 1e-5
assert score <= 1.0 + 1e-5
assert source is not None
if i == 0:
assert text == 'fruit'
assert np.isclose(score, 1.0)
print(results)
def test_vectorretriever_retrieve_from_db(tmpdir):
conf = defaults.Config()
conf.db = os.path.join(tmpdir, 'test.db')
embedding_frontend = conf.embedding_frontend
api_key = conf[f'{embedding_frontend}_api_key']
if api_key.startswith('your-') and api_key.endswith('-key'):
pytest.skip(f'API Key for {embedding_frontend} not configured')
retriever = retrieval.VectorRetriever(conf)
# insert some documents
vectors = retriever.batch_add(
['temp'] * 7,
['fruit', 'sky', 'orange', 'dog', 'cat', 'apple', 'banana'])
assert len(vectors) == 7
# retrieve from db
query = 'fruit'
results = retriever.retrieve_from_db(query, topk=3)
assert len(results) == 3
for i, result in enumerate(results):
score, source, text = result
assert text in ['fruit', 'orange', 'apple', 'banana']
assert score >= 0.0 - 1e-5
assert score <= 1.0 + 1e-5
assert source is not None
if i == 0:
assert text == 'fruit'
assert np.isclose(score, 1.0)
print(results)
def test_retrieval_main(tmpdir):
common_args = ['--db', os.path.join(tmpdir, 'test.db'), '-E', 'random']
retrieval.main([*common_args, 'add', 'x'])
retrieval.main([*common_args, 'add', 'y'])
retrieval.main([*common_args, 'add', 'z'])
retrieval.main([*common_args, 'ret', 'w'])
retrieval.main([*common_args, 'retrieve', 'w'])
|
