synapse-memory-engine 1.3.0__py3-none-any.whl

synapse_memory/core/encryption.py

@@ -0,0 +1,226 @@
+import logging
+import os
+from typing import Any, List, Optional
+
+logger = logging.getLogger("SynapseEncryption")
+
+try:
+    from cryptography.fernet import Fernet
+    HAS_CRYPTOGRAPHY = True
+except ImportError:
+    HAS_CRYPTOGRAPHY = False
+    Fernet = None  # type: ignore[misc,assignment]
+
+
+class KeyManagementProvider:
+    """
+    Durable encryption key manager with team-safe loading order.
+
+    Resolution order (first hit wins):
+      1. `SYNAPSE_ENCRYPTION_KEY` — raw 32-byte urlsafe-base64 Fernet key.
+      2. `SYNAPSE_ENCRYPTION_KEY_B64_FILE` — path to a file containing ONLY the key.
+         (The correct way to inject a secret in Docker/K8s without baking it
+         into the image or leaving it in process args.)
+      3. KMS/Vault provider hook: if `SYNAPSE_KMS_PROVIDER` is set, the matching
+         loader below is invoked with `SYNAPSE_KMS_KEY_REF`. Providers are
+         best-effort optional — unknown or misconfigured providers raise, never
+         silently fall through, so a prod boot fails loudly instead of
+         encrypting with a throwaway key.
+      4. Local keyfile (dev fallback): `SYNAPSE_KEY_FILE` (default `.synapse_key`).
+
+    Production deployments MUST use (1), (2), or (3). The local keyfile (4) is a
+    convenience that must never back team/production data.
+    """
+    @staticmethod
+    def get_encryption_keys(key_file: Optional[str] = None) -> List[bytes]:
+        """Return the configured encryption keys, primary (encrypting) key first.
+
+        Resolution order (first configured source wins):
+          1. `SYNAPSE_ENCRYPTION_KEYS` — comma-separated list of urlsafe-base64
+             Fernet keys. The first key encrypts new writes; every listed key can
+             decrypt, so older keys can be retired without re-encrypting existing
+             ciphertext (MultiFernet rotation).
+          2. `SYNAPSE_ENCRYPTION_KEY` — single key (backward compatible).
+          3. `SYNAPSE_ENCRYPTION_KEY_B64_FILE` — path to a file containing ONE key.
+          4. KMS/Vault provider hook (`SYNAPSE_KMS_PROVIDER` + `SYNAPSE_KMS_KEY_REF`).
+          5. Local keyfile (dev fallback): `SYNAPSE_KEY_FILE` (default `.synapse_key`).
+          6. Auto-generate and persist one durable key (dev only).
+
+        Production deployments MUST use (1), (2), (3) or (4).
+        """
+        keys: List[bytes] = []
+
+        # 1. Direct env key(s) — comma separated; first listed key encrypts.
+        env_keys = os.environ.get("SYNAPSE_ENCRYPTION_KEYS")
+        if env_keys:
+            for k in env_keys.split(","):
+                k = k.strip()
+                if k:
+                    keys.append(k.encode("utf-8"))
+
+        # 2. Backward compatible single key.
+        if not keys:
+            env_key = os.environ.get("SYNAPSE_ENCRYPTION_KEY")
+            if env_key:
+                keys.append(env_key.strip().encode("utf-8"))
+
+        # 3. Key-from-file (Docker secrets / K8s secret volume mount).
+        if not keys:
+            b64_file = os.environ.get("SYNAPSE_ENCRYPTION_KEY_B64_FILE")
+            if b64_file:
+                try:
+                    with open(b64_file, "rb") as f:
+                        file_key = f.read().strip()
+                        if file_key:
+                            logger.info("Loaded encryption key from SYNAPSE_ENCRYPTION_KEY_B64_FILE.")
+                            keys.append(file_key)
+                except Exception as e:
+                    logger.warning(f"Could not read key file {b64_file}: {e}")
+
+        # 4. KMS/Vault provider hook.
+        if not keys:
+            kms_provider = os.environ.get("SYNAPSE_KMS_PROVIDER")
+            if kms_provider:
+                kms_ref = os.environ.get("SYNAPSE_KMS_KEY_REF", "")
+                loader = _KMS_LOADERS.get(kms_provider.lower())
+                if loader is None:
+                    raise RuntimeError(
+                        f"Unknown SYNAPSE_KMS_PROVIDER '{kms_provider}'. "
+                        f"Supported: {sorted(_KMS_LOADERS)}."
+                    )
+                logger.info(f"Loading encryption key via KMS provider '{kms_provider}'.")
+                keys.append(loader(kms_ref))
+
+        # 5. Local dev fallback keyfile.
+        if not keys:
+            target_file = key_file or os.environ.get("SYNAPSE_KEY_FILE", ".synapse_key")
+            if os.path.exists(target_file):
+                try:
+                    with open(target_file, "rb") as f:
+                        file_key = f.read().strip()
+                        if file_key:
+                            keys.append(file_key)
+                except Exception as e:
+                    logger.warning(f"Could not read key file {target_file}: {e}")
+
+        # 6. Generate a durable 32-byte urlsafe base64 key (dev only).
+        if not keys:
+            new_key = FernetEncryptionProvider.generate_key()
+            target_file = key_file or os.environ.get("SYNAPSE_KEY_FILE", ".synapse_key")
+            try:
+                with open(target_file, "wb") as f:
+                    f.write(new_key)
+                os.chmod(target_file, 0o600)
+            except Exception as e:
+                logger.warning(f"Could not persist encryption key to {target_file}: {e}")
+            keys.append(new_key)
+
+        return keys
+
+    @staticmethod
+    def get_encryption_key(key_file: Optional[str] = None) -> bytes:
+        return KeyManagementProvider.get_encryption_keys(key_file)[0]
+
+
+def _kms_loader_aws(ref: str) -> bytes:
+    """AWS KMS: decrypts a base64 ciphertext blob via kms:Decrypt (boto3 optional)."""
+    try:
+        import base64 as _b64
+        import boto3  # type: ignore
+    except ImportError as e:
+        raise RuntimeError("SYNAPSE_KMS_PROVIDER=aws requires boto3 to be installed.") from e
+    region = os.environ.get("SYNAPSE_KMS_REGION") or os.environ.get("AWS_REGION") or os.environ.get("AWS_DEFAULT_REGION")
+    client = boto3.client("kms", region_name=region) if region else boto3.client("kms")
+    resp = client.decrypt(CiphertextBlob=_b64.b64decode(ref))
+    return _b64.urlsafe_b64encode(resp["Plaintext"][:32].ljust(32, b"\0"))
+
+
+def _kms_loader_vault(ref: str) -> bytes:
+    """HashiCorp Vault KV: reads the key field over HTTPS (hvac optional)."""
+    try:
+        import hvac  # type: ignore
+    except ImportError as e:
+        raise RuntimeError("SYNAPSE_KMS_PROVIDER=vault requires hvac to be installed.") from e
+    addr = os.environ.get("VAULT_ADDR", "")
+    token = os.environ.get("VAULT_TOKEN", "")
+    field = os.environ.get("SYNAPSE_KMS_FIELD", "key")
+    if not addr or not token or not ref:
+        raise RuntimeError("vault KMS needs VAULT_ADDR, VAULT_TOKEN and SYNAPSE_KMS_KEY_REF (KV path).")
+    client = hvac.Client(url=addr, token=token)
+    secret = client.secrets.kv.v2.read_secret_version(path=ref)
+    value = (secret.get("data", {}).get("data", {}) or {}).get(field, "")
+    if not value:
+        raise RuntimeError(f"Vault KV path '{ref}' has no '{field}' field.")
+    return str(value).strip().encode("utf-8")
+
+
+def _kms_loader_file(ref: str) -> bytes:
+    """Alias of the key-file loader using an explicit path as the ref."""
+    with open(ref, "rb") as f:
+        return f.read().strip()
+
+
+_KMS_LOADERS = {
+    "aws": _kms_loader_aws,
+    "vault": _kms_loader_vault,
+    "file": _kms_loader_file,
+}
+
+
+# Aliases for backwards compatibility
+FernetKeyManager = KeyManagementProvider
+
+
+class FernetEncryptionProvider:
+    """
+    Encryption provider using Fernet symmetric authenticated encryption.
+    Underlying specification: 128-bit AES in CBC mode with PKCS7 padding
+    and HMAC-SHA256 authenticated encryption.
+    """
+    def __init__(self, key: Optional[bytes] = None):
+        if key is None:
+            self.key = KeyManagementProvider.get_encryption_key()
+        elif isinstance(key, str):
+            self.key = key.encode("utf-8")
+        else:
+            self.key = key
+
+        self._cipher: Any = None
+        if HAS_CRYPTOGRAPHY and Fernet is not None:
+            self._cipher = Fernet(self.key)  # type: ignore[operator]
+
+    @staticmethod
+    def generate_key() -> bytes:
+        if HAS_CRYPTOGRAPHY and Fernet is not None:
+            return Fernet.generate_key()
+        raise RuntimeError(
+            "Key generation requires the `cryptography` package (Fernet). "
+            "Install it with: pip install cryptography"
+        )
+
+    def encrypt(self, data: str) -> str:
+        """Encrypts a string and returns a base64 encoded string."""
+        if self._cipher is None:
+            raise RuntimeError(
+                "Encryption requires the `cryptography` package (Fernet). "
+                "Install it with: pip install cryptography"
+            )
+        return self._cipher.encrypt(data.encode("utf-8")).decode("utf-8")
+
+    def decrypt(self, encrypted_data: str) -> str:
+        """Decrypts a base64 encoded string and returns the original string."""
+        if self._cipher is None:
+            raise RuntimeError(
+                "Decryption requires the `cryptography` package (Fernet). "
+                "Install it with: pip install cryptography"
+            )
+        try:
+            return self._cipher.decrypt(
+                encrypted_data.encode("utf-8")
+            ).decode("utf-8")
+        except Exception as e:
+            raise ValueError("Decryption failed: invalid key or payload") from e
+
+
+# Aliases for backwards compatibility
+EncryptedMemoryStore = FernetEncryptionProvider

