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

memos/configs/llm.py

@@ -27,6 +27,40 @@ class OpenAILLMConfig(BaseLLMConfig):
     extra_body: Any = Field(default=None, description="extra body")
 
 
+class QwenLLMConfig(BaseLLMConfig):
+    api_key: str = Field(..., description="API key for DashScope (Qwen)")
+    api_base: str = Field(
+        default="https://dashscope-intl.aliyuncs.com/compatible-mode/v1",
+        description="Base URL for Qwen OpenAI-compatible API",
+    )
+    extra_body: Any = Field(default=None, description="extra body")
+    model_name_or_path: str = Field(..., description="Model name for Qwen, e.g., 'qwen-plus'")
+
+
+class DeepSeekLLMConfig(BaseLLMConfig):
+    api_key: str = Field(..., description="API key for DeepSeek")
+    api_base: str = Field(
+        default="https://api.deepseek.com",
+        description="Base URL for DeepSeek OpenAI-compatible API",
+    )
+    extra_body: Any = Field(default=None, description="Extra options for API")
+    model_name_or_path: str = Field(
+        ..., description="Model name: 'deepseek-chat' or 'deepseek-reasoner'"
+    )
+
+
+class AzureLLMConfig(BaseLLMConfig):
+    base_url: str = Field(
+        default="https://api.openai.azure.com/",
+        description="Base URL for Azure OpenAI API",
+    )
+    api_version: str = Field(
+        default="2024-03-01-preview",
+        description="API version for Azure OpenAI",
+    )
+    api_key: str = Field(..., description="API key for Azure OpenAI")
+
+
 class OllamaLLMConfig(BaseLLMConfig):
     api_base: str = Field(
         default="http://localhost:11434",
@@ -45,6 +79,14 @@ class HFLLMConfig(BaseLLMConfig):
     )
 
 
+class VLLMLLMConfig(BaseLLMConfig):
+    api_key: str = Field(default="", description="API key for vLLM (optional for local server)")
+    api_base: str = Field(
+        default="http://localhost:8088/v1",
+        description="Base URL for vLLM API",
+    )
+
+
 class LLMConfigFactory(BaseConfig):
     """Factory class for creating LLM configurations."""
 
@@ -54,7 +96,12 @@ class LLMConfigFactory(BaseConfig):
     backend_to_class: ClassVar[dict[str, Any]] = {
         "openai": OpenAILLMConfig,
         "ollama": OllamaLLMConfig,
+        "azure": AzureLLMConfig,
         "huggingface": HFLLMConfig,
+        "vllm": VLLMLLMConfig,
+        "huggingface_singleton": HFLLMConfig,  # Add singleton support
+        "qwen": QwenLLMConfig,
+        "deepseek": DeepSeekLLMConfig,
     }
 
     @field_validator("backend")

memos/configs/mem_cube.py

@@ -70,7 +70,7 @@ class GeneralMemCubeConfig(BaseMemCubeConfig):
     @classmethod
     def validate_act_mem(cls, act_mem: MemoryConfigFactory) -> MemoryConfigFactory:
         """Validate the act_mem field."""
-        allowed_backends = ["kv_cache", "uninitialized"]
+        allowed_backends = ["kv_cache", "vllm_kv_cache", "uninitialized"]
         if act_mem.backend not in allowed_backends:
             raise ConfigurationError(
                 f"GeneralMemCubeConfig requires act_mem backend to be one of {allowed_backends}, got '{act_mem.backend}'"

memos/configs/mem_scheduler.py

@@ -1,13 +1,17 @@
+import os
+
+from pathlib import Path
 from typing import Any, ClassVar
 
 from pydantic import ConfigDict, Field, field_validator, model_validator
 
 from memos.configs.base import BaseConfig
 from memos.mem_scheduler.modules.schemas import (
+    BASE_DIR,
     DEFAULT_ACT_MEM_DUMP_PATH,
-    DEFAULT_ACTIVATION_MEM_SIZE,
     DEFAULT_CONSUME_INTERVAL_SECONDS,
     DEFAULT_THREAD__POOL_MAX_WORKERS,
+    DictConversionMixin,
 )
 
 
@@ -33,6 +37,10 @@ class BaseSchedulerConfig(BaseConfig):
         le=60,
         description=f"Interval for consuming messages from queue in seconds (default: {DEFAULT_CONSUME_INTERVAL_SECONDS})",
     )
