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

memos/hello_world.py

@@ -0,0 +1,97 @@
+from memos import log
+
+
+logger = log.get_logger(__name__)
+
+
+def memos_hello_world() -> str:
+    logger.info("memos_hello_world function called.")
+    return "Hello world from memos!"
+
+
+def memos_chend_hello_world() -> str:
+    logger.info("memos_chend_hello_world function called.")
+    return "Hello world from memos-chend!"
+
+
+def memos_wanghy_hello_world() -> str:
+    logger.info("memos_wanghy_hello_world function called.")
+    return "Hello world from memos-wanghy!"
+
+
+def memos_niusm_hello_world() -> str:
+    logger.info("memos_niusm_hello_world function called.")
+    return "Hello world from memos-niusm!"
+
+
+def memos_huojh_hello_world(arr: list) -> list:
+    logger.info("memos_huojh_hello_world function called.")
+    if len(arr) <= 1:
+        return arr
+    else:
+        pivot = arr[0]
+        left = [x for x in arr[1:] if x < pivot]
+        right = [x for x in arr[1:] if x >= pivot]
+        return [*memos_huojh_hello_world(left), pivot, *memos_huojh_hello_world(right)]
+
+
+def memos_dany_hello_world(para_1: int, para_2: str) -> str:
+    logger.info(f"logger.info: para_1 is {para_1}")
+    logger.debug(f"logger.debug: para_2 is {para_2}")
+    return f"return_value_{para_1}"
+
+
+def memos_wangyzh_hello_world() -> str:
+    logger.info("memos_wangyzh_hello_world function called.")
+    return "Hello world from memos-wangyzh!"
+
+
+def memos_zhaojihao_hello_world() -> str:
+    logger.info("memos_zhaojihao_hello_world function called.")
+    return "Hello world from memos-zhaojihao!"
+
+
+def memos_yuqingchen_hello_world() -> str:
+    logger.info("memos_yuqingchen_hello_world function called.")
+    return "Hello world from memos-yuqingchen!"
+
+
+def memos_chentang_hello_world(user_id: str = "locomo_exp_user_1", version: str = "default"):
+    import os
+
+    from memos.configs.memory import MemoryConfigFactory
+    from memos.memories.factory import MemoryFactory
+
+    config = MemoryConfigFactory(
+        backend="general_text",
+        config={
+            "extractor_llm": {
+                "backend": "openai",
+                "config": {
+                    "model_name_or_path": os.getenv("MODEL"),
+                    "temperature": 0,
+                    "max_tokens": 8192,
+                    "api_key": os.getenv("OPENAI_API_KEY"),
+                    "api_base": os.getenv("OPENAI_BASE_URL"),
+                },
+            },
+            "vector_db": {
+                "backend": "qdrant",
+                "config": {
+                    "path": f"outputs/locomo/memos-{version}/storages/{user_id}/qdrant",
+                    "collection_name": "test_textual_memory",
+                    "distance_metric": "cosine",
+                    "vector_dimension": 768,  # nomic-embed-text model's embedding dimension is 768
+                },
+            },
+            "embedder": {
+                "backend": "ollama",
+                "config": {
+                    "model_name_or_path": os.getenv("EMBEDDING_MODEL"),
+                },
+            },
+        },
+    )
+    memory = MemoryFactory.from_config(config)
+
+    return memory

memos/llms/base.py

@@ -0,0 +1,16 @@
+from abc import ABC, abstractmethod
+
+from memos.configs.llm import BaseLLMConfig
+from memos.types import MessageList
+
+
+class BaseLLM(ABC):
+    """Base class for all LLMs."""
+
+    @abstractmethod
+    def __init__(self, config: BaseLLMConfig):
+        """Initialize the LLM with the given configuration."""
+
+    @abstractmethod
+    def generate(self, messages: MessageList, **kwargs) -> str:
+        """Generate a response from the LLM."""

memos/llms/factory.py

