Adding a New Embedding Model

Relevant source files

The following files were used as context for generating this wiki page:

• monitor/embeddings/init.py
• monitor/embeddings/ollama.py
• monitor/config.py
• monitor/main.py
• test_basic.py
• WARP.md

The technical-blog-monitor features a pluggable architecture for embedding generation, allowing developers to easily integrate new embedding providers. The system is designed with a clear separation of concerns, using an abstract base class and a factory function to decouple the main application logic from the specific implementation details of any given embedding client. This allows for swapping between providers like OpenAI, HuggingFace, or local models via Ollama by simply changing configuration settings.

This guide outlines the steps required to add a new embedding model to the system. The process involves updating the configuration, implementing a client class that adheres to the BaseEmbeddingClient interface, and registering the new client in the system's factory function.

Sources: WARP.md, monitor/embeddings/init.py

Architecture Overview

The embedding subsystem is built around three key components: the EmbeddingConfig model, the BaseEmbeddingClient abstract class, and the get_embedding_client factory function.

1. EmbeddingConfig: A Pydantic model in monitor/config.py that defines all configuration options related to embeddings, such as model type, name, and API keys.
2. BaseEmbeddingClient: An abstract base class in monitor/embeddings/__init__.py that defines the contract for all embedding clients. It ensures that every client provides a consistent API for generating text and image embeddings.
3. get_embedding_client: A factory function in monitor/embeddings/__init__.py that reads the EmbeddingConfig and returns an initialized instance of the corresponding client.

Sources: monitor/config.py:100-132, monitor/embeddings/init.py:22-77

Class Diagram

This diagram illustrates the relationship between the abstract base class and a concrete implementation, using OllamaEmbeddingClient as an example. A new client must inherit from BaseEmbeddingClient.

Step 1: Update Configuration

The first step is to make the application aware of the new model type by updating the configuration files.

1.1. Extend EmbeddingModelType Enum

Add your new model's identifier to the EmbeddingModelType enum in monitor/config.py. This string will be used in environment variables to select your new client.

# monitor/config.py

class EmbeddingModelType(str, Enum):
    """Types of embedding models supported."""
    OPENAI = "openai"
    HUGGINGFACE = "huggingface"
    SENTENCE_TRANSFORMERS = "sentence_transformers"
    OLLAMA = "ollama"
    CUSTOM = "custom"
    NEW_MODEL = "new_model" # Add your new model type here

Sources: monitor/config.py:93-101

1.2. Add to EmbeddingConfig

If your new model requires specific configuration parameters, such as an API key or a service URL, add them to the EmbeddingConfig Pydantic model. Use SecretStr for sensitive values.

Field	Type	Description
text_model_type	EmbeddingModelType	The type of text embedding model to use.
text_model_name	str	The specific name of the model (e.g., text-embedding-ada-002).
image_model_type	Optional[EmbeddingModelType]	The type of image embedding model to use.
image_model_name	Optional[str]	The specific name of the image model.
openai_api_key	Optional[SecretStr]	API key for OpenAI services.
huggingface_api_key	Optional[SecretStr]	API key for HuggingFace services.
embedding_dimensions	Optional[int]	The dimension of the text embedding vectors.
batch_size	int	The number of items to process in a single batch.

You may also need to add a model_validator to ensure that the required credentials for your new model type are provided at runtime.

Sources: monitor/config.py:104-132

Step 2: Implement the BaseEmbeddingClient

Create a new Python file in the monitor/embeddings/ directory for your client (e.g., new_model_client.py). In this file, define a class that inherits from BaseEmbeddingClient and implements its abstract methods.

Key Methods to Implement

