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

memos/llms/hf_singleton.py

@@ -0,0 +1,114 @@
+import threading
+
+from typing import ClassVar
+
+from memos.configs.llm import HFLLMConfig
+from memos.llms.hf import HFLLM
+from memos.log import get_logger
+
+
+logger = get_logger(__name__)
+
+
+class HFSingletonLLM(HFLLM):
+    """
+    Singleton version of HFLLM that prevents multiple loading of the same model.
+    This class inherits from HFLLM and adds singleton behavior.
+    """
+
+    _instances: ClassVar[dict[str, "HFSingletonLLM"]] = {}
+    _lock: ClassVar[threading.Lock] = threading.Lock()
+
+    def __new__(cls, config: HFLLMConfig):
+        """
+        Singleton pattern implementation.
+        Returns existing instance if config already exists, otherwise creates new one.
+        """
+        config_key = cls._get_config_key(config)
+
+        if config_key in cls._instances:
+            logger.debug(f"Reusing existing HF model: {config.model_name_or_path}")
+            return cls._instances[config_key]
+
+        with cls._lock:
+            # Double-check pattern to prevent race conditions
+            if config_key in cls._instances:
+                logger.debug(f"Reusing existing HF model: {config.model_name_or_path}")
+                return cls._instances[config_key]
+
+            logger.info(f"Creating new HF model: {config.model_name_or_path}")
+            instance = super().__new__(cls)
+            cls._instances[config_key] = instance
+            return instance
+
+    def __init__(self, config: HFLLMConfig):
+        """
+        Initialize the singleton HFLLM instance.
+        Only initializes if this is a new instance.
+        """
+        # Check if already initialized
+        if hasattr(self, "_initialized"):
+            return
+
+        # Call parent constructor
+        super().__init__(config)
+        self._initialized = True
+
+    @classmethod
+    def _get_config_key(cls, config: HFLLMConfig) -> str:
+        """
+        Generate a unique key for the HF model configuration.
+
+        Args:
+            config: The HFLLM configuration
+
+        Returns:
+            A unique string key representing the configuration
+        """
+        # Create a unique key based on model path and key parameters
+        key_parts = [config.model_name_or_path]
+        return "|".join(key_parts)
+
+    @classmethod
+    def get_instance_count(cls) -> int:
+        """
+        Get the number of unique HF model instances currently managed.
+
+        Returns:
+            Number of HF model instances
+        """
+        return len(cls._instances)
+
+    @classmethod
+    def get_instance_info(cls) -> dict[str, str]:
+        """
+        Get information about all managed HF model instances.
+
+        Returns:
+            Dictionary mapping config keys to model paths
+        """
+        return {key: instance.config.model_name_or_path for key, instance in cls._instances.items()}
+
+    @classmethod
+    def clear_all(cls) -> None:
+        """
+        Clear all HF model instances from memory.
+        This should be used carefully as it will force reloading of models.
+        """
+        with cls._lock:
+            cls._instances.clear()
+            logger.info("All HF model instances cleared from singleton manager")
+
+
+# Convenience function to get singleton manager info
+def get_hf_singleton_info() -> dict[str, int]:
+    """
+    Get information about the HF singleton manager.
+
+    Returns:
+        Dictionary with instance count and info
+    """
+    return {
+        "instance_count": HFSingletonLLM.get_instance_count(),
+        "instance_info": HFSingletonLLM.get_instance_info(),
+    }

memos/llms/ollama.py

@@ -1,3 +1,4 @@
+from collections.abc import Generator
 from typing import Any
 
 from ollama import Client
@@ -80,3 +81,6 @@ class OllamaLLM(BaseLLM):
             return remove_thinking_tags(str_response)
         else:
             return str_response
+
+    def generate_stream(self, messages: MessageList, **kwargs) -> Generator[str, None, None]:
+        raise NotImplementedError

memos/llms/openai.py

@@ -1,6 +1,8 @@
+from collections.abc import Generator
+
 import openai
 
