diff --git a/dspy/clients/embedding.py b/dspy/clients/embedding.py index 01bfc1e641..02d831c507 100644 --- a/dspy/clients/embedding.py +++ b/dspy/clients/embedding.py @@ -37,6 +37,15 @@ class Embedder: model. batch_size (int, optional): The default batch size for processing inputs in batches. Defaults to 200. caching (bool, optional): Whether to cache the embedding response when using a hosted model. Defaults to True. + model_id (str | None, optional): A stable, user-supplied identifier for the underlying callable model, used only + to disambiguate the embedding cache key (it is never forwarded to the model). Set this when caching is + enabled and the callable ``model`` may collide in the cache with a *different* callable model that shares + the same source (e.g. two instances/checkpoints of the same ``SentenceTransformer`` class, whose bound + ``encode`` methods are byte-identical) or the same ``__name__`` (e.g. two callable instances or + ``functools.partial`` wrappers, which would both collapse to ``""``). A stable string + (such as the checkpoint name) survives across processes, so the same ``model_id`` reuses the on-disk cache + while distinct ``model_id`` values never share cached results. When unset (``None``), this argument is + omitted from the cache key entirely, preserving existing cache keys. Defaults to None. **kwargs: Additional default keyword arguments to pass to the embedding model. Examples: @@ -81,15 +90,42 @@ def my_embedder(texts): assert embeddings.shape == (2, 10) ``` + + Example 4: Avoiding cache collisions across distinct callable models of the same class. + + The cache key for a callable ``model`` is derived from the callable's source (or ``__name__``), which is + identical across two instances/checkpoints of the same class (e.g. two ``SentenceTransformer`` checkpoints) and + collapses to ``""`` for callable instances / ``functools.partial``. With ``caching=True`` + (the default) and a shared cache namespace (e.g. the default on-disk ``~/.dspy_cache`` reused by every run), + such callables would otherwise collide and return the wrong cached embeddings. Pass a stable, checkpoint-level + ``model_id`` to disambiguate the cache key while still reusing the cache across processes for the same model: + + ```python + import dspy + from sentence_transformers import SentenceTransformer + + # Switching checkpoints across runs that share ~/.dspy_cache: pass the checkpoint + # name as ``model_id`` so each checkpoint keeps its own cached vectors. + embedder = dspy.Embedder(SentenceTransformer("paraphrase-MiniLM-L6-v2").encode, model_id="paraphrase-MiniLM-L6-v2") + embeddings = embedder(["hello", "world"], batch_size=1) + ``` """ - def __init__(self, model: str | Callable, batch_size: int = 200, caching: bool = True, **kwargs: dict[str, Any]): + def __init__( + self, + model: str | Callable, + batch_size: int = 200, + caching: bool = True, + model_id: str | None = None, + **kwargs: dict[str, Any], + ): self.model = model self.batch_size = batch_size self.caching = caching + self.model_id = model_id self.default_kwargs = kwargs - def _preprocess(self, inputs, batch_size=None, caching=None, **kwargs): + def _preprocess(self, inputs, batch_size=None, caching=None, model_id=None, **kwargs): if isinstance(inputs, str): is_single_input = True inputs = [inputs] @@ -101,6 +137,7 @@ def _preprocess(self, inputs, batch_size=None, caching=None, **kwargs): batch_size = batch_size or self.batch_size caching = caching if caching is not None else self.caching + model_id = model_id if model_id is not None else self.model_id merged_kwargs = self.default_kwargs.copy() merged_kwargs.update(kwargs) @@ -108,7 +145,7 @@ def _preprocess(self, inputs, batch_size=None, caching=None, **kwargs): for i in range(0, len(inputs), batch_size): input_batches.append(inputs[i : i + batch_size]) - return input_batches, caching, merged_kwargs, is_single_input + return input_batches, caching, model_id, merged_kwargs, is_single_input def _postprocess(self, embeddings_list, is_single_input): embeddings = np.array(embeddings_list, dtype=np.float32) @@ -117,7 +154,13 @@ def _postprocess(self, embeddings_list, is_single_input): else: return np.array(embeddings, dtype=np.float32) - def __call__(self, inputs: str | list[str], batch_size: int | None = None, caching: bool | None = None, **kwargs: dict[str, Any]) -> np.ndarray: + def __call__( + self, + inputs: str | list[str], + batch_size: int | None = None, + caching: bool | None = None, + **kwargs: dict[str, Any], + ) -> np.ndarray: """Compute embeddings for the given inputs. Args: @@ -126,6 +169,9 @@ def __call__(self, inputs: str | list[str], batch_size: int | None = None, cachi during initialization. caching (bool, optional): Whether to cache the embedding response when using a hosted model. If None, defaults to the caching setting from initialization. + model_id (str | None, optional): If provided, overrides the ``model_id`` set during initialization for this + call only and is folded into the embedding cache key (it is never forwarded to the model). See + ``Embedder`` for the full cache-safety rationale. kwargs: Additional keyword arguments to pass to the embedding model. These will override the default kwargs provided during initialization. @@ -133,28 +179,41 @@ def __call__(self, inputs: str | list[str], batch_size: int | None = None, cachi numpy.ndarray: If the input is a single string, returns a 1D numpy array representing the embedding. If the input is a list of strings, returns a 2D numpy array of embeddings, one embedding per row. """ - input_batches, caching, kwargs, is_single_input = self._preprocess(inputs, batch_size, caching, **kwargs) + input_batches, caching, model_id, kwargs, is_single_input = self._preprocess( + inputs, batch_size, caching, **kwargs + ) compute_embeddings = _cached_compute_embeddings if caching else _compute_embeddings + call_kwargs = {"caching": caching} + if model_id is not None: + call_kwargs["model_id"] = model_id + embeddings_list = [] for batch in input_batches: - embeddings_list.extend(compute_embeddings(self.model, batch, caching=caching, **kwargs)) + embeddings_list.extend(compute_embeddings(self.model, batch, **call_kwargs, **kwargs)) return self._postprocess(embeddings_list, is_single_input) async def acall(self, inputs, batch_size=None, caching=None, **kwargs): - input_batches, caching, kwargs, is_single_input = self._preprocess(inputs, batch_size, caching, **kwargs) + input_batches, caching, model_id, kwargs, is_single_input = self._preprocess( + inputs, batch_size, caching, **kwargs + ) - embeddings_list = [] acompute_embeddings = _cached_acompute_embeddings if caching else _acompute_embeddings + call_kwargs = {"caching": caching} + if model_id is not None: + call_kwargs["model_id"] = model_id + + embeddings_list = [] + for batch in input_batches: - embeddings_list.extend(await acompute_embeddings(self.model, batch, caching=caching, **kwargs)) + embeddings_list.extend(await acompute_embeddings(self.model, batch, **call_kwargs, **kwargs)) return self._postprocess(embeddings_list, is_single_input) -def _compute_embeddings(model, batch_inputs, caching=False, **kwargs): +def _compute_embeddings(model, batch_inputs, caching=False, model_id=None, **kwargs): if isinstance(model, str): caching = caching and _get_litellm().cache is not None embedding_response = _get_litellm().embedding(model=model, input=batch_inputs, caching=caching, **kwargs) @@ -166,11 +225,11 @@ def _compute_embeddings(model, batch_inputs, caching=False, **kwargs): @request_cache(ignored_args_for_cache_key=["api_key", "api_base", "base_url"]) -def _cached_compute_embeddings(model, batch_inputs, caching=True, **kwargs): - return _compute_embeddings(model, batch_inputs, caching=caching, **kwargs) +def _cached_compute_embeddings(model, batch_inputs, caching=True, model_id=None, **kwargs): + return _compute_embeddings(model, batch_inputs, caching=caching, model_id=model_id, **kwargs) -async def _acompute_embeddings(model, batch_inputs, caching=False, **kwargs): +async def _acompute_embeddings(model, batch_inputs, caching=False, model_id=None, **kwargs): if isinstance(model, str): caching = caching and _get_litellm().cache is not None embedding_response = await _get_litellm().aembedding(model=model, input=batch_inputs, caching=caching, **kwargs) @@ -182,5 +241,5 @@ async def _acompute_embeddings(model, batch_inputs, caching=False, **kwargs): @request_cache(ignored_args_for_cache_key=["api_key", "api_base", "base_url"]) -async def _cached_acompute_embeddings(model, batch_inputs, caching=True, **kwargs): - return await _acompute_embeddings(model, batch_inputs, caching=caching, **kwargs) +async def _cached_acompute_embeddings(model, batch_inputs, caching=True, model_id=None, **kwargs): + return await _acompute_embeddings(model, batch_inputs, caching=caching, model_id=model_id, **kwargs) diff --git a/dspy/predict/knn.py b/dspy/predict/knn.py index 68f07b3a63..b28f87ccd4 100644 --- a/dspy/predict/knn.py +++ b/dspy/predict/knn.py @@ -36,6 +36,13 @@ def __init__(self, k: int, trainset: list[Example], vectorizer: Embedder): # Find similar examples similar_examples = knn(input="hello") ``` + + Note: when caching is enabled (the default) and you switch the checkpoint while sharing the default on-disk + cache (``~/.dspy_cache``), pass a checkpoint-level ``model_id`` to the ``Embedder`` so each checkpoint keeps + its own cached vectors -- two ``Embedder`` instances backed by different checkpoints of the same class would + otherwise collide in the cache. For example:: + + vectorizer=dspy.Embedder(SentenceTransformer("paraphrase-MiniLM-L6-v2").encode, model_id="paraphrase-MiniLM-L6-v2") """ self.k = k self.trainset = trainset diff --git a/dspy/teleprompt/knn_fewshot.py b/dspy/teleprompt/knn_fewshot.py index 20124afef5..2bccf6374f 100644 --- a/dspy/teleprompt/knn_fewshot.py +++ b/dspy/teleprompt/knn_fewshot.py @@ -9,7 +9,9 @@ class KNNFewShot(Teleprompter): - def __init__(self, k: int, trainset: list[Example], vectorizer: Embedder, **few_shot_bootstrap_args: dict[str, Any]): + def __init__( + self, k: int, trainset: list[Example], vectorizer: Embedder, **few_shot_bootstrap_args: dict[str, Any] + ): """ KNNFewShot is an optimizer that uses an in-memory KNN retriever to find the k nearest neighbors in a trainset at test time. For each input example in a forward call, it identifies the k most @@ -48,6 +50,13 @@ def __init__(self, k: int, trainset: list[Example], vectorizer: Embedder, **few_ # Use the compiled module result = compiled_qa("What is the capital of Belgium?") ``` + + Note: when caching is enabled (the default) and you switch the checkpoint while sharing the default on-disk + cache (``~/.dspy_cache``), pass a checkpoint-level ``model_id`` to the ``Embedder`` so each checkpoint keeps + its own cached vectors -- two ``Embedder`` instances backed by different checkpoints of the same class would + otherwise collide in the cache. For example:: + + vectorizer=dspy.Embedder(SentenceTransformer("paraphrase-MiniLM-L6-v2").encode, model_id="paraphrase-MiniLM-L6-v2") """ self.KNN = KNN(k, trainset, vectorizer=vectorizer) self.few_shot_bootstrap_args = few_shot_bootstrap_args diff --git a/tests/clients/test_embedding.py b/tests/clients/test_embedding.py index fae0ea42c1..d9e7e076f3 100644 --- a/tests/clients/test_embedding.py +++ b/tests/clients/test_embedding.py @@ -1,3 +1,4 @@ +import inspect from unittest.mock import MagicMock, patch import numpy as np @@ -197,3 +198,219 @@ async def test_acall_caching_true_overrides_instance_false(cache): await embedding.acall(inputs, caching=True) await embedding.acall(inputs, caching=True) assert mock_litellm.call_count == 1 + + +# --- Callable-model cache-collision fixtures and tests --- +# +# The following helper classes are intentionally defined at module top-level (not inside a +# test function) so that ``inspect.getsource`` can retrieve their source. ``FakeSTSource`` +# mirrors ``SentenceTransformer``: ``encode`` is a *class-level* method, so +# ``inspect.getsource(instance.encode)`` returns byte-identical source across two distinct +# instances and triggers the ```` branch of ``_transform_value``. +# ``FakeCallableModel`` defines ``__call__`` but is a callable *instance* (no ``__name__``), +# which triggers the ```` fallback branch. + + +class FakeSTSource: + """Mimics ``SentenceTransformer.encode`` as a class-level bound method (source branch).""" + + def __init__(self, offset=0): + self.offset = offset + self.call_count = 0 + self.last_kwargs = None + + def encode(self, texts, **kwargs): + self.call_count += 1 + self.last_kwargs = dict(kwargs) + return [[float(self.offset + i)] for i in range(len(texts))] + + +class FakeCallableModel: + """A callable instance (defines ``__call__``); triggers the fallback branch.""" + + def __init__(self, offset=0): + self.offset = offset + self.call_count = 0 + self.last_kwargs = None + + def __call__(self, texts, **kwargs): + self.call_count += 1 + self.last_kwargs = dict(kwargs) + return [[float(self.offset + i)] for i in range(len(texts))] + + +def test_transform_value_branches_behave_as_reported(): + """Confirm the two buggy branches (source + fallback) collapse distinct callables.""" + from dspy.clients.cache import _transform_value + + a, b = FakeSTSource(0), FakeSTSource(1) + # Source branch: identical source across same-class bound methods. + assert inspect.getsource(a.encode) == inspect.getsource(b.encode) + assert _transform_value(a.encode) == _transform_value(b.encode) + + ia, ib = FakeCallableModel(0), FakeCallableModel(1) + # Fallback branch: instances lack __name__ -> both collapse to "". + assert _transform_value(ia) == _transform_value(ib) == "" + + +def test_model_id_disambiguates_source_branch(cache): + """Documented pattern ``Embedder(model.encode)`` + ``model_id`` no longer collides.""" + inst_a = FakeSTSource(offset=0) + inst_b = FakeSTSource(offset=1000) + emb_a = Embedder(inst_a.encode, model_id="ckpt-a") + emb_b = Embedder(inst_b.encode, model_id="ckpt-b") + inputs = ["x", "y"] + + r_a = emb_a(inputs) + r_b = emb_b(inputs) + + assert inst_a.call_count == 1 + assert inst_b.call_count == 1 # computed its own, did not reuse emb_a's cache + assert not np.array_equal(r_a, r_b) + np.testing.assert_allclose(r_a, [[0.0], [1.0]]) + np.testing.assert_allclose(r_b, [[1000.0], [1001.0]]) + + +def test_model_id_disambiguates_fallback_branch_callable_instance(cache): + """Distinct callable instances (fallback branch) no longer collide with ``model_id``.""" + inst_a = FakeCallableModel(offset=0) + inst_b = FakeCallableModel(offset=1000) + emb_a = Embedder(inst_a, model_id="instance-a") + emb_b = Embedder(inst_b, model_id="instance-b") + inputs = ["x", "y"] + + r_a = emb_a(inputs) + r_b = emb_b(inputs) + + assert inst_a.call_count == 1 + assert inst_b.call_count == 1 + assert not np.array_equal(r_a, r_b) + np.testing.assert_allclose(r_a, [[0.0], [1.0]]) + np.testing.assert_allclose(r_b, [[1000.0], [1001.0]]) + + +def test_same_model_id_reuses_cache_across_instances(cache, tmp_path): + """A stable ``model_id`` still reuses the cache for the same model across instances/sessions.""" + dspy.configure_cache(disk_cache_dir=tmp_path / ".dspy_cache_reuse") + + inst_a = FakeSTSource(offset=0) + inst_b = FakeSTSource(offset=0) # same vectors -> models intended to be "the same" + emb_a = Embedder(inst_a.encode, model_id="same-ckpt") + emb_b = Embedder(inst_b.encode, model_id="same-ckpt") + inputs = ["x"] + + r_a = emb_a(inputs) + assert inst_a.call_count == 1 + + # Clear memory cache to simulate a fresh process reusing the on-disk cache. + dspy.cache.reset_memory_cache() + + r_b = emb_b(inputs) + # Same model_id + same input -> on-disk cache hit; inst_b never invoked. + assert inst_b.call_count == 0 + np.testing.assert_allclose(r_b, r_a) + + +def test_model_id_prevents_cross_cache_collision(cache, tmp_path): + """Simulate the cross-process scenario from the report (E3): distinct ``model_id`` values + keep two same-class bound methods from sharing cached vectors through the on-disk cache.""" + dspy.configure_cache(disk_cache_dir=tmp_path / ".dspy_cache_xproc") + + inst_a = FakeSTSource(offset=0) + emb_a = Embedder(inst_a.encode, model_id="checkpoint-A") + inputs = ["hello", "world"] + + r_a = emb_a(inputs) + assert inst_a.call_count == 1 + + # Fresh process: clear memory cache, leaving only the shared on-disk cache. + dspy.cache.reset_memory_cache() + + inst_b = FakeSTSource(offset=1000) + emb_b = Embedder(inst_b.encode, model_id="checkpoint-B") + r_b = emb_b(inputs) + + # inst_b computes its own vectors; it does NOT read inst_a's on-disk cache entry. + assert inst_b.call_count == 1 + assert not np.array_equal(r_b, r_a) + np.testing.assert_allclose(r_b, [[1000.0], [1001.0]]) + + +def test_model_id_none_preserves_existing_cache_keys(cache): + """``model_id=None`` (explicit or unset) must not add a key dimension, so existing cache + entries are still reused (no silent global cache invalidation).""" + inst = FakeSTSource(offset=0) + emb_no_id = Embedder(inst.encode) # unset + emb_none = Embedder(inst.encode, model_id=None) # explicit None + inputs = ["hello"] + + r1 = emb_no_id(inputs) + assert inst.call_count == 1 + + r2 = emb_none(inputs) + # Same key (no model_id field) -> cache hit, callable not invoked again. + assert inst.call_count == 1 + np.testing.assert_allclose(r2, r1) + + +def test_model_id_not_forwarded_to_callable(cache): + """``model_id`` folds into the cache key but is never passed to the callable model.""" + inst = FakeSTSource(offset=0) + emb = Embedder(inst.encode, model_id="my-ckpt", caching=True) + emb(["hello"]) + + assert inst.call_count == 1 + assert inst.last_kwargs is not None + assert "model_id" not in inst.last_kwargs + assert "caching" not in inst.last_kwargs + + +def test_model_id_not_forwarded_to_litellm(cache): + """``model_id`` is not forwarded to litellm for hosted models.""" + model = "text-embedding-ada-002" + inputs = ["hello"] + with patch("litellm.embedding") as mock_litellm: + mock_litellm.return_value = MockEmbeddingResponse([[0.1, 0.2, 0.3]]) + embedding = Embedder(model, model_id="my-ckpt", caching=True) + result = embedding(inputs) + + mock_litellm.assert_called_once() + assert "model_id" not in mock_litellm.call_args.kwargs + np.testing.assert_allclose(result, [[0.1, 0.2, 0.3]]) + + +def test_per_call_model_id_overrides_instance_model_id(cache): + """A per-call ``model_id`` overrides the instance-level one for the cache key.""" + inst_a = FakeSTSource(offset=0) + inst_b = FakeSTSource(offset=1000) + emb = Embedder(inst_a.encode, model_id="ckpt") + + # First call uses the instance model_id "ckpt". + emb(["x"]) + assert inst_a.call_count == 1 + + # Second call with the same callable but a different per-call model_id must recompute + # (distinct cache key), and may bind to a different underlying model. + emb.model = inst_b.encode + r2 = emb(["x"], model_id="ckpt-b") + assert inst_b.call_count == 1 + np.testing.assert_allclose(r2, [[1000.0]]) + + +@pytest.mark.asyncio +async def test_model_id_disambiguates_async(cache): + """The async path also disambiguates distinct callable models via ``model_id``.""" + inst_a = FakeSTSource(offset=0) + inst_b = FakeSTSource(offset=1000) + emb_a = Embedder(inst_a.encode, model_id="ckpt-a") + emb_b = Embedder(inst_b.encode, model_id="ckpt-b") + inputs = ["x", "y"] + + r_a = await emb_a.acall(inputs) + r_b = await emb_b.acall(inputs) + + assert inst_a.call_count == 1 + assert inst_b.call_count == 1 + assert not np.array_equal(r_a, r_b) + np.testing.assert_allclose(r_a, [[0.0], [1.0]]) + np.testing.assert_allclose(r_b, [[1000.0], [1001.0]])