memplex 3.2.0__py3-none-any.whl

memnex/llm/fallback_chain.py

@@ -0,0 +1,87 @@
+"""Chain-of-responsibility fallback for LLM providers."""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING
+
+from memnex.models import IntentType
+
+if TYPE_CHECKING:
+    from memnex.llm.provider import LLMProvider
+
+logger = logging.getLogger(__name__)
+
+
+class FallbackChain:
+    """Try providers in order; first success wins, final fallback to RuleBasedProvider.
+
+    Parameters
+    ----------
+    providers:
+        Ordered list of LLMProvider implementations.  They are tried
+        sequentially; the first one that returns without raising wins.
+    """
+
+    def __init__(self, providers: list[LLMProvider] | None = None) -> None:
+        self._providers: list[LLMProvider] = providers or []
+
+    def _fallback(self) -> LLMProvider:
+        """Lazily create a RuleBasedProvider as the ultimate fallback."""
+        from memnex.llm.providers.rule_based import RuleBasedProvider
+
+        return RuleBasedProvider()
+
+    # -- LLMProvider interface ------------------------------------------
+
+    async def classify_intent(
+        self, query: str, context: dict | None = None
+    ) -> IntentType:
+        errors: list[str] = []
+        for p in self._providers:
+            try:
+                return await p.classify_intent(query, context)
+            except Exception as exc:
+                errors.append(f"{p.__class__.__name__}: {exc}")
+                logger.debug("classify_intent fallback: %s", errors[-1])
+        return await self._fallback().classify_intent(query, context)
+
+    async def summarize(self, content: str, max_tokens: int = 256) -> str:
+        for p in self._providers:
+            try:
+                return await p.summarize(content, max_tokens)
+            except Exception as exc:
+                logger.debug("summarize fallback: %s: %s", p.__class__.__name__, exc)
+        return content[:max_tokens]
+
+    async def extract_structured(self, prompt: str, schema: dict) -> dict:
+        for p in self._providers:
+            try:
+                return await p.extract_structured(prompt, schema)
+            except Exception as exc:
+                logger.debug("extract_structured fallback: %s: %s", p.__class__.__name__, exc)
+        return {}
+
+    async def generate_hypothetical(self, query: str) -> str:
+        for p in self._providers:
+            try:
+                return await p.generate_hypothetical(query)
+            except Exception as exc:
+                logger.debug("generate_hypothetical fallback: %s: %s", p.__class__.__name__, exc)
+        return query
+
+    async def complete(self, prompt: str) -> str:
+        for p in self._providers:
+            try:
+                return await p.complete(prompt)
+            except Exception as exc:
+                logger.debug("complete fallback: %s: %s", p.__class__.__name__, exc)
+        return ""
+
+    async def complete_json(self, prompt: str) -> dict:
+        for p in self._providers:
+            try:
+                return await p.complete_json(prompt)
+            except Exception as exc:
+                logger.debug("complete_json fallback: %s: %s", p.__class__.__name__, exc)
+        return {}

memnex/llm/injection_guard.py

