Merge branch 'main' into fix/chore-fix

This commit is contained in:
Yeuoly
2024-12-16 14:08:18 +08:00
370 changed files with 11022 additions and 6468 deletions

View File

@ -7,9 +7,10 @@ env = environs.Env()
class Config:
SEARCH_ENDPOINT = env.str("SEARCH_ENDPOINT", "http://ld-*************-proxy-search-pub.lindorm.aliyuncs.com:30070")
SEARCH_ENDPOINT = env.str("SEARCH_ENDPOINT", "http://ld-************-proxy-search-pub.lindorm.aliyuncs.com:30070")
SEARCH_USERNAME = env.str("SEARCH_USERNAME", "ADMIN")
SEARCH_PWD = env.str("SEARCH_PWD", "PWD")
SEARCH_PWD = env.str("SEARCH_PWD", "ADMIN")
USING_UGC = env.bool("USING_UGC", True)
class TestLindormVectorStore(AbstractVectorTest):
@ -31,5 +32,27 @@ class TestLindormVectorStore(AbstractVectorTest):
assert ids[0] == self.example_doc_id
def test_lindorm_vector(setup_mock_redis):
class TestLindormVectorStoreUGC(AbstractVectorTest):
def __init__(self):
super().__init__()
self.vector = LindormVectorStore(
collection_name="ugc_index_test",
config=LindormVectorStoreConfig(
hosts=Config.SEARCH_ENDPOINT,
username=Config.SEARCH_USERNAME,
password=Config.SEARCH_PWD,
using_ugc=Config.USING_UGC,
),
routing_value=self.collection_name,
)
def get_ids_by_metadata_field(self):
ids = self.vector.get_ids_by_metadata_field(key="doc_id", value=self.example_doc_id)
assert ids is not None
assert len(ids) == 1
assert ids[0] == self.example_doc_id
def test_lindorm_vector_ugc(setup_mock_redis):
TestLindormVectorStore().run_all_tests()
TestLindormVectorStoreUGC().run_all_tests()

View File

@ -12,11 +12,11 @@ def tidb_vector():
return TiDBVector(
collection_name="test_collection",
config=TiDBVectorConfig(
host="xxx.eu-central-1.xxx.aws.tidbcloud.com",
port="4000",
user="xxx.root",
password="xxxxxx",
database="dify",
host="localhost",
port=4000,
user="root",
password="",
database="test",
program_name="langgenius/dify",
),
)
@ -27,35 +27,14 @@ class TiDBVectorTest(AbstractVectorTest):
super().__init__()
self.vector = vector
def text_exists(self):
exist = self.vector.text_exists(self.example_doc_id)
assert exist == False
def search_by_vector(self):
hits_by_vector: list[Document] = self.vector.search_by_vector(query_vector=self.example_embedding)
assert len(hits_by_vector) == 0
def search_by_full_text(self):
hits_by_full_text: list[Document] = self.vector.search_by_full_text(query=get_example_text())
assert len(hits_by_full_text) == 0
def get_ids_by_metadata_field(self):
ids = self.vector.get_ids_by_metadata_field(key="document_id", value=self.example_doc_id)
assert len(ids) == 0
ids = self.vector.get_ids_by_metadata_field(key="doc_id", value=self.example_doc_id)
assert len(ids) == 1
def test_tidb_vector(setup_mock_redis, setup_tidbvector_mock, tidb_vector, mock_session):
def test_tidb_vector(setup_mock_redis, tidb_vector):
TiDBVectorTest(vector=tidb_vector).run_all_tests()
@pytest.fixture
def mock_session():
with patch("core.rag.datasource.vdb.tidb_vector.tidb_vector.Session", new_callable=MagicMock) as mock_session:
yield mock_session
@pytest.fixture
def setup_tidbvector_mock(tidb_vector, mock_session):
with patch("core.rag.datasource.vdb.tidb_vector.tidb_vector.create_engine"):
with patch.object(tidb_vector._engine, "connect"):
yield tidb_vector