build: optimize dependencies management

Author: Ki-Seki

poetry.lock

@@ -12,31 +12,58 @@ keywords = ["memory", "llm", "language model", "memoryOS", "agent"]
 packages = [{include = "memos", from = "src"}]
 
 [tool.poetry.dependencies]
+
+# Core dependencies. These should not be optional.
 python = "^3.10"
 openai = "^1.77.0"
 ollama = "^0.4.8"
-qdrant-client = "^1.14.2"
+tenacity = "^9.1.2"  # Error handling and retrying library
 transformers = "^4.51.3"
-markitdown = {extras = ["docx", "pdf", "pptx", "xls", "xlsx"], version = "^0.1.1"}
-chonkie = "^1.0.7"
-tenacity = "^9.1.2"
-neo4j = "^5.28.1"
-accelerate = "^1.7.0"
 fastapi = {extras = ["all"], version = "^0.115.12"}
-sentence-transformers = "^4.1.0"
-sqlalchemy = "^2.0.41"
-redis = "^6.2.0"
-pika = "^1.3.2"
-schedule = "^1.2.2"
-volcengine-python-sdk = "^4.0.4"
+sqlalchemy = "^2.0.41"  # SQL toolkit
+
+# GeneralTextualMemory dependencies
+qdrant-client = {version = "^1.14.2", optional = true}  # Vector database
+
+# TreeTextualMemory dependencies
+neo4j = {version = "^5.28.1", optional = true}  # Graph database
+schedule = {version = "^1.2.2", optional = true}  # Task scheduling library
+scikit-learn = {version = "^1.7.0", optional = true}  # Machine learning library
+
+# MemScheduler dependencies
+redis = {version = "^6.2.0", optional = true}  # Key-value store
+pika = {version = "^1.3.2", optional = true}  # RabbitMQ client for Python
+
+[tool.poetry.extras]
+general-mem = ["qdrant-client"]
+tree-mem = ["neo4j", "schedule", "scikit-learn"]
+mem-scheduler = ["redis", "pika"]
+all = [
+    "qdrant-client",
+    "neo4j",
+    "schedule",
+    "scikit-learn",
+    "redis",
+    "pika",
+]
 
 [tool.poetry.group.dev]
 optional = false
 
 [tool.poetry.group.dev.dependencies]
+# Core development dependencies
 pre-commit = "^4.2.0"
 ruff = "^0.11.8"
 
+# Hard-to-avoid dependencies
+torch = "^2.0.0"  # CPU-version of PyTorch, only used for testing
+
+# Infrequently used dependencies
+markitdown = {extras = ["docx", "pdf", "pptx", "xls", "xlsx"], version = "^0.1.1"}  # MarkItDown parser for various file formats
+chonkie = "^1.0.7"  # Sentence chunking library
+sentence-transformers = "^4.1.0"  # Text Embedding
+volcengine-python-sdk = "^4.0.4"  # ByteDance's AI Service SDK (namely, Volcano Ark)
+
 [tool.poetry.group.test]
 optional = false
 

pyproject.toml

@@ -1,6 +1,5 @@
-from chonkie import SentenceChunker as ChonkieSentenceChunker
-
 from memos.configs.chunker import SentenceChunkerConfig
+from memos.dependency import require_python_package
 from memos.log import get_logger
 
 from .base import BaseChunker, Chunk
@@ -12,7 +11,14 @@
 class SentenceChunker(BaseChunker):
     """Sentence-based text chunker."""
 
+    @require_python_package(
+        import_name="chonkie",
+        install_command="pip install chonkie",
+        install_link="https://docs.chonkie.ai/python-sdk/getting-started/installation",
+    )
     def __init__(self, config: SentenceChunkerConfig):
+        from chonkie import SentenceChunker as ChonkieSentenceChunker
+
         self.config = config
         self.chunker = ChonkieSentenceChunker(
             tokenizer_or_token_counter=config.tokenizer_or_token_counter,

src/memos/chunkers/sentence_chunker.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

src/memos/dependency.py

@@ -1,11 +1,5 @@
-from volcenginesdkarkruntime import Ark
-from volcenginesdkarkruntime.types.multimodal_embedding import (
-    EmbeddingInputParam,
-    MultimodalEmbeddingContentPartTextParam,
-    MultimodalEmbeddingResponse,
-)
-
 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
 
@@ -16,7 +10,14 @@
 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:
@@ -44,6 +45,10 @@ def embed(self, texts: list[str]) -> list[list[float]]:
         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
@@ -66,8 +71,12 @@ def text_embedding(self, inputs: list[str], chunk_size: int | None = None) -> li
         return embeddings
 
     def multimodal_embeddings(
-        self, inputs: list[EmbeddingInputParam], chunk_size: int | None = None
+        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]] = []
 

src/memos/embedders/ark.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 @@
 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

src/memos/embedders/sentence_transformer.py