@@ -0,0 +1,178 @@
+"""Indirect prompt injection guard for memory recall contexts."""
+
+from __future__ import annotations
+
+import logging
+import re
+from typing import TYPE_CHECKING, List
+
+if TYPE_CHECKING:
+    from memnex.models import SearchResult
+
+logger = logging.getLogger(__name__)
+
+
+class IndirectInjectionGuard:
+    """Detect and mitigate indirect prompt injection in recalled memories.
+
+    While ``LLMPromptSanitizer`` protects against *direct* injection (user
+    input), this class handles *indirect* injection: an attacker embeds
+    malicious instructions in a source document that survives extraction
+    and gets injected into the LLM context when recalled via RAG.
+
+    Defense layers:
+    1. Content scanning -- regex-based detection of system-role keywords.
+    2. Protective wrapping -- memory content is wrapped in ``[MEMORY ...]``
+       tags so the LLM treats them as data, not instructions.
+    3. Trust-level labelling -- each memory is annotated with a trust level
+       derived from its ``source_type``.
+    """
+
+    # Multi-language injection patterns (compiled once at class load)
+    INJECTION_PATTERNS: list[str] = [
+        r"ignore\s+(all\s+)?previous\s+instructions?",
+        r"disregard\s+(all\s+)?prior\s+instructions?",
+        r"system\s*:\s*you\s+are",
+        r"<\|endoftext\|>",
+        r"<\|im_start\|>",
+        r"忽略(之前|前面|上面)的(所有|全部)?指令",
+        r"忽略系统提示",
+        r"你现在是",
+        r"新的系统提示",
+        r"assistant\s*:\s*sure",
+    ]
+
+    _compiled: list[re.Pattern] = [re.compile(p, re.IGNORECASE) for p in INJECTION_PATTERNS]
+
+    # Trust level mapping: source_type value -> trust level label
+    TRUST_LEVELS: dict[str, str] = {
+        "requirement": "HIGH",
+        "meeting": "MEDIUM",
+        "code": "MEDIUM",
+        "wiki": "LOW",
+    }
+
+    @classmethod
+    def scan(cls, content: str) -> bool:
+        """Scan content for suspected injection attacks.
+
+        Returns
+        -------
+        True if the content is suspected to contain an injection payload;
+        the caller should discard or isolate the memory entry.
+        """
+        for pattern in cls._compiled:
+            if pattern.search(content):
+                return True
+        return False
+
+    @classmethod
+    def wrap_for_context(
+        cls,
+        memories: list[SearchResult],
+        store: object,
+    ) -> str:
+        """Wrap recalled memories in protective tags for LLM context injection.
+
+        Each memory is enclosed in ``[MEMORY START | trust=LEVEL | id=...]``
+        / ``[MEMORY END]`` markers so the LLM treats the content as
+        data, not as system instructions.
+
+        Parameters
+        ----------
+        memories:
+            Search results to be injected into the LLM context.
+        store:
+            A MemoryStore-like object with a ``get(id)`` method that
+            returns a MemoryNode (or None).
+        """
+        parts: list[str] = []
+        for r in memories:
+            func = store.get(r.func_id) if store else None
+            if not func:
+                continue
+            source_type_val = (
+                func.source_type.value
+                if hasattr(func.source_type, "value")
+                else (func.source_type or "wiki")
+            )
+            trust = cls.TRUST_LEVELS.get(source_type_val, "LOW")
+            summary = r.summary or func.name
+            parts.append(
+                f"[MEMORY START | trust={trust} | id={r.func_id}]\n"
+                f"{summary}\n"
+                f"[MEMORY END]"
+            )
+        return "\n\n".join(parts)
+
+    @classmethod
+    def filter_and_wrap(
+        cls,
+        memories: list[SearchResult],
+        store: object,
+    ) -> str:
+        """Filter out injection-suspected memories, then wrap the rest.
+
+        Memories that trigger the injection scanner are logged as warnings
+        and excluded from the output.
+        """
+        safe: list[SearchResult] = []
+        for r in memories:
+            func = store.get(r.func_id) if store else None
+            if func:
+                memory_type = getattr(func, "memory_type", "function")
+                text = cls._extract_scan_text(func, memory_type)
+                if cls.scan(text):
+                    logger.warning(
+                        "Indirect injection detected in memory %s (type=%s), skipped.",
+                        r.func_id,
+                        memory_type,
+                    )
+                    continue
+            safe.append(r)
+        return cls.wrap_for_context(safe, store)
+
+    @classmethod
+    def _extract_scan_text(cls, func: object, memory_type: str) -> str:
+        """Extract the relevant text fields for injection scanning by memory type."""
+        if memory_type == "function":
+            return " ".join(
+                fv.desc
+                for role in ("trigger", "condition", "action", "benefit")
+                for fv in getattr(func, role, [])
+            )
+        if memory_type == "fact":
+            return " ".join(
+                filter(
+                    None,
+                    [
+                        getattr(func, "subject", ""),
+                        getattr(func, "predicate", ""),
+                        getattr(func, "object_", ""),
+                    ],
+                )
+            )
+        if memory_type == "preference":
+            return " ".join(
+                filter(
+                    None,
+                    [
+                        getattr(func, "aspect", ""),
+                        getattr(func, "preference", ""),
+                    ],
+                )
+            )
+        if memory_type == "observation":
+            return " ".join(
+                filter(
+                    None,
+                    [
+                        getattr(func, "event", ""),
+                        getattr(func, "context", ""),
+                    ],
+                )
+            )
+        # Unknown type: scan all string attributes
+        return " ".join(
+            str(v) for v in vars(func).values() if isinstance(v, str)
+        )