@@ -0,0 +1,25 @@
+from typing import Any, ClassVar
+
+from memos.configs.llm import LLMConfigFactory
+from memos.llms.base import BaseLLM
+from memos.llms.hf import HFLLM
+from memos.llms.ollama import OllamaLLM
+from memos.llms.openai import OpenAILLM
+
+
+class LLMFactory(BaseLLM):
+    """Factory class for creating LLM instances."""
+
+    backend_to_class: ClassVar[dict[str, Any]] = {
+        "openai": OpenAILLM,
+        "ollama": OllamaLLM,
+        "huggingface": HFLLM,
+    }
+
+    @classmethod
+    def from_config(cls, config_factory: LLMConfigFactory) -> BaseLLM:
+        backend = config_factory.backend
+        if backend not in cls.backend_to_class:
+            raise ValueError(f"Invalid backend: {backend}")
+        llm_class = cls.backend_to_class[backend]
+        return llm_class(config_factory.config)

memos/llms/hf.py

@@ -0,0 +1,231 @@
+import torch
+
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    DynamicCache,
+    LogitsProcessorList,
+    TemperatureLogitsWarper,
+    TopKLogitsWarper,
+    TopPLogitsWarper,
+)
+
+from memos.configs.llm import HFLLMConfig
+from memos.llms.base import BaseLLM
+from memos.llms.utils import remove_thinking_tags
+from memos.log import get_logger
+from memos.types import MessageList
+
+
+logger = get_logger(__name__)
+
+
+class HFLLM(BaseLLM):
+    """
+    HFLLM: Transformers LLM class supporting cache-augmented generation (CAG) and sampling.
+    """
+
+    def __init__(self, config: HFLLMConfig):
+        """
+        Initialize the HFLLM model and tokenizer, and set up logits processors for sampling.
+        """
+        self.config = config
+
+        # Default model if not specified
+        if not self.config.model_name_or_path:
+            self.config.model_name_or_path = "Qwen/Qwen3-1.7B"
+
+        # Initialize hf model
+        self.model = AutoModelForCausalLM.from_pretrained(
+            self.config.model_name_or_path, torch_dtype="auto", device_map="auto"
+        )
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            self.config.model_name_or_path, use_fast=True
+        )
+
+        # Logits processors for sampling
+        processors = []
+        if getattr(self.config, "temperature", 1.0) != 1.0:
+            processors.append(TemperatureLogitsWarper(self.config.temperature))
+        if getattr(self.config, "top_k", 0) > 0:
+            processors.append(TopKLogitsWarper(self.config.top_k))
+        if 0.0 < getattr(self.config, "top_p", 1.0) < 1.0:
+            processors.append(TopPLogitsWarper(self.config.top_p))
+        self.logits_processors = LogitsProcessorList(processors)
+
+    def generate(self, messages: MessageList, past_key_values: DynamicCache | None = None):
+        """
+        Generate a response from the model. If past_key_values is provided, use cache-augmented generation.
+        Args:
+            messages (MessageList): Chat messages for prompt construction.
+            past_key_values (DynamicCache | None): Optional KV cache for fast generation.
+        Returns:
+            str: Model response.
+        """
+        prompt = self.tokenizer.apply_chat_template(
+            messages, tokenize=False, add_generation_prompt=self.config.add_generation_prompt
+        )
+        logger.info(f"HFLLM prompt: {prompt}")
+        if past_key_values is None:
+            return self._generate_full(prompt)
+        else:
+            return self._generate_with_cache(prompt, past_key_values)
+
+    def _generate_full(self, prompt: str) -> str:
+        """
+        Generate output from scratch using the full prompt.
+        Args:
+            prompt (str): The input prompt string.
+        Returns:
+            str: Model response.
+        """
+        inputs = self.tokenizer([prompt], return_tensors="pt").to(self.model.device)
+        gen_kwargs = {
+            "max_new_tokens": getattr(self.config, "max_tokens", 128),
+            "do_sample": getattr(self.config, "do_sample", True),
+        }
+        if self.config.do_sample:
+            gen_kwargs["temperature"] = self.config.temperature
+            gen_kwargs["top_k"] = self.config.top_k
+            gen_kwargs["top_p"] = self.config.top_p
+        gen_ids = self.model.generate(
+            **inputs,
+            **gen_kwargs,
+        )
+        new_ids = [
+            out_ids[len(src_ids) :]
+            for src_ids, out_ids in zip(inputs.input_ids, gen_ids, strict=False)
+        ]
+        response = self.tokenizer.batch_decode(new_ids, skip_special_tokens=True)[0]
+        logger.info(f"Full-gen raw response: {response}")
+        return (
+            remove_thinking_tags(response)
+            if getattr(self.config, "remove_think_prefix", False)
+            else response
+        )
+
+    def _generate_with_cache(self, query: str, kv: DynamicCache) -> str:
+        """
+        Generate output incrementally using an existing KV cache.
+        Args:
+            query (str): The new user query string.
+            kv (DynamicCache): The prefilled KV cache.
+        Returns:
+            str: Model response.
+        """
+        query_ids = self.tokenizer(
+            query, return_tensors="pt", add_special_tokens=False
+        ).input_ids.to(self.model.device)
+        logits, kv = self._prefill(query_ids, kv)
+        next_token = self._select_next_token(logits)
+        generated = [next_token]
+        for _ in range(getattr(self.config, "max_tokens", 128) - 1):
+            if self._should_stop(next_token):
+                break
+            logits, kv = self._prefill(next_token, kv)
+            next_token = self._select_next_token(logits)
+            generated.append(next_token)
+        if generated:
+            concat = torch.cat(generated, dim=-1)
+            response = self.tokenizer.decode(concat[0], skip_special_tokens=True)
+        else:
+            response = ""
+        logger.info(f"Cache-gen raw response: {response}")
+        return (
+            remove_thinking_tags(response)
+            if getattr(self.config, "remove_think_prefix", False)
+            else response
+        )
+
+    @torch.no_grad()
+    def _prefill(
+        self, input_ids: torch.Tensor, kv: DynamicCache
+    ) -> tuple[torch.Tensor, DynamicCache]:
+        """
+        Forward the model once, returning last-step logits and updated KV cache.
+        Args:
+            input_ids (torch.Tensor): Input token IDs.
+            kv (DynamicCache): Existing KV cache.
+        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,
+        )
+        return out.logits[:, -1, :], out.past_key_values
+
+    def _select_next_token(self, logits: torch.Tensor) -> torch.Tensor:
+        """
+        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).
+        """
+        if getattr(self.config, "do_sample", True):
+            batch_size, _ = logits.size()
+            dummy_ids = torch.zeros((batch_size, 1), dtype=torch.long, device=logits.device)
+            filtered = self.logits_processors(dummy_ids, logits)
+            probs = torch.softmax(filtered, dim=-1)
+            return torch.multinomial(probs, num_samples=1)
+        return torch.argmax(logits, dim=-1, keepdim=True)
+
+    def _should_stop(self, token: torch.Tensor) -> bool:
+        """
+        Check if the given token is the EOS (end-of-sequence) token.
+        Args:
+            token (torch.Tensor): Token ID to check.
+        Returns:
+            bool: True if token is EOS, else False.
+        """
+        eos_id = self.tokenizer.eos_token_id
+        return eos_id is not None and token.item() == eos_id
+
+    def build_kv_cache(self, messages) -> DynamicCache:
+        """
+        Build a KV cache from chat messages via one forward pass.
+        Supports the following input types:
+            - str: Used as a system prompt.
+            - list[str]: Concatenated and used as a system prompt.
+            - list[dict]: Used directly as chat messages.
+        The messages are always converted to a standard chat template.
+        Raises:
+            ValueError: If the resulting prompt is empty after template processing.
+        Returns:
+            DynamicCache: The constructed KV cache object.
+        """
+        # Accept multiple input types and convert to standard chat messages
+        if isinstance(messages, str):
+            messages = [
+                {
+                    "role": "system",
+                    "content": f"Below is some information about the user.\n{messages}",
+                }
+            ]
+        elif isinstance(messages, list) and messages and isinstance(messages[0], str):
+            messages = [
+                {
+                    "role": "system",
+                    "content": f"Below is some information about the user.\n{' '.join(messages)}",
+                }
+            ]
+        prompt = self.tokenizer.apply_chat_template(
+            messages, tokenize=False, add_generation_prompt=False
+        )
+        inputs = self.tokenizer(prompt, return_tensors="pt")
+        inputs["input_ids"] = inputs["input_ids"].to(self.model.device, dtype=torch.long)
+        seq_len = inputs["input_ids"].size(-1)
+        if seq_len == 0:
+            raise ValueError(
+                "Prompt after chat template is empty, cannot build KV cache. Check your messages input."
+            )
+        kv = DynamicCache()
+        with torch.no_grad():
+            self.model(**inputs, use_cache=True, past_key_values=kv)
+        for i, (k, v) in enumerate(zip(kv.key_cache, kv.value_cache, strict=False)):
+            kv.key_cache[i] = k[:, :, :seq_len, :]
+            kv.value_cache[i] = v[:, :, :seq_len, :]
+        return kv

memos/llms/ollama.py

@@ -0,0 +1,82 @@
+from typing import Any
+
+from ollama import Client
+
+from memos.configs.llm import OllamaLLMConfig
+from memos.llms.base import BaseLLM
+from memos.llms.utils import remove_thinking_tags
+from memos.log import get_logger
+from memos.types import MessageList
+
+
+logger = get_logger(__name__)
+
+
+class OllamaLLM(BaseLLM):
+    """Ollama LLM class."""
+
+    def __init__(self, config: OllamaLLMConfig):
+        self.config = config
+        self.api_base = config.api_base
+
+        # Default model if not specified
+        if not self.config.model_name_or_path:
+            self.config.model_name_or_path = "llama3.1: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 generate(self, messages: MessageList) -> Any:
+        """
+        Generate a response from Ollama LLM.
+
+        Args:
+            messages: List of message dicts containing 'role' and 'content'.
+
+        Returns:
+            str: The generated response.
+        """
+        response = self.client.chat(
+            model=self.config.model_name_or_path,
+            messages=messages,
+            options={
+                "temperature": self.config.temperature,
+                "num_predict": self.config.max_tokens,
+                "top_p": self.config.top_p,
+                "top_k": self.config.top_k,
+            },
+        )
+        logger.info(f"Raw response from Ollama: {response.model_dump_json()}")
+
+        str_response = response["message"]["content"] or ""
+        if self.config.remove_think_prefix:
+            return remove_thinking_tags(str_response)
+        else:
+            return str_response

memos/llms/openai.py

@@ -0,0 +1,34 @@
+import openai
+
+from memos.configs.llm import OpenAILLMConfig
+from memos.llms.base import BaseLLM
+from memos.llms.utils import remove_thinking_tags
+from memos.log import get_logger
+from memos.types import MessageList
+
+
+logger = get_logger(__name__)
+
+
+class OpenAILLM(BaseLLM):
+    """OpenAI LLM class."""
+
+    def __init__(self, config: OpenAILLMConfig):
+        self.config = config
+        self.client = openai.Client(api_key=config.api_key, base_url=config.api_base)
+
+    def generate(self, messages: MessageList) -> str:
+        """Generate a response from 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 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

