Removes Pinecone and Weaviate retrievers from __init__.py (neo4j#59)

* Add try/catch in __init__.py for optional dependencies

* Removed pinecone and weaviate retrievers from init

* Resolve type error resulting from sentence_tranformers .encode method

* Updated CHANGELOG

Author: willtai

CHANGELOG.md

@@ -20,3 +20,4 @@
 ### Fixed
 - Updated documentation to include new custom exceptions.
 - Improved the use of Pydantic for input data validation for retriever objects.
+- Removed Pinecone and Weaviate retrievers from __init__.py to prevent ImportError when optional dependencies are not installed.

poetry.lock

@@ -13,18 +13,15 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
-from .retrievers.external.pinecone import PineconeNeo4jRetriever
-from .retrievers.external.weaviate import WeaviateNeo4jRetriever
 from .retrievers.hybrid import HybridCypherRetriever, HybridRetriever
 from .retrievers.text2cypher import Text2CypherRetriever
 from .retrievers.vector import VectorCypherRetriever, VectorRetriever
 
+
 __all__ = [
     "VectorRetriever",
     "VectorCypherRetriever",
     "HybridRetriever",
     "HybridCypherRetriever",
     "Text2CypherRetriever",
-    "WeaviateNeo4jRetriever",
-    "PineconeNeo4jRetriever",
 ]

src/neo4j_genai/__init__.py

@@ -1,5 +1,7 @@
 from neo4j_genai.embedder import Embedder
 from typing import Any
+import torch
+import numpy as np
 
 
 class SentenceTransformerEmbeddings(Embedder):
@@ -17,4 +19,12 @@ def __init__(
         self.model = SentenceTransformer(model, *args, **kwargs)
 
     def embed_query(self, text: str) -> Any:
-        return self.model.encode([text]).flatten().tolist()
+        result = self.model.encode([text])
+        if isinstance(result, torch.Tensor) or isinstance(result, np.ndarray):
+            return result.flatten().tolist()
+        elif isinstance(result, list) and all(
+            isinstance(x, torch.Tensor) for x in result
+        ):
+            return [item for tensor in result for item in tensor.flatten().tolist()]
+        else:
+            raise ValueError("Unexpected return type from model encoding")

src/neo4j_genai/embeddings/sentence_transformers.py

@@ -52,7 +52,7 @@ class PineconeNeo4jRetriever(ExternalRetriever):
     .. code-block:: python
 
       from neo4j import GraphDatabase
-      from neo4j_genai import PineconeNeo4jRetriever
+      from neo4j_genai.retrievers.external.pinecone 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 import PineconeNeo4jRetriever
+          from neo4j_genai.retrievers.external.pinecone import PineconeNeo4jRetriever
           from pinecone import Pinecone
 
           with GraphDatabase.driver(NEO4J_URL, auth=NEO4J_AUTH) as neo4j_driver:

src/neo4j_genai/retrievers/external/pinecone/pinecone.py

@@ -45,7 +45,7 @@ class WeaviateNeo4jRetriever(ExternalRetriever):
     .. code-block:: python
 
       from neo4j import GraphDatabase
-      from neo4j_genai import WeaviateNeo4jRetriever
+      from neo4j_genai.retrievers.external.weaviate 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 import WeaviateNeo4jRetriever
+          from neo4j_genai.retrievers.external.weaviate import WeaviateNeo4jRetriever
 
           driver = neo4j.GraphDatabase.driver(URI, auth=AUTH)
 

src/neo4j_genai/retrievers/external/weaviate/weaviate.py

@@ -19,7 +19,7 @@
 
 import pytest
 from langchain_community.embeddings import HuggingFaceEmbeddings
-from neo4j_genai import PineconeNeo4jRetriever
+from neo4j_genai.retrievers.external.pinecone import PineconeNeo4jRetriever
 from neo4j_genai.embedder import Embedder
 from neo4j_genai.types import RetrieverResult, RetrieverResultItem
 from pinecone import Pinecone