Skip to content

Conversation

@CaralHsi
Copy link
Collaborator

@CaralHsi CaralHsi commented Jan 28, 2026

Description

  1. In playground, we use list index in empty list, thus new user with no memories can not chat with the system.
  2. Using an IN clause in the filter for search_by_embedding can lead to significant performance degradation. Before further optimization, such calls should be removed to avoid heavy database load during deduplication in the add stage.

Related Issue (Required): Fixes #969

Type of change

Please delete options that are not relevant.

  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • Refactor (does not change functionality, e.g. code style improvements, linting)
  • Documentation update

How Has This Been Tested?

import asyncio
import json
import os
import random as _random
import socket
import time

from memos.api import handlers
from memos.api.handlers.add_handler import AddHandler
from memos.api.handlers.base_handler import HandlerDependencies
from memos.api.handlers.chat_handler import ChatHandler
from memos.api.handlers.search_handler import SearchHandler
from memos.api.product_models import ChatPlaygroundRequest
from memos.log import get_logger

logger = get_logger(__name__)

# Instance ID for identifying this server instance in logs and responses
INSTANCE_ID = f"{socket.gethostname()}:{os.getpid()}:{_random.randint(1000, 9999)}"

# Initialize all server components
print("=" * 80)
print("初始化服务组件...")
print("=" * 80)
components = handlers.init_server()

# 获取关键组件
mem_reader = components["mem_reader"]
graph_db = components["graph_db"]
embedder = components["embedder"]
chat_llms = components.get("chat_llms")

if chat_llms is None:
    raise ValueError(
        "chat_llms is not initialized. Please set ENABLE_CHAT_API=true environment variable."
    )

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

# Initialize all handlers with dependency injection
search_handler = SearchHandler(dependencies)
add_handler = AddHandler(dependencies)
chat_handler = ChatHandler(
    dependencies,
    chat_llms,
    search_handler,
    add_handler,
    online_bot=components.get("online_bot"),
)


def chat_stream_playground(chat_req: ChatPlaygroundRequest):
    """
    Chat with MemOS playground for a specific user.
    This endpoint uses the class-based ChatHandler for better code organization.
    Returns SSE stream.
    """
    return chat_handler.handle_chat_stream_playground(chat_req)


async def process_stream_response(response):
    buffer = ""
    full_text = ""
    
    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: "):
                    json_str = line[6:]
                    try:
                        data = json.loads(json_str)
                        msg_type = data.get("type")
                        
                        if msg_type == "status":
                            print(f"[Status] {data.get('data')}")
                        elif msg_type == "reference":
                            print(f"[Reference] {len(data.get('data', []))} memories found")
                        elif msg_type == "pref_md_string":
                            print(f"[Preference] {data.get('data', '')[:100]}...")
                        elif msg_type == "text":
                            text_chunk = data.get("data", "")
                            full_text += text_chunk
                            print(text_chunk, end="", flush=True)
                        elif msg_type == "reasoning":
                            print(f"[Reasoning] {data.get('data', '')[:50]}...")
                        elif msg_type == "time":
                            time_data = data.get("data", {})
                            print(f"\n[Time] Total: {time_data.get('total_time')}s, Speed improvement: {time_data.get('speed_improvement')}")
                        elif msg_type == "suggestion":
                            suggestions = data.get("data", [])
                            print(f"\n[Suggestions] {suggestions}")
                        elif msg_type == "end":
                            print("\n[End] Stream completed")
                        elif msg_type == "error":
                            print(f"\n[Error] {data.get('content', data.get('data', 'Unknown error'))}")
                    except json.JSONDecodeError:
                        pass  # 忽略无效的 JSON
    
    return full_text


if __name__ == '__main__':
    user_id = "user_id_test_chat_stream_playground"
    cube_id = 'cube_id_test_chat_stream_playground'
    req = ChatPlaygroundRequest.model_validate({
        "user_id": user_id,
        "query": "你好,请介绍一下你自己",
        "readable_cube_ids": [cube_id],
        "writable_cube_ids": [cube_id],
        "session_id": "default_session",
        "history": [],
        "internet_search": False,
        "include_preference": True,
        "pref_top_k": 6,
        "beginner_guide_step": None,
    })
    init_time = time.time()
    res = chat_stream_playground(req)
    
    print(f"Response type: {type(res)}")
    print("=" * 80)
    print("开始接收流式响应...")
    print("=" * 80)

    full_response = asyncio.run(process_stream_response(res))
    
    print("\n" + "=" * 80)
    print(f"完整响应长度: {len(full_response)} 字符")
    print(f"总耗时: {time.time() - init_time:.2f} 秒")
    print("=" * 80)
  • Unit Test
  • Test Script Or Test Steps (please provide)
  • Pipeline Automated API Test (please provide)

Checklist

  • I have performed a self-review of my own code | 我已自行检查了自己的代码
  • I have commented my code in hard-to-understand areas | 我已在难以理解的地方对代码进行了注释
  • I have added tests that prove my fix is effective or that my feature works | 我已添加测试以证明我的修复有效或功能正常
  • I have created related documentation issue/PR in MemOS-Docs (if applicable) | 我已在 MemOS-Docs 中创建了相关的文档 issue/PR(如果适用)
  • I have linked the issue to this PR (if applicable) | 我已将 issue 链接到此 PR(如果适用)
  • I have mentioned the person who will review this PR | 我已提及将审查此 PR 的人

Reviewer Checklist

  • closes #xxxx (Replace xxxx with the GitHub issue number)
  • Made sure Checks passed
  • Tests have been provided

@CaralHsi CaralHsi marked this pull request as ready for review January 28, 2026 13:04
@CaralHsi CaralHsi closed this Jan 28, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

fix: system busy in playground

1 participant