MemoryOS 0.2.0__py3-none-any.whl → 0.2.1__py3-none-any.whl

memos/mem_user/persistent_user_manager.py

@@ -0,0 +1,260 @@
+"""Persistent user management system for MemOS with configuration storage.
+
+This module extends the base UserManager to provide persistent storage
+for user configurations and MOS instances.
+"""
+
+import json
+
+from datetime import datetime
+from typing import Any
+
+from sqlalchemy import Column, String, Text
+
+from memos.configs.mem_os import MOSConfig
+from memos.log import get_logger
+from memos.mem_user.user_manager import Base, UserManager
+
+
+logger = get_logger(__name__)
+
+
+class UserConfig(Base):
+    """User configuration model for the database."""
+
+    __tablename__ = "user_configs"
+
+    user_id = Column(String, primary_key=True)
+    config_data = Column(Text, nullable=False)  # JSON string of MOSConfig
+    created_at = Column(String, nullable=False)  # ISO format timestamp
+    updated_at = Column(String, nullable=False)  # ISO format timestamp
+
+    def __repr__(self):
+        return f"<UserConfig(user_id='{self.user_id}')>"
+
+
+class PersistentUserManager(UserManager):
+    """Extended UserManager with configuration persistence."""
+
+    def __init__(self, db_path: str | None = None, user_id: str = "root"):
+        """Initialize the persistent user manager.
+
+        Args:
+            db_path (str, optional): Path to the SQLite database file.
+                If None, uses default path in MEMOS_DIR.
+            user_id (str, optional): User ID. If None, uses default user ID.
+        """
+        super().__init__(db_path, user_id)
+
+        # Create user_configs table
+        Base.metadata.create_all(bind=self.engine)
+        logger.info("PersistentUserManager initialized with configuration storage")
+
+    def _convert_datetime_strings(self, obj: Any) -> Any:
+        """Recursively convert datetime strings back to datetime objects in config dict.
+
+        Args:
+            obj: The object to process (dict, list, or primitive type)
+
+        Returns:
+            The object with datetime strings converted to datetime objects
+        """
+        if isinstance(obj, dict):
+            result = {}
+            for key, value in obj.items():
+                if key == "created_at" and isinstance(value, str):
+                    try:
+                        result[key] = datetime.fromisoformat(value)
+                    except ValueError:
+                        # If parsing fails, keep the original string
+                        result[key] = value
+                else:
+                    result[key] = self._convert_datetime_strings(value)
+            return result
+        elif isinstance(obj, list):
+            return [self._convert_datetime_strings(item) for item in obj]
+        else:
+            return obj
+
+    def save_user_config(self, user_id: str, config: MOSConfig) -> bool:
+        """Save user configuration to database.
+
+        Args:
+            user_id (str): The user ID.
+            config (MOSConfig): The user's MOS configuration.
+
+        Returns:
+            bool: True if successful, False otherwise.
+        """
+        session = self._get_session()
+        try:
+            # Convert config to JSON string with proper datetime handling
+            config_dict = config.model_dump(mode="json")
+            config_json = json.dumps(config_dict, indent=2)
+
+            from datetime import datetime
+
+            now = datetime.now().isoformat()
+
+            # Check if config already exists
+            existing_config = (
+                session.query(UserConfig).filter(UserConfig.user_id == user_id).first()
+            )
+
+            if existing_config:
+                # Update existing config
+                existing_config.config_data = config_json
+                existing_config.updated_at = now
+                logger.info(f"Updated configuration for user {user_id}")
+            else:
+                # Create new config
+                user_config = UserConfig(
+                    user_id=user_id, config_data=config_json, created_at=now, updated_at=now
+                )
+                session.add(user_config)
+                logger.info(f"Saved new configuration for user {user_id}")
+
+            session.commit()
+            return True
+
+        except Exception as e:
+            session.rollback()
+            logger.error(f"Error saving user config for {user_id}: {e}")
+            return False
+        finally:
+            session.close()
+
+    def get_user_config(self, user_id: str) -> MOSConfig | None:
+        """Get user configuration from database.
+
+        Args:
+            user_id (str): The user ID.
+
+        Returns:
+            MOSConfig | None: The user's configuration or None if not found.
+        """
+        session = self._get_session()
+        try:
+            user_config = session.query(UserConfig).filter(UserConfig.user_id == user_id).first()
+
+            if user_config:
+                config_dict = json.loads(user_config.config_data)
+                # Convert datetime strings back to datetime objects
+                config_dict = self._convert_datetime_strings(config_dict)
+                return MOSConfig(**config_dict)
+            return None
+
+        except Exception as e:
+            logger.error(f"Error loading user config for {user_id}: {e}")
+            return None
+        finally:
+            session.close()
+
+    def delete_user_config(self, user_id: str) -> bool:
+        """Delete user configuration from database.
+
+        Args:
+            user_id (str): The user ID.
+
+        Returns:
+            bool: True if successful, False otherwise.
+        """
+        session = self._get_session()
+        try:
+            user_config = session.query(UserConfig).filter(UserConfig.user_id == user_id).first()
+
+            if user_config:
+                session.delete(user_config)
+                session.commit()
+                logger.info(f"Deleted configuration for user {user_id}")
+                return True
+            return False
+
+        except Exception as e:
+            session.rollback()
+            logger.error(f"Error deleting user config for {user_id}: {e}")
+            return False
+        finally:
+            session.close()
+
+    def list_user_configs(self) -> dict[str, MOSConfig]:
+        """List all user configurations.
+
+        Returns:
+            Dict[str, MOSConfig]: Dictionary mapping user_id to MOSConfig.
+        """
+        session = self._get_session()
+        try:
+            user_configs = session.query(UserConfig).all()
+            result = {}
+
+            for user_config in user_configs:
+                try:
+                    config_dict = json.loads(user_config.config_data)
+                    # Convert datetime strings back to datetime objects
+                    config_dict = self._convert_datetime_strings(config_dict)
+                    result[user_config.user_id] = MOSConfig(**config_dict)
+                except Exception as e:
+                    logger.error(f"Error parsing config for user {user_config.user_id}: {e}")
+                    continue
+
+            return result
+
+        except Exception as e:
+            logger.error(f"Error listing user configs: {e}")
+            return {}
+        finally:
+            session.close()
+
+    def create_user_with_config(
+        self, user_name: str, config: MOSConfig, role=None, user_id: str | None = None
+    ) -> str:
+        """Create a new user with configuration.
+
+        Args:
+            user_name (str): Name of the user.
+            config (MOSConfig): The user's configuration.
+            role: User role (optional, uses default from UserManager).
+            user_id (str, optional): Custom user ID.
+
+        Returns:
+            str: The created user ID.
+
+        Raises:
+            ValueError: If user_name already exists.
+        """
+        # Create user using parent method
+        created_user_id = self.create_user(user_name, role, user_id)
+
+        # Save configuration
+        if not self.save_user_config(created_user_id, config):
+            logger.error(f"Failed to save configuration for user {created_user_id}")
+
+        return created_user_id
+
+    def delete_user(self, user_id: str) -> bool:
+        """Delete a user and their configuration.
+
+        Args:
+            user_id (str): The user ID.
+
+        Returns:
+            bool: True if successful, False otherwise.
+        """
+        # Delete configuration first
+        self.delete_user_config(user_id)
+
+        # Delete user using parent method
+        return super().delete_user(user_id)
+
+    def get_user_cube_access(self, user_id: str) -> list[str]:
+        """Get list of cube IDs that a user has access to.
+
+        Args:
+            user_id (str): The user ID.
+
+        Returns:
+            list[str]: List of cube IDs the user can access.
+        """
+        cubes = self.get_user_cubes(user_id)
+        return [cube.cube_id for cube in cubes]

