MemoryOS 1.0.1__py3-none-any.whl → 1.1.1__py3-none-any.whl

memos/types.py

@@ -22,11 +22,14 @@ MessageRole: TypeAlias = Literal["user", "assistant", "system"]
 
 
 # Message structure
-class MessageDict(TypedDict):
+class MessageDict(TypedDict, total=False):
     """Typed dictionary for chat message dictionaries."""
 
     role: MessageRole
     content: str
+    chat_time: str | None  # Optional timestamp for the message, format is not
+    # restricted, it can be any vague or precise time string.
+    message_id: str | None  # Optional unique identifier for the message
 
 
 # Message collections

memos/api/context/context.py

@@ -1,147 +0,0 @@
-"""
-Global request context management for trace_id and request-scoped data.
-
-This module provides optional trace_id functionality that can be enabled
-when using the API components. It uses ContextVar to ensure thread safety
-and request isolation.
-"""
-
-import uuid
-
-from collections.abc import Callable
-from contextvars import ContextVar
-from typing import Any
-
-
-# Global context variable for request-scoped data
-_request_context: ContextVar[dict[str, Any] | None] = ContextVar("request_context", default=None)
-
-
-class RequestContext:
-    """
-    Request-scoped context object that holds trace_id and other request data.
-
-    This provides a Flask g-like object for FastAPI applications.
-    """
-
-    def __init__(self, trace_id: str | None = None):
-        self.trace_id = trace_id or str(uuid.uuid4())
-        self._data: dict[str, Any] = {}
-
-    def set(self, key: str, value: Any) -> None:
-        """Set a value in the context."""
-        self._data[key] = value
-
-    def get(self, key: str, default: Any | None = None) -> Any:
-        """Get a value from the context."""
-        return self._data.get(key, default)
-
-    def __setattr__(self, name: str, value: Any) -> None:
-        if name.startswith("_") or name == "trace_id":
-            super().__setattr__(name, value)
-        else:
-            if not hasattr(self, "_data"):
-                super().__setattr__(name, value)
-            else:
-                self._data[name] = value
-
-    def __getattr__(self, name: str) -> Any:
-        if hasattr(self, "_data") and name in self._data:
-            return self._data[name]
-        raise AttributeError(f"'{self.__class__.__name__}' object has no attribute '{name}'")
-
-    def to_dict(self) -> dict[str, Any]:
-        """Convert context to dictionary."""
-        return {"trace_id": self.trace_id, "data": self._data.copy()}
-
-
-def set_request_context(context: RequestContext) -> None:
-    """
-    Set the current request context.
-
-    This is typically called by the API dependency injection system.
-    """
-    _request_context.set(context.to_dict())
-
-
-def get_current_trace_id() -> str | None:
-    """
-    Get the current request's trace_id.
-
-    Returns:
-        The trace_id if available, None otherwise.
-    """
-    context = _request_context.get()
-    if context:
-        return context.get("trace_id")
-    return None
-
-
-def get_current_context() -> RequestContext | None:
-    """
-    Get the current request context.
-
-    Returns:
-        The current RequestContext if available, None otherwise.
-    """
-    context_dict = _request_context.get()
-    if context_dict:
-        ctx = RequestContext(trace_id=context_dict.get("trace_id"))
-        ctx._data = context_dict.get("data", {}).copy()
-        return ctx
-    return None
-
-
-def require_context() -> RequestContext:
-    """
-    Get the current request context, raising an error if not available.
-
-    Returns:
-        The current RequestContext.
-
-    Raises:
-        RuntimeError: If called outside of a request context.
-    """
-    context = get_current_context()
-    if context is None:
-        raise RuntimeError(
-            "No request context available. This function must be called within a request handler."
-        )
-    return context
-
-
-# Type for trace_id getter function
-TraceIdGetter = Callable[[], str | None]
-
-# Global variable to hold the trace_id getter function
-_trace_id_getter: TraceIdGetter | None = None
-
-
-def set_trace_id_getter(getter: TraceIdGetter) -> None:
-    """
-    Set a custom trace_id getter function.
-
-    This allows the logging system to retrieve trace_id without importing
-    API-specific general_modules.
-    """
-    global _trace_id_getter
-    _trace_id_getter = getter
-
-
-def get_trace_id_for_logging() -> str | None:
-    """
-    Get trace_id for logging purposes.
-
-    This function is used by the logging system and will use either
-    the custom getter function or fall back to the default context.
-    """
-    if _trace_id_getter:
-        try:
-            return _trace_id_getter()
-        except Exception:
-            pass
-    return get_current_trace_id()
-
-
-# Initialize the default trace_id getter
-set_trace_id_getter(get_current_trace_id)