+    auth_config_path: str | None = Field(
+        default=None,
+        description="Path to the authentication configuration file containing private credentials",
+    )
 
 
 class GeneralSchedulerConfig(BaseSchedulerConfig):
@@ -42,14 +50,13 @@ class GeneralSchedulerConfig(BaseSchedulerConfig):
     context_window_size: int | None = Field(
         default=5, description="Size of the context window for conversation history"
     )
-    activation_mem_size: int | None = Field(
-        default=DEFAULT_ACTIVATION_MEM_SIZE,  # Assuming DEFAULT_ACTIVATION_MEM_SIZE is 1000
-        description="Maximum size of the activation memory",
-    )
     act_mem_dump_path: str | None = Field(
         default=DEFAULT_ACT_MEM_DUMP_PATH,  # Replace with DEFAULT_ACT_MEM_DUMP_PATH
         description="File path for dumping activation memory",
     )
+    enable_act_memory_update: bool = Field(
+        default=False, description="Whether to enable automatic activation memory updates"
+    )
 
 
 class SchedulerConfigFactory(BaseConfig):
@@ -76,3 +83,82 @@ class SchedulerConfigFactory(BaseConfig):
         config_class = self.backend_to_class[self.backend]
         self.config = config_class(**self.config)
         return self
+
+
+# ************************* Auth *************************
+class RabbitMQConfig(
+    BaseConfig,
+):
+    host_name: str = Field(default="", description="Endpoint for RabbitMQ instance access")
+    user_name: str = Field(default="", description="Static username for RabbitMQ instance")
+    password: str = Field(default="", description="Password for the static username")
+    virtual_host: str = Field(default="", description="Vhost name for RabbitMQ instance")
+    erase_on_connect: bool = Field(
+        default=True, description="Whether to clear connection state or buffers upon connecting"
+    )
+    port: int = Field(
+        default=5672,
+        description="Port number for RabbitMQ instance access",
+        ge=1,  # Port must be >= 1
+        le=65535,  # Port must be <= 65535
+    )
+
+
+class GraphDBAuthConfig(BaseConfig):
+    uri: str = Field(default="localhost", description="URI for graph database access")
+
+
+class OpenAIConfig(BaseConfig):
+    api_key: str = Field(default="", description="API key for OpenAI service")
+    base_url: str = Field(default="", description="Base URL for API endpoint")
+    default_model: str = Field(default="", description="Default model to use")
+
+
+class AuthConfig(BaseConfig, DictConversionMixin):
+    rabbitmq: RabbitMQConfig
+    openai: OpenAIConfig
+    graph_db: GraphDBAuthConfig
+    default_config_path: ClassVar[str] = (
+        f"{BASE_DIR}/examples/data/config/mem_scheduler/scheduler_auth.yaml"
+    )
+
+    @classmethod
+    def from_local_yaml(cls, config_path: str | None = None) -> "AuthConfig":
+        """
+        Load configuration from YAML file
+
+        Args:
+            config_path: Path to YAML configuration file
+
+        Returns:
+            AuthConfig instance
+
+        Raises:
+            FileNotFoundError: If config file doesn't exist
+            ValueError: If YAML parsing or validation fails
+        """
+
+        if config_path is None:
+            config_path = cls.default_config_path
+
+        # Check file exists
+        if not Path(config_path).exists():
+            raise FileNotFoundError(f"Config file not found: {config_path}")
+
+        return cls.from_yaml_file(yaml_path=config_path)
+
+    def set_openai_config_to_environment(self):
+        # Set environment variables
+        os.environ["OPENAI_API_KEY"] = self.openai.api_key
+        os.environ["OPENAI_BASE_URL"] = self.openai.base_url
+        os.environ["MODEL"] = self.openai.default_model
+
+    @classmethod
+    def default_config_exists(cls) -> bool:
+        """
+        Check if the default configuration file exists.
+
+        Returns:
+            bool: True if the default config file exists, False otherwise
+        """
+        return Path(cls.default_config_path).exists()