memos/memories/activation/item.py

@@ -1,5 +1,6 @@
 import uuid
 
+from datetime import datetime
 from typing import Any
 
 from pydantic import BaseModel, ConfigDict, Field
@@ -12,6 +13,16 @@ class ActivationMemoryItem(BaseModel):
     metadata: dict = {}
 
 
+class KVCacheRecords(BaseModel):
+    text_memories: list[str] = Field(
+        default=[],
+        description="The list of text memories transformed to the activation memory.",
+    )
+    timestamp: datetime = Field(
+        default_factory=datetime.now, description="submit time for schedule_messages"
+    )
+
+
 class KVCacheItem(ActivationMemoryItem):
     id: str = Field(default_factory=lambda: str(uuid.uuid4()))
     memory: DynamicCache = Field(
@@ -23,3 +34,17 @@ class KVCacheItem(ActivationMemoryItem):
     )
 
     model_config = ConfigDict(arbitrary_types_allowed=True)  # To allow DynamicCache as a field type
+    records: KVCacheRecords = KVCacheRecords()
+
+
+class VLLMKVCacheItem(KVCacheItem):
+    """
+    VLLM KV Cache Item that stores prompt strings instead of DynamicCache objects.
+    This is because vLLM handles KV cache on the server side via preloading.
+    """
+
+    # Override memory field to store prompt string instead of DynamicCache
+    memory: str = Field(
+        default="",
+        description="Prompt string used to preload KV cache in vLLM server",
+    )

