membrane-aol 0.2.0__py3-none-any.whl

membrane/__init__.py

@@ -0,0 +1,209 @@
+"""Membrane — prompt budgeting, U-shaped layout, caching, and retrieval hooks."""
+
+from membrane.async_preflight import (
+    get_preflight_cache,
+    get_preflight_cache_key_prefix,
+    get_preflight_cache_ttl_default,
+    run_async_preflight,
+    set_preflight_cache,
+)
+from membrane.cache_format import (
+    CacheableBlock,
+    anthropic_system_blocks,
+    assemble_system_content,
+    build_system_messages,
+    is_anthropic_model,
+    make_cacheable_block,
+)
+from membrane.compiler import (
+    PREFLIGHT_ANALYST_SYSTEM,
+    INTERRUPT_INSTRUCTION,
+    MEMORY_PLACEMENT_END,
+    MEMORY_PLACEMENT_SYSTEM,
+    CompileQuote,
+    CompileResult,
+    ContextCompiler,
+    SecurityFlag,
+    compile_anonymous,
+    compile_prompt_async,
+    compile_prompt_payload,
+    compile_quote_from_result,
+)
+from membrane.environment import (
+    EnvironmentSnapshot,
+    capture_cpu_load,
+    clear_environment_snapshot,
+    get_current_environment_state,
+    set_environment_snapshot,
+    set_environment_snapshot_from_state,
+    snapshot_to_stress_inputs,
+)
+from membrane.firewall import (
+    TIER_EPISODIC,
+    TIER_PROCEDURAL,
+    TIER_SEMANTIC,
+    PreFlightContextFirewall,
+    count_tokens,
+)
+from membrane.glaciator import GlaciateResult, glaciate
+from membrane.interrupt import (
+    INTERRUPT_TAG_END,
+    INTERRUPT_TAG_START,
+    contains_interrupt,
+    parse_interrupt,
+    strip_all_interrupts,
+)
+from membrane.layout import ContextLayoutEngine
+from membrane.exceptions import (
+    FeatureNotAvailable,
+    MembraneError,
+    MemoryCollisionError,
+)
+from membrane.llm_adapters import (
+    CallablePreFlightAnalyst,
+    PreFlightAnalyst,
+    as_preflight_analyst,
+)
+from membrane.lod import (
+    compute_attention_shift,
+    compute_attention_shift_with_embeddings,
+    cosine_distance,
+    embed_texts,
+)
+from membrane.metrics import (
+    CompileMetrics,
+    LoggingMetricsCallback,
+    MetricsCallback,
+    NoOpMetricsCallback,
+    build_compile_metrics,
+)
+from membrane.memory_protocol import (
+    EpisodicMemory,
+    KERNEL_MEMORY_SCHEMA_ID,
+    KERNEL_MEMORY_SCHEMA_VERSION,
+    KernelMemoryObject,
+    KernelMemoryRecord,
+    MemoryItem,
+    ProceduralMemory,
+    ProceduralStep,
+    SemanticMemory,
+    infer_memory_kind,
+    kernel_memory_to_memory_item,
+    load_kernel_memory_json_schema,
+    memory_item_to_kernel_memory,
+    normalize_compile_memories,
+)
+from membrane.provider_registry import (
+    get_capabilities_for_model,
+    get_context_cap_for_model,
+    lookup_context_cap,
+    get_full_capability_map,
+    get_full_context_caps,
+    has_prompt_caching,
+    has_streaming_interrupt,
+    reload_capability_maps,
+)
+from membrane.ranking_diagnostics import RankingWeights, compute_ranking_diagnostics
+from membrane.ref_map import get_ref_map_for_conversation, set_ref_map_for_conversation
+from membrane.serializers import DenseSerializer
+from membrane.stress import (
+    compute_stress_from_environment,
+    get_recent_stress_for_user,
+    set_recent_stress_for_user,
+)
+
+PreFlightFirewall = PreFlightContextFirewall
+LayoutEngine = ContextLayoutEngine
+
+__all__ = [
+    "INTERRUPT_INSTRUCTION",
+    "INTERRUPT_TAG_END",
+    "INTERRUPT_TAG_START",
+    "KERNEL_MEMORY_SCHEMA_ID",
+    "KERNEL_MEMORY_SCHEMA_VERSION",
+    "MEMORY_PLACEMENT_END",
+    "MEMORY_PLACEMENT_SYSTEM",
+    "PREFLIGHT_ANALYST_SYSTEM",
+    "TIER_EPISODIC",
+    "TIER_PROCEDURAL",
+    "TIER_SEMANTIC",
+    "CacheableBlock",
+    "CallablePreFlightAnalyst",
+    "CompileMetrics",
+    "CompileQuote",
+    "CompileResult",
+    "ContextCompiler",
+    "ContextLayoutEngine",
+    "DenseSerializer",
+    "EnvironmentSnapshot",
+    "EpisodicMemory",
+    "FeatureNotAvailable",
+    "GlaciateResult",
+    "KernelMemoryObject",
+    "KernelMemoryRecord",
+    "LayoutEngine",
+    "LoggingMetricsCallback",
+    "MembraneError",
+    "MemoryCollisionError",
+    "MemoryItem",
+    "MetricsCallback",
+    "NoOpMetricsCallback",
+    "PreFlightAnalyst",
+    "PreFlightContextFirewall",
+    "PreFlightFirewall",
+    "ProceduralMemory",
+    "ProceduralStep",
+    "RankingWeights",
+    "SecurityFlag",
+    "SemanticMemory",
+    "anthropic_system_blocks",
+    "as_preflight_analyst",
+    "assemble_system_content",
+    "build_compile_metrics",
+    "build_system_messages",
+    "capture_cpu_load",
+    "clear_environment_snapshot",
+    "compile_anonymous",
+    "compile_prompt_async",
+    "compile_prompt_payload",
+    "compile_quote_from_result",
+    "compute_attention_shift",
+    "compute_attention_shift_with_embeddings",
+    "compute_ranking_diagnostics",
+    "compute_stress_from_environment",
+    "contains_interrupt",
+    "cosine_distance",
+    "count_tokens",
+    "embed_texts",
+    "get_capabilities_for_model",
+    "get_context_cap_for_model",
+    "get_current_environment_state",
+    "get_full_capability_map",
+    "get_full_context_caps",
+    "get_preflight_cache",
+    "get_preflight_cache_key_prefix",
+    "get_preflight_cache_ttl_default",
+    "get_recent_stress_for_user",
+    "get_ref_map_for_conversation",
+    "glaciate",
+    "has_prompt_caching",
+    "has_streaming_interrupt",
+    "infer_memory_kind",
+    "is_anthropic_model",
+    "kernel_memory_to_memory_item",
+    "load_kernel_memory_json_schema",
+    "lookup_context_cap",
+    "make_cacheable_block",
+    "memory_item_to_kernel_memory",
+    "normalize_compile_memories",
+    "parse_interrupt",
+    "reload_capability_maps",
+    "run_async_preflight",
+    "set_environment_snapshot",
+    "set_environment_snapshot_from_state",
+    "set_preflight_cache",
+    "set_recent_stress_for_user",
+    "set_ref_map_for_conversation",
+    "snapshot_to_stress_inputs",
+    "strip_all_interrupts",
+]