memos/configs/memory.py

@@ -52,9 +52,9 @@ class KVCacheMemoryConfig(BaseActMemoryConfig):
     @classmethod
     def validate_extractor_llm(cls, extractor_llm: LLMConfigFactory) -> LLMConfigFactory:
         """Validate the extractor_llm field."""
-        if extractor_llm.backend != "huggingface":
+        if extractor_llm.backend not in ["huggingface", "huggingface_singleton", "vllm"]:
             raise ConfigurationError(
-                f"KVCacheMemoryConfig requires extractor_llm backend to be 'huggingface', got '{extractor_llm.backend}'"
+                f"KVCacheMemoryConfig requires extractor_llm backend to be 'huggingface' or 'huggingface_singleton', got '{extractor_llm.backend}'"
             )
         return extractor_llm
 
@@ -84,9 +84,9 @@ class LoRAMemoryConfig(BaseParaMemoryConfig):
     @classmethod
     def validate_extractor_llm(cls, extractor_llm: LLMConfigFactory) -> LLMConfigFactory:
         """Validate the extractor_llm field."""
-        if extractor_llm.backend not in ["huggingface"]:
+        if extractor_llm.backend not in ["huggingface", "huggingface_singleton"]:
             raise ConfigurationError(
-                f"LoRAMemoryConfig requires extractor_llm backend to be 'huggingface', got '{extractor_llm.backend}'"
+                f"LoRAMemoryConfig requires extractor_llm backend to be 'huggingface' or 'huggingface_singleton', got '{extractor_llm.backend}'"
             )
         return extractor_llm
 
@@ -181,6 +181,7 @@ class MemoryConfigFactory(BaseConfig):
         "general_text": GeneralTextMemoryConfig,
         "tree_text": TreeTextMemoryConfig,
         "kv_cache": KVCacheMemoryConfig,
+        "vllm_kv_cache": KVCacheMemoryConfig,  # Use same config as kv_cache
         "lora": LoRAMemoryConfig,
         "uninitialized": UninitializedMemoryConfig,
     }

memos/dependency.py

@@ -0,0 +1,52 @@
+"""
+This utility provides tools for managing dependencies in MemOS.
+"""
+
+import functools
+import importlib
+
+
+def require_python_package(
+    import_name: str, install_command: str | None = None, install_link: str | None = None
+):
+    """Check if a package is available and provide installation hints on import failure.
+
+    Args:
+        import_name (str): The top-level importable module name a package provides.
+        install_command (str, optional): Installation command.
+        install_link (str, optional): URL link to installation guide.
+
+    Returns:
+        Callable: A decorator function that wraps the target function with package availability check.
+
+    Raises:
+        ImportError: When the specified package is not available, with installation
+            instructions included in the error message.
+
+    Example:
+        >>> @require_python_package(
+        ...     import_name='faiss',
+        ...     install_command='pip install faiss-cpu',
+        ...     install_link='https://github.com/facebookresearch/faiss/blob/main/INSTALL.md'
+        ... )
+        ... def create_faiss_index():
+        ...     from faiss import IndexFlatL2  # Actual import in function
+        ...     return IndexFlatL2(128)
+    """
+
+    def decorator(func):
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            try:
+                importlib.import_module(import_name)
+            except ImportError:
+                error_msg = f"Missing required module - '{import_name}'\n"
+                error_msg += f"💡 Install command: {install_command}\n" if install_command else ""
+                error_msg += f"💡 Install guide:   {install_link}\n" if install_link else ""
+
+                raise ImportError(error_msg) from None
+            return func(*args, **kwargs)
+
+        return wrapper
+
+    return decorator