memos/memories/activation/kv.py

@@ -3,11 +3,10 @@ import pickle
 
 from datetime import datetime
 
-import torch
-
 from transformers import DynamicCache
 
 from memos.configs.memory import KVCacheMemoryConfig
+from memos.dependency import require_python_package
 from memos.llms.factory import LLMFactory
 from memos.memories.activation.base import BaseActMemory
 from memos.memories.activation.item import KVCacheItem
@@ -20,6 +19,10 @@ class KVCacheMemory(BaseActMemory):
     This memory type is designed to store and retrieve key-value caches.
     """
 
+    @require_python_package(
+        import_name="torch",
+        install_link="https://pytorch.org/get-started/locally/",
+    )
     def __init__(self, config: KVCacheMemoryConfig) -> None:
         """Initialize the KV Cache Memory with a configuration."""
         self.config = config
@@ -139,6 +142,8 @@ class KVCacheMemory(BaseActMemory):
         Args:
             dir (str): The directory containing the memory files.
         """
+        import torch
+
         file_path = os.path.join(dir, self.config.memory_filename)
 
         if not os.path.exists(file_path):
@@ -197,6 +202,8 @@ class KVCacheMemory(BaseActMemory):
         Faster concat merge: for each layer, gather all caches' tensors
         and do a single torch.cat per layer.
         """
+        import torch
+
         assert caches, "Need at least one cache"
         if len(caches) == 1:
             return caches[0]
@@ -215,7 +222,7 @@ class KVCacheMemory(BaseActMemory):
         return merged
 
 
-def move_dynamic_cache_htod(dynamic_cache: DynamicCache, device: torch.device) -> DynamicCache:
+def move_dynamic_cache_htod(dynamic_cache: DynamicCache, device: str) -> DynamicCache:
     """
     In SimpleMemChat.run(), if self.config.enable_activation_memory is enabled,
     we load serialized kv cache from a [class KVCacheMemory] object, which has a kv_cache_memories on CPU.

memos/memories/activation/vllmkv.py

