MemoryOS 2.0.3__py3-none-any.whl

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/llms/vllm.py

@@ -0,0 +1,218 @@
+import json
+
+from typing import Any, cast
+
+import openai
+
+from openai.types.chat.chat_completion_message_tool_call import ChatCompletionMessageToolCall
+
+from memos.configs.llm import VLLMLLMConfig
+from memos.llms.base import BaseLLM
+from memos.llms.utils import remove_thinking_tags
+from memos.log import get_logger
+from memos.types import MessageDict
+
+
+logger = get_logger(__name__)
+
+
+class VLLMLLM(BaseLLM):
+    """
+    VLLM LLM class for connecting to existing vLLM servers.
+    """
+
+    def __init__(self, config: VLLMLLMConfig):
+        """
+        Initialize the VLLM LLM to connect to an existing vLLM server.
+        """
+        self.config = config
+
+        # Initialize OpenAI client for API calls
+        self.client = None
+        api_key = getattr(self.config, "api_key", "dummy")
+        if not api_key:
+            api_key = "dummy"
+
+        self.client = openai.Client(
+            api_key=api_key,
+            base_url=getattr(self.config, "api_base", "http://localhost:8088/v1"),
+            default_headers=self.config.default_headers,
+        )
+
+    def build_vllm_kv_cache(self, messages: Any) -> str:
+        """
+        Build a KV cache from chat messages via one vLLM request.
+        Handles str, list[str], and MessageList formats.
+        """
+        # 1. Normalize input to a MessageList
+        processed_messages: list[MessageDict] = []
+        if isinstance(messages, str):
+            processed_messages = [
+                {
+                    "role": "system",
+                    "content": f"Below is some information about the user.\n{messages}",
+                }
+            ]
+        elif isinstance(messages, list):
+            if not messages:
+                pass  # Empty list
+            elif isinstance(messages[0], str):
+                str_content = " ".join(str(msg) for msg in messages)
+                processed_messages = [
+                    {
+                        "role": "system",
+                        "content": f"Below is some information about the user.\n{str_content}",
+                    }
+                ]
+            elif isinstance(messages[0], dict):
+                processed_messages = cast("list[MessageDict]", messages)
+
+        # 2. Convert to prompt for logging/return value.
+        prompt = self._messages_to_prompt(processed_messages)
+
+        if not prompt.strip():
+            raise ValueError("Prompt is empty, cannot build KV cache.")
+
+        # 3. Send request to vLLM server to preload the KV cache
+        if self.client:
+            try:
+                # Use the processed messages for the API call
+                prefill_kwargs = {
+                    "model": self.config.model_name_or_path,
+                    "messages": processed_messages,
+                    "max_tokens": 2,
+                    "temperature": 0.0,
+                    "top_p": 1.0,
+                }
+                self.client.chat.completions.create(**prefill_kwargs)
+                logger.info(f"vLLM KV cache prefill completed for prompt: '{prompt[:100]}...'")
+            except Exception as e:
+                logger.warning(f"Failed to prefill vLLM KV cache: {e}")
+
+        return prompt
+
+    def generate(self, messages: list[MessageDict], **kwargs) -> str:
+        """
+        Generate a response from the model.
+        """
+        if self.client:
+            return self._generate_with_api_client(messages, **kwargs)
+        else:
+            raise RuntimeError("API client is not available")
+
+    def _generate_with_api_client(self, messages: list[MessageDict], **kwargs) -> str:
+        """
+        Generate response using vLLM API client. detail view https://docs.vllm.ai/en/latest/features/reasoning_outputs/
+        """
+        if self.client:
+            completion_kwargs = {
+                "model": kwargs.get("model_name_or_path", self.config.model_name_or_path),
+                "messages": messages,
+                "temperature": kwargs.get("temperature", self.config.temperature),
+                "max_tokens": kwargs.get("max_tokens", self.config.max_tokens),
+                "top_p": kwargs.get("top_p", self.config.top_p),
+                "extra_body": {
+                    "chat_template_kwargs": {
+                        "enable_thinking": kwargs.get(
+                            "enable_thinking", self.config.enable_thinking
+                        )
+                    }
+                },
+            }
+            if kwargs.get("tools"):
+                completion_kwargs["tools"] = kwargs.get("tools")
+                completion_kwargs["tool_choice"] = kwargs.get("tool_choice", "auto")
+
+            response = self.client.chat.completions.create(**completion_kwargs)
+
+            if not response.choices:
+                logger.warning("VLLM response has no choices")
+                return ""
+
+            if response.choices[0].message.tool_calls:
+                return self.tool_call_parser(response.choices[0].message.tool_calls)
+
+            reasoning_content = (
+                f"<think>{response.choices[0].message.reasoning}</think>"
+                if hasattr(response.choices[0].message, "reasoning")
+                else ""
+            )
+            response_text = response.choices[0].message.content or ""
+            logger.info(f"VLLM API response: {response_text}")
+            return (
+                remove_thinking_tags(response_text)
+                if getattr(self.config, "remove_think_prefix", False)
+                else reasoning_content + response_text
+            )
+        else:
+            raise RuntimeError("API client is not available")
+
+    def _messages_to_prompt(self, messages: list[MessageDict]) -> str:
+        """
+        Convert messages to prompt string.
+        """
+        prompt_parts = []
+        for msg in messages:
+            role = msg["role"]
+            content = msg["content"]
+            prompt_parts.append(f"{role.capitalize()}: {content}")
+        return "\n".join(prompt_parts)
+
+    def generate_stream(self, messages: list[MessageDict], **kwargs):
+        """
+        Generate a response from the model using streaming.
+        Yields content chunks as they are received.
+        """
+        if kwargs.get("tools"):
+            logger.info("stream api not support tools")
+            return
+
+        if self.client:
+            completion_kwargs = {
+                "model": self.config.model_name_or_path,
+                "messages": messages,
+                "temperature": kwargs.get("temperature", self.config.temperature),
+                "max_tokens": kwargs.get("max_tokens", self.config.max_tokens),
+                "top_p": kwargs.get("top_p", self.config.top_p),
+                "stream": True,
+                "extra_body": {
+                    "chat_template_kwargs": {
+                        "enable_thinking": kwargs.get(
+                            "enable_thinking", self.config.enable_thinking
+                        )
+                    }
+                },
+            }
+
+            stream = self.client.chat.completions.create(**completion_kwargs)
+
+            reasoning_started = False
+            for chunk in stream:
+                if not chunk.choices:
+                    continue
+                delta = chunk.choices[0].delta
+                if hasattr(delta, "reasoning") and delta.reasoning:
+                    if not reasoning_started and not self.config.remove_think_prefix:
+                        yield "<think>"
+                        reasoning_started = True
+                    yield delta.reasoning
+
+            if hasattr(delta, "content") and delta.content:
+                if reasoning_started and not self.config.remove_think_prefix:
+                    yield "</think>"
+                    reasoning_started = False
+                yield delta.content
+
+        else:
+            raise RuntimeError("API client is not available")
+
+    def tool_call_parser(self, tool_calls: list[ChatCompletionMessageToolCall]) -> list[dict]:
+        """Parse tool calls from OpenAI response."""
+        return [
+            {
+                "tool_call_id": tool_call.id,
+                "function_name": tool_call.function.name,
+                "arguments": json.loads(tool_call.function.arguments),
+            }
+            for tool_call in tool_calls
+        ]