memos/api/context/context_thread.py

@@ -1,96 +0,0 @@
-import functools
-import threading
-
-from collections.abc import Callable
-from concurrent.futures import ThreadPoolExecutor
-from typing import Any, TypeVar
-
-from memos.api.context.context import (
-    RequestContext,
-    get_current_context,
-    get_current_trace_id,
-    set_request_context,
-)
-
-
-T = TypeVar("T")
-
-
-class ContextThread(threading.Thread):
-    """
-    Thread class that automatically propagates the main thread's trace_id to child threads.
-    """
-
-    def __init__(self, target, args=(), kwargs=None, **thread_kwargs):
-        super().__init__(**thread_kwargs)
-        self.target = target
-        self.args = args
-        self.kwargs = kwargs or {}
-
-        self.main_trace_id = get_current_trace_id()
-        self.main_context = get_current_context()
-
-    def run(self):
-        # Create a new RequestContext with the main thread's trace_id
-        if self.main_context:
-            # Copy the context data
-            child_context = RequestContext(trace_id=self.main_trace_id)
-            child_context._data = self.main_context._data.copy()
-
-            # Set the context in the child thread
-            set_request_context(child_context)
-
-        # Run the target function
-        self.target(*self.args, **self.kwargs)
-
-
-class ContextThreadPoolExecutor(ThreadPoolExecutor):
-    """
-    ThreadPoolExecutor that automatically propagates the main thread's trace_id to worker threads.
-    """
-
-    def submit(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> Any:
-        """
-        Submit a callable to be executed with the given arguments.
-        Automatically propagates the current thread's context to the worker thread.
-        """
-        main_trace_id = get_current_trace_id()
-        main_context = get_current_context()
-
-        @functools.wraps(fn)
-        def wrapper(*args: Any, **kwargs: Any) -> Any:
-            if main_context:
-                # Create and set new context in worker thread
-                child_context = RequestContext(trace_id=main_trace_id)
-                child_context._data = main_context._data.copy()
-                set_request_context(child_context)
-
-            return fn(*args, **kwargs)
-
-        return super().submit(wrapper, *args, **kwargs)
-
-    def map(
-        self,
-        fn: Callable[..., T],
-        *iterables: Any,
-        timeout: float | None = None,
-        chunksize: int = 1,
-    ) -> Any:
-        """
-        Returns an iterator equivalent to map(fn, iter).
-        Automatically propagates the current thread's context to worker threads.
-        """
-        main_trace_id = get_current_trace_id()
-        main_context = get_current_context()
-
-        @functools.wraps(fn)
-        def wrapper(*args: Any, **kwargs: Any) -> Any:
-            if main_context:
-                # Create and set new context in worker thread
-                child_context = RequestContext(trace_id=main_trace_id)
-                child_context._data = main_context._data.copy()
-                set_request_context(child_context)
-
-            return fn(*args, **kwargs)
-
-        return super().map(wrapper, *iterables, timeout=timeout, chunksize=chunksize)

memos/mem_scheduler/mos_for_test_scheduler.py

@@ -1,146 +0,0 @@
-from datetime import datetime
-
-from memos.configs.mem_os import MOSConfig
-from memos.log import get_logger
-from memos.mem_os.main import MOS
-from memos.mem_scheduler.schemas.general_schemas import (
-    ANSWER_LABEL,
-    MONITOR_WORKING_MEMORY_TYPE,
-    QUERY_LABEL,
-)
-from memos.mem_scheduler.schemas.message_schemas import ScheduleMessageItem
-
-
-logger = get_logger(__name__)
-
-
-class MOSForTestScheduler(MOS):
-    """This class is only to test abilities of mem scheduler"""
-
-    def __init__(self, config: MOSConfig):
-        super().__init__(config)
-
-    def _str_memories(self, memories: list[str]) -> str:
-        """Format memories for display."""
-        if not memories:
-            return "No memories."
-        return "\n".join(f"{i + 1}. {memory}" for i, memory in enumerate(memories))
-
-    def chat(self, query: str, user_id: str | None = None) -> str:
-        """
-        Chat with the MOS.
-
-        Args:
-            query (str): The user's query.
-
-        Returns:
-            str: The response from the MOS.
-        """
-        target_user_id = user_id if user_id is not None else self.user_id
-        accessible_cubes = self.user_manager.get_user_cubes(target_user_id)
-        user_cube_ids = [cube.cube_id for cube in accessible_cubes]
-        if target_user_id not in self.chat_history_manager:
-            self._register_chat_history(target_user_id)
-
-        chat_history = self.chat_history_manager[target_user_id]
-
-        topk_for_scheduler = 2
-
-        if self.config.enable_textual_memory and self.mem_cubes:
-            memories_all = []
-            for mem_cube_id, mem_cube in self.mem_cubes.items():
-                if mem_cube_id not in user_cube_ids:
-                    continue
-                if not mem_cube.text_mem:
-                    continue
-
-                message_item = ScheduleMessageItem(
-                    user_id=target_user_id,
-                    mem_cube_id=mem_cube_id,
-                    mem_cube=mem_cube,
-                    label=QUERY_LABEL,
-                    content=query,
-                    timestamp=datetime.now(),
-                )
-                cur_working_memories = [m.memory for m in mem_cube.text_mem.get_working_memory()]
-                print(f"Working memories before schedule: {cur_working_memories}")
-
-                # --- force to run mem_scheduler ---
-                self.mem_scheduler.monitor.query_trigger_interval = 0
-                self.mem_scheduler._query_message_consumer(messages=[message_item])
-
-                # from scheduler
-                scheduler_memories = self.mem_scheduler.monitor.get_monitor_memories(
-                    user_id=target_user_id,
-                    mem_cube_id=mem_cube_id,
-                    memory_type=MONITOR_WORKING_MEMORY_TYPE,
-                    top_k=topk_for_scheduler,
-                )
-                print(f"Working memories after schedule: {scheduler_memories}")
-                memories_all.extend(scheduler_memories)
-
-                # from mem_cube
-                memories = mem_cube.text_mem.search(
-                    query,
-                    top_k=self.config.top_k - topk_for_scheduler,
-                    info={
-                        "user_id": target_user_id,
-                        "session_id": self.session_id,
-                        "chat_history": chat_history.chat_history,
-                    },
-                )
-                text_memories = [m.memory for m in memories]
-                print(f"Search results with new working memories: {text_memories}")
-                memories_all.extend(text_memories)
-
-                memories_all = list(set(memories_all))
-
-            logger.info(f"🧠 [Memory] Searched memories:\n{self._str_memories(memories_all)}\n")
-            system_prompt = self._build_system_prompt(memories_all)
-        else:
-            system_prompt = self._build_system_prompt()
-        current_messages = [
-            {"role": "system", "content": system_prompt},
-            *chat_history.chat_history,
-            {"role": "user", "content": query},
-        ]
-        past_key_values = None
-
-        if self.config.enable_activation_memory:
-            assert self.config.chat_model.backend == "huggingface", (
-                "Activation memory only used for huggingface backend."
-            )
-            # TODO this only one cubes
-            for mem_cube_id, mem_cube in self.mem_cubes.items():
-                if mem_cube_id not in user_cube_ids:
-                    continue
-                if mem_cube.act_mem:
-                    kv_cache = next(iter(mem_cube.act_mem.get_all()), None)
-                    past_key_values = (
-                        kv_cache.memory if (kv_cache and hasattr(kv_cache, "memory")) else None
-                    )
-                    break
-            # Generate response
-            response = self.chat_llm.generate(current_messages, past_key_values=past_key_values)
-        else:
-            response = self.chat_llm.generate(current_messages)
-        logger.info(f"🤖 [Assistant] {response}\n")
-        chat_history.chat_history.append({"role": "user", "content": query})
-        chat_history.chat_history.append({"role": "assistant", "content": response})
-        self.chat_history_manager[user_id] = chat_history
-
-        # submit message to scheduler
-        for accessible_mem_cube in accessible_cubes:
-            mem_cube_id = accessible_mem_cube.cube_id
-            mem_cube = self.mem_cubes[mem_cube_id]
-            if self.enable_mem_scheduler and self.mem_scheduler is not None:
-                message_item = ScheduleMessageItem(
-                    user_id=target_user_id,
-                    mem_cube_id=mem_cube_id,
-                    mem_cube=mem_cube,
-                    label=ANSWER_LABEL,
-                    content=response,
-                    timestamp=datetime.now(),
-                )
-                self.mem_scheduler.submit_messages(messages=[message_item])
-        return response