Rename variables (neo4j#65)

willtai · web-flow · commit 24964a2bfcfa · 2024-06-20T13:04:19.000+01:00
* Renamed embedding property variable to be more explicit

* Rename format_record_function to result_formatter

* Changed default_format_record to default_record_formatter

* Add tox to gitignore

* Rebase

* Rename propertyKey and textProperty to vectorProperty
diff --git a/.gitignore b/.gitignore
@@ -11,3 +11,4 @@ docs/build/
 .python-version
 .DS_Store
 .venv
+.tox/
diff --git a/README.md b/README.md
@@ -89,7 +89,7 @@ create_vector_index(
     driver,
     INDEX_NAME,
     label="Document",
-    property="textProperty",
+    embedding_property="vectorProperty",
     dimensions=1536,
     similarity_fn="euclidean",
 )
@@ -118,7 +118,7 @@ vector = [random() for _ in range(DIMENSION)]
 insert_query = (
     "MERGE (n:Document {id: $id})"
     "WITH n "
-    "CALL db.create.setNodeVectorProperty(n, 'textProperty', $vector)"
+    "CALL db.create.setNodeVectorProperty(n, 'vectorProperty', $vector)"
     "RETURN n"
 )
 parameters = {
diff --git a/docs/source/api.rst b/docs/source/api.rst
@@ -52,14 +52,14 @@ This section includes retrievers that integrate with databases external to Neo4j
 WeaviateNeo4jRetriever
 ======================
 
-.. autoclass:: neo4j_genai.retrievers.external.weaviate.WeaviateNeo4jRetriever
+.. autoclass:: neo4j_genai.retrievers.external.weaviate.weaviate.WeaviateNeo4jRetriever
    :members:
 
 
 PineconeNeo4jRetriever
 ======================
 
-.. autoclass:: neo4j_genai.retrievers.external.pinecone.PineconeNeo4jRetriever
+.. autoclass:: neo4j_genai.retrievers.external.pinecone.pinecone.PineconeNeo4jRetriever
    :members:
 
 
diff --git a/docs/source/index.rst b/docs/source/index.rst
@@ -112,7 +112,7 @@ When creating a vector index, make sure you match the number of dimensions in th
         driver,
         INDEX_NAME,
         label="Document",
-        property="textProperty",
+        embedding_property="vectorProperty",
         dimensions=1536,
         similarity_fn="euclidean",
     )
@@ -144,7 +144,7 @@ See below for how to write using Cypher via the Neo4j driver.
     insert_query = (
         "MERGE (n:Document {id: $id})"
         "WITH n "
-        "CALL db.create.setNodeVectorProperty(n, 'textProperty', $vector)"
+        "CALL db.create.setNodeVectorProperty(n, 'vectorProperty', $vector)"
         "RETURN n"
     )
     parameters = {
diff --git a/examples/graphrag.py b/examples/graphrag.py
@@ -44,7 +44,7 @@ def formatter(record: neo4j.Record) -> RetrieverResultItem:
     driver,
     index_name=INDEX,
     retrieval_query="with node, score return node.title as title, node.plot as plot",
-    format_record_function=formatter,
+    result_formatter=formatter,
     embedder=embedder,
 )
 
diff --git a/examples/graphrag_custom_prompt.py b/examples/graphrag_custom_prompt.py
@@ -45,7 +45,7 @@ def formatter(record: neo4j.Record) -> RetrieverResultItem:
     driver,
     index_name=INDEX,
     retrieval_query="with node, score return node.title as title, node.plot as plot",
-    format_record_function=formatter,
+    result_formatter=formatter,
     embedder=embedder,
 )
 
diff --git a/examples/hybrid_cypher_search.py b/examples/hybrid_cypher_search.py
@@ -29,12 +29,12 @@ def embed_query(self, text: str) -> list[float]:
     driver,
     INDEX_NAME,
     label="Document",
