diff --git a/datastew/embedding/base.py b/datastew/embedding/base.py index 648045e..7f791b8 100644 --- a/datastew/embedding/base.py +++ b/datastew/embedding/base.py @@ -12,6 +12,8 @@ _GLOBAL_LOCKS = {} _INIT_LOCK = Lock() +_WHITESPACE_RE = re.compile(r"\s+") + class EmbeddingModel(ABC): def __init__(self, model_name: str, cache: bool = False, cache_size: int = 10000): @@ -37,25 +39,93 @@ def __init__(self, model_name: str, cache: bool = False, cache_size: int = 10000 self._cache = None self._cache_lock = None - @abstractmethod def get_embedding(self, text: str) -> Sequence[float]: """Retrieve the embedding vector for a single text input. :param text: The input text to embed. + :raises TypeError: If the input is not a string. + :raises ValueError: If the input text is empty or only whitespace. :return: A sequence of floats representing the embedding. """ - pass + if not isinstance(text, str): + raise TypeError(f"Expected string for 'text', got {type(text).__name__}.") + + if not text.strip(): + raise ValueError("Input 'text' cannot be empty or solely whitespace.") + + text = self._sanitize(text) + + if self._cache is not None: + cached = self._get_from_cache(text) + if cached: + return cached + + embedding = self._generate_embedding(text) + self._add_to_cache(text, embedding) + return embedding - @abstractmethod def get_embeddings(self, messages: List[str]) -> Sequence[Sequence[float]]: """Retrieve embeddings for a list of text messages :param messages: A list of text messages to embed. + :raises TypeError: If 'messages' is not a list. + :raises ValueError: If 'messages' is empty. :return: A sequence of embedding vectors. """ + if not isinstance(messages, list): + raise TypeError(f"Expected a list for 'messages', got {type(messages).__name__}") + + if not messages: + raise ValueError("The 'messages' list cannot be empty.") + + sanitized_messages = [self._sanitize(msg) for msg in messages if msg and isinstance(msg, str)] + + if self._cache is not None: + embeddings, uncached_indices, uncached_messages = self._get_cached_embeddings(sanitized_messages) + + if uncached_messages: + new_embeddings = self._generate_embeddings(uncached_messages) + self._add_batch_to_cache(uncached_messages, new_embeddings) + for idx, embedding in zip(uncached_indices, new_embeddings): + embeddings[idx] = embedding + + return [emb for emb in embeddings if emb is not None] + + return self._generate_embeddings(sanitized_messages) + + @abstractmethod + def _generate_embedding(self, text: str) -> Sequence[float]: + """Generate an embedding vector for a single sanitized text input + + This is an abstract method that must be implemented by subclasses to interface + with the specific model's API or inference engine. + + :param text: The sanitized input text to embed. + :return: A sequence of floats respresenting the generated embedding. + """ + pass + + @abstractmethod + def _generate_embeddings(self, messages: list[str]) -> Sequence[Sequence[float]]: + """Generate embedding vectors for a batch of sanitized text messages. + + This is an abstract method that must be implemented by subclasses to interface + with the specific model's API or inference engine. + + :param messages: A list of sanitized text messages to embed. + :return: A sequence of embedding vectors corresponding to the input messages. + """ pass - def add_to_cache(self, text: str, embedding: Sequence[float]): + def _sanitize(self, message: str) -> str: + """Clean up the input text by trimming and converting to lowercase. + + :param message: The input text message. + :return: Sanitized text. + """ + return _WHITESPACE_RE.sub(" ", message).strip().lower() + + def _add_to_cache(self, text: str, embedding: Sequence[float]): """Add a text-embedding pair to cache. :param text: The input text to be cached. @@ -65,14 +135,20 @@ def add_to_cache(self, text: str, embedding: Sequence[float]): with self._cache_lock: self._cache[text] = embedding - def add_batch_to_cache(self, texts: Sequence[str], embeddings: Sequence[Sequence[float]]): - """Acquire lock once for the entire batch update.""" + def _add_batch_to_cache(self, texts: Sequence[str], embeddings: Sequence[Sequence[float]]): + """Add a batch of text-embedding pairs to the cache. + + Acquires the thread lock once for the entire batch update to minimize overhead. + + :param texts: A sequence of input texts to be cached. + :param embeddings: A sequence of corresponding embedding vectors. + """ if self._cache_lock is not None and self._cache is not None: with self._cache_lock: for text, emb in zip(texts, embeddings): self._cache[text] = emb - def get_from_cache(self, text: str) -> Optional[Sequence[float]]: + def _get_from_cache(self, text: str) -> Optional[Sequence[float]]: """Retrieve an embedding from the cache. :param text: Cached input text. @@ -83,7 +159,7 @@ def get_from_cache(self, text: str) -> Optional[Sequence[float]]: return self._cache.get(text, None) return None - def get_cached_embeddings( + def _get_cached_embeddings( self, messages: List[str] ) -> Tuple[List[Optional[Sequence[float]]], List[int], List[str]]: """Retrieve cached embeddings and identify uncached messages. @@ -95,8 +171,7 @@ def get_cached_embeddings( - A list of uncached messages. """ if self._cache_lock is None or self._cache is None: - empty_embeddings: List[Optional[Sequence[float]]] = [None] * len(messages) - return empty_embeddings, list(range(len(messages))), messages + return [None] * len(messages), list(range(len(messages))), messages embeddings: List[Optional[Sequence[float]]] = [] uncached_indices: List[int] = [] @@ -111,11 +186,3 @@ def get_cached_embeddings( uncached_messages.append(msg) return embeddings, uncached_indices, uncached_messages - - def sanitize(self, message: str) -> str: - """Clean up the input text by trimming and converting to lowercase. - - :param message: The input text message. - :return: Sanitized text. - """ - return re.sub(r"\s+", " ", message).strip().lower() diff --git a/datastew/embedding/hugging_face.py b/datastew/embedding/hugging_face.py index 0383fc1..f62e9c4 100644 --- a/datastew/embedding/hugging_face.py +++ b/datastew/embedding/hugging_face.py @@ -1,68 +1,33 @@ import logging -from typing import List, Sequence +from typing import Sequence +from cachetools import LRUCache from sentence_transformers import SentenceTransformer from datastew.embedding.base import EmbeddingModel class HuggingFaceAdapter(EmbeddingModel): - _model_cache = {} - _load_count = 0 + _model_cache = LRUCache(maxsize=3) def __init__(self, model_name: str = "sentence-transformers/all-MiniLM-L6-v2", cache: bool = False): super().__init__(model_name, cache) if model_name not in self._model_cache: - HuggingFaceAdapter._load_count += 1 self._model_cache[model_name] = SentenceTransformer(model_name) self.model = self._model_cache[model_name] - def get_embedding(self, text: str) -> Sequence[float]: - if not text or not isinstance(text, str): - logging.warning("Empty or invalid text passed to get_embedding") - return [] - text = self.sanitize(text) - - if self._cache is not None: - cached = self.get_from_cache(text) - if cached: - return cached - + def _generate_embedding(self, text: str) -> Sequence[float]: try: - embedding = self.model.encode(text) - embedding = [float(x) for x in embedding] - self.add_to_cache(text, embedding) - return embedding + return self.model.encode(text).tolist() except Exception as e: logging.error(f"Error getting embedding for {text}: {e}") raise - def get_embeddings(self, messages: List[str]) -> Sequence[Sequence[float]]: - sanitized_messages = [self.sanitize(msg) for msg in messages] - if self._cache is not None: - embeddings, uncached_indices, uncached_messages = self.get_cached_embeddings(sanitized_messages) - - if uncached_messages: - try: - new_embeddings = self.model.encode(uncached_messages, show_progress_bar=True) - flattened_embeddings = [ - [float(element) for element in row] for row in new_embeddings if row is not None - ] - self.add_batch_to_cache(uncached_messages, flattened_embeddings) - for idx, embedding in zip(uncached_indices, flattened_embeddings): - embeddings[idx] = embedding - except Exception as e: - logging.error(f"Failed processing messages: {e}") - raise - - return [emb for emb in embeddings if emb is not None] - + def _generate_embeddings(self, messages: list[str]) -> Sequence[Sequence[float]]: try: - embeddings = self.model.encode(sanitized_messages, show_progress_bar=True) - flattened_embeddings = [[float(element) for element in row] for row in embeddings] - return flattened_embeddings + return self.model.encode(messages, show_progress_bar=True).tolist() except Exception as e: logging.error(f"Failed processing messages: {e}") raise diff --git a/datastew/embedding/ollama.py b/datastew/embedding/ollama.py index 12b49e5..ed0c0be 100644 --- a/datastew/embedding/ollama.py +++ b/datastew/embedding/ollama.py @@ -1,5 +1,5 @@ import logging -from typing import List, Sequence +from typing import Sequence from ollama import Client @@ -13,44 +13,23 @@ def __init__( super().__init__(model_name, cache) self.client = Client(host) - def get_embedding(self, text: str) -> Sequence[float]: - if not text: - logging.warning("Empty text passed to get_embedding") - return [] - text = self.sanitize(text) - - if self._cache is not None: - cached = self.get_from_cache(text) - if cached: - return cached + def _generate_embedding(self, text: str) -> Sequence[float]: try: - embedding = self.client.embed(self.model_name, text).get("embeddings")[0] - self.add_to_cache(text, embedding) - return embedding + return self.client.embed(self.model_name, text).get("embeddings")[0] except Exception as e: logging.error(f"Error getting embedding for {text}: {e}") raise - def get_embeddings(self, messages: List[str]) -> Sequence[Sequence[float]]: - sanitized_messages = [self.sanitize(msg) for msg in messages] - - if self._cache is not None: - embeddings, uncached_indices, uncached_messages = self.get_cached_embeddings(sanitized_messages) - - if uncached_messages: - try: - new_embeddings = self.client.embed(self.model_name, uncached_messages).get("embeddings") - self.add_batch_to_cache(uncached_messages, new_embeddings) - for idx, embedding in zip(uncached_indices, new_embeddings): - embeddings[idx] = embedding - except Exception as e: - logging.error(f"Failed processing messages: {e}") - raise - - return [emb for emb in embeddings if emb is not None] + def _generate_embeddings(self, messages: list[str]) -> Sequence[Sequence[float]]: + chunk_size = 500 + all_embeddings = [] try: - return self.client.embed(self.model_name, sanitized_messages).get("embeddings") + for i in range(0, len(messages), chunk_size): + chunk = messages[i : i + chunk_size] + response = self.client.embed(self.model_name, chunk) + all_embeddings.extend(response.get("embeddings", [])) + return all_embeddings except Exception as e: logging.error(f"Failed processing messages: {e}") raise diff --git a/datastew/embedding/openai.py b/datastew/embedding/openai.py index cce3380..a83e75a 100644 --- a/datastew/embedding/openai.py +++ b/datastew/embedding/openai.py @@ -1,18 +1,13 @@ import logging -from typing import List, Sequence +from typing import Sequence -import openai +from openai import OpenAI from datastew.embedding.base import EmbeddingModel class GPT4Adapter(EmbeddingModel): - def __init__( - self, - api_key: str, - model_name: str = "text-embedding-ada-002", - cache: bool = False, - ): + def __init__(self, api_key: str, model_name: str = "text-embedding-ada-002", cache: bool = False): """Initialize the GPT-4 adapter with OpenAI API key and model name. :param api_key: The API key for accessing OpenAI services. @@ -20,52 +15,26 @@ def __init__( :param cache: Enable or disable caching, defaults to False. """ super().__init__(model_name, cache) - self.api_key = api_key - openai.api_key = api_key - - def get_embedding(self, text: str) -> Sequence[float]: - if not text: - logging.warning("Empty text passed to get_embedding") - return [] - text = self.sanitize(text) - - if self._cache is not None: - cached = self.get_from_cache(text) - if cached: - return cached + self.client = OpenAI(api_key=api_key) + def _generate_embedding(self, text: str) -> Sequence[float]: try: - response = openai.embeddings.create(input=[text], model=self.model_name) - embedding = response.data[0].embedding - self.add_to_cache(text, embedding) - return embedding + response = self.client.embeddings.create(input=[text], model=self.model_name) + return response.data[0].embedding except Exception as e: logging.error(f"Error getting embedding for {text}: {e}") raise - def get_embeddings(self, messages: List[str]) -> Sequence[Sequence[float]]: - sanitized_messages = [self.sanitize(msg) for msg in messages] - - if self._cache is not None: - embeddings, uncached_indices, uncached_messages = self.get_cached_embeddings(sanitized_messages) - - if uncached_messages: - try: - response = openai.embeddings.create(model=self.model_name, input=uncached_messages) - new_embeddings = [item.embedding for item in response.data] - self.add_batch_to_cache(uncached_messages, new_embeddings) - for idx, embedding in zip(uncached_indices, new_embeddings): - embeddings[idx] = embedding - except Exception as e: - logging.error(f"Error in processing chunk: {e}") - raise - - return [emb for emb in embeddings if emb is not None] + def _generate_embeddings(self, messages: list[str]) -> Sequence[Sequence[float]]: + chunk_size = 1000 + all_embeddings = [] try: - response = openai.embeddings.create(model=self.model_name, input=sanitized_messages) - embeddings = [item.embedding for item in response.data] - return embeddings + for i in range(0, len(messages), chunk_size): + chunk = messages[i : i + chunk_size] + response = self.client.embeddings.create(model=self.model_name, input=chunk) + all_embeddings.extend([item.embedding for item in response.data]) + return all_embeddings except Exception as e: logging.error(f"Failed processing messages: {e}") raise diff --git a/datastew/embedding/vectorizer.py b/datastew/embedding/vectorizer.py index 26a78d3..7348fc8 100644 --- a/datastew/embedding/vectorizer.py +++ b/datastew/embedding/vectorizer.py @@ -1,4 +1,4 @@ -from typing import List, Literal, Optional +from typing import Literal, Optional, Sequence SupportedModel = Literal[ "sentence-transformers/all-MiniLM-L6-v2", @@ -12,6 +12,21 @@ class Vectorizer: + """Factory class to initialize and route requests to the appropriate embedding model adapter. + + Acts as a unified interface for local Hugging Face models, OpenAI API models, and locally hosted Ollama models. + """ + + _MODEL_REGISTRY = { + "sentence-transformers/all-MiniLM-L6-v2": "hugging_face", + "sentence-transformers/all-mpnet-base-v2": "hugging_face", + "FremyCompany/BioLORD-2023": "hugging_face", + "text-embedding-ada-002": "openai", + "text-embedding-3-large": "openai", + "text-embedding-3-small": "openai", + "nomic-embed-text": "ollama", + } + def __init__( self, model: SupportedModel = "sentence-transformers/all-mpnet-base-v2", @@ -23,45 +38,57 @@ def __init__( :param model: The model to use for generating embeddings, defaults to sentence-transformers/all-mpnet-base-v2. :param api_key: The API key for GPT-based models, defaults to None. - :param host: The host URL for locally hosted Ollama models. defaults to http://localhost:11434. + :param host: The host URL for locally hosted Ollama models, defaults to http://localhost:11434. :param cache: Whether to enable caching for embeddings, defaults to False. """ - self.model = self.initialize_model(model, api_key, host, cache) + self.model = self._initialize_model(model, api_key, host, cache) self.model_name = self.model.model_name - def initialize_model(self, model: SupportedModel, api_key: Optional[str], host: str, cache: bool): - if model in [ - "sentence-transformers/all-MiniLM-L6-v2", - "sentence-transformers/all-mpnet-base-v2", - "FremyCompany/BioLORD-2023", - ]: + def get_embedding(self, text: str) -> Sequence[float]: + """Retrieve the embedding vector for a single text input. + + :param text: The input text to embed. + :return: A sequence of floats representing the embedding. + """ + return self.model.get_embedding(text) + + def get_embeddings(self, messages: list[str]) -> Sequence[Sequence[float]]: + """Retrieve embedding vectors for a batch of text messages. + + :param messages: A list of text messages to embed. + :return: A sequence of embedding vectors corresponding to the input messages. + """ + return self.model.get_embeddings(messages) + + def _initialize_model(self, model: SupportedModel, api_key: Optional[str], host: str, cache: bool): + """Resolve the provider from the registry and instantiate the corresponding adapter. + + :param model: The specific embedding model identifier. + :param api_key: The API key required for OpenAI models. + :param host: The host URL for Ollama models. + :param cache: Flag to enable LRU caching within the adapter. + :raises ValueError: If an OpenAI model is required but no API key is provided. + :raises NotImplementedError: If the requested model is not found in the registry. + :return: An initialized subclass of EmbeddingModel. + """ + provider = self._MODEL_REGISTRY.get(model) + + if provider == "hugging_face": from datastew.embedding.hugging_face import HuggingFaceAdapter return HuggingFaceAdapter(model, cache) - elif ( - model - in [ - "text-embedding-ada-002", - "text-embedding-3-large", - "text-embedding-3-small", - ] - and api_key - ): + elif provider == "openai": + if not api_key: + raise ValueError(f"API key is required for OpenAI model '{model}'.") from datastew.embedding.openai import GPT4Adapter return GPT4Adapter(api_key, model, cache) - elif model == "nomic-embed-text": + elif provider == "ollama": from datastew.embedding.ollama import OllamaAdapter return OllamaAdapter(model, host, cache) else: - raise NotImplementedError(f"The model '{model}' is not supported or missing API key.") - - def get_embedding(self, text: str): - return self.model.get_embedding(text) - - def get_embeddings(self, messages: List[str]): - return self.model.get_embeddings(messages) + raise NotImplementedError(f"The model '{model}' is not supported in the registry.") diff --git a/tests/embedding/__init__.py b/tests/embedding/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/embedding/test_base.py b/tests/embedding/test_base.py new file mode 100644 index 0000000..453e4d2 --- /dev/null +++ b/tests/embedding/test_base.py @@ -0,0 +1,76 @@ +import unittest +from typing import Sequence +from unittest.mock import patch + +from datastew.embedding.base import _GLOBAL_CACHES, _GLOBAL_LOCKS, EmbeddingModel + + +class DummyAdapter(EmbeddingModel): + """A minimal implementation of EmbeddingModel to strictly test base class logic.""" + + def _generate_embedding(self, text: str) -> Sequence[float]: + """Generate a dummy embedding vector based on text length.""" + val = float(len(text)) + return [val, val] + + def _generate_embeddings(self, messages: list[str]) -> Sequence[Sequence[float]]: + """Generate a batch of dummy embedding vectors based on text lengths.""" + return [[float(len(msg)), float(len(msg))] for msg in messages] + + +class TestEmbeddingBase(unittest.TestCase): + def setUp(self): + """Reset global caches and initialize a fresh DummyAdapter before each test.""" + _GLOBAL_CACHES.clear() + _GLOBAL_LOCKS.clear() + self.adapter = DummyAdapter(model_name="dummy-model", cache=True) + + def test_sanitization(self): + """Verify that text is correctly trimmed, lowercased, and whitespace-normalized.""" + text1 = " Test" + text2 = "test " + embedding1 = self.adapter.get_embedding(text1) + embedding2 = self.adapter.get_embedding(text2) + self.assertSequenceEqual(embedding1, embedding2) + self.assertEqual(self.adapter._sanitize(" Multiple Spaces "), "multiple spaces") + + def test_caching_get_embedding_deterministic(self): + """Ensure single embeddings are retrieved from cache on subsequent calls.""" + text = "deterministic test sentence" + + with patch.object(self.adapter, "_generate_embedding", wraps=self.adapter._generate_embedding) as spy: + emb1 = self.adapter.get_embedding(text) + spy.assert_called_once() + emb2 = self.adapter.get_embedding(text) + self.assertEqual(spy.call_count, 1) + self.assertSequenceEqual(emb1, emb2) + + def test_caching_get_embeddings_batch(self): + """Ensure batch embeddings are retrieved from cache on subsequent calls.""" + messages = [f"Batch message {i}." for i in range(5)] + + with patch.object(self.adapter, "_generate_embeddings", wraps=self.adapter._generate_embeddings) as spy: + embeddings1 = self.adapter.get_embeddings(messages) + self.assertEqual(spy.call_count, 1) + + embeddings2 = self.adapter.get_embeddings(messages) + self.assertEqual(spy.call_count, 1) + + for emb1, emb2 in zip(embeddings1, embeddings2): + self.assertSequenceEqual(emb1, emb2) + + def test_partial_cache_hit_batch(self): + """Verify that batch processing only generates embeddings for uncached items.""" + messages = ["A", "B", "C"] + self.adapter._add_batch_to_cache(["a"], [[0.1, 0.2]]) + + with patch.object(self.adapter, "_generate_embeddings", return_value=[[0.3, 0.4], [0.5, 0.6]]) as spy: + embeddings = self.adapter.get_embeddings(messages) + spy.assert_called_once_with(["b", "c"]) + self.assertEqual(len(embeddings), 3) + self.assertEqual(embeddings[0], [0.1, 0.2]) + self.assertEqual(embeddings[1], [0.3, 0.4]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/embedding/test_hugging_face.py b/tests/embedding/test_hugging_face.py new file mode 100644 index 0000000..1691385 --- /dev/null +++ b/tests/embedding/test_hugging_face.py @@ -0,0 +1,60 @@ +import unittest +from unittest.mock import patch + +import numpy as np + +from datastew.embedding.hugging_face import HuggingFaceAdapter + + +class TestHuggingFaceAdapter(unittest.TestCase): + + @patch("datastew.embedding.hugging_face.SentenceTransformer") + def setUp(self, mock_transformer): + """Reset the class-level model cache and initialize a mocked adapter.""" + HuggingFaceAdapter._model_cache.clear() + self.mock_model_instance = mock_transformer.return_value + self.adapter = HuggingFaceAdapter(cache=False) + + def test_generate_embedding(self): + """Verify single embedding generation and NumPy to list conversion.""" + self.mock_model_instance.encode.return_value = np.array([0.1, 0.2]) + result = self.adapter._generate_embedding("test") + self.assertEqual(result, [0.1, 0.2]) + self.mock_model_instance.encode.assert_called_once_with("test") + + def test_generate_embeddings(self): + """Verify batch embedding generation and proper conversion of nested NumPy arrays.""" + messages = ["msg1", "msg2"] + self.mock_model_instance.encode.return_value = np.array([[0.1], [0.2]]) + result = self.adapter._generate_embeddings(messages) + self.assertEqual(result, [[0.1], [0.2]]) + self.mock_model_instance.encode.assert_called_once_with(messages, show_progress_bar=True) + + def test_exception_handling_aborts_silence(self): + """Ensure that exceptions from the underlying model are logged and reraised.""" + messages = ["Fail 1", "Fail 2"] + self.mock_model_instance.encode.side_effect = Exception("Model Down") + with self.assertRaisesRegex(Exception, "Model Down"): + self.adapter._generate_embeddings(messages) + + @patch("datastew.embedding.hugging_face.SentenceTransformer") + def test_huggingface_model_loaded_once(self, mock_transformer): + """Verify LRUCache singleton behavior and maxsize enforcement.""" + HuggingFaceAdapter._model_cache.clear() + adapter1 = HuggingFaceAdapter(model_name="sentence-transformers/all-MiniLM-L6-v2") + adapter2 = HuggingFaceAdapter(model_name="sentence-transformers/all-MiniLM-L6-v2") + self.assertIs(adapter1.model, adapter2.model) + self.assertEqual(len(HuggingFaceAdapter._model_cache), 1) + self.assertEqual(mock_transformer.call_count, 1) + + # Force eviction + HuggingFaceAdapter(model_name="model_A") + HuggingFaceAdapter(model_name="model_B") + HuggingFaceAdapter(model_name="model_C") + self.assertEqual(len(HuggingFaceAdapter._model_cache), 3) + self.assertNotIn("sentence-transformers/all-MiniLM-L6-v2", HuggingFaceAdapter._model_cache) + self.assertEqual(mock_transformer.call_count, 4) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/embedding/test_ollama.py b/tests/embedding/test_ollama.py new file mode 100644 index 0000000..82ca160 --- /dev/null +++ b/tests/embedding/test_ollama.py @@ -0,0 +1,30 @@ +import unittest +from unittest.mock import patch + +from datastew.embedding.ollama import OllamaAdapter + + +class TestOllamaAdapter(unittest.TestCase): + @patch("datastew.embedding.ollama.Client") + def setUp(self, mock_client): + """Initialize the adapter with a mocked Ollama client.""" + self.mock_client_instance = mock_client.return_value + self.adapter = OllamaAdapter(cache=False) + + def test_generate_embedding(self): + """Verify single embedding generation via the Ollama client.""" + self.mock_client_instance.embed.return_value = {"embeddings": [[0.3, 0.4]]} + result = self.adapter._generate_embedding("test") + self.assertEqual(result, [0.3, 0.4]) + self.mock_client_instance.embed.assert_called_once_with("nomic-embed-text", "test") + + def test_generate_embeddings_chunking(self): + """Verify that large batches are split into smaller chunks.""" + messages = [f"msg{i}" for i in range(1200)] + self.mock_client_instance.embed.return_value = {"embeddings": [[0.5]]} + self.adapter._generate_embeddings(messages) + self.assertEqual(self.mock_client_instance.embed.call_count, 3) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/embedding/test_openai.py b/tests/embedding/test_openai.py new file mode 100644 index 0000000..f647ce4 --- /dev/null +++ b/tests/embedding/test_openai.py @@ -0,0 +1,40 @@ +import unittest +from unittest.mock import MagicMock, patch + +from datastew.embedding.openai import GPT4Adapter + + +class TestOpenAIAdapter(unittest.TestCase): + @patch("datastew.embedding.openai.OpenAI") + def setUp(self, mock_openai): + """Initialize the adapter with a mocked OpenAI client.""" + self.mock_client = mock_openai.return_value + self.adapter = GPT4Adapter(api_key="test-key", cache=False) + + def test_generate_embedding(self): + """Verify single embedding generation and response data extraction.""" + mock_response = MagicMock() + mock_response.data = [MagicMock(embedding=[0.1, 0.2])] + self.mock_client.embeddings.create.return_value = mock_response + result = self.adapter._generate_embedding("test") + self.assertEqual(result, [0.1, 0.2]) + self.mock_client.embeddings.create.assert_called_once_with(input=["test"], model="text-embedding-ada-002") + + def test_generate_embeddings_chunking(self): + """Verify batch request chunking for OpenAI API calls.""" + messages = [f"msg{i}" for i in range(2500)] + mock_response = MagicMock() + mock_response.data = [MagicMock(embedding=[0.5])] + self.mock_client.embeddings.create.return_value = mock_response + self.adapter._generate_embeddings(messages) + self.assertEqual(self.mock_client.embeddings.create.call_count, 3) + + def test_api_failure_propagation(self): + """Ensure that OpenAI API exceptions are correctly reraised.""" + self.mock_client.embeddings.create.side_effect = Exception("API Error") + with self.assertRaisesRegex(Exception, "API Error"): + self.adapter._generate_embedding("test") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/embedding/test_vectorizer.py b/tests/embedding/test_vectorizer.py new file mode 100644 index 0000000..aad3718 --- /dev/null +++ b/tests/embedding/test_vectorizer.py @@ -0,0 +1,37 @@ +import unittest + +from datastew.embedding.hugging_face import HuggingFaceAdapter +from datastew.embedding.ollama import OllamaAdapter +from datastew.embedding.openai import GPT4Adapter +from datastew.embedding.vectorizer import Vectorizer + + +class TestVectorizer(unittest.TestCase): + def test_initialize_hugging_face(self): + """Verify initialization of Hugging Face models.""" + vec = Vectorizer(model="sentence-transformers/all-MiniLM-L6-v2") + self.assertIsInstance(vec.model, HuggingFaceAdapter) + + def test_initialize_openai_valid(self): + """Verify initialization of OpenAI models with valid credentials.""" + vec = Vectorizer(model="text-embedding-ada-002", api_key="test_key") + self.assertIsInstance(vec.model, GPT4Adapter) + + def test_initialize_openai_missing_key(self): + """Ensure a ValueError is raised when OpenAI models lack an API key.""" + with self.assertRaisesRegex(ValueError, "API key is required"): + Vectorizer(model="text-embedding-ada-002") + + def test_initialize_ollama(self): + """Verify initialization of Ollama local API models.""" + vec = Vectorizer(model="nomic-embed-text") + self.assertIsInstance(vec.model, OllamaAdapter) + + def test_initialize_unsupported_model(self): + """Ensure NotImplementedError is raised for unregistered model strings.""" + with self.assertRaisesRegex(NotImplementedError, "not supported in the registry"): + Vectorizer(model="invalid-model-string") # type: ignore + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_embedding.py b/tests/test_embedding.py deleted file mode 100644 index 39fbba1..0000000 --- a/tests/test_embedding.py +++ /dev/null @@ -1,90 +0,0 @@ -import unittest -from typing import Sequence -from unittest.mock import patch - -from datastew.embedding.base import _GLOBAL_CACHES, _GLOBAL_LOCKS -from datastew.embedding.hugging_face import HuggingFaceAdapter - - -class TestEmbedding(unittest.TestCase): - - def setUp(self): - _GLOBAL_CACHES.clear() - _GLOBAL_LOCKS.clear() - self.adapter = HuggingFaceAdapter(cache=True) - - def test_hugging_face_adapter_get_embedding(self): - text = "This is a test sentence." - embedding = self.adapter.get_embedding(text) - self.assertIsInstance(embedding, Sequence) - - def test_hugging_face_adapter_get_embeddings(self): - messages = [f"This is message {i}." for i in range(20)] - embeddings = self.adapter.get_embeddings(messages) - self.assertIsInstance(embeddings, Sequence) - self.assertEqual(len(embeddings), len(messages)) - - def test_sanitization(self): - text1 = " Test" - text2 = "test " - embedding1 = self.adapter.get_embedding(text1) - embedding2 = self.adapter.get_embedding(text2) - self.assertSequenceEqual(embedding1, embedding2) - - def test_caching_get_embedding_deterministic(self): - text = "deterministic test sentence" - - with patch.object(self.adapter.model, "encode", wraps=self.adapter.model.encode) as spy_encode: - emb1 = self.adapter.get_embedding(text) - spy_encode.assert_called_once() - - emb2 = self.adapter.get_embedding(text) - self.assertEqual(spy_encode.call_count, 1) - self.assertSequenceEqual(emb1, emb2) - - def test_caching_get_embeddings_batch(self): - messages = [f"Batch message {i}." for i in range(5)] - - with patch.object(self.adapter.model, "encode", wraps=self.adapter.model.encode) as spy_encode: - embeddings1 = self.adapter.get_embeddings(messages) - self.assertEqual(spy_encode.call_count, 1) - embeddings2 = self.adapter.get_embeddings(messages) - self.assertEqual(spy_encode.call_count, 1) - - for emb1, emb2 in zip(embeddings1, embeddings2): - self.assertSequenceEqual(emb1, emb2) - - def test_partial_cache_hit_batch(self): - messages = ["A", "B", "C"] - # Cache "A" manually - self.adapter.add_batch_to_cache([self.adapter.sanitize("A")], [[0.1, 0.2]]) - - with patch.object(self.adapter.model, "encode", wraps=self.adapter.model.encode) as spy_encode: - embeddings = self.adapter.get_embeddings(messages) - # encode should only be called once, and only for ["b", "c"] - spy_encode.assert_called_once_with(["b", "c"], show_progress_bar=True) - self.assertEqual(len(embeddings), 3) - self.assertEqual(embeddings[0], [0.1, 0.2]) - - def test_exception_handling_aborts_silence(self): - messages = ["Fail 1", "Fail 2"] - - with patch.object(self.adapter.model, "encode", side_effect=Exception("API Down")): - with self.assertRaises(Exception) as context: - self.adapter.get_embeddings(messages) - - self.assertTrue("API Down" in str(context.exception)) - - def test_huggingface_model_loaded_once(self): - HuggingFaceAdapter._model_cache.clear() - HuggingFaceAdapter._load_count = 0 - - adapter1 = HuggingFaceAdapter(cache=True) - adapter2 = HuggingFaceAdapter(cache=True) - - self.assertIs(adapter1.model, adapter2.model) - self.assertEqual(HuggingFaceAdapter._load_count, 1) - - -if __name__ == "__main__": - unittest.main()