test_pgvector.py 913 B

12345678910111213141516171819202122232425262728293031
  1. from core.rag.datasource.vdb.pgvector.pgvector import PGVector, PGVectorConfig
  2. from core.rag.models.document import Document
  3. from tests.integration_tests.vdb.test_vector_store import (
  4. AbstractVectorTest,
  5. get_example_text,
  6. setup_mock_redis,
  7. )
  8. class TestPGVector(AbstractVectorTest):
  9. def __init__(self):
  10. super().__init__()
  11. self.vector = PGVector(
  12. collection_name=self.collection_name,
  13. config=PGVectorConfig(
  14. host="localhost",
  15. port=5433,
  16. user="postgres",
  17. password="difyai123456",
  18. database="dify",
  19. ),
  20. )
  21. def search_by_full_text(self):
  22. hits_by_full_text: list[Document] = self.vector.search_by_full_text(query=get_example_text())
  23. assert len(hits_by_full_text) == 0
  24. def test_pgvector(setup_mock_redis):
  25. TestPGVector().run_all_tests()