/cognee/infrastructure/databases/vector/embeddings/OllamaEmbeddingEngine.py

import asyncio
import math
from cognee.shared.logging_utils import get_logger
import aiohttp
from typing import List, Optional
import os
import litellm
import logging
import aiohttp.http_exceptions
import numpy as np
from tenacity import (
    retry,
    stop_after_delay,
    wait_exponential_jitter,
    retry_if_not_exception_type,
    before_sleep_log,
)

from cognee.infrastructure.databases.vector.embeddings.EmbeddingEngine import EmbeddingEngine
from cognee.infrastructure.databases.exceptions import EmbeddingException
from cognee.infrastructure.llm.tokenizer.resolver import resolve_embedding_tokenizer
from cognee.shared.rate_limiting import embedding_rate_limiter_context_manager
from cognee.shared.utils import create_secure_ssl_context
from cognee.infrastructure.databases.vector.embeddings.utils import (
    sanitize_embedding_text_inputs,
    handle_embedding_response,
)

logger = get_logger("OllamaEmbeddingEngine")


class OllamaEmbeddingEngine(EmbeddingEngine):
    """
    Implements an embedding engine using the Ollama embedding model.

    Public methods:
    - embed_text
    - get_vector_size
    - get_tokenizer

    Instance variables:
    - model
    - dimensions
    - max_completion_tokens
    - endpoint
    - mock
    - huggingface_tokenizer_name
    - tokenizer
    """

    model: str
    dimensions: int
    max_completion_tokens: int
    endpoint: str
    mock: bool
    huggingface_tokenizer_name: str

    MAX_RETRIES = 5

    def __init__(
        self,
        model: Optional[str] = "avr/sfr-embedding-mistral:latest",
        dimensions: Optional[int] = 1024,
        max_completion_tokens: int = 512,
        endpoint: Optional[str] = "http://localhost:11434/api/embed",
        huggingface_tokenizer: str = "Salesforce/SFR-Embedding-Mistral",
        batch_size: int = 100,
    ):
        self.model = model
        self.dimensions = dimensions
        self.max_completion_tokens = max_completion_tokens
        self.endpoint = endpoint
        self.huggingface_tokenizer_name = huggingface_tokenizer
        self.batch_size = batch_size
        self.tokenizer = self.get_tokenizer()

        enable_mocking = os.getenv("MOCK_EMBEDDING", "false")
        if isinstance(enable_mocking, bool):
            enable_mocking = str(enable_mocking).lower()
        self.mock = enable_mocking in ("true", "1", "yes")

    async def embed_text(self, text: List[str]) -> List[List[float]]:
        """
        Generate embedding vectors for a list of text prompts.

        If mocking is enabled, returns a list of zero vectors instead of actual embeddings.

        Parameters:
        -----------

            - text (List[str]): A list of text prompts for which to generate embeddings.

        Returns:
        --------

            - List[List[float]]: A list of embedding vectors corresponding to the text prompts.
        """
        original_texts = text if isinstance(text, list) else [text]
        sanitized_text = sanitize_embedding_text_inputs(original_texts)

        if self.mock:
            embeddings = [[0.0] * self.dimensions for _ in sanitized_text]
            return handle_embedding_response(original_texts, embeddings, self.dimensions)

        try:
            embeddings = await asyncio.gather(
                *[self._get_embedding(prompt) for prompt in sanitized_text]
            )
            return handle_embedding_response(original_texts, embeddings, self.dimensions)
        except Exception as error:
            error_str = str(error).lower()
            context_error_patterns = (
                "context length",
                "context window",
                "input length",
                "too long",
                "maximum context",
                "maximum tokens",
                "max tokens",
            )
            if any(pattern in error_str for pattern in context_error_patterns):
                if len(original_texts) > 1:
                    mid = math.ceil(len(original_texts) / 2)
                    left_vecs, right_vecs = await asyncio.gather(
                        self.embed_text(original_texts[:mid]),
                        self.embed_text(original_texts[mid:]),
                    )
                    embeddings = left_vecs + right_vecs
                    return handle_embedding_response(original_texts, embeddings, self.dimensions)

                if len(original_texts) == 1:
                    s = original_texts[0]
                    third = len(s) // 3
                    if third == 0:
                        raise EmbeddingException(
                            "Text is too short to split further but exceeds context window."
                        ) from error
                    left_part, right_part = s[: third * 2], s[third:]
                    (left_vec,), (right_vec,) = await asyncio.gather(
                        self.embed_text([left_part]),
                        self.embed_text([right_part]),
                    )
                    pooled = (np.array(left_vec) + np.array(right_vec)) / 2
                    embeddings = [pooled.tolist()]
                    return handle_embedding_response(original_texts, embeddings, self.dimensions)

                return handle_embedding_response(original_texts, embeddings, self.dimensions)

            logger.error(f"Embedding error in OllamaEmbeddingEngine: {str(error)}")
            raise EmbeddingException(
                f"Failed to index data points using model {self.model}"
            ) from error

    @retry(
        stop=stop_after_delay(128),
        wait=wait_exponential_jitter(8, 128),
        retry=retry_if_not_exception_type(
            (litellm.exceptions.NotFoundError, ValueError, asyncio.CancelledError)
        ),
        before_sleep=before_sleep_log(logger, logging.WARNING),
        reraise=True,
    )
    async def _get_embedding(self, prompt: str) -> List[float]:
        """
        Internal method to call the Ollama embeddings endpoint for a single prompt.
        """

        payload = {
            "model": self.model,
            "input": prompt,
            "dimensions": self.dimensions,
        }

        headers = {}
        api_key = os.getenv("LLM_API_KEY")
        if api_key:
            headers["Authorization"] = f"Bearer {api_key}"

        ssl_context = create_secure_ssl_context()
        connector = aiohttp.TCPConnector(ssl=ssl_context)
        async with aiohttp.ClientSession(connector=connector) as session:
            async with embedding_rate_limiter_context_manager():
                async with session.post(
                    self.endpoint, json=payload, headers=headers, timeout=60.0
                ) as response:
                    data = await response.json()

                    if "error" in data:
                        error_msg = data["error"]
                        logger.error(f"Ollama embedding error: {error_msg}")
                        if "context length" in error_msg or "input length" in error_msg:
                            raise ValueError(f"Text too long for embedding model: {error_msg}")
                        raise RuntimeError(f"Ollama embedding API error: {error_msg}")

                    if "embeddings" in data:
                        return data["embeddings"][0]
                    elif "embedding" in data:
                        return data["embedding"]
                    elif "data" in data and len(data["data"]) > 0:
                        return data["data"][0]["embedding"]
                    else:
                        raise ValueError(f"Unexpected response format from Ollama: {data}")

    def get_vector_size(self) -> int:
        """
        Retrieve the size of the embedding vectors.

        Returns:
        --------

            - int: The dimension of the embedding vectors.
        """
        return self.dimensions

    def get_batch_size(self) -> int:
        """
        Return the desired batch size for embedding calls

        Returns:

        """
        return self.batch_size

    def get_tokenizer(self):
        """
        Load and return the tokenizer for the embedding engine.

        An Ollama model id is not a HuggingFace repo, so the configured
        HUGGINGFACE_TOKENIZER override selects the tokenizer; resolution warns on
        mismatch and falls back safely to TikToken (issue #3646).

        Returns:
        --------

            The tokenizer used by the embedding engine.
        """
        logger.debug("Loading tokenizer for OllamaEmbeddingEngine...")
        tokenizer = resolve_embedding_tokenizer(
            provider="ollama",
            model=self.model,
            max_completion_tokens=self.max_completion_tokens,
            huggingface_tokenizer=self.huggingface_tokenizer_name,
        )
        logger.debug("Tokenizer loaded for OllamaEmbeddingEngine")
        return tokenizer