alma-memory 0.5.1__py3-none-any.whl → 0.7.0__py3-none-any.whl

alma/retrieval/embeddings.py

@@ -1,202 +1,202 @@
-"""
-ALMA Embedding Providers.
-
-Supports local (sentence-transformers) and Azure OpenAI embeddings.
-"""
-
-import logging
-from abc import ABC, abstractmethod
-from typing import List, Optional
-
-logger = logging.getLogger(__name__)
-
-
-class EmbeddingProvider(ABC):
-    """Abstract base class for embedding providers."""
-
-    @abstractmethod
-    def encode(self, text: str) -> List[float]:
-        """Generate embedding for text."""
-        pass
-
-    @abstractmethod
-    def encode_batch(self, texts: List[str]) -> List[List[float]]:
-        """Generate embeddings for multiple texts."""
-        pass
-
-    @property
-    @abstractmethod
-    def dimension(self) -> int:
-        """Return embedding dimension."""
-        pass
-
-
-class LocalEmbedder(EmbeddingProvider):
-    """
-    Local embeddings using sentence-transformers.
-
-    Default model: all-MiniLM-L6-v2 (384 dimensions, fast, good quality)
-    """
-
-    def __init__(self, model_name: str = "all-MiniLM-L6-v2"):
-        """
-        Initialize local embedder.
-
-        Args:
-            model_name: Sentence-transformers model name
-        """
-        self.model_name = model_name
-        self._model = None
-        self._dimension: Optional[int] = None
-
-    def _load_model(self):
-        """Lazy load the model."""
-        if self._model is None:
-            try:
-                from sentence_transformers import SentenceTransformer
-
-                logger.info(f"Loading embedding model: {self.model_name}")
-                self._model = SentenceTransformer(self.model_name)
-                self._dimension = self._model.get_sentence_embedding_dimension()
-                logger.info(f"Model loaded, dimension: {self._dimension}")
-            except ImportError as err:
-                raise ImportError(
-                    "sentence-transformers is required for local embeddings. "
-                    "Install with: pip install sentence-transformers"
-                ) from err
-
-    def encode(self, text: str) -> List[float]:
-        """Generate embedding for text."""
-        self._load_model()
-        embedding = self._model.encode(text, normalize_embeddings=True)
-        return embedding.tolist()
-
-    def encode_batch(self, texts: List[str]) -> List[List[float]]:
-        """Generate embeddings for multiple texts."""
-        self._load_model()
-        embeddings = self._model.encode(texts, normalize_embeddings=True)
-        return [emb.tolist() for emb in embeddings]
-
-    @property
-    def dimension(self) -> int:
-        """Return embedding dimension."""
-        if self._dimension is None:
-            self._load_model()
-        return self._dimension or 384  # Default for all-MiniLM-L6-v2
-
-
-class AzureEmbedder(EmbeddingProvider):
-    """
-    Azure OpenAI embeddings.
-
-    Uses text-embedding-3-small by default (1536 dimensions).
-    """
-
-    def __init__(
-        self,
-        endpoint: Optional[str] = None,
-        api_key: Optional[str] = None,
-        deployment: str = "text-embedding-3-small",
-        api_version: str = "2024-02-01",
-    ):
-        """
-        Initialize Azure OpenAI embedder.
-
-        Args:
-            endpoint: Azure OpenAI endpoint (or use AZURE_OPENAI_ENDPOINT env var)
-            api_key: Azure OpenAI API key (or use AZURE_OPENAI_KEY env var)
-            deployment: Deployment name for embedding model
-            api_version: API version
-        """
-        import os
-
-        self.endpoint = endpoint or os.environ.get("AZURE_OPENAI_ENDPOINT")
-        self.api_key = api_key or os.environ.get("AZURE_OPENAI_KEY")
-        self.deployment = deployment
-        self.api_version = api_version
-        self._client = None
-        self._dimension = 1536  # Default for text-embedding-3-small
-
-        if not self.endpoint:
-            raise ValueError(
-                "Azure OpenAI endpoint required. Set AZURE_OPENAI_ENDPOINT env var "
-                "or pass endpoint parameter."
-            )
-
-    def _get_client(self):
-        """Get or create Azure OpenAI client."""
-        if self._client is None:
-            try:
-                from openai import AzureOpenAI
-
-                self._client = AzureOpenAI(
-                    azure_endpoint=self.endpoint,
-                    api_key=self.api_key,
-                    api_version=self.api_version,
-                )
-            except ImportError as err:
-                raise ImportError(
-                    "openai is required for Azure embeddings. "
-                    "Install with: pip install openai"
-                ) from err
-        return self._client
-
-    def encode(self, text: str) -> List[float]:
-        """Generate embedding for text."""
-        client = self._get_client()
-        response = client.embeddings.create(
-            input=text,
-            model=self.deployment,
-        )
-        return response.data[0].embedding
-
-    def encode_batch(self, texts: List[str]) -> List[List[float]]:
-        """Generate embeddings for multiple texts."""
-        client = self._get_client()
-        response = client.embeddings.create(
-            input=texts,
-            model=self.deployment,
-        )
-        # Sort by index to ensure order matches input
-        sorted_data = sorted(response.data, key=lambda x: x.index)
-        return [item.embedding for item in sorted_data]
-
-    @property
-    def dimension(self) -> int:
-        """Return embedding dimension."""
-        return self._dimension
-
-
-class MockEmbedder(EmbeddingProvider):
-    """
-    Mock embedder for testing.
-
-    Generates deterministic fake embeddings based on text hash.
-    """
-
-    def __init__(self, dimension: int = 384):
-        """Initialize mock embedder."""
-        self._dimension = dimension
-
-    def encode(self, text: str) -> List[float]:
-        """Generate fake embedding based on text hash."""
-        import hashlib
-
-        # Create deterministic embedding from text hash
-        hash_bytes = hashlib.sha256(text.encode()).digest()
-        # Use first N bytes to create float values
-        embedding = []
-        for i in range(self._dimension):
-            byte_val = hash_bytes[i % len(hash_bytes)]
-            # Normalize to [-1, 1] range
-            embedding.append((byte_val / 127.5) - 1.0)
-        return embedding
-
-    def encode_batch(self, texts: List[str]) -> List[List[float]]:
-        """Generate fake embeddings for multiple texts."""
-        return [self.encode(text) for text in texts]
-
-    @property
-    def dimension(self) -> int:
-        """Return embedding dimension."""
-        return self._dimension
+"""
+ALMA Embedding Providers.
+
+Supports local (sentence-transformers) and Azure OpenAI embeddings.
+"""
+
+import logging
+from abc import ABC, abstractmethod
+from typing import List, Optional
+
+logger = logging.getLogger(__name__)
+
+
+class EmbeddingProvider(ABC):
+    """Abstract base class for embedding providers."""
+
+    @abstractmethod
+    def encode(self, text: str) -> List[float]:
+        """Generate embedding for text."""
+        pass
+
+    @abstractmethod
+    def encode_batch(self, texts: List[str]) -> List[List[float]]:
+        """Generate embeddings for multiple texts."""
+        pass
+
+    @property
+    @abstractmethod
+    def dimension(self) -> int:
+        """Return embedding dimension."""
+        pass
+
+
+class LocalEmbedder(EmbeddingProvider):
+    """
+    Local embeddings using sentence-transformers.
+
+    Default model: all-MiniLM-L6-v2 (384 dimensions, fast, good quality)
+    """
+
+    def __init__(self, model_name: str = "all-MiniLM-L6-v2"):
+        """
+        Initialize local embedder.
+
+        Args:
+            model_name: Sentence-transformers model name
+        """
+        self.model_name = model_name
+        self._model = None
+        self._dimension: Optional[int] = None
+
+    def _load_model(self):
+        """Lazy load the model."""
+        if self._model is None:
+            try:
+                from sentence_transformers import SentenceTransformer
+
+                logger.info(f"Loading embedding model: {self.model_name}")
+                self._model = SentenceTransformer(self.model_name)
+                self._dimension = self._model.get_sentence_embedding_dimension()
+                logger.info(f"Model loaded, dimension: {self._dimension}")
+            except ImportError as err:
+                raise ImportError(
+                    "sentence-transformers is required for local embeddings. "
+                    "Install with: pip install sentence-transformers"
+                ) from err
+
+    def encode(self, text: str) -> List[float]:
+        """Generate embedding for text."""
+        self._load_model()
+        embedding = self._model.encode(text, normalize_embeddings=True)
+        return embedding.tolist()
+
+    def encode_batch(self, texts: List[str]) -> List[List[float]]:
+        """Generate embeddings for multiple texts."""
+        self._load_model()
+        embeddings = self._model.encode(texts, normalize_embeddings=True)
+        return [emb.tolist() for emb in embeddings]
+
+    @property
+    def dimension(self) -> int:
+        """Return embedding dimension."""
+        if self._dimension is None:
+            self._load_model()
+        return self._dimension or 384  # Default for all-MiniLM-L6-v2
+
+
+class AzureEmbedder(EmbeddingProvider):
+    """
+    Azure OpenAI embeddings.
+
+    Uses text-embedding-3-small by default (1536 dimensions).
+    """
+
+    def __init__(
+        self,
+        endpoint: Optional[str] = None,
+        api_key: Optional[str] = None,
+        deployment: str = "text-embedding-3-small",
+        api_version: str = "2024-02-01",
+    ):
+        """
+        Initialize Azure OpenAI embedder.
+
+        Args:
+            endpoint: Azure OpenAI endpoint (or use AZURE_OPENAI_ENDPOINT env var)
+            api_key: Azure OpenAI API key (or use AZURE_OPENAI_KEY env var)
+            deployment: Deployment name for embedding model
+            api_version: API version
+        """
+        import os
+
+        self.endpoint = endpoint or os.environ.get("AZURE_OPENAI_ENDPOINT")
+        self.api_key = api_key or os.environ.get("AZURE_OPENAI_KEY")
+        self.deployment = deployment
+        self.api_version = api_version
+        self._client = None
+        self._dimension = 1536  # Default for text-embedding-3-small
+
+        if not self.endpoint:
+            raise ValueError(
+                "Azure OpenAI endpoint required. Set AZURE_OPENAI_ENDPOINT env var "
+                "or pass endpoint parameter."
+            )
+
+    def _get_client(self):
+        """Get or create Azure OpenAI client."""
+        if self._client is None:
+            try:
+                from openai import AzureOpenAI
+
+                self._client = AzureOpenAI(
+                    azure_endpoint=self.endpoint,
+                    api_key=self.api_key,
+                    api_version=self.api_version,
+                )
+            except ImportError as err:
+                raise ImportError(
+                    "openai is required for Azure embeddings. "
+                    "Install with: pip install openai"
+                ) from err
+        return self._client
+
+    def encode(self, text: str) -> List[float]:
+        """Generate embedding for text."""
+        client = self._get_client()
+        response = client.embeddings.create(
+            input=text,
+            model=self.deployment,
+        )
+        return response.data[0].embedding
+
+    def encode_batch(self, texts: List[str]) -> List[List[float]]:
+        """Generate embeddings for multiple texts."""
+        client = self._get_client()
+        response = client.embeddings.create(
+            input=texts,
+            model=self.deployment,
+        )
+        # Sort by index to ensure order matches input
+        sorted_data = sorted(response.data, key=lambda x: x.index)
+        return [item.embedding for item in sorted_data]
+
+    @property
+    def dimension(self) -> int:
+        """Return embedding dimension."""
+        return self._dimension
+
+
+class MockEmbedder(EmbeddingProvider):
+    """
+    Mock embedder for testing.
+
+    Generates deterministic fake embeddings based on text hash.
+    """
+
+    def __init__(self, dimension: int = 384):
+        """Initialize mock embedder."""
+        self._dimension = dimension
+
+    def encode(self, text: str) -> List[float]:
+        """Generate fake embedding based on text hash."""
+        import hashlib
+
+        # Create deterministic embedding from text hash
+        hash_bytes = hashlib.sha256(text.encode()).digest()
+        # Use first N bytes to create float values
+        embedding = []
+        for i in range(self._dimension):
+            byte_val = hash_bytes[i % len(hash_bytes)]
+            # Normalize to [-1, 1] range
+            embedding.append((byte_val / 127.5) - 1.0)
+        return embedding
+
+    def encode_batch(self, texts: List[str]) -> List[List[float]]:
+        """Generate fake embeddings for multiple texts."""
+        return [self.encode(text) for text in texts]
+
+    @property
+    def dimension(self) -> int:
+        """Return embedding dimension."""
+        return self._dimension