memos/llms/utils.py

@@ -0,0 +1,14 @@
+import re
+
+
+def remove_thinking_tags(text: str) -> str:
+    """
+    Remove thinking tags from the generated text.
+
+    Args:
+        text: The generated text.
+
+    Returns:
+        str: The cleaned text.
+    """
+    return re.sub(r"^<think>.*?</think>\s*", "", text, flags=re.DOTALL).strip()

memos/log.py

@@ -0,0 +1,78 @@
+import logging
+
+from logging.config import dictConfig
+from pathlib import Path
+from sys import stdout
+
+from memos import settings
+
+
+selected_log_level = logging.DEBUG if settings.DEBUG else logging.WARNING
+
+
+def _setup_logfile() -> Path:
+    """ensure the logger filepath is in place
+
+    Returns: the logfile Path
+    """
+    logfile = Path(settings.MEMOS_DIR / "logs" / "memos.log")
+    logfile.parent.mkdir(parents=True, exist_ok=True)
+    logfile.touch(exist_ok=True)
+    return logfile
+
+
+LOGGING_CONFIG = {
+    "version": 1,
+    "disable_existing_loggers": False,
+    "formatters": {
+        "standard": {
+            "format": "%(asctime)s - %(name)s - %(levelname)s - %(filename)s:%(lineno)d - %(funcName)s - %(message)s"
+        },
+        "no_datetime": {
+            "format": "%(name)s - %(levelname)s - %(filename)s:%(lineno)d - %(funcName)s - %(message)s"
+        },
+    },
+    "filters": {
+        "package_tree_filter": {"()": "logging.Filter", "name": settings.LOG_FILTER_TREE_PREFIX}
+    },
+    "handlers": {
+        "console": {
+            "level": selected_log_level,
+            "class": "logging.StreamHandler",
+            "stream": stdout,
+            "formatter": "no_datetime",
+            "filters": ["package_tree_filter"],
+        },
+        "file": {
+            "level": "DEBUG",
+            "class": "logging.handlers.RotatingFileHandler",
+            "filename": _setup_logfile(),
+            "maxBytes": 1024**2 * 10,
+            "backupCount": 3,
+            "formatter": "standard",
+        },
+    },
+    "root": {  # Root logger handles all logs
+        "level": logging.DEBUG if settings.DEBUG else logging.INFO,
+        "handlers": ["console", "file"],
+    },
+    "loggers": {
+        "memos": {
+            "level": logging.DEBUG if settings.DEBUG else logging.INFO,
+            "propagate": True,  # Let logs bubble up to root
+        },
+    },
+}
+
+
+def get_logger(name: str | None = None) -> logging.Logger:
+    """returns the project logger, scoped to a child name if provided
+    Args:
+        name: will define a child logger
+    """
+    dictConfig(LOGGING_CONFIG)
+
+    parent_logger = logging.getLogger("")
+    if name:
+        return parent_logger.getChild(name)
+    return parent_logger

