MemScheduler

from memos.api import handlers
from memos.api.handlers.base_handler import HandlerDependencies
from memos.mem_scheduler.base_scheduler import BaseScheduler
from memos.mem_scheduler.utils.status_tracker import TaskStatusTracker

# ... 其他导入 ...

# 1. 初始化所有服务器组件 (包括 DB, LLM, Memory, Scheduler)
# init_server() 会读取环境变量并初始化全局单例组件
components = handlers.init_server()

# Create dependency container for handlers
dependencies = HandlerDependencies.from_init_server(components)

# Initialize handlers...
# search_handler = SearchHandler(dependencies)
# ...

# 2. 从组件字典中获取调度器实例
# 调度器在 init_server 内部已经被初始化并启动（如果启用了的话）
mem_scheduler: BaseScheduler = components["mem_scheduler"]

# 3. 用户还可以在components中获取其他调度相关组件 (可选，用于自定义任务处理)
# redis_client 用于直接操作 Redis 或监控任务状态
redis_client = components["redis_client"]
# ...

import uuid
from datetime import datetime

# 1. 导入必要的类型定义和调度器实例
# 注意：mem_scheduler 需要从 server_router 导入，因为它是一个全局单例
from memos.api.routers.server_router import mem_scheduler
from memos.mem_scheduler.schemas.message_schemas import ScheduleMessageItem

# 定义一个自定义的任务标签
MY_TASK_LABEL = "MY_CUSTOM_TASK"


# 定义处理器函数
def my_task_handler(messages: list[ScheduleMessageItem]):
    """
    处理自定义任务的函数
    """
    for msg in messages:
        print(f"⚡️ [Handler] 收到任务: {msg.item_id}")
        print(f"📦 内容: {msg.content}")
        # 在这里执行你的业务逻辑，例如：调用 LLM、写数据库、触发其他任务等


# 2. 注册处理器到调度器
# 这一步将您的自定义逻辑挂载到调度系统中
mem_scheduler.register_handlers({
    MY_TASK_LABEL: my_task_handler
})

# 3. 提交任务
task = ScheduleMessageItem(
    item_id=str(uuid.uuid4()),
    user_id="user_123",
    mem_cube_id="cube_001",
    label=MY_TASK_LABEL,
    content="这是一条测试消息",
    timestamp=datetime.now()
)

# 如果调度器未启动，这里会直接放入队列等待处理（如果是 Redis 队列）
# 或者在本地队列模式下可能需要先调用 mem_scheduler.start()
mem_scheduler.submit_messages([task])

print(f"Task submitted: {task.item_id}")

# 防止调度器主进程提前退出
time.sleep(10)

import asyncio
import json
import os
import sys
import time
from pathlib import Path

# --- 环境准备 ---
# 1. 设置项目根目录到 sys.path，确保能导入 memos 模块
FILE_PATH = Path(__file__).absolute()
BASE_DIR = FILE_PATH.parent.parent.parent
sys.path.insert(0, str(BASE_DIR))

# 2. 设置必要的环境变量 (模拟 .env 配置)
os.environ["ENABLE_CHAT_API"] = "true"
os.environ["MOS_ENABLE_SCHEDULER"] = "true"
# 决定使用 Redis 还是 Local 队列
os.environ["MOS_SCHEDULER_USE_REDIS_QUEUE"] = "false" 

# --- 导入组件 ---
# 注意：导入 server_router 会触发组件初始化，确保环境变量在此之前设置好
from memos.api.product_models import APIADDRequest, ChatPlaygroundRequest
from memos.api.routers.server_router import (
    add_handler,
    chat_stream_playground,
    mem_scheduler,  # 这里的 mem_scheduler 已经是初始化好的单例
)
from memos.log import get_logger
from memos.mem_scheduler.schemas.message_schemas import ScheduleMessageItem
from memos.mem_scheduler.schemas.task_schemas import (
    MEM_UPDATE_TASK_LABEL,
    QUERY_TASK_LABEL,
)

logger = get_logger(__name__)

# 全局变量用于演示记忆检索结果
working_memories = []

# --- 自定义处理器 ---