-    property="propertyKey",
+    embedding_property="vectorProperty",
     dimensions=DIMENSION,
     similarity_fn="euclidean",
 )
 create_fulltext_index(
-    driver, FULLTEXT_INDEX_NAME, label="Document", node_properties=["propertyKey"]
+    driver, FULLTEXT_INDEX_NAME, label="Document", node_properties=["vectorProperty"]
 )
 
 # Initialize the retriever
@@ -48,7 +48,7 @@ def embed_query(self, text: str) -> list[float]:
 insert_query = (
     "MERGE (n:Document {id: $id})"
     "WITH n "
-    "CALL db.create.setNodeVectorProperty(n, 'propertyKey', $vector)"
+    "CALL db.create.setNodeVectorProperty(n, 'vectorProperty', $vector)"
     "RETURN n"
 )
 parameters = {
diff --git a/examples/hybrid_search.py b/examples/hybrid_search.py
@@ -29,12 +29,12 @@ def embed_query(self, text: str) -> list[float]:
     driver,
     INDEX_NAME,
     label="Document",
-    property="propertyKey",
+    embedding_property="vectorProperty",
     dimensions=DIMENSION,
     similarity_fn="euclidean",
 )
 create_fulltext_index(
-    driver, FULLTEXT_INDEX_NAME, label="Document", node_properties=["propertyKey"]
+    driver, FULLTEXT_INDEX_NAME, label="Document", node_properties=["vectorProperty"]
 )
 
 # Initialize the retriever
@@ -45,7 +45,7 @@ def embed_query(self, text: str) -> list[float]:
 insert_query = (
     "MERGE (n:Document {id: $id})"
     "WITH n "
-    "CALL db.create.setNodeVectorProperty(n, 'propertyKey', $vector)"
+    "CALL db.create.setNodeVectorProperty(n, 'vectorProperty', $vector)"
     "RETURN n"
 )
 parameters = {
diff --git a/examples/openai_search.py b/examples/openai_search.py
@@ -27,7 +27,7 @@
     driver,
     INDEX_NAME,
     label="Document",
-    property="propertyKey",
+    embedding_property="vectorProperty",
     dimensions=DIMENSION,
     similarity_fn="cosine",
 )
@@ -38,7 +38,7 @@
 insert_query = (
     "MERGE (n:Document {id: $id})"
     "WITH n "
-    "CALL db.create.setNodeVectorProperty(n, 'propertyKey', $vector)"
+    "CALL db.create.setNodeVectorProperty(n, 'vectorProperty', $vector)"
     "RETURN n"
 )
 parameters = {
diff --git a/examples/similarity_search_for_text.py b/examples/similarity_search_for_text.py
@@ -28,7 +28,7 @@ def embed_query(self, text: str) -> list[float]:
     driver,
     INDEX_NAME,
     label="Document",
-    property="propertyKey",
+    embedding_property="vectorProperty",
     dimensions=DIMENSION,
     similarity_fn="euclidean",
 )
@@ -41,7 +41,7 @@ def embed_query(self, text: str) -> list[float]:
 insert_query = (
     "MERGE (n:Document {id: $id})"
     "WITH n "
-    "CALL db.create.setNodeVectorProperty(n, 'propertyKey', $vector)"
+    "CALL db.create.setNodeVectorProperty(n, 'vectorProperty', $vector)"
     "RETURN n"
 )
 parameters = {
diff --git a/examples/similarity_search_for_vector.py b/examples/similarity_search_for_vector.py
@@ -19,7 +19,7 @@
     driver,
     INDEX_NAME,
     label="Document",
-    property="propertyKey",
+    embedding_property="vectorProperty",
     dimensions=DIMENSION,
     similarity_fn="euclidean",
 )
@@ -32,7 +32,7 @@
 insert_query = (
     "MERGE (n:Document {id: $id})"
     "WITH n "
-    "CALL db.create.setNodeVectorProperty(n, 'propertyKey', $vector)"
+    "CALL db.create.setNodeVectorProperty(n, 'vectorProperty', $vector)"
     "RETURN n"
 )
 parameters = {
diff --git a/examples/vector_cypher_retrieval.py b/examples/vector_cypher_retrieval.py
@@ -35,7 +35,7 @@ def random_str(n: int) -> str:
     driver,
     INDEX_NAME,
     label="Document",
-    property="propertyKey",
+    embedding_property="vectorProperty",
     dimensions=DIMENSION,
     similarity_fn="euclidean",
 )
