stonks-oracle/services/intelligence_pipeline_v3/novelty/embeddings.py

"""Replaceable compact embedding backend for novelty retrieval.

Provides:
- EmbeddingBackend protocol for pluggable embedding models
- MockEmbeddingBackend for deterministic testing
- SentenceTransformerBackend stub for production (all-MiniLM-L6-v2)
- cosine_similarity utility function
"""

from __future__ import annotations

import hashlib
import math
import struct
from typing import Protocol, runtime_checkable


@runtime_checkable
class EmbeddingBackend(Protocol):
    """Protocol for embedding backends.

    Implementations must produce fixed-dimension vectors for a batch of texts.
    The backend is designed to be replaceable: swap between mock, local model,
    and remote API backends without changing scoring logic.
    """

    @property
    def dimension(self) -> int:
        """Return the embedding dimension produced by this backend."""
        ...

    def embed(self, texts: list[str]) -> list[list[float]]:
        """Embed a batch of texts into dense vectors.

        Args:
            texts: List of text strings to embed.

        Returns:
            List of embedding vectors, one per input text.
            Each vector has length == self.dimension.
        """
        ...


class MockEmbeddingBackend:
    """Deterministic hash-based embedding backend for testing.

    Produces consistent embeddings using text hashing. Useful for
    unit tests and integration tests that need repeatable results
    without loading a real model.
    """

    def __init__(self, dimension: int = 384) -> None:
        self._dimension = dimension

    @property
    def dimension(self) -> int:
        return self._dimension

    def embed(self, texts: list[str]) -> list[list[float]]:
        """Generate deterministic embeddings from text hashes.

        Uses SHA-256 expanded to fill the dimension. Normalizes to unit length
        for compatibility with cosine similarity.
        """
        results = []
        for text in texts:
            raw = self._hash_to_vector(text)
            norm = math.sqrt(sum(x * x for x in raw))
            if norm > 0:
                normalized = [x / norm for x in raw]
            else:
                normalized = raw
            results.append(normalized)
        return results

    def _hash_to_vector(self, text: str) -> list[float]:
        """Expand text hash into a vector of the target dimension."""
        vector = []
        # Generate enough hash bytes to fill dimension
        chunk_idx = 0
        while len(vector) < self._dimension:
            data = f"{text}:{chunk_idx}".encode("utf-8")
            digest = hashlib.sha256(data).digest()
            # Convert 32 bytes to 8 floats (4 bytes each)
            for i in range(0, 32, 4):
                if len(vector) >= self._dimension:
                    break
                # Unpack as float in [-1, 1] range
                raw_int = struct.unpack("<I", digest[i : i + 4])[0]
                value = (raw_int / (2**32 - 1)) * 2.0 - 1.0
                vector.append(value)
            chunk_idx += 1
        return vector[: self._dimension]


class SentenceTransformerBackend:
    """Production embedding backend using sentence-transformers.

    Wraps all-MiniLM-L6-v2 (384-dimensional) for compact, fast embeddings.
    The model loads lazily on first call to avoid startup cost when not needed.

    Note: Requires `sentence-transformers` package to be installed.
    This is a stub—actual model loading is deferred to production deployment.
    """

    MODEL_NAME = "all-MiniLM-L6-v2"

    def __init__(self) -> None:
        self._model = None
        self._dimension = 384

    @property
    def dimension(self) -> int:
        return self._dimension

    def embed(self, texts: list[str]) -> list[list[float]]:
        """Embed texts using the sentence-transformers model.

        Lazily loads the model on first invocation.

        Raises:
            ImportError: If sentence-transformers is not installed.
        """
        if self._model is None:
            self._load_model()
        embeddings = self._model.encode(texts, normalize_embeddings=True)
        return [emb.tolist() for emb in embeddings]

    def _load_model(self) -> None:
        """Load the sentence-transformers model."""
        try:
            from sentence_transformers import SentenceTransformer
        except ImportError as e:
            raise ImportError(
                "sentence-transformers package is required for SentenceTransformerBackend. "
                "Install with: pip install sentence-transformers"
            ) from e
        self._model = SentenceTransformer(self.MODEL_NAME)


def cosine_similarity(a: list[float], b: list[float]) -> float:
    """Compute cosine similarity between two embedding vectors.

    Args:
        a: First embedding vector.
        b: Second embedding vector (must be same dimension as a).

    Returns:
        Cosine similarity in range [-1, 1]. Returns 0.0 for zero vectors.

    Raises:
        ValueError: If vectors have different dimensions.
    """
    if len(a) != len(b):
        raise ValueError(f"Vectors must have same dimension: {len(a)} != {len(b)}")

    dot = sum(x * y for x, y in zip(a, b))
    norm_a = math.sqrt(sum(x * x for x in a))
    norm_b = math.sqrt(sum(x * x for x in b))

    if norm_a == 0.0 or norm_b == 0.0:
        return 0.0

    return dot / (norm_a * norm_b)