membrane/adapters/__init__.py

@@ -0,0 +1,5 @@
+"""Framework adapters (optional extras)."""
+
+from membrane.exceptions import MemoryCollisionError
+
+__all__ = ["MemoryCollisionError"]

membrane/adapters/crewai.py

@@ -0,0 +1,51 @@
+"""CrewAI: agent-level context hook pattern."""
+
+from __future__ import annotations
+
+from typing import Any, Literal, cast
+from collections.abc import Sequence
+
+from membrane.memory_protocol import MemoryItem
+
+MemoryType = Literal["semantic", "episodic", "procedural"]
+
+
+def _memory_type(value: Any) -> MemoryType:
+    text = str(value)
+    if text in ("semantic", "episodic", "procedural"):
+        return cast(MemoryType, text)
+    return "semantic"
+
+
+def document_to_memory(doc: Any) -> MemoryItem:
+    """Map a generic CrewAI / dict-like document to
+    :class:`~membrane.memory_protocol.MemoryItem`.
+    """
+    if isinstance(doc, dict):
+        return MemoryItem(
+            id=str(doc.get("id", "doc")),
+            content=str(doc.get("content", "")),
+            memory_type=_memory_type(doc.get("memory_type", "semantic")),
+            priority=int(doc.get("priority", 2)),
+        )
+    content = getattr(doc, "content", None) or getattr(doc, "raw", str(doc))
+    return MemoryItem(
+        id=str(getattr(doc, "id", "doc")),
+        content=str(content),
+        memory_type="semantic",
+        priority=2,
+    )
+
+
+def compile_for_agent(
+    compiler: Any,
+    *,
+    memories: Sequence[Any],
+    user_message: str,
+    profile: Any = None,
+) -> Any:
+    """Run :meth:`~membrane.compiler.ContextCompiler.compile` for one agent turn."""
+    fn = getattr(compiler, "compile", None)
+    if fn is None:
+        raise TypeError("compiler must have compile()")
+    return fn(memories=memories, user_message=user_message, profile=profile)

