Added comment about Embedder type in docstring

neo4j · Apr 22, 2024 · 7b5ac29 · 7b5ac29
1 parent 2a21e0d
commit 7b5ac29
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 4 deletions.
diff --git a/examples/similarity_search_for_text.py b/examples/similarity_search_for_text.py
@@ -3,7 +3,7 @@
 
 from random import random
 from neo4j_genai.embedder import Embedder
-from neo4j_genai.indexes import create_vector_index, drop_index
+from neo4j_genai.indexes import create_vector_index
 
 URI = "neo4j://localhost:7687"
 AUTH = ("neo4j", "password")
@@ -15,7 +15,7 @@
 driver = GraphDatabase.driver(URI, auth=AUTH)
 
 
-# Create Embedder object
+# Create CustomEmbedder object with the required Embedder type
 class CustomEmbedder(Embedder):
     def embed_query(self, text: str) -> list[float]:
         return [random() for _ in range(DIMENSION)]

diff --git a/src/neo4j_genai/retrievers.py b/src/neo4j_genai/retrievers.py
@@ -66,7 +66,8 @@ def search(self, *args, **kwargs) -> Any:
 
 class VectorRetriever(Retriever):
     """
-    Provides retrieval method using vector search over embeddings
+    Provides retrieval method using vector search over embeddings.
+    If an embedder is provided, it needs to have the required Embedder type.
     """
 
     def __init__(
@@ -160,7 +161,8 @@ def search(
 
 class VectorCypherRetriever(Retriever):
     """
-    Provides retrieval method using vector similarity and custom Cypher query
+    Provides retrieval method using vector similarity and custom Cypher query.
+    If an embedder is provided, it needs to have the required Embedder type.
     """
 
     def __init__(