memos/mem_chat/base.py

@@ -0,0 +1,30 @@
+from abc import ABC, abstractmethod
+
+from memos.configs.mem_chat import BaseMemChatConfig
+from memos.mem_cube.base import BaseMemCube
+
+
+class BaseMemChat(ABC):
+    """Base class for all MemChat."""
+
+    @abstractmethod
+    def __init__(self, config: BaseMemChatConfig):
+        """Initialize the MemChat with the given configuration."""
+
+    @property
+    @abstractmethod
+    def mem_cube(self) -> BaseMemCube:
+        """The memory cube associated with this MemChat."""
+
+    @mem_cube.setter
+    @abstractmethod
+    def mem_cube(self, value: BaseMemCube) -> None:
+        """The memory cube associated with this MemChat."""
+
+    @abstractmethod
+    def run(self) -> None:
+        """Run the MemChat.
+
+        This `run` method can represent the core logic of a MemChat.
+        It could be an iterative chat process.
+        """

memos/mem_chat/factory.py

@@ -0,0 +1,21 @@
+from typing import Any, ClassVar
+
+from memos.configs.mem_chat import MemChatConfigFactory
+from memos.mem_chat.base import BaseMemChat
+from memos.mem_chat.simple import SimpleMemChat
+
+
+class MemChatFactory(BaseMemChat):
+    """Factory class for creating MemChat instances."""
+
+    backend_to_class: ClassVar[dict[str, Any]] = {
+        "simple": SimpleMemChat,
+    }
+
+    @classmethod
+    def from_config(cls, config_factory: MemChatConfigFactory) -> BaseMemChat:
+        backend = config_factory.backend
+        if backend not in cls.backend_to_class:
+            raise ValueError(f"Invalid backend: {backend}")
+        mem_chat_class = cls.backend_to_class[backend]
+        return mem_chat_class(config_factory.config)