diff --git a/python/python/lancedb/embeddings/ollama.py b/python/python/lancedb/embeddings/ollama.py index 0f4a9cdc1..871a04247 100644 --- a/python/python/lancedb/embeddings/ollama.py +++ b/python/python/lancedb/embeddings/ollama.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright The LanceDB Authors from functools import cached_property -from typing import TYPE_CHECKING, List, Optional, Sequence, Union +from typing import TYPE_CHECKING, Any, List, Optional, Sequence, Union import numpy as np @@ -56,6 +56,16 @@ class OllamaEmbeddings(TextEmbeddingFunction): embeddings = self._compute_embedding(texts) return list(embeddings) + def __getstate__(self) -> dict[str, Any]: + state = super().__getstate__() + state["__dict__"] = { + k: v for k, v in state["__dict__"].items() if k != "_ollama_client" + } + return state + + def __setstate__(self, state: dict[str, Any]) -> None: + super().__setstate__(state) + @cached_property def _ollama_client(self) -> "ollama.Client": ollama = attempt_import_or_raise("ollama") diff --git a/python/python/tests/test_embeddings.py b/python/python/tests/test_embeddings.py index c0cdf1c0d..5efb7d98a 100644 --- a/python/python/tests/test_embeddings.py +++ b/python/python/tests/test_embeddings.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright The LanceDB Authors import os +import pickle from typing import List, Optional, Union from unittest.mock import MagicMock, patch @@ -591,6 +592,22 @@ def test_openai_no_retry_on_401(mock_sleep): assert mock_sleep.call_count == 0 +def test_ollama_embeddings_pickle(): + """OllamaEmbeddings must pickle even after the cached client is created.""" + registry = get_registry() + model = registry.get("ollama").create(name="nomic-embed-text") + + # Simulate accessing the cached client, which stores it on the instance. + model.__dict__["_ollama_client"] = MagicMock() + + pickled = pickle.dumps(model) + restored = pickle.loads(pickled) + + assert restored.name == "nomic-embed-text" + assert restored.host == "http://localhost:11434" + assert "_ollama_client" not in restored.__dict__ + + def test_url_retrieve_downloads_image(): """ Embedding functions like open-clip, siglip, and jinaai use url_retrieve()