adl-toolkit 0.0.0-stage → 0.1.1

package/adl_drift/cortex/__init__.py

@@ -0,0 +1 @@
+"""Memory layer - persistent pattern storage."""

package/adl_drift/cortex/embeddings.py

@@ -0,0 +1,96 @@
+"""Vector embedding generation for semantic search.
+
+Uses sentence-transformers to convert code patterns and knowledge into
+ dense vectors for similarity search.
+"""
+
+from __future__ import annotations
+
+import os
+
+# Loading model weights emits a tqdm progress bar ("Loading weights:  0%|...")
+# straight to stdout the first time a model is materialized. Fatal wherever
+# stdout is a strict channel — MCP JSON-RPC framing, or a `--json`/`--baseline`
+# CLI payload — since the bar interleaves with the real output and breaks
+# parsing. `transformers`/`huggingface_hub` only consult this env var once, at
+# import time, so it must be set before either is ever imported — hence
+# module-top, ahead of the sentence-transformers import below.
+os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1")
+
+from typing import List, Optional
+import numpy as np
+
+try:
+    from sentence_transformers import SentenceTransformer
+    HAS_SENTENCE_TRANSFORMERS = True
+except ImportError:
+    HAS_SENTENCE_TRANSFORMERS = False
+
+
+class EmbeddingEngine:
+    """Generates and compares embeddings for code patterns.
+
+    Uses 'all-MiniLM-L6-v2' by default (small, fast, good for code).
+    For better code understanding, use 'intfloat/e5-base-v2' or
+    a code-specific model like 'jinaai/jina-embeddings-v2-base-code'.
+    """
+
+    DEFAULT_MODEL = "all-MiniLM-L6-v2"
+
+    def __init__(self, model_name: Optional[str] = None):
+        self.model_name = model_name or self.DEFAULT_MODEL
+        self._model: Optional[SentenceTransformer] = None
+
+        if HAS_SENTENCE_TRANSFORMERS:
+            self._load_model()
+
+    def _load_model(self) -> None:
+        """Lazy-load the embedding model."""
+        if self._model is not None or not HAS_SENTENCE_TRANSFORMERS:
+            return
+        try:
+            # An already-cached model still pays a network round-trip here
+            # by default (Hugging Face Hub checking for a newer revision on
+            # every single load) — slow, and occasionally stalls hard enough
+            # to blow past a caller's timeout (e.g. a fresh MCP server
+            # subprocess per tool call). Try local-only first; only reach
+            # the network if the model truly isn't cached yet.
+            self._model = SentenceTransformer(self.model_name, local_files_only=True)
+        except Exception:
+            self._model = SentenceTransformer(self.model_name)
+
+    def embed(self, texts: List[str]) -> np.ndarray:
+        """Generate embeddings for a list of texts.
+
+        Returns:
+            numpy array of shape (len(texts), embedding_dim)
+        """
+        if not HAS_SENTENCE_TRANSFORMERS or self._model is None:
+            # Fallback: random embeddings for testing
+            return np.random.randn(len(texts), 384).astype(np.float32)
+
+        return self._model.encode(texts, convert_to_numpy=True, show_progress_bar=False)
+
+    def embed_single(self, text: str) -> np.ndarray:
+        """Generate embedding for a single text."""
+        return self.embed([text])[0]
+
+    def similarity(self, a: np.ndarray, b: np.ndarray) -> float:
+        """Compute cosine similarity between two embeddings."""
+        return float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b)))
+
+    def search(self, query: str, corpus: List[str], top_k: int = 5) -> List[tuple[int, float]]:
+        """Find most similar texts to query.
+
+        Returns:
+            List of (index, similarity_score) tuples, sorted by score.
+        """
+        query_emb = self.embed_single(query)
+        corpus_emb = self.embed(corpus)
+
+        similarities = np.dot(corpus_emb, query_emb) / (
+            np.linalg.norm(corpus_emb, axis=1) * np.linalg.norm(query_emb)
+        )
+
+        top_indices = np.argsort(similarities)[::-1][:top_k]
+        return [(int(i), float(similarities[i])) for i in top_indices]