MemoryOS 2.0.3__py3-none-any.whl

memos/embedders/ark.py

@@ -0,0 +1,95 @@
+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,
+        )
+
+        # Truncate texts if max_tokens is configured
+        texts = self._truncate_texts(texts)
+
+        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/base.py

@@ -0,0 +1,106 @@
+import re
+
+from abc import ABC, abstractmethod
+
+from memos.configs.embedder import BaseEmbedderConfig
+
+
+def _count_tokens_for_embedding(text: str) -> int:
+    """
+    Count tokens in text for embedding truncation.
+    Uses tiktoken if available, otherwise falls back to heuristic.
+
+    Args:
+        text: Text to count tokens for.
+
+    Returns:
+        Number of tokens.
+    """
+    try:
+        import tiktoken
+
+        try:
+            enc = tiktoken.encoding_for_model("gpt-4o-mini")
+        except Exception:
+            enc = tiktoken.get_encoding("cl100k_base")
+        return len(enc.encode(text or "", disallowed_special=()))
+    except Exception:
+        # Heuristic fallback: zh chars ~1 token, others ~1 token per ~4 chars
+        if not text:
+            return 0
+        zh_chars = re.findall(r"[\u4e00-\u9fff]", text)
+        zh = len(zh_chars)
+        rest = len(text) - zh
+        return zh + max(1, rest // 4)
+
+
+def _truncate_text_to_tokens(text: str, max_tokens: int) -> str:
+    """
+    Truncate text to fit within max_tokens limit.
+    Uses binary search to find the optimal truncation point.
+
+    Args:
+        text: Text to truncate.
+        max_tokens: Maximum number of tokens allowed.
+
+    Returns:
+        Truncated text.
+    """
+    if not text or max_tokens is None or max_tokens <= 0:
+        return text
+
+    current_tokens = _count_tokens_for_embedding(text)
+    if current_tokens <= max_tokens:
+        return text
+
+    # Binary search for the right truncation point
+    low, high = 0, len(text)
+    best_text = ""
+
+    while low < high:
+        mid = (low + high + 1) // 2  # Use +1 to avoid infinite loop
+        truncated = text[:mid]
+        tokens = _count_tokens_for_embedding(truncated)
+
+        if tokens <= max_tokens:
+            best_text = truncated
+            low = mid
+        else:
+            high = mid - 1
+
+    return best_text if best_text else text[:1]  # Fallback to at least one character
+
+
+class BaseEmbedder(ABC):
+    """Base class for all Embedding models."""
+
+    @abstractmethod
+    def __init__(self, config: BaseEmbedderConfig):
+        """Initialize the embedding model with the given configuration."""
+        self.config = config
+
+    def _truncate_texts(self, texts: list[str], approx_char_per_token=1.0) -> (list)[str]:
+        """
+        Truncate texts to fit within max_tokens limit if configured.
+
+        Args:
+            texts: List of texts to truncate.
+
+        Returns:
+            List of truncated texts.
+        """
+        if not hasattr(self, "config") or self.config.max_tokens is None:
+            return texts
+        max_tokens = self.config.max_tokens
+
+        truncated = []
+        for t in texts:
+            if len(t) < max_tokens * approx_char_per_token:
+                truncated.append(t)
+            else:
+                truncated.append(t[:max_tokens])
+        return truncated
+
+    @abstractmethod
+    def embed(self, texts: list[str]) -> list[list[float]]:
+        """Generate embeddings for the given texts."""

memos/embedders/factory.py

@@ -0,0 +1,29 @@
+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
+from memos.memos_tools.singleton import singleton_factory
+
+
+class EmbedderFactory(BaseEmbedder):
+    """Factory class for creating embedder instances."""
+
+    backend_to_class: ClassVar[dict[str, Any]] = {
+        "ollama": OllamaEmbedder,
+        "sentence_transformer": SenTranEmbedder,
+        "ark": ArkEmbedder,
+        "universal_api": UniversalAPIEmbedder,
+    }
+
+    @classmethod
+    @singleton_factory()
+    def from_config(cls, config_factory: EmbedderConfigFactory) -> BaseEmbedder:
+        backend = config_factory.backend
+        if backend not in cls.backend_to_class:
+            raise ValueError(f"Invalid backend: {backend}")
+        embedder_class = cls.backend_to_class[backend]
+        return embedder_class(config_factory.config)

memos/embedders/ollama.py

@@ -0,0 +1,77 @@
+from ollama import Client
+
+from memos.configs.embedder import OllamaEmbedderConfig
+from memos.embedders.base import BaseEmbedder
+from memos.log import get_logger
+
+
+logger = get_logger(__name__)
+
+
+class OllamaEmbedder(BaseEmbedder):
+    """Ollama Embedder class."""
+
+    def __init__(self, config: OllamaEmbedderConfig):
+        self.config = config
+        self.api_base = config.api_base
+
+        if self.config.embedding_dims is not None:
+            logger.warning(
+                "Ollama 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 = "nomic-embed-text:latest"
+
+        # Initialize ollama client
+        self.client = Client(host=self.api_base)
+
+        # Ensure the model exists locally
+        self._ensure_model_exists()
+
+    def _list_models(self) -> list[str]:
+        """
+        List all models available in the Ollama client.
+
+        Returns:
+            List of model names.
+        """
+        local_models = self.client.list()["models"]
+        return [model.model for model in local_models]
+
+    def _ensure_model_exists(self):
+        """
+        Ensure the specified model exists locally. If not, pull it from Ollama.
+        """
+        try:
+            local_models = self._list_models()
+            if self.config.model_name_or_path not in local_models:
+                logger.warning(
+                    f"Model {self.config.model_name_or_path} not found locally. Pulling from Ollama..."
+                )
+                self.client.pull(self.config.model_name_or_path)
+        except Exception as e:
+            logger.warning(f"Could not verify model existence: {e}")
+
+    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.
+        """
+        # Truncate texts if max_tokens is configured
+        texts = self._truncate_texts(texts)
+
+        response = self.client.embed(
+            model=self.config.model_name_or_path,
+            input=texts,
+        )
+        return response.embeddings

memos/embedders/sentence_transformer.py

@@ -0,0 +1,49 @@
+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
+
+
+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
+        )
+
+        if self.config.embedding_dims is not None:
+            logger.warning(
+                "SentenceTransformer does not support specifying embedding dimensions directly. "
+                "The embedding dimension is determined by the model."
+                "`embedding_dims` will be ignored."
+            )
+            # Get embedding dimensions from the model
+            self.config.embedding_dims = self.model.get_sentence_embedding_dimension()
+
+    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.
+        """
+        # Truncate texts if max_tokens is configured
+        texts = self._truncate_texts(texts)
+
+        embeddings = self.model.encode(texts, convert_to_numpy=True)
+        return embeddings.tolist()

memos/embedders/universal_api.py

@@ -0,0 +1,51 @@
+from openai import AzureOpenAI as AzureClient
+from openai import OpenAI as OpenAIClient
+
+from memos.configs.embedder import UniversalAPIEmbedderConfig
+from memos.embedders.base import BaseEmbedder
+from memos.log import get_logger
+from memos.utils import timed_with_status
+
+
+logger = get_logger(__name__)
+
+
+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,
+                default_headers=config.headers_extra if config.headers_extra else None,
+            )
+        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"Embeddings unsupported provider: {self.provider}")
+
+    @timed_with_status(
+        log_prefix="model_timed_embedding",
+        log_extra_args={"model_name_or_path": "text-embedding-3-large"},
+    )
+    def embed(self, texts: list[str]) -> list[list[float]]:
+        # Truncate texts if max_tokens is configured
+        texts = self._truncate_texts(texts)
+
+        if self.provider == "openai" or self.provider == "azure":
+            try:
+                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]
+            except Exception as e:
+                raise Exception(f"Embeddings request ended with error: {e}") from e
+        else:
+            raise ValueError(f"Embeddings unsupported provider: {self.provider}")

memos/exceptions.py

@@ -0,0 +1,30 @@
+"""Custom exceptions for the MemOS library.
+
+This module defines all custom exceptions used throughout the MemOS project.
+All exceptions inherit from a base MemOSError class to provide a consistent
+error handling interface.
+"""
+
+
+class MemOSError(Exception): ...
+
+
+class ConfigurationError(MemOSError): ...
+
+
+class MemoryError(MemOSError): ...
+
+
+class MemCubeError(MemOSError): ...
+
+
+class VectorDBError(MemOSError): ...
+
+
+class LLMError(MemOSError): ...
+
+
+class EmbedderError(MemOSError): ...
+
+
+class ParserError(MemOSError): ...