synapse_memory/core/hybrid_search.py

@@ -0,0 +1,199 @@
+import math
+import re
+from typing import List, Dict, Any, Optional
+
+
+def _tokenize(text: str) -> List[str]:
+    return [w for w in re.findall(r'[a-zA-Z0-9_]+', text.lower()) if len(w) > 1]
+
+
+class HybridSearch:
+    """
+    State-of-the-art Hybrid Information Retrieval Engine.
+    Combines Lexical Sparse Search (real BM25 algorithm) with Dense Cosine
+    vector similarity, fused via Reciprocal Rank Fusion (RRF).
+    """
+
+    def __init__(self, k1: float = 1.5, b: float = 0.75, rrf_k: int = 60, alpha: float = 0.5, k: Optional[int] = None, **kwargs: Any):
+        """
+        Args:
+            k1: BM25 term frequency saturation parameter.
+            b: BM25 document length normalization parameter.
+            rrf_k: Constant used to smooth Reciprocal Rank Fusion ranks.
+            alpha: Hybrid fusion weight between dense (alpha) and sparse (1 - alpha).
+            k: Alias for rrf_k or top_k fallback.
+        """
+        self.k1 = k1
+        self.b = b
+        self.rrf_k = k if k is not None else rrf_k
+        self.alpha = alpha
+
+    def _calculate_bm25_scores(self, query_terms: List[str], corpus: List[Dict[str, Any]]) -> Dict[str, float]:
+        """
+        Calculates exact Okapi BM25 scores for all documents in the corpus.
+        """
+        scores = {doc["id"]: 0.0 for doc in corpus}
+        if not query_terms or not corpus:
+            return scores
+
+        N = len(corpus)
+        tokenized_corpus = {doc["id"]: _tokenize(doc.get("content", "")) for doc in corpus}
+        doc_lengths = {doc_id: len(words) for doc_id, words in tokenized_corpus.items()}
+        avgdl = sum(doc_lengths.values()) / max(1, N)
+
+        # 1. Compute Document Frequency (df) for each query term
+        df = {}
+        for term in query_terms:
+            df[term] = sum(1 for words in tokenized_corpus.values() if term in words)
+
+        # 2. Compute BM25 scores for each document
+        for term in query_terms:
+            n_q = df.get(term, 0)
+            # Standard IDF with floor to avoid negative values
+            idf = math.log((N - n_q + 0.5) / (n_q + 0.5) + 1.0)
+            idf = max(0.0001, idf)
+
+            for doc in corpus:
+                doc_id = doc["id"]
+                words = tokenized_corpus[doc_id]
+                tf = words.count(term)
+
+                if tf > 0:
+                    dl = doc_lengths[doc_id]
+                    numerator = tf * (self.k1 + 1.0)
+                    denominator = tf + self.k1 * (1.0 - self.b + self.b * (dl / max(1.0, avgdl)))
+                    scores[doc_id] += idf * (numerator / denominator)
+
+        return scores
+
+    def _calculate_dense_similarities(
+        self,
+        query_vector: List[float],
+        corpus: List[Dict[str, Any]]
+    ) -> Dict[str, float]:
+        """Calculates true Cosine similarity between query vector and document embeddings."""
+        similarities = {doc["id"]: 0.0 for doc in corpus}
+        if not query_vector:
+            return similarities
+
+        norm_q = math.sqrt(sum(q * q for q in query_vector))
+        if norm_q == 0.0:
+            return similarities
+
+        for doc in corpus:
+            doc_id = doc["id"]
+            doc_vec = doc.get("embedding", [])
+
+            if not doc_vec or len(doc_vec) != len(query_vector):
+                continue
+
+            dot_product = sum(q * d for q, d in zip(query_vector, doc_vec))
+            norm_d = math.sqrt(sum(d * d for d in doc_vec))
+
+            if norm_d > 0.0:
+                similarities[doc_id] = dot_product / (norm_q * norm_d)
+
+        return similarities
+
+    def fused_search(
+        self,
+        query: str,
+        documents: List[Dict[str, Any]],
+        query_vector: Optional[List[float]] = None,
+        top_k: int = 10,
+        category_filter: Optional[str] = None,
+        tenant_filter: Optional[str] = None
+    ) -> List[Dict[str, Any]]:
+        """
+        Executes unified Reciprocal Rank Fusion (RRF) combining Sparse BM25 and Dense Cosine ranks.
+        Provides retrieval explainability via detailed score breakdown.
+
+        Formula:
+            RRF_Score(d) = 1 / (rrf_k + rank_bm25(d)) + 1 / (rrf_k + rank_dense(d))
+
+        Args:
+            query: Raw query string.
+            documents: List of memory documents.
+            query_vector: Dense vector array representing the query.
+            top_k: Maximum elements to return.
+            category_filter: Filter by memory category.
+            tenant_filter: Filter by tenant_id.
+        """
+        # 1. Apply Filtering
+        filtered_docs = documents
+        if category_filter:
+            filtered_docs = [d for d in filtered_docs if d.get("category") == category_filter]
+        if tenant_filter:
+            filtered_docs = [d for d in filtered_docs if d.get("tenant_id") == tenant_filter]
+
+        if not filtered_docs:
+            return []
+
+        query_terms = _tokenize(query)
+
+        # 2. Compute pure Lexical BM25 scores and sort to get Sparse ranks
+        bm25_scores = self._calculate_bm25_scores(query_terms, filtered_docs)
+        sparse_ranking = sorted(filtered_docs, key=lambda x: bm25_scores[x["id"]], reverse=True)
+        sparse_ranks = {doc["id"]: idx + 1 for idx, doc in enumerate(sparse_ranking)}
+
+        # 3. Compute pure Cosine Vector similarities and sort to get Dense ranks
+        dense_scores = self._calculate_dense_similarities(query_vector or [], filtered_docs)
+        dense_ranking = sorted(filtered_docs, key=lambda x: dense_scores[x["id"]], reverse=True)
+        dense_ranks = {doc["id"]: idx + 1 for idx, doc in enumerate(dense_ranking)}
+
+        # 4. Apply Reciprocal Rank Fusion (RRF)
+        results = []
+        for doc in filtered_docs:
+            doc_id = doc["id"]
+
+            # Reciprocal sparse rank score
+            rank_s = sparse_ranks.get(doc_id, len(filtered_docs))
+            rrf_sparse = 1.0 / (self.rrf_k + rank_s)
+
+            # Reciprocal dense rank score
+            rank_d = dense_ranks.get(doc_id, len(filtered_docs))
+            rrf_dense = 1.0 / (self.rrf_k + rank_d)
+
+            fused_score = rrf_sparse + rrf_dense
+
+            # Structure Explainability Data
+            explanation = {
+                "lexical_score": bm25_scores[doc_id],
+                "semantic_score": dense_scores[doc_id],
+                "fused_score": fused_score,
+                "sparse_rank": rank_s,
+                "dense_rank": rank_d
+            }
+
+            scored_doc = doc.copy()
+            scored_doc["relevance_score"] = fused_score
+            scored_doc["explanation"] = explanation
+            results.append(scored_doc)
+
+        # 5. Final Ranking
+        results.sort(key=lambda x: x["relevance_score"], reverse=True)
+        return results[:top_k]
+
+    def search(
+        self,
+        query: str,
+        corpus: Optional[List[Dict[str, Any]]] = None,
+        documents: Optional[List[Dict[str, Any]]] = None,
+        query_vector: Optional[List[float]] = None,
+        top_k: int = 10,
+        category_filter: Optional[str] = None,
+        tenant_filter: Optional[str] = None,
+        **kwargs: Any
+    ) -> List[Dict[str, Any]]:
+        """
+        Intuitive alias for fused_search supporting both 'corpus' and 'documents' arguments.
+        """
+        target_docs = corpus if corpus is not None else (documents or [])
+        return self.fused_search(
+            query=query,
+            documents=target_docs,
+            query_vector=query_vector,
+            top_k=top_k,
+            category_filter=category_filter,
+            tenant_filter=tenant_filter
+        )

