From 5b342f89f14aaa48d48830695ffb0468ad1aa1d8 Mon Sep 17 00:00:00 2001 From: dingliang <2650876010@qq.com> Date: Wed, 8 Jul 2026 16:13:46 +0800 Subject: [PATCH 01/15] feat(api): align OpenMem v1 SDK with cloud API --- src/memos/api/client.py | 263 ++++++++++++++++++++++--- tests/api/test_client.py | 406 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 647 insertions(+), 22 deletions(-) create mode 100644 tests/api/test_client.py diff --git a/src/memos/api/client.py b/src/memos/api/client.py index 818ce5e0d..1ace6ec8f 100644 --- a/src/memos/api/client.py +++ b/src/memos/api/client.py @@ -65,6 +65,32 @@ def _validate_required_params(self, **params): if not param_value: raise ValueError(f"{param_name} is required") + def _validate_profile_subject(self, user_id: str | None, agent_id: str | None) -> None: + if bool(user_id) == bool(agent_id): + raise ValueError("exactly one of user_id or agent_id is required") + + def _post_json_dict( + self, endpoint: str, payload: dict[str, Any], operation: str + ) -> dict[str, Any] | None: + url = f"{self.base_url}/{endpoint}" + for retry in range(MAX_RETRY_COUNT): + try: + response = requests.post( + url, data=json.dumps(payload), headers=self.headers, timeout=30 + ) + response.raise_for_status() + return response.json() + except Exception as e: + logger.error( + "Failed to %s (retry %s/%s): %s", + operation, + retry + 1, + MAX_RETRY_COUNT, + e, + ) + if retry == MAX_RETRY_COUNT - 1: + raise + def get_message( self, user_id: str, @@ -102,7 +128,7 @@ def get_message( def add_message( self, messages: list[dict[str, Any]], - user_id: str, + user_id: str | list[str], conversation_id: str, info: dict[str, Any] | None = None, source: str | None = None, @@ -112,6 +138,7 @@ def add_message( tags: list[str] | None = None, allow_public: bool = False, allow_knowledgebase_ids: list[str] | None = None, + allow_memory_view: list[str] | None = None, ) -> MemOSAddResponse | None: """Add message""" # Validate required parameters @@ -130,8 +157,9 @@ def add_message( "agent_id": agent_id, "allow_public": allow_public, "allow_knowledgebase_ids": allow_knowledgebase_ids, + "allow_memory_view": allow_memory_view, "tags": tags, - "asyncMode": async_mode, + "async_mode": async_mode, } for retry in range(MAX_RETRY_COUNT): try: @@ -151,7 +179,8 @@ def search_memory( self, query: str, user_id: str, - conversation_id: str, + conversation_id: str | None = None, + agent_id: str | None = None, memory_limit_number: int = 6, include_preference: bool = True, knowledgebase_ids: list[str] | None = None, @@ -160,6 +189,11 @@ def search_memory( include_tool_memory: bool = False, preference_limit_number: int = 6, tool_memory_limit_number: int = 6, + relativity: float | None = None, + include_skill: bool = False, + skill_limit_number: int = 6, + include_memory_view: list[str] | None = None, + context_format: str = "memory", ) -> MemOSSearchResponse | None: """Search memories""" # Validate required parameters @@ -170,12 +204,18 @@ def search_memory( "query": query, "user_id": user_id, "conversation_id": conversation_id, + "agent_id": agent_id, "memory_limit_number": memory_limit_number, "include_preference": include_preference, "knowledgebase_ids": knowledgebase_ids, "filter": filter, "preference_limit_number": preference_limit_number, "tool_memory_limit_number": tool_memory_limit_number, + "relativity": relativity, + "include_skill": include_skill, + "skill_limit_number": skill_limit_number, + "include_memory_view": include_memory_view, + "context_format": context_format, "source": source, "include_tool_memory": include_tool_memory, } @@ -195,16 +235,28 @@ def search_memory( raise def get_memory( - self, user_id: str, include_preference: bool = True, page: int = 1, size: int = 10 + self, + user_id: str | None = None, + include_preference: bool = True, + page: int = 1, + size: int = 10, + agent_id: str | None = None, + include_tool_memory: bool = True, + include_memory_view: list[str] | None = None, + filter: dict[str, Any] | None = None, ) -> MemOSGetMemoryResponse | None: """get memories""" # Validate required parameters - self._validate_required_params(include_preference=include_preference, user_id=user_id) + self._validate_profile_subject(user_id, agent_id) url = f"{self.base_url}/get/memory" payload = { "include_preference": include_preference, "user_id": user_id, + "agent_id": agent_id, + "include_tool_memory": include_tool_memory, + "include_memory_view": include_memory_view, + "filter": filter, "page": page, "size": size, } @@ -313,7 +365,7 @@ def add_knowledgebase_file_json( raise def add_knowledgebase_file_form( - self, knowledgebase_id: str, files: list[str] + self, knowledgebase_id: str, files: list[str], type: str | None = None ) -> MemOSAddKnowledgebaseFileResponse | None: """ add knowledgebase-file from form @@ -321,12 +373,12 @@ def add_knowledgebase_file_form( # Validate required parameters self._validate_required_params(knowledgebase_id=knowledgebase_id, files=files) - def build_file_form_param(file_path): + def build_file_form_param(file_path: str): """ form-Automatically generate the structure required for the `files` parameter in requests based on the local file path """ if not os.path.isfile(file_path): - logger.warning(f"File {file_path} does not exist") + logger.warning("File %s does not exist", file_path) return None filename = os.path.basename(file_path) @@ -335,31 +387,47 @@ def build_file_form_param(file_path): mime_type = "application/octet-stream" return ("file", (filename, open(file_path, "rb"), mime_type)) + def build_file_form_params() -> list: + file_params = [ + file_param + for file_path in files + if (file_param := build_file_form_param(file_path)) is not None + ] + if not file_params: + raise ValueError("files must contain at least one valid file path") + return file_params + url = f"{self.base_url}/add/knowledgebase-file" payload = { "knowledgebase_id": knowledgebase_id, } + if type is not None: + payload["type"] = type headers = { "Authorization": f"Token {self.api_key}", } for retry in range(MAX_RETRY_COUNT): + file_params = [] try: + file_params = build_file_form_params() response = requests.post( url, params=payload, headers=headers, timeout=30, - files=[build_file_form_param(file_path) for file_path in files], + files=file_params, ) response.raise_for_status() response_data = response.json() - print(response_data) return MemOSAddKnowledgebaseFileResponse(**response_data) except Exception as e: logger.error(f"Failed to add knowledgebase-file form (retry {retry + 1}/3): {e}") if retry == MAX_RETRY_COUNT - 1: raise + finally: + for file_param in file_params: + file_param[1][1].close() def delete_knowledgebase_file( self, file_ids: list[str] @@ -390,17 +458,27 @@ def delete_knowledgebase_file( raise def get_knowledgebase_file( - self, file_ids: list[str] + self, + file_ids: list[str] | None = None, + knowledgebase_id: str | None = None, + type: str | None = None, + page: int | None = None, + page_size: int | None = None, ) -> MemOSGetKnowledgebaseFileResponse | None: """ get knowledgebase-file """ # Validate required parameters - self._validate_required_params(file_ids=file_ids) + if bool(file_ids) == bool(knowledgebase_id): + raise ValueError("exactly one of file_ids or knowledgebase_id is required") url = f"{self.base_url}/get/knowledgebase-file" payload = { "file_ids": file_ids, + "knowledgebase_id": knowledgebase_id, + "type": type, + "page": page, + "page_size": page_size, } for retry in range(MAX_RETRY_COUNT): @@ -486,17 +564,43 @@ def add_feedback( raise def delete_memory( - self, user_ids: list[str], memory_ids: list[str] + self, + user_ids: list[str] | None = None, + memory_ids: list[str] | None = None, + *, + user_id: str | None = None, + agent_id: str | None = None, + filter: dict[str, Any] | None = None, + memory_type: str | None = None, ) -> MemOSDeleteMemoryResponse | None: """delete_memory memories""" - # Validate required parameters - self._validate_required_params(user_ids=user_ids, memory_ids=memory_ids) + if user_id is None and user_ids: + if len(user_ids) != 1 and not memory_ids: + raise ValueError("current API supports a single user_id, not multiple user_ids") + if not memory_ids: + user_id = user_ids[0] + + delete_modes = [ + bool(memory_ids), + bool(user_id), + bool(agent_id), + filter is not None, + ] + if sum(delete_modes) != 1: + raise ValueError("exactly one delete condition is required") url = f"{self.base_url}/delete/memory" - payload = { - "user_ids": user_ids, - "memory_ids": memory_ids, - } + payload: dict[str, Any] = {} + if memory_ids: + payload["memory_ids"] = memory_ids + if user_id: + payload["user_id"] = user_id + if agent_id: + payload["agent_id"] = agent_id + if filter is not None: + payload["filter"] = filter + if memory_type is not None: + payload["memory_type"] = memory_type for retry in range(MAX_RETRY_COUNT): try: @@ -512,6 +616,111 @@ def delete_memory( if retry == MAX_RETRY_COUNT - 1: raise + def update_memory( + self, + memory_id: str, + content: str | None = None, + title: str | None = None, + status: str | None = None, + ) -> dict[str, Any] | None: + """Update an existing memory.""" + self._validate_required_params(memory_id=memory_id) + if not content and not title and not status: + raise ValueError("content, title or status is required") + + payload = { + "memory_id": memory_id, + "content": content, + "title": title, + "status": status, + } + return self._post_json_dict("update/memory", payload, "update memory") + + def extract_memory( + self, + messages: list[dict[str, Any]], + extraction_types: list[str] | None = None, + model: str | None = None, + ) -> dict[str, Any] | None: + """Extract memory candidates from conversation messages.""" + self._validate_required_params(messages=messages) + + payload = { + "messages": messages, + "extraction_types": extraction_types, + "model": model, + } + return self._post_json_dict("extract/memory", payload, "extract memory") + + def rerank( + self, + query: str, + documents: list[str], + model: str | None = None, + top_n: int | None = None, + ) -> dict[str, Any] | None: + """Rerank documents for a query.""" + self._validate_required_params(query=query, documents=documents) + if top_n is not None and top_n <= 0: + raise ValueError("top_n must be greater than 0") + + payload = { + "query": query, + "documents": documents, + "model": model, + "top_n": top_n, + } + return self._post_json_dict("rerank", payload, "rerank documents") + + def bind_profile_template(self, bind_list: list[dict[str, Any]]) -> dict[str, Any] | None: + """Bind profile templates to user or agent subjects.""" + self._validate_required_params(bind_list=bind_list) + + payload = { + "bind_list": bind_list, + } + return self._post_json_dict("bind/profile_template", payload, "bind profile template") + + def edit_profile( + self, + profile_template_id: str, + user_id: str | None = None, + agent_id: str | None = None, + metadata: dict[str, Any] | None = None, + remove_fields: list[str] | None = None, + ) -> dict[str, Any] | None: + """Edit a profile instance.""" + self._validate_required_params(profile_template_id=profile_template_id) + self._validate_profile_subject(user_id, agent_id) + if metadata is None and not remove_fields: + raise ValueError("metadata or remove_fields is required") + + payload = { + "user_id": user_id, + "agent_id": agent_id, + "profile_template_id": profile_template_id, + "metadata": metadata, + "remove_fields": remove_fields, + } + return self._post_json_dict("edit/profile", payload, "edit profile") + + def delete_profile( + self, + profile_template_id: str, + user_id: str | None = None, + agent_id: str | None = None, + ) -> dict[str, Any] | None: + """Delete a profile instance.""" + self._validate_required_params(profile_template_id=profile_template_id) + self._validate_profile_subject(user_id, agent_id) + + payload = { + "user_id": user_id, + "agent_id": agent_id, + "profile_template_id": profile_template_id, + } + return self._post_json_dict("delete/profile", payload, "delete profile") + def chat( self, user_id: str, @@ -524,20 +733,25 @@ def chat( system_prompt: str | None = None, model_name: str | None = None, knowledgebase_ids: list[str] | None = None, - filter: dict[str:Any] | None = None, - add_message_on_answer: bool = False, + filter: dict[str, Any] | None = None, + add_message_on_answer: bool = True, app_id: str | None = None, agent_id: str | None = None, async_mode: bool = True, tags: list[str] | None = None, - info: dict[str:Any] | None = None, + info: dict[str, Any] | None = None, allow_public: bool = False, + allow_knowledgebase_ids: list[str] | None = None, max_tokens: int = 8192, temperature: float | None = None, top_p: float | None = None, include_preference: bool = True, preference_limit_number: int = 6, memory_limit_number: int = 6, + stream: bool = False, + include_tool_memory: bool = False, + tool_memory_limit_number: int = 6, + relativity: float | None = None, ) -> MemOSChatResponse | None: """chat""" # Validate required parameters @@ -565,12 +779,17 @@ def chat( "tags": tags, "info": info, "allow_public": allow_public, + "allow_knowledgebase_ids": allow_knowledgebase_ids, "max_tokens": max_tokens, "temperature": temperature, "top_p": top_p, "include_preference": include_preference, "preference_limit_number": preference_limit_number, "memory_limit_number": memory_limit_number, + "stream": stream, + "include_tool_memory": include_tool_memory, + "tool_memory_limit_number": tool_memory_limit_number, + "relativity": relativity, } for retry in range(MAX_RETRY_COUNT): diff --git a/tests/api/test_client.py b/tests/api/test_client.py new file mode 100644 index 000000000..61724b0d4 --- /dev/null +++ b/tests/api/test_client.py @@ -0,0 +1,406 @@ +import json +import sys +import types + +from pathlib import Path +from typing import Any + +import pytest + + +SRC_DIR = Path(__file__).resolve().parents[2] / "src" / "memos" + + +def _install_memos_package_stub() -> None: + if "memos" not in sys.modules: + memos_pkg = types.ModuleType("memos") + memos_pkg.__path__ = [str(SRC_DIR)] + sys.modules["memos"] = memos_pkg + + if "memos.api" not in sys.modules: + api_pkg = types.ModuleType("memos.api") + api_pkg.__path__ = [str(SRC_DIR / "api")] + sys.modules["memos.api"] = api_pkg + sys.modules["memos"].api = api_pkg + + +def _load_client_module() -> Any: + _install_memos_package_stub() + + import memos.api.client as client_module + + return client_module + + +class DummyResponse: + def __init__(self, payload: dict): + self.payload = payload + + def raise_for_status(self) -> None: + return None + + def json(self) -> dict: + return self.payload + + +def _response_for(url: str) -> dict: + if url.endswith("/add/message"): + return { + "code": 200, + "message": "ok", + "data": {"success": True, "task_id": "task-1", "status": "completed"}, + } + if url.endswith("/search/memory"): + return {"code": 200, "message": "ok", "data": {"memory_detail_list": []}} + if url.endswith("/get/memory"): + return {"code": 200, "message": "ok", "data": {"memory_detail_list": []}} + if url.endswith("/get/knowledgebase-file"): + return {"code": 200, "message": "ok", "data": {"file_detail_list": []}} + if url.endswith("/delete/memory"): + return {"code": 200, "message": "ok", "data": {"success": True}} + if url.endswith("/chat"): + return {"code": 200, "message": "ok", "data": {"response": "answer"}} + if url.endswith("/add/knowledgebase-file"): + return {"code": 200, "message": "ok", "data": []} + if url.endswith("/update/memory"): + return {"code": 200, "message": "ok", "data": {"success": True}} + if url.endswith("/extract/memory"): + return { + "code": 200, + "message": "ok", + "data": { + "success": True, + "memory_detail_list": [], + "preference_detail_list": [], + }, + } + if url.endswith("/rerank"): + return {"code": 200, "message": "ok", "data": {"id": "rerank-1", "results": []}} + if url.endswith("/bind/profile_template"): + return {"code": 200, "message": "ok", "data": {"success": True}} + if url.endswith("/edit/profile"): + return {"code": 200, "message": "ok", "data": {"success": True}} + if url.endswith("/delete/profile"): + return {"code": 200, "message": "ok", "data": {"success": True}} + raise AssertionError(f"Unexpected URL: {url}") + + +@pytest.fixture +def client_module() -> Any: + return _load_client_module() + + +@pytest.fixture +def posted_requests(monkeypatch, client_module): + calls: list[dict] = [] + + def fake_post(url: str, **kwargs): + calls.append({"url": url, **kwargs}) + return DummyResponse(_response_for(url)) + + monkeypatch.setattr(client_module.requests, "post", fake_post) + return calls + + +@pytest.fixture +def client(client_module) -> Any: + return client_module.MemOSClient(api_key="test-key", base_url="https://example.test/openmem/v1") + + +def _json_payload(call: dict) -> dict: + return json.loads(call["data"]) + + +def test_add_message_uses_snake_case_async_mode_and_memory_view( + client: Any, posted_requests: list[dict] +) -> None: + client.add_message( + messages=[{"role": "user", "content": "hello"}], + user_id="user-1", + conversation_id="conversation-1", + async_mode=False, + allow_memory_view=["kb-1"], + ) + + payload = _json_payload(posted_requests[0]) + + assert payload["async_mode"] is False + assert "asyncMode" not in payload + assert payload["allow_memory_view"] == ["kb-1"] + + +def test_search_memory_sends_updated_existing_request_fields( + client: Any, posted_requests: list[dict] +) -> None: + client.search_memory( + query="hello", + user_id="user-1", + agent_id="agent-1", + relativity=0.2, + include_skill=True, + skill_limit_number=4, + include_memory_view=["kb-1"], + context_format="json", + ) + + payload = _json_payload(posted_requests[0]) + + assert payload["conversation_id"] is None + assert payload["agent_id"] == "agent-1" + assert payload["relativity"] == 0.2 + assert payload["include_skill"] is True + assert payload["skill_limit_number"] == 4 + assert payload["include_memory_view"] == ["kb-1"] + assert payload["context_format"] == "json" + + +def test_get_memory_can_scope_by_agent_and_include_updated_filters( + client: Any, posted_requests: list[dict] +) -> None: + memory_filter = {"and": [{"memory_type": "LongTermMemory"}]} + + client.get_memory( + user_id=None, + agent_id="agent-1", + include_tool_memory=False, + include_memory_view=["kb-1"], + filter=memory_filter, + page=2, + size=20, + ) + + payload = _json_payload(posted_requests[0]) + + assert payload["user_id"] is None + assert payload["agent_id"] == "agent-1" + assert payload["include_tool_memory"] is False + assert payload["include_memory_view"] == ["kb-1"] + assert payload["filter"] == memory_filter + assert payload["page"] == 2 + assert payload["size"] == 20 + + +def test_get_memory_rejects_multiple_subjects(client: Any) -> None: + with pytest.raises(ValueError, match="exactly one of user_id or agent_id"): + client.get_memory(user_id="user-1", agent_id="agent-1") + + +def test_get_knowledgebase_file_supports_listing_by_knowledgebase( + client: Any, posted_requests: list[dict] +) -> None: + client.get_knowledgebase_file( + knowledgebase_id="kb-1", + type="doc", + page=2, + page_size=50, + ) + + payload = _json_payload(posted_requests[0]) + + assert payload == { + "file_ids": None, + "knowledgebase_id": "kb-1", + "type": "doc", + "page": 2, + "page_size": 50, + } + + +def test_delete_memory_keeps_legacy_memory_id_call_but_sends_current_contract( + client: Any, posted_requests: list[dict] +) -> None: + client.delete_memory(user_ids=["legacy-user"], memory_ids=["memory-1"]) + + payload = _json_payload(posted_requests[0]) + + assert payload == {"memory_ids": ["memory-1"]} + + +def test_delete_memory_supports_quick_delete_by_user_id( + client: Any, posted_requests: list[dict] +) -> None: + client.delete_memory(user_id="user-1") + + payload = _json_payload(posted_requests[0]) + + assert payload == {"user_id": "user-1"} + + +def test_chat_sends_updated_existing_request_fields( + client: Any, posted_requests: list[dict] +) -> None: + client.chat( + user_id="user-1", + conversation_id="conversation-1", + query="hello", + stream=True, + allow_knowledgebase_ids=["kb-1"], + include_tool_memory=True, + tool_memory_limit_number=3, + relativity=0.1, + ) + + payload = _json_payload(posted_requests[0]) + + assert payload["stream"] is True + assert payload["allow_knowledgebase_ids"] == ["kb-1"] + assert payload["include_tool_memory"] is True + assert payload["tool_memory_limit_number"] == 3 + assert payload["relativity"] == 0.1 + assert payload["add_message_on_answer"] is True + + +def test_add_knowledgebase_file_form_sends_type_and_closes_files( + client: Any, posted_requests: list[dict], tmp_path +) -> None: + file_path = tmp_path / "note.txt" + file_path.write_text("hello", encoding="utf-8") + + client.add_knowledgebase_file_form( + knowledgebase_id="kb-1", + files=[str(file_path)], + type="doc", + ) + + call = posted_requests[0] + uploaded_file = call["files"][0][1][1] + + assert call["params"] == {"knowledgebase_id": "kb-1", "type": "doc"} + assert uploaded_file.closed + + +def test_update_memory_sends_selected_fields(client: Any, posted_requests: list[dict]) -> None: + response = client.update_memory( + memory_id="memory-1", + content="new content", + title="new title", + status="activated", + ) + + payload = _json_payload(posted_requests[0]) + + assert posted_requests[0]["url"].endswith("/update/memory") + assert payload == { + "memory_id": "memory-1", + "content": "new content", + "title": "new title", + "status": "activated", + } + assert response["data"]["success"] is True + + +def test_update_memory_requires_a_change(client: Any) -> None: + with pytest.raises(ValueError, match="content, title or status is required"): + client.update_memory(memory_id="memory-1") + + +def test_extract_memory_sends_messages_and_options( + client: Any, posted_requests: list[dict] +) -> None: + messages = [{"role": "user", "content": "I like tea", "chat_time": "2026-07-06"}] + + client.extract_memory( + messages=messages, + extraction_types=["memory", "preference"], + model="extract-model", + ) + + payload = _json_payload(posted_requests[0]) + + assert posted_requests[0]["url"].endswith("/extract/memory") + assert payload == { + "messages": messages, + "extraction_types": ["memory", "preference"], + "model": "extract-model", + } + + +def test_rerank_sends_query_documents_and_options(client: Any, posted_requests: list[dict]) -> None: + client.rerank( + query="memory query", + documents=["doc a", "doc b"], + model="rerank-model", + top_n=1, + ) + + payload = _json_payload(posted_requests[0]) + + assert posted_requests[0]["url"].endswith("/rerank") + assert payload == { + "query": "memory query", + "documents": ["doc a", "doc b"], + "model": "rerank-model", + "top_n": 1, + } + + +def test_rerank_rejects_non_positive_top_n(client: Any) -> None: + with pytest.raises(ValueError, match="top_n must be greater than 0"): + client.rerank(query="memory query", documents=["doc a"], top_n=0) + + +def test_bind_profile_template_sends_bind_list(client: Any, posted_requests: list[dict]) -> None: + bind_list = [{"profile_template_id": "profile-template-1", "user_id": "user-1"}] + + client.bind_profile_template(bind_list=bind_list) + + payload = _json_payload(posted_requests[0]) + + assert posted_requests[0]["url"].endswith("/bind/profile_template") + assert payload == {"bind_list": bind_list} + + +def test_edit_profile_sends_metadata_and_remove_fields( + client: Any, posted_requests: list[dict] +) -> None: + metadata = {"basic": {"city": "Hangzhou"}} + + client.edit_profile( + profile_template_id="profile-template-1", + user_id="user-1", + metadata=metadata, + remove_fields=["basic.job"], + ) + + payload = _json_payload(posted_requests[0]) + + assert posted_requests[0]["url"].endswith("/edit/profile") + assert payload == { + "user_id": "user-1", + "agent_id": None, + "profile_template_id": "profile-template-1", + "metadata": metadata, + "remove_fields": ["basic.job"], + } + + +def test_edit_profile_requires_metadata_or_remove_fields(client: Any) -> None: + with pytest.raises(ValueError, match="metadata or remove_fields is required"): + client.edit_profile(profile_template_id="profile-template-1", user_id="user-1") + + +def test_delete_profile_sends_profile_template_and_subject( + client: Any, posted_requests: list[dict] +) -> None: + client.delete_profile(profile_template_id="profile-template-1", agent_id="agent-1") + + payload = _json_payload(posted_requests[0]) + + assert posted_requests[0]["url"].endswith("/delete/profile") + assert payload == { + "user_id": None, + "agent_id": "agent-1", + "profile_template_id": "profile-template-1", + } + + +def test_profile_subject_requires_exactly_one_user_or_agent(client: Any) -> None: + with pytest.raises(ValueError, match="exactly one of user_id or agent_id is required"): + client.delete_profile(profile_template_id="profile-template-1") + + with pytest.raises(ValueError, match="exactly one of user_id or agent_id is required"): + client.delete_profile( + profile_template_id="profile-template-1", + user_id="user-1", + agent_id="agent-1", + ) From f46c02ed7b9199e812edc03cd3f4fd7dae459dc9 Mon Sep 17 00:00:00 2001 From: dingliang <2650876010@qq.com> Date: Tue, 14 Jul 2026 17:18:25 +0800 Subject: [PATCH 02/15] feat(api): align OpenMem v1 SDK with cloud API --- src/memos/api/client.py | 89 +++++++--- src/memos/api/product_models.py | 46 +++++- tests/api/test_client.py | 282 +++++++++++++++++++++++++++++++- 3 files changed, 387 insertions(+), 30 deletions(-) diff --git a/src/memos/api/client.py b/src/memos/api/client.py index 1ace6ec8f..b8055b5b6 100644 --- a/src/memos/api/client.py +++ b/src/memos/api/client.py @@ -2,7 +2,9 @@ import mimetypes import os +from collections.abc import Iterator from typing import Any +from urllib.parse import quote import requests @@ -95,13 +97,13 @@ def get_message( self, user_id: str, conversation_id: str | None = None, - conversation_limit_number: int = 6, - message_limit_number: int = 6, + conversation_limit_number: int | None = None, + message_limit_number: int | None = None, source: str | None = None, ) -> MemOSGetMessagesResponse | None: """Get message""" # Validate required parameters - self._validate_required_params(user_id=user_id) + self._validate_required_params(user_id=user_id, conversation_id=conversation_id) url = f"{self.base_url}/get/message" payload = { @@ -128,12 +130,12 @@ def get_message( def add_message( self, messages: list[dict[str, Any]], - user_id: str | list[str], - conversation_id: str, + user_id: str | list[str] | None = None, + conversation_id: str | None = None, info: dict[str, Any] | None = None, source: str | None = None, app_id: str | None = None, - agent_id: str | None = None, + agent_id: str | list[str] | None = None, async_mode: bool = True, tags: list[str] | None = None, allow_public: bool = False, @@ -142,9 +144,9 @@ def add_message( ) -> MemOSAddResponse | None: """Add message""" # Validate required parameters - self._validate_required_params( - messages=messages, user_id=user_id, conversation_id=conversation_id - ) + self._validate_required_params(messages=messages) + if not user_id and not agent_id: + raise ValueError("user_id or agent_id is required") url = f"{self.base_url}/add/message" payload = { @@ -178,7 +180,7 @@ def add_message( def search_memory( self, query: str, - user_id: str, + user_id: str | None = None, conversation_id: str | None = None, agent_id: str | None = None, memory_limit_number: int = 6, @@ -197,7 +199,8 @@ def search_memory( ) -> MemOSSearchResponse | None: """Search memories""" # Validate required parameters - self._validate_required_params(query=query, user_id=user_id) + self._validate_required_params(query=query) + self._validate_profile_subject(user_id, agent_id) url = f"{self.base_url}/search/memory" payload = { @@ -248,6 +251,8 @@ def get_memory( """get memories""" # Validate required parameters self._validate_profile_subject(user_id, agent_id) + if size > 50: + raise ValueError("size must be less than or equal to 50") url = f"{self.base_url}/get/memory" payload = { @@ -275,17 +280,47 @@ def get_memory( if retry == MAX_RETRY_COUNT - 1: raise + @staticmethod + def _iter_sse_data(response: requests.Response) -> Iterator[str]: + """Yield decoded data payloads from a Server-Sent Events response.""" + try: + for line in response.iter_lines(decode_unicode=True): + if isinstance(line, bytes): + line = line.decode("utf-8") + if not line or not line.startswith("data:"): + continue + yield line.removeprefix("data:").lstrip() + finally: + response.close() + + def get_memory_by_id(self, memid: str) -> dict[str, Any] | None: + """Get one memory detail by its memory ID.""" + self._validate_required_params(memid=memid) + + url = f"{self.base_url}/get/memory/{quote(memid, safe='')}" + for retry in range(MAX_RETRY_COUNT): + try: + response = requests.get(url, headers=self.headers, timeout=30) + response.raise_for_status() + return response.json() + except Exception as e: + logger.error( + "Failed to get memory by ID (retry %s/%s): %s", + retry + 1, + MAX_RETRY_COUNT, + e, + ) + if retry == MAX_RETRY_COUNT - 1: + raise + def create_knowledgebase( - self, knowledgebase_name: str, knowledgebase_description: str + self, knowledgebase_name: str, knowledgebase_description: str | None = None ) -> MemOSCreateKnowledgebaseResponse | None: """ Create knowledgebase """ # Validate required parameters - self._validate_required_params( - knowledgebase_name=knowledgebase_name, - knowledgebase_description=knowledgebase_description, - ) + self._validate_required_params(knowledgebase_name=knowledgebase_name) url = f"{self.base_url}/create/knowledgebase" payload = { @@ -524,8 +559,8 @@ def get_task_status(self, task_id: str) -> MemOSGetTaskStatusResponse | None: def add_feedback( self, user_id: str, - conversation_id: str, - feedback_content: str, + conversation_id: str | None = None, + feedback_content: str | None = None, agent_id: str | None = None, app_id: str | None = None, feedback_time: str | None = None, @@ -534,9 +569,7 @@ def add_feedback( ) -> MemOSAddFeedBackResponse | None: """Add feedback""" # Validate required parameters - self._validate_required_params( - feedback_content=feedback_content, user_id=user_id, conversation_id=conversation_id - ) + self._validate_required_params(feedback_content=feedback_content, user_id=user_id) url = f"{self.base_url}/add/feedback" payload = { @@ -743,8 +776,8 @@ def chat( allow_public: bool = False, allow_knowledgebase_ids: list[str] | None = None, max_tokens: int = 8192, - temperature: float | None = None, - top_p: float | None = None, + temperature: float | None = 0.7, + top_p: float | None = 0.95, include_preference: bool = True, preference_limit_number: int = 6, memory_limit_number: int = 6, @@ -752,7 +785,7 @@ def chat( include_tool_memory: bool = False, tool_memory_limit_number: int = 6, relativity: float | None = None, - ) -> MemOSChatResponse | None: + ) -> MemOSChatResponse | Iterator[str] | None: """chat""" # Validate required parameters self._validate_required_params( @@ -795,9 +828,15 @@ def chat( for retry in range(MAX_RETRY_COUNT): try: response = requests.post( - url, data=json.dumps(payload), headers=self.headers, timeout=30 + url, + data=json.dumps(payload), + headers=self.headers, + timeout=30, + stream=stream, ) response.raise_for_status() + if stream: + return self._iter_sse_data(response) response_data = response.json() return MemOSChatResponse(**response_data) diff --git a/src/memos/api/product_models.py b/src/memos/api/product_models.py index ff9d94859..86f6fb80c 100644 --- a/src/memos/api/product_models.py +++ b/src/memos/api/product_models.py @@ -1048,6 +1048,15 @@ class SearchMemoryData(BaseModel): alias="tool_memory_detail_list", description="List of tool_memor details (usually None)", ) + skill_detail_list: list[MemoryDetail] | None = Field( + None, alias="skill_detail_list", description="List of skill memory details" + ) + profile_detail_list: list[MemoryDetail] | None = Field( + None, alias="profile_detail_list", description="List of profile memory details" + ) + event_detail_list: list[MemoryDetail] | None = Field( + None, alias="event_detail_list", description="List of event memory details" + ) preference_note: str = Field( None, alias="preference_note", description="String of preference_note" ) @@ -1059,6 +1068,9 @@ class GetKnowledgebaseFileData(BaseModel): file_detail_list: list[FileDetail] = Field( default_factory=list, alias="file_detail_list", description="List of files details" ) + total: int | None = Field(None, description="Total number of matching files") + page: int | None = Field(None, description="Current page number") + page_size: int | None = Field(None, alias="page_size", description="Page size") class GetMemoryData(BaseModel): @@ -1070,6 +1082,22 @@ class GetMemoryData(BaseModel): preference_detail_list: list[MessageDetail] | None = Field( None, alias="preference_detail_list", description="List of preference detail" ) + tool_memory_detail_list: list[MemoryDetail] | None = Field( + None, alias="tool_memory_detail_list", description="List of tool memory details" + ) + profile_detail_list: list[MemoryDetail] | None = Field( + None, alias="profile_detail_list", description="List of profile memory details" + ) + event_detail_list: list[MemoryDetail] | None = Field( + None, alias="event_detail_list", description="List of event memory details" + ) + skill_detail_list: list[MemoryDetail] | None = Field( + None, alias="skill_detail_list", description="List of skill memory details" + ) + total: int | None = Field(None, description="Total number of memories") + size: int | None = Field(None, description="Page size") + current: int | None = Field(None, description="Current page number") + pages: int | None = Field(None, description="Total number of pages") class AddMessageData(BaseModel): @@ -1098,6 +1126,16 @@ class GetTaskStatusMessageData(BaseModel): status: str = Field(..., description="Operation task status") +class GetTaskStatusData(BaseModel): + """Current OpenMem task status response data.""" + + task_id: str = Field(..., description="Task identifier") + status: str = Field(..., description="Operation task status") + memory_views: dict[str, Any] | None = Field( + None, alias="memory_views", description="Memory view changes produced by the task" + ) + + # ─── MemOS Response Models (Similar to OpenAI ChatCompletion) ────────────────── @@ -1181,12 +1219,12 @@ class MemOSGetTaskStatusResponse(BaseModel): code: int = Field(..., description="Response status code") message: str = Field(..., description="Response message") - data: list[GetTaskStatusMessageData] = Field(..., description="Task status data") + data: GetTaskStatusData = Field(..., description="Task status data") @property - def messages(self) -> list[GetTaskStatusMessageData]: - """Convenient access to task status messages.""" - return self.data + def messages(self) -> list[GetTaskStatusData]: + """Backward-compatible list access to task status data.""" + return [self.data] class MemOSCreateKnowledgebaseResponse(BaseModel): diff --git a/tests/api/test_client.py b/tests/api/test_client.py index 61724b0d4..2e911e4f4 100644 --- a/tests/api/test_client.py +++ b/tests/api/test_client.py @@ -43,7 +43,30 @@ def json(self) -> dict: return self.payload +class DummyStreamResponse: + def __init__(self, lines: list[str]): + self.lines = lines + self.closed = False + self.json_called = False + + def raise_for_status(self) -> None: + return None + + def json(self) -> dict: + self.json_called = True + raise AssertionError("streaming responses must not be parsed as JSON") + + def iter_lines(self, decode_unicode: bool = False): + assert decode_unicode is True + yield from self.lines + + def close(self) -> None: + self.closed = True + + def _response_for(url: str) -> dict: + if url.endswith("/get/message"): + return {"code": 200, "message": "ok", "data": {"message_detail_list": []}} if url.endswith("/add/message"): return { "code": 200, @@ -54,10 +77,18 @@ def _response_for(url: str) -> dict: return {"code": 200, "message": "ok", "data": {"memory_detail_list": []}} if url.endswith("/get/memory"): return {"code": 200, "message": "ok", "data": {"memory_detail_list": []}} + if url.endswith("/create/knowledgebase"): + return {"code": 200, "message": "ok", "data": {"id": "kb-1"}} if url.endswith("/get/knowledgebase-file"): return {"code": 200, "message": "ok", "data": {"file_detail_list": []}} if url.endswith("/delete/memory"): return {"code": 200, "message": "ok", "data": {"success": True}} + if url.endswith("/add/feedback"): + return { + "code": 200, + "message": "ok", + "data": {"success": True, "task_id": "task-1", "status": "running"}, + } if url.endswith("/chat"): return {"code": 200, "message": "ok", "data": {"response": "answer"}} if url.endswith("/add/knowledgebase-file"): @@ -102,6 +133,24 @@ def fake_post(url: str, **kwargs): return calls +@pytest.fixture +def fetched_requests(monkeypatch, client_module): + calls: list[dict] = [] + + def fake_get(url: str, **kwargs): + calls.append({"url": url, **kwargs}) + return DummyResponse( + { + "code": 200, + "message": "ok", + "data": {"id": "memory-1", "memory_type": "LongTermMemory"}, + } + ) + + monkeypatch.setattr(client_module.requests, "get", fake_get) + return calls + + @pytest.fixture def client(client_module) -> Any: return client_module.MemOSClient(api_key="test-key", base_url="https://example.test/openmem/v1") @@ -134,7 +183,7 @@ def test_search_memory_sends_updated_existing_request_fields( ) -> None: client.search_memory( query="hello", - user_id="user-1", + user_id=None, agent_id="agent-1", relativity=0.2, include_skill=True, @@ -404,3 +453,234 @@ def test_profile_subject_requires_exactly_one_user_or_agent(client: Any) -> None user_id="user-1", agent_id="agent-1", ) + + +def test_task_status_response_parses_current_object_shape(client_module: Any) -> None: + response = client_module.MemOSGetTaskStatusResponse( + code=200, + message="ok", + data={ + "task_id": "task-1", + "status": "running", + "memory_views": {"added": 1}, + }, + ) + + assert response.data.task_id == "task-1" + assert response.data.status == "running" + assert response.data.memory_views == {"added": 1} + + +def test_search_response_keeps_all_current_memory_view_lists(client_module: Any) -> None: + response = client_module.MemOSSearchResponse( + code=200, + message="ok", + data={ + "memory_detail_list": [], + "skill_detail_list": [{"id": "skill-1"}], + "profile_detail_list": [{"id": "profile-1"}], + "event_detail_list": [{"id": "event-1"}], + }, + ) + + assert response.data.skill_detail_list[0].id == "skill-1" + assert response.data.profile_detail_list[0].id == "profile-1" + assert response.data.event_detail_list[0].id == "event-1" + + +def test_get_memory_response_keeps_views_and_pagination(client_module: Any) -> None: + response = client_module.MemOSGetMemoryResponse( + code=200, + message="ok", + data={ + "memory_detail_list": [], + "tool_memory_detail_list": [{"id": "tool-1"}], + "profile_detail_list": [{"id": "profile-1"}], + "event_detail_list": [{"id": "event-1"}], + "skill_detail_list": [{"id": "skill-1"}], + "total": 21, + "size": 10, + "current": 2, + "pages": 3, + }, + ) + + assert response.data.tool_memory_detail_list[0].id == "tool-1" + assert response.data.profile_detail_list[0].id == "profile-1" + assert response.data.event_detail_list[0].id == "event-1" + assert response.data.skill_detail_list[0].id == "skill-1" + assert response.data.total == 21 + assert response.data.size == 10 + assert response.data.current == 2 + assert response.data.pages == 3 + + +def test_get_knowledgebase_file_response_keeps_pagination(client_module: Any) -> None: + response = client_module.MemOSGetKnowledgebaseFileResponse( + code=200, + message="ok", + data={ + "file_detail_list": [], + "total": 8, + "page": 2, + "page_size": 5, + }, + ) + + assert response.data.total == 8 + assert response.data.page == 2 + assert response.data.page_size == 5 + + +def test_get_message_requires_conversation_id(client: Any, posted_requests: list[dict]) -> None: + with pytest.raises(ValueError, match="conversation_id is required"): + client.get_message(user_id="user-1") + + assert posted_requests == [] + + +def test_get_message_uses_playground_default_limits( + client: Any, posted_requests: list[dict] +) -> None: + client.get_message(user_id="user-1", conversation_id="conversation-1") + + payload = _json_payload(posted_requests[0]) + + assert payload["conversation_limit_number"] is None + assert payload["message_limit_number"] is None + + +def test_add_message_allows_agent_only_and_generated_conversation( + client: Any, posted_requests: list[dict] +) -> None: + client.add_message( + messages=[{"role": "user", "content": "hello"}], + user_id=None, + agent_id="agent-1", + conversation_id=None, + ) + + payload = _json_payload(posted_requests[0]) + + assert payload["user_id"] is None + assert payload["agent_id"] == "agent-1" + assert payload["conversation_id"] is None + + +def test_search_memory_allows_agent_only(client: Any, posted_requests: list[dict]) -> None: + client.search_memory(query="hello", user_id=None, agent_id="agent-1") + + payload = _json_payload(posted_requests[0]) + + assert payload["user_id"] is None + assert payload["agent_id"] == "agent-1" + + +def test_search_memory_rejects_multiple_subjects(client: Any, posted_requests: list[dict]) -> None: + with pytest.raises(ValueError, match="exactly one of user_id or agent_id"): + client.search_memory(query="hello", user_id="user-1", agent_id="agent-1") + + assert posted_requests == [] + + +def test_create_knowledgebase_allows_empty_description( + client: Any, posted_requests: list[dict] +) -> None: + client.create_knowledgebase(knowledgebase_name="Knowledge Base") + + payload = _json_payload(posted_requests[0]) + + assert payload == { + "knowledgebase_name": "Knowledge Base", + "knowledgebase_description": None, + } + + +def test_add_feedback_allows_generated_conversation( + client: Any, posted_requests: list[dict] +) -> None: + client.add_feedback(user_id="user-1", feedback_content="helpful") + + payload = _json_payload(posted_requests[0]) + + assert payload["conversation_id"] is None + assert payload["feedback_content"] == "helpful" + + +def test_chat_uses_playground_sampling_defaults(client: Any, posted_requests: list[dict]) -> None: + client.chat(user_id="user-1", conversation_id="conversation-1", query="hello") + + payload = _json_payload(posted_requests[0]) + + assert payload["temperature"] == 0.7 + assert payload["top_p"] == 0.95 + + +def test_get_memory_rejects_size_above_playground_limit( + client: Any, posted_requests: list[dict] +) -> None: + with pytest.raises(ValueError, match="size must be less than or equal to 50"): + client.get_memory(user_id="user-1", size=51) + + assert posted_requests == [] + + +def test_get_memory_by_id_uses_detail_get_endpoint( + client: Any, fetched_requests: list[dict] +) -> None: + response = client.get_memory_by_id("memory-1") + + assert fetched_requests == [ + { + "url": "https://example.test/openmem/v1/get/memory/memory-1", + "headers": client.headers, + "timeout": 30, + } + ] + assert response == { + "code": 200, + "message": "ok", + "data": {"id": "memory-1", "memory_type": "LongTermMemory"}, + } + + +def test_get_memory_by_id_requires_memid(client: Any, fetched_requests: list[dict]) -> None: + with pytest.raises(ValueError, match="memid is required"): + client.get_memory_by_id("") + + assert fetched_requests == [] + + +def test_chat_stream_yields_sse_data_and_closes_response(monkeypatch, client_module: Any) -> None: + calls: list[dict] = [] + stream_response = DummyStreamResponse( + [ + "event: message", + 'data: {"response":"first"}', + "", + "data: [DONE]", + ] + ) + + def fake_post(url: str, **kwargs): + calls.append({"url": url, **kwargs}) + return stream_response + + monkeypatch.setattr(client_module.requests, "post", fake_post) + client = client_module.MemOSClient( + api_key="test-key", base_url="https://example.test/openmem/v1" + ) + + chunks = list( + client.chat( + user_id="user-1", + conversation_id="conversation-1", + query="hello", + stream=True, + ) + ) + + assert calls[0]["stream"] is True + assert chunks == ['{"response":"first"}', "[DONE]"] + assert stream_response.json_called is False + assert stream_response.closed is True From 9d2c9b75d2f462e8e05f7750e2266d7eea0a7a48 Mon Sep 17 00:00:00 2001 From: de1ty <7804799+de1tydev@users.noreply.github.com> Date: Tue, 14 Jul 2026 22:27:08 +0800 Subject: [PATCH 03/15] fix(memos-local): include json hint in user messages (#1756) Co-authored-by: shinetata <149466187+shinetata@users.noreply.github.com> --- apps/memos-local-plugin/core/llm/client.ts | 19 +++++++++++++++++-- .../tests/unit/llm/client.test.ts | 8 ++++++-- 2 files changed, 23 insertions(+), 4 deletions(-) diff --git a/apps/memos-local-plugin/core/llm/client.ts b/apps/memos-local-plugin/core/llm/client.ts index 41f31a439..6bedafa70 100644 --- a/apps/memos-local-plugin/core/llm/client.ts +++ b/apps/memos-local-plugin/core/llm/client.ts @@ -260,6 +260,21 @@ export function createLlmClientWithProvider( return [{ role: "system", content: systemInsert }, ...messages]; } + function ensureJsonWordInUserMessage(messages: LlmMessage[]): LlmMessage[] { + const lastUserIdx = messages.map((m) => m.role).lastIndexOf("user"); + if (lastUserIdx < 0) return [...messages, { role: "user", content: "Return valid json only." }]; + + const msg = messages[lastUserIdx]; + if (/\bjson\b/i.test(msg.content)) return messages; + + const out = messages.slice(); + out[lastUserIdx] = { + ...msg, + content: `${msg.content}\n\nReturn valid json only.`, + }; + return out; + } + function buildCallInput(opts: LlmCallOptions | undefined, jsonMode: boolean): ProviderCallInput { return { temperature: opts?.temperature ?? config.temperature, @@ -463,7 +478,7 @@ export function createLlmClientWithProvider( ): Promise { const messages = normalizeMessages(input); const msgsWithJsonHint = opts?.jsonMode - ? inject(messages, buildJsonSystemHint()) + ? ensureJsonWordInUserMessage(inject(messages, buildJsonSystemHint())) : messages; const call = buildCallInput(opts, opts?.jsonMode === true); const { completion } = await callWithFallback(msgsWithJsonHint, call, opts, opts?.op ?? "complete"); @@ -476,7 +491,7 @@ export function createLlmClientWithProvider( ): Promise> { const messages = normalizeMessages(input); const systemHint = buildJsonSystemHint(opts.schemaHint); - const msgs = inject(messages, systemHint); + const msgs = ensureJsonWordInUserMessage(inject(messages, systemHint)); const call = buildCallInput(opts, true); const op = opts.op ?? "complete.json"; const maxMalformedRetries = Math.max(0, opts.malformedRetries ?? 1); diff --git a/apps/memos-local-plugin/tests/unit/llm/client.test.ts b/apps/memos-local-plugin/tests/unit/llm/client.test.ts index 7e904a2c6..dee0de228 100644 --- a/apps/memos-local-plugin/tests/unit/llm/client.test.ts +++ b/apps/memos-local-plugin/tests/unit/llm/client.test.ts @@ -96,12 +96,14 @@ describe("llm/client", () => { expect(fake.lastMessages).toEqual([{ role: "user", content: "hi there" }]); }); - it("injects a json system hint when jsonMode=true", async () => { + it("injects json hints into system and user messages when jsonMode=true", async () => { const fake = new FakeProvider("openai_compatible", () => ({ text: '{"ok":1}', durationMs: 1 })); const client = createLlmClientWithProvider(cfg(), fake); await client.complete("do it", { jsonMode: true }); expect(fake.lastMessages?.[0]?.role).toBe("system"); expect(fake.lastMessages?.[0]?.content).toMatch(/single valid JSON value/i); + expect(fake.lastMessages?.at(-1)?.role).toBe("user"); + expect(fake.lastMessages?.at(-1)?.content).toMatch(/valid json only/i); expect(fake.lastInput?.jsonMode).toBe(true); }); @@ -270,7 +272,9 @@ describe("llm/client", () => { expect(fake.lastMessages?.[0]?.role).toBe("system"); expect(fake.lastMessages?.[0]?.content).toMatch(/You are strict\./); expect(fake.lastMessages?.[0]?.content).toMatch(/single valid JSON value/); - expect(fake.lastMessages?.[1]).toEqual({ role: "user", content: "go" }); + expect(fake.lastMessages?.[1]?.role).toBe("user"); + expect(fake.lastMessages?.[1]?.content).toMatch(/^go/); + expect(fake.lastMessages?.[1]?.content).toMatch(/valid json only/i); }); it("rejects empty messages array", async () => { From a000a1c7341fdfbe3d13ac091b83e5aec4cf9abd Mon Sep 17 00:00:00 2001 From: shinetata <149466187+shinetata@users.noreply.github.com> Date: Thu, 16 Jul 2026 14:35:44 +0800 Subject: [PATCH 04/15] docs: fix scheduler API examples to match server endpoints (#2113) Rewrite scheduler Quick Start examples to call the real /product/scheduler REST endpoints with requests instead of non-existent MemOSClient methods; align response fields with the actual handler output; rename the mis-named ' wait.md' to 'wait.md'. Closes #2083 Co-authored-by: sunqi Co-authored-by: Cursor --- .../open_source_api/scheduler/get_status.md | 48 ++++++++++------- .../scheduler/{ wait.md => wait.md} | 51 +++++++++++-------- .../open_source_api/scheduler/get_status.md | 48 ++++++++++------- .../open_source_api/scheduler/wait.md | 51 +++++++++++-------- 4 files changed, 124 insertions(+), 74 deletions(-) rename docs/cn/open_source/open_source_api/scheduler/{ wait.md => wait.md} (61%) diff --git a/docs/cn/open_source/open_source_api/scheduler/get_status.md b/docs/cn/open_source/open_source_api/scheduler/get_status.md index 87e1a4a5a..0a60d6fac 100644 --- a/docs/cn/open_source/open_source_api/scheduler/get_status.md +++ b/docs/cn/open_source/open_source_api/scheduler/get_status.md @@ -64,34 +64,48 @@ desc: 监控 MemOS 异步任务的生命周期,提供包括任务进度、队 ## 4. 快速上手示例 -使用 SDK 轮询任务状态直至完成: +这些接口由开源版 Server(`server_api`,路由前缀 `/product`)直接提供,使用标准 HTTP 请求即可访问。以下示例轮询任务状态直至完成: ```python -from memos.api.client import MemOSClient import time -client = MemOSClient(api_key="...", base_url="...") +import requests + +# 自部署 MemOS Server 的地址(如启用了鉴权,请自行补充 Authorization 请求头) +base_url = "http://localhost:8000" # 1. 系统级概览:查看整个 MemOS 系统的运行健康度 -global_res = client.get_all_scheduler_status() -if global_res: - print(f"系统运行概况: {global_res.data['scheduler_summary']}") +resp = requests.get(f"{base_url}/product/scheduler/allstatus", timeout=10) +resp.raise_for_status() +global_res = resp.json() +print(f"系统运行概况: {global_res['data']['scheduler_summary']}") # 2. 队列指标监控:检查特定用户的任务积压情况 -queue_res = client.get_task_queue_status(user_id="dev_user_01") -if queue_res: - print(f"待处理任务数: {queue_res.data['remaining_tasks_count']}") - print(f"已下发未完成任务数: {queue_res.data['pending_tasks_count']}") +resp = requests.get( + f"{base_url}/product/scheduler/task_queue_status", + params={"user_id": "dev_user_01"}, + timeout=10, +) +resp.raise_for_status() +queue_res = resp.json() +print(f"排队中任务数: {queue_res['data']['remaining_tasks_count']}") +print(f"已下发未确认任务数: {queue_res['data']['pending_tasks_count']}") # 3. 任务进度追踪:轮询特定任务直至结束 task_id = "task_888999" +active_states = {"waiting", "pending", "in_progress"} while True: - res = client.get_task_status(user_id="dev_user_01", task_id=task_id) - if res and res.code == 200: - current_status = res.data[0]['status'] # data 为状态列表 - print(f"任务 {task_id} 当前状态: {current_status}") - - if current_status in ['completed', 'failed', 'cancelled']: - break + resp = requests.get( + f"{base_url}/product/scheduler/status", + params={"user_id": "dev_user_01", "task_id": task_id}, + timeout=10, + ) + resp.raise_for_status() + items = resp.json().get("data", []) # data 为状态列表:[{"task_id": ..., "status": ...}] + statuses = {item["status"] for item in items} + print(f"任务 {task_id} 当前状态: {statuses or '空'}") + + if not statuses or statuses.isdisjoint(active_states): + break time.sleep(2) ``` diff --git a/docs/cn/open_source/open_source_api/scheduler/ wait.md b/docs/cn/open_source/open_source_api/scheduler/wait.md similarity index 61% rename from docs/cn/open_source/open_source_api/scheduler/ wait.md rename to docs/cn/open_source/open_source_api/scheduler/wait.md index 9849ffe68..52b0d79f6 100644 --- a/docs/cn/open_source/open_source_api/scheduler/ wait.md +++ b/docs/cn/open_source/open_source_api/scheduler/wait.md @@ -42,36 +42,47 @@ desc: 提供阻塞等待与流式进度观测能力,确保在执行后续操 ## 4. 快速上手示例 -使用开源版 SDK 进行阻塞式等待: +这些接口由开源版 Server(`server_api`,路由前缀 `/product`)直接提供,使用标准 HTTP 请求即可访问。注意:`user_name`、`timeout_seconds`、`poll_interval` 均为查询参数(Query),而非请求体(Body)。以下示例进行阻塞式等待: ```python -from memos.api.client import MemOSClient +import json -client = MemOSClient(api_key="...", base_url="...") +import requests + +# 自部署 MemOS Server 的地址(如启用了鉴权,请自行补充 Authorization 请求头) +base_url = "http://localhost:8000" user_name = "dev_user_01" # --- 场景 A:同步阻塞等待 (常用于 Python 自动化脚本) --- print(f"正在等待用户 {user_name} 的任务队列清空...") -res = client.wait_until_idle( - user_name=user_name, - timeout_seconds=300, - poll_interval=2 +resp = requests.post( + f"{base_url}/product/scheduler/wait", + params={"user_name": user_name, "timeout_seconds": 300, "poll_interval": 2}, + timeout=310, # HTTP 超时应大于 timeout_seconds ) -if res and res.code == 200: +resp.raise_for_status() +result = resp.json() # {"message": "idle" | "timeout", "data": {...}} +if result["message"] == "idle": print("✅ 任务已全部完成。") +else: + print(f"⚠️ 等待超时,仍有 {result['data']['running_tasks']} 个任务在执行。") # --- 场景 B:流式进度观测 (常用于前端进度条渲染) --- print("开始监听任务实时进度流...") -# 注意:SSE 接口在 SDK 中通常返回一个生成器 (Generator) -progress_stream = client.stream_scheduler_progress( - user_name=user_name, - timeout_seconds=300 -) - -for event in progress_stream: - # 实时打印剩余任务数 - print(f"当前排队任务数: {event['remaining_tasks_count']}") - if event['status'] == 'idle': - print("🎉 调度器已空闲") - break +with requests.get( + f"{base_url}/product/scheduler/wait/stream", + params={"user_name": user_name, "timeout_seconds": 300}, + stream=True, + timeout=310, +) as resp: + resp.raise_for_status() + for line in resp.iter_lines(decode_unicode=True): + if not line or not line.startswith("data:"): + continue + event = json.loads(line.removeprefix("data:").strip()) + # 实时打印仍在执行的任务数 + print(f"当前活跃任务数: {event['active_tasks']},状态: {event['status']}") + if event["status"] in ("idle", "timeout"): + print("🎉 调度器已空闲" if event["status"] == "idle" else "⚠️ 监听超时") + break ``` diff --git a/docs/en/open_source/open_source_api/scheduler/get_status.md b/docs/en/open_source/open_source_api/scheduler/get_status.md index f2014d9e5..7d1565858 100644 --- a/docs/en/open_source/open_source_api/scheduler/get_status.md +++ b/docs/en/open_source/open_source_api/scheduler/get_status.md @@ -66,34 +66,48 @@ When you send a status request, **SchedulerHandler** performs the following oper ## 4. Quick Start -Poll task status with the SDK until completion: +These endpoints are served directly by the open-source Server (`server_api`, router prefix `/product`) and can be called with plain HTTP requests. The example below polls task status until completion: ```python -from memos.api.client import MemOSClient import time -client = MemOSClient(api_key="...", base_url="...") +import requests + +# Address of your self-hosted MemOS Server (add an Authorization header if auth is enabled) +base_url = "http://localhost:8000" # 1. System overview: inspect overall MemOS health. -global_res = client.get_all_scheduler_status() -if global_res: - print(f"System summary: {global_res.data['scheduler_summary']}") +resp = requests.get(f"{base_url}/product/scheduler/allstatus", timeout=10) +resp.raise_for_status() +global_res = resp.json() +print(f"System summary: {global_res['data']['scheduler_summary']}") # 2. Queue metrics: inspect backlog for a specific user. -queue_res = client.get_task_queue_status(user_id="dev_user_01") -if queue_res: - print(f"Remaining tasks: {queue_res.data['remaining_tasks_count']}") - print(f"Pending tasks: {queue_res.data['pending_tasks_count']}") +resp = requests.get( + f"{base_url}/product/scheduler/task_queue_status", + params={"user_id": "dev_user_01"}, + timeout=10, +) +resp.raise_for_status() +queue_res = resp.json() +print(f"Remaining tasks: {queue_res['data']['remaining_tasks_count']}") +print(f"Pending tasks: {queue_res['data']['pending_tasks_count']}") # 3. Task progress: poll a specific task until it finishes. task_id = "task_888999" +active_states = {"waiting", "pending", "in_progress"} while True: - res = client.get_task_status(user_id="dev_user_01", task_id=task_id) - if res and res.code == 200: - current_status = res.data[0]['status'] # data is a status list - print(f"Task {task_id} status: {current_status}") - - if current_status in ['completed', 'failed', 'cancelled']: - break + resp = requests.get( + f"{base_url}/product/scheduler/status", + params={"user_id": "dev_user_01", "task_id": task_id}, + timeout=10, + ) + resp.raise_for_status() + items = resp.json().get("data", []) # data is a status list: [{"task_id": ..., "status": ...}] + statuses = {item["status"] for item in items} + print(f"Task {task_id} status: {statuses or 'empty'}") + + if not statuses or statuses.isdisjoint(active_states): + break time.sleep(2) ``` diff --git a/docs/en/open_source/open_source_api/scheduler/wait.md b/docs/en/open_source/open_source_api/scheduler/wait.md index 9de0ff4be..6f356341a 100644 --- a/docs/en/open_source/open_source_api/scheduler/wait.md +++ b/docs/en/open_source/open_source_api/scheduler/wait.md @@ -41,36 +41,47 @@ Both endpoints share the following query parameters: ## 4. Quick Start -Use the open-source SDK for a blocking wait: +These endpoints are served directly by the open-source Server (`server_api`, router prefix `/product`) and can be called with plain HTTP requests. Note that `user_name`, `timeout_seconds`, and `poll_interval` are query parameters, not a request body. The example below performs a blocking wait: ```python -from memos.api.client import MemOSClient +import json -client = MemOSClient(api_key="...", base_url="...") +import requests + +# Address of your self-hosted MemOS Server (add an Authorization header if auth is enabled) +base_url = "http://localhost:8000" user_name = "dev_user_01" # Scenario A: blocking wait, commonly used in Python automation scripts. print(f"Waiting for user {user_name}'s task queue to drain...") -res = client.wait_until_idle( - user_name=user_name, - timeout_seconds=300, - poll_interval=2 +resp = requests.post( + f"{base_url}/product/scheduler/wait", + params={"user_name": user_name, "timeout_seconds": 300, "poll_interval": 2}, + timeout=310, # HTTP timeout should be larger than timeout_seconds ) -if res and res.code == 200: +resp.raise_for_status() +result = resp.json() # {"message": "idle" | "timeout", "data": {...}} +if result["message"] == "idle": print("All tasks have completed.") +else: + print(f"Timed out with {result['data']['running_tasks']} task(s) still running.") # Scenario B: streaming progress, commonly used by frontend progress bars. print("Listening to the live task progress stream...") -# The SSE endpoint usually returns a generator from the SDK. -progress_stream = client.stream_scheduler_progress( - user_name=user_name, - timeout_seconds=300 -) - -for event in progress_stream: - # Print the remaining queued tasks in real time. - print(f"Remaining queued tasks: {event['remaining_tasks_count']}") - if event['status'] == 'idle': - print("Scheduler is idle") - break +with requests.get( + f"{base_url}/product/scheduler/wait/stream", + params={"user_name": user_name, "timeout_seconds": 300}, + stream=True, + timeout=310, +) as resp: + resp.raise_for_status() + for line in resp.iter_lines(decode_unicode=True): + if not line or not line.startswith("data:"): + continue + event = json.loads(line.removeprefix("data:").strip()) + # Print the number of active tasks in real time. + print(f"Active tasks: {event['active_tasks']}, status: {event['status']}") + if event["status"] in ("idle", "timeout"): + print("Scheduler is idle" if event["status"] == "idle" else "Stream timed out") + break ``` From 2c60267e9672cf9e0e7af53e8e348be8c2c1d99a Mon Sep 17 00:00:00 2001 From: shinetata <149466187+shinetata@users.noreply.github.com> Date: Thu, 16 Jul 2026 14:36:07 +0800 Subject: [PATCH 05/15] fix(api): show structured add example in /docs for /product/add (#2112) fix(api): show structured add example in /docs (messages not string) Swagger UI rendered APIADDRequest.messages as "string" because it picks the leading `str` branch of the `str | MessageList | RawMessageList` union when generating the example. Add a model-level json_schema_extra example so /docs shows a copy-paste-ready payload with a structured messages list. Schema-only change: field types, validation and runtime behaviour are unchanged, so the core add pipeline is unaffected. Closes #1505 Co-authored-by: sunqi Co-authored-by: Cursor --- src/memos/api/product_models.py | 15 ++++++++++++ tests/api/test_product_models.py | 41 ++++++++++++++++++++++++++++++++ 2 files changed, 56 insertions(+) create mode 100644 tests/api/test_product_models.py diff --git a/src/memos/api/product_models.py b/src/memos/api/product_models.py index a2c83f622..1769b0ee6 100644 --- a/src/memos/api/product_models.py +++ b/src/memos/api/product_models.py @@ -608,6 +608,21 @@ def _convert_deprecated_fields(self) -> "APISearchRequest": class APIADDRequest(BaseRequest): """Request model for creating memories.""" + # Model-level example so the interactive docs (/docs) show a copy-paste-ready + # payload. Without it, Swagger UI renders the leading `str` branch of the + # `messages` union as `"string"` (see issue #1505). This only affects the + # generated OpenAPI schema, not validation or runtime behaviour. + model_config = { + "json_schema_extra": { + "example": { + "user_id": "8736b16e-1d20-4163-980b-a5063c3facdc", + "writable_cube_ids": ["b32d0977-435d-4828-a86f-4f47f8b55bca"], + "messages": [{"role": "user", "content": "I am learning ggplot2 in R."}], + "async_mode": "async", + } + } + } + # ==== Basic identifiers ==== user_id: str = Field(None, description="User ID") session_id: str | None = Field( diff --git a/tests/api/test_product_models.py b/tests/api/test_product_models.py new file mode 100644 index 000000000..177f8713d --- /dev/null +++ b/tests/api/test_product_models.py @@ -0,0 +1,41 @@ +"""Unit tests for API request-model OpenAPI schemas. + +These tests lock the OpenAPI schema behaviour of request models so the +interactive docs (``/docs``) stay consistent with the documented contract. + +Regression guard for issue #1505: the ``/product/add`` example must render +``messages`` as a structured message list instead of a bare ``"string"``. +Because ``messages`` is typed as ``str | MessageList | RawMessageList``, Swagger +UI would otherwise pick the leading ``str`` branch of the ``anyOf`` and show +``"messages": "string"``, which misleads users into sending plain text. +""" + +from memos.api.product_models import APIADDRequest + + +def test_add_request_exposes_model_level_example(): + """APIADDRequest must ship a model-level example for the interactive docs.""" + schema = APIADDRequest.model_json_schema() + + assert "example" in schema, "APIADDRequest should define a model-level example" + + +def test_add_request_example_messages_is_structured_list(): + """The example's ``messages`` must be a non-empty list of role/content items.""" + example = APIADDRequest.model_json_schema()["example"] + + messages = example.get("messages") + assert isinstance(messages, list), "messages example must be a list, not a bare string" + assert messages, "messages example should not be empty" + + first = messages[0] + assert first.get("role"), "each example message needs a role" + assert first.get("content"), "each example message needs content" + + +def test_add_request_example_covers_core_fields(): + """The example should be a copy-paste-ready payload for the core add flow.""" + example = APIADDRequest.model_json_schema()["example"] + + assert "user_id" in example + assert "writable_cube_ids" in example From c0f6c2ccee73ff56f05e38a5f94ee91f06aaf5ec Mon Sep 17 00:00:00 2001 From: Memtensor-AI Date: Thu, 16 Jul 2026 17:41:57 +0800 Subject: [PATCH 06/15] fix(memos-local-plugin): revert half-merged MemosHttpClient usages (#2096) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The hermes adapter's __init__.py imported MemosHttpClient from bridge_client, but the class was never committed — every plugin python test failed at collection and any real Hermes host loading memos_provider hit ImportError at module load. The phantom usages arrived via the pr/may27-fixes merge (0398e0ed) as a half-committed HTTP-bridge feature. Revert the four usages and the HTTP-first branches in initialize / _reconnect_bridge, remove the _connect_http_bridge helper and the now- unused probe_viewer_status / startup_lock_active imports (their definitions in daemon_manager.py are untouched for a future HTTP PR). Add test_module_imports_cleanly regression test asserting MemTensorProvider / MemosBridgeClient / BridgeError are present on the real contract surfaces and MemosHttpClient is absent on both memos_provider and bridge_client, so this class of import-level breakage cannot silently return. Restores the adapter to stdio-only behavior. Net diff: 57 additions / 54 deletions across the adapter plus the regression test. Fixes #2096 --- .../hermes/memos_provider/__init__.py | 66 ++++--------------- .../python/test_hermes_provider_pipeline.py | 45 +++++++++++++ 2 files changed, 57 insertions(+), 54 deletions(-) diff --git a/apps/memos-local-plugin/adapters/hermes/memos_provider/__init__.py b/apps/memos-local-plugin/adapters/hermes/memos_provider/__init__.py index 25546dfff..47392b8a7 100644 --- a/apps/memos-local-plugin/adapters/hermes/memos_provider/__init__.py +++ b/apps/memos-local-plugin/adapters/hermes/memos_provider/__init__.py @@ -65,13 +65,11 @@ if str(_PLUGIN_DIR) not in sys.path: sys.path.insert(0, str(_PLUGIN_DIR)) -from bridge_client import BridgeError, MemosBridgeClient, MemosHttpClient # noqa: E402 +from bridge_client import BridgeError, MemosBridgeClient # noqa: E402 from daemon_manager import ( # noqa: E402 ensure_bridge_running, ensure_viewer_daemon, kill_zombie_bridges, - probe_viewer_status, - startup_lock_active, ) @@ -279,7 +277,7 @@ class MemTensorProvider(MemoryProvider): """ def __init__(self) -> None: - self._bridge: MemosBridgeClient | MemosHttpClient | None = None + self._bridge: MemosBridgeClient | None = None self._reconnect_lock = threading.Lock() self._session_id: str = "" self._episode_id: str = "" @@ -329,23 +327,6 @@ def is_available(self) -> bool: # type: ignore[override] # ─── Lifecycle ──────────────────────────────────────────────────────── - def _connect_http_bridge(self, session_id: str, *, timeout: float = 60.0) -> bool: - """Try to connect via HTTP bridge. Sets self._bridge on success.""" - http_bridge: MemosHttpClient | None = None - try: - http_bridge = MemosHttpClient() - http_bridge.register_host_handler("host.llm.complete", self._handle_host_llm_complete) - self._bridge = http_bridge - self._open_session(session_id, timeout=timeout) - return True - except Exception as err: - logger.warning("MemOS: HTTP bridge failed, falling back to stdio — %s", err) - if http_bridge is not None: - with contextlib.suppress(Exception): - http_bridge.close() - self._bridge = None - return False - def initialize(self, session_id: str, **kwargs: Any) -> None: # type: ignore[override] """Called once at agent startup. @@ -414,32 +395,13 @@ def initialize(self, session_id: str, **kwargs: Any) -> None: # type: ignore[ov except Exception: pass - # If the daemon is already running on the viewer port, connect - # to it over HTTP instead of spawning a new stdio bridge. This - # eliminates zombie bridge accumulation. - viewer_status = probe_viewer_status() - if viewer_status == "running_memos": - if self._connect_http_bridge(session_id): - logger.info( - "MemOS: bridge ready (HTTP) session=%s platform=%s (episode deferred)", - self._session_id, - self._platform, - ) - else: - viewer_status = "free" # force stdio fallback below - elif viewer_status == "free": - # Re-probe after a short wait only when another process may be - # mid-startup (startup lock is held). On a cold first-launch the - # lock doesn't exist, so we skip the delay entirely. - if startup_lock_active(): - time.sleep(1.0) - viewer_status = probe_viewer_status() - if viewer_status == "running_memos" and self._connect_http_bridge(session_id): - logger.info( - "MemOS: bridge ready (HTTP, late probe) session=%s platform=%s (episode deferred)", - self._session_id, - self._platform, - ) + # NOTE: An HTTP bridge path used to live here that connected to a + # running viewer daemon over HTTP instead of spawning a stdio + # subprocess. It depended on ``MemosHttpClient`` which was never + # committed — the class was referenced by name only. Issue #2096 + # reverts the half-merged HTTP feature; the stdio path below is + # the sole connection mechanism until the HTTP client lands as a + # complete change. if self._bridge is None: try: @@ -1958,13 +1920,9 @@ def _reconnect_bridge(self, session_id: str = "", *, timeout: float = 30.0) -> N logger.info("MemOS: old bridge closed (pid=%s)", old_pid) ensure_bridge_running() - # Try HTTP first if daemon is running - viewer_status = probe_viewer_status() - if viewer_status == "running_memos" and self._connect_http_bridge( - session_id, timeout=timeout - ): - logger.info("MemOS: reconnected via HTTP") - return + # NOTE: HTTP bridge reconnect path was removed alongside issue + # #2096. See ``initialize`` for the rationale. Reconnect always + # spawns a fresh stdio bridge. try: ensure_viewer_daemon() diff --git a/apps/memos-local-plugin/tests/python/test_hermes_provider_pipeline.py b/apps/memos-local-plugin/tests/python/test_hermes_provider_pipeline.py index 1c8cb6a6c..71e3b347b 100644 --- a/apps/memos-local-plugin/tests/python/test_hermes_provider_pipeline.py +++ b/apps/memos-local-plugin/tests/python/test_hermes_provider_pipeline.py @@ -68,6 +68,51 @@ def request(self, method: str, params: dict | None = None, **_kwargs: object) -> class HermesProviderPipelineTests(unittest.TestCase): + def test_module_imports_cleanly(self) -> None: + """Regression guard for #2096: asserts that ``MemosHttpClient`` is + NOT present in ``memos_provider``, since the class was referenced + before it was ever committed (see issue #2096). + + Note: the import itself is already validated at collection time — + the ``import memos_provider`` at the top of this file will raise + ``ImportError`` if a dangling reference is reintroduced, causing + the entire test file to fail to load. This test body only adds: + + * the explicit negative guard on ``MemosHttpClient`` below (unique + to this test), which covers both the ``memos_provider`` + re-export surface *and* ``bridge_client`` itself so a partial + re-add of the class only in ``bridge_client`` (with no + matching re-export) still fails the guard, and + * positive checks on ``MemTensorProvider`` (the class the Hermes + host actually instantiates) and on ``bridge_client``'s real + contract (``MemosBridgeClient`` / ``BridgeError``), rather than + on their incidental re-exports through ``memos_provider`` — the + latter only appear on the package namespace because + ``__init__.py`` uses a bare ``from bridge_client import ...``, + which is an implementation detail we don't want the test to + lock in. + """ + import importlib + + # MemTensorProvider is the class hermes-agent host instantiates. + self.assertTrue(hasattr(memos_provider, "MemTensorProvider")) + + # Assert the actual contract on bridge_client directly rather + # than on its re-exports through memos_provider. + bc = importlib.import_module("bridge_client") + self.assertTrue(hasattr(bc, "MemosBridgeClient")) + self.assertTrue(hasattr(bc, "BridgeError")) + + # ``MemosHttpClient`` was referenced by name in a half-merged HTTP + # bridge feature (see #2096). It must not reappear until the class + # itself is committed in ``bridge_client``. Guard both the + # ``memos_provider`` re-export (which is what the original + # ImportError travelled through) and ``bridge_client`` itself — + # otherwise a partial re-add of the class in ``bridge_client`` + # without a matching re-export would slip past this test. + self.assertFalse(hasattr(memos_provider, "MemosHttpClient")) + self.assertFalse(hasattr(bc, "MemosHttpClient")) + def test_lifecycle_persists_turn_and_closes_real_episode(self) -> None: bridge = FakeBridge() with ( From 698c30f7b8fd36ba6beeba02f6d95b0d11efde73 Mon Sep 17 00:00:00 2001 From: bittergreen Date: Thu, 16 Jul 2026 20:50:37 +0800 Subject: [PATCH 07/15] fix: configure logging once per process --- src/memos/log.py | 23 ++++++++++++++++++++++- tests/test_log.py | 31 +++++++++++++++++++++++++++++++ 2 files changed, 53 insertions(+), 1 deletion(-) diff --git a/src/memos/log.py b/src/memos/log.py index c0bb5bf31..ede5e8280 100644 --- a/src/memos/log.py +++ b/src/memos/log.py @@ -27,6 +27,8 @@ load_dotenv() selected_log_level = logging.DEBUG if settings.DEBUG else logging.WARNING +_LOGGING_CONFIG_LOCK = threading.RLock() +_LOGGING_CONFIGURED_PID: int | None = None def _setup_logfile() -> Path: @@ -224,12 +226,31 @@ def close(self): } +def configure_logging(force: bool = False) -> None: + """Configure process-local logging once. + + Re-running dictConfig replaces and closes existing handlers. Guarding it avoids races with + background threads that may be emitting log records while other modules import loggers. + """ + global _LOGGING_CONFIGURED_PID + + current_pid = os.getpid() + if not force and current_pid == _LOGGING_CONFIGURED_PID: + return + + with _LOGGING_CONFIG_LOCK: + current_pid = os.getpid() + if force or current_pid != _LOGGING_CONFIGURED_PID: + dictConfig(LOGGING_CONFIG) + _LOGGING_CONFIGURED_PID = current_pid + + def get_logger(name: str | None = None) -> logging.Logger: """returns the project logger, scoped to a child name if provided Args: name: will define a child logger """ - dictConfig(LOGGING_CONFIG) + configure_logging() parent_logger = logging.getLogger("") if name: diff --git a/tests/test_log.py b/tests/test_log.py index fbd8791ee..e799ff1c4 100644 --- a/tests/test_log.py +++ b/tests/test_log.py @@ -28,3 +28,34 @@ def test_get_logger_returns_logger(): assert any(isinstance(h, logging.StreamHandler) for h in logger.parent.handlers) or any( isinstance(h, logging.FileHandler) for h in logger.parent.handlers ) + + +def test_get_logger_configures_logging_once_per_process(monkeypatch): + calls = [] + + monkeypatch.setattr(log, "_LOGGING_CONFIGURED_PID", None) + monkeypatch.setattr(log.os, "getpid", lambda: 123) + monkeypatch.setattr(log, "dictConfig", lambda config: calls.append(config)) + + log.get_logger("first") + log.get_logger("second") + + assert len(calls) == 1 + + +def test_get_logger_reconfigures_after_process_fork(monkeypatch): + calls = [] + pid = 123 + + def getpid(): + return pid + + monkeypatch.setattr(log, "_LOGGING_CONFIGURED_PID", None) + monkeypatch.setattr(log.os, "getpid", getpid) + monkeypatch.setattr(log, "dictConfig", lambda config: calls.append(config)) + + log.get_logger("parent") + pid = 456 + log.get_logger("child") + + assert len(calls) == 2 From 2017cd5ecb67c5a960d0dfedfe877614c69e15b0 Mon Sep 17 00:00:00 2001 From: bittergreen Date: Thu, 16 Jul 2026 20:50:53 +0800 Subject: [PATCH 08/15] fix: close scheduler resources on API shutdown --- src/memos/api/lifecycle.py | 26 ++++++++++++++++++++++++++ src/memos/api/server_api.py | 10 ++++++++-- src/memos/api/server_api_ext.py | 12 +++++++++--- tests/api/test_lifecycle.py | 31 +++++++++++++++++++++++++++++++ 4 files changed, 74 insertions(+), 5 deletions(-) create mode 100644 src/memos/api/lifecycle.py create mode 100644 tests/api/test_lifecycle.py diff --git a/src/memos/api/lifecycle.py b/src/memos/api/lifecycle.py new file mode 100644 index 000000000..e05a2f777 --- /dev/null +++ b/src/memos/api/lifecycle.py @@ -0,0 +1,26 @@ +from collections.abc import Mapping +from typing import Any + +from memos.log import get_logger + + +logger = get_logger(__name__) + + +def shutdown_components(components: Mapping[str, Any] | None) -> None: + """Release long-lived API components before the logging system shuts down.""" + if not components: + return + + mem_scheduler = components.get("mem_scheduler") + if mem_scheduler is None: + return + + for method_name in ("stop", "rabbitmq_close"): + method = getattr(mem_scheduler, method_name, None) + if not callable(method): + continue + try: + method() + except Exception: + logger.exception("Failed to run mem_scheduler.%s during API shutdown", method_name) diff --git a/src/memos/api/server_api.py b/src/memos/api/server_api.py index a9afe554c..b05e8ac7c 100644 --- a/src/memos/api/server_api.py +++ b/src/memos/api/server_api.py @@ -7,8 +7,9 @@ from starlette.staticfiles import StaticFiles from memos.api.exceptions import APIExceptionHandler +from memos.api.lifecycle import shutdown_components from memos.api.middleware.request_context import RequestContextMiddleware -from memos.api.routers.server_router import router as server_router +from memos.api.routers import server_router as server_router_module from memos.plugins.manager import plugin_manager @@ -35,7 +36,7 @@ app.add_middleware(RequestContextMiddleware, source="server_api") # Include routers -app.include_router(server_router) +app.include_router(server_router_module.router) @app.get("/health") @@ -48,6 +49,11 @@ def health_check(): } +@app.on_event("shutdown") +def shutdown_server_components() -> None: + shutdown_components(server_router_module.components) + + # Request validation failed app.exception_handler(RequestValidationError)(APIExceptionHandler.validation_error_handler) # Invalid business code parameters diff --git a/src/memos/api/server_api_ext.py b/src/memos/api/server_api_ext.py index 8c457e362..e28505332 100644 --- a/src/memos/api/server_api_ext.py +++ b/src/memos/api/server_api_ext.py @@ -26,11 +26,12 @@ from starlette.responses import Response # Import Krolik extensions +from memos.api.lifecycle import shutdown_components from memos.api.middleware.rate_limit import RateLimitMiddleware -from memos.api.routers.admin_router import router as admin_router # Import base routers from MemOS -from memos.api.routers.server_router import router as server_router +from memos.api.routers import server_router as server_router_module +from memos.api.routers.admin_router import router as admin_router # Try to import exception handlers (may vary between MemOS versions) @@ -94,7 +95,7 @@ async def dispatch(self, request: Request, call_next) -> Response: logger.info("Rate limiting enabled") # Include routers -app.include_router(server_router) +app.include_router(server_router_module.router) app.include_router(admin_router) # Exception handlers @@ -118,6 +119,11 @@ async def health_check(): } +@app.on_event("shutdown") +def shutdown_server_components() -> None: + shutdown_components(server_router_module.components) + + if __name__ == "__main__": import uvicorn diff --git a/tests/api/test_lifecycle.py b/tests/api/test_lifecycle.py new file mode 100644 index 000000000..1daa7a22e --- /dev/null +++ b/tests/api/test_lifecycle.py @@ -0,0 +1,31 @@ +from memos.api.lifecycle import shutdown_components + + +class SchedulerStub: + def __init__(self, fail_stop: bool = False): + self.fail_stop = fail_stop + self.calls = [] + + def stop(self): + self.calls.append("stop") + if self.fail_stop: + raise RuntimeError("stop failed") + + def rabbitmq_close(self): + self.calls.append("rabbitmq_close") + + +def test_shutdown_components_stops_scheduler_before_rabbitmq_close(): + scheduler = SchedulerStub() + + shutdown_components({"mem_scheduler": scheduler}) + + assert scheduler.calls == ["stop", "rabbitmq_close"] + + +def test_shutdown_components_still_closes_rabbitmq_when_stop_fails(): + scheduler = SchedulerStub(fail_stop=True) + + shutdown_components({"mem_scheduler": scheduler}) + + assert scheduler.calls == ["stop", "rabbitmq_close"] From 8aefda7b571f8d4991a5c7b6b5513e5c6daca91d Mon Sep 17 00:00:00 2001 From: bittergreen Date: Thu, 16 Jul 2026 21:09:16 +0800 Subject: [PATCH 09/15] fix: configure logging once per process Avoid re-running dictConfig on every get_logger call, since repeated configuration can close and replace active handlers while background threads are emitting records. Keep logging configuration process-local and reconfigure only after a PID change. --- src/memos/log.py | 10 +++++----- tests/test_log.py | 4 ++-- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/src/memos/log.py b/src/memos/log.py index ede5e8280..c18bd2118 100644 --- a/src/memos/log.py +++ b/src/memos/log.py @@ -226,6 +226,10 @@ def close(self): } +def _get_current_pid() -> int: + return os.getpid() + + def configure_logging(force: bool = False) -> None: """Configure process-local logging once. @@ -234,12 +238,8 @@ def configure_logging(force: bool = False) -> None: """ global _LOGGING_CONFIGURED_PID - current_pid = os.getpid() - if not force and current_pid == _LOGGING_CONFIGURED_PID: - return - with _LOGGING_CONFIG_LOCK: - current_pid = os.getpid() + current_pid = _get_current_pid() if force or current_pid != _LOGGING_CONFIGURED_PID: dictConfig(LOGGING_CONFIG) _LOGGING_CONFIGURED_PID = current_pid diff --git a/tests/test_log.py b/tests/test_log.py index e799ff1c4..5387adc34 100644 --- a/tests/test_log.py +++ b/tests/test_log.py @@ -34,7 +34,7 @@ def test_get_logger_configures_logging_once_per_process(monkeypatch): calls = [] monkeypatch.setattr(log, "_LOGGING_CONFIGURED_PID", None) - monkeypatch.setattr(log.os, "getpid", lambda: 123) + monkeypatch.setattr(log, "_get_current_pid", lambda: 123) monkeypatch.setattr(log, "dictConfig", lambda config: calls.append(config)) log.get_logger("first") @@ -51,7 +51,7 @@ def getpid(): return pid monkeypatch.setattr(log, "_LOGGING_CONFIGURED_PID", None) - monkeypatch.setattr(log.os, "getpid", getpid) + monkeypatch.setattr(log, "_get_current_pid", getpid) monkeypatch.setattr(log, "dictConfig", lambda config: calls.append(config)) log.get_logger("parent") From 9c7dd2fdbf132a75e9be0eea73b096e41d1ee385 Mon Sep 17 00:00:00 2001 From: bittergreen Date: Thu, 16 Jul 2026 21:09:22 +0800 Subject: [PATCH 10/15] fix: close scheduler resources via API lifespan Use FastAPI lifespan cleanup to stop scheduler and RabbitMQ resources before logging handlers are torn down during worker or pod shutdown. --- src/memos/api/server_api.py | 16 +++++++++++----- src/memos/api/server_api_ext.py | 15 ++++++++++----- tests/api/test_lifecycle.py | 7 +++++++ 3 files changed, 28 insertions(+), 10 deletions(-) diff --git a/src/memos/api/server_api.py b/src/memos/api/server_api.py index b05e8ac7c..1c6d93b9f 100644 --- a/src/memos/api/server_api.py +++ b/src/memos/api/server_api.py @@ -1,6 +1,9 @@ import logging import os +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager + from dotenv import load_dotenv from fastapi import FastAPI, HTTPException from fastapi.exceptions import RequestValidationError @@ -26,10 +29,18 @@ os.getenv("MEMSCHEDULER_REDIS_STREAM_KEY_PREFIX"), ) + +@asynccontextmanager +async def lifespan(_app: FastAPI) -> AsyncIterator[None]: + yield + shutdown_components(server_router_module.components) + + app = FastAPI( title="MemOS Server REST APIs", description="A REST API for managing multiple users with MemOS Server.", version="1.0.1", + lifespan=lifespan, ) app.mount("/download", StaticFiles(directory=os.getenv("FILE_LOCAL_PATH")), name="static_mapping") @@ -49,11 +60,6 @@ def health_check(): } -@app.on_event("shutdown") -def shutdown_server_components() -> None: - shutdown_components(server_router_module.components) - - # Request validation failed app.exception_handler(RequestValidationError)(APIExceptionHandler.validation_error_handler) # Invalid business code parameters diff --git a/src/memos/api/server_api_ext.py b/src/memos/api/server_api_ext.py index e28505332..7b5aafc46 100644 --- a/src/memos/api/server_api_ext.py +++ b/src/memos/api/server_api_ext.py @@ -18,6 +18,9 @@ import logging import os +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager + from fastapi import FastAPI from fastapi.exceptions import RequestValidationError from fastapi.middleware.cors import CORSMiddleware @@ -59,11 +62,18 @@ async def dispatch(self, request: Request, call_next) -> Response: return response +@asynccontextmanager +async def lifespan(_app: FastAPI) -> AsyncIterator[None]: + yield + shutdown_components(server_router_module.components) + + # Create FastAPI app app = FastAPI( title="MemOS Server REST APIs (Krolik Extended)", description="MemOS API with authentication, rate limiting, and admin endpoints.", version="2.0.3-krolik", + lifespan=lifespan, ) # CORS configuration @@ -119,11 +129,6 @@ async def health_check(): } -@app.on_event("shutdown") -def shutdown_server_components() -> None: - shutdown_components(server_router_module.components) - - if __name__ == "__main__": import uvicorn diff --git a/tests/api/test_lifecycle.py b/tests/api/test_lifecycle.py index 1daa7a22e..b2f7f7064 100644 --- a/tests/api/test_lifecycle.py +++ b/tests/api/test_lifecycle.py @@ -29,3 +29,10 @@ def test_shutdown_components_still_closes_rabbitmq_when_stop_fails(): shutdown_components({"mem_scheduler": scheduler}) assert scheduler.calls == ["stop", "rabbitmq_close"] + + +def test_shutdown_components_skips_missing_methods(): + class MinimalScheduler: + pass + + shutdown_components({"mem_scheduler": MinimalScheduler()}) From e0726760e7c30c43cc5b7a6782622dcada2d229f Mon Sep 17 00:00:00 2001 From: dingliang <2650876010@qq.com> Date: Fri, 17 Jul 2026 11:57:10 +0800 Subject: [PATCH 11/15] feat(api):update SDK version number --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index c7297c7d1..00d240ba8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ ############################################################################## name = "MemoryOS" -version = "2.0.23" +version = "2.0.24" description = "Intelligence Begins with Memory" license = {text = "Apache-2.0"} readme = "README.md" From d0802b67bb42f13d859d250a0ff3f6ccaf4a1744 Mon Sep 17 00:00:00 2001 From: harvey_xiang Date: Sun, 19 Jul 2026 11:39:07 +0800 Subject: [PATCH 12/15] chore: update version to v2.0.24 --- src/memos/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/memos/__init__.py b/src/memos/__init__.py index 4b1858cf0..c42c7b8a5 100644 --- a/src/memos/__init__.py +++ b/src/memos/__init__.py @@ -1,4 +1,4 @@ -__version__ = "2.0.23" +__version__ = "2.0.24" from memos.configs.mem_cube import GeneralMemCubeConfig from memos.configs.mem_os import MOSConfig From d9130a7f9c5ea710ae36f032e830dba3892448d8 Mon Sep 17 00:00:00 2001 From: Erick Wipprecht <53452309+chiefmojo@users.noreply.github.com> Date: Thu, 23 Jul 2026 11:32:13 +0200 Subject: [PATCH 13/15] feat: chunk batch reflection scoring (#1957) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Fix #1540: fix: (#1824) docs(memos-local-plugin): clarify install path and stale dir names (#1540) The README's 'Quick start' section told users to use install.sh instead of npm install, but the warning was buried and users still tried 'npm install -g @memtensor/memos-local-plugin' first. The reporter in #1540 encountered this on a Hermes deployment. This change: - Promotes the 'do not run npm install -g' notice to a prominent IMPORTANT callout explaining why global install is wrong (no agent-home deploy, no config.yaml, no bridge/viewer) and that the tarball intentionally ships built artifacts only. - Adds a Troubleshooting subsection covering the two specific symptoms in the bug report: the 'package not found' misread, and the stale web/ and site/ directory names (web/ is now viewer/, site/ was removed by commit 26e7e3db). - Mentions install.ps1 for Windows alongside install.sh. - CHANGELOG: record the docs fix and reference #1540. Documentation-only change; no code or runtime behavior touched. Co-authored-by: MemOS AutoDev Co-authored-by: Matthew * Fix #1888: [Bug] test_system_parser.py: SystemParser.__init__() got an unexpected keyword a (#1889) fix: remove invalid chunker parameter from SystemParser test instantiation - SystemParser.__init__() signature changed to (embedder, llm=None) - Test was still passing chunker=None causing TypeError - Fixes all 5 failing tests in test_system_parser.py Fixes #1888 Co-authored-by: MemOS AutoDev Co-authored-by: Matthew * Fix #1525: [Bug] clean_json_response crashes with cryptic AttributeError when given None (#1884) * test: add comprehensive tests for clean_json_response (issue #1525) - Add test suite in tests/mem_os/test_format_utils.py - Cover None input ValueError with diagnostic message - Cover markdown removal, whitespace stripping, edge cases - Verify fix for AttributeError when LLM returns None * style: format clean_json_response tests --------- Co-authored-by: MemOS AutoDev Co-authored-by: Matthew * Fix #1901: share_cube_with_user passes swapped args to _validate_cube_access — fails for ev (#1903) fix: validate current user not target in share_cube_with_user (#1901) share_cube_with_user(cube_id, target_user_id) called _validate_cube_access(cube_id, target_user_id), but the validator signature is (user_id, cube_id). The cube_id therefore landed in the user_id slot and _validate_user_exists raised "User '' does not exist or is inactive" for every well-formed call, making the API unusable. The in-code comment "Validate current user has access to this cube" already documented the correct intent: the sharing user (self.user_id) must have access to the cube being shared, not the target. Switch the call to self._validate_cube_access(self.user_id, cube_id). The target user's existence is independently checked on the next line via validate_user(target_user_id), so that path is unchanged. Add regression tests in tests/mem_os/test_memos_core.py that pin down: - validate_user_cube_access is consulted with (self.user_id, cube_id), - add_user_to_cube is called with (target_user_id, cube_id) on success, - a missing target raises "Target user '' does not exist". Closes #1901 Co-authored-by: MemOS AutoDev Bot Co-authored-by: Matthew * feat(memos): chunk batch reflection scoring * Fix #1897: bug(memos-local-plugin): Hermes background model status checks can trigger paid (#1899) * Fix #1897: fix(memos-local-plugin): add LLM circuit breaker for terminal provider errors Issue #1897 reported ~12,900 paid LLM requests in 24 h on Hermes against a DeepSeek key with insufficient balance. The local `system_model_status` row count (12,900) closely tracked the provider-side `request_count` (11,344) for the same billing window. The naming is misleading: `system_model_status` is not a health probe; it is the audit row written once per LLM call (ok / fallback / error) inside `core/llm/client.ts`. With no circuit breaker, every pipeline subscriber (capture / session-relation / reward / L2 / L3 / skill / retrieval LLM filter / world-model) kept firing on every turn / closed episode / induction, generating one paid request each. Add a per-`LlmClient` circuit breaker: - Trips on terminal errors: HTTP 401/402/403 or messages containing `insufficient balance` / `invalid api key` / `unauthorized` / `account suspended` / `billing`. - Open: short-circuits subsequent calls inside the facade without contacting the provider. Throws `MemosError(LLM_UNAVAILABLE)` with `details.circuitOpen=true` so existing catch blocks still work. - Half-open after cool-down (default 5 min, configurable, min 30 s): next call probes the provider; success closes the breaker, terminal failure re-opens it for another cool-down. - Host fallback rescues a call without tripping the breaker — fallback exists precisely to keep going when the primary is down. - Coalesces `system_model_status="circuit_open"` audit rows to at most one per ~25 s while the breaker stays open, so we don't replace paid spam with audit-row spam. - Exposes `circuitOpen` / `circuitOpenUntil` / `circuitOpenedReason` via `LlmClientStats` for the Overview viewer card. - Enabled by default; legacy behaviour available via `circuitBreaker.enabled = false`. Tests: 9 new vitest cases under `tests/unit/llm/client.test.ts` covering trip on 402, trip on "insufficient balance" message, no trip on generic transient, coalescing, half-open close on success, host-fallback rescues without trip, disabled mode, stats fields, and re-open on terminal probe failure. All 59 LLM and 28 pipeline tests pass; `tsc --noEmit` clean. Out of scope (tracked separately): 429 `Retry-After` handling (issue #1620), per-tool rate limits, daily budget caps. * Fix LLM breaker with host fallback --------- Co-authored-by: autodev-bot Co-authored-by: Jiang <33757498+hijzy@users.noreply.github.com> Co-authored-by: GU TIANCHUN <96930846+TianchunGu@users.noreply.github.com> Co-authored-by: Dubberman <48425266+whipser030@users.noreply.github.com> --------- Co-authored-by: Memtensor-AI Co-authored-by: MemOS AutoDev Co-authored-by: Matthew Co-authored-by: MemOS AutoDev Co-authored-by: MemOS AutoDev Bot Co-authored-by: Jiang <33757498+hijzy@users.noreply.github.com> Co-authored-by: GU TIANCHUN <96930846+TianchunGu@users.noreply.github.com> Co-authored-by: Dubberman <48425266+whipser030@users.noreply.github.com> --- .../core/capture/ALGORITHMS.md | 13 +- .../core/capture/batch-scorer.ts | 21 +- .../core/capture/capture.ts | 75 +++++-- .../tests/helpers/fake-llm.ts | 4 +- .../tests/unit/capture/capture-batch.test.ts | 183 +++++++++++++++--- 5 files changed, 238 insertions(+), 58 deletions(-) diff --git a/apps/memos-local-plugin/core/capture/ALGORITHMS.md b/apps/memos-local-plugin/core/capture/ALGORITHMS.md index 15a15ac64..49c98a105 100644 --- a/apps/memos-local-plugin/core/capture/ALGORITHMS.md +++ b/apps/memos-local-plugin/core/capture/ALGORITHMS.md @@ -78,7 +78,8 @@ priority once reward arrives. ## V7 §3.2 batched variant — `batch-scorer.ts` The per-step path (`reflection-synth.ts` + `alpha-scorer.ts`) issues 2N -LLM calls per N-step episode. `batch-scorer.ts` collapses them into ONE: +LLM calls per N-step episode. `batch-scorer.ts` collapses up to +`batchThreshold` steps into one call: ``` inputs = [{idx, state, action, outcome, reflection, synth_allowed}, …] @@ -91,8 +92,8 @@ Dispatch (in `capture.ts`): | `cfg.batchMode` | `cfg.batchThreshold` | behavior | |-------------------|----------------------|----------| | `per_step` | (ignored) | legacy: 2N calls | -| `per_episode` | (ignored) | always batch | -| `auto` (default) | `12` | batch when `N ≤ 12`; else per-step | +| `per_episode` | chunk size | batch when `N ≤ threshold`; else chunk-batch | +| `auto` (default) | `12` | batch when `N ≤ 12`; else chunk-batch | The dispatcher also refuses to batch when no LLM is wired — same fallback path as missing-LLM in per-step mode. @@ -107,15 +108,15 @@ Failure handling: - LLM throws / facade gives up after `malformedRetries=1` → capture catches in `runBatchScoring`, surfaces a `{stage: "batch"}` warning, - and the per-step path runs as a fallback. + and the per-step path runs as a fallback for that chunk. - Validator rejects on length mismatch, missing/non-numeric `alpha`, non-boolean `usable`, non-string `reflection_text`. Same fallback. Bookkeeping (`CaptureResult.llmCalls`): -- `batchedReflection`: 0 or 1 per episode (1 on a successful batch). +- `batchedReflection`: number of successful batch/chunk calls. - `reflectionSynth` / `alphaScoring`: only nonzero when the per-step path - ran (either selected directly, or as fallback after a batch failure). + ran (either selected directly, or as fallback after a chunk failure). Stable prompt fingerprint: diff --git a/apps/memos-local-plugin/core/capture/batch-scorer.ts b/apps/memos-local-plugin/core/capture/batch-scorer.ts index e7b8ab50f..da434c3b2 100644 --- a/apps/memos-local-plugin/core/capture/batch-scorer.ts +++ b/apps/memos-local-plugin/core/capture/batch-scorer.ts @@ -16,11 +16,12 @@ * `transferability` axes benefit directly. * * Trade-offs (encoded in capture.ts dispatch): - * - Prompt grows linearly with N steps. Capped via `batchThreshold`; - * long episodes degrade to the per-step path automatically. - * - One bad output value forces a single batched retry instead of N - * isolated retries — but the facade already does `malformedRetries` - * for us, and on hard failure capture.ts falls back to per-step. + * - Prompt grows linearly with N steps. Each call is capped at + * `batchThreshold`; long episodes run as several bounded chunks. + * - One bad chunk forces a single batched retry for that chunk instead + * of N isolated retries — but the facade already does + * `malformedRetries` for us, and on hard failure capture.ts falls + * back to per-step for that chunk only. * * Wire format ↔ prompt: * Send `{ host_context?, task_context?, steps: [{idx, state, action, outcome, reflection, synth_allowed}] }`. @@ -170,6 +171,7 @@ export async function batchScoreReflections( validate: (v) => validateBatchPayload(v, inputs.length), malformedRetries: 1, temperature: 0, + maxTokens: batchMaxTokens(inputs.length), }, ); @@ -321,6 +323,15 @@ function validateBatchPayload(v: unknown, expected: number): void { } } +function batchMaxTokens(stepCount: number): number { + // Batch output scales with step count; keep a per-step budget but cap below + // the 16k range that triggered avoidable reasoning spend on mimo replay. + const perStepOutputBudget = 512; + const baseBudget = 768; + const ceiling = 8_192; + return Math.min(ceiling, baseBudget + Math.max(1, stepCount) * perStepOutputBudget); +} + function lastToolOutcome(step: NormalizedStep, max: number): string { const last = step.toolCalls[step.toolCalls.length - 1]; if (!last) return "(assistant-only step)"; diff --git a/apps/memos-local-plugin/core/capture/capture.ts b/apps/memos-local-plugin/core/capture/capture.ts index d5b62aa38..bc9c7a812 100644 --- a/apps/memos-local-plugin/core/capture/capture.ts +++ b/apps/memos-local-plugin/core/capture/capture.ts @@ -463,14 +463,14 @@ export function createCaptureRunner(deps: CaptureDeps): CaptureRunner { } // Batch reflection + α across every step of the now-closed - // episode. Falls back to per-step scoring when over the threshold - // or when batching fails / no LLM is wired. The reflect pass uses + // episode. Long episodes are chunk-batched at `batchThreshold`; + // failed chunks fall back to per-step scoring. The reflect pass uses // `reflectLlm` (skill-evolver model when configured) for higher // quality reflections; per-turn lite capture still uses `llm`. const reflectStart = now(); const rLlm = deps.reflectLlm ?? deps.llm; - const useBatch = shouldBatch(deps.cfg, normalized.length, rLlm !== null); - const contextEnabled = contextModeFor(deps.cfg, useBatch, normalized.length); + const scoringPlan = planScoring(deps.cfg, normalized.length, rLlm !== null); + const contextEnabled = contextModeFor(deps.cfg, scoringPlan, normalized.length); const taskSummary = contextEnabled.includeTask ? buildTaskReflectionSummary(input.episode, normalized, deps.cfg.taskContextMaxChars) : null; @@ -481,7 +481,10 @@ export function createCaptureRunner(deps: CaptureDeps): CaptureRunner { episodeId: input.episode.id, sessionId: input.episode.sessionId, steps: normalized.length, - mode: useBatch ? "batch" : contextEnabled.includeDownstream ? "per_step_downstream" : "per_step", + mode: scoringPlan === "per_step" && contextEnabled.includeDownstream ? "per_step_downstream" : scoringPlan, + chunks: scoringPlan === "chunk_batch" + ? Math.ceil(normalized.length / Math.max(1, deps.cfg.batchThreshold)) + : undefined, reflectionContextMode: deps.cfg.reflectionContextMode, downstreamPreview: contextEnabled.includeDownstream, provider: rLlm?.provider ?? "none", @@ -489,10 +492,13 @@ export function createCaptureRunner(deps: CaptureDeps): CaptureRunner { taskSummary: taskSummary ? taskSummary.slice(0, 240) : null, }); let scored: ScoredStep[] = []; - if (useBatch) { + if (scoringPlan === "batch") { scored = await runBatchScoring(normalized, rLlm!, deps, warnings, llmCalls, input.episode.id, taskSummary); } - if (!useBatch || scored.length === 0) { + if (scoringPlan === "chunk_batch") { + scored = await runChunkedBatchScoring(normalized, rLlm!, deps, warnings, llmCalls, input.episode.id, taskSummary); + } + if (scoringPlan === "per_step" || scored.length === 0) { scored = await runPerStepScoring( normalized, rLlm, @@ -1062,30 +1068,30 @@ export function createCaptureRunner(deps: CaptureDeps): CaptureRunner { // ─── helpers ──────────────────────────────────────────────────────────────── /** - * Decide whether to use the batched reflection+α path. + * Decide which reflection+α path to use. * * `per_step` → never (legacy path). - * `per_episode` → always, when an LLM is available. - * `auto` → batch when step count fits inside `batchThreshold`. + * `per_episode` → batch up to threshold, then chunk-batch. + * `auto` → batch up to threshold, then chunk-batch. */ -function shouldBatch(cfg: CaptureConfig, stepCount: number, hasLlm: boolean): boolean { - if (!hasLlm) return false; - if (stepCount === 0) return false; - if (cfg.batchMode === "per_step") return false; - if (cfg.batchMode === "per_episode") return true; - // "auto" - return stepCount <= cfg.batchThreshold; +type ScoringPlan = "per_step" | "batch" | "chunk_batch"; + +function planScoring(cfg: CaptureConfig, stepCount: number, hasLlm: boolean): ScoringPlan { + if (!hasLlm) return "per_step"; + if (stepCount === 0) return "per_step"; + if (cfg.batchMode === "per_step") return "per_step"; + return stepCount <= Math.max(1, cfg.batchThreshold) ? "batch" : "chunk_batch"; } function contextModeFor( cfg: CaptureConfig, - useBatch: boolean, + scoringPlan: ScoringPlan, stepCount: number, ): { includeTask: boolean; includeDownstream: boolean } { const mode = cfg.reflectionContextMode; const includeTask = mode === "task" || mode === "task_downstream"; const wantsDownstream = mode === "downstream" || mode === "task_downstream"; - const longPerStep = !useBatch && stepCount > cfg.batchThreshold; + const longPerStep = scoringPlan === "per_step" && stepCount > cfg.batchThreshold; const includeDownstream = wantsDownstream && cfg.longEpisodeReflectMode === "per_step_downstream" && @@ -1145,6 +1151,37 @@ async function runBatchScoring( } } +async function runChunkedBatchScoring( + normalized: NormalizedStep[], + llm: LlmClient, + deps: CaptureDeps, + warnings: CaptureResult["warnings"], + llmCalls: { reflectionSynth: number; alphaScoring: number; batchedReflection: number }, + episodeId: string, + taskSummary: string | null, +): Promise { + const chunkSize = Math.max(1, deps.cfg.batchThreshold); + const chunks: NormalizedStep[][] = []; + for (let start = 0; start < normalized.length; start += chunkSize) { + chunks.push(normalized.slice(start, start + chunkSize)); + } + const concurrency = Math.max(1, deps.cfg.llmConcurrency); + const scoredChunks = await runConcurrently(chunks, concurrency, async (chunk): Promise => { + const scored = await runBatchScoring(chunk, llm, deps, warnings, llmCalls, episodeId, taskSummary); + if (scored.length > 0) return scored; + return runPerStepScoring( + chunk, + llm, + deps, + warnings, + llmCalls, + episodeId, + buildReflectionContexts(chunk, taskSummary, chunk.map(() => [])), + ); + }); + return scoredChunks.flat(); +} + async function runPerStepScoring( normalized: NormalizedStep[], llm: LlmClient | null, diff --git a/apps/memos-local-plugin/tests/helpers/fake-llm.ts b/apps/memos-local-plugin/tests/helpers/fake-llm.ts index 22d9fc1a3..ec2c00261 100644 --- a/apps/memos-local-plugin/tests/helpers/fake-llm.ts +++ b/apps/memos-local-plugin/tests/helpers/fake-llm.ts @@ -20,7 +20,7 @@ export interface FakeLlmScript { complete?: Record string | Promise)>; completeJson?: Record< string, - unknown | ((input: unknown) => unknown | Promise) + unknown | ((input: unknown, opts?: unknown) => unknown | Promise) >; /** Override the served-by identifier. */ servedBy?: LlmProviderName | "host_fallback"; @@ -64,7 +64,7 @@ export function fakeLlm(script: FakeLlmScript = {}): LlmClient { throw new Error(`fakeLlm: no completeJson mock for op="${op}"`); } const value = (typeof entry === "function" - ? await (entry as (x: unknown) => unknown)(input) + ? await (entry as (x: unknown, o?: unknown) => unknown)(input, opts) : entry) as T; if (o?.validate) o.validate(value); return { diff --git a/apps/memos-local-plugin/tests/unit/capture/capture-batch.test.ts b/apps/memos-local-plugin/tests/unit/capture/capture-batch.test.ts index d86290517..bc4b76f28 100644 --- a/apps/memos-local-plugin/tests/unit/capture/capture-batch.test.ts +++ b/apps/memos-local-plugin/tests/unit/capture/capture-batch.test.ts @@ -7,14 +7,15 @@ * 2. existing reflections are preserved verbatim; * 3. synth-disabled steps stay at α=0 even when the LLM tries to write * one for them; - * 4. `auto` mode falls back to per-step when stepCount > batchThreshold; - * 5. a malformed batched response degrades into the per-step path - * instead of crashing capture. + * 4. `auto` mode chunk-batches when stepCount > batchThreshold; + * 5. a malformed chunk degrades only that chunk into the per-step path + * instead of dropping the whole episode to per-step. */ import { afterEach, beforeAll, beforeEach, describe, expect, it } from "vitest"; import { createCaptureRunner, type CaptureRunner } from "../../../core/capture/capture.js"; +import { batchScoreReflections } from "../../../core/capture/batch-scorer.js"; import { createCaptureEventBus } from "../../../core/capture/events.js"; import { BATCH_REFLECTION_PROMPT, @@ -312,20 +313,31 @@ describe("capture/pipeline (batched ρ+α path)", () => { expect(t.alpha).toBe(0); // V7 disabledScore semantics }); - it("auto mode falls back to per-step when stepCount > batchThreshold", async () => { + it("auto mode chunk-batches when stepCount > batchThreshold", async () => { + const batchStates: string[][] = []; const llm = fakeLlm({ completeJson: { - // ONLY per-step alpha mock; if batched gets called, the test fails - // with "no completeJson mock for op=...batch...". - [alphaOp]: { alpha: 0.5, usable: true, reason: "ok" }, - }, - complete: { - "capture.reflection.synth": "I made this decision deliberately.", + [batchOp]: (input) => { + const messages = input as Array<{ role: string; content: string }>; + const payload = JSON.parse(messages[messages.length - 1]!.content) as { + steps: Array<{ idx: number; state: string }>; + }; + batchStates.push(payload.steps.map((s) => s.state)); + return { + scores: payload.steps.map((step) => ({ + idx: step.idx, + reflection_text: `reflection ${step.state}`, + alpha: step.idx === 0 ? 0.2 : 0.4, + usable: true, + reason: "ok", + })), + }; + }, }, }); const runner = buildRunner({ batchMode: "auto", batchThreshold: 2 }, llm); - // 3 steps → above threshold → per-step path. + // 3 steps → above threshold → two bounded batch chunks. const ep = episodeSnapshot({ id: "ep_1", sessionId: "se_1", @@ -341,10 +353,17 @@ describe("capture/pipeline (batched ρ+α path)", () => { const result = await runCapture(runner, ep); expect(result.traceIds).toHaveLength(3); - expect(result.llmCalls.batchedReflection).toBe(0); - // 3 synth + 3 alpha calls in per-step mode. - expect(result.llmCalls.reflectionSynth).toBe(3); - expect(result.llmCalls.alphaScoring).toBe(3); + expect(batchStates).toEqual([["a", "b"], ["c"]]); + expect(result.llmCalls.batchedReflection).toBe(2); + expect(result.llmCalls.reflectionSynth).toBe(0); + expect(result.llmCalls.alphaScoring).toBe(0); + + const rows = result.traceIds.map((id) => tmp.repos.traces.getById(id)!); + expect(rows.map((row) => row.reflection)).toEqual([ + "reflection a", + "reflection b", + "reflection c", + ]); }); it("long per-step downstream mode injects up to three following steps", async () => { @@ -368,7 +387,7 @@ describe("capture/pipeline (batched ρ+α path)", () => { }); const runner = buildRunner( { - batchMode: "auto", + batchMode: "per_step", batchThreshold: 2, reflectionContextMode: "task_downstream", longEpisodeReflectMode: "per_step_downstream", @@ -415,16 +434,27 @@ describe("capture/pipeline (batched ρ+α path)", () => { expect(step3Prompt).not.toContain("[step+3]"); }); - it("per_episode mode batches even when step count is large", async () => { - const scores = Array.from({ length: 5 }, (_, i) => ({ - idx: i, - reflection_text: `reflection #${i}`, - alpha: 0.4, - usable: true, - reason: "ok", - })); + it("per_episode mode chunk-batches when step count is large", async () => { + const chunkSizes: number[] = []; const llm = fakeLlm({ - completeJson: { [batchOp]: { scores } }, + completeJson: { + [batchOp]: (input) => { + const messages = input as Array<{ role: string; content: string }>; + const payload = JSON.parse(messages[messages.length - 1]!.content) as { + steps: Array<{ idx: number; state: string }>; + }; + chunkSizes.push(payload.steps.length); + return { + scores: payload.steps.map((step) => ({ + idx: step.idx, + reflection_text: `reflection ${step.state}`, + alpha: 0.4, + usable: true, + reason: "ok", + })), + }; + }, + }, }); const runner = buildRunner({ batchMode: "per_episode", batchThreshold: 2 }, llm); @@ -436,10 +466,111 @@ describe("capture/pipeline (batched ρ+α path)", () => { const ep = episodeSnapshot({ id: "ep_1", sessionId: "se_1", turns }); const result = await runCapture(runner, ep); expect(result.traceIds).toHaveLength(5); - expect(result.llmCalls.batchedReflection).toBe(1); + expect(chunkSizes).toEqual([2, 2, 1]); + expect(result.llmCalls.batchedReflection).toBe(3); expect(result.llmCalls.alphaScoring).toBe(0); }); + it("chunk-batch falls back to per-step only for the failed chunk", async () => { + const llm = fakeLlm({ + completeJson: { + [batchOp]: (input) => { + const messages = input as Array<{ role: string; content: string }>; + const payload = JSON.parse(messages[messages.length - 1]!.content) as { + steps: Array<{ idx: number; state: string }>; + }; + if (payload.steps[0]?.state === "q2") { + throw new Error("chunk failed"); + } + return { + scores: payload.steps.map((step) => ({ + idx: step.idx, + reflection_text: `batch ${step.state}`, + alpha: step.state === "q4" ? 0.5 : 0.2, + usable: true, + reason: "ok", + })), + }; + }, + [alphaOp]: { alpha: 0.9, usable: true, reason: "fallback" }, + }, + complete: { + "capture.reflection.synth": "per-step fallback reflection", + }, + }); + const runner = buildRunner({ batchMode: "auto", batchThreshold: 2 }, llm); + + const turns: EpisodeTurn[] = []; + for (let i = 0; i < 5; i++) { + turns.push(turn("user", `q${i}`, 1_000 + i * 100)); + turns.push(turn("assistant", `a${i}`, 1_050 + i * 100)); + } + const ep = episodeSnapshot({ id: "ep_1", sessionId: "se_1", turns }); + + const result = await runCapture(runner, ep); + expect(result.traceIds).toHaveLength(5); + expect(result.llmCalls.batchedReflection).toBe(2); + expect(result.llmCalls.reflectionSynth).toBe(2); + expect(result.llmCalls.alphaScoring).toBe(2); + expect(result.warnings.filter((w) => w.stage === "batch")).toHaveLength(1); + + const rows = result.traceIds.map((id) => tmp.repos.traces.getById(id)!); + expect(rows.map((row) => row.reflection)).toEqual([ + "batch q0", + "batch q1", + "per-step fallback reflection", + "per-step fallback reflection", + "batch q4", + ]); + expect(rows.map((row) => row.alpha)).toEqual([0.2, 0.2, 0.9, 0.9, 0.5]); + }); + + it("batch scorer passes an explicit maxTokens budget", async () => { + let seenMaxTokens: number | undefined; + const llm = fakeLlm({ + completeJson: { + [batchOp]: (_input, opts) => { + seenMaxTokens = (opts as { maxTokens?: number }).maxTokens; + return { + scores: [ + { + idx: 0, + reflection_text: "I made a useful choice.", + alpha: 0.5, + usable: true, + reason: "ok", + }, + ], + }; + }, + }, + }); + + await batchScoreReflections( + llm, + [ + { + step: { + key: "s1", + ts: 1_000 as EpochMs, + type: "text", + userText: "q", + agentText: "a", + agentThinking: null, + toolCalls: [], + rawReflection: null, + meta: {}, + }, + existingReflection: null, + }, + ], + { synthReflections: true }, + ); + + expect(seenMaxTokens).toBeGreaterThan(0); + expect(seenMaxTokens).toBeLessThan(16_384); + }); + it("malformed batched response → falls back to per-step + emits warning", async () => { const llm = fakeLlm({ completeJson: { From df41fe7b1e2188286a7ed407ddc08395f60505c0 Mon Sep 17 00:00:00 2001 From: zhaxi Date: Thu, 23 Jul 2026 17:50:48 +0800 Subject: [PATCH 14/15] Revert "feat: chunk batch reflection scoring (#1957)" (#2145) This reverts commit 21a27ed19dcbbfcf830a3ce1253aed5b2ee66898. Co-authored-by: jiachengzhen --- .../core/capture/ALGORITHMS.md | 13 +- .../core/capture/batch-scorer.ts | 21 +- .../core/capture/capture.ts | 75 ++----- .../tests/helpers/fake-llm.ts | 4 +- .../tests/unit/capture/capture-batch.test.ts | 183 +++--------------- 5 files changed, 58 insertions(+), 238 deletions(-) diff --git a/apps/memos-local-plugin/core/capture/ALGORITHMS.md b/apps/memos-local-plugin/core/capture/ALGORITHMS.md index 49c98a105..15a15ac64 100644 --- a/apps/memos-local-plugin/core/capture/ALGORITHMS.md +++ b/apps/memos-local-plugin/core/capture/ALGORITHMS.md @@ -78,8 +78,7 @@ priority once reward arrives. ## V7 §3.2 batched variant — `batch-scorer.ts` The per-step path (`reflection-synth.ts` + `alpha-scorer.ts`) issues 2N -LLM calls per N-step episode. `batch-scorer.ts` collapses up to -`batchThreshold` steps into one call: +LLM calls per N-step episode. `batch-scorer.ts` collapses them into ONE: ``` inputs = [{idx, state, action, outcome, reflection, synth_allowed}, …] @@ -92,8 +91,8 @@ Dispatch (in `capture.ts`): | `cfg.batchMode` | `cfg.batchThreshold` | behavior | |-------------------|----------------------|----------| | `per_step` | (ignored) | legacy: 2N calls | -| `per_episode` | chunk size | batch when `N ≤ threshold`; else chunk-batch | -| `auto` (default) | `12` | batch when `N ≤ 12`; else chunk-batch | +| `per_episode` | (ignored) | always batch | +| `auto` (default) | `12` | batch when `N ≤ 12`; else per-step | The dispatcher also refuses to batch when no LLM is wired — same fallback path as missing-LLM in per-step mode. @@ -108,15 +107,15 @@ Failure handling: - LLM throws / facade gives up after `malformedRetries=1` → capture catches in `runBatchScoring`, surfaces a `{stage: "batch"}` warning, - and the per-step path runs as a fallback for that chunk. + and the per-step path runs as a fallback. - Validator rejects on length mismatch, missing/non-numeric `alpha`, non-boolean `usable`, non-string `reflection_text`. Same fallback. Bookkeeping (`CaptureResult.llmCalls`): -- `batchedReflection`: number of successful batch/chunk calls. +- `batchedReflection`: 0 or 1 per episode (1 on a successful batch). - `reflectionSynth` / `alphaScoring`: only nonzero when the per-step path - ran (either selected directly, or as fallback after a chunk failure). + ran (either selected directly, or as fallback after a batch failure). Stable prompt fingerprint: diff --git a/apps/memos-local-plugin/core/capture/batch-scorer.ts b/apps/memos-local-plugin/core/capture/batch-scorer.ts index da434c3b2..e7b8ab50f 100644 --- a/apps/memos-local-plugin/core/capture/batch-scorer.ts +++ b/apps/memos-local-plugin/core/capture/batch-scorer.ts @@ -16,12 +16,11 @@ * `transferability` axes benefit directly. * * Trade-offs (encoded in capture.ts dispatch): - * - Prompt grows linearly with N steps. Each call is capped at - * `batchThreshold`; long episodes run as several bounded chunks. - * - One bad chunk forces a single batched retry for that chunk instead - * of N isolated retries — but the facade already does - * `malformedRetries` for us, and on hard failure capture.ts falls - * back to per-step for that chunk only. + * - Prompt grows linearly with N steps. Capped via `batchThreshold`; + * long episodes degrade to the per-step path automatically. + * - One bad output value forces a single batched retry instead of N + * isolated retries — but the facade already does `malformedRetries` + * for us, and on hard failure capture.ts falls back to per-step. * * Wire format ↔ prompt: * Send `{ host_context?, task_context?, steps: [{idx, state, action, outcome, reflection, synth_allowed}] }`. @@ -171,7 +170,6 @@ export async function batchScoreReflections( validate: (v) => validateBatchPayload(v, inputs.length), malformedRetries: 1, temperature: 0, - maxTokens: batchMaxTokens(inputs.length), }, ); @@ -323,15 +321,6 @@ function validateBatchPayload(v: unknown, expected: number): void { } } -function batchMaxTokens(stepCount: number): number { - // Batch output scales with step count; keep a per-step budget but cap below - // the 16k range that triggered avoidable reasoning spend on mimo replay. - const perStepOutputBudget = 512; - const baseBudget = 768; - const ceiling = 8_192; - return Math.min(ceiling, baseBudget + Math.max(1, stepCount) * perStepOutputBudget); -} - function lastToolOutcome(step: NormalizedStep, max: number): string { const last = step.toolCalls[step.toolCalls.length - 1]; if (!last) return "(assistant-only step)"; diff --git a/apps/memos-local-plugin/core/capture/capture.ts b/apps/memos-local-plugin/core/capture/capture.ts index bc9c7a812..d5b62aa38 100644 --- a/apps/memos-local-plugin/core/capture/capture.ts +++ b/apps/memos-local-plugin/core/capture/capture.ts @@ -463,14 +463,14 @@ export function createCaptureRunner(deps: CaptureDeps): CaptureRunner { } // Batch reflection + α across every step of the now-closed - // episode. Long episodes are chunk-batched at `batchThreshold`; - // failed chunks fall back to per-step scoring. The reflect pass uses + // episode. Falls back to per-step scoring when over the threshold + // or when batching fails / no LLM is wired. The reflect pass uses // `reflectLlm` (skill-evolver model when configured) for higher // quality reflections; per-turn lite capture still uses `llm`. const reflectStart = now(); const rLlm = deps.reflectLlm ?? deps.llm; - const scoringPlan = planScoring(deps.cfg, normalized.length, rLlm !== null); - const contextEnabled = contextModeFor(deps.cfg, scoringPlan, normalized.length); + const useBatch = shouldBatch(deps.cfg, normalized.length, rLlm !== null); + const contextEnabled = contextModeFor(deps.cfg, useBatch, normalized.length); const taskSummary = contextEnabled.includeTask ? buildTaskReflectionSummary(input.episode, normalized, deps.cfg.taskContextMaxChars) : null; @@ -481,10 +481,7 @@ export function createCaptureRunner(deps: CaptureDeps): CaptureRunner { episodeId: input.episode.id, sessionId: input.episode.sessionId, steps: normalized.length, - mode: scoringPlan === "per_step" && contextEnabled.includeDownstream ? "per_step_downstream" : scoringPlan, - chunks: scoringPlan === "chunk_batch" - ? Math.ceil(normalized.length / Math.max(1, deps.cfg.batchThreshold)) - : undefined, + mode: useBatch ? "batch" : contextEnabled.includeDownstream ? "per_step_downstream" : "per_step", reflectionContextMode: deps.cfg.reflectionContextMode, downstreamPreview: contextEnabled.includeDownstream, provider: rLlm?.provider ?? "none", @@ -492,13 +489,10 @@ export function createCaptureRunner(deps: CaptureDeps): CaptureRunner { taskSummary: taskSummary ? taskSummary.slice(0, 240) : null, }); let scored: ScoredStep[] = []; - if (scoringPlan === "batch") { + if (useBatch) { scored = await runBatchScoring(normalized, rLlm!, deps, warnings, llmCalls, input.episode.id, taskSummary); } - if (scoringPlan === "chunk_batch") { - scored = await runChunkedBatchScoring(normalized, rLlm!, deps, warnings, llmCalls, input.episode.id, taskSummary); - } - if (scoringPlan === "per_step" || scored.length === 0) { + if (!useBatch || scored.length === 0) { scored = await runPerStepScoring( normalized, rLlm, @@ -1068,30 +1062,30 @@ export function createCaptureRunner(deps: CaptureDeps): CaptureRunner { // ─── helpers ──────────────────────────────────────────────────────────────── /** - * Decide which reflection+α path to use. + * Decide whether to use the batched reflection+α path. * * `per_step` → never (legacy path). - * `per_episode` → batch up to threshold, then chunk-batch. - * `auto` → batch up to threshold, then chunk-batch. + * `per_episode` → always, when an LLM is available. + * `auto` → batch when step count fits inside `batchThreshold`. */ -type ScoringPlan = "per_step" | "batch" | "chunk_batch"; - -function planScoring(cfg: CaptureConfig, stepCount: number, hasLlm: boolean): ScoringPlan { - if (!hasLlm) return "per_step"; - if (stepCount === 0) return "per_step"; - if (cfg.batchMode === "per_step") return "per_step"; - return stepCount <= Math.max(1, cfg.batchThreshold) ? "batch" : "chunk_batch"; +function shouldBatch(cfg: CaptureConfig, stepCount: number, hasLlm: boolean): boolean { + if (!hasLlm) return false; + if (stepCount === 0) return false; + if (cfg.batchMode === "per_step") return false; + if (cfg.batchMode === "per_episode") return true; + // "auto" + return stepCount <= cfg.batchThreshold; } function contextModeFor( cfg: CaptureConfig, - scoringPlan: ScoringPlan, + useBatch: boolean, stepCount: number, ): { includeTask: boolean; includeDownstream: boolean } { const mode = cfg.reflectionContextMode; const includeTask = mode === "task" || mode === "task_downstream"; const wantsDownstream = mode === "downstream" || mode === "task_downstream"; - const longPerStep = scoringPlan === "per_step" && stepCount > cfg.batchThreshold; + const longPerStep = !useBatch && stepCount > cfg.batchThreshold; const includeDownstream = wantsDownstream && cfg.longEpisodeReflectMode === "per_step_downstream" && @@ -1151,37 +1145,6 @@ async function runBatchScoring( } } -async function runChunkedBatchScoring( - normalized: NormalizedStep[], - llm: LlmClient, - deps: CaptureDeps, - warnings: CaptureResult["warnings"], - llmCalls: { reflectionSynth: number; alphaScoring: number; batchedReflection: number }, - episodeId: string, - taskSummary: string | null, -): Promise { - const chunkSize = Math.max(1, deps.cfg.batchThreshold); - const chunks: NormalizedStep[][] = []; - for (let start = 0; start < normalized.length; start += chunkSize) { - chunks.push(normalized.slice(start, start + chunkSize)); - } - const concurrency = Math.max(1, deps.cfg.llmConcurrency); - const scoredChunks = await runConcurrently(chunks, concurrency, async (chunk): Promise => { - const scored = await runBatchScoring(chunk, llm, deps, warnings, llmCalls, episodeId, taskSummary); - if (scored.length > 0) return scored; - return runPerStepScoring( - chunk, - llm, - deps, - warnings, - llmCalls, - episodeId, - buildReflectionContexts(chunk, taskSummary, chunk.map(() => [])), - ); - }); - return scoredChunks.flat(); -} - async function runPerStepScoring( normalized: NormalizedStep[], llm: LlmClient | null, diff --git a/apps/memos-local-plugin/tests/helpers/fake-llm.ts b/apps/memos-local-plugin/tests/helpers/fake-llm.ts index ec2c00261..22d9fc1a3 100644 --- a/apps/memos-local-plugin/tests/helpers/fake-llm.ts +++ b/apps/memos-local-plugin/tests/helpers/fake-llm.ts @@ -20,7 +20,7 @@ export interface FakeLlmScript { complete?: Record string | Promise)>; completeJson?: Record< string, - unknown | ((input: unknown, opts?: unknown) => unknown | Promise) + unknown | ((input: unknown) => unknown | Promise) >; /** Override the served-by identifier. */ servedBy?: LlmProviderName | "host_fallback"; @@ -64,7 +64,7 @@ export function fakeLlm(script: FakeLlmScript = {}): LlmClient { throw new Error(`fakeLlm: no completeJson mock for op="${op}"`); } const value = (typeof entry === "function" - ? await (entry as (x: unknown, o?: unknown) => unknown)(input, opts) + ? await (entry as (x: unknown) => unknown)(input) : entry) as T; if (o?.validate) o.validate(value); return { diff --git a/apps/memos-local-plugin/tests/unit/capture/capture-batch.test.ts b/apps/memos-local-plugin/tests/unit/capture/capture-batch.test.ts index bc4b76f28..d86290517 100644 --- a/apps/memos-local-plugin/tests/unit/capture/capture-batch.test.ts +++ b/apps/memos-local-plugin/tests/unit/capture/capture-batch.test.ts @@ -7,15 +7,14 @@ * 2. existing reflections are preserved verbatim; * 3. synth-disabled steps stay at α=0 even when the LLM tries to write * one for them; - * 4. `auto` mode chunk-batches when stepCount > batchThreshold; - * 5. a malformed chunk degrades only that chunk into the per-step path - * instead of dropping the whole episode to per-step. + * 4. `auto` mode falls back to per-step when stepCount > batchThreshold; + * 5. a malformed batched response degrades into the per-step path + * instead of crashing capture. */ import { afterEach, beforeAll, beforeEach, describe, expect, it } from "vitest"; import { createCaptureRunner, type CaptureRunner } from "../../../core/capture/capture.js"; -import { batchScoreReflections } from "../../../core/capture/batch-scorer.js"; import { createCaptureEventBus } from "../../../core/capture/events.js"; import { BATCH_REFLECTION_PROMPT, @@ -313,31 +312,20 @@ describe("capture/pipeline (batched ρ+α path)", () => { expect(t.alpha).toBe(0); // V7 disabledScore semantics }); - it("auto mode chunk-batches when stepCount > batchThreshold", async () => { - const batchStates: string[][] = []; + it("auto mode falls back to per-step when stepCount > batchThreshold", async () => { const llm = fakeLlm({ completeJson: { - [batchOp]: (input) => { - const messages = input as Array<{ role: string; content: string }>; - const payload = JSON.parse(messages[messages.length - 1]!.content) as { - steps: Array<{ idx: number; state: string }>; - }; - batchStates.push(payload.steps.map((s) => s.state)); - return { - scores: payload.steps.map((step) => ({ - idx: step.idx, - reflection_text: `reflection ${step.state}`, - alpha: step.idx === 0 ? 0.2 : 0.4, - usable: true, - reason: "ok", - })), - }; - }, + // ONLY per-step alpha mock; if batched gets called, the test fails + // with "no completeJson mock for op=...batch...". + [alphaOp]: { alpha: 0.5, usable: true, reason: "ok" }, + }, + complete: { + "capture.reflection.synth": "I made this decision deliberately.", }, }); const runner = buildRunner({ batchMode: "auto", batchThreshold: 2 }, llm); - // 3 steps → above threshold → two bounded batch chunks. + // 3 steps → above threshold → per-step path. const ep = episodeSnapshot({ id: "ep_1", sessionId: "se_1", @@ -353,17 +341,10 @@ describe("capture/pipeline (batched ρ+α path)", () => { const result = await runCapture(runner, ep); expect(result.traceIds).toHaveLength(3); - expect(batchStates).toEqual([["a", "b"], ["c"]]); - expect(result.llmCalls.batchedReflection).toBe(2); - expect(result.llmCalls.reflectionSynth).toBe(0); - expect(result.llmCalls.alphaScoring).toBe(0); - - const rows = result.traceIds.map((id) => tmp.repos.traces.getById(id)!); - expect(rows.map((row) => row.reflection)).toEqual([ - "reflection a", - "reflection b", - "reflection c", - ]); + expect(result.llmCalls.batchedReflection).toBe(0); + // 3 synth + 3 alpha calls in per-step mode. + expect(result.llmCalls.reflectionSynth).toBe(3); + expect(result.llmCalls.alphaScoring).toBe(3); }); it("long per-step downstream mode injects up to three following steps", async () => { @@ -387,7 +368,7 @@ describe("capture/pipeline (batched ρ+α path)", () => { }); const runner = buildRunner( { - batchMode: "per_step", + batchMode: "auto", batchThreshold: 2, reflectionContextMode: "task_downstream", longEpisodeReflectMode: "per_step_downstream", @@ -434,27 +415,16 @@ describe("capture/pipeline (batched ρ+α path)", () => { expect(step3Prompt).not.toContain("[step+3]"); }); - it("per_episode mode chunk-batches when step count is large", async () => { - const chunkSizes: number[] = []; + it("per_episode mode batches even when step count is large", async () => { + const scores = Array.from({ length: 5 }, (_, i) => ({ + idx: i, + reflection_text: `reflection #${i}`, + alpha: 0.4, + usable: true, + reason: "ok", + })); const llm = fakeLlm({ - completeJson: { - [batchOp]: (input) => { - const messages = input as Array<{ role: string; content: string }>; - const payload = JSON.parse(messages[messages.length - 1]!.content) as { - steps: Array<{ idx: number; state: string }>; - }; - chunkSizes.push(payload.steps.length); - return { - scores: payload.steps.map((step) => ({ - idx: step.idx, - reflection_text: `reflection ${step.state}`, - alpha: 0.4, - usable: true, - reason: "ok", - })), - }; - }, - }, + completeJson: { [batchOp]: { scores } }, }); const runner = buildRunner({ batchMode: "per_episode", batchThreshold: 2 }, llm); @@ -466,111 +436,10 @@ describe("capture/pipeline (batched ρ+α path)", () => { const ep = episodeSnapshot({ id: "ep_1", sessionId: "se_1", turns }); const result = await runCapture(runner, ep); expect(result.traceIds).toHaveLength(5); - expect(chunkSizes).toEqual([2, 2, 1]); - expect(result.llmCalls.batchedReflection).toBe(3); + expect(result.llmCalls.batchedReflection).toBe(1); expect(result.llmCalls.alphaScoring).toBe(0); }); - it("chunk-batch falls back to per-step only for the failed chunk", async () => { - const llm = fakeLlm({ - completeJson: { - [batchOp]: (input) => { - const messages = input as Array<{ role: string; content: string }>; - const payload = JSON.parse(messages[messages.length - 1]!.content) as { - steps: Array<{ idx: number; state: string }>; - }; - if (payload.steps[0]?.state === "q2") { - throw new Error("chunk failed"); - } - return { - scores: payload.steps.map((step) => ({ - idx: step.idx, - reflection_text: `batch ${step.state}`, - alpha: step.state === "q4" ? 0.5 : 0.2, - usable: true, - reason: "ok", - })), - }; - }, - [alphaOp]: { alpha: 0.9, usable: true, reason: "fallback" }, - }, - complete: { - "capture.reflection.synth": "per-step fallback reflection", - }, - }); - const runner = buildRunner({ batchMode: "auto", batchThreshold: 2 }, llm); - - const turns: EpisodeTurn[] = []; - for (let i = 0; i < 5; i++) { - turns.push(turn("user", `q${i}`, 1_000 + i * 100)); - turns.push(turn("assistant", `a${i}`, 1_050 + i * 100)); - } - const ep = episodeSnapshot({ id: "ep_1", sessionId: "se_1", turns }); - - const result = await runCapture(runner, ep); - expect(result.traceIds).toHaveLength(5); - expect(result.llmCalls.batchedReflection).toBe(2); - expect(result.llmCalls.reflectionSynth).toBe(2); - expect(result.llmCalls.alphaScoring).toBe(2); - expect(result.warnings.filter((w) => w.stage === "batch")).toHaveLength(1); - - const rows = result.traceIds.map((id) => tmp.repos.traces.getById(id)!); - expect(rows.map((row) => row.reflection)).toEqual([ - "batch q0", - "batch q1", - "per-step fallback reflection", - "per-step fallback reflection", - "batch q4", - ]); - expect(rows.map((row) => row.alpha)).toEqual([0.2, 0.2, 0.9, 0.9, 0.5]); - }); - - it("batch scorer passes an explicit maxTokens budget", async () => { - let seenMaxTokens: number | undefined; - const llm = fakeLlm({ - completeJson: { - [batchOp]: (_input, opts) => { - seenMaxTokens = (opts as { maxTokens?: number }).maxTokens; - return { - scores: [ - { - idx: 0, - reflection_text: "I made a useful choice.", - alpha: 0.5, - usable: true, - reason: "ok", - }, - ], - }; - }, - }, - }); - - await batchScoreReflections( - llm, - [ - { - step: { - key: "s1", - ts: 1_000 as EpochMs, - type: "text", - userText: "q", - agentText: "a", - agentThinking: null, - toolCalls: [], - rawReflection: null, - meta: {}, - }, - existingReflection: null, - }, - ], - { synthReflections: true }, - ); - - expect(seenMaxTokens).toBeGreaterThan(0); - expect(seenMaxTokens).toBeLessThan(16_384); - }); - it("malformed batched response → falls back to per-step + emits warning", async () => { const llm = fakeLlm({ completeJson: { From 90c53de9d5970cfa82a90b7580b9434685c214f3 Mon Sep 17 00:00:00 2001 From: Zhe-SH-CN <1968988211@qq.com> Date: Thu, 23 Jul 2026 18:07:00 +0800 Subject: [PATCH 15/15] fix: ensure gateway restarts on install failure via ERR trap Fixes #1391 When causes the install script to exit early (e.g. npm install fails, tarball extraction fails, etc.), the OpenClaw gateway was already stopped but never restarted, leaving the service down. Add an ERR trap after in both install scripts that restarts the gateway on failure. The trap is cleared once the normal step succeeds. --- apps/memos-local-openclaw/install.sh | 6 ++++++ apps/memos-local-plugin/install.sh | 19 +++++++++++++++++++ 2 files changed, 25 insertions(+) diff --git a/apps/memos-local-openclaw/install.sh b/apps/memos-local-openclaw/install.sh index 696ff94e5..ff050d645 100644 --- a/apps/memos-local-openclaw/install.sh +++ b/apps/memos-local-openclaw/install.sh @@ -318,6 +318,10 @@ NODE info "Stop OpenClaw Gateway, 停止 OpenClaw Gateway..." "${OPENCLAW_BIN}" gateway stop >/dev/null 2>&1 || true +# Ensure the gateway is restarted if the install fails after this point. +# The trap is cleared once the normal start step at the end succeeds. +trap '"${OPENCLAW_BIN}" gateway start >/dev/null 2>&1 || true' ERR + if command -v lsof >/dev/null 2>&1; then PIDS="$(lsof -i :"${PORT}" -t 2>/dev/null || true)" if [[ -n "$PIDS" ]]; then @@ -406,6 +410,8 @@ info "Install OpenClaw Gateway service, 安装 OpenClaw Gateway 服务..." success "Start OpenClaw Gateway service, 启动 OpenClaw Gateway 服务..." "${OPENCLAW_BIN}" gateway start 2>&1 +# Gateway is back up - clear the error recovery trap. +trap - ERR info "Starting Memory Viewer, 正在启动记忆面板..." VIEWER_URL="http://127.0.0.1:18799" diff --git a/apps/memos-local-plugin/install.sh b/apps/memos-local-plugin/install.sh index 20135e3f7..40298da7d 100755 --- a/apps/memos-local-plugin/install.sh +++ b/apps/memos-local-plugin/install.sh @@ -419,6 +419,18 @@ wait_for_viewer() { return 1 } +# ─── Gateway recovery helper ────────────────────────────────────────────── +# If the install fails after stopping the OpenClaw gateway (set -e / die), +# this trap ensures the gateway is restarted so the user is not left with +# a downed service. +_gateway_recovery() { + local oc_bin + oc_bin="$(find_openclaw_cli 2>/dev/null || true)" + if [[ -n "${oc_bin}" ]]; then + "${oc_bin}" gateway start >/dev/null 2>&1 || true + fi +} + # ─── OpenClaw install ───────────────────────────────────────────────────── install_openclaw() { STEP_CURRENT=0 @@ -429,11 +441,16 @@ install_openclaw() { mkdir -p "${HOME}/.openclaw" local oc_bin="" + local _gateway_was_stopped="false" if oc_bin="$(find_openclaw_cli)"; then step "Stopping OpenClaw gateway" "${oc_bin}" gateway stop >/dev/null 2>&1 || true sleep 1 success "Gateway stopped" + _gateway_was_stopped="true" + # Ensure the gateway is restarted if the install fails after this point. + # The trap is cleared once the normal start step at the end succeeds. + trap '_gateway_recovery' ERR fi deploy_tarball_to_prefix "${prefix}" @@ -616,6 +633,8 @@ NODE else success "OpenClaw gateway started" fi + # Gateway is back up — clear the error recovery trap. + trap - ERR step "Waiting for Memory Viewer" if wait_for_viewer "${OPENCLAW_PORT}"; then