membrane/adapters/haystack.py

@@ -0,0 +1,51 @@
+"""Haystack: document conversion and prompt assembly hook."""
+
+from __future__ import annotations
+
+from typing import Any, Literal, cast
+from collections.abc import Sequence
+
+from membrane.memory_protocol import MemoryItem
+
+MemoryType = Literal["semantic", "episodic", "procedural"]
+
+
+def _memory_type(value: Any) -> MemoryType:
+    text = str(value)
+    if text in ("semantic", "episodic", "procedural"):
+        return cast(MemoryType, text)
+    return "semantic"
+
+
+def haystack_document_to_memory(doc: Any) -> MemoryItem:
+    """Convert Haystack ``Document`` to :class:`~membrane.memory_protocol.MemoryItem`."""
+    try:
+        from haystack import Document  # type: ignore
+    except ImportError as e:
+        raise ImportError("membrane-aol[haystack] requires haystack-ai") from e
+    if not isinstance(doc, Document):
+        raise TypeError("expected haystack.Document")
+    meta = doc.meta or {}
+    return MemoryItem(
+        id=str(meta.get("id", doc.id or "doc")),
+        content=doc.content or "",
+        memory_type=_memory_type(meta.get("memory_type", "semantic")),
+        priority=int(meta.get("priority", 2)),
+    )
+
+
+class MembranePromptBuilder:
+    """Callable that returns Membrane-compiled system text for Haystack pipelines."""
+
+    def __init__(self, compiler: Any, memories: Sequence[Any], *, profile: Any = None) -> None:
+        self._compiler = compiler
+        self._memories = memories
+        self._profile = profile
+
+    def __call__(self, query: str) -> str:
+        fn = getattr(self._compiler, "compile", None)
+        if fn is None:
+            raise TypeError("compiler must have compile()")
+        res = fn(memories=self._memories, user_message=query, profile=self._profile)
+        msgs = getattr(res, "messages", []) or []
+        return "\n".join(str(m.get("content", "")) for m in msgs if isinstance(m, dict))

membrane/adapters/langchain.py