@@ -0,0 +1,219 @@
+import os
+import pickle
+
+from datetime import datetime
+
+from memos.configs.memory import KVCacheMemoryConfig
+from memos.dependency import require_python_package
+from memos.llms.factory import LLMFactory
+from memos.memories.activation.base import BaseActMemory
+from memos.memories.activation.item import VLLMKVCacheItem
+from memos.memories.textual.item import TextualMemoryItem
+
+
+class VLLMKVCacheMemory(BaseActMemory):
+    """
+    VLLM Key-Value Cache Memory for activation memories.
+    This memory type is designed to store and retrieve prompt strings for vLLM KV cache preloading.
+    Unlike traditional KV cache that stores DynamicCache objects, vLLM handles cache on server side.
+    """
+
+    @require_python_package(
+        import_name="torch",
+        install_link="https://pytorch.org/get-started/locally/",
+    )
+    def __init__(self, config: KVCacheMemoryConfig) -> None:
+        """Initialize the VLLM KV Cache Memory with a configuration."""
+        self.config = config
+        self.llm = LLMFactory.from_config(config.extractor_llm)
+        self.kv_cache_memories: dict[str, VLLMKVCacheItem] = {}
+
+    def extract(self, text: str) -> VLLMKVCacheItem:
+        """Extract memory based on the text.
+
+        Uses the LLM to build vLLM KV cache from the provided text.
+        For vLLM, this means preloading the KV cache on the server side.
+
+        Args:
+            text: Input text to extract memory from
+
+        Returns:
+            Extracted VLLM KV cache item with prompt string
+        """
+        # Build vLLM KV cache from the text using the LLM
+        # This preloads the cache on the vLLM server and returns the prompt
+        prompt = self.llm.build_vllm_kv_cache(text)
+
+        # Create a VLLMKVCacheItem with the extracted prompt
+        cache_item = VLLMKVCacheItem(
+            memory=prompt,
+            metadata={"source_text": text, "extracted_at": datetime.now().isoformat()},
+        )
+
+        return cache_item
+
+    def add(self, memories: list[VLLMKVCacheItem]) -> None:
+        """Add memories to the VLLM KV cache memory.
+
+        Args:
+            memories: List of VLLMKVCacheItem to add
+        """
+        for memory in memories:
+            self.kv_cache_memories[memory.id] = memory
+
+    def get_cache(self, cache_ids: list[str]) -> str | None:
+        """Get the prompt string for the most recent cache.
+
+        Since vLLM handles KV cache on server side, we return the prompt string
+        that can be used for generation. For multiple caches, we return the most recent one.
+
+        Args:
+            cache_ids: List of cache IDs to consider
+
+        Returns:
+            Prompt string for the most recent cache or None if no caches found
+        """
+        if not cache_ids:
+            return None
+
+        # For vLLM, we typically want the most recent cache
+        # Return the prompt from the last cache ID in the list
+        latest_cache_id = cache_ids[-1]
+        cache_item = self.kv_cache_memories.get(latest_cache_id)
+
+        if cache_item and cache_item.memory:
+            return cache_item.memory
+
+        return None
+
+    def get(self, memory_id: str) -> VLLMKVCacheItem | None:
+        """Get a memory by its ID.
+
+        Args:
+            memory_id: ID of the memory to retrieve
+
+        Returns:
+            VLLMKVCacheItem or None if not found
+        """
+        return self.kv_cache_memories.get(memory_id)
+
+    def get_by_ids(self, memory_ids: list[str]) -> list[VLLMKVCacheItem | None]:
+        """Get memories by their IDs.
+
+        Args:
+            memory_ids: List of memory IDs to retrieve
+
+        Returns:
+            List of VLLMKVCacheItem or None for missing ones
+        """
+        results = []
+        for memory_id in memory_ids:
+            memory = self.get(memory_id)
+            results.append(memory)
+        return results
+
+    def get_all(self) -> list[VLLMKVCacheItem]:
+        """Get all memories.
+
+        Returns:
+            List of all VLLMKVCacheItems in the memory
+        """
+        return list(self.kv_cache_memories.values())
+
+    def delete(self, memory_ids: list[str]) -> None:
+        """Delete memories by their IDs.
+
+        Args:
+            memory_ids: List of memory IDs to delete
+        """
+        for memory_id in memory_ids:
+            self.kv_cache_memories.pop(memory_id, None)
+
+    def delete_all(self) -> None:
+        """Delete all memories."""
+        self.kv_cache_memories.clear()
+
+    def from_textual_memory(self, mem: TextualMemoryItem) -> VLLMKVCacheItem:
+        """
+        Convert a TextualMemoryItem to a VLLMKVCacheItem.
+        This method extracts the prompt string from the textual memory.
+        """
+        # Build vLLM KV cache from the textual memory content
+        prompt = self.llm.build_vllm_kv_cache(mem.memory)
+        return VLLMKVCacheItem(memory=prompt, metadata=mem.metadata.model_dump())
+
+    def load(self, dir: str) -> None:
+        """Load memories from os.path.join(dir, self.config.memory_filename)
+
+        Args:
+            dir (str): The directory containing the memory files.
+        """
+        file_path = os.path.join(dir, self.config.memory_filename)
+
+        if not os.path.exists(file_path):
+            # If file doesn't exist, start with empty memories
+            return
+
+        try:
+            # Allow loading VLLMKVCacheItem types
+            import torch
+
+            torch.serialization.add_safe_globals([VLLMKVCacheItem])
+
+            with open(file_path, "rb") as f:
+                data = pickle.load(f)
+
+            if isinstance(data, dict):
+                # Load memories, handle both old and new formats
+                if "kv_cache_memories" in data:
+                    memories = data["kv_cache_memories"]
+                    if isinstance(memories, list):
+                        # Convert list to dict format
+                        self.kv_cache_memories = {item.id: item for item in memories}
+                    else:
+                        self.kv_cache_memories = memories
+                else:
+                    # Reset to empty if no memories in data
+                    self.kv_cache_memories = {}
+            elif isinstance(data, list):
+                # Backward compatibility: convert list to dict
+                self.kv_cache_memories = {item.id: item for item in data}
+            else:
+                # Reset to empty if data format is unexpected
+                self.kv_cache_memories = {}
+
+        except (EOFError, pickle.UnpicklingError, Exception):
+            # If loading fails, start with empty memories
+            self.kv_cache_memories = {}
+
+    def dump(self, dir: str) -> None:
+        """Dump memories to os.path.join(dir, self.config.memory_filename)
+
+        Args:
+            dir (str): The directory where the memory files will be saved.
+        """
+        file_path = os.path.join(dir, self.config.memory_filename)
+
+        # Create directory if it doesn't exist
+        os.makedirs(dir, exist_ok=True)
+
+        # Prepare data to save (only memories)
+        data = {"kv_cache_memories": self.kv_cache_memories}
+
+        with open(file_path, "wb") as f:
+            pickle.dump(data, f, protocol=pickle.HIGHEST_PROTOCOL)
+
+    def preload_kv_cache(self, cache_ids: list[str]) -> None:
+        """
+        Preload KV cache on vLLM server for the given cache IDs.
+        This method calls build_vllm_kv_cache for each cache to ensure
+        the KV cache is loaded on the server side.
+
+        Args:
+            cache_ids: List of cache IDs to preload
+        """
+        for cache_id in cache_ids:
+            cache_item = self.kv_cache_memories.get(cache_id)
+            if cache_item and cache_item.memory:
+                # Re-preload the KV cache on the server
+                self.llm.build_vllm_kv_cache(cache_item.memory)