memos/embedders/ark.py

@@ -0,0 +1,92 @@
+from memos.configs.embedder import ArkEmbedderConfig
+from memos.dependency import require_python_package
+from memos.embedders.base import BaseEmbedder
+from memos.log import get_logger
+
+
+logger = get_logger(__name__)
+
+
+class ArkEmbedder(BaseEmbedder):
+    """Ark Embedder class."""
+
+    @require_python_package(
+        import_name="volcenginesdkarkruntime",
+        install_command="pip install 'volcengine-python-sdk[ark]'",
+        install_link="https://www.volcengine.com/docs/82379/1541595",
+    )
+    def __init__(self, config: ArkEmbedderConfig):
+        from volcenginesdkarkruntime import Ark
+
+        self.config = config
+
+        if self.config.embedding_dims is not None:
+            logger.warning(
+                "Ark does not support specifying embedding dimensions. "
+                "The embedding dimensions is determined by the model."
+                "`embedding_dims` will be set to None."
+            )
+            self.config.embedding_dims = None
+
+        # Default model if not specified
+        if not self.config.model_name_or_path:
+            self.config.model_name_or_path = "doubao-embedding-vision-250615"
+
+        # Initialize ark client
+        self.client = Ark(api_key=self.config.api_key, base_url=self.config.api_base)
+
+    def embed(self, texts: list[str]) -> list[list[float]]:
+        """
+        Generate embeddings for the given texts.
+
+        Args:
+            texts: List of texts to embed.
+
+        Returns:
+            List of embeddings, each represented as a list of floats.
+        """
+        from volcenginesdkarkruntime.types.multimodal_embedding import (
+            MultimodalEmbeddingContentPartTextParam,
+        )
+
+        if self.config.multi_modal:
+            texts_input = [
+                MultimodalEmbeddingContentPartTextParam(text=text, type="text") for text in texts
+            ]
+            return self.multimodal_embeddings(inputs=texts_input, chunk_size=self.config.chunk_size)
+        return self.text_embedding(texts, chunk_size=self.config.chunk_size)
+
+    def text_embedding(self, inputs: list[str], chunk_size: int | None = None) -> list[list[float]]:
+        chunk_size_ = chunk_size or self.config.chunk_size
+        embeddings: list[list[float]] = []
+        for i in range(0, len(inputs), chunk_size_):
+            response = self.client.embeddings.create(
+                model=self.config.model_name_or_path,
+                input=inputs[i : i + chunk_size_],
+            )
+
+            data = [response.data] if isinstance(response.data, dict) else response.data
+            embeddings.extend(r.embedding for r in data)
+
+        return embeddings
+
+    def multimodal_embeddings(
+        self, inputs: list, chunk_size: int | None = None
+    ) -> list[list[float]]:
+        from volcenginesdkarkruntime.types.multimodal_embedding import (
+            MultimodalEmbeddingResponse,  # noqa: TC002
+        )
+
+        chunk_size_ = chunk_size or self.config.chunk_size
+        embeddings: list[list[float]] = []
+
+        for i in range(0, len(inputs), chunk_size_):
+            response: MultimodalEmbeddingResponse = self.client.multimodal_embeddings.create(
+                model=self.config.model_name_or_path,
+                input=inputs[i : i + chunk_size_],
+            )
+
+            data = [response.data] if isinstance(response.data, dict) else response.data
+            embeddings.extend(r["embedding"] for r in data)
+
+        return embeddings

memos/embedders/factory.py

@@ -1,9 +1,11 @@
 from typing import Any, ClassVar
 
 from memos.configs.embedder import EmbedderConfigFactory
+from memos.embedders.ark import ArkEmbedder
 from memos.embedders.base import BaseEmbedder
 from memos.embedders.ollama import OllamaEmbedder
 from memos.embedders.sentence_transformer import SenTranEmbedder
+from memos.embedders.universal_api import UniversalAPIEmbedder
 
 
 class EmbedderFactory(BaseEmbedder):