memnex/llm/provider.py

@@ -0,0 +1,130 @@
+"""LLM Provider protocol definition."""
+
+from typing import Protocol, runtime_checkable
+
+from memnex.models import IntentType
+
+
+@runtime_checkable
+class LLMProvider(Protocol):
+    """LLM Provider standard protocol.
+
+    All LLM provider implementations must satisfy this protocol.
+    Used for intent classification, summarization, structured extraction,
+    HyDE generation, and general-purpose completion.
+    """
+
+    async def classify_intent(
+        self, query: str, context: dict | None = None
+    ) -> IntentType:
+        """Classify user query intent.
+
+        Returns one of: IMMEDIATE, SYNTHESIS, RELATION, ALL.
+        """
+        ...
+
+    async def summarize(self, content: str, max_tokens: int = 256) -> str:
+        """Summarize content into a concise representation."""
+        ...
+
+    async def extract_structured(self, prompt: str, schema: dict) -> dict:
+        """Extract structured data according to the provided JSON schema."""
+        ...
+
+    async def generate_hypothetical(self, query: str) -> str:
+        """Generate a hypothetical answer for HyDE (Hypothetical Document Embeddings)."""
+        ...
+
+    async def complete(self, prompt: str) -> str:
+        """General-purpose text completion."""
+        ...
+
+    async def complete_json(self, prompt: str) -> dict:
+        """Complete and parse response as JSON."""
+        ...
+
+
+def create_provider(
+    provider: str = "auto",
+    *,
+    anthropic_api_key: str | None = None,
+    anthropic_model: str = "claude-sonnet-4-6",
+    local_endpoint: str = "http://localhost:11434/v1",
+    local_model: str = "qwen2.5",
+    fallback_chain: list[str] | None = None,
+) -> LLMProvider:
+    """Factory: create an LLM provider instance by name.
+
+    Parameters
+    ----------
+    provider:
+        One of "auto", "anthropic", "local", "rule-based".
+    anthropic_api_key:
+        Anthropic API key (required for anthropic provider).
+    anthropic_model:
+        Model name for Anthropic.
+    local_endpoint:
+        OpenAI-compatible endpoint URL (e.g., Ollama / LM Studio).
+    local_model:
+        Model name for the local provider.
+    fallback_chain:
+        Ordered list of provider names for FallbackChain.
+        Defaults to ["anthropic", "local", "rule-based"] when *provider*
+        is "auto".
+
+    Returns
+    -------
+    An object satisfying the LLMProvider protocol.
+    """
+    if provider == "auto":
+        chain = fallback_chain or ["anthropic", "local", "rule-based"]
+        return create_provider(
+            provider=None,
+            anthropic_api_key=anthropic_api_key,
+            anthropic_model=anthropic_model,
+            local_endpoint=local_endpoint,
+            local_model=local_model,
+            fallback_chain=chain,
+        )
+
+    # If provider is None, build a FallbackChain from fallback_chain list.
+    if provider is None:
+        from memnex.llm.fallback_chain import FallbackChain
+
+        chain_names = fallback_chain or ["anthropic", "local", "rule-based"]
+        providers = []
+        for name in chain_names:
+            try:
+                p = create_provider(
+                    provider=name,
+                    anthropic_api_key=anthropic_api_key,
+                    anthropic_model=anthropic_model,
+                    local_endpoint=local_endpoint,
+                    local_model=local_model,
+                )
+                providers.append(p)
+            except Exception:
+                continue
+        return FallbackChain(providers)
+
+    if provider == "anthropic":
+        from memnex.llm.providers.anthropic import AnthropicProvider
+
+        if not anthropic_api_key:
+            raise ValueError("anthropic_api_key is required for the Anthropic provider")
+        return AnthropicProvider(api_key=anthropic_api_key, model=anthropic_model)
+
+    if provider == "local":
+        from memnex.llm.providers.local import LocalProvider
+
+        return LocalProvider(endpoint=local_endpoint, model=local_model)
+
+    if provider == "rule-based":
+        from memnex.llm.providers.rule_based import RuleBasedProvider
+
+        return RuleBasedProvider()
+
+    raise ValueError(
+        f"Unknown provider: {provider!r}. "
+        f"Choose from: auto, anthropic, local, rule-based"
+    )