memos/memories/factory.py

@@ -3,6 +3,7 @@ from typing import Any, ClassVar
 from memos.configs.memory import MemoryConfigFactory
 from memos.memories.activation.base import BaseActMemory
 from memos.memories.activation.kv import KVCacheMemory
+from memos.memories.activation.vllmkv import VLLMKVCacheMemory
 from memos.memories.base import BaseMemory
 from memos.memories.parametric.base import BaseParaMemory
 from memos.memories.parametric.lora import LoRAMemory
@@ -20,6 +21,7 @@ class MemoryFactory(BaseMemory):
         "general_text": GeneralTextMemory,
         "tree_text": TreeTextMemory,
         "kv_cache": KVCacheMemory,
+        "vllm_kv_cache": VLLMKVCacheMemory,
         "lora": LoRAMemory,
     }
 

memos/memories/textual/general.py

@@ -7,8 +7,8 @@ from typing import Any
 from tenacity import retry, retry_if_exception_type, stop_after_attempt
 
 from memos.configs.memory import GeneralTextMemoryConfig
-from memos.embedders.factory import EmbedderFactory, OllamaEmbedder
-from memos.llms.factory import LLMFactory, OllamaLLM, OpenAILLM
+from memos.embedders.factory import ArkEmbedder, EmbedderFactory, OllamaEmbedder
+from memos.llms.factory import AzureLLM, LLMFactory, OllamaLLM, OpenAILLM
 from memos.log import get_logger
 from memos.memories.textual.base import BaseTextMemory
 from memos.memories.textual.item import TextualMemoryItem
@@ -26,9 +26,11 @@ class GeneralTextMemory(BaseTextMemory):
     def __init__(self, config: GeneralTextMemoryConfig):
         """Initialize memory with the given configuration."""
         self.config: GeneralTextMemoryConfig = config
-        self.extractor_llm: OpenAILLM | OllamaLLM = LLMFactory.from_config(config.extractor_llm)
+        self.extractor_llm: OpenAILLM | OllamaLLM | AzureLLM = LLMFactory.from_config(
+            config.extractor_llm
+        )
         self.vector_db: QdrantVecDB = VecDBFactory.from_config(config.vector_db)