memos/log.py

@@ -0,0 +1,237 @@
+import atexit
+import logging
+import os
+import threading
+import time
+
+from concurrent.futures import ThreadPoolExecutor
+from logging.config import dictConfig
+from pathlib import Path
+from sys import stdout
+
+import requests
+
+from dotenv import load_dotenv
+
+from memos import settings
+from memos.context.context import (
+    get_current_api_path,
+    get_current_env,
+    get_current_trace_id,
+    get_current_user_name,
+    get_current_user_type,
+)
+
+
+# Load environment variables
+load_dotenv()
+
+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
+
+
+class ContextFilter(logging.Filter):
+    """add context to the log record"""
+
+    def filter(self, record):
+        try:
+            trace_id = get_current_trace_id()
+            record.trace_id = trace_id if trace_id else "trace-id"
+            record.env = get_current_env()
+            record.user_type = get_current_user_type()
+            record.user_name = get_current_user_name()
+            record.api_path = get_current_api_path()
+        except Exception:
+            record.api_path = "unknown"
+            record.trace_id = "trace-id"
+            record.env = "prod"
+            record.user_type = "normal"
+            record.user_name = "unknown"
+        return True
+
+
+class CustomLoggerRequestHandler(logging.Handler):
+    _instance = None
+    _lock = threading.Lock()
+
+    def __new__(cls):
+        if cls._instance is None:
+            with cls._lock:
+                if cls._instance is None:
+                    cls._instance = super().__new__(cls)
+                    cls._instance._initialized = False
+                    cls._instance._executor = None
+                    cls._instance._session = None
+                    cls._instance._is_shutting_down = None
+        return cls._instance
+
+    def __init__(self):
+        """Initialize handler with minimal setup"""
+        if not self._initialized:
+            super().__init__()
+            workers = int(os.getenv("CUSTOM_LOGGER_WORKERS", "2"))
+            self._executor = ThreadPoolExecutor(
+                max_workers=workers, thread_name_prefix="log_sender"
+            )
+            self._is_shutting_down = threading.Event()
+            self._session = requests.Session()
+            self._initialized = True
+            atexit.register(self._cleanup)
+
+    def emit(self, record):
+        """Process log records of INFO or ERROR level (non-blocking)"""
+        if os.getenv("CUSTOM_LOGGER_URL") is None or self._is_shutting_down.is_set():
+            return
+
+        # Only process INFO and ERROR level logs
+        if record.levelno < logging.INFO:  # Skip DEBUG and lower
+            return
+
+        try:
+            trace_id = get_current_trace_id() or "trace-id"
+            api_path = get_current_api_path()
+            env = get_current_env()
+            user_type = get_current_user_type()
+            user_name = get_current_user_name()
+            if api_path is not None:
+                self._executor.submit(
+                    self._send_log_sync,
+                    record.getMessage(),
+                    trace_id,
+                    api_path,
+                    env,
+                    user_type,
+                    user_name,
+                )
+        except Exception as e:
+            if not self._is_shutting_down.is_set():
+                print(f"Error sending log: {e}")
+
+    def _send_log_sync(self, message, trace_id, api_path, env, user_type, user_name):
+        """Send log message synchronously in a separate thread"""
+        try:
+            logger_url = os.getenv("CUSTOM_LOGGER_URL")
+            token = os.getenv("CUSTOM_LOGGER_TOKEN")
+
+            headers = {"Content-Type": "application/json"}
+            post_content = {
+                "message": message,
+                "trace_id": trace_id,
+                "action": api_path,
+                "current_time": round(time.time(), 3),
+                "env": env,
+                "user_type": user_type,
+                "user_name": user_name,
+            }
+
+            # Add auth token if exists
+            if token:
+                headers["Authorization"] = f"Bearer {token}"
+
+            # Add traceId to headers for consistency
+            headers["traceId"] = trace_id
+
+            # Add custom attributes from env
+            for key, value in os.environ.items():
+                if key.startswith("CUSTOM_LOGGER_ATTRIBUTE_"):
+                    attribute_key = key[len("CUSTOM_LOGGER_ATTRIBUTE_") :].lower()
+                    post_content[attribute_key] = value
+
+            self._session.post(logger_url, headers=headers, json=post_content, timeout=5)
+        except Exception:
+            # Silently ignore errors to avoid affecting main application
+            pass
+
+    def _cleanup(self):
+        """Clean up resources during program exit"""
+        if not self._initialized:
+            return
+
+        self._is_shutting_down.set()
+        try:
+            self._executor.shutdown(wait=False)
+            self._session.close()
+        except Exception as e:
+            print(f"Error during cleanup: {e}")
+
+    def close(self):
+        """Override close to prevent premature shutdown"""
+
+
+LOGGING_CONFIG = {
+    "version": 1,
+    "disable_existing_loggers": False,
+    "formatters": {
+        "standard": {
+            "format": "%(asctime)s | %(trace_id)s | path=%(api_path)s | env=%(env)s | user_type=%(user_type)s | user_name=%(user_name)s | %(name)s - %(levelname)s - %(filename)s:%(lineno)d - %(funcName)s - %(message)s"
+        },
+        "no_datetime": {
+            "format": "%(trace_id)s | path=%(api_path)s | %(name)s - %(levelname)s - %(filename)s:%(lineno)d - %(funcName)s - %(message)s"
+        },
+        "simplified": {
+            "format": "%(asctime)s | %(trace_id)s | path=%(api_path)s | % %(levelname)s | %(filename)s:%(lineno)d: %(funcName)s | %(message)s"
+        },
+    },
+    "filters": {
+        "package_tree_filter": {"()": "logging.Filter", "name": settings.LOG_FILTER_TREE_PREFIX},
+        "context_filter": {"()": "memos.log.ContextFilter"},
+    },
+    "handlers": {
+        "console": {
+            "level": selected_log_level,
+            "class": "logging.StreamHandler",
+            "stream": stdout,
+            "formatter": "no_datetime",
+            "filters": ["package_tree_filter", "context_filter"],
+        },
+        "file": {
+            "level": "INFO",
+            "class": "concurrent_log_handler.ConcurrentTimedRotatingFileHandler",
+            "when": "midnight",
+            "interval": 1,
+            "backupCount": 3,
+            "filename": _setup_logfile(),
+            "formatter": "standard",
+            "filters": ["context_filter"],
+        },
+        "custom_logger": {
+            "level": "INFO",
+            "class": "memos.log.CustomLoggerRequestHandler",
+            "formatter": "simplified",
+        },
+    },
+    "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_agent/base.py

@@ -0,0 +1,19 @@
+from abc import ABC, abstractmethod
+
+from memos.configs.mem_agent import BaseAgentConfig
+
+
+class BaseMemAgent(ABC):
+    """
+    Base class for all agents.
+    """
+
+    def __init__(self, config: BaseAgentConfig):
+        """Initialize the BaseMemAgent with the given configuration."""
+        self.config = config
+
+    @abstractmethod
+    def run(self, input: str) -> str:
+        """
+        Run the agent.
+        """