Save embedding_model_name in the Chunk node

CHANGELOG.md

@@ -5,10 +5,12 @@
 ### Added
 
 - `AnthropicLLM` now supports structured output via the `response_format` argument, accepting a Pydantic model or an Anthropic `output_config` dict, alongside `OpenAILLM` and `VertexAILLM`.
+- Chunk node metadata now contains `embedding_model_name` and `embedding_dimensions
 
 ### Changed
 
 - (**breaking**) `AnthropicLLM.supports_structured_output` is now `True`. As a result, `SchemaFromTextExtractor` and `LLMEntityRelationExtractor` (and `SimpleKGPipeline`, which enables structured output automatically when the LLM supports it) now use structured output by default with `AnthropicLLM`. This requires a Claude 4.5+ model (e.g. `claude-sonnet-4-5`); using `AnthropicLLM` with an older Claude model in these components will now raise an error where it previously worked. To keep the previous behavior, use a Claude 4.5+ model, or construct `LLMEntityRelationExtractor` / `SchemaFromTextExtractor` directly with `use_structured_output=False`.
+- (**breaking**) `Embedder` subclasses now require a `model` argument and optionally a `dimensions` argument.
 
 - Preparation for 2.0:
   - All `Components` have been moved out of the `experimental` namespace

docs/source/user_guide_kg_builder.rst

@@ -420,7 +420,7 @@ The same principles apply to `embedder_config`:
         "embedder_config": {
             "class_": "OpenAIEmbeddings",
             "params_": {
-                "mode": "text-embedding-ada-002",
+                "model": "text-embedding-ada-002",
                 "api_key": {
                     "resolver_": "ENV",
                     "var_": "OPENAI_API_KEY",

examples/customize/embeddings/custom_embeddings.py

@@ -5,14 +5,13 @@
 
 
 class CustomEmbeddings(Embedder):
-    def __init__(self, dimension: int = 10, **kwargs: Any):
-        super().__init__(**kwargs)
-        self.dimension = dimension
+    def __init__(self, model: str, dimensions: int = 10, **kwargs: Any):
+        super().__init__(model, dimensions, **kwargs)
 
     def embed_query(self, input: str) -> list[float]:
-        return [random.random() for _ in range(self.dimension)]
+        return [random.random() for _ in range(self.dimensions)]
 
 
-llm = CustomEmbeddings(dimensions=1024)
+llm = CustomEmbeddings("", dimensions=1024)
 res = llm.embed_query("text")
 print(res[:10])

examples/customize/retrievers/hybrid_retrievers/hybrid_cypher_search.py

@@ -20,6 +20,9 @@
 
 # Create Embedder object
 class CustomEmbedder(Embedder):
+    def __init__(self) -> None:
+        super().__init__(model="custom")
+
     def embed_query(self, text: str) -> list[float]:
         return [random() for _ in range(DIMENSION)]
 

examples/customize/retrievers/hybrid_retrievers/hybrid_search.py

@@ -20,6 +20,9 @@
 
 # Create Embedder object
 class CustomEmbedder(Embedder):
+    def __init__(self) -> None:
+        super().__init__(model="custom")
+
     def embed_query(self, text: str) -> list[float]:
         return [random() for _ in range(DIMENSION)]
 

pyproject.toml

@@ -50,7 +50,7 @@ google-genai = ["google-genai>=1.62.0,<2.0.0"]
 cohere = ["cohere>=5.9.0,<6.0.0"]
 anthropic = ["anthropic>=0.77.0,<1.0.0"]
 bedrock = ["boto3>=1.35.0,<2.0.0"]
-ollama = ["ollama>=0.4.4,<0.5.0"]
+ollama = ["ollama>=0.6.0,<0.10.0"]
 openai = ["openai>=1.51.1,<2.0.0"]
 mistralai = ["mistralai>=1.0.3,<2.0.0"]
 litellm = ["litellm>=1.0.0,<2.0.0"]

src/neo4j_graphrag/components/embedder.py

@@ -83,6 +83,7 @@ async def _async_embed_chunk(
             embedding = await self._embedder.async_embed_query(text_chunk.text)
             metadata = text_chunk.metadata if text_chunk.metadata else {}
             metadata["embedding"] = embedding
+            metadata["embedding_model_name"] = getattr(self._embedder, "model", None)
             return TextChunk(
                 text=text_chunk.text,
                 index=text_chunk.index,

src/neo4j_graphrag/embeddings/base.py

@@ -14,6 +14,7 @@
 #  limitations under the License.
 from __future__ import annotations
 
+import warnings
 from abc import ABC, abstractmethod
 from typing import Optional
 
@@ -29,10 +30,32 @@ class Embedder(ABC):
     An embedder passed into a retriever must implement this interface.
 
     Args:
+        model (str): The model name. Subclasses must pass this to super().__init__(); omitting it is deprecated and will be required in 2.0.
+        dimensions (Optional[int]): The number of dimensions of the embeddings. Defaults to None.
         rate_limit_handler (Optional[RateLimitHandler]): Handler for rate limiting. Defaults to retry with exponential backoff.
     """
 
-    def __init__(self, rate_limit_handler: Optional[RateLimitHandler] = None):
+    def __init__(
+        self,
+        model: str = "",
+        dimensions: int | None = None,
+        rate_limit_handler: Optional[RateLimitHandler] = None,
+    ):
+        if not isinstance(model, str):
+            raise TypeError(
+                f"Embedder.__init__() 'model' must be a str, got {type(model).__name__}. "
+                "If you are passing a rate_limit_handler, use the keyword argument: "
+                "super().__init__(model='my-model', rate_limit_handler=handler)."
+            )
+        if not model:
+            warnings.warn(
+                "Embedder subclasses should pass 'model' to super().__init__(). "
+                "This will be required in 2.0.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+        self.model = model
+        self.dimensions = dimensions
         if rate_limit_handler is not None:
             self._rate_limit_handler = rate_limit_handler
         else:

src/neo4j_graphrag/embeddings/bedrock.py

@@ -17,6 +17,7 @@
 import asyncio
 import json
 import os
+import warnings
 from typing import Any, Optional
 
 from neo4j_graphrag.embeddings.base import Embedder
@@ -82,15 +83,21 @@ def __init__(
                 "Could not import boto3 python client. "
                 'Please install it with `pip install "neo4j-graphrag[bedrock]"`.'
             )
-        super().__init__(rate_limit_handler)
-        self.model_id = model_id
-        self.dimensions = dimensions
+        super().__init__(model_id, dimensions, rate_limit_handler)
         self.normalize = normalize
         client_kwargs: dict[str, Any] = {**kwargs}
         if region_name:
             client_kwargs["region_name"] = region_name
         self.client = boto3.client("bedrock-runtime", **client_kwargs)
 
+    @property
+    def model_id(self) -> str | None:
+        warnings.warn(
+            "model_id is deprecated. Use model instead.",
+            DeprecationWarning,
+        )
+        return self.model
+
     def _invoke_embedding(self, text: str) -> list[float]:
         """Invoke the Bedrock embedding model and return the embedding vector."""
         body = json.dumps(
@@ -102,7 +109,7 @@ def _invoke_embedding(self, text: str) -> list[float]:
         )
         response = self.client.invoke_model(
             body=body,
-            modelId=self.model_id,
+            modelId=self.model,
             accept="application/json",
             contentType="application/json",
         )

src/neo4j_graphrag/embeddings/cohere.py

@@ -14,6 +14,7 @@
 #  limitations under the License.
 from __future__ import annotations
 
+import warnings
 from typing import Any, Optional
 
 from neo4j_graphrag.embeddings.base import Embedder
@@ -30,6 +31,7 @@ class CohereEmbeddings(Embedder):
     def __init__(
         self,
         model: str = "",
+        dimensions: int | None = None,
         rate_limit_handler: Optional[RateLimitHandler] = None,
         **kwargs: Any,
     ) -> None:
@@ -38,8 +40,11 @@ def __init__(
                 """Could not import cohere python client.
                 Please install it with `pip install "neo4j-graphrag[cohere]"`."""
             )
-        super().__init__(rate_limit_handler)
-        self.model = model
+        super().__init__(model, dimensions, rate_limit_handler)
+        if self.dimensions:
+            warnings.warn(
+                "Dimension parameter is ignored in CohereEmbeddings.", UserWarning
+            )
         self.client = cohere.Client(**kwargs)
 
     @rate_limit_handler

src/neo4j_graphrag/embeddings/google_genai.py

@@ -16,6 +16,7 @@
 
 # built-in dependencies
 import os
+import warnings
 from typing import Any, Optional
 
 # project dependencies
@@ -60,20 +61,24 @@ def __init__(
                 "Could not import google-genai python client. "
                 'Please install it with `pip install "neo4j-graphrag[google-genai]"`.'
             )
-        super().__init__(rate_limit_handler)
-        self.model = model
-        self.embedding_dim = embedding_dim
+        super().__init__(model, embedding_dim, rate_limit_handler)
         self.client = genai.Client(**kwargs)
 
+    @property
+    def embedding_dim(self) -> int | None:
+        warnings.warn(
+            "embedding_dim is deprecated. Use dimensions instead",
+            DeprecationWarning,
+        )
+        return self.dimensions
+
     @rate_limit_handler
     def embed_query(self, text: str, **kwargs: Any) -> list[float]:
         try:
             result = self.client.models.embed_content(
                 model=self.model,
                 contents=[text],  # type: ignore[arg-type]
-                config=types.EmbedContentConfig(
-                    output_dimensionality=self.embedding_dim
-                ),
+                config=types.EmbedContentConfig(output_dimensionality=self.dimensions),
                 **kwargs,
             )
             if not result or not result.embeddings or not result.embeddings[0].values:
@@ -90,9 +95,7 @@ async def async_embed_query(self, text: str, **kwargs: Any) -> list[float]:
             result = await self.client.aio.models.embed_content(
                 model=self.model,
                 contents=[text],  # type: ignore[arg-type]
-                config=types.EmbedContentConfig(
-                    output_dimensionality=self.embedding_dim
-                ),
+                config=types.EmbedContentConfig(output_dimensionality=self.dimensions),
                 **kwargs,
             )
             if not result or not result.embeddings or not result.embeddings[0].values:

src/neo4j_graphrag/embeddings/mistral.py

@@ -40,6 +40,7 @@ class MistralAIEmbeddings(Embedder):
     def __init__(
         self,
         model: str = "mistral-embed",
+        dimensions: int | None = None,
         rate_limit_handler: Optional[RateLimitHandler] = None,
         **kwargs: Any,
     ) -> None:
@@ -48,11 +49,10 @@ def __init__(
                 """Could not import mistralai.
                 Please install it with `pip install "neo4j-graphrag[mistralai]"`."""
             )
-        super().__init__(rate_limit_handler)
+        super().__init__(model, dimensions, rate_limit_handler)
         api_key = kwargs.pop("api_key", None)
         if api_key is None:
             api_key = os.getenv("MISTRAL_API_KEY", "")
-        self.model = model
         self.mistral_client = Mistral(api_key=api_key, **kwargs)
 
     @rate_limit_handler
@@ -64,9 +64,13 @@ def embed_query(self, text: str, **kwargs: Any) -> list[float]:
             text (str): The text to generate an embedding for.
             **kwargs (Any): Additional keyword arguments to pass to the Mistral AI client.
         """
+        params = {**kwargs}
+        if "output_dimension" not in params:
+            params["output_dimension"] = self.dimensions
+
         try:
             embeddings_batch_response = self.mistral_client.embeddings.create(
-                model=self.model, inputs=[text], **kwargs
+                model=self.model, inputs=[text], **params
             )
         except Exception as e:
             raise EmbeddingsGenerationError(

src/neo4j_graphrag/embeddings/ollama.py

@@ -38,6 +38,7 @@ class OllamaEmbeddings(Embedder):
     def __init__(
         self,
         model: str,
+        dimensions: int | None = None,
         rate_limit_handler: Optional[RateLimitHandler] = None,
         **kwargs: Any,
     ) -> None:
@@ -48,8 +49,7 @@ def __init__(
                 """Could not import ollama python client.
                 Please install it with `pip install "neo4j_graphrag[ollama]"`."""
             )
-        super().__init__(rate_limit_handler)
-        self.model = model
+        super().__init__(model, dimensions, rate_limit_handler)
         self.client = ollama.Client(**kwargs)
         self.async_client = ollama.AsyncClient(**kwargs)
 
@@ -62,10 +62,14 @@ def embed_query(self, text: str, **kwargs: Any) -> list[float]:
             text (str): The text to generate an embedding for.
             **kwargs (Any): Additional keyword arguments to pass to the Ollama client.
         """
+        params = {**kwargs}
+        if "dimensions" not in params:
+            params["dimensions"] = self.dimensions
+
         embeddings_response = self.client.embed(
             model=self.model,
             input=text,
-            **kwargs,
+            **params,
         )
 
         if embeddings_response is None or not embeddings_response.embeddings:
@@ -88,10 +92,13 @@ async def async_embed_query(self, text: str, **kwargs: Any) -> list[float]:
             text (str): The text to generate an embedding for.
             **kwargs (Any): Additional keyword arguments to pass to the Ollama client.
         """
+        params = {**kwargs}
+        if "dimensions" not in params:
+            params["dimensions"] = self.dimensions
         embeddings_response = await self.async_client.embed(
             model=self.model,
             input=text,
-            **kwargs,
+            **params,
         )
 
         if embeddings_response is None or not embeddings_response.embeddings:

src/neo4j_graphrag/embeddings/openai.py

@@ -36,6 +36,7 @@ class BaseOpenAIEmbeddings(Embedder, abc.ABC):
     def __init__(
         self,
         model: str = "text-embedding-ada-002",
+        dimensions: int | None = None,
         rate_limit_handler: Optional[RateLimitHandler] = None,
         **kwargs: Any,
     ) -> None:
@@ -46,9 +47,8 @@ def __init__(
                 """Could not import openai python client.
                 Please install it with `pip install "neo4j-graphrag[openai]"`."""
             )
-        super().__init__(rate_limit_handler)
+        super().__init__(model, dimensions, rate_limit_handler)
         self.openai = openai
-        self.model = model
         self.client = self._initialize_client(**kwargs)
 
     @abc.abstractmethod
@@ -68,9 +68,13 @@ def embed_query(self, text: str, **kwargs: Any) -> list[float]:
             text (str): The text to generate an embedding for.
             **kwargs (Any): Additional arguments to pass to the OpenAI embedding generation function.
         """
+        params = {**kwargs}
+        if "dimensions" not in params:
+            params["dimensions"] = self.dimensions
+
         try:
             response = self.client.embeddings.create(
-                input=text, model=self.model, **kwargs
+                input=text, model=self.model, **params
             )
             embedding: list[float] = response.data[0].embedding
             return embedding