estel/dify
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

test: run vdb tests on TiDB Vector with docker in CI tests (#11645)

Browse Source
This commit is contained in:
Bowen Liang
2024-12-15 17:16:40 +08:00
committed by GitHub
parent 7e154a467b
commit 924b4fe742
5 changed files with 21 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
api/core/rag/datasource/vdb/tidb_vector/tidb_vector.py
View File

@@ -37,8 +37,6 @@ class TiDBVectorConfig(BaseModel):
raise ValueError("config TIDB_VECTOR_PORT is required")
if not values["user"]:
raise ValueError("config TIDB_VECTOR_USER is required")
if not values["password"]:
raise ValueError("config TIDB_VECTOR_PASSWORD is required")
if not values["database"]:
raise ValueError("config TIDB_VECTOR_DATABASE is required")
if not values["program_name"]:

37
api/tests/integration_tests/vdb/tidb_vector/test_tidb_vector.py
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