-from memos.configs.llm import OpenAILLMConfig
+from memos.configs.llm import AzureLLMConfig, OpenAILLMConfig
 from memos.llms.base import BaseLLM
 from memos.llms.utils import remove_thinking_tags
 from memos.log import get_logger
@@ -22,6 +24,7 @@ class OpenAILLM(BaseLLM):
         response = self.client.chat.completions.create(
             model=self.config.model_name_or_path,
             messages=messages,
+            extra_body=self.config.extra_body,
             temperature=self.config.temperature,
             max_tokens=self.config.max_tokens,
             top_p=self.config.top_p,
@@ -32,3 +35,67 @@ class OpenAILLM(BaseLLM):
             return remove_thinking_tags(response_content)
         else:
             return response_content
+
+    def generate_stream(self, messages: MessageList, **kwargs) -> Generator[str, None, None]:
+        """Stream response from OpenAI LLM with optional reasoning support."""
+        response = self.client.chat.completions.create(
+            model=self.config.model_name_or_path,
+            messages=messages,
+            stream=True,
+            temperature=self.config.temperature,
+            max_tokens=self.config.max_tokens,
+            top_p=self.config.top_p,
+            extra_body=self.config.extra_body,
+        )
+
+        reasoning_started = False
+
+        for chunk in response:
+            delta = chunk.choices[0].delta
+
+            # Support for custom 'reasoning_content' (if present in OpenAI-compatible models like Qwen)
+            if hasattr(delta, "reasoning_content") and delta.reasoning_content:
+                if not reasoning_started and not self.config.remove_think_prefix:
+                    yield "<think>"
+                    reasoning_started = True
+                yield delta.reasoning_content
+            elif hasattr(delta, "content") and delta.content:
+                if reasoning_started and not self.config.remove_think_prefix:
+                    yield "</think>"
+                    reasoning_started = False
+                yield delta.content
+
+        # Ensure we close the <think> block if not already done
+        if reasoning_started and not self.config.remove_think_prefix:
+            yield "</think>"
+
+
+class AzureLLM(BaseLLM):
+    """Azure OpenAI LLM class."""
+
+    def __init__(self, config: AzureLLMConfig):
+        self.config = config
+        self.client = openai.AzureOpenAI(
+            azure_endpoint=config.base_url,
+            api_version=config.api_version,
+            api_key=config.api_key,
+        )
+
+    def generate(self, messages: MessageList) -> str:
+        """Generate a response from Azure OpenAI LLM."""
+        response = self.client.chat.completions.create(
+            model=self.config.model_name_or_path,
+            messages=messages,
+            temperature=self.config.temperature,
+            max_tokens=self.config.max_tokens,
+            top_p=self.config.top_p,
+        )
+        logger.info(f"Response from Azure OpenAI: {response.model_dump_json()}")
+        response_content = response.choices[0].message.content
+        if self.config.remove_think_prefix:
+            return remove_thinking_tags(response_content)
+        else:
+            return response_content
+
+    def generate_stream(self, messages: MessageList, **kwargs) -> Generator[str, None, None]:
+        raise NotImplementedError

memos/llms/qwen.py

@@ -0,0 +1,63 @@
+from collections.abc import Generator
+
+from memos.configs.llm import QwenLLMConfig
+from memos.llms.openai import OpenAILLM
+from memos.llms.utils import remove_thinking_tags
+from memos.log import get_logger
+from memos.types import MessageList
+
+
+logger = get_logger(__name__)
+
+
+class QwenLLM(OpenAILLM):
+    """Qwen (DashScope) LLM class via OpenAI-compatible API."""
+
+    def __init__(self, config: QwenLLMConfig):
+        super().__init__(config)
+
+    def generate(self, messages: MessageList) -> str:
+        """Generate a response from Qwen LLM."""
+        response = self.client.chat.completions.create(
+            model=self.config.model_name_or_path,
+            messages=messages,
+            extra_body=self.config.extra_body,
+            temperature=self.config.temperature,
+            max_tokens=self.config.max_tokens,
+            top_p=self.config.top_p,
+        )
+        logger.info(f"Response from Qwen: {response.model_dump_json()}")
+        response_content = response.choices[0].message.content
+        if self.config.remove_think_prefix:
+            return remove_thinking_tags(response_content)
+        else:
+            return response_content
+
+    def generate_stream(self, messages: MessageList, **kwargs) -> Generator[str, None, None]:
+        """Stream response from Qwen LLM."""
+        response = self.client.chat.completions.create(
+            model=self.config.model_name_or_path,
+            messages=messages,
+            stream=True,
+            temperature=self.config.temperature,
+            max_tokens=self.config.max_tokens,
+            top_p=self.config.top_p,
+            extra_body=self.config.extra_body,
+        )
+
+        reasoning_started = False
+        for chunk in response:
+            delta = chunk.choices[0].delta
+
+            # Some models may have separate `reasoning_content` vs `content`
+            # For Qwen (DashScope), likely only `content` is used
+            if hasattr(delta, "reasoning_content") and delta.reasoning_content:
+                if not reasoning_started and not self.config.remove_think_prefix:
+                    yield "<think>"
+                    reasoning_started = True
+                yield delta.reasoning_content
+            elif hasattr(delta, "content") and delta.content:
+                if reasoning_started and not self.config.remove_think_prefix:
+                    yield "</think>"
+                    reasoning_started = False
+                yield delta.content

memos/llms/vllm.py

@@ -0,0 +1,153 @@
+from typing import Any, cast
+
+from memos.configs.llm import VLLMLLMConfig
+from memos.llms.base import BaseLLM
+from memos.llms.utils import remove_thinking_tags
+from memos.log import get_logger
+from memos.types import MessageDict
+
+
+logger = get_logger(__name__)
+
+
+class VLLMLLM(BaseLLM):
+    """
+    VLLM LLM class for connecting to existing vLLM servers.
+    """
+
+    def __init__(self, config: VLLMLLMConfig):
+        """
+        Initialize the VLLM LLM to connect to an existing vLLM server.
+        """
+        self.config = config
+
+        # Initialize OpenAI client for API calls
+        self.client = None
+        api_key = getattr(self.config, "api_key", "dummy")
+        if not api_key:
+            api_key = "dummy"
+
+        import openai
+
+        self.client = openai.Client(
+            api_key=api_key, base_url=getattr(self.config, "api_base", "http://localhost:8088/v1")
+        )
+
+    def build_vllm_kv_cache(self, messages: Any) -> str:
+        """
+        Build a KV cache from chat messages via one vLLM request.
+        Handles str, list[str], and MessageList formats.
+        """
+        # 1. Normalize input to a MessageList
+        processed_messages: list[MessageDict] = []
+        if isinstance(messages, str):
+            processed_messages = [
+                {
+                    "role": "system",
+                    "content": f"Below is some information about the user.\n{messages}",
+                }
+            ]
+        elif isinstance(messages, list):
+            if not messages:
+                pass  # Empty list
+            elif isinstance(messages[0], str):
+                str_content = " ".join(str(msg) for msg in messages)
+                processed_messages = [
+                    {
+                        "role": "system",
+                        "content": f"Below is some information about the user.\n{str_content}",
+                    }
+                ]
+            elif isinstance(messages[0], dict):
+                processed_messages = cast("list[MessageDict]", messages)
+
+        # 2. Convert to prompt for logging/return value.
+        prompt = self._messages_to_prompt(processed_messages)
+
+        if not prompt.strip():
+            raise ValueError("Prompt is empty, cannot build KV cache.")
+
+        # 3. Send request to vLLM server to preload the KV cache
+        if self.client:
+            try:
+                # Use the processed messages for the API call
+                prefill_kwargs = {
+                    "model": self.config.model_name_or_path,
+                    "messages": processed_messages,
+                    "max_tokens": 2,
+                    "temperature": 0.0,
+                    "top_p": 1.0,
+                }
+                self.client.chat.completions.create(**prefill_kwargs)
+                logger.info(f"vLLM KV cache prefill completed for prompt: '{prompt[:100]}...'")
+            except Exception as e:
+                logger.warning(f"Failed to prefill vLLM KV cache: {e}")
+
+        return prompt
+
+    def generate(self, messages: list[MessageDict]) -> str:
+        """
+        Generate a response from the model.
+        """
+        if self.client:
+            return self._generate_with_api_client(messages)
+        else:
+            raise RuntimeError("API client is not available")
+
+    def _generate_with_api_client(self, messages: list[MessageDict]) -> str:
+        """
+        Generate response using vLLM API client.
+        """
+        if self.client:
+            completion_kwargs = {
+                "model": self.config.model_name_or_path,
+                "messages": messages,
+                "temperature": float(getattr(self.config, "temperature", 0.8)),
+                "max_tokens": int(getattr(self.config, "max_tokens", 1024)),
+                "top_p": float(getattr(self.config, "top_p", 0.9)),
+            }
+
+            response = self.client.chat.completions.create(**completion_kwargs)
+            response_text = response.choices[0].message.content or ""
+            logger.info(f"VLLM API response: {response_text}")
+            return (
+                remove_thinking_tags(response_text)
+                if getattr(self.config, "remove_think_prefix", False)
+                else response_text
+            )
+        else:
+            raise RuntimeError("API client is not available")
+
+    def _messages_to_prompt(self, messages: list[MessageDict]) -> str:
+        """
+        Convert messages to prompt string.
+        """
+        prompt_parts = []
+        for msg in messages:
+            role = msg["role"]
+            content = msg["content"]
+            prompt_parts.append(f"{role.capitalize()}: {content}")
+        return "\n".join(prompt_parts)
+
+    def generate_stream(self, messages: list[MessageDict]):
+        """
+        Generate a response from the model using streaming.
+        Yields content chunks as they are received.
+        """
+        if self.client:
+            completion_kwargs = {
+                "model": self.config.model_name_or_path,
+                "messages": messages,
+                "temperature": float(getattr(self.config, "temperature", 0.8)),
+                "max_tokens": int(getattr(self.config, "max_tokens", 1024)),
+                "top_p": float(getattr(self.config, "top_p", 0.9)),
+                "stream": True,  # Enable streaming
+            }
+
+            stream = self.client.chat.completions.create(**completion_kwargs)
+            for chunk in stream:
+                content = chunk.choices[0].delta.content
+                if content:
+                    yield content
+        else:
+            raise RuntimeError("API client is not available")

memos/mem_cube/general.py

@@ -1,11 +1,13 @@
 import os
 
+from typing import Literal
+
 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.mem_cube.utils import download_repo, merge_config_with_default
 from memos.memories.activation.base import BaseActMemory
 from memos.memories.factory import MemoryFactory
 from memos.memories.parametric.base import BaseParaMemory
@@ -37,10 +39,15 @@ class GeneralMemCube(BaseMemCube):
             else None
         )
 
