MemOS-MCP

import shutil import sys from datetime import datetime from pathlib import Path from queue import Queue from memos.configs.mem_cube import GeneralMemCubeConfig from memos.configs.mem_os import MOSConfig from memos.configs.mem_scheduler import SchedulerConfigFactory from memos.log import get_logger from memos.mem_cube.general import GeneralMemCube from memos.mem_os.main import MOS from memos.mem_scheduler.general_scheduler import GeneralScheduler from memos.mem_scheduler.modules.schemas import ( ANSWER_LABEL, QUERY_LABEL, ScheduleLogForWebItem, ScheduleMessageItem, ) from memos.mem_scheduler.scheduler_factory import SchedulerFactory from memos.mem_scheduler.utils import parse_yaml FILE_PATH = Path(__file__).absolute() BASE_DIR = FILE_PATH.parent.parent.parent sys.path.insert(0, str(BASE_DIR)) # Enable execution from any working directory logger = get_logger(__name__) def init_task(): conversations = [ {"role": "user", "content": "I just adopted a golden retriever puppy yesterday."}, {"role": "assistant", "content": "Congratulations! What did you name your new puppy?"}, { "role": "user", "content": "His name is Max. I live near Central Park in New York where we'll walk daily.", }, {"role": "assistant", "content": "Max will love those walks! Any favorite treats for him?"}, { "role": "user", "content": "He loves peanut butter biscuits. Personally, I'm allergic to nuts though.", }, {"role": "assistant", "content": "Good to know about your allergy. I'll note that."}, # Question 1 (Pet) - Name {"role": "user", "content": "What's my dog's name again?"}, {"role": "assistant", "content": "Your dog is named Max."}, # Question 2 (Pet) - Breed {"role": "user", "content": "Can you remind me what breed Max is?"}, {"role": "assistant", "content": "Max is a golden retriever."}, # Question 3 (Pet) - Treat {"role": "user", "content": "What treats does Max like?"}, {"role": "assistant", "content": "He loves peanut butter biscuits."}, # Question 4 (Address) {"role": "user", "content": "Where did I say I live?"}, {"role": "assistant", "content": "You live near Central Park in New York."}, # Question 5 (Allergy) {"role": "user", "content": "What food should I avoid due to allergy?"}, {"role": "assistant", "content": "You're allergic to nuts."}, {"role": "user", "content": "Perfect, just wanted to check what you remembered."}, {"role": "assistant", "content": "Happy to help! Let me know if you need anything else."}, ] questions = [ {"question": "What's my dog's name again?", "category": "Pet"}, {"question": "Can you remind me what breed Max is?", "category": "Pet"}, {"question": "What treats does Max like?", "category": "Pet"}, {"question": "Where did I say I live?", "category": "Address"}, {"question": "What food should I avoid due to allergy?", "category": "Allergy"}, ] return conversations, questions def show_web_logs(mem_scheduler: GeneralScheduler): """Display all web log entries from the scheduler's log queue. Args: mem_scheduler: The scheduler instance containing web logs to display """ if mem_scheduler._web_log_message_queue.empty(): print("Web log queue is currently empty.") return print("\n" + "=" * 50 + " WEB LOGS " + "=" * 50) # Create a temporary queue to preserve the original queue contents temp_queue = Queue() log_count = 0 while not mem_scheduler._web_log_message_queue.empty(): log_item: ScheduleLogForWebItem = mem_scheduler._web_log_message_queue.get() temp_queue.put(log_item) log_count += 1 # Print log entry details print(f"\nLog Entry #{log_count}:") print(f"- log: {log_item}") print("-" * 50) # Restore items back to the original queue while not temp_queue.empty(): mem_scheduler._web_log_message_queue.put(temp_queue.get()) print(f"\nTotal {log_count} web log entries displayed.") print("=" * 110 + "\n") def run_with_automatic_scheduler_init(): print("==== run_with_automatic_scheduler_init ====") conversations, questions = init_task() config = parse_yaml( f"{BASE_DIR}/examples/data/config/mem_scheduler/memos_config_w_scheduler.yaml" ) mos_config = MOSConfig(**config) mos = MOS(mos_config) user_id = "user_1" mos.create_user(user_id) config = GeneralMemCubeConfig.from_yaml_file( f"{BASE_DIR}/examples/data/config/mem_scheduler/mem_cube_config.yaml" ) mem_cube_id = "mem_cube_5" mem_cube_name_or_path = f"{BASE_DIR}/outputs/mem_scheduler/{user_id}/{mem_cube_id}" if Path(mem_cube_name_or_path).exists(): shutil.rmtree(mem_cube_name_or_path) print(f"{mem_cube_name_or_path} is not empty, and has been removed.") mem_cube = GeneralMemCube(config) mem_cube.dump(mem_cube_name_or_path) mos.register_mem_cube( mem_cube_name_or_path=mem_cube_name_or_path, mem_cube_id=mem_cube_id, user_id=user_id ) mos.add(conversations, user_id=user_id, mem_cube_id=mem_cube_id) for item in questions: query = item["question"] response = mos.chat(query, user_id=user_id) print(f"Query:\n {query}\n\nAnswer:\n {response}") show_web_logs(mos.mem_scheduler) mos.mem_scheduler.stop() def run_with_manual_scheduler_init(): print("==== run_with_manual_scheduler_init ====") conversations, questions = init_task() config = parse_yaml( f"{BASE_DIR}/examples/data/config/mem_scheduler/memos_config_wo_scheduler.yaml" ) mos_config = MOSConfig(**config) mos = MOS(mos_config) user_id = "user_1" mos.create_user(user_id) config = GeneralMemCubeConfig.from_yaml_file( f"{BASE_DIR}/examples/data/config/mem_scheduler/mem_cube_config.yaml" ) mem_cube_id = "mem_cube_5" mem_cube_name_or_path = f"{BASE_DIR}/outputs/mem_scheduler/{user_id}/{mem_cube_id}" if Path(mem_cube_name_or_path).exists(): shutil.rmtree(mem_cube_name_or_path) print(f"{mem_cube_name_or_path} is not empty, and has been removed.") mem_cube = GeneralMemCube(config) mem_cube.dump(mem_cube_name_or_path) mos.register_mem_cube( mem_cube_name_or_path=mem_cube_name_or_path, mem_cube_id=mem_cube_id, user_id=user_id ) example_scheduler_config_path = ( f"{BASE_DIR}/examples/data/config/mem_scheduler/general_scheduler_config.yaml" ) scheduler_config = SchedulerConfigFactory.from_yaml_file( yaml_path=example_scheduler_config_path ) mem_scheduler = SchedulerFactory.from_config(scheduler_config) mem_scheduler.initialize_modules(chat_llm=mos.chat_llm) mos.mem_scheduler = mem_scheduler mos.mem_scheduler.start() mos.add(conversations, user_id=user_id, mem_cube_id=mem_cube_id) for item in questions: query = item["question"] message_item = ScheduleMessageItem( user_id=user_id, mem_cube_id=mem_cube_id, label=QUERY_LABEL, mem_cube=mos.mem_cubes[mem_cube_id], content=query, timestamp=datetime.now(), ) mos.mem_scheduler.submit_messages(messages=message_item) response = mos.chat(query, user_id=user_id) message_item = ScheduleMessageItem( user_id=user_id, mem_cube_id=mem_cube_id, label=ANSWER_LABEL, mem_cube=mos.mem_cubes[mem_cube_id], content=response, timestamp=datetime.now(), ) mos.mem_scheduler.submit_messages(messages=message_item) print(f"Query:\n {query}\n\nAnswer:\n {response}") show_web_logs(mos.mem_scheduler) mos.mem_scheduler.stop() if __name__ == "__main__": run_with_automatic_scheduler_init() run_with_manual_scheduler_init()