• async def _embed_text_batch(self, texts: List[str]) -> List[List[float]]: This is the core method for processing a batch of texts and returning their vector representations.
• async def _embed_image_batch(self, image_paths: List[str]) -> List[List[float]]: Implement this if your model supports image embeddings. If not, you can return dummy vectors or raise a NotImplementedError, as seen in the OllamaEmbeddingClient.
• async def get_dimension(self) -> int: Returns the dimensionality of the embedding vectors produced by the model. This can be a fixed value or determined dynamically.
• async def close(self) -> None: Implement this to gracefully release any resources, such as closing network connections.

The OllamaEmbeddingClient serves as a good reference. It uses an httpx.AsyncClient for making API calls and includes retry logic using the tenacity library for robustness.

# monitor/embeddings/ollama.py

class OllamaEmbeddingClient(BaseEmbeddingClient):
    # ... (initialization)

    async def _embed_text_batch(self, texts: List[str]) -> List[List[float]]:
        # ...
        tasks = [self._embed_one(text) for text in texts]
        results = await asyncio.gather(*tasks)
        # ...
        return results

    @retry(...)
    async def _embed_one(self, text: str) -> List[float]:
        payload = {
            "model": self.model,
            "prompt": text,
            "keep_alive": self.keep_alive,
        }
        url = f"{self.base_url}/api/embeddings"
        resp = await self._client.post(url, json=payload)
        # ...
        return embedding

    async def _embed_image_batch(self, image_paths: List[str]) -> List[List[float]]:
        logger.warning(
            "Ollama text models don't support image embeddings, returning dummy vectors",
            count=len(image_paths),
        )
        from monitor.embeddings import DummyEmbeddingClient
        dummy_client = DummyEmbeddingClient(self.config)
        return await dummy_client._embed_image_batch(image_paths)

    async def get_dimension(self) -> int:
        if self._dim is None:
            probe = await self._embed_one("dimension probe")
            self._dim = len(probe)
        return self._dim

    async def close(self) -> None:
        await self._client.aclose()
        await super().close()

Sources: monitor/embeddings/ollama.py:22-192, monitor/embeddings/init.py:22-77

Step 3: Register in the Factory Function

After implementing the client, you must register it in the get_embedding_client factory function located in monitor/embeddings/__init__.py.

1. Import your new client class.
2. Add an elif block to the function to handle the EmbeddingModelType you defined in Step 1.
3. Instantiate and return your new client.

# monitor/embeddings/__init__.py

async def get_embedding_client(config: EmbeddingConfig) -> "BaseEmbeddingClient":
    # ...
    if config.text_model_type == EmbeddingModelType.OLLAMA:
        from monitor.embeddings.ollama import OllamaEmbeddingClient
        return OllamaEmbeddingClient(config, model=config.text_model_name)
    
    # Add your new client here
    elif config.text_model_type == EmbeddingModelType.NEW_MODEL:
        from .new_model_client import NewModelClient
        return NewModelClient(config)
    # ...
    else:
        raise NotImplementedError(f"Unsupported embedding model type: {config.text_model_type}")

Sources: monitor/embeddings/init.py:80-120

Factory Logic Flow

The following diagram shows how the get_embedding_client function selects and initializes the correct client based on the configuration.

Step 4: System Integration and Usage

Once registered, the new embedding client is automatically used by the application wherever embeddings are needed. The main application logic in monitor/main.py calls get_embedding_client during startup and uses the returned instance throughout its lifecycle.

Initialization and Usage Sequence

This diagram shows the sequence of events from application startup to generating an embedding for a blog post.

For testing, the DummyEmbeddingClient is often used to avoid dependencies on external services. You can follow the pattern in test_basic.py to write tests for your new client, either by mocking its external dependencies or by running it against a test instance of its service.

Sources: test_basic.py:103-107

Conclusion

The embedding generation system in technical-blog-monitor is designed for extensibility. By following the steps of updating the configuration, implementing the BaseEmbeddingClient interface, and registering the new class in the factory function, developers can seamlessly integrate new embedding models with minimal changes to the core application logic. This modular design ensures the system can adapt to new technologies and services in the machine learning ecosystem.