MemoryOS 0.0.1__py3-none-any.whl → 0.1.13__py3-none-any.whl

memos/mem_chat/simple.py

@@ -0,0 +1,200 @@
+import os
+
+from typing import Literal
+
+from memos.configs.mem_chat import SimpleMemChatConfig
+from memos.llms.factory import LLMFactory
+from memos.log import get_logger
+from memos.mem_chat.base import BaseMemChat
+from memos.mem_cube.base import BaseMemCube
+from memos.memories.activation.kv import move_dynamic_cache_htod
+from memos.memories.textual.item import TextualMemoryItem
+from memos.types import ChatHistory, MessageList
+
+
+logger = get_logger(__name__)
+
+
+class SimpleMemChat(BaseMemChat):
+    """Simple MemChat class."""
+
+    def __init__(self, config: SimpleMemChatConfig):
+        """Initialize the MemChat with the given configuration."""
+        self.config = config
+        self.chat_llm = LLMFactory.from_config(config.chat_llm)
+        self._mem_cube = None
+
+    @property
+    def mem_cube(self) -> BaseMemCube:
+        """The memory cube associated with this MemChat."""
+        return self._mem_cube
+
+    @mem_cube.setter
+    def mem_cube(self, value: BaseMemCube) -> None:
+        """The memory cube associated with this MemChat."""
+        self._mem_cube = value
+
+    def run(self) -> None:
+        """Run the MemChat."""
+
+        # Start MemChat
+
+        print(
+            "\n📢 [System] " + "Simple MemChat is running.\n"
+            "Commands: 'bye' to quit, 'clear' to clear chat history, 'mem' to show all memories, 'export' to export chat history\n",
+        )
+
+        messages = []
+        while True:
+            # Get user input
+
+            user_input = input("👤 [You] ").strip()
+            print()
+
+            if user_input.lower() == "bye":
+                break
+            elif user_input.lower() == "clear":
+                messages = []
+                print("📢 [System] Chat history cleared.")
+                continue
+            elif user_input.lower() == "mem":
+                if self.config.enable_textual_memory:
+                    all_memories = self.mem_cube.text_mem.get_all()
+                    print(f"🧠 [Memory] \n{self._str_memories(all_memories)}\n")
+                else:
+                    print("📢 [System] Textual memory is not enabled.\n")
+                continue
+            elif user_input.lower() == "export":
+                if messages:
+                    filepath = self._export_chat_history(messages)
+                    print(f"📢 [System] Chat history exported to: {filepath}\n")
+                else:
+                    print("📢 [System] No chat history to export.\n")
+                continue
+            elif user_input == "":
+                continue
+
+            # Get memories
+
+            if self.config.enable_textual_memory:
+                memories = self.mem_cube.text_mem.search(user_input, top_k=self.config.top_k)
+                print(
+                    f"🧠 [Memory] Searched memories:\n{self._str_memories(memories, mode='concise')}\n"
+                )
+                system_prompt = self._build_system_prompt(memories)
+            else:
+                system_prompt = self._build_system_prompt()
+            current_messages = [
+                {"role": "system", "content": system_prompt},
+                *messages,
+                {"role": "user", "content": user_input},
+            ]
+
+            if self.config.enable_activation_memory:
+                past_key_values = None
+                loaded_kv_cache_item = next(
+                    iter(self.mem_cube.act_mem.kv_cache_memories.values()), None
+                )
+                if loaded_kv_cache_item is not None:
+                    # If has loaded kv cache, we move it to device before inferring.
+                    # Currently, we move only single kv cache item
+                    past_key_values = loaded_kv_cache_item
+                    past_key_values.kv_cache = move_dynamic_cache_htod(
+                        past_key_values.kv_cache, self.chat_llm.model.device
+                    )
+
+                # Generate response
+                response = self.chat_llm.generate(
+                    current_messages,
+                    past_key_values=past_key_values.kv_cache if past_key_values else None,
+                )
+            else:
+                # Generate response without activation memory
+                response = self.chat_llm.generate(current_messages)
+
+            print(f"🤖 [Assistant] {response}\n")
+            messages.append({"role": "user", "content": user_input})
+            messages.append({"role": "assistant", "content": response})
+            messages = messages[
+                -self.config.max_turns_window :
+            ]  # Keep only recent messages to avoid context overflow
+
+            # Extract memories
+
+            if self.config.enable_textual_memory:
+                new_memories = self.mem_cube.text_mem.extract(messages[-2:])
+                for memory in new_memories:
+                    memory.metadata.user_id = self.config.user_id
+                    memory.metadata.session_id = self.config.session_id
+                    memory.metadata.status = "activated"
+                self.mem_cube.text_mem.add(new_memories)
+                print(
+                    f"🧠 [Memory] Stored {len(new_memories)} new memory(ies):\n"
+                    f"{self._str_memories(new_memories, 'concise')}\n"
+                )
+
+        # Stop MemChat
+
+        print("📢 [System] MemChat has stopped.")
+
+    def _build_system_prompt(self, memories: list | None = None) -> str:
+        """Build system prompt with optional memories context."""
+        base_prompt = (
+            "You are a knowledgeable and helpful AI assistant. "
+            "You have access to conversation memories that help you provide more personalized responses. "
+            "Use the memories to understand the user's context, preferences, and past interactions. "
+            "If memories are provided, reference them naturally when relevant, but don't explicitly mention having memories."
+        )
+
+        if memories:
+            memory_context = "\n\n## Memories:\n"
+            for i, memory in enumerate(memories, 1):
+                memory_context += f"{i}. ({memory.metadata.memory_time}) {memory.memory}\n"
+            return base_prompt + memory_context
+
+        return base_prompt
+
+    def _str_memories(
+        self, memories: list[TextualMemoryItem], mode: Literal["concise", "full"] = "full"
+    ) -> str:
+        """Format memories for display."""
+        if not memories:
+            return "No memories."
+        if mode == "concise":
+            return "\n".join(f"{i + 1}. {memory.memory}" for i, memory in enumerate(memories))
+        elif mode == "full":
+            return "\n".join(f"{i + 1}. {memory}" for i, memory in enumerate(memories))
+
+    def _export_chat_history(self, messages: MessageList, output_dir: str = "chat_exports") -> str:
+        """Export chat history to JSON file.
+
+        Args:
+            messages: List of chat messages
+            output_dir: Directory to save the export file
+
+        Returns:
+            Path to the exported JSON file
+        """
+        # Create output directory if it doesn't exist
+        os.makedirs(output_dir, exist_ok=True)
+
+        # Generate filename with user_id and timestamp
+        timestamp = self.config.created_at.strftime("%Y%m%d_%H%M%S")
+        filename = f"{self.config.user_id}_{timestamp}_chat_history.json"
+        filepath = os.path.join(output_dir, filename)
+
+        # Prepare export data
+        export_data = ChatHistory(
+            user_id=self.config.user_id,
+            session_id=self.config.session_id,
+            created_at=self.config.created_at,
+            total_messages=len(messages),
+            chat_history=messages,
+        )
+
+        # Write to JSON file
+        with open(filepath, "w", encoding="utf-8") as f:
+            f.write(export_data.model_dump_json(indent=4, exclude_none=True, warnings="none"))
+
+        logger.info(f"Chat history exported to {filepath}")
+        return filepath

