MemOS-MCP

import asyncio
import threading

from collections.abc import Callable

import redis

from redis import Redis

from memos.log import get_logger
from memos.mem_scheduler.modules.base import BaseSchedulerModule


logger = get_logger(__name__)


class RedisSchedulerModule(BaseSchedulerModule):
    def __init__(self):
        """
        intent_detector: Object used for intent recognition (such as the above IntentDetector)
        scheduler: The actual scheduling module/interface object
        trigger_intents: The types of intents that need to be triggered (list)
        """
        super().__init__()

        # settings for redis
        self.redis_host: str = None
        self.redis_port: int = None
        self.redis_db: int = None
        self._redis_conn = None
        self.query_list_capacity = 1000

        self._redis_listener_running = False
        self._redis_listener_thread: threading.Thread | None = None
        self._redis_listener_loop: asyncio.AbstractEventLoop | None = None

    @property
    def redis(self) -> Redis:
        return self._redis_conn

    @redis.setter
    def redis(self, value: Redis) -> None:
        self._redis_conn = value

    def initialize_redis(
        self, redis_host: str = "localhost", redis_port: int = 6379, redis_db: int = 0
    ):
        self.redis_host = redis_host
        self.redis_port = redis_port
        self.redis_db = redis_db

        try:
            logger.debug(f"Connecting to Redis at {redis_host}:{redis_port}/{redis_db}")
            self._redis_conn = Redis(
                host=self.redis_host, port=self.redis_port, db=self.redis_db, decode_responses=True
            )
            # test conn
            if not self._redis_conn.ping():
                logger.error("Redis connection failed")
        except redis.ConnectionError as e:
            self._redis_conn = None
            logger.error(f"Redis connection error: {e}")
        self._redis_conn.xtrim("user:queries:stream", self.query_list_capacity)
        return self._redis_conn

    async def add_message_stream(self, message: dict):
        logger.debug(f"add_message_stream: {message}")
        return self._redis_conn.xadd("user:queries:stream", message)

    async def consume_message_stream(self, message: dict):
        logger.debug(f"consume_message_stream: {message}")

    def _run_listener_async(self, handler: Callable):
        """Run the async listener in a separate thread"""
        self._redis_listener_loop = asyncio.new_event_loop()
        asyncio.set_event_loop(self._redis_listener_loop)

        async def listener_wrapper():
            try:
                await self._listen_query_stream(handler)
            except Exception as e:
                logger.error(f"Listener thread error: {e}")
            finally:
                self._redis_listener_running = False

        self._redis_listener_loop.run_until_complete(listener_wrapper())

    async def _listen_query_stream(self, handler=None, last_id: str = "$", block_time: int = 2000):
        """Internal async stream listener"""
        self._redis_listener_running = True
        while self._redis_listener_running:
            try:
                # Blocking read for new messages
                messages = self.redis.xread(
                    {"user:queries:stream": last_id}, count=1, block=block_time
                )

                if messages:
                    for _, stream_messages in messages:
                        for message_id, message_data in stream_messages:
                            try:
                                await handler(message_data)
                                last_id = message_id
                            except Exception as e:
                                logger.error(f"Error processing message {message_id}: {e}")

            except redis.ConnectionError as e:
                logger.error(f"Redis connection error: {e}")
                await asyncio.sleep(5)  # Wait before reconnecting
                self._redis_conn = None  # Force reconnection
            except Exception as e:
                logger.error(f"Unexpected error: {e}")
                await asyncio.sleep(1)

    def start_listening(self, handler: Callable | None = None):
        """Start the Redis stream listener in a background thread"""
        if self._redis_listener_thread and self._redis_listener_thread.is_alive():
            logger.warning("Listener is already running")
            return

        if handler is None:
            handler = self.consume_message_stream

        self._redis_listener_thread = threading.Thread(
            target=self._run_listener_async,
            args=(handler,),
            daemon=True,
            name="RedisListenerThread",
        )
        self._redis_listener_thread.start()
        logger.info("Started Redis stream listener thread")

    def close(self):
        """Close Redis connection"""
        if self._redis_conn is not None:
            self._redis_conn.close()
            self._redis_conn = None

    def stop_listening(self):
        """Stop the listener thread gracefully"""
        self._redis_listener_running = False
        if self._redis_listener_thread and self._redis_listener_thread.is_alive():
            self._redis_listener_thread.join(timeout=5.0)
            if self._redis_listener_thread.is_alive():
                logger.warning("Listener thread did not stop gracefully")
        logger.info("Redis stream listener stopped")