fix(qwen): isolate and refresh indexes with code-search prompts
This commit is contained in:
1 parent
5dbb1b21b4
commit
2079792a74
2 files changed
+171
-32
No files matched your search
@@ -1,6 +1,13 @@
|
||||
import importlib.util
|
||||
from array import array
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
import os
|
||||
from pathlib import Path
|
||||
from tempfile import TemporaryDirectory
|
||||
from types import SimpleNamespace
|
||||
import time
|
||||
import unittest
|
||||
from unittest.mock import patch
|
||||
|
||||
|
||||
ROOT = Path(__file__).resolve().parent.parent
|
||||
@@ -10,6 +17,112 @@ SPEC.loader.exec_module(MODULE)
|
||||
|
||||
|
||||
class RetrievalServerTests(unittest.TestCase):
|
||||
@unittest.skipUnless(os.environ.get('QWEN_TEST_MODEL_CACHE'), 'Set QWEN_TEST_MODEL_CACHE to run the real-model check.')
|
||||
def test_real_models(self):
|
||||
os.environ['HF_HOME'] = os.environ['QWEN_TEST_MODEL_CACHE']
|
||||
with TemporaryDirectory() as temporary:
|
||||
root = Path(temporary)
|
||||
(root / 'authentication.md').write_text('Password reset tokens are created by create_reset_token. The reset_password endpoint sends the user a password reset email.', encoding='utf-8')
|
||||
(root / 'weather.md').write_text('The weather service fetches temperature and rainfall forecasts for a city.', encoding='utf-8')
|
||||
index = MODULE.Index(root, root / 'data')
|
||||
try:
|
||||
for query, expected in [('Where are password reset tokens created?', 'authentication.md'),
|
||||
('Which service provides rainfall forecasts?', 'weather.md')]:
|
||||
started = time.monotonic()
|
||||
results = index.search(query, 2)
|
||||
self.assertEqual(results[0]['path'], expected)
|
||||
self.assertEqual(len(results), 2)
|
||||
self.assertGreater(results[0]['score'], results[1]['score'])
|
||||
print(f'Qwen real search: {expected}, {time.monotonic() - started:.2f}s, embedding device {index.encoder().device}, reranker device {index.reranker().device}', flush=True)
|
||||
finally:
|
||||
index.db.close()
|
||||
|
||||
def test_isolation_freshness_and_incremental_encoding(self):
|
||||
with TemporaryDirectory() as temporary:
|
||||
root = Path(temporary)
|
||||
first, second = root / 'first', root / 'second'
|
||||
first.mkdir()
|
||||
second.mkdir()
|
||||
file = first / 'a.md'
|
||||
file.write_text('alpha', encoding='utf-8')
|
||||
(second / 'a.md').write_text('other', encoding='utf-8')
|
||||
indexes = [MODULE.Index(repository, root / 'data') for repository in (first, second)]
|
||||
encoded = []
|
||||
|
||||
def encode(values, **kwargs):
|
||||
self.assertEqual(kwargs.get('prompt_name'), 'query' if values == ['query'] else None)
|
||||
encoded.extend(values)
|
||||
return [array('f', [1.0]) for _ in values]
|
||||
|
||||
for index in indexes:
|
||||
index.encoder = lambda: SimpleNamespace(encode=encode)
|
||||
index.rerank = lambda query, rows, limit: rows[:limit]
|
||||
numpy = SimpleNamespace(dot=lambda a, b: sum(x * y for x, y in zip(a, b)),
|
||||
frombuffer=lambda blob, **kwargs: array('f', blob), float32=None)
|
||||
try:
|
||||
with patch.dict('sys.modules', numpy=numpy):
|
||||
self.assertEqual(indexes[0].search('query', 5)[0]['text'], 'alpha')
|
||||
self.assertEqual(indexes[1].search('query', 5)[0]['text'], 'other')
|
||||
encoded.clear()
|
||||
with ThreadPoolExecutor(max_workers=1) as executor:
|
||||
self.assertEqual(executor.submit(indexes[0].search, 'query', 5).result()[0]['text'], 'alpha')
|
||||
self.assertEqual(encoded, ['query'])
|
||||
timestamp = file.stat()
|
||||
file.write_text('bravo', encoding='utf-8')
|
||||
os.utime(file, ns=(timestamp.st_atime_ns, timestamp.st_mtime_ns))
|
||||
self.assertEqual(indexes[0].search('query', 5)[0]['text'], 'bravo')
|
||||
file.rename(first / 'renamed.md')
|
||||
self.assertEqual(indexes[0].search('query', 5)[0]['path'], 'renamed.md')
|
||||
(first / 'renamed.md').unlink()
|
||||
encoded.clear()
|
||||
self.assertEqual(indexes[0].search('query', 5), [])
|
||||
self.assertEqual(indexes[0].search('query', 5), [])
|
||||
self.assertEqual(encoded, [])
|
||||
self.assertEqual(indexes[1].search('query', 5)[0]['text'], 'other')
|
||||
finally:
|
||||
for index in indexes:
|
||||
index.db.close()
|
||||
|
||||
def test_failed_refresh_preserves_previous_index(self):
|
||||
with TemporaryDirectory() as temporary:
|
||||
root = Path(temporary)
|
||||
(root / 'a.md').write_text('before', encoding='utf-8')
|
||||
index = MODULE.Index(root, root / 'data')
|
||||
index.encoder = lambda: SimpleNamespace(encode=lambda values, **kwargs: [array('f', [1.0]) for _ in values])
|
||||
try:
|
||||
index.rebuild()
|
||||
(root / 'a.md').write_text('after', encoding='utf-8')
|
||||
(root / 'b.md').write_text('fail', encoding='utf-8')
|
||||
|
||||
def encode(values, **kwargs):
|
||||
if values == ['fail']:
|
||||
raise RuntimeError('encoder failed')
|
||||
return [array('f', [1.0]) for _ in values]
|
||||
|
||||
index.encoder = lambda: SimpleNamespace(encode=encode)
|
||||
with self.assertRaisesRegex(RuntimeError, 'encoder failed'):
|
||||
index.rebuild()
|
||||
self.assertEqual(index.db.execute('select path, text from chunks').fetchall(), [('a.md', 'before')])
|
||||
finally:
|
||||
index.db.close()
|
||||
|
||||
def test_scan_prunes_ignored_directories_and_stops_at_limit(self):
|
||||
with TemporaryDirectory() as temporary:
|
||||
root = Path(temporary)
|
||||
index = MODULE.Index(root, root / 'data')
|
||||
(root / 'visible.md').write_text('visible', encoding='utf-8')
|
||||
for name in ('node_modules', '.git', 'generated', 'data'):
|
||||
directory = root / name
|
||||
directory.mkdir(exist_ok=True)
|
||||
(directory / 'hidden.md').write_text('hidden', encoding='utf-8')
|
||||
try:
|
||||
self.assertEqual(index.files(), [(root / 'visible.md').resolve()])
|
||||
filenames = ['visible.md'] * 2001
|
||||
with patch.object(MODULE.os, 'walk', return_value=iter([(str(root), [], filenames)])):
|
||||
self.assertEqual(len(index.files()), 2000)
|
||||
finally:
|
||||
index.db.close()
|
||||
|
||||
def test_rerank_uses_cross_encoder_scores(self):
|
||||
index = MODULE.Index.__new__(MODULE.Index)
|
||||
|
||||
|
||||
Reference in new issue
Block a user