Moved import for retrievers down a folder (neo4j#63)

willtai · web-flow · commit dd60687e2340 · 2024-06-19T11:16:59.000+01:00
* Moved import for retrievers down a folder

* Included Pinecone and Weaviate retrievers in neo4j_genai.retrievers init file

* Update imports in docs and examples

* Update CHANGELOG
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -6,6 +6,10 @@
 -   Stopped embeddings from being returned when searching with `VectorRetriever`. Added `nodeLabels` and `id` to the metadata of `VectorRetriever` results.
 -   Added `upsert_vector` utility function for attaching vectors to node properties.
 -   Introduced `Neo4jInsertionError` for handling insertion failures in Neo4j.
+-   Included Pinecone and Weaviate retrievers in neo4j_genai.retrievers.
+
+### Changed
+-   Refactored import paths for retrievers to neo4j_genai.retrievers.
 
 ## 0.2.0
 
diff --git a/README.md b/README.md
@@ -30,7 +30,7 @@ Assumption: Neo4j running with populated vector index in place.
 
 ```python
 from neo4j import GraphDatabase
-from neo4j_genai import VectorRetriever
+from neo4j_genai.retrievers import VectorRetriever
 from langchain_openai import OpenAIEmbeddings
 
 URI = "neo4j://localhost:7687"
diff --git a/docs/source/api.rst b/docs/source/api.rst
@@ -140,7 +140,7 @@ Neo4jIndexError
 
 
 Neo4jInsertionError
-===============
+===================
 
 .. autoclass:: neo4j_genai.exceptions.Neo4jInsertionError
    :show-inheritance:
diff --git a/docs/source/index.rst b/docs/source/index.rst
@@ -61,7 +61,7 @@ Performing a similarity search
 .. code:: python
 
     from neo4j import GraphDatabase
-    from neo4j_genai import VectorRetriever
+    from neo4j_genai.retrievers import VectorRetriever
     from langchain_openai import OpenAIEmbeddings
 
     URI = "neo4j://localhost:7687"
diff --git a/examples/hybrid_cypher_search.py b/examples/hybrid_cypher_search.py
@@ -1,7 +1,7 @@
 from neo4j import GraphDatabase
 
 from random import random
-from neo4j_genai import HybridCypherRetriever
+from neo4j_genai.retrievers import HybridCypherRetriever
 from neo4j_genai.embedder import Embedder
 from neo4j_genai.indexes import create_vector_index, create_fulltext_index
 
diff --git a/examples/hybrid_search.py b/examples/hybrid_search.py
@@ -1,7 +1,7 @@
 from neo4j import GraphDatabase
 
 from random import random
-from neo4j_genai import HybridRetriever
+from neo4j_genai.retrievers import HybridRetriever
 from neo4j_genai.embedder import Embedder
 from neo4j_genai.indexes import create_vector_index, create_fulltext_index
 
diff --git a/examples/openai_search.py b/examples/openai_search.py
@@ -1,5 +1,5 @@
 from neo4j import GraphDatabase
-from neo4j_genai import VectorRetriever
+from neo4j_genai.retrievers import VectorRetriever
 
 from random import random
 from neo4j_genai.indexes import create_vector_index
diff --git a/examples/pinecone/text_search.py b/examples/pinecone/text_search.py
@@ -1,6 +1,6 @@
 from langchain_huggingface.embeddings import HuggingFaceEmbeddings
 from neo4j import GraphDatabase
-from neo4j_genai.retrievers.external.pinecone import PineconeNeo4jRetriever
+from neo4j_genai.retrievers import PineconeNeo4jRetriever
 from pinecone import Pinecone
 
 NEO4J_AUTH = ("neo4j", "password")
diff --git a/examples/pinecone/vector_search.py b/examples/pinecone/vector_search.py
@@ -1,5 +1,5 @@
 from neo4j import GraphDatabase
-from neo4j_genai.retrievers.external.pinecone import PineconeNeo4jRetriever
+from neo4j_genai.retrievers import PineconeNeo4jRetriever
 from pinecone import Pinecone
 
 from examples.embedding_biology import EMBEDDING_BIOLOGY
diff --git a/examples/similarity_search_for_text.py b/examples/similarity_search_for_text.py
@@ -1,5 +1,5 @@
 from neo4j import GraphDatabase
-from neo4j_genai import VectorRetriever
+from neo4j_genai.retrievers import VectorRetriever
 
 from random import random
 from neo4j_genai.embedder import Embedder
diff --git a/examples/similarity_search_for_vector.py b/examples/similarity_search_for_vector.py
@@ -1,5 +1,5 @@
 from neo4j import GraphDatabase
-from neo4j_genai import VectorRetriever
+from neo4j_genai.retrievers import VectorRetriever
 
 from random import random
 
diff --git a/examples/vector_cypher_retrieval.py b/examples/vector_cypher_retrieval.py
@@ -1,5 +1,5 @@
 from neo4j import GraphDatabase
-from neo4j_genai import VectorCypherRetriever
+from neo4j_genai.retrievers import VectorCypherRetriever
 
 import random
 import string
diff --git a/examples/vector_search_with_filters.py b/examples/vector_search_with_filters.py
@@ -1,5 +1,5 @@
 from neo4j import GraphDatabase
-from neo4j_genai import VectorRetriever
+from neo4j_genai.retrievers import VectorRetriever
 
 import random
 import string
diff --git a/examples/weaviate/text_search_local_embedder.py b/examples/weaviate/text_search_local_embedder.py
@@ -1,6 +1,6 @@
 from langchain_huggingface.embeddings import HuggingFaceEmbeddings
 from neo4j import GraphDatabase
-from neo4j_genai.retrievers.external.weaviate import WeaviateNeo4jRetriever
+from neo4j_genai.retrievers import WeaviateNeo4jRetriever
 from weaviate.connect.helpers import connect_to_local
 
 NEO4J_URL = "neo4j://localhost:7687"
diff --git a/examples/weaviate/text_search_remote_embedder.py b/examples/weaviate/text_search_remote_embedder.py
@@ -1,4 +1,4 @@
-from neo4j_genai.retrievers.external.weaviate import WeaviateNeo4jRetriever
+from neo4j_genai.retrievers import WeaviateNeo4jRetriever
 from neo4j import GraphDatabase
 from weaviate.connect.helpers import connect_to_local
 
diff --git a/examples/weaviate/vector_search.py b/examples/weaviate/vector_search.py
@@ -1,5 +1,5 @@
 from neo4j import GraphDatabase
-from neo4j_genai.retrievers.external.weaviate import WeaviateNeo4jRetriever
+from neo4j_genai.retrievers import WeaviateNeo4jRetriever
 from weaviate.connect.helpers import connect_to_local
 
 from examples.embedding_biology import EMBEDDING_BIOLOGY
diff --git a/poetry.lock b/poetry.lock
diff --git a/src/neo4j_genai/__init__.py b/src/neo4j_genai/__init__.py
@@ -12,16 +12,3 @@
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
-
-from .retrievers.hybrid import HybridCypherRetriever, HybridRetriever
-from .retrievers.text2cypher import Text2CypherRetriever
-from .retrievers.vector import VectorCypherRetriever, VectorRetriever
-
-
-__all__ = [
-    "VectorRetriever",
-    "VectorCypherRetriever",
-    "HybridRetriever",
-    "HybridCypherRetriever",
-    "Text2CypherRetriever",
-]
diff --git a/src/neo4j_genai/retrievers/__init__.py b/src/neo4j_genai/retrievers/__init__.py
@@ -12,3 +12,32 @@
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
+
+from .hybrid import HybridCypherRetriever, HybridRetriever
+from .text2cypher import Text2CypherRetriever
+from .vector import VectorCypherRetriever, VectorRetriever
+
+
+__all__ = [
+    "VectorRetriever",
+    "VectorCypherRetriever",
+    "HybridRetriever",
+    "HybridCypherRetriever",
+    "Text2CypherRetriever",
+]
+
+
+try:
+    from .external.pinecone.pinecone import PineconeNeo4jRetriever  # noqa: F401
+
+    __all__.append("PineconeNeo4jRetriever")
+except ImportError:
+    pass
+
+
+try:
+    from .external.weaviate.weaviate import WeaviateNeo4jRetriever  # noqa: F401
+
+    __all__.append("WeaviateNeo4jRetriever")
+except ImportError:
+    pass
diff --git a/src/neo4j_genai/retrievers/external/pinecone/__init__.py b/src/neo4j_genai/retrievers/external/pinecone/__init__.py
@@ -12,7 +12,3 @@
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
-
-from .pinecone import PineconeNeo4jRetriever
-
-__all__ = ["PineconeNeo4jRetriever"]
diff --git a/src/neo4j_genai/retrievers/external/pinecone/pinecone.py b/src/neo4j_genai/retrievers/external/pinecone/pinecone.py
@@ -52,7 +52,7 @@ class PineconeNeo4jRetriever(ExternalRetriever):
     .. code-block:: python
 
       from neo4j import GraphDatabase
-      from neo4j_genai.retrievers.external.pinecone import PineconeNeo4jRetriever
+      from neo4j_genai.retrievers import PineconeNeo4jRetriever
       from pinecone import Pinecone
 
       with GraphDatabase.driver(NEO4J_URL, auth=NEO4J_AUTH) as neo4j_driver:
@@ -153,7 +153,7 @@ def _get_search_results(
         .. code-block:: python
 
           from neo4j import GraphDatabase
-          from neo4j_genai.retrievers.external.pinecone import PineconeNeo4jRetriever
+          from neo4j_genai.retrievers import PineconeNeo4jRetriever
           from pinecone import Pinecone
 
           with GraphDatabase.driver(NEO4J_URL, auth=NEO4J_AUTH) as neo4j_driver:
diff --git a/src/neo4j_genai/retrievers/external/weaviate/__init__.py b/src/neo4j_genai/retrievers/external/weaviate/__init__.py
@@ -12,7 +12,3 @@
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
-
-from .weaviate import WeaviateNeo4jRetriever
-
-__all__ = ["WeaviateNeo4jRetriever"]
diff --git a/src/neo4j_genai/retrievers/external/weaviate/types.py b/src/neo4j_genai/retrievers/external/weaviate/types.py
@@ -25,7 +25,7 @@
 from weaviate.client import WeaviateClient
 from weaviate.collections.classes.filters import _Filters
 
-from neo4j_genai.retrievers.utils import validate_search_query_input
+from neo4j_genai.utils import validate_search_query_input
 from neo4j_genai.types import Neo4jDriverModel, EmbedderModel
 
 
diff --git a/src/neo4j_genai/retrievers/external/weaviate/weaviate.py b/src/neo4j_genai/retrievers/external/weaviate/weaviate.py
@@ -45,7 +45,7 @@ class WeaviateNeo4jRetriever(ExternalRetriever):
     .. code-block:: python
 
       from neo4j import GraphDatabase
-      from neo4j_genai.retrievers.external.weaviate import WeaviateNeo4jRetriever
+      from neo4j_genai.retrievers import WeaviateNeo4jRetriever
       from weaviate.connect.helpers import connect_to_local
 
       with GraphDatabase.driver(NEO4J_URL, auth=NEO4J_AUTH) as neo4j_driver:
@@ -136,7 +136,7 @@ def _get_search_results(
         .. code-block:: python
 
           import neo4j
-          from neo4j_genai.retrievers.external.weaviate import WeaviateNeo4jRetriever
+          from neo4j_genai.retrievers import WeaviateNeo4jRetriever
 
           driver = neo4j.GraphDatabase.driver(URI, auth=AUTH)
 
diff --git a/src/neo4j_genai/retrievers/hybrid.py b/src/neo4j_genai/retrievers/hybrid.py
@@ -52,7 +52,7 @@ class HybridRetriever(Retriever):
     .. code-block:: python
 
       import neo4j
-      from neo4j_genai import HybridRetriever
+      from neo4j_genai.retrievers import HybridRetriever
 
       driver = neo4j.GraphDatabase.driver(URI, auth=AUTH)
 
@@ -187,7 +187,7 @@ class HybridCypherRetriever(Retriever):
     .. code-block:: python
 
       import neo4j
-      from neo4j_genai import HybridCypherRetriever
+      from neo4j_genai.retrievers import HybridCypherRetriever
 
       driver = neo4j.GraphDatabase.driver(URI, auth=AUTH)
 
diff --git a/src/neo4j_genai/retrievers/vector.py b/src/neo4j_genai/retrievers/vector.py
@@ -52,7 +52,7 @@ class VectorRetriever(Retriever):
     .. code-block:: python
 
       import neo4j
-      from neo4j_genai import VectorRetriever
+      from neo4j_genai.retrievers import VectorRetriever
 
       driver = neo4j.GraphDatabase.driver(URI, auth=AUTH)
 
@@ -192,7 +192,7 @@ class VectorCypherRetriever(Retriever):
     .. code-block:: python
 
       import neo4j
-      from neo4j_genai import VectorCypherRetriever
+      from neo4j_genai.retrievers import VectorCypherRetriever
 
       driver = neo4j.GraphDatabase.driver(URI, auth=AUTH)
 
diff --git a/src/neo4j_genai/types.py b/src/neo4j_genai/types.py
@@ -24,7 +24,7 @@
     model_validator,
 )
 
-from neo4j_genai.retrievers.utils import validate_search_query_input
+from neo4j_genai.utils import validate_search_query_input
 
 
 class RawSearchResult(BaseModel):
diff --git a/src/neo4j_genai/utils.py b/src/neo4j_genai/utils.py
diff --git a/tests/e2e/pinecone_e2e/test_pinecone_e2e.py b/tests/e2e/pinecone_e2e/test_pinecone_e2e.py
@@ -19,8 +19,8 @@
 
 import pytest
 from langchain_huggingface.embeddings import HuggingFaceEmbeddings
+from neo4j_genai.retrievers import PineconeNeo4jRetriever
 from neo4j_genai.embedder import Embedder
-from neo4j_genai.retrievers.external.pinecone import PineconeNeo4jRetriever
 from neo4j_genai.types import RetrieverResult, RetrieverResultItem
 from pinecone import Pinecone
 
diff --git a/tests/e2e/test_hybrid_e2e.py b/tests/e2e/test_hybrid_e2e.py
@@ -15,7 +15,7 @@
 
 import pytest
 from neo4j import Driver
-from neo4j_genai import (
+from neo4j_genai.retrievers import (
     HybridCypherRetriever,
     HybridRetriever,
 )
diff --git a/tests/e2e/test_text2cypher_e2e.py b/tests/e2e/test_text2cypher_e2e.py
@@ -1,5 +1,5 @@
 import pytest
-from neo4j_genai import Text2CypherRetriever
+from neo4j_genai.retrievers import Text2CypherRetriever
 from unittest.mock import MagicMock
 from neo4j_genai.types import RetrieverResult, RetrieverResultItem
 
diff --git a/tests/e2e/test_vector_e2e.py b/tests/e2e/test_vector_e2e.py
@@ -15,7 +15,7 @@
 
 import pytest
 from neo4j import Driver
-from neo4j_genai import VectorCypherRetriever, VectorRetriever
+from neo4j_genai.retrievers import VectorCypherRetriever, VectorRetriever
 from neo4j_genai.embedder import Embedder
 from neo4j_genai.types import RetrieverResult, RetrieverResultItem
 
diff --git a/tests/e2e/weaviate_e2e/test_weaviate_e2e.py b/tests/e2e/weaviate_e2e/test_weaviate_e2e.py
@@ -20,7 +20,7 @@
 from langchain_huggingface.embeddings import HuggingFaceEmbeddings
 from neo4j import Driver
 from neo4j_genai.embedder import Embedder
-from neo4j_genai.retrievers.external.weaviate import WeaviateNeo4jRetriever
+from neo4j_genai.retrievers import WeaviateNeo4jRetriever
 from neo4j_genai.types import RetrieverResult, RetrieverResultItem
 from weaviate.client import Client
 from weaviate.connect.helpers import connect_to_local
diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py
@@ -17,7 +17,7 @@
 
 import neo4j
 import pytest
-from neo4j_genai import (
+from neo4j_genai.retrievers import (
     HybridRetriever,
     Text2CypherRetriever,
     VectorCypherRetriever,
@@ -43,15 +43,15 @@ def llm() -> MagicMock:
 
 
 @pytest.fixture(scope="function")
-@patch("neo4j_genai.VectorRetriever._verify_version")
+@patch("neo4j_genai.retrievers.VectorRetriever._verify_version")
 def vector_retriever(
     _verify_version_mock: MagicMock, driver: MagicMock
 ) -> VectorRetriever:
     return VectorRetriever(driver, "my-index")
 
 
 @pytest.fixture(scope="function")
-@patch("neo4j_genai.VectorCypherRetriever._verify_version")
+@patch("neo4j_genai.retrievers.VectorCypherRetriever._verify_version")
 def vector_cypher_retriever(
     _verify_version_mock: MagicMock, driver: MagicMock
 ) -> VectorCypherRetriever:
@@ -62,15 +62,15 @@ def vector_cypher_retriever(
 
 
 @pytest.fixture(scope="function")
-@patch("neo4j_genai.HybridRetriever._verify_version")
+@patch("neo4j_genai.retrievers.HybridRetriever._verify_version")
 def hybrid_retriever(
     _verify_version_mock: MagicMock, driver: MagicMock
 ) -> HybridRetriever:
     return HybridRetriever(driver, "my-index", "my-fulltext-index")
 
 
 @pytest.fixture(scope="function")
-@patch("neo4j_genai.Text2CypherRetriever._verify_version")
+@patch("neo4j_genai.retrievers.Text2CypherRetriever._verify_version")
 def t2c_retriever(
     _verify_version_mock: MagicMock, driver: MagicMock, llm: MagicMock
 ) -> Text2CypherRetriever:
diff --git a/tests/unit/retrievers/external/test_pinecone.py b/tests/unit/retrievers/external/test_pinecone.py
@@ -18,7 +18,7 @@
 import neo4j
 import pytest
 from neo4j_genai.exceptions import RetrieverInitializationError
-from neo4j_genai.retrievers.external.pinecone import PineconeNeo4jRetriever
+from neo4j_genai.retrievers import PineconeNeo4jRetriever
 from neo4j_genai.retrievers.external.utils import get_match_query
 from neo4j_genai.types import RetrieverResult, RetrieverResultItem
 from pinecone import Pinecone
diff --git a/tests/unit/retrievers/external/test_weaviate.py b/tests/unit/retrievers/external/test_weaviate.py
@@ -19,7 +19,7 @@
 
 import neo4j
 from neo4j_genai.retrievers.external.utils import get_match_query
-from neo4j_genai.retrievers.external.weaviate import WeaviateNeo4jRetriever
+from neo4j_genai.retrievers import WeaviateNeo4jRetriever
 from neo4j_genai.types import RetrieverResult, RetrieverResultItem
 from weaviate.client import WeaviateClient
 
diff --git a/tests/unit/retrievers/test_hybrid.py b/tests/unit/retrievers/test_hybrid.py
diff --git a/tests/unit/retrievers/test_text2cypher.py b/tests/unit/retrievers/test_text2cypher.py
diff --git a/tests/unit/retrievers/test_vector.py b/tests/unit/retrievers/test_vector.py
