alma-memory 0.2.0__py3-none-any.whl

alma/retrieval/embeddings.py

@@ -0,0 +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:
+                raise ImportError(
+                    "sentence-transformers is required for local embeddings. "
+                    "Install with: pip install sentence-transformers"
+                )
+
+    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:
+                raise ImportError(
+                    "openai is required for Azure embeddings. "
+                    "Install with: pip install openai"
+                )
+        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/retrieval/engine.py

@@ -0,0 +1,287 @@
+"""
+ALMA Retrieval Engine.
+
+Handles semantic search and memory retrieval with scoring and caching.
+"""
+
+import time
+import logging
+from typing import Optional, List, Dict, Any
+
+from alma.types import MemorySlice, MemoryScope
+from alma.storage.base import StorageBackend
+from alma.retrieval.scoring import MemoryScorer, ScoringWeights, ScoredItem
+from alma.retrieval.cache import RetrievalCache, NullCache
+
+logger = logging.getLogger(__name__)
+
+
+class RetrievalEngine:
+    """
+    Retrieves relevant memories for task context injection.
+
+    Features:
+    - Semantic search via embeddings
+    - Recency weighting (newer memories preferred)
+    - Success rate weighting (proven strategies ranked higher)
+    - Caching for repeated queries
+    - Configurable scoring weights
+    """
+
+    def __init__(
+        self,
+        storage: StorageBackend,
+        embedding_provider: str = "local",
+        cache_ttl_seconds: int = 300,
+        enable_cache: bool = True,
+        max_cache_entries: int = 1000,
+        scoring_weights: Optional[ScoringWeights] = None,
+        recency_half_life_days: float = 30.0,
+        min_score_threshold: float = 0.2,
+    ):
+        """
+        Initialize retrieval engine.
+
+        Args:
+            storage: Storage backend to query
+            embedding_provider: "local" (sentence-transformers) or "azure" (Azure OpenAI)
+            cache_ttl_seconds: How long to cache query results
+            enable_cache: Whether to enable caching
+            max_cache_entries: Maximum cache entries before eviction
+            scoring_weights: Custom weights for similarity/recency/success/confidence
+            recency_half_life_days: Days after which recency score halves
+            min_score_threshold: Minimum score to include in results
+        """
+        self.storage = storage
+        self.embedding_provider = embedding_provider
+        self.min_score_threshold = min_score_threshold
+        self._embedder = None
+
+        # Initialize scorer
+        self.scorer = MemoryScorer(
+            weights=scoring_weights or ScoringWeights(),
+            recency_half_life_days=recency_half_life_days,
+        )
+
+        # Initialize cache
+        if enable_cache:
+            self.cache = RetrievalCache(
+                ttl_seconds=cache_ttl_seconds,
+                max_entries=max_cache_entries,
+            )
+        else:
+            self.cache = NullCache()
+
+    def retrieve(
+        self,
+        query: str,
+        agent: str,
+        project_id: str,
+        user_id: Optional[str] = None,
+        top_k: int = 5,
+        scope: Optional[MemoryScope] = None,
+        bypass_cache: bool = False,
+    ) -> MemorySlice:
+        """
+        Retrieve relevant memories for a task.
+
+        Args:
+            query: Task description to find relevant memories for
+            agent: Agent requesting memories
+            project_id: Project context
+            user_id: Optional user for preference retrieval
+            top_k: Max items per memory type
+            scope: Agent's learning scope for filtering
+            bypass_cache: Skip cache lookup/storage
+
+        Returns:
+            MemorySlice with relevant memories, scored and ranked
+        """
+        start_time = time.time()
+
+        # Check cache first
+        if not bypass_cache:
+            cached = self.cache.get(query, agent, project_id, user_id, top_k)
+            if cached is not None:
+                cached.retrieval_time_ms = int((time.time() - start_time) * 1000)
+                logger.debug(f"Cache hit for query: {query[:50]}...")
+                return cached
+
+        # Generate embedding for query
+        query_embedding = self._get_embedding(query)
+
+        # Retrieve raw items from storage (with vector search)
+        raw_heuristics = self.storage.get_heuristics(
+            project_id=project_id,
+            agent=agent,
+            embedding=query_embedding,
+            top_k=top_k * 2,  # Get extra for scoring/filtering
+            min_confidence=0.0,  # Let scorer handle filtering
+        )
+
+        raw_outcomes = self.storage.get_outcomes(
+            project_id=project_id,
+            agent=agent,
+            embedding=query_embedding,
+            top_k=top_k * 2,
+            success_only=False,
+        )
+
+        raw_domain_knowledge = self.storage.get_domain_knowledge(
+            project_id=project_id,
+            agent=agent,
+            embedding=query_embedding,
+            top_k=top_k * 2,
+        )
+
+        raw_anti_patterns = self.storage.get_anti_patterns(
+            project_id=project_id,
+            agent=agent,
+            embedding=query_embedding,
+            top_k=top_k * 2,
+        )
+
+        # Score and rank each type
+        scored_heuristics = self.scorer.score_heuristics(raw_heuristics)
+        scored_outcomes = self.scorer.score_outcomes(raw_outcomes)
+        scored_knowledge = self.scorer.score_domain_knowledge(raw_domain_knowledge)
+        scored_anti_patterns = self.scorer.score_anti_patterns(raw_anti_patterns)
+
+        # Apply threshold and limit
+        final_heuristics = self._extract_top_k(scored_heuristics, top_k)
+        final_outcomes = self._extract_top_k(scored_outcomes, top_k)
+        final_knowledge = self._extract_top_k(scored_knowledge, top_k)
+        final_anti_patterns = self._extract_top_k(scored_anti_patterns, top_k)
+
+        # Get user preferences (not scored, just retrieved)
+        preferences = []
+        if user_id:
+            preferences = self.storage.get_user_preferences(user_id=user_id)
+
+        retrieval_time_ms = int((time.time() - start_time) * 1000)
+
+        result = MemorySlice(
+            heuristics=final_heuristics,
+            outcomes=final_outcomes,
+            preferences=preferences,
+            domain_knowledge=final_knowledge,
+            anti_patterns=final_anti_patterns,
+            query=query,
+            agent=agent,
+            retrieval_time_ms=retrieval_time_ms,
+        )
+
+        # Cache result
+        if not bypass_cache:
+            self.cache.set(query, agent, project_id, result, user_id, top_k)
+
+        logger.info(
+            f"Retrieved {result.total_items} memories for '{query[:50]}...' "
+            f"in {retrieval_time_ms}ms"
+        )
+
+        return result
+
+    def _extract_top_k(
+        self,
+        scored_items: List[ScoredItem],
+        top_k: int,
+    ) -> List[Any]:
+        """
+        Extract top-k items after filtering by score threshold.
+
+        Args:
+            scored_items: Scored and sorted items
+            top_k: Maximum number to return
+
+        Returns:
+            List of original items (unwrapped from ScoredItem)
+        """
+        filtered = self.scorer.apply_score_threshold(
+            scored_items, self.min_score_threshold
+        )
+        return [item.item for item in filtered[:top_k]]
+
+    def _get_embedding(self, text: str) -> List[float]:
+        """
+        Generate embedding for text.
+
+        Uses lazy initialization of embedding model.
+        """
+        if self._embedder is None:
+            self._embedder = self._init_embedder()
+
+        return self._embedder.encode(text)
+
+    def _init_embedder(self):
+        """Initialize the embedding model based on provider config."""
+        if self.embedding_provider == "azure":
+            from alma.retrieval.embeddings import AzureEmbedder
+            return AzureEmbedder()
+        elif self.embedding_provider == "mock":
+            from alma.retrieval.embeddings import MockEmbedder
+            return MockEmbedder()
+        else:
+            from alma.retrieval.embeddings import LocalEmbedder
+            return LocalEmbedder()
+
+    def invalidate_cache(
+        self,
+        agent: Optional[str] = None,
+        project_id: Optional[str] = None,
+    ):
+        """
+        Invalidate cache entries.
+
+        Should be called after memory updates to ensure fresh results.
+
+        Args:
+            agent: Invalidate entries for this agent
+            project_id: Invalidate entries for this project
+        """
+        self.cache.invalidate(agent=agent, project_id=project_id)
+
+    def get_cache_stats(self) -> Dict[str, Any]:
+        """Get cache performance statistics."""
+        stats = self.cache.get_stats()
+        return stats.to_dict()
+
+    def clear_cache(self):
+        """Clear all cached results."""
+        self.cache.clear()
+
+    def get_scorer_weights(self) -> Dict[str, float]:
+        """Get current scoring weights."""
+        w = self.scorer.weights
+        return {
+            "similarity": w.similarity,
+            "recency": w.recency,
+            "success_rate": w.success_rate,
+            "confidence": w.confidence,
+        }
+
+    def update_scorer_weights(
+        self,
+        similarity: Optional[float] = None,
+        recency: Optional[float] = None,
+        success_rate: Optional[float] = None,
+        confidence: Optional[float] = None,
+    ):
+        """
+        Update scoring weights (will be normalized to sum to 1.0).
+
+        Args:
+            similarity: Weight for semantic similarity
+            recency: Weight for recency
+            success_rate: Weight for success rate
+            confidence: Weight for stored confidence
+        """
+        current = self.scorer.weights
+        self.scorer.weights = ScoringWeights(
+            similarity=similarity if similarity is not None else current.similarity,
+            recency=recency if recency is not None else current.recency,
+            success_rate=success_rate if success_rate is not None else current.success_rate,
+            confidence=confidence if confidence is not None else current.confidence,
+        )
+        # Clear cache since scoring changed
+        self.cache.clear()