memnex/llm/providers/__init__.py

@@ -0,0 +1,22 @@
+"""LLM provider implementations."""
+
+from memnex.llm.providers.rule_based import RuleBasedProvider
+
+# Anthropic and Local providers are optional -- they require external
+# packages that may not be installed.  Import them lazily or guard with
+# try/except at the call site.
+
+__all__ = [
+    "RuleBasedProvider",
+]
+
+
+def __getattr__(name: str):
+    """Lazy-load optional providers that depend on external packages."""
+    if name == "AnthropicProvider":
+        from memnex.llm.providers.anthropic import AnthropicProvider
+        return AnthropicProvider
+    if name == "LocalProvider":
+        from memnex.llm.providers.local import LocalProvider
+        return LocalProvider
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

memnex/llm/providers/anthropic.py

@@ -0,0 +1,135 @@
+"""Anthropic LLM provider implementation."""
+
+import json
+import logging
+
+from memnex.models import IntentType
+
+logger = logging.getLogger(__name__)
+
+try:
+    import anthropic
+except ImportError:
+    raise ImportError(
+        "The 'anthropic' package is required for AnthropicProvider. "
+        "Install it with: pip install anthropic"
+    )
+
+
+class AnthropicProvider:
+    """LLM provider backed by the Anthropic SDK (Claude).
+
+    Parameters
+    ----------
+    api_key:
+        Anthropic API key.
+    model:
+        Model identifier (default: claude-sonnet-4-6).
+    max_tokens:
+        Default maximum tokens for completions.
+    """
+
+    def __init__(
+        self,
+        api_key: str,
+        model: str = "claude-sonnet-4-6",
+        max_tokens: int = 1024,
+    ) -> None:
+        self._client = anthropic.AsyncAnthropic(api_key=api_key)
+        self._model = model
+        self._max_tokens = max_tokens
+
+    # -- helpers --------------------------------------------------------
+
+    async def _raw_complete(self, prompt: str, max_tokens: int | None = None) -> str:
+        """Send a single-turn user message and return the assistant text."""
+        resp = await self._client.messages.create(
+            model=self._model,
+            max_tokens=max_tokens or self._max_tokens,
+            messages=[{"role": "user", "content": prompt}],
+        )
+        return resp.content[0].text
+
+    async def _raw_complete_json(self, prompt: str) -> dict:
+        """Complete with JSON response expectation and parse the result."""
+        text = await self._raw_complete(prompt)
+        return self._parse_json(text)
+
+    @staticmethod
+    def _parse_json(text: str) -> dict:
+        """Best-effort JSON extraction from LLM output."""
+        text = text.strip()
+        # Try direct parse
+        try:
+            return json.loads(text)
+        except json.JSONDecodeError:
+            pass
+        # Try to extract JSON block from markdown code fence
+        import re
+        m = re.search(r"```(?:json)?\s*\n?(.*?)```", text, re.DOTALL)
+        if m:
+            try:
+                return json.loads(m.group(1).strip())
+            except json.JSONDecodeError:
+                pass
+        # Try to find first { ... } block
+        start = text.find("{")
+        end = text.rfind("}")
+        if start != -1 and end > start:
+            try:
+                return json.loads(text[start : end + 1])
+            except json.JSONDecodeError:
+                pass
+        logger.warning("Failed to parse JSON from LLM response, returning empty dict")
+        return {}
+
+    # -- LLMProvider interface ------------------------------------------
+
+    async def classify_intent(
+        self, query: str, context: dict | None = None
+    ) -> IntentType:
+        """Classify user query intent using Claude."""
+        result = await self.complete_json(
+            f"Classify the intent of the following query. "
+            f'Respond with a JSON object: {{"intent": "search|understand|compare|relation"}}\n\nQuery: {query}'
+        )
+        intent_str = result.get("intent", "search")
+        mapping = {
+            "search": IntentType.IMMEDIATE,
+            "understand": IntentType.SYNTHESIS,
+            "compare": IntentType.RELATION,
+            "relation": IntentType.RELATION,
+        }
+        return mapping.get(intent_str, IntentType.IMMEDIATE)
+
+    async def summarize(self, content: str, max_tokens: int = 256) -> str:
+        """Summarize content using Claude."""
+        prompt = (
+            f"Summarize the following content concisely in at most {max_tokens} tokens:\n\n{content}"
+        )
+        return await self._raw_complete(prompt, max_tokens=max_tokens)
+
+    async def extract_structured(self, prompt: str, schema: dict) -> dict:
+        """Extract structured data according to a JSON schema."""
+        full_prompt = (
+            f"{prompt}\n\n"
+            f"Respond with valid JSON matching this schema:\n"
+            f"{json.dumps(schema, ensure_ascii=False)}"
+        )
+        return await self._raw_complete_json(full_prompt)
+
+    async def generate_hypothetical(self, query: str) -> str:
+        """Generate a hypothetical answer for HyDE."""
+        prompt = (
+            f"Given the query below, write a brief hypothetical answer (2-3 sentences) "
+            f"as if a comprehensive knowledge base entry existed for it.\n\nQuery: {query}"
+        )
+        return await self._raw_complete(prompt, max_tokens=256)
+
+    async def complete(self, prompt: str) -> str:
+        """General-purpose text completion."""
+        return await self._raw_complete(prompt)
+
+    async def complete_json(self, prompt: str) -> dict:
+        """Complete and parse response as JSON."""
+        return await self._raw_complete_json(prompt)

