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

memos/mem_scheduler/modules/redis_service.py

@@ -0,0 +1,146 @@
+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")

memos/mem_scheduler/modules/retriever.py

@@ -0,0 +1,41 @@
+from memos.log import get_logger
+from memos.mem_scheduler.modules.base import BaseSchedulerModule
+
+
+logger = get_logger(__name__)
+
+
+class SchedulerRetriever(BaseSchedulerModule):
+    def __init__(self, chat_llm, context_window_size=5):
+        """
+        monitor: Object used to acquire monitoring information
+        mem_cube: Object/interface for querying the underlying database
+        context_window_size: Size of the context window for conversation history
+        """
+        super().__init__()
+
+        self.monitors = {}
+        self.context_window_size = context_window_size
+
+        self._chat_llm = chat_llm
+        self._current_mem_cube = None
+
+    @property
+    def memory_texts(self) -> list[str]:
+        """The memory cube associated with this MemChat."""
+        return self._memory_text_list
+
+    @memory_texts.setter
+    def memory_texts(self, value: list[str]) -> None:
+        """The memory cube associated with this MemChat."""
+        self._memory_text_list = value
+
+    def fetch_context(self):
+        """
+        Extract the context window from the current conversation
+        conversation_history: a list (in chronological order)
+        """
+        return self._memory_text_list[-self.context_window_size :]
+
+    def retrieve(self, query: str, memory_texts: list[str], top_k: int = 5) -> list[str]:
+        return None

memos/mem_scheduler/modules/schemas.py

@@ -0,0 +1,146 @@
+from datetime import datetime
+from pathlib import Path
+from typing import ClassVar, TypeVar
+from uuid import uuid4
+
+from pydantic import BaseModel, Field
+from typing_extensions import TypedDict
+
+from memos.mem_cube.general import GeneralMemCube
+
+
+FILE_PATH = Path(__file__).absolute()
+BASE_DIR = FILE_PATH.parent.parent.parent.parent.parent
+
+QUERY_LABEL = "query"
+ANSWER_LABEL = "answer"
+
+TreeTextMemory_SEARCH_METHOD = "tree_text_memory_search"
+TextMemory_SEARCH_METHOD = "text_memory_search"
+DEFAULT_ACTIVATION_MEM_SIZE = 5
+DEFAULT_ACT_MEM_DUMP_PATH = f"{BASE_DIR}/outputs/mem_scheduler/mem_cube_scheduler_test.kv_cache"
+DEFAULT_THREAD__POOL_MAX_WORKERS = 5
+DEFAULT_CONSUME_INTERVAL_SECONDS = 3
+NOT_INITIALIZED = -1
+BaseModelType = TypeVar("T", bound="BaseModel")
+
+
+class DictConversionMixin:
+    def to_dict(self) -> dict:
+        """Convert the instance to a dictionary."""
+        return {
+            **self.dict(),
+            "timestamp": self.timestamp.isoformat() if hasattr(self, "timestamp") else None,
+        }
+
+    @classmethod
+    def from_dict(cls: type[BaseModelType], data: dict) -> BaseModelType:
+        """Create an instance from a dictionary."""
+        if "timestamp" in data:
+            data["timestamp"] = datetime.fromisoformat(data["timestamp"])
+        return cls(**data)
+
+    class Config:
+        json_encoders: ClassVar[dict[type, object]] = {datetime: lambda v: v.isoformat()}
+
+
+class ScheduleMessageItem(BaseModel, DictConversionMixin):
+    item_id: str = Field(description="uuid", default_factory=lambda: str(uuid4()))
+    user_id: str = Field(..., description="user id")
+    mem_cube_id: str = Field(..., description="memcube id")
+    label: str = Field(..., description="Label of the schedule message")
+    mem_cube: GeneralMemCube | str = Field(..., description="memcube for schedule")
+    content: str = Field(..., description="Content of the schedule message")
+    timestamp: datetime = Field(
+        default_factory=datetime.now, description="submit time for schedule_messages"
+    )
+
+    class Config:
+        arbitrary_types_allowed = True
+        json_encoders: ClassVar[dict[type, object]] = {
+            datetime: lambda v: v.isoformat(),
+            GeneralMemCube: lambda v: f"<GeneralMemCube:{id(v)}>",
+        }
+
+    def to_dict(self) -> dict:
+        """Convert model to dictionary suitable for Redis Stream"""
+        return {
+            "item_id": self.item_id,
+            "user_id": self.user_id,
+            "cube_id": self.mem_cube_id,
+            "message_id": self.message_id,
+            "label": self.label,
+            "cube": "Not Applicable",  # Custom cube serialization
+            "content": self.content,
+            "timestamp": self.timestamp.isoformat(),
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict) -> "ScheduleMessageItem":
+        """Create model from Redis Stream dictionary"""
+        return cls(
+            item_id=data.get("item_id", str(uuid4())),
+            user_id=data["user_id"],
+            cube_id=data["cube_id"],
+            message_id=data.get("message_id", str(uuid4())),
+            label=data["label"],
+            cube="Not Applicable",  # Custom cube deserialization
+            content=data["content"],
+            timestamp=datetime.fromisoformat(data["timestamp"]),
+        )
+
+
+class MemorySizes(TypedDict):
+    long_term_memory_size: int
+    user_memory_size: int
+    working_memory_size: int
+    transformed_act_memory_size: int
+
+
+class MemoryCapacities(TypedDict):
+    long_term_memory_capacity: int
+    user_memory_capacity: int
+    working_memory_capacity: int
+    transformed_act_memory_capacity: int
+
+
+DEFAULT_MEMORY_SIZES = {
+    "long_term_memory_size": NOT_INITIALIZED,
+    "user_memory_size": NOT_INITIALIZED,
+    "working_memory_size": NOT_INITIALIZED,
+    "transformed_act_memory_size": NOT_INITIALIZED,
+    "parameter_memory_size": NOT_INITIALIZED,
+}
+
+DEFAULT_MEMORY_CAPACITIES = {
+    "long_term_memory_capacity": 10000,
+    "user_memory_capacity": 10000,
+    "working_memory_capacity": 20,
+    "transformed_act_memory_capacity": NOT_INITIALIZED,
+    "parameter_memory_capacity": NOT_INITIALIZED,
+}
+
+
+class ScheduleLogForWebItem(BaseModel, DictConversionMixin):
+    item_id: str = Field(
+        description="Unique identifier for the log entry", default_factory=lambda: str(uuid4())
+    )
+    user_id: str = Field(..., description="Identifier for the user associated with the log")
+    mem_cube_id: str = Field(
+        ..., description="Identifier for the memcube associated with this log entry"
+    )
+    label: str = Field(..., description="Label categorizing the type of log")
+    log_title: str = Field(..., description="Title or brief summary of the log content")
+    log_content: str = Field(..., description="Detailed content of the log entry")
+    current_memory_sizes: MemorySizes = Field(
+        default_factory=lambda: dict(DEFAULT_MEMORY_SIZES),
+        description="Current utilization of memory partitions",
+    )
+    memory_capacities: MemoryCapacities = Field(
+        default_factory=lambda: dict(DEFAULT_MEMORY_CAPACITIES),
+        description="Maximum capacities of memory partitions",
+    )
+    timestamp: datetime = Field(
+        default_factory=datetime.now,
+        description="Timestamp indicating when the log entry was created",
+    )