-        self.embedder: OllamaEmbedder = EmbedderFactory.from_config(config.embedder)
+        self.embedder: OllamaEmbedder | ArkEmbedder = EmbedderFactory.from_config(config.embedder)
 
     @retry(
         stop=stop_after_attempt(3),
@@ -202,7 +204,7 @@ class GeneralTextMemory(BaseTextMemory):
 
     def _embed_one_sentence(self, sentence: str) -> list[float]:
         """Embed a single sentence."""
-        return self.embedder.embed(sentence)[0]
+        return self.embedder.embed([sentence])[0]
 
 
 EXTRACTION_PROMPT_PART_1 = f"""You are a memory extractor. Your task is to extract memories from the given messages.

memos/memories/textual/tree.py

@@ -10,7 +10,7 @@ from typing import Any
 from memos.configs.memory import TreeTextMemoryConfig
 from memos.embedders.factory import EmbedderFactory, OllamaEmbedder
 from memos.graph_dbs.factory import GraphStoreFactory, Neo4jGraphDB
-from memos.llms.factory import LLMFactory, OllamaLLM, OpenAILLM
+from memos.llms.factory import AzureLLM, LLMFactory, OllamaLLM, OpenAILLM
 from memos.log import get_logger
 from memos.memories.textual.base import BaseTextMemory
 from memos.memories.textual.item import TextualMemoryItem, TreeNodeTextualMemoryMetadata
@@ -31,8 +31,12 @@ class TreeTextMemory(BaseTextMemory):
     def __init__(self, config: TreeTextMemoryConfig):
         """Initialize memory with the given configuration."""
         self.config: TreeTextMemoryConfig = config
-        self.extractor_llm: OpenAILLM | OllamaLLM = LLMFactory.from_config(config.extractor_llm)
-        self.dispatcher_llm: OpenAILLM | OllamaLLM = LLMFactory.from_config(config.dispatcher_llm)
+        self.extractor_llm: OpenAILLM | OllamaLLM | AzureLLM = LLMFactory.from_config(
+            config.extractor_llm
+        )
+        self.dispatcher_llm: OpenAILLM | OllamaLLM | AzureLLM = LLMFactory.from_config(
+            config.dispatcher_llm
+        )
         self.embedder: OllamaEmbedder = EmbedderFactory.from_config(config.embedder)
         self.graph_store: Neo4jGraphDB = GraphStoreFactory.from_config(config.graph_db)
         self.is_reorganize = config.reorganize
@@ -53,7 +57,7 @@ class TreeTextMemory(BaseTextMemory):
         else:
             logger.info("No internet retriever configured")
 
-    def add(self, memories: list[TextualMemoryItem | dict[str, Any]]) -> None:
+    def add(self, memories: list[TextualMemoryItem | dict[str, Any]]) -> list[str]:
         """Add memories.
         Args:
             memories: List of TextualMemoryItem objects or dictionaries to add.
@@ -63,7 +67,7 @@ class TreeTextMemory(BaseTextMemory):
             plan = plan_memory_operations(memory_items, metadata, self.graph_store)
             execute_plan(memory_items, metadata, plan, self.graph_store)
         """
-        self.memory_manager.add(memories)
+        return self.memory_manager.add(memories)
 
     def replace_working_memory(self, memories: list[TextualMemoryItem]) -> None:
         self.memory_manager.replace_working_memory(memories)

memos/memories/textual/tree_text_memory/organize/conflict.py

@@ -167,10 +167,12 @@ class ConflictHandler:
             if not self.graph_store.edge_exists(new_from, new_to, edge["type"], direction="ANY"):
                 self.graph_store.add_edge(new_from, new_to, edge["type"])
 
-        self.graph_store.delete_node(conflict_a.id)
-        self.graph_store.delete_node(conflict_b.id)
+        self.graph_store.update_node(conflict_a.id, {"status": "archived"})
+        self.graph_store.update_node(conflict_b.id, {"status": "archived"})
+        self.graph_store.add_edge(conflict_a.id, merged.id, type="MERGED_TO")
+        self.graph_store.add_edge(conflict_b.id, merged.id, type="MERGED_TO")
         logger.debug(
-            f"Remove {conflict_a.id} and {conflict_b.id}, and inherit their edges to {merged.id}."
+            f"Archive {conflict_a.id} and {conflict_b.id}, and inherit their edges to {merged.id}."
         )
 
     def _merge_metadata(