-    def load(self, dir: str) -> None:
+    def load(
+        self, dir: str, memory_types: list[Literal["text_mem", "act_mem", "para_mem"]] | None = None
+    ) -> None:
         """Load memories.
         Args:
             dir (str): The directory containing the memory files.
+            memory_types (list[str], optional): List of memory types to load.
+                If None, loads all available memory types.
+                Options: ["text_mem", "act_mem", "para_mem"]
         """
         loaded_schema = get_json_file_model_schema(os.path.join(dir, self.config.config_filename))
         if loaded_schema != self.config.model_schema:
@@ -48,60 +55,114 @@ class GeneralMemCube(BaseMemCube):
                 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}")
+        # If no specific memory types specified, load all
+        if memory_types is None:
+            memory_types = ["text_mem", "act_mem", "para_mem"]
+
+        # Load specified memory types
+        if "text_mem" in memory_types and self.text_mem:
+            self.text_mem.load(dir)
+            logger.debug(f"Loaded text_mem from {dir}")
+
+        if "act_mem" in memory_types and self.act_mem:
+            self.act_mem.load(dir)
+            logger.info(f"Loaded act_mem from {dir}")
 
-    def dump(self, dir: str) -> None:
+        if "para_mem" in memory_types and self.para_mem:
+            self.para_mem.load(dir)
+            logger.info(f"Loaded para_mem from {dir}")
+
+        logger.info(f"MemCube loaded successfully from {dir} (types: {memory_types})")
+
+    def dump(
+        self, dir: str, memory_types: list[Literal["text_mem", "act_mem", "para_mem"]] | None = None
+    ) -> None:
         """Dump memories.
         Args:
             dir (str): The directory where the memory files will be saved.
+            memory_types (list[str], optional): List of memory types to dump.
+                If None, dumps all available memory types.
+                Options: ["text_mem", "act_mem", "para_mem"]
         """
         if os.path.exists(dir) and os.listdir(dir):
             raise MemCubeError(
                 f"Directory {dir} is not empty. Please provide an empty directory for dumping."
             )
 