@@ -3,10 +3,8 @@
 from datetime import datetime
 from typing import Any, Literal
 
-from neo4j import GraphDatabase
-from neo4j.exceptions import ClientError
-
 from memos.configs.graph_db import Neo4jGraphDBConfig
+from memos.dependency import require_python_package
 from memos.graph_dbs.base import BaseGraphDB
 from memos.log import get_logger
 
@@ -57,6 +55,11 @@ def _prepare_node_metadata(metadata: dict[str, Any]) -> dict[str, Any]:
 class Neo4jGraphDB(BaseGraphDB):
     """Neo4j-based implementation of a graph memory store."""
 
+    @require_python_package(
+        import_name="neo4j",
+        install_command="pip install neo4j",
+        install_link="https://neo4j.com/docs/python-manual/current/install/",
+    )
     def __init__(self, config: Neo4jGraphDBConfig):
         """Neo4j-based implementation of a graph memory store.
 
@@ -75,6 +78,7 @@ def __init__(self, config: Neo4jGraphDBConfig):
             All node queries will enforce `user_name` in WHERE conditions and store it in metadata,
             but it will be removed automatically before returning to external consumers.
         """
+        from neo4j import GraphDatabase
 
         self.config = config
         self.driver = GraphDatabase.driver(config.uri, auth=(config.user, config.password))
@@ -994,6 +998,8 @@ def drop_database(self) -> None:
             )
 
     def _ensure_database_exists(self):
+        from neo4j.exceptions import ClientError
+
         try:
             with self.driver.session(database="system") as session:
                 session.run(f"CREATE DATABASE `{self.db_name}` IF NOT EXISTS")

src/memos/graph_dbs/neo4j.py

@@ -1,6 +1,5 @@
 from collections.abc import Generator
-
-import torch
+from typing import Any
 
 from transformers import (
     AutoModelForCausalLM,
@@ -134,6 +133,8 @@ def _generate_full_stream(self, prompt: str) -> Generator[str, None, None]:
         Yields:
             str: Streaming response chunks.
         """
+        import torch
+
         inputs = self.tokenizer([prompt], return_tensors="pt").to(self.model.device)
 
         # Get generation parameters
@@ -200,6 +201,8 @@ def _generate_with_cache(self, query: str, kv: DynamicCache) -> str:
         Returns:
             str: Model response.
         """
+        import torch
+
         query_ids = self.tokenizer(
             query, return_tensors="pt", add_special_tokens=False
         ).input_ids.to(self.model.device)
@@ -287,10 +290,7 @@ def _generate_with_cache_stream(
 
             generated.append(next_token)
 
-    @torch.no_grad()
-    def _prefill(
-        self, input_ids: torch.Tensor, kv: DynamicCache
-    ) -> tuple[torch.Tensor, DynamicCache]:
+    def _prefill(self, input_ids: Any, kv: DynamicCache) -> tuple[Any, DynamicCache]:
         """
         Forward the model once, returning last-step logits and updated KV cache.
         Args:
@@ -299,22 +299,27 @@ def _prefill(
         Returns:
             tuple[torch.Tensor, DynamicCache]: (last-step logits, updated KV cache)
         """
-        out = self.model(
-            input_ids=input_ids,
-            use_cache=True,
-            past_key_values=kv,
-            return_dict=True,
-        )
+        import torch
+
+        with torch.no_grad():
+            out = self.model(
+                input_ids=input_ids,
+                use_cache=True,
+                past_key_values=kv,
+                return_dict=True,
+            )
         return out.logits[:, -1, :], out.past_key_values
 
-    def _select_next_token(self, logits: torch.Tensor) -> torch.Tensor:
+    def _select_next_token(self, logits: Any) -> Any:
         """
         Select the next token from logits using sampling or argmax, depending on config.
         Args:
             logits (torch.Tensor): Logits for the next token.
         Returns:
             torch.Tensor: Selected token ID(s).
         """
+        import torch
+
         if getattr(self.config, "do_sample", True):
             batch_size, _ = logits.size()
             dummy_ids = torch.zeros((batch_size, 1), dtype=torch.long, device=logits.device)
@@ -323,7 +328,7 @@ def _select_next_token(self, logits: torch.Tensor) -> torch.Tensor:
             return torch.multinomial(probs, num_samples=1)
         return torch.argmax(logits, dim=-1, keepdim=True)
 
-    def _should_stop(self, token: torch.Tensor) -> bool:
+    def _should_stop(self, token: Any) -> bool:
         """
         Check if the given token is the EOS (end-of-sequence) token.
         Args:
@@ -347,6 +352,8 @@ def build_kv_cache(self, messages) -> DynamicCache:
         Returns:
             DynamicCache: The constructed KV cache object.
         """
+        import torch
+
         # Accept multiple input types and convert to standard chat messages
         if isinstance(messages, str):
             messages = [