pdm-memory 0.1.2__py3-none-any.whl

pdm_memory/integrations/anthropic_adapter.py

@@ -0,0 +1,179 @@
+"""
+Anthropic Adapter — Task 4.1
+
+Wraps an Anthropic client for automatic PDM memory injection and saving.
+
+Usage:
+    from pdm_memory.integrations import wrap_anthropic
+    from pdm_memory import Memory
+
+    mem = Memory(store="./my_app.db")
+    client = wrap_anthropic(api_key="sk-ant-...", memory=mem)
+    reply = client.chat("What units should I use?")
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any, Optional
+
+logger = logging.getLogger(__name__)
+
+
+class PDMAnthropicClient:
+    """
+    Memory-augmented Anthropic client.
+
+    All chat() calls automatically:
+      - Recall relevant memories before the API call.
+      - Inject them via the system prompt.
+      - Save both turns as new memories.
+    """
+
+    def __init__(
+        self,
+        anthropic_client: Any,
+        memory: Any,
+        model: str = "claude-3-haiku-20240307",
+        max_memory_tokens: int = 1500,
+        recall_k: int = 5,
+        auto_save: bool = True,
+        system_prompt: str = "You are a helpful AI assistant.",
+    ) -> None:
+        self._client = anthropic_client
+        self._memory = memory
+        self._model = model
+        self._max_memory_tokens = max_memory_tokens
+        self._recall_k = recall_k
+        self._auto_save = auto_save
+        self._system_prompt = system_prompt
+
+        from pdm_memory.integrations.context_manager import ContextWindowManager
+        self._ctx_manager = ContextWindowManager(
+            max_tokens=max_memory_tokens,
+            model="gpt-4o-mini",  # Use cl100k_base approximation for Anthropic
+        )
+
+    def chat(
+        self,
+        message: str,
+        model: Optional[str] = None,
+        system_prompt: Optional[str] = None,
+        recall_k: Optional[int] = None,
+        save_reply: Optional[bool] = None,
+        max_tokens: int = 1024,
+        **anthropic_kwargs: Any,
+    ) -> str:
+        """
+        Send a message to Anthropic Claude with automatic PDM memory injection.
+
+        Args:
+            message:      The user's message.
+            model:        Override the default model.
+            system_prompt: Override the base system prompt.
+            recall_k:     Override the number of memories to recall.
+            save_reply:   Override whether to save this exchange to memory.
+            max_tokens:   Max tokens for Claude's response.
+            **anthropic_kwargs: Passed to anthropic.messages.create().
+
+        Returns:
+            The assistant's reply text.
+        """
+        k = recall_k if recall_k is not None else self._recall_k
+        should_save = save_reply if save_reply is not None else self._auto_save
+        base_system = system_prompt or self._system_prompt
+
+        # Step 1: Recall relevant memories
+        hits = self._memory.recall(query=message, k=k)
+        trimmed = self._ctx_manager.fit(hits)
+
+        # Step 2: Build system message with memory block
+        memory_block = self._ctx_manager.format_for_prompt(trimmed)
+        full_system = base_system
+        if memory_block:
+            full_system = f"{base_system}\n\n{memory_block}"
+
+        # Step 3: Call Anthropic
+        response = self._client.messages.create(
+            model=model or self._model,
+            max_tokens=max_tokens,
+            system=full_system,
+            messages=[{"role": "user", "content": message}],
+            **anthropic_kwargs,
+        )
+        reply_text = response.content[0].text if response.content else ""
+
+        # Step 4: Save this turn to memory
+        if should_save:
+            self._save_turn(message, reply_text)
+
+        logger.debug(
+            "[PDM-Anthropic] chat() | recalled=%d injected=%d model=%s",
+            len(hits), len(trimmed), model or self._model,
+        )
+        return reply_text
+
+    def _save_turn(self, user_msg: str, assistant_reply: str) -> None:
+        try:
+            if user_msg.strip():
+                self._memory.save(
+                    text=user_msg[:500],
+                    source="anthropic_chat",
+                    tags=["conversation", "user_input"],
+                    p_magnitude=40.0,
+                )
+        except Exception as e:
+            logger.warning("[PDM-Anthropic] Failed to save user message: %s", e)
+
+        try:
+            if assistant_reply.strip():
+                self._memory.save(
+                    text=assistant_reply[:500],
+                    source="anthropic_chat",
+                    tags=["conversation", "ai_reply"],
+                    p_magnitude=35.0,
+                )
+        except Exception as e:
+            logger.warning("[PDM-Anthropic] Failed to save assistant reply: %s", e)
+
+
+def wrap_anthropic(
+    api_key: str,
+    memory: Any,
+    model: str = "claude-3-haiku-20240307",
+    max_memory_tokens: int = 1500,
+    recall_k: int = 5,
+    system_prompt: str = "You are a helpful AI assistant.",
+    **anthropic_init_kwargs: Any,
+) -> PDMAnthropicClient:
+    """
+    Create a memory-augmented Anthropic client.
+
+    Args:
+        api_key:           Your Anthropic API key.
+        memory:            pdm_memory.Memory instance.
+        model:             Default Claude model.
+        max_memory_tokens: Token budget for injected memories.
+        recall_k:          Memories to recall per turn.
+        system_prompt:     Base system prompt.
+        **anthropic_init_kwargs: Passed to anthropic.Anthropic().
+
+    Returns:
+        PDMAnthropicClient wrapping an anthropic.Anthropic client.
+    """
+    try:
+        import anthropic
+    except ImportError:
+        raise ImportError(
+            "anthropic package is required. Install it: pip install pdm-memory[anthropic]"
+        )
+
+    client = anthropic.Anthropic(api_key=api_key, **anthropic_init_kwargs)
+    return PDMAnthropicClient(
+        anthropic_client=client,
+        memory=memory,
+        model=model,
+        max_memory_tokens=max_memory_tokens,
+        recall_k=recall_k,
+        system_prompt=system_prompt,
+    )

pdm_memory/integrations/context_manager.py

@@ -0,0 +1,134 @@
+"""
+Context Window Manager — Task 4.2
+
+Trims recalled memories to fit within a model's token budget.
+
+Usage:
+    manager = ContextWindowManager(max_tokens=1500, model="gpt-4o")
+    trimmed_hits = manager.fit(hits)
+    system_block = manager.format_for_prompt(trimmed_hits)
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import List, Optional
+
+from pdm_memory.core.signature import MemoryHit
+
+logger = logging.getLogger(__name__)
+
+
+class ContextWindowManager:
+    """
+    Manages the token budget for PDM memories injected into LLM context.
+
+    Trimming strategy: drop memories with the lowest p_effective first,
+    until the remaining set fits within max_tokens.
+
+    Args:
+        max_tokens:     Maximum tokens the memory block may consume.
+        model:          Model name (used to load the correct tokeniser).
+        chars_per_token: Approximation fallback if tiktoken is unavailable (default 4).
+    """
+
+    def __init__(
+        self,
+        max_tokens: int = 1500,
+        model: str = "gpt-4o-mini",
+        chars_per_token: float = 4.0,
+    ) -> None:
+        self.max_tokens = max_tokens
+        self.model = model
+        self._chars_per_token = chars_per_token
+        self._encoder = self._load_encoder(model)
+
+    # ------------------------------------------------------------------
+    # Public API
+    # ------------------------------------------------------------------
+
+    def fit(self, hits: List[MemoryHit]) -> List[MemoryHit]:
+        """
+        Return the highest-pressure subset of hits that fits within max_tokens.
+
+        Memories are sorted by p_effective descending.  The lowest-p ones
+        are dropped first to respect the token budget.
+
+        Args:
+            hits: List of MemoryHit objects from mem.recall().
+
+        Returns:
+            Trimmed list of MemoryHit, still sorted by p_effective desc.
+        """
+        if not hits:
+            return hits
+
+        # Sort highest pressure first (most important memories stay)
+        sorted_hits = sorted(hits, key=lambda h: h.p_effective, reverse=True)
+
+        kept: List[MemoryHit] = []
+        used_tokens = 0
+
+        for hit in sorted_hits:
+            tokens = self.count_tokens(hit.text)
+            if used_tokens + tokens <= self.max_tokens:
+                kept.append(hit)
+                used_tokens += tokens
+            else:
+                logger.debug(
+                    "[CTX] Trimmed memory %s (P=%.1f, tokens=%d): budget exhausted",
+                    hit.id[:8], hit.p_effective, tokens,
+                )
+
+        if len(kept) < len(hits):
+            logger.info(
+                "[CTX] Token budget %d: kept %d/%d memories (used %d tokens)",
+                self.max_tokens, len(kept), len(hits), used_tokens,
+            )
+
+        return kept
+
+    def format_for_prompt(self, hits: List[MemoryHit]) -> str:
+        """
+        Format a list of MemoryHit objects into a system prompt block.
+
+        Returns an empty string if no hits.
+        """
+        if not hits:
+            return ""
+
+        lines = ["[MEMORY CONTEXT — Pressure-Driven Memory]"]
+        for i, hit in enumerate(hits, 1):
+            tags = ", ".join(hit.intent_tags) if hit.intent_tags else hit.domain
+            lines.append(
+                f"{i}. [{hit.drawer}] {hit.text}  "
+                f"(relevance: {hit.p_effective:.0f}/100, tags: {tags})"
+            )
+        lines.append("[END MEMORY CONTEXT]")
+        return "\n".join(lines)
+
+    def count_tokens(self, text: str) -> int:
+        """Estimate token count for a string."""
+        if self._encoder is not None:
+            try:
+                return len(self._encoder.encode(text))
+            except Exception:
+                pass
+        # Fallback: character approximation
+        return max(1, int(len(text) / self._chars_per_token))
+
+    # ------------------------------------------------------------------
+    # Private
+    # ------------------------------------------------------------------
+
+    @staticmethod
+    def _load_encoder(model: str) -> Optional[object]:
+        try:
+            import tiktoken
+            try:
+                return tiktoken.encoding_for_model(model)
+            except KeyError:
+                return tiktoken.get_encoding("cl100k_base")
+        except ImportError:
+            logger.debug("[CTX] tiktoken not installed; using char approximation")
+            return None

