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

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).
         """
 
@@ -146,7 +146,7 @@ class BaseGraphDB(ABC):
         """
 
     @abstractmethod
-    def get_by_metadata(self, filters: dict[str, Any]) -> list[str]:
+    def get_by_metadata(self, filters: list[dict[str, Any]]) -> list[str]:
         """
         Retrieve node IDs that match given metadata filters.
 
@@ -162,6 +162,14 @@ class BaseGraphDB(ABC):
             - Can be used for faceted recall or prefiltering before embedding rerank.
         """
 
+    @abstractmethod
+    def get_structure_optimization_candidates(self, scope: str) -> list[dict]:
+        """
+        Find nodes that are likely candidates for structure optimization:
+        - Isolated nodes, nodes with empty background, or nodes with exactly one child.
+        - Plus: the child of any parent node that has exactly one child.
+        """
+
     # Structure Maintenance
     @abstractmethod
     def deduplicate_nodes(self) -> None:

memos/graph_dbs/factory.py

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