@@ -0,0 +1,113 @@
+"""
+LangChain integration.
+
+**Collision prevention:** Membrane replaces LangChain's conversation memory when you use
+:class:`MembranePromptTemplate`. Do not combine with ``ConversationBufferMemory`` /
+``ConversationSummaryMemory`` on the same chain — call :func:`disable_langchain_memory` on the
+chain or raise :class:`membrane.exceptions.MemoryCollisionError` if both are detected.
+"""
+
+from __future__ import annotations
+
+import contextlib
+
+from typing import Any, Literal, cast
+from collections.abc import Sequence
+
+from membrane.exceptions import MemoryCollisionError
+from membrane.memory_protocol import MemoryItem
+
+MemoryType = Literal["semantic", "episodic", "procedural"]
+
+
+def _memory_type(value: Any) -> MemoryType:
+    text = str(value)
+    if text in ("semantic", "episodic", "procedural"):
+        return cast(MemoryType, text)
+    return "semantic"
+
+
+def _priority(value: Any) -> int:
+    try:
+        return max(1, min(3, int(value)))
+    except (TypeError, ValueError):
+        return 2
+
+
+def langchain_doc_to_memory(doc: Any) -> MemoryItem:
+    """Convert a LangChain ``Document`` to :class:`~membrane.memory_protocol.MemoryItem`."""
+    try:
+        from langchain_core.documents import Document  # type: ignore
+    except ImportError as e:
+        raise ImportError("membrane-aol[langchain] requires langchain-core") from e
+    if not isinstance(doc, Document):
+        raise TypeError("expected langchain_core.documents.Document")
+    return MemoryItem(
+        id=str(doc.metadata.get("id", doc.metadata.get("source", "doc"))),
+        content=doc.page_content,
+        memory_type=_memory_type(doc.metadata.get("memory_type", "semantic")),
+        priority=_priority(doc.metadata.get("priority", 2)),
+    )
+
+
+class MembranePromptTemplate:
+    """
+    Wraps formatting so ``ContextCompiler.compile`` runs before prompt construction.
+
+    Pass ``compiler`` (:class:`~membrane.compiler.ContextCompiler`), ``memories``, and
+    ``profile`` / kwargs supported by :meth:`~membrane.compiler.ContextCompiler.compile`.
+    """
+
+    def __init__(
+        self,
+        compiler: Any,
+        template: str,
+        *,
+        memories: Sequence[Any],
+        profile: Any = None,
+        **compile_kwargs: Any,
+    ) -> None:
+        self._compiler = compiler
+        self._template = template
+        self._memories = memories
+        self._profile = profile
+        self._compile_kwargs = compile_kwargs
+
+    def format(self, **kwargs: Any) -> str:
+        user_message = kwargs.get("user_message") or kwargs.get("input") or ""
+        if not user_message:
+            raise ValueError("MembranePromptTemplate.format requires user_message or input")
+        compile_fn = getattr(self._compiler, "compile", None)
+        if compile_fn is None:
+            raise TypeError("compiler must have compile()")
+        res = compile_fn(
+            memories=self._memories,
+            user_message=str(user_message),
+            profile=self._profile,
+            **self._compile_kwargs,
+        )
+        msgs = getattr(res, "messages", []) or []
+        core = "\n".join(str(m.get("content", "")) for m in msgs if isinstance(m, dict))
+        fmt_kwargs = dict(kwargs)
+        fmt_kwargs["compiled_context"] = core
+        return self._template.format(**fmt_kwargs)
+
+
+def disable_langchain_memory(chain: Any) -> Any:
+    """Remove common LangChain memory components from a chain (best-effort)."""
+    mem_attrs = ("memory", "memory_key")
+    for a in mem_attrs:
+        if hasattr(chain, a):
+            with contextlib.suppress(Exception):
+                setattr(chain, a, None)
+    return chain
+
+
+def detect_memory_collision(chain: Any, membrane_active: bool) -> None:
+    if not membrane_active:
+        return
+    if hasattr(chain, "memory") and getattr(chain, "memory", None) is not None:
+        raise MemoryCollisionError(
+            "LangChain memory is set while Membrane manages context. "
+            "Remove ConversationBufferMemory or disable Membrane."
+        )

membrane/adapters/llamaindex.py