+        # Always dump config
         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}")
+        # If no specific memory types specified, dump all
+        if memory_types is None:
+            memory_types = ["text_mem", "act_mem", "para_mem"]
+
+        # Dump specified memory types
+        if "text_mem" in memory_types and self.text_mem:
+            self.text_mem.dump(dir)
+            logger.info(f"Dumped text_mem to {dir}")
+
+        if "act_mem" in memory_types and self.act_mem:
+            self.act_mem.dump(dir)
+            logger.info(f"Dumped act_mem to {dir}")
+
+        if "para_mem" in memory_types and self.para_mem:
+            self.para_mem.dump(dir)
+            logger.info(f"Dumped para_mem to {dir}")
+
+        logger.info(f"MemCube dumped successfully to {dir} (types: {memory_types})")
 
     @staticmethod
-    def init_from_dir(dir: str) -> "GeneralMemCube":
+    def init_from_dir(
+        dir: str,
+        memory_types: list[Literal["text_mem", "act_mem", "para_mem"]] | None = None,
+        default_config: GeneralMemCubeConfig | None = None,
+    ) -> "GeneralMemCube":
         """Create a MemCube instance from a MemCube directory.
 
         Args:
             dir (str): The directory containing the memory files.
+            memory_types (list[str], optional): List of memory types to load.
+                If None, loads all available memory types.
+            default_config (GeneralMemCubeConfig, optional): Default configuration to merge with existing config.
+                If provided, will merge general settings while preserving critical user-specific fields.
 
         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)
+
+        # Merge with default config if provided
+        if default_config is not None:
+            config = merge_config_with_default(config, default_config)
+            logger.info(f"Applied default config to cube {config.cube_id}")
+
         mem_cube = GeneralMemCube(config)
-        mem_cube.load(dir)
+        mem_cube.load(dir, memory_types)
         return mem_cube
 
     @staticmethod
     def init_from_remote_repo(
-        cube_id: str, base_url: str = "https://huggingface.co/datasets"
+        cube_id: str,
+        base_url: str = "https://huggingface.co/datasets",
+        memory_types: list[Literal["text_mem", "act_mem", "para_mem"]] | None = None,
+        default_config: GeneralMemCubeConfig | None = None,
     ) -> "GeneralMemCube":
         """Create a MemCube instance from a remote repository.
 
         Args:
-            repo (str): The repository name.
+            cube_id (str): The repository name.
             base_url (str): The base URL of the remote repository.
+            memory_types (list[str], optional): List of memory types to load.
+                If None, loads all available memory types.
+            default_config (GeneralMemCubeConfig, optional): Default configuration to merge with existing config.
 
         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)
+        return GeneralMemCube.init_from_dir(dir, memory_types, default_config)
 
     @property
     def text_mem(self) -> "BaseTextMemory | None":