def custom_query_handler(messages: list[ScheduleMessageItem]):
    """
    处理用户查询消息：
    1. 打印查询内容
    2. 将消息转换为 MEM_UPDATE 任务，触发记忆检索/更新流程
    """
    for msg in messages:
        print(f"\n[Scheduler 🟢] 收到用户查询: {msg.content}")
        
        # 复制消息并将标签改为 MEM_UPDATE，这是一种常见的“任务链”模式
        new_msg = msg.model_copy(update={"label": MEM_UPDATE_TASK_LABEL})
        
        # 提交新任务回调度器
        mem_scheduler.submit_messages([new_msg])


def custom_mem_update_handler(messages: list[ScheduleMessageItem]):
    """
    处理记忆更新任务：
    1. 使用检索器 (Retriever) 查找相关记忆
    2. 更新全局的工作记忆列表
    """
    global working_memories
    search_args = {}
    top_k = 2
    
    for msg in messages:
        print(f"[Scheduler 🔵] 正在为查询检索记忆...")
        # 调用核心检索功能
        results = mem_scheduler.retriever.search(
            query=msg.content,
            user_id=msg.user_id,
            mem_cube_id=msg.mem_cube_id,
            mem_cube=mem_scheduler.current_mem_cube,
            top_k=top_k,
            method=mem_scheduler.search_method,
            search_args=search_args,
        )
        
        # 模拟工作记忆的更新
        working_memories.extend(results)
        working_memories = working_memories[-5:] # 保持最新的5条
        
        for mem in results:
            # 打印检索到的记忆片段
            print(f"  ↳ [Memory Found]: {mem.memory[:50]}...")

# --- 模拟业务数据 ---

def get_mock_data():
    """生成模拟对话数据"""
    conversations = [
        {"role": "user", "content": "I just adopted a golden retriever puppy named Max."},
        {"role": "assistant", "content": "That's exciting! Max is a great name."},
        {"role": "user", "content": "He loves peanut butter treats but I am allergic to nuts."},
        {"role": "assistant", "content": "Noted. Peanut butter for Max, no nuts for you."},
    ]
    
    questions = [
        {"question": "What is my dog's name?", "category": "Pet"},
        {"question": "What am I allergic to?", "category": "Allergy"},
    ]
    return conversations, questions

# --- 主流程 ---

async def run_demo():
    print("==== MemScheduler Demo Start ====")
    conversations, questions = get_mock_data()

    user_id = "demo_user_001"
    mem_cube_id = "cube_demo_001"

    print(f"1. 初始化用户记忆库 ({user_id})...")
    # 使用 API Handler 添加初始记忆 (同步模式)
    add_req = APIADDRequest(
        user_id=user_id,
        writable_cube_ids=[mem_cube_id],
        messages=conversations,
        async_mode="sync", 
    )
    add_handler.handle_add_memories(add_req)
    print("   记忆添加完成。")

    print("\n2. 开始对话测试 (并在后台触发调度任务)...")
    for item in questions:
        query = item["question"]
        print(f"\n>> User: {query}")

        # 发起聊天请求
        chat_req = ChatPlaygroundRequest(
            user_id=user_id,
            query=query,
            readable_cube_ids=[mem_cube_id],
            writable_cube_ids=[mem_cube_id],
        )
        
        # 获取流式响应
        response = chat_stream_playground(chat_req)
        
        # 处理流式输出 (简化版)
        full_answer = ""
        buffer = ""
        async for chunk in response.body_iterator:
            if isinstance(chunk, bytes):
                chunk = chunk.decode("utf-8")
            buffer += chunk
            while "\n\n" in buffer:
                msg, buffer = buffer.split("\n\n", 1)
                for line in msg.split("\n"):
                    if line.startswith("data: "):
                        try:
                            data = json.loads(line[6:])
                            if data.get("type") == "text":
                                full_answer += data["data"]
                        except: pass
                        
        print(f">> AI: {full_answer}")
        
        # 等待一小会儿让后台调度器处理任务并打印日志
        await asyncio.sleep(1)

