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

* Moved Embedder class to under neo4j_graphrag.embeddings

* Update documentation

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

* Update CHANGELOG

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.

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:
 
 **********

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",

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):

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
 

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
 

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
 

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
 

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
 

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):

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):

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

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
 

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,

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,

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,

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,

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,

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

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,

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 (

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
 

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

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,

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,
 )