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 () => { diff --git a/src/memos/api/client.py b/src/memos/api/client.py index 818ce5e0d..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 @@ -65,17 +67,43 @@ 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, 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 = { @@ -102,22 +130,23 @@ def get_message( def add_message( self, messages: list[dict[str, Any]], - user_id: 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, allow_knowledgebase_ids: list[str] | None = None, + allow_memory_view: list[str] | None = None, ) -> 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 = { @@ -130,8 +159,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: @@ -150,8 +180,9 @@ def add_message( def search_memory( self, query: str, - user_id: str, - conversation_id: str, + user_id: str | None = None, + 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,22 +191,34 @@ 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 - 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 = { "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 +238,30 @@ 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) + if size > 50: + raise ValueError("size must be less than or equal to 50") 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, } @@ -223,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 = { @@ -313,7 +400,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 +408,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 +422,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 +493,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): @@ -446,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, @@ -456,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 = { @@ -486,17 +597,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 +649,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,21 +766,26 @@ 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, + 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, - ) -> MemOSChatResponse | None: + stream: bool = False, + include_tool_memory: bool = False, + tool_memory_limit_number: int = 6, + relativity: float | None = None, + ) -> MemOSChatResponse | Iterator[str] | None: """chat""" # Validate required parameters self._validate_required_params( @@ -565,20 +812,31 @@ 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): 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 2db6d0f75..a2c83f622 100644 --- a/src/memos/api/product_models.py +++ b/src/memos/api/product_models.py @@ -1055,6 +1055,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" ) @@ -1066,6 +1075,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): @@ -1077,6 +1089,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): @@ -1105,6 +1133,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) ────────────────── @@ -1188,12 +1226,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/src/memos/multi_mem_cube/single_cube.py b/src/memos/multi_mem_cube/single_cube.py index f84fc60e1..b24823427 100644 --- a/src/memos/multi_mem_cube/single_cube.py +++ b/src/memos/multi_mem_cube/single_cube.py @@ -1,6 +1,7 @@ from __future__ import annotations import json +import logging import time import traceback @@ -37,6 +38,69 @@ logger = get_logger(__name__) +# Memory partitions in `MOSSearchResult` that may carry `metadata.embedding` +# lists we want to redact from INFO logs. Kept as a module constant so a new +# partition (e.g. `notif_mem`) can be added by one edit. +_MEM_PARTITIONS_WITH_EMBEDDING = ("text_mem", "pref_mem", "tool_mem", "skill_mem") + + +def _redact_embeddings_for_log(memories_result: dict[str, Any]) -> dict[str, Any]: + """Return a log-safe shallow copy of ``memories_result`` with each + ``metadata.embedding`` list replaced by a length placeholder. + + Motivation (GitHub issue #2103): + ``format_memory_item(..., include_embedding=True)`` keeps the full + embedding vector on each memory when ``dedup`` is ``mmr`` or + ``sim``. ``APISearchRequest.dedup`` defaults to ``mmr``, so the + raw ``logger.info(memories_result)`` call would dump every high + dimensional vector into INFO logs on every search — MBs of noise + plus a potential privacy / compliance concern. + + The returned structure is a shallow view suitable for interpolation + into a log message; the caller's live ``memories_result`` object is + NOT mutated (business logic must still see the real embedding for + dedup / downstream consumers). + """ + if not isinstance(memories_result, dict): + return memories_result + + redacted: dict[str, Any] = dict(memories_result) + for partition in _MEM_PARTITIONS_WITH_EMBEDDING: + groups = memories_result.get(partition) + if not groups: + continue + new_groups = [] + for group in groups: + if not isinstance(group, dict): + new_groups.append(group) + continue + new_group = dict(group) + new_memories = [] + # `or []` intentionally handles the `{"memories": None}` case — + # `dict.get("memories")` returns `None` (not `[]`) when the value + # is explicitly None, and downstream `format_memory_item` / + # post-processing may leave that key null on empty partitions. + for mem in group.get("memories") or []: + if not isinstance(mem, dict): + new_memories.append(mem) + continue + metadata = mem.get("metadata") + if isinstance(metadata, dict) and "embedding" in metadata: + embedding = metadata.get("embedding") + length = len(embedding) if hasattr(embedding, "__len__") else 0 + new_metadata = dict(metadata) + new_metadata["embedding"] = f"" + new_mem = dict(mem) + new_mem["metadata"] = new_metadata + new_memories.append(new_mem) + else: + new_memories.append(mem) + new_group["memories"] = new_memories + new_groups.append(new_group) + redacted[partition] = new_groups + return redacted + + if TYPE_CHECKING: from memos.api.product_models import APIADDRequest, APIFeedbackRequest, APISearchRequest from memos.mem_cube.navie import NaiveMemCube @@ -129,7 +193,10 @@ def search_memories(self, search_req: APISearchRequest) -> dict[str, Any]: self.cube_id, ) - self.logger.info(f"Search memories result: {memories_result}") + if self.logger.isEnabledFor(logging.INFO): + self.logger.info( + "Search memories result: %s", _redact_embeddings_for_log(memories_result) + ) self.logger.info(f"Search {len(memories_result)} memories.") return memories_result diff --git a/tests/api/test_client.py b/tests/api/test_client.py new file mode 100644 index 000000000..2e911e4f4 --- /dev/null +++ b/tests/api/test_client.py @@ -0,0 +1,686 @@ +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 + + +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, + "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("/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"): + 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 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") + + +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=None, + 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", + ) + + +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 diff --git a/tests/multi_mem_cube/test_search_log_redaction.py b/tests/multi_mem_cube/test_search_log_redaction.py new file mode 100644 index 000000000..8e423ea23 --- /dev/null +++ b/tests/multi_mem_cube/test_search_log_redaction.py @@ -0,0 +1,335 @@ +"""Regression tests for GitHub issue #2103. + +`SingleCubeView.search_memories` used to `logger.info(f"Search memories +result: {memories_result}")`, which — when `dedup` is `mmr` or `sim` +(the default is `mmr`) — leaks full embedding vectors into INFO logs. + +These tests verify: + 1. The pure redaction helper `_redact_embeddings_for_log` replaces + `metadata.embedding` lists with a `` placeholder, + across every memory partition (text_mem / pref_mem / tool_mem / + skill_mem), while preserving every other field. + 2. `search_memories` never prints the raw embedding values into + the INFO log, even when `dedup=mmr` (the default). + +NOTE: heavy mocking mirrors `tests/test_add_stage_logging.py`; the +`SingleCubeView` is imported lazily to work around a circular import in +`memos.api.handlers.__init__`. +""" + +from __future__ import annotations + +import logging + +from typing import Any +from unittest.mock import MagicMock + +import pytest + +# Prime the `memos.api.handlers` package before any test tries to import +# `memos.multi_mem_cube.single_cube` — this side-steps the known circular +# import path (`api.handlers.__init__` → `add_handler` → `single_cube`) +# by ensuring the handlers package finishes initializing before the +# single_cube module is pulled in on-demand from a test fixture. +import memos.api.handlers # noqa: F401 + + +# --------------------------------------------------------------------------- +# _redact_embeddings_for_log — pure helper unit tests +# --------------------------------------------------------------------------- + + +class TestRedactEmbeddingsForLog: + def test_replaces_embedding_in_text_mem(self): + from memos.multi_mem_cube.single_cube import _redact_embeddings_for_log + + result = { + "text_mem": [ + { + "cube_id": "c1", + "memories": [ + { + "id": "m1", + "memory": "hello", + "metadata": { + "embedding": [0.1, 0.2, 0.3, 0.4], + "relativity": 0.9, + }, + } + ], + "total_nodes": 1, + } + ], + "act_mem": [], + "para_mem": [], + "pref_mem": [], + "tool_mem": [], + "skill_mem": [], + "pref_note": "", + } + + redacted = _redact_embeddings_for_log(result) + + emb = redacted["text_mem"][0]["memories"][0]["metadata"]["embedding"] + assert emb == "" + # Non-embedding fields preserved verbatim + assert redacted["text_mem"][0]["memories"][0]["memory"] == "hello" + assert redacted["text_mem"][0]["memories"][0]["metadata"]["relativity"] == 0.9 + assert redacted["text_mem"][0]["total_nodes"] == 1 + assert redacted["text_mem"][0]["cube_id"] == "c1" + + def test_replaces_embedding_across_all_partitions(self): + from memos.multi_mem_cube.single_cube import _redact_embeddings_for_log + + def _mk(vec: list[float]) -> dict[str, Any]: + return { + "id": "x", + "memory": "t", + "metadata": {"embedding": vec, "memory_type": "WorkingMemory"}, + } + + result: dict[str, Any] = { + "text_mem": [{"cube_id": "c", "memories": [_mk([1.0, 2.0])]}], + "pref_mem": [{"cube_id": "c", "memories": [_mk([1.0, 2.0, 3.0])]}], + "tool_mem": [{"cube_id": "c", "memories": [_mk([9.9])]}], + "skill_mem": [{"cube_id": "c", "memories": [_mk([0.1, 0.2, 0.3, 0.4, 0.5])]}], + "act_mem": [], + "para_mem": [], + "pref_note": "", + } + + redacted = _redact_embeddings_for_log(result) + + assert ( + redacted["text_mem"][0]["memories"][0]["metadata"]["embedding"] == "" + ) + assert ( + redacted["pref_mem"][0]["memories"][0]["metadata"]["embedding"] == "" + ) + assert ( + redacted["tool_mem"][0]["memories"][0]["metadata"]["embedding"] == "" + ) + assert ( + redacted["skill_mem"][0]["memories"][0]["metadata"]["embedding"] == "" + ) + + def test_empty_embedding_reports_zero(self): + from memos.multi_mem_cube.single_cube import _redact_embeddings_for_log + + result = { + "text_mem": [ + { + "cube_id": "c", + "memories": [ + {"id": "m", "memory": "t", "metadata": {"embedding": []}}, + ], + } + ], + "pref_mem": [], + "tool_mem": [], + "skill_mem": [], + "act_mem": [], + "para_mem": [], + "pref_note": "", + } + + redacted = _redact_embeddings_for_log(result) + assert ( + redacted["text_mem"][0]["memories"][0]["metadata"]["embedding"] == "" + ) + + def test_missing_metadata_or_embedding_is_untouched(self): + from memos.multi_mem_cube.single_cube import _redact_embeddings_for_log + + result = { + "text_mem": [ + { + "cube_id": "c", + "memories": [ + {"id": "m1", "memory": "no-meta"}, + {"id": "m2", "memory": "meta-no-emb", "metadata": {"relativity": 0.5}}, + ], + } + ], + "pref_mem": [], + "tool_mem": [], + "skill_mem": [], + "act_mem": [], + "para_mem": [], + "pref_note": "", + } + + redacted = _redact_embeddings_for_log(result) + # First item: no metadata → untouched + assert "metadata" not in redacted["text_mem"][0]["memories"][0] + # Second item: has metadata but no embedding → relativity preserved, + # no accidental embedding key added + assert redacted["text_mem"][0]["memories"][1]["metadata"] == {"relativity": 0.5} + + def test_does_not_mutate_input(self): + from memos.multi_mem_cube.single_cube import _redact_embeddings_for_log + + original_vec = [0.1, 0.2, 0.3] + result = { + "text_mem": [ + { + "cube_id": "c", + "memories": [ + {"id": "m", "memory": "t", "metadata": {"embedding": original_vec}}, + ], + } + ], + "pref_mem": [], + "tool_mem": [], + "skill_mem": [], + "act_mem": [], + "para_mem": [], + "pref_note": "", + } + + _redact_embeddings_for_log(result) + + # Caller's live dict must still hold the original vector + assert result["text_mem"][0]["memories"][0]["metadata"]["embedding"] == original_vec + assert result["text_mem"][0]["memories"][0]["metadata"]["embedding"] is original_vec + + def test_handles_missing_partition_keys(self): + """`memories_result` in `search_memories` is fully populated, but + we should still tolerate partial dicts defensively. + """ + from memos.multi_mem_cube.single_cube import _redact_embeddings_for_log + + result = { + "text_mem": [ + { + "cube_id": "c", + "memories": [ + {"id": "m", "memory": "t", "metadata": {"embedding": [1.0, 2.0]}}, + ], + } + ], + # Every other partition omitted + } + + redacted = _redact_embeddings_for_log(result) + assert ( + redacted["text_mem"][0]["memories"][0]["metadata"]["embedding"] == "" + ) + + +# --------------------------------------------------------------------------- +# Integration — SingleCubeView.search_memories INFO log must not leak +# --------------------------------------------------------------------------- + + +@pytest.fixture() +def search_view_with_embedding(): + """Build a SingleCubeView whose _search_text returns memories with + populated embedding vectors (matching the `dedup=mmr` code path). + """ + from memos.multi_mem_cube.single_cube import SingleCubeView + + view = SingleCubeView( + cube_id="cube_test", + naive_mem_cube=MagicMock(), + mem_reader=MagicMock(), + mem_scheduler=MagicMock(), + logger=logging.getLogger("test.search_log_redaction"), + searcher=MagicMock(), + ) + + fake_embedding = [0.12345, -0.6789, 0.999888, -0.111222] + + def _fake_search_text(_req, _ctx, _mode): + return [ + { + "id": "mem-1", + "memory": "user likes dark mode", + "ref_id": "[mem]", + "metadata": { + "memory_type": "WorkingMemory", + "embedding": fake_embedding, + "relativity": 0.87, + "ref_id": "[mem]", + "id": "mem-1", + "memory": "user likes dark mode", + "usage": [], + }, + } + ] + + view._search_text = _fake_search_text # type: ignore[method-assign] + return view, fake_embedding + + +def _make_search_req(**overrides): + from memos.api.product_models import APISearchRequest + + defaults = {"user_id": "u1", "query": "test query"} + defaults.update(overrides) + return APISearchRequest(**defaults) + + +class TestSearchMemoriesLoggingRedaction: + def test_default_dedup_mmr_does_not_leak_embedding(self, search_view_with_embedding, caplog): + """Default `APISearchRequest.dedup = "mmr"` → embedding stays in + the returned result, but MUST NOT appear in the INFO log.""" + view, fake_embedding = search_view_with_embedding + req = _make_search_req() # defaults: mode=fast, dedup=mmr + + with caplog.at_level(logging.INFO, logger="test.search_log_redaction"): + result = view.search_memories(req) + + # Sanity: business result still contains the embedding untouched + assert result["text_mem"][0]["memories"][0]["metadata"]["embedding"] == fake_embedding + + # The 'Search memories result:' INFO line must NOT include any of + # the embedding float digits. + search_log_lines = [ + r.message for r in caplog.records if "Search memories result" in r.message + ] + assert len(search_log_lines) == 1, ( + f"expected exactly one summary log line, got: {search_log_lines}" + ) + summary = search_log_lines[0] + + for val in fake_embedding: + # Each float, formatted straight into a Python list repr, + # would embed its digits into the log. None must appear. + assert str(val) not in summary, ( + f"embedding value {val!r} leaked into INFO log: {summary}" + ) + + # And the redaction placeholder should be present + assert "" in summary, ( + f"expected redacted placeholder in log, got: {summary}" + ) + + def test_summary_still_includes_useful_fields(self, search_view_with_embedding, caplog): + view, _ = search_view_with_embedding + req = _make_search_req() + + with caplog.at_level(logging.INFO, logger="test.search_log_redaction"): + view.search_memories(req) + + matched = [r.message for r in caplog.records if "Search memories result" in r.message] + assert len(matched) == 1, f"expected exactly one summary log line, got: {matched}" + summary = matched[0] + # Useful fields still logged — this is a log hygiene fix, not a log removal. + assert "mem-1" in summary + assert "user likes dark mode" in summary + assert "0.87" in summary # relativity preserved + + def test_count_log_still_emitted(self, search_view_with_embedding, caplog): + view, _ = search_view_with_embedding + req = _make_search_req() + + with caplog.at_level(logging.INFO, logger="test.search_log_redaction"): + view.search_memories(req) + + count_lines = [ + r.message + for r in caplog.records + if r.message.startswith("Search ") and "memories." in r.message + ] + assert count_lines, "expected 'Search N memories.' count log to remain"