@@ -49,7 +49,7 @@ def random_str(n: int) -> str:
 insert_query = (
     "MERGE (doc:Document {id: $id})"
     "WITH doc "
-    "CALL db.create.setNodeVectorProperty(doc, 'propertyKey', $vector)"
+    "CALL db.create.setNodeVectorProperty(doc, 'vectorProperty', $vector)"
     "WITH doc "
     "MERGE (author:Author {name: $authorName})"
     "MERGE (doc)-[:AUTHORED_BY]->(author)"
diff --git a/examples/vector_search_with_filters.py b/examples/vector_search_with_filters.py
@@ -35,7 +35,7 @@ def random_str(n: int) -> str:
     driver,
     INDEX_NAME,
     label="Document",
-    property="propertyKey",
+    embedding_property="vectorProperty",
     dimensions=DIMENSION,
     similarity_fn="euclidean",
 )
@@ -50,7 +50,7 @@ def random_str(n: int) -> str:
     "ON CREATE SET  doc.int_property = $id, "
     "               doc.short_text_property = toString($id)"
     "WITH doc "
-    "CALL db.create.setNodeVectorProperty(doc, 'propertyKey', $vector)"
+    "CALL db.create.setNodeVectorProperty(doc, 'vectorProperty', $vector)"
     "WITH doc "
     "MERGE (author:Author {name: $authorName})"
     "MERGE (doc)-[:AUTHORED_BY]->(author)"
diff --git a/src/neo4j_genai/indexes.py b/src/neo4j_genai/indexes.py
@@ -29,7 +29,7 @@ def create_vector_index(
     driver: neo4j.Driver,
     name: str,
     label: str,
-    property: str,
+    embedding_property: str,
     dimensions: int,
     similarity_fn: Literal["euclidean", "cosine"],
 ) -> None:
