add pgvecto_rs support and upgrade SQLAlchemy (#3833)

This commit is contained in:
Jyong
2024-04-29 11:58:17 +08:00
committed by GitHub
parent 975b2fb79e
commit 3e9dbe3e0a
26 changed files with 584 additions and 220 deletions

View File

@@ -0,0 +1,37 @@
from core.rag.datasource.vdb.pgvecto_rs.pgvecto_rs import PGVectoRS, PgvectoRSConfig
from tests.integration_tests.vdb.test_vector_store import (
AbstractVectorTest,
get_example_text,
setup_mock_redis,
)
class TestPgvectoRSVector(AbstractVectorTest):
def __init__(self):
super().__init__()
self.vector = PGVectoRS(
collection_name=self.collection_name.lower(),
config=PgvectoRSConfig(
host='localhost',
port=5431,
user='postgres',
password='difyai123456',
database='dify',
),
dim=128
)
def search_by_full_text(self):
# pgvecto rs only support english text search, So its not open for now
hits_by_full_text = self.vector.search_by_full_text(query=get_example_text())
assert len(hits_by_full_text) == 0
def delete_by_document_id(self):
self.vector.delete_by_document_id(document_id=self.example_doc_id)
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) == 1
def test_pgvecot_rs(setup_mock_redis):
TestPgvectoRSVector().run_all_tests()

View File

@@ -45,7 +45,7 @@ class AbstractVectorTest:
def __init__(self):
self.vector = None
self.dataset_id = str(uuid.uuid4())
self.collection_name = Dataset.gen_collection_name_by_id(self.dataset_id)
self.collection_name = Dataset.gen_collection_name_by_id(self.dataset_id) + '_test'
self.example_doc_id = str(uuid.uuid4())
self.example_embedding = [1.001 * i for i in range(128)]