Coverage for haystack/components/embedders/mock_document_embedder.py: 100%
58 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-21 13:53 +0000
1# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
2#
3# SPDX-License-Identifier: Apache-2.0
5from __future__ import annotations
7from dataclasses import replace
8from typing import Any
10from haystack import Document, component, default_from_dict, default_to_dict
11from haystack.components.embedders.mock_utils import (
12 EmbeddingFn,
13 _coerce_embedding,
14 _deterministic_embedding,
15 _estimate_usage,
16)
17from haystack.utils import deserialize_callable, serialize_callable
20@component
21class MockDocumentEmbedder:
22 """
23 A Document Embedder that returns deterministic embeddings without calling any API.
25 It is a drop-in replacement for real Document Embedders (such as `OpenAIDocumentEmbedder`) in tests, smoke tests,
26 and quick prototypes. It implements the same interface (`run`, `run_async`, serialization) but never contacts an
27 external service, so it is fully deterministic and free to run.
29 The embedding is selected based on how the component is configured:
31 - **Deterministic (default)**: with no configuration, each document's embedding is derived from a hash of its
32 (prepared) text. The same text always yields the same embedding, and different texts yield different
33 embeddings, so the mock works in retrieval pipelines and is reproducible across runs and processes.
34 - **Fixed embedding**: pass an `embedding` vector. The same vector is assigned to every document.
35 - **Dynamic embedding**: pass an `embedding_fn` callable that receives the (prepared) text of a document and
36 returns the embedding. This is useful when the embedding should depend on the input in a custom way.
38 Like real Document Embedders, the metadata fields listed in `meta_fields_to_embed` are concatenated with the
39 document content before embedding, so the deterministic embedding reflects the embedded metadata.
41 ### Usage example
43 ```python
44 from haystack import Document
45 from haystack.components.embedders import MockDocumentEmbedder
47 embedder = MockDocumentEmbedder(dimension=8)
48 result = embedder.run([Document(content="I love pizza!")])
49 print(result["documents"][0].embedding) # a deterministic list of 8 floats
50 ```
51 """
53 def __init__(
54 self,
55 embedding: list[float] | None = None,
56 *,
57 embedding_fn: EmbeddingFn | None = None,
58 dimension: int = 768,
59 model: str = "mock-model",
60 meta: dict[str, Any] | None = None,
61 prefix: str = "",
62 suffix: str = "",
63 meta_fields_to_embed: list[str] | None = None,
64 embedding_separator: str = "\n",
65 progress_bar: bool = False,
66 ) -> None:
67 """
68 Creates an instance of MockDocumentEmbedder.
70 :param embedding: An optional fixed embedding assigned to every document. Mutually exclusive with
71 `embedding_fn`. If neither is provided, a deterministic embedding is derived from each document's text.
72 :param embedding_fn: An optional callable that receives the prepared text of a document and returns the
73 embedding as a list of floats. Mutually exclusive with `embedding`. To support serialization, pass a
74 named function (lambdas and nested functions cannot be serialized).
75 :param dimension: The number of dimensions of the deterministic embedding. Ignored when `embedding` or
76 `embedding_fn` is provided, since their length is determined by the value or callable.
77 :param model: The model name reported in the metadata. Purely cosmetic; no model is loaded.
78 :param meta: Additional metadata merged into the output `meta`.
79 :param prefix: A string to add at the beginning of each text before embedding.
80 :param suffix: A string to add at the end of each text before embedding.
81 :param meta_fields_to_embed: List of metadata fields to embed along with the document text.
82 :param embedding_separator: Separator used to concatenate the metadata fields to the document text.
83 :param progress_bar: Accepted for interface compatibility with real Document Embedders and ignored.
84 :raises ValueError: If both `embedding` and `embedding_fn` are provided, if `dimension` is not positive, or
85 if `embedding` is an empty list.
86 :raises TypeError: If `embedding` is not a sequence of numbers.
87 """
88 if embedding is not None and embedding_fn is not None:
89 raise ValueError("Pass either 'embedding' or 'embedding_fn', not both.")
90 if dimension <= 0:
91 raise ValueError("'dimension' must be a positive integer.")
93 self.embedding = _coerce_embedding(embedding, name="'embedding'") if embedding is not None else None
94 self.embedding_fn = embedding_fn
95 self.dimension = dimension
96 self.model = model
97 self.meta = meta or {}
98 self.prefix = prefix
99 self.suffix = suffix
100 self.meta_fields_to_embed = meta_fields_to_embed or []
101 self.embedding_separator = embedding_separator
102 self.progress_bar = progress_bar
103 self._is_warmed_up = False
105 def to_dict(self) -> dict[str, Any]:
106 """Serialize the component to a dictionary."""
107 embedding_fn = serialize_callable(self.embedding_fn) if self.embedding_fn is not None else None
108 return default_to_dict(
109 self,
110 embedding=self.embedding,
111 embedding_fn=embedding_fn,
112 dimension=self.dimension,
113 model=self.model,
114 meta=self.meta,
115 prefix=self.prefix,
116 suffix=self.suffix,
117 meta_fields_to_embed=self.meta_fields_to_embed,
118 embedding_separator=self.embedding_separator,
119 progress_bar=self.progress_bar,
120 )
122 @classmethod
123 def from_dict(cls, data: dict[str, Any]) -> MockDocumentEmbedder:
124 """Deserialize the component from a dictionary."""
125 init_params = data.get("init_parameters", {})
126 embedding_fn = init_params.get("embedding_fn")
127 if embedding_fn:
128 init_params["embedding_fn"] = deserialize_callable(embedding_fn)
129 return default_from_dict(cls, data)
131 def warm_up(self) -> None:
132 """No-op warm up, provided for interface compatibility with real Embedders."""
133 self._is_warmed_up = True
135 def _prepare_text_to_embed(self, document: Document) -> str:
136 """Concatenate the document content with the metadata fields to embed, mirroring real Document Embedders."""
137 meta_values_to_embed = [
138 str(document.meta[key])
139 for key in self.meta_fields_to_embed
140 if key in document.meta and document.meta[key] is not None
141 ]
142 return (
143 self.prefix + self.embedding_separator.join([*meta_values_to_embed, document.content or ""]) + self.suffix
144 )
146 def _embed(self, text: str) -> list[float]:
147 """Produce the embedding for the prepared text according to the configured mode."""
148 if self.embedding_fn is not None:
149 return _coerce_embedding(self.embedding_fn(text), name="the return value of 'embedding_fn'")
150 if self.embedding is not None:
151 return list(self.embedding)
152 return _deterministic_embedding(text, self.dimension)
154 @component.output_types(documents=list[Document], meta=dict[str, Any])
155 def run(self, documents: list[Document]) -> dict[str, Any]:
156 """
157 Return the input documents with deterministic embeddings added, without calling any API.
159 :param documents: A list of documents to embed.
160 :returns: A dictionary with the following keys:
161 - `documents`: A list of documents with embeddings.
162 - `meta`: Metadata about the (mock) model.
163 :raises TypeError: If `documents` is not a list of `Document` objects.
164 """
165 self.warm_up()
167 if not isinstance(documents, list) or (documents and not isinstance(documents[0], Document)):
168 raise TypeError(
169 "MockDocumentEmbedder expects a list of Documents as input. "
170 "In case you want to embed a string, please use the MockTextEmbedder."
171 )
173 texts_to_embed = [self._prepare_text_to_embed(document) for document in documents]
174 new_documents = [
175 replace(document, embedding=self._embed(text))
176 for document, text in zip(documents, texts_to_embed, strict=True)
177 ]
179 meta: dict[str, Any] = {"model": self.model, "usage": _estimate_usage(texts_to_embed)}
180 meta.update(self.meta)
181 return {"documents": new_documents, "meta": meta}
183 @component.output_types(documents=list[Document], meta=dict[str, Any])
184 async def run_async(self, documents: list[Document]) -> dict[str, Any]:
185 """
186 Asynchronously return the input documents with deterministic embeddings added, without calling any API.
188 :param documents: A list of documents to embed.
189 :returns: A dictionary with the following keys:
190 - `documents`: A list of documents with embeddings.
191 - `meta`: Metadata about the (mock) model.
192 :raises TypeError: If `documents` is not a list of `Document` objects.
193 """
194 return self.run(documents=documents)