if __name__ == "__main__":
    # 1. 注册我们的自定义 Handler
    # 这会覆盖或添加到默认的调度逻辑中
    mem_scheduler.register_handlers(
        {
            QUERY_TASK_LABEL: custom_query_handler,
            MEM_UPDATE_TASK_LABEL: custom_mem_update_handler,
        }
    )
    
    # 2. 确保调度器已启动
    if not mem_scheduler._running:
        mem_scheduler.start()

    try:
        asyncio.run(run_demo())
    except KeyboardInterrupt:
        pass
    finally:
        # 防止调度器主进程提前退出
        time.sleep(10)

        print("\n==== 停止调度器 ====")
        mem_scheduler.stop()

from pathlib import Path
from time import sleep

from memos.api.routers.server_router import mem_scheduler
from memos.mem_scheduler.schemas.message_schemas import ScheduleMessageItem


# 调试：打印调度器配置
print("=== Scheduler Configuration Debug ===")
print(f"Scheduler type: {type(mem_scheduler).__name__}")
print(f"Config: {mem_scheduler.config}")
print(f"use_redis_queue: {mem_scheduler.use_redis_queue}")
print(f"Queue type: {type(mem_scheduler.memos_message_queue).__name__}")
print(f"Queue maxsize: {getattr(mem_scheduler.memos_message_queue, 'maxsize', 'N/A')}")
print("=====================================\n")

queue = mem_scheduler.memos_message_queue


# 定义处理函数
def my_test_handler(messages: list[ScheduleMessageItem]):
    print(f"My test handler received {len(messages)} messages: {[one.item_id for one in messages]}")
    for msg in messages:
        # 根据 task_id 创建文件（使用 item_id 作为数字 ID 0..99）
        task_id = str(msg.item_id)
        file_path = tmp_dir / f"{task_id}.txt"
        try:
            sleep(5)
            file_path.write_text(f"Task {task_id} processed.\n")
            print(f"writing {file_path} done")
        except Exception as e:
            print(f"Failed to write {file_path}: {e}")


def submit_tasks():
    mem_scheduler.memos_message_queue.clear()

    # 创建 100 条消息（task_id 0..99）
    users = ["user_A", "user_B"]
    messages_to_send = [
        ScheduleMessageItem(
            item_id=str(i),
            user_id=users[i % 2],
            mem_cube_id="test_mem_cube",
            label=TEST_HANDLER_LABEL,
            content=f"Create file for task {i}",
        )
        for i in range(100)
    ]
    # 批量提交消息并打印完成信息
    print(f"Submitting {len(messages_to_send)} messages to the scheduler...")
    mem_scheduler.memos_message_queue.submit_messages(messages_to_send)
    print(f"Task submission done! tasks in queue: {mem_scheduler.get_tasks_status()}")


# 注册处理函数
TEST_HANDLER_LABEL = "test_handler"
mem_scheduler.register_handlers({TEST_HANDLER_LABEL: my_test_handler})

# 5秒重启
mem_scheduler.orchestrator.tasks_min_idle_ms[TEST_HANDLER_LABEL] = 5_000

tmp_dir = Path("./tmp")
tmp_dir.mkdir(exist_ok=True)

# 测试停止并重启：如果 tmp 中已有 >1 个文件，跳过提交并打印信息
existing_count = len(list(Path("tmp").glob("*.txt"))) if Path("tmp").exists() else 0
if existing_count > 1:
    print(f"Skip submission: found {existing_count} files in tmp (>1), continue processing")
else:
    submit_tasks()

# 6. 等待直到 tmp 有 100 个文件或超时
poll_interval = 1
expected = 100
tmp_dir = Path("tmp")
tasks_status = mem_scheduler.get_tasks_status()
mem_scheduler.print_tasks_status(tasks_status=tasks_status)
while (
    mem_scheduler.get_tasks_status()["remaining"] != 0
    or mem_scheduler.get_tasks_status()["running"] != 0
):
    count = len(list(tmp_dir.glob("*.txt"))) if tmp_dir.exists() else 0
    tasks_status = mem_scheduler.get_tasks_status()
    mem_scheduler.print_tasks_status(tasks_status=tasks_status)
    print(f"[Monitor] Files in tmp: {count}/{expected}")
    sleep(poll_interval)
print(f"[Result] Final files in tmp: {len(list(tmp_dir.glob('*.txt')))})")

# 7. 停止调度器
sleep(20)
print("Stopping the scheduler...")
mem_scheduler.stop()