diff --git a/integrations/supabase/src/haystack_integrations/components/retrievers/supabase/__init__.py b/integrations/supabase/src/haystack_integrations/components/retrievers/supabase/__init__.py new file mode 100644 index 000000000..fe2320e47 --- /dev/null +++ b/integrations/supabase/src/haystack_integrations/components/retrievers/supabase/__init__.py @@ -0,0 +1,3 @@ +from .retriever import SupabaseEmbeddingRetriever + +__all__ = ["SupabaseEmbeddingRetriever"] diff --git a/integrations/supabase/src/supabase_haystack/retriever.py b/integrations/supabase/src/haystack_integrations/components/retrievers/supabase/retriever.py similarity index 95% rename from integrations/supabase/src/supabase_haystack/retriever.py rename to integrations/supabase/src/haystack_integrations/components/retrievers/supabase/retriever.py index 5fb706d50..675be03db 100644 --- a/integrations/supabase/src/supabase_haystack/retriever.py +++ b/integrations/supabase/src/haystack_integrations/components/retrievers/supabase/retriever.py @@ -5,8 +5,7 @@ from haystack import component from haystack.dataclasses import Document - -from src.supabase_haystack.document_store import SupabaseDocumentStore +from haystack_integrations.document_stores.supabase import SupabaseDocumentStore @component diff --git a/integrations/supabase/src/supabase_haystack/__init__.py b/integrations/supabase/src/haystack_integrations/document_stores/supabase/__init__.py similarity index 66% rename from integrations/supabase/src/supabase_haystack/__init__.py rename to integrations/supabase/src/haystack_integrations/document_stores/supabase/__init__.py index dca62f953..53939bdff 100644 --- a/integrations/supabase/src/supabase_haystack/__init__.py +++ b/integrations/supabase/src/haystack_integrations/document_stores/supabase/__init__.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: 2023-present John Doe # # SPDX-License-Identifier: Apache-2.0 -from src.supabase_haystack.document_store import SupabaseDocumentStore +from .document_store import SupabaseDocumentStore __all__ = ["SupabaseDocumentStore"] diff --git a/integrations/supabase/src/supabase_haystack/document_store.py b/integrations/supabase/src/haystack_integrations/document_stores/supabase/document_store.py similarity index 98% rename from integrations/supabase/src/supabase_haystack/document_store.py rename to integrations/supabase/src/haystack_integrations/document_stores/supabase/document_store.py index 0492079bd..8ecd3660b 100644 --- a/integrations/supabase/src/supabase_haystack/document_store.py +++ b/integrations/supabase/src/haystack_integrations/document_stores/supabase/document_store.py @@ -10,7 +10,7 @@ from haystack.document_stores.types import DuplicatePolicy from haystack.utils.filters import convert -from src.supabase_haystack.filters import _normalize_filters +from .filters import _normalize_filters logger = logging.getLogger(__name__) @@ -44,6 +44,7 @@ def __init__( self.dimension = dimension self._collection_name = collection_name self._dummy_vector = [0.0]*dimension + self.collection_creation_kwargs = collection_creation_kwargs db_connection = f"postgresql://{user}:{password}@{host}:{port}/{db_name}" self._pgvector_client = vecs.create_client(db_connection) self._collection = self._pgvector_client.get_or_create_collection(name=collection_name, dimension=dimension, **collection_creation_kwargs) diff --git a/integrations/supabase/src/supabase_haystack/filters.py b/integrations/supabase/src/haystack_integrations/document_stores/supabase/filters.py similarity index 100% rename from integrations/supabase/src/supabase_haystack/filters.py rename to integrations/supabase/src/haystack_integrations/document_stores/supabase/filters.py diff --git a/integrations/supabase/tests/conftest.py b/integrations/supabase/tests/conftest.py index 9b37d5c50..2c91e50e1 100644 --- a/integrations/supabase/tests/conftest.py +++ b/integrations/supabase/tests/conftest.py @@ -1,7 +1,7 @@ import pytest from haystack.document_stores.types import DuplicatePolicy -from src.supabase_haystack.document_store import SupabaseDocumentStore +from src.haystack_integrations.document_stores.supabase import SupabaseDocumentStore # This is the approximate time it takes for the documents to be available SLEEP_TIME = 20 diff --git a/integrations/supabase/tests/test_document_store.py b/integrations/supabase/tests/test_document_store.py index 0883b6b22..4f9519369 100644 --- a/integrations/supabase/tests/test_document_store.py +++ b/integrations/supabase/tests/test_document_store.py @@ -10,7 +10,7 @@ from haystack.document_stores.types import DocumentStore, DuplicatePolicy from haystack.testing.document_store import CountDocumentsTest, DeleteDocumentsTest, WriteDocumentsTest -from src.supabase_haystack.document_store import SupabaseDocumentStore +from src.haystack_integrations.document_stores.supabase import SupabaseDocumentStore class TestDocumentStore(CountDocumentsTest, DeleteDocumentsTest, WriteDocumentsTest): @@ -78,7 +78,7 @@ def test_embedding_retrieval(self, document_store: SupabaseDocumentStore): assert results[0].content == "Most similar document" assert results[1].content == "2nd best document" - @patch("src.supabase_haystack.document_store.vecs") + @patch("src.haystack_integrations.document_stores.supabase.document_store.vecs") def test_init(self, mock_supabase): document_store = SupabaseDocumentStore(