From 1b7e0bc380c0f5d0d866a577f356fca046bd7977 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Boschi?= Date: Mon, 3 Nov 2025 20:57:10 +0100 Subject: [PATCH] emb abstraction --- memora/embeddings.py | 101 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 101 insertions(+) create mode 100644 memora/embeddings.py diff --git a/memora/embeddings.py b/memora/embeddings.py new file mode 100644 index 00000000..607c0386 --- /dev/null +++ b/memora/embeddings.py @@ -0,0 +1,101 @@ +""" +Embeddings abstraction for the memory system. + +Provides an interface for generating embeddings with different backends. + +IMPORTANT: All embeddings must produce 384-dimensional vectors to match +the database schema (pgvector column defined as vector(384)). +""" +from abc import ABC, abstractmethod +from typing import List +import logging + +logger = logging.getLogger(__name__) + +# Fixed embedding dimension required by database schema +EMBEDDING_DIMENSION = 384 + + +class Embeddings(ABC): + """ + Abstract base class for embedding generation. + + All implementations MUST generate 384-dimensional embeddings to match + the database schema. + """ + + @abstractmethod + def encode(self, texts: List[str]) -> List[List[float]]: + """ + Generate 384-dimensional embeddings for a list of texts. + + Args: + texts: List of text strings to encode + + Returns: + List of 384-dimensional embedding vectors (each is a list of floats) + """ + pass + + +class SentenceTransformersEmbeddings(Embeddings): + """ + Embeddings implementation using SentenceTransformers. + + Uses lazy import so sentence-transformers is not required if another + embedding backend is used. + + Default model is BAAI/bge-small-en-v1.5 which produces 384-dimensional + embeddings matching the database schema. + """ + + def __init__(self, model_name: str = "BAAI/bge-small-en-v1.5"): + """ + Initialize SentenceTransformers embeddings. + + Args: + model_name: Name of the SentenceTransformer model to use. + Must produce 384-dimensional embeddings. + Default: BAAI/bge-small-en-v1.5 + """ + self.model_name = model_name + self._model = None + + def _load_model(self): + """Lazy load and validate the SentenceTransformer model.""" + if self._model is None: + try: + from sentence_transformers import SentenceTransformer + except ImportError: + raise ImportError( + "sentence-transformers is required for SentenceTransformersEmbeddings. " + "Install it with: pip install sentence-transformers" + ) + + logger.info(f"Loading embedding model: {self.model_name}...") + self._model = SentenceTransformer(self.model_name) + + # Validate dimension matches database schema + model_dim = self._model.get_sentence_embedding_dimension() + if model_dim != EMBEDDING_DIMENSION: + raise ValueError( + f"Model {self.model_name} produces {model_dim}-dimensional embeddings, " + f"but database schema requires {EMBEDDING_DIMENSION} dimensions. " + f"Use a model that produces {EMBEDDING_DIMENSION}-dimensional embeddings." + ) + + logger.info(f"Model loaded (embedding dim: {model_dim})") + + def encode(self, texts: List[str]) -> List[List[float]]: + """ + Generate 384-dimensional embeddings for a list of texts. + + Args: + texts: List of text strings to encode + + Returns: + List of 384-dimensional embedding vectors + """ + self._load_model() + embeddings = self._model.encode(texts, convert_to_numpy=True, show_progress_bar=False) + return [emb.tolist() for emb in embeddings]