@@ -0,0 +1,90 @@
+"""LlamaIndex: convert retrieved nodes to Membrane memories and post-process."""
+
+from __future__ import annotations
+
+import contextlib
+from typing import Any
+
+from membrane.memory_protocol import KernelMemoryObject
+
+
+def node_to_memory(node: Any) -> KernelMemoryObject:
+    """Convert a LlamaIndex ``NodeWithScore`` (or similar) to
+    :class:`~membrane.memory_protocol.KernelMemoryObject`.
+
+    The retrieval score becomes ``relevance_score``; the firewall tier is always 2 (semantic).
+    """
+    text = getattr(node, "text", None) or getattr(node, "get_content", lambda: "")()
+    if callable(text):
+        text = text()
+    nid = getattr(node, "id_", None) or getattr(node, "node_id", None) or "node"
+    score = getattr(node, "score", None)
+    if score is None and hasattr(node, "metadata"):
+        score = node.metadata.get("score")
+    relevance = 1.0
+    if score is not None:
+        with contextlib.suppress(TypeError, ValueError):
+            relevance = max(0.0, float(score))
+    return KernelMemoryObject(
+        id=str(nid),
+        payload=str(text),
+        relevance_score=relevance,
+        semantic_tier=2,
+    )
+
+
+try:
+    from llama_index.core.bridge.pydantic import PrivateAttr  # type: ignore
+    from llama_index.core.postprocessor.types import BaseNodePostprocessor  # type: ignore
+    from llama_index.core.schema import NodeWithScore, TextNode  # type: ignore
+
+    _HAS_LLAMA_INDEX = True
+except ImportError:  # pragma: no cover - exercised only without the extra
+    _HAS_LLAMA_INDEX = False
+    BaseNodePostprocessor = object  # type: ignore[assignment,misc]
+
+    def PrivateAttr(*args: Any, **kwargs: Any) -> Any:  # type: ignore[misc]
+        return None
+
+
+class MembraneNodePostprocessor(BaseNodePostprocessor):  # type: ignore[misc,valid-type]
+    """
+    Postprocessor that runs Membrane's compiler on retrieved nodes.
+
+    Requires ``llama-index-core`` and a configured :class:`~membrane.compiler.ContextCompiler`.
+    """
+
+    _compiler: Any = PrivateAttr()
+    _max_tokens: int = PrivateAttr()
+    _profile: Any = PrivateAttr()
+
+    def __init__(self, compiler: Any, *, max_tokens: int = 8192, profile: Any = None) -> None:
+        if not _HAS_LLAMA_INDEX:
+            raise ImportError("membrane-aol[llamaindex] requires llama-index-core")
+        super().__init__()
+        self._compiler = compiler
+        self._max_tokens = max_tokens
+        self._profile = profile
+
+    @classmethod
+    def class_name(cls) -> str:
+        return "MembraneNodePostprocessor"
+
+    def _postprocess_nodes(self, nodes: list[Any], query_bundle: Any = None) -> list[Any]:
+        memories = [node_to_memory(n) for n in nodes]
+        q = ""
+        if query_bundle is not None:
+            q = str(getattr(query_bundle, "query_str", None) or getattr(query_bundle, "query", ""))
+        compile_fn = getattr(self._compiler, "compile", None)
+        if compile_fn is None:
+            return nodes
+        res = compile_fn(
+            memories=memories,
+            user_message=q or "context",
+            profile=self._profile,
+        )
+        # Return a single synthetic node carrying compiled text (caller interprets)
+        text = "\n".join(
+            str(m.get("content", "")) for m in getattr(res, "messages", []) or []
+        )
+        return [NodeWithScore(node=TextNode(text=text, metadata={"membrane": True}))]

membrane/adapters/mcp_resources.py

@@ -0,0 +1,35 @@
+"""MCP resource payloads for profiles, last compile, and metrics (session state)."""
+
+from __future__ import annotations
+
+import json
+from typing import Any
+
+from membrane.profile import PRESET_NAMES
+
+
+_PROFILE_NAMES = PRESET_NAMES
+
+
+def list_profile_resource_text() -> str:
+    return json.dumps({"presets": list(_PROFILE_NAMES)}, indent=2)
+
+
+def last_compile_resource_text(last: Any | None) -> str:
+    if last is None:
+        return json.dumps({"error": "no_compilation_yet"})
+    if hasattr(last, "to_debug_dict"):
+        return json.dumps(last.to_debug_dict(), default=str, indent=2)
+    return json.dumps({"raw": str(last)})
+
+
+def metrics_resource_text(snapshot: dict[str, Any] | None) -> str:
+    return json.dumps(snapshot or {"session": "membrane-mcp", "metrics": "none"}, indent=2)
+
+
+def resource_uris() -> list[str]:
+    return [
+        "membrane://profiles",
+        "membrane://last_compile",
+        "membrane://metrics",
+    ]