Moved Embedder class to under neo4j_graphrag.embeddings (#138)

willtai · web-flow · commit fc7d319c6390 · 2024-09-18T13:29:19.000+01:00
* Moved Embedder class to under neo4j_graphrag.embeddings

* Update documentation

* Rename to base and add await to run example in docs

* Update CHANGELOG
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,6 +1,7 @@
 # @neo4j/neo4j-graphrag-python
 
 ## Next
+
 ### Added
 - Add `template` validation in `PromptTemplate` class upon construction.
 - `custom_prompt` arg is now converted to `Text2CypherTemplate` class within the `Text2CypherRetriever.get_search_results` method.
@@ -9,6 +10,9 @@
 - Add feature to include kwargs in `Text2CypherRetriever.search()` that will be injected into a custom prompt, if provided.
 - Add validation to `custom_prompt` parameter of `Text2CypherRetriever` to ensure that `query_text` placeholder exists in prompt.
 
+### Changed
+- Moved the Embedder class to the neo4j_graphrag.embeddings directory for better organization alongside other custom embedders.
+
 ## 0.6.3
 ### Changed
 - Updated documentation links in README.
@@ -128,4 +132,4 @@
 ### Fixed
 
 -   Updated documentation to include new custom exceptions.
--   Improved the use of Pydantic for input data validation for retriever objects.
+-   Improved the use of Pydantic for input data validation for retriever objects.
diff --git a/docs/source/api.rst b/docs/source/api.rst
@@ -135,7 +135,7 @@ PineconeNeo4jRetriever
 Embedder
 ********
 
-.. autoclass:: neo4j_graphrag.embedder.Embedder
+.. autoclass:: neo4j_graphrag.embeddings.base.Embedder
     :members:
 
 SentenceTransformerEmbeddings
@@ -147,13 +147,13 @@ SentenceTransformerEmbeddings
 OpenAIEmbeddings
 ================
 
-.. autoclass:: neo4j_genai.embeddings.openai.OpenAIEmbeddings
+.. autoclass:: neo4j_graphrag.embeddings.openai.OpenAIEmbeddings
     :members:
 
 VertexAIEmbeddings
 ==================
 
-.. autoclass:: neo4j_genai.embeddings.vertexai.VertexAIEmbeddings
+.. autoclass:: neo4j_graphrag.embeddings.vertexai.VertexAIEmbeddings
     :members:
 
 **********
diff --git a/docs/source/user_guide_kg_builder.rst b/docs/source/user_guide_kg_builder.rst
@@ -188,7 +188,7 @@ Here is a code block illustrating these concepts:
 
     schema_builder = SchemaBuilder()
 
-    schema_builder.run(
+    await schema_builder.run(
         entities=[
             SchemaEntity(
                 label="Person",
diff --git a/docs/source/user_guide_rag.rst b/docs/source/user_guide_rag.rst
@@ -277,7 +277,7 @@ the following implementation of an embedder that wraps the `OllamaEmbedding` mod
 .. code:: python
 
     from llama_index.embeddings.ollama import OllamaEmbedding
-    from neo4j_graphrag.embedder import Embedder
+    from neo4j_graphrag.embeddings.base import Embedder
 
     class OllamaEmbedder(Embedder):
         def __init__(self, ollama_embedding):
diff --git a/examples/hybrid_cypher_search.py b/examples/hybrid_cypher_search.py
@@ -3,7 +3,7 @@
 from random import random
 
 from neo4j import GraphDatabase
-from neo4j_graphrag.embedder import Embedder
+from neo4j_graphrag.embeddings.base import Embedder
 from neo4j_graphrag.indexes import create_fulltext_index, create_vector_index
 from neo4j_graphrag.retrievers import HybridCypherRetriever
 
diff --git a/examples/hybrid_search.py b/examples/hybrid_search.py
@@ -3,7 +3,7 @@
 from random import random
 
 from neo4j import GraphDatabase
-from neo4j_graphrag.embedder import Embedder
+from neo4j_graphrag.embeddings.base import Embedder
 from neo4j_graphrag.indexes import create_fulltext_index, create_vector_index
 from neo4j_graphrag.retrievers import HybridRetriever
 
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 import GraphDatabase
-from neo4j_graphrag.embedder import Embedder
+from neo4j_graphrag.embeddings.base import Embedder
 from neo4j_graphrag.indexes import create_vector_index
 from neo4j_graphrag.retrievers import VectorRetriever
 
diff --git a/examples/vector_cypher_retrieval.py b/examples/vector_cypher_retrieval.py
@@ -4,7 +4,7 @@
 import string
 
 from neo4j import GraphDatabase
-from neo4j_graphrag.embedder import Embedder
+from neo4j_graphrag.embeddings.base import Embedder
 from neo4j_graphrag.indexes import create_vector_index
 from neo4j_graphrag.retrievers import VectorCypherRetriever
 
diff --git a/examples/vector_search_with_filters.py b/examples/vector_search_with_filters.py
@@ -4,7 +4,7 @@
 import string
 
 from neo4j import GraphDatabase
-from neo4j_graphrag.embedder import Embedder
+from neo4j_graphrag.embeddings.base import Embedder
 from neo4j_graphrag.indexes import create_vector_index
 from neo4j_graphrag.retrievers import VectorRetriever
 
diff --git a/src/neo4j_graphrag/embeddings/base.py b/src/neo4j_graphrag/embeddings/base.py
diff --git a/src/neo4j_graphrag/embeddings/openai.py b/src/neo4j_graphrag/embeddings/openai.py
@@ -17,7 +17,7 @@
 
 from typing import Any
 
-from neo4j_graphrag.embedder import Embedder
+from neo4j_graphrag.embeddings.base import Embedder
 
 
 class OpenAIEmbeddings(Embedder):
diff --git a/src/neo4j_graphrag/embeddings/sentence_transformers.py b/src/neo4j_graphrag/embeddings/sentence_transformers.py
@@ -18,7 +18,7 @@
 import numpy as np
 import torch
 
-from neo4j_graphrag.embedder import Embedder
+from neo4j_graphrag.embeddings.base import Embedder
 
 
 class SentenceTransformerEmbeddings(Embedder):
diff --git a/src/neo4j_graphrag/embeddings/vertexai.py b/src/neo4j_graphrag/embeddings/vertexai.py
@@ -17,7 +17,7 @@
 
 from typing import Any
 
-from neo4j_graphrag.embedder import Embedder
+from neo4j_graphrag.embeddings.base import Embedder
 
 try:
     from vertexai.language_models import TextEmbeddingInput, TextEmbeddingModel
diff --git a/src/neo4j_graphrag/experimental/components/embedder.py b/src/neo4j_graphrag/experimental/components/embedder.py
@@ -14,7 +14,7 @@
 #  limitations under the License.
 from pydantic import validate_call
 
-from neo4j_graphrag.embedder import Embedder
+from neo4j_graphrag.embeddings.base import Embedder
 from neo4j_graphrag.experimental.components.types import TextChunk, TextChunks
 from neo4j_graphrag.experimental.pipeline.component import Component
 
diff --git a/src/neo4j_graphrag/retrievers/external/pinecone/pinecone.py b/src/neo4j_graphrag/retrievers/external/pinecone/pinecone.py
@@ -21,7 +21,7 @@
 from pinecone import Pinecone
 from pydantic import ValidationError
 
-from neo4j_graphrag.embedder import Embedder
+from neo4j_graphrag.embeddings.base import Embedder
 from neo4j_graphrag.exceptions import (
     EmbeddingRequiredError,
     RetrieverInitializationError,
diff --git a/src/neo4j_graphrag/retrievers/external/weaviate/weaviate.py b/src/neo4j_graphrag/retrievers/external/weaviate/weaviate.py
@@ -22,7 +22,7 @@
 from pydantic import ValidationError
 from weaviate.client import WeaviateClient
 
-from neo4j_graphrag.embedder import Embedder
+from neo4j_graphrag.embeddings.base import Embedder
 from neo4j_graphrag.exceptions import (
     RetrieverInitializationError,
     SearchValidationError,
diff --git a/src/neo4j_graphrag/retrievers/hybrid.py b/src/neo4j_graphrag/retrievers/hybrid.py
@@ -20,7 +20,7 @@
 import neo4j
 from pydantic import ValidationError
 
-from neo4j_graphrag.embedder import Embedder
+from neo4j_graphrag.embeddings.base import Embedder
 from neo4j_graphrag.exceptions import (
     EmbeddingRequiredError,
     RetrieverInitializationError,
diff --git a/src/neo4j_graphrag/retrievers/vector.py b/src/neo4j_graphrag/retrievers/vector.py
@@ -20,7 +20,7 @@
 import neo4j
 from pydantic import ValidationError
 
-from neo4j_graphrag.embedder import Embedder
+from neo4j_graphrag.embeddings.base import Embedder
 from neo4j_graphrag.exceptions import (
     EmbeddingRequiredError,
     RetrieverInitializationError,
diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py
@@ -22,7 +22,7 @@
 
 import pytest
 from neo4j import Driver, GraphDatabase
-from neo4j_graphrag.embedder import Embedder
+from neo4j_graphrag.embeddings.base import Embedder
 from neo4j_graphrag.indexes import (
     create_fulltext_index,
     create_vector_index,
diff --git a/tests/e2e/pinecone_e2e/test_pinecone_e2e.py b/tests/e2e/pinecone_e2e/test_pinecone_e2e.py
@@ -19,7 +19,7 @@
 
 import pytest
 from langchain_huggingface.embeddings import HuggingFaceEmbeddings
-from neo4j_graphrag.embedder import Embedder
+from neo4j_graphrag.embeddings.base import Embedder
 from neo4j_graphrag.retrievers import PineconeNeo4jRetriever
 from neo4j_graphrag.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_graphrag.embedder import Embedder
+from neo4j_graphrag.embeddings.base import Embedder
 from neo4j_graphrag.retrievers import (
     HybridCypherRetriever,
     HybridRetriever,
diff --git a/tests/e2e/test_kg_builder_pipeline_e2e.py b/tests/e2e/test_kg_builder_pipeline_e2e.py
@@ -21,7 +21,7 @@
 import neo4j
 import pytest
 from langchain_text_splitters import CharacterTextSplitter
-from neo4j_graphrag.embedder import Embedder
+from neo4j_graphrag.embeddings.base import Embedder
 from neo4j_graphrag.exceptions import LLMGenerationError
 from neo4j_graphrag.experimental.components.embedder import TextChunkEmbedder
 from neo4j_graphrag.experimental.components.entity_relation_extractor import (
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_graphrag.embedder import Embedder
+from neo4j_graphrag.embeddings.base import Embedder
 from neo4j_graphrag.retrievers import VectorCypherRetriever, VectorRetriever
 from neo4j_graphrag.types import RetrieverResult, RetrieverResultItem
 
diff --git a/tests/e2e/weaviate_e2e/test_weaviate_e2e.py b/tests/e2e/weaviate_e2e/test_weaviate_e2e.py
@@ -19,7 +19,7 @@
 import pytest
 from langchain_huggingface.embeddings import HuggingFaceEmbeddings
 from neo4j import Driver
-from neo4j_graphrag.embedder import Embedder
+from neo4j_graphrag.embeddings.base import Embedder
 from neo4j_graphrag.retrievers import WeaviateNeo4jRetriever
 from neo4j_graphrag.types import RetrieverResult, RetrieverResultItem
 from weaviate.client import WeaviateClient
diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py
@@ -18,7 +18,7 @@
 
 import neo4j
 import pytest
-from neo4j_graphrag.embedder import Embedder
+from neo4j_graphrag.embeddings.base import Embedder
 from neo4j_graphrag.llm import LLMInterface
 from neo4j_graphrag.retrievers import (
     HybridRetriever,
diff --git a/tests/unit/embeddings/test_sentence_transformers.py b/tests/unit/embeddings/test_sentence_transformers.py
@@ -2,7 +2,7 @@
 
 import numpy as np
 import pytest
-from neo4j_graphrag.embedder import Embedder
+from neo4j_graphrag.embeddings.base import Embedder
 from neo4j_graphrag.embeddings.sentence_transformers import (
     SentenceTransformerEmbeddings,
 )