pdm_memory/integrations/gemini_adapter.py

@@ -0,0 +1,215 @@
+"""
+Gemini Adapter
+
+Wraps Google Gemini client for automatic PDM memory injection and saving.
+Supports both the new google-genai Client and legacy google-generativeai GenerativeModel.
+
+Usage:
+    from pdm_memory.integrations import wrap_gemini
+    from pdm_memory import Memory
+    from google import genai
+
+    mem = Memory(store="./my_app.db")
+    client = genai.Client(api_key="...")
+    wrapped_client = wrap_gemini(client, memory=mem)
+    reply = wrapped_client.chat("What are my preferences?")
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any, Optional
+
+logger = logging.getLogger(__name__)
+
+
+class PDMGeminiClient:
+    """
+    Memory-augmented Gemini client.
+
+    All chat() calls automatically:
+      - Recall relevant memories before the API call.
+      - Inject them into the system instruction/prompt.
+      - Save both user message and AI reply as new memories.
+    """
+
+    def __init__(
+        self,
+        gemini_client: Any,
+        memory: Any,                  # pdm_memory.Memory instance
+        model: str = "gemini-2.5-flash",
+        max_memory_tokens: int = 1500,
+        recall_k: int = 5,
+        auto_save: bool = True,
+        system_prompt: str = "You are a helpful AI assistant.",
+    ) -> None:
+        self._client = gemini_client
+        self._memory = memory
+        self._model = model
+        self._max_memory_tokens = max_memory_tokens
+        self._recall_k = recall_k
+        self._auto_save = auto_save
+        self._system_prompt = system_prompt
+
+        from pdm_memory.integrations.context_manager import ContextWindowManager
+        self._ctx_manager = ContextWindowManager(
+            max_tokens=max_memory_tokens,
+            model="gpt-4o-mini",      # Use cl100k_base approximation
+        )
+
+    def chat(
+        self,
+        message: str,
+        model: Optional[str] = None,
+        system_prompt: Optional[str] = None,
+        recall_k: Optional[int] = None,
+        save_reply: Optional[bool] = None,
+        **gemini_kwargs: Any,
+    ) -> str:
+        """
+        Send a message to Gemini with automatic PDM memory injection.
+
+        Args:
+            message:        The user's message.
+            model:          Override the default model name.
+            system_prompt:   Override the base system prompt.
+            recall_k:       Override the number of memories to recall.
+            save_reply:     Override whether to save this exchange to memory.
+            **gemini_kwargs: Passed through to generate_content API calls.
+
+        Returns:
+            The assistant's reply text.
+        """
+        k = recall_k if recall_k is not None else self._recall_k
+        should_save = save_reply if save_reply is not None else self._auto_save
+        base_system = system_prompt or self._system_prompt
+
+        # Step 1: Recall relevant memories
+        hits = self._memory.recall(query=message, k=k)
+        trimmed = self._ctx_manager.fit(hits)
+
+        # Step 2: Build system message with memory block
+        memory_block = self._ctx_manager.format_for_prompt(trimmed)
+        full_system = base_system
+        if memory_block:
+            full_system = f"{base_system}\n\n{memory_block}"
+
+        model_name = model or self._model
+
+        # Step 3: Call Gemini API
+        if hasattr(self._client, "models") and hasattr(self._client.models, "generate_content"):
+            # New google-genai Client SDK
+            config = gemini_kwargs.pop("config", None)
+            if config is None:
+                try:
+                    from google.genai import types
+                    config = types.GenerateContentConfig(system_instruction=full_system)
+                except ImportError:
+                    config = {"system_instruction": full_system}
+            else:
+                if hasattr(config, "system_instruction"):
+                    config.system_instruction = full_system
+                elif isinstance(config, dict):
+                    config["system_instruction"] = full_system
+
+            response = self._client.models.generate_content(
+                model=model_name,
+                contents=message,
+                config=config,
+                **gemini_kwargs,
+            )
+            reply_text = response.text or ""
+        elif hasattr(self._client, "generate_content"):
+            # Legacy google-generativeai SDK (GenerativeModel object)
+            # Recreate model with dynamic system_instruction or prepend to content
+            try:
+                import google.generativeai as genai
+                temp_model = genai.GenerativeModel(
+                    model_name=getattr(self._client, "model_name", model_name),
+                    system_instruction=full_system,
+                )
+                response = temp_model.generate_content(
+                    contents=message,
+                    **gemini_kwargs,
+                )
+                reply_text = response.text or ""
+            except ImportError:
+                # Fallback: prepend system prompt as context if package is not importable
+                # (unlikely if legacy client was passed)
+                contents_with_ctx = f"System Instruction:\n{full_system}\n\nUser Input:\n{message}"
+                response = self._client.generate_content(
+                    contents=contents_with_ctx,
+                    **gemini_kwargs,
+                )
+                reply_text = response.text or ""
+        else:
+            raise TypeError(
+                "Unsupported Gemini client type. Expected a Client from google-genai "
+                "or GenerativeModel from google-generativeai."
+            )
+
+        # Step 4: Save this turn to memory
+        if should_save:
+            self._save_turn(message, reply_text)
+
+        logger.debug(
+            "[PDM-Gemini] chat() | recalled=%d injected=%d model=%s",
+            len(hits), len(trimmed), model_name,
+        )
+        return reply_text
+
+    def _save_turn(self, user_msg: str, assistant_reply: str) -> None:
+        """Save user message and assistant reply as memories."""
+        try:
+            if user_msg.strip():
+                self._memory.save(
+                    text=user_msg[:500],
+                    source="gemini_chat",
+                    tags=["conversation", "user_input"],
+                    p_magnitude=40.0,
+                )
+        except Exception as e:
+            logger.warning("[PDM-Gemini] Failed to save user message: %s", e)
+
+        try:
+            if assistant_reply.strip():
+                self._memory.save(
+                    text=assistant_reply[:500],
+                    source="gemini_chat",
+                    tags=["conversation", "ai_reply"],
+                    p_magnitude=35.0,
+                )
+        except Exception as e:
+            logger.warning("[PDM-Gemini] Failed to save assistant reply: %s", e)
+
+
+def wrap_gemini(
+    client: Any,
+    memory: Any,
+    model: str = "gemini-2.5-flash",
+    max_memory_tokens: int = 1500,
+    recall_k: int = 5,
+    system_prompt: str = "You are a helpful AI assistant.",
+) -> PDMGeminiClient:
+    """
+    Create a memory-augmented Gemini client.
+
+    Args:
+        client:            google-genai Client or google-generativeai GenerativeModel instance.
+        memory:            pdm_memory.Memory instance.
+        model:             Default model (default: gemini-2.5-flash).
+        max_memory_tokens: Token budget for injected memories.
+        recall_k:          Number of memories to recall per turn.
+        system_prompt:     Base system prompt.
+
+    Returns:
+        PDMGeminiClient wrapping the provided client.
+    """
+    return PDMGeminiClient(
+        gemini_client=client,
+        memory=memory,
+        model=model,
+        max_memory_tokens=max_memory_tokens,
+        recall_k=recall_k,
+        system_prompt=system_prompt,
+    )

pdm_memory/integrations/groq_adapter.py

@@ -0,0 +1,171 @@
+"""
+Groq Adapter
+
+Wraps Groq client for automatic PDM memory injection and saving.
+
+Usage:
+    from pdm_memory.integrations import wrap_groq
+    from pdm_memory import Memory
+    from groq import Groq
+
+    mem = Memory(store="./my_app.db")
+    client = Groq(api_key="gsk_...")
+    wrapped_client = wrap_groq(client, memory=mem)
+    reply = wrapped_client.chat("What are my preferences?")
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any, Optional
+
+logger = logging.getLogger(__name__)
+
+
+class PDMGroqClient:
+    """
+    Memory-augmented Groq client.
+
+    All chat() calls automatically:
+      - Recall relevant memories before the API call.
+      - Inject them into the system prompt.
+      - Save both user message and AI reply as new memories.
+    """
+
+    def __init__(
+        self,
+        groq_client: Any,
+        memory: Any,                  # pdm_memory.Memory instance
+        model: str = "llama-3.1-70b-versatile",
+        max_memory_tokens: int = 1500,
+        recall_k: int = 5,
+        auto_save: bool = True,
+        system_prompt: str = "You are a helpful AI assistant.",
+    ) -> None:
+        self._client = groq_client
+        self._memory = memory
+        self._model = model
+        self._max_memory_tokens = max_memory_tokens
+        self._recall_k = recall_k
+        self._auto_save = auto_save
+        self._system_prompt = system_prompt
+
+        from pdm_memory.integrations.context_manager import ContextWindowManager
+        self._ctx_manager = ContextWindowManager(
+            max_tokens=max_memory_tokens,
+            model="gpt-4o-mini",      # Use cl100k_base approximation
+        )
+
+    def chat(
+        self,
+        message: str,
+        model: Optional[str] = None,
+        system_prompt: Optional[str] = None,
+        recall_k: Optional[int] = None,
+        save_reply: Optional[bool] = None,
+        **groq_kwargs: Any,
+    ) -> str:
+        """
+        Send a message to Groq with automatic PDM memory injection.
+
+        Args:
+            message:      The user's message.
+            model:        Override the default model.
+            system_prompt: Override the base system prompt.
+            recall_k:     Override the number of memories to recall.
+            save_reply:   Override whether to save this exchange to memory.
+            **groq_kwargs: Passed through to groq.chat.completions.create().
+
+        Returns:
+            The assistant's reply text.
+        """
+        k = recall_k if recall_k is not None else self._recall_k
+        should_save = save_reply if save_reply is not None else self._auto_save
+        base_system = system_prompt or self._system_prompt
+
+        # Step 1: Recall relevant memories
+        hits = self._memory.recall(query=message, k=k)
+        trimmed = self._ctx_manager.fit(hits)
+
+        # Step 2: Build system message with memory block
+        memory_block = self._ctx_manager.format_for_prompt(trimmed)
+        full_system = base_system
+        if memory_block:
+            full_system = f"{base_system}\n\n{memory_block}"
+
+        # Step 3: Call Groq
+        response = self._client.chat.completions.create(
+            model=model or self._model,
+            messages=[
+                {"role": "system", "content": full_system},
+                {"role": "user", "content": message},
+            ],
+            **groq_kwargs,
+        )
+        reply_text = response.choices[0].message.content or ""
+
+        # Step 4: Save this turn to memory
+        if should_save:
+            self._save_turn(message, reply_text)
+
+        logger.debug(
+            "[PDM-Groq] chat() | recalled=%d injected=%d model=%s",
+            len(hits), len(trimmed), model or self._model,
+        )
+        return reply_text
+
+    def _save_turn(self, user_msg: str, assistant_reply: str) -> None:
+        """Save user message and assistant reply as memories."""
+        try:
+            if user_msg.strip():
+                self._memory.save(
+                    text=user_msg[:500],
+                    source="groq_chat",
+                    tags=["conversation", "user_input"],
+                    p_magnitude=40.0,
+                )
+        except Exception as e:
+            logger.warning("[PDM-Groq] Failed to save user message: %s", e)
+
+        try:
+            if assistant_reply.strip():
+                self._memory.save(
+                    text=assistant_reply[:500],
+                    source="groq_chat",
+                    tags=["conversation", "ai_reply"],
+                    p_magnitude=35.0,
+                )
+        except Exception as e:
+            logger.warning("[PDM-Groq] Failed to save assistant reply: %s", e)
+
+
+def wrap_groq(
+    client: Any,
+    memory: Any,
+    model: str = "llama-3.1-70b-versatile",
+    max_memory_tokens: int = 1500,
+    recall_k: int = 5,
+    system_prompt: str = "You are a helpful AI assistant.",
+) -> PDMGroqClient:
+    """
+    Create a memory-augmented Groq client.
+
+    Args:
+        client:            groq.Groq client instance.
+        memory:            pdm_memory.Memory instance.
+        model:             Default model (default: llama-3.1-70b-versatile).
+        max_memory_tokens: Token budget for injected memories.
+        recall_k:          Number of memories to recall per turn.
+        system_prompt:     Base system prompt.
+
+    Returns:
+        PDMGroqClient wrapping the provided groq client.
+    """
+    return PDMGroqClient(
+        groq_client=client,
+        memory=memory,
+        model=model,
+        max_memory_tokens=max_memory_tokens,
+        recall_k=recall_k,
+        system_prompt=system_prompt,
+    )