@@ -12,6 +14,8 @@ class EmbedderFactory(BaseEmbedder):
     backend_to_class: ClassVar[dict[str, Any]] = {
         "ollama": OllamaEmbedder,
         "sentence_transformer": SenTranEmbedder,
+        "ark": ArkEmbedder,
+        "universal_api": UniversalAPIEmbedder,
     }
 
     @classmethod

memos/embedders/sentence_transformer.py

@@ -1,6 +1,5 @@
-from sentence_transformers import SentenceTransformer
-
 from memos.configs.embedder import SenTranEmbedderConfig
+from memos.dependency import require_python_package
 from memos.embedders.base import BaseEmbedder
 from memos.log import get_logger
 
@@ -11,7 +10,14 @@ logger = get_logger(__name__)
 class SenTranEmbedder(BaseEmbedder):
     """Sentence Transformer Embedder class."""
 
+    @require_python_package(
+        import_name="sentence_transformers",
+        install_command="pip install sentence-transformers",
+        install_link="https://www.sbert.net/docs/installation.html",
+    )
     def __init__(self, config: SenTranEmbedderConfig):
+        from sentence_transformers import SentenceTransformer
+
         self.config = config
         self.model = SentenceTransformer(
             self.config.model_name_or_path, trust_remote_code=self.config.trust_remote_code

memos/embedders/universal_api.py

@@ -0,0 +1,32 @@
+from openai import AzureOpenAI as AzureClient
+from openai import OpenAI as OpenAIClient
+
+from memos.configs.embedder import UniversalAPIEmbedderConfig
+from memos.embedders.base import BaseEmbedder
+
+
+class UniversalAPIEmbedder(BaseEmbedder):
+    def __init__(self, config: UniversalAPIEmbedderConfig):
+        self.provider = config.provider
+        self.config = config
+
+        if self.provider == "openai":
+            self.client = OpenAIClient(api_key=config.api_key, base_url=config.base_url)
+        elif self.provider == "azure":
+            self.client = AzureClient(
+                azure_endpoint=config.base_url,
+                api_version="2024-03-01-preview",
+                api_key=config.api_key,
+            )
+        else:
+            raise ValueError(f"Unsupported provider: {self.provider}")
+
+    def embed(self, texts: list[str]) -> list[list[float]]:
+        if self.provider == "openai" or self.provider == "azure":
+            response = self.client.embeddings.create(
+                model=getattr(self.config, "model_name_or_path", "text-embedding-3-large"),
+                input=texts,
+            )
+            return [r.embedding for r in response.data]
+        else:
+            raise ValueError(f"Unsupported provider: {self.provider}")

memos/graph_dbs/base.py

@@ -9,12 +9,12 @@ class BaseGraphDB(ABC):
 
     # Node (Memory) Management
     @abstractmethod
-    def add_node(self, id: str, content: str, metadata: dict[str, Any]) -> None:
+    def add_node(self, id: str, memory: str, metadata: dict[str, Any]) -> None:
         """
         Add a memory node to the graph.
         Args:
             id: Unique identifier for the memory node.
-            content: Raw memory content (e.g., text).
+            memory: Raw memory content (e.g., text).
             metadata: Dictionary of metadata (e.g., timestamp, tags, source).
         """
 

memos/graph_dbs/factory.py

@@ -3,6 +3,7 @@ from typing import Any, ClassVar
 from memos.configs.graph_db import GraphDBConfigFactory
 from memos.graph_dbs.base import BaseGraphDB
 from memos.graph_dbs.neo4j import Neo4jGraphDB
+from memos.graph_dbs.neo4j_community import Neo4jCommunityGraphDB
 
 
 class GraphStoreFactory(BaseGraphDB):
@@ -10,6 +11,7 @@ class GraphStoreFactory(BaseGraphDB):
 
     backend_to_class: ClassVar[dict[str, Any]] = {
         "neo4j": Neo4jGraphDB,
+        "neo4j-community": Neo4jCommunityGraphDB,
     }
 
     @classmethod