@@ -46,7 +46,7 @@ def create_vector_index(
         driver (neo4j.Driver): Neo4j Python driver instance.
         name (str): The unique name of the index.
         label (str): The node label to be indexed.
-        property (str): The property key of a node which contains embedding values.
+        embedding_property (str): The property key of a node which contains embedding values.
         dimensions (int): Vector embedding dimension
         similarity_fn (str): case-insensitive values for the vector similarity function:
             ``euclidean`` or ``cosine``.
@@ -60,7 +60,7 @@ def create_vector_index(
             driver=driver,
             name=name,
             label=label,
-            property=property,
+            embedding_property=embedding_property,
             dimensions=dimensions,
             similarity_fn=similarity_fn,
         )
@@ -69,7 +69,7 @@ def create_vector_index(
 
     try:
         query = (
-            f"CREATE VECTOR INDEX $name FOR (n:{label}) ON n.{property} OPTIONS "
+            f"CREATE VECTOR INDEX $name FOR (n:{label}) ON n.{embedding_property} OPTIONS "
             "{ indexConfig: { `vector.dimensions`: toInteger($dimensions), `vector.similarity_function`: $similarity_fn } }"
         )
         logger.info(f"Creating vector index named '{name}'")
@@ -149,7 +149,7 @@ def drop_index_if_exists(driver: neo4j.Driver, name: str) -> None:
 def upsert_vector(
     driver: neo4j.Driver,
     node_id: int,
-    vector_prop: str,
+    embedding_property: str,
     vector: list[float],
 ) -> None:
     """
@@ -158,7 +158,7 @@ def upsert_vector(
     Args:
         driver (neo4j.Driver): Neo4j Python driver instance.
         node_id (int): The id of the node.
-        vector_prop (str): The name of the property to store the vector in.
+        embedding_property (str): The name of the property to store the vector in.
         vector (list[float]): The vector to store.
 
     Raises:
@@ -169,12 +169,12 @@ def upsert_vector(
         MATCH (n)
         WHERE elementId(n) = $id
         WITH n
-        CALL db.create.setNodeVectorProperty(n, $vector_prop, $vector)
+        CALL db.create.setNodeVectorProperty(n, $embedding_property, $vector)
         RETURN n
         """
         parameters = {
             "id": node_id,
-            "vector_prop": vector_prop,
+            "embedding_property": embedding_property,
             "vector": vector,
         }
         driver.execute_query(query, parameters)
diff --git a/src/neo4j_genai/retrievers/base.py b/src/neo4j_genai/retrievers/base.py
@@ -99,11 +99,11 @@ def get_result_formatter(self) -> Callable[[neo4j.Record], RetrieverResultItem]:
         """
         Returns the function to use to transform a neo4j.Record to a RetrieverResultItem.
         """
-        if hasattr(self, "format_record_function"):
-            return self.format_record_function or self.default_format_record
-        return self.default_format_record
+        if hasattr(self, "result_formatter"):
+            return self.result_formatter or self.default_record_formatter
+        return self.default_record_formatter
 
-    def default_format_record(self, record: neo4j.Record) -> RetrieverResultItem:
+    def default_record_formatter(self, record: neo4j.Record) -> RetrieverResultItem:
         """
         Best effort to guess the node to text method. Inherited classes
         can override this method to implement custom text formatting.
diff --git a/src/neo4j_genai/retrievers/external/pinecone/pinecone.py b/src/neo4j_genai/retrievers/external/pinecone/pinecone.py
@@ -77,7 +77,7 @@ class PineconeNeo4jRetriever(ExternalRetriever):
         id_property_neo4j (str): The name of the Neo4j node property that's used as the identifier for relating matches from Weaviate to Neo4j nodes.
         embedder (Optional[Embedder]): Embedder object to embed query text.
         return_properties (Optional[list[str]]): List of node properties to return.
-        format_record_function (Optional[Callable[[Any], Any]]): Function to transform a neo4j.Record to a RetrieverResultItem.
+        result_formatter (Optional[Callable[[Any], Any]]): Function to transform a neo4j.Record to a RetrieverResultItem.
 
     Raises:
         RetrieverInitializationError: If validation of the input arguments fail.
@@ -92,7 +92,7 @@ def __init__(
         embedder: Optional[Embedder] = None,
         return_properties: Optional[list[str]] = None,
         retrieval_query: Optional[str] = None,
-        format_record_function: Optional[Callable[[Any], Any]] = None,
+        result_formatter: Optional[Callable[[Any], Any]] = None,
     ):
         try:
             driver_model = Neo4jDriverModel(driver=driver)
@@ -106,7 +106,7 @@ def __init__(
                 embedder_model=embedder_model,
                 return_properties=return_properties,
                 retrieval_query=retrieval_query,
-                format_record_function=format_record_function,
+                result_formatter=result_formatter,
             )
         except ValidationError as e:
             raise RetrieverInitializationError(e.errors())
@@ -127,7 +127,7 @@ def __init__(
         )
         self.return_properties = validated_data.return_properties
         self.retrieval_query = validated_data.retrieval_query
-        self.format_record_function = validated_data.format_record_function
+        self.result_formatter = validated_data.result_formatter
 
     def _get_search_results(
         self,
diff --git a/src/neo4j_genai/retrievers/external/pinecone/types.py b/src/neo4j_genai/retrievers/external/pinecone/types.py
@@ -48,4 +48,4 @@ class PineconeNeo4jRetrieverModel(BaseModel):
     embedder_model: Optional[EmbedderModel] = None
     return_properties: Optional[list[str]] = None
     retrieval_query: Optional[str] = None
-    format_record_function: Optional[Callable[[neo4j.Record], str]] = None
+    result_formatter: Optional[Callable[[neo4j.Record], str]] = None
diff --git a/src/neo4j_genai/retrievers/external/weaviate/types.py b/src/neo4j_genai/retrievers/external/weaviate/types.py
@@ -51,7 +51,7 @@ class WeaviateNeo4jRetrieverModel(BaseModel):
     embedder_model: Optional[EmbedderModel]
     return_properties: Optional[list[str]] = None
     retrieval_query: Optional[str] = None
-    format_record_function: Optional[Callable[[neo4j.Record], str]] = None
+    result_formatter: Optional[Callable[[neo4j.Record], str]] = None
 
 
 class WeaviateNeo4jSearchModel(BaseModel):
diff --git a/src/neo4j_genai/retrievers/external/weaviate/weaviate.py b/src/neo4j_genai/retrievers/external/weaviate/weaviate.py
@@ -68,7 +68,7 @@ class WeaviateNeo4jRetriever(ExternalRetriever):
         id_property_neo4j (str): The name of the Neo4j node property that's used as the identifier for relating matches from Weaviate to Neo4j nodes.
         embedder (Optional[Embedder]): Embedder object to embed query text.
         return_properties (Optional[list[str]]): List of node properties to return.
-        format_record_function (Optional[Callable[[Any], Any]]): Function to transform a neo4j.Record to a RetrieverResultItem.
+        result_formatter (Optional[Callable[[Any], Any]]): Function to transform a neo4j.Record to a RetrieverResultItem.
 
     Raises:
         RetrieverInitializationError: If validation of the input arguments fail.
@@ -84,7 +84,7 @@ def __init__(
         embedder: Optional[Embedder] = None,
         return_properties: Optional[list[str]] = None,
         retrieval_query: Optional[str] = None,
-        format_record_function: Optional[Callable[[Any], Any]] = None,
+        result_formatter: Optional[Callable[[Any], Any]] = None,
     ):
         try:
             driver_model = Neo4jDriverModel(driver=driver)
@@ -99,7 +99,7 @@ def __init__(
                 embedder_model=embedder_model,
                 return_properties=return_properties,
                 retrieval_query=retrieval_query,
-                format_record_function=format_record_function,
+                result_formatter=result_formatter,
             )
         except ValidationError as e:
             raise RetrieverInitializationError(e.errors())
@@ -115,7 +115,7 @@ def __init__(
         )
         self.return_properties = validated_data.return_properties
         self.retrieval_query = validated_data.retrieval_query
-        self.format_record_function = validated_data.format_record_function
+        self.result_formatter = validated_data.result_formatter
 
     def _get_search_results(
         self,
diff --git a/src/neo4j_genai/retrievers/hybrid.py b/src/neo4j_genai/retrievers/hybrid.py
diff --git a/src/neo4j_genai/retrievers/vector.py b/src/neo4j_genai/retrievers/vector.py
diff --git a/src/neo4j_genai/types.py b/src/neo4j_genai/types.py
diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py
diff --git a/tests/e2e/utils.py b/tests/e2e/utils.py
diff --git a/tests/unit/retrievers/test_hybrid.py b/tests/unit/retrievers/test_hybrid.py
diff --git a/tests/unit/retrievers/test_vector.py b/tests/unit/retrievers/test_vector.py
diff --git a/tests/unit/test_indexes.py b/tests/unit/test_indexes.py

Original file line number	Diff line number	Diff line change
`@@ -44,7 +44,7 @@ def formatter(record: neo4j.Record) -> RetrieverResultItem:`
`44`	`44`	`driver,`
`45`	`45`	`index_name=INDEX,`
`46`	`46`	`retrieval_query="with node, score return node.title as title, node.plot as plot",`
`47`		`- format_record_function=formatter,`
	`47`	`+ result_formatter=formatter,`
`48`	`48`	`embedder=embedder,`
`49`	`49`	`)`
`50`	`50`
Original file line number	Diff line number	Diff line change
`@@ -45,7 +45,7 @@ def formatter(record: neo4j.Record) -> RetrieverResultItem:`
`45`	`45`	`driver,`
`46`	`46`	`index_name=INDEX,`
`47`	`47`	`retrieval_query="with node, score return node.title as title, node.plot as plot",`
`48`		`- format_record_function=formatter,`
	`48`	`+ result_formatter=formatter,`
`49`	`49`	`embedder=embedder,`
`50`	`50`	`)`
`51`	`51`