memnex/llm/providers/local.py

@@ -0,0 +1,135 @@
+"""Local LLM provider using OpenAI-compatible API (Ollama / LM Studio)."""
+
+import json
+import logging
+
+from memnex.models import IntentType
+
+logger = logging.getLogger(__name__)
+
+try:
+    from openai import AsyncOpenAI
+except ImportError:
+    raise ImportError(
+        "The 'openai' package is required for LocalProvider. "
+        "Install it with: pip install openai"
+    )
+
+
+class LocalProvider:
+    """LLM provider backed by an OpenAI-compatible local API.
+
+    Works with Ollama, LM Studio, vLLM, and any server that exposes
+    the ``/v1/chat/completions`` endpoint.
+
+    Parameters
+    ----------
+    endpoint:
+        Base URL of the OpenAI-compatible API.
+    model:
+        Model identifier served by the local endpoint.
+    max_tokens:
+        Default maximum tokens for completions.
+    """
+
+    def __init__(
+        self,
+        endpoint: str = "http://localhost:11434/v1",
+        model: str = "qwen2.5",
+        max_tokens: int = 1024,
+    ) -> None:
+        self._client = AsyncOpenAI(base_url=endpoint, api_key="not-needed")
+        self._model = model
+        self._max_tokens = max_tokens
+
+    # -- helpers --------------------------------------------------------
+
+    async def _raw_complete(self, prompt: str, max_tokens: int | None = None) -> str:
+        """Send a single-turn chat completion request."""
+        resp = await self._client.chat.completions.create(
+            model=self._model,
+            max_tokens=max_tokens or self._max_tokens,
+            messages=[{"role": "user", "content": prompt}],
+        )
+        return resp.choices[0].message.content or ""
+
+    async def _raw_complete_json(self, prompt: str) -> dict:
+        """Complete with JSON response expectation and parse the result."""
+        text = await self._raw_complete(prompt)
+        return self._parse_json(text)
+
+    @staticmethod
+    def _parse_json(text: str) -> dict:
+        """Best-effort JSON extraction from LLM output."""
+        text = text.strip()
+        try:
+            return json.loads(text)
+        except json.JSONDecodeError:
+            pass
+        import re
+        m = re.search(r"```(?:json)?\s*\n?(.*?)```", text, re.DOTALL)
+        if m:
+            try:
+                return json.loads(m.group(1).strip())
+            except json.JSONDecodeError:
+                pass
+        start = text.find("{")
+        end = text.rfind("}")
+        if start != -1 and end > start:
+            try:
+                return json.loads(text[start : end + 1])
+            except json.JSONDecodeError:
+                pass
+        logger.warning("Failed to parse JSON from local LLM response, returning empty dict")
+        return {}
+
+    # -- LLMProvider interface ------------------------------------------
+
+    async def classify_intent(
+        self, query: str, context: dict | None = None
+    ) -> IntentType:
+        """Classify user query intent using local LLM."""
+        result = await self.complete_json(
+            f"Classify the intent of the following query. "
+            f'Respond with a JSON object: {{"intent": "search|understand|compare|relation"}}\n\nQuery: {query}'
+        )
+        intent_str = result.get("intent", "search")
+        mapping = {
+            "search": IntentType.IMMEDIATE,
+            "understand": IntentType.SYNTHESIS,
+            "compare": IntentType.RELATION,
+            "relation": IntentType.RELATION,
+        }
+        return mapping.get(intent_str, IntentType.IMMEDIATE)
+
+    async def summarize(self, content: str, max_tokens: int = 256) -> str:
+        """Summarize content using local LLM."""
+        prompt = (
+            f"Summarize the following content concisely in at most {max_tokens} tokens:\n\n{content}"
+        )
+        return await self._raw_complete(prompt, max_tokens=max_tokens)
+
+    async def extract_structured(self, prompt: str, schema: dict) -> dict:
+        """Extract structured data according to a JSON schema."""
+        full_prompt = (
+            f"{prompt}\n\n"
+            f"Respond with valid JSON matching this schema:\n"
+            f"{json.dumps(schema, ensure_ascii=False)}"
+        )
+        return await self._raw_complete_json(full_prompt)
+
+    async def generate_hypothetical(self, query: str) -> str:
+        """Generate a hypothetical answer for HyDE."""
+        prompt = (
+            f"Given the query below, write a brief hypothetical answer (2-3 sentences) "
+            f"as if a comprehensive knowledge base entry existed for it.\n\nQuery: {query}"
+        )
+        return await self._raw_complete(prompt, max_tokens=256)
+
+    async def complete(self, prompt: str) -> str:
+        """General-purpose text completion."""
+        return await self._raw_complete(prompt)
+
+    async def complete_json(self, prompt: str) -> dict:
+        """Complete and parse response as JSON."""
+        return await self._raw_complete_json(prompt)