memos/mem_cube/base.py

@@ -0,0 +1,29 @@
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING
+
+from memos.configs.mem_cube import BaseMemCubeConfig
+
+
+if TYPE_CHECKING:
+    from memos.memories.activation.base import BaseActMemory
+    from memos.memories.parametric.base import BaseParaMemory
+    from memos.memories.textual.base import BaseTextMemory
+
+
+class BaseMemCube(ABC):
+    """Base class for all MemCube implementations."""
+
+    @abstractmethod
+    def __init__(self, config: BaseMemCubeConfig):
+        """Initialize the MemCube with the given configuration."""
+        self.text_mem: BaseTextMemory
+        self.act_mem: BaseActMemory
+        self.para_mem: BaseParaMemory
+
+    @abstractmethod
+    def load(self, dir: str) -> None:
+        """Load memories from a directory."""
+
+    @abstractmethod
+    def dump(self, dir: str) -> None:
+        """Dump memories to a directory."""

memos/mem_cube/general.py

@@ -0,0 +1,146 @@
+import os
+
+from memos.configs.mem_cube import GeneralMemCubeConfig
+from memos.configs.utils import get_json_file_model_schema
+from memos.exceptions import ConfigurationError, MemCubeError
+from memos.log import get_logger
+from memos.mem_cube.base import BaseMemCube
+from memos.mem_cube.utils import download_repo
+from memos.memories.activation.base import BaseActMemory
+from memos.memories.factory import MemoryFactory
+from memos.memories.parametric.base import BaseParaMemory
+from memos.memories.textual.base import BaseTextMemory
+
+
+logger = get_logger(__name__)
+
+
+class GeneralMemCube(BaseMemCube):
+    """MemCube is a box for loading and dumping three types of memories."""
+
+    def __init__(self, config: GeneralMemCubeConfig):
+        """Initialize the MemCube with a configuration."""
+        self.config = config
+        self._text_mem: BaseTextMemory | None = (
+            MemoryFactory.from_config(config.text_mem)
+            if config.text_mem.backend != "uninitialized"
+            else None
+        )
+        self._act_mem: BaseActMemory | None = (
+            MemoryFactory.from_config(config.act_mem)
+            if config.act_mem.backend != "uninitialized"
+            else None
+        )
+        self._para_mem: BaseParaMemory | None = (
+            MemoryFactory.from_config(config.para_mem)
+            if config.para_mem.backend != "uninitialized"
+            else None
+        )
+
+    def load(self, dir: str) -> None:
+        """Load memories.
+        Args:
+            dir (str): The directory containing the memory files.
+        """
+        loaded_schema = get_json_file_model_schema(os.path.join(dir, self.config.config_filename))
+        if loaded_schema != self.config.model_schema:
+            raise ConfigurationError(
+                f"Configuration schema mismatch. Expected {self.config.model_schema}, "
+                f"but found {loaded_schema}."
+            )
+        self.text_mem.load(dir) if self.text_mem else None
+        self.act_mem.load(dir) if self.act_mem else None
+        self.para_mem.load(dir) if self.para_mem else None
+
+        logger.info(f"MemCube loaded successfully from {dir}")
+
+    def dump(self, dir: str) -> None:
+        """Dump memories.
+        Args:
+            dir (str): The directory where the memory files will be saved.
+        """
+        if os.path.exists(dir) and os.listdir(dir):
+            raise MemCubeError(
+                f"Directory {dir} is not empty. Please provide an empty directory for dumping."
+            )
+
+        self.config.to_json_file(os.path.join(dir, self.config.config_filename))
+        self.text_mem.dump(dir) if self.text_mem else None
+        self.act_mem.dump(dir) if self.act_mem else None
+        self.para_mem.dump(dir) if self.para_mem else None
+
+        logger.info(f"MemCube dumped successfully to {dir}")
+
+    @staticmethod
+    def init_from_dir(dir: str) -> "GeneralMemCube":
+        """Create a MemCube instance from a MemCube directory.
+
+        Args:
+            dir (str): The directory containing the memory files.
+
+        Returns:
+            MemCube: An instance of MemCube loaded with memories from the specified directory.
+        """
+        config_path = os.path.join(dir, "config.json")
+        config = GeneralMemCubeConfig.from_json_file(config_path)
+        mem_cube = GeneralMemCube(config)
+        mem_cube.load(dir)
+        return mem_cube
+
+    @staticmethod
+    def init_from_remote_repo(
+        cube_id: str, base_url: str = "https://huggingface.co/datasets"
+    ) -> "GeneralMemCube":
+        """Create a MemCube instance from a remote repository.
+
+        Args:
+            repo (str): The repository name.
+            base_url (str): The base URL of the remote repository.
+
+        Returns:
+            MemCube: An instance of MemCube loaded with memories from the specified remote repository.
+        """
+        dir = download_repo(cube_id, base_url)
+        return GeneralMemCube.init_from_dir(dir)
+
+    @property
+    def text_mem(self) -> "BaseTextMemory | None":
+        """Get the textual memory."""
+        if self._text_mem is None:
+            logger.warning("Textual memory is not initialized. Returning None.")
+        return self._text_mem
+
+    @text_mem.setter
+    def text_mem(self, value: BaseTextMemory) -> None:
+        """Set the textual memory."""
+        if not isinstance(value, BaseTextMemory):
+            raise TypeError(f"Expected BaseTextMemory, got {type(value).__name__}")
+        self._text_mem = value
+
+    @property
+    def act_mem(self) -> "BaseActMemory | None":
+        """Get the activation memory."""
+        if self._act_mem is None:
+            logger.warning("Activation memory is not initialized. Returning None.")
+        return self._act_mem
+
+    @act_mem.setter
+    def act_mem(self, value: BaseActMemory) -> None:
+        """Set the activation memory."""
+        if not isinstance(value, BaseActMemory):
+            raise TypeError(f"Expected BaseActMemory, got {type(value).__name__}")
+        self._act_mem = value
+
+    @property
+    def para_mem(self) -> "BaseParaMemory | None":
+        """Get the parametric memory."""
+        if self._para_mem is None:
+            logger.warning("Parametric memory is not initialized. Returning None.")
+        return self._para_mem
+
+    @para_mem.setter
+    def para_mem(self, value: BaseParaMemory) -> None:
+        """Set the parametric memory."""
+        if not isinstance(value, BaseParaMemory):
+            raise TypeError(f"Expected BaseParaMemory, got {type(value).__name__}")
+        self._para_mem = value

memos/mem_cube/utils.py

@@ -0,0 +1,24 @@
+import subprocess
+import tempfile
+
+
+def download_repo(repo: str, base_url: str, dir: str | None = None) -> str:
+    """Download a repository from a remote source.
+
+    Args:
+        repo (str): The repository name.
+        base_url (str): The base URL of the remote repository.
+        dir (str, optional): The directory where the repository will be downloaded. If None, a temporary directory will be created.
+    If a directory is provided, it will be used instead of creating a temporary one.
+
+    Returns:
+        str: The local directory where the repository is downloaded.
+    """
+    if dir is None:
+        dir = tempfile.mkdtemp()
+    repo_url = f"{base_url}/{repo}"
+
+    # Clone the repo
+    subprocess.run(["git", "clone", repo_url, dir], check=True)
+
+    return dir

memos/mem_os/client.py

@@ -0,0 +1,5 @@
+# TODO: @Li Ji
+
+
+class ClientMOS:
+    pass