synapse_memory/core/knapsack_packer.py

@@ -0,0 +1,103 @@
+from typing import List, Dict, Any, Optional
+import logging
+
+logger = logging.getLogger("KnapsackPacker")
+
+
+class MemoryNodeResult(dict):
+    """
+    Dual-contract result node that acts both as a dictionary with keys
+    ('id', 'content', etc.) and equals its string ID for backwards-compatible
+    membership checks.
+    """
+    def __eq__(self, other: Any) -> bool:
+        if isinstance(other, str):
+            return self.get("id") == other
+        return super().__eq__(other)
+
+    def __hash__(self) -> int:  # type: ignore[override]
+        return hash(self.get("id", ""))
+
+
+class KnapsackPacker:
+    """
+    Robust 0/1 Knapsack optimizer for packing memory nodes into a token budget.
+    Implements a memory-efficient Dynamic Programming solution with a greedy fallback.
+    Auto-switches to greedy when n*budget exceeds threshold to bound latency.
+    """
+    MAX_BUDGET = 8000
+    DP_OPS_THRESHOLD = 2_000_000
+
+    def __init__(self, token_budget: int = 4096):
+        self.token_budget = min(token_budget, self.MAX_BUDGET)
+
+    def pack(
+        self,
+        memories: List[Dict[str, Any]],
+        token_budget: Optional[int] = None
+    ) -> List[MemoryNodeResult]:
+        """
+        Packs memories into the token budget using a DP approach.
+
+        Args:
+            memories: List of memory nodes with 'id', 'token_cost', and 'relevance_score'.
+            token_budget: Optional override for budget. If None, uses self.token_budget.
+
+        Returns:
+            List of packed memory nodes that maximize total relevance score within budget.
+        """
+        budget = token_budget if token_budget is not None else self.token_budget
+        budget = min(budget, self.MAX_BUDGET)
+        if not memories or budget <= 0:
+            return []
+
+        # Guard: DP is O(n*budget). Fall back to greedy for huge inputs to keep p99 low.
+        if len(memories) * budget > self.DP_OPS_THRESHOLD:
+            logger.warning(
+                "Knapsack DP input too large (n=%d, budget=%d); using greedy fallback.",
+                len(memories),
+                budget,
+            )
+            return self.greedy_pack(memories, token_budget=budget)
+
+        n = len(memories)
+        weights = [max(1, int(m.get("token_cost", 1))) for m in memories]
+        values = [max(0, int(float(m.get("relevance_score", 0.0)) * 1000)) for m in memories]
+
+        dp = [0] * (budget + 1)
+        items_included: List[List[int]] = [[] for _ in range(budget + 1)]
+
+        for i in range(n):
+            weight = weights[i]
+            value = values[i]
+            if weight > budget:
+                continue
+            for w in range(budget, weight - 1, -1):
+                if dp[w - weight] + value > dp[w]:
+                    dp[w] = dp[w - weight] + value
+                    items_included[w] = items_included[w - weight] + [i]
+
+        selected_indices = items_included[budget]
+        return [MemoryNodeResult(memories[idx]) for idx in selected_indices]
+
+    def greedy_pack(
+        self,
+        memories: List[Dict[str, Any]],
+        token_budget: Optional[int] = None
+    ) -> List[MemoryNodeResult]:
+        """Greedy fallback: Pack by highest relevance/cost ratio."""
+        budget = token_budget if token_budget is not None else self.token_budget
+        sorted_memories = sorted(
+            memories,
+            key=lambda m: float(m.get("relevance_score", 0.0)) / max(1, m.get("token_cost", 1)),
+            reverse=True
+        )
+
+        selected: List[MemoryNodeResult] = []
+        current_cost = 0
+        for m in sorted_memories:
+            cost = m.get("token_cost", 1)
+            if current_cost + cost <= budget:
+                selected.append(MemoryNodeResult(m))
+                current_cost += cost
+        return selected