memos/mem_scheduler/scheduler_factory.py

@@ -0,0 +1,21 @@
+from typing import Any, ClassVar
+
+from memos.configs.mem_scheduler import SchedulerConfigFactory
+from memos.mem_scheduler.base_scheduler import BaseScheduler
+from memos.mem_scheduler.general_scheduler import GeneralScheduler
+
+
+class SchedulerFactory(BaseScheduler):
+    """Factory class for creating scheduler instances."""
+
+    backend_to_class: ClassVar[dict[str, Any]] = {
+        "general_scheduler": GeneralScheduler,
+    }
+
+    @classmethod
+    def from_config(cls, config_factory: SchedulerConfigFactory) -> GeneralScheduler:
+        backend = config_factory.backend
+        if backend not in cls.backend_to_class:
+            raise ValueError(f"Invalid backend: {backend}")
+        mem_scheduler_class = cls.backend_to_class[backend]
+        return mem_scheduler_class(config_factory.config)

memos/mem_scheduler/utils.py

@@ -0,0 +1,26 @@
+import json
+
+from pathlib import Path
+
+import yaml
+
+
+def extract_json_dict(text: str):
+    text = text.strip()
+    patterns_to_remove = ["json'''", "latex'''", "'''"]
+    for pattern in patterns_to_remove:
+        text = text.replace(pattern, "")
+    res = json.loads(text)
+    return res
+
+
+def parse_yaml(yaml_file):
+    yaml_path = Path(yaml_file)
+    yaml_path = Path(yaml_file)
+    if not yaml_path.is_file():
+        raise FileNotFoundError(f"No such file: {yaml_file}")
+
+    with yaml_path.open("r", encoding="utf-8") as fr:
+        data = yaml.safe_load(fr)
+
+    return data