From 46564ae2831470b74244777ddac99b2d56de6adc Mon Sep 17 00:00:00 2001 From: zhiar <172391900+epinethrone@users.noreply.github.com> Date: Tue, 18 Aug 2026 16:43:16 +0200 Subject: [PATCH] Group chunked memories in dashboard --- mempalace_dashboard/server.py | 122 ++++++++++++++++++++++++++++------ tests/test_chunk_identity.py | 113 +++++++++++++++++++++++++++++++ 2 files changed, 216 insertions(+), 19 deletions(-) create mode 100644 tests/test_chunk_identity.py diff --git a/mempalace_dashboard/server.py b/mempalace_dashboard/server.py index 4ec072e..bdf846b 100644 --- a/mempalace_dashboard/server.py +++ b/mempalace_dashboard/server.py @@ -794,6 +794,65 @@ def _row_dicts(cursor: sqlite3.Cursor) -> list[dict]: return [dict(row) for row in cursor.fetchall()] +_CHUNKED_DRAWER_ID_RE = re.compile(r"^(?P.+)_chunk_(?P\d{6})$") + + +def logical_drawer_id(drawer_id: str) -> tuple[str, int | None]: + """Return the public drawer id and chunk ordinal for a Chroma row. + + MemPalace stores oversized drawers as ``_chunk_000000`` + records. They are one drawer from the user's perspective: exposing the + physical IDs makes a delete or notification dismissal affect only one + fragment. IDs that do not use the exact six-digit suffix remain normal + one-row drawers. + """ + match = _CHUNKED_DRAWER_ID_RE.match(drawer_id) + if not match: + return drawer_id, None + return match.group("parent"), int(match.group("ordinal")) + + +def _like_prefix(value: str) -> str: + """Escape a SQLite LIKE prefix while retaining its final wildcard.""" + return value.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_") + "%" + + +def _group_physical_drawers(by_id: dict[int, dict], *, light: bool) -> list[dict]: + """Collapse Chroma chunk rows into their parent drawer representation.""" + groups: dict[str, list[tuple[int, int, dict]]] = {} + for position, item in enumerate(by_id.values()): + parent_id, ordinal = logical_drawer_id(item["drawer_id"]) + # Normal rows retain their existing ordering. Chunk ordering comes + # from the explicit ordinal rather than Chroma's internal numeric id. + sort_key = ordinal if ordinal is not None else -1 + groups.setdefault(parent_id, []).append((position, sort_key, item)) + + drawers: list[dict] = [] + for parent_id, members in groups.items(): + members.sort(key=lambda member: (member[1], member[0])) + first = members[0][2] + if len(members) == 1 and members[0][1] == -1: + drawers.append(first) + continue + + # Metadata is duplicated on every physical Chroma chunk. Take it + # from the first logical chunk, while concatenating chunk documents + # exactly in their authored ordinal order. + item = dict(first) + item["drawer_id"] = parent_id + item["content"] = "".join(member[2].get("content", "") for member in members) + if light: + item["truncated"] = len(item["content"]) > PALACE_PREVIEW_CHARS + item["content"] = item["content"][:PALACE_PREVIEW_CHARS] + else: + item["etag"] = content_etag(item["content"]) + # Light-mode titles must be derived from the same preview body the + # client receives; full-mode titles keep the complete body behaviour. + item["title"] = extract_title(item["content"], parent_id) + drawers.append(item) + return drawers + + def count_drawers() -> int: """Cheap COUNT(*) of drawers — used to decide light vs full payload without materializing every body. Returns 0 if the DB is missing or @@ -838,8 +897,34 @@ def read_drawers(light: bool = False, ids: list[int] | None = None) -> list[dict id_filter = "" params: tuple = () if ids is not None: - id_filter = f" and e.id in ({','.join('?' for _ in ids)})" - params = tuple(ids) + # A full-text hit on one chunk must materialize the whole logical + # drawer, otherwise lazy-open/search would show only its matching + # fragment. The numeric id branch preserves normal-row behaviour. + parent_ids: set[str] = set() + try: + lookup = sqlite3.connect(PALACE_DB) + try: + placeholders = ",".join("?" for _ in ids) + selected = lookup.execute( + f"select embedding_id from embeddings where id in ({placeholders})", + tuple(ids), + ).fetchall() + finally: + lookup.close() + parent_ids = { + parent for (embedding_id,) in selected + for parent, ordinal in [logical_drawer_id(str(embedding_id))] + if ordinal is not None + } + except sqlite3.OperationalError: + return [] + clauses = [f"e.id in ({','.join('?' for _ in ids)})"] + query_params: list[object] = list(ids) + if parent_ids: + clauses.extend("e.embedding_id like ? escape '\\'" for _ in parent_ids) + query_params.extend(_like_prefix(f"{parent}_chunk_") for parent in parent_ids) + id_filter = " and (" + " or ".join(clauses) + ")" + params = tuple(query_params) try: con = sqlite3.connect(PALACE_DB) con.row_factory = sqlite3.Row @@ -907,7 +992,7 @@ def read_drawers(light: bool = False, ids: list[int] | None = None) -> list[dict item.setdefault("truncated", item.get("truncated", False)) else: item["etag"] = content_etag(item.get("content", "")) - drawers = list(by_id.values()) + drawers = _group_physical_drawers(by_id, light=light) enrich_drawers_with_updated_at(drawers) return drawers @@ -929,39 +1014,38 @@ def read_single_drawer(drawer_id: str) -> dict | None: from embeddings e join embedding_metadata em on em.id = e.id where e.embedding_id = ? + or e.embedding_id like ? escape '\\' """, - (drawer_id,), + (drawer_id, _like_prefix(f"{drawer_id}_chunk_")), ) ) con.close() except sqlite3.OperationalError: return None + rows = [ + row for row in rows + if logical_drawer_id(str(row["embedding_id"]))[0] == drawer_id + ] if not rows: return None - item = { - "id": rows[0]["id"], - "drawer_id": drawer_id, - "wing": "unknown", - "room": "unknown", - "title": "Untitled", - "content": "", - "source_file": "", - "filed_at": "", - "added_by": "", - "truncated": False, - "metadata": {}, - } + by_id: dict[int, dict] = {} for row in rows: + item = by_id.setdefault(row["id"], { + "id": row["id"], "drawer_id": row["embedding_id"], + "wing": "unknown", "room": "unknown", "title": "Untitled", + "content": "", "source_file": "", "filed_at": "", + "added_by": "", "truncated": False, "metadata": {}, + }) key = row["key"] value = row["value"] if key == "chroma:document": item["content"] = value or "" - item["title"] = extract_title(item["content"], item["drawer_id"]) elif key in item: item[key] = value or "" else: item["metadata"][key] = value - item["etag"] = content_etag(item.get("content", "")) + item = _group_physical_drawers(by_id, light=False)[0] + item["truncated"] = False enrich_drawers_with_updated_at([item]) return item diff --git a/tests/test_chunk_identity.py b/tests/test_chunk_identity.py new file mode 100644 index 0000000..a2b6209 --- /dev/null +++ b/tests/test_chunk_identity.py @@ -0,0 +1,113 @@ +from __future__ import annotations + +import sqlite3 +import tempfile +import unittest +from contextlib import ExitStack +from pathlib import Path +from unittest import mock + +from mempalace_dashboard import server + + +class ChunkIdentityTests(unittest.TestCase): + """Keep Chroma's physical chunk representation out of dashboard state.""" + + def setUp(self) -> None: + self._tmp = tempfile.TemporaryDirectory() + self.addCleanup(self._tmp.cleanup) + self._patches = ExitStack() + self.addCleanup(self._patches.close) + self.db = Path(self._tmp.name) / "chroma.sqlite3" + self._patches.enter_context(mock.patch.object(server, "PALACE_DB", self.db)) + self._patches.enter_context(mock.patch.object(server, "PALACE_PREVIEW_CHARS", 8)) + self._patches.enter_context(mock.patch.object(server, "enrich_drawers_with_updated_at")) + self._create_db() + + def _create_db(self) -> None: + con = sqlite3.connect(self.db) + con.executescript(""" + create table embeddings (id integer primary key, embedding_id text not null); + create table embedding_metadata ( + id integer not null, + key text not null, + string_value text, + int_value integer, + float_value real, + bool_value integer + ); + """) + # Deliberately store chunk 1 before chunk 0: dashboard order must use + # the chunk suffix, not Chroma's internal row id. + con.executemany("insert into embeddings values (?, ?)", [ + (10, "drawer_story_chunk_000001"), + (11, "drawer_story_chunk_000000"), + (12, "drawer_normal"), + ]) + metadata = [ + (10, "chroma:document", "world", None, None, None), + (10, "wing", "archive", None, None, None), + (10, "room", "stories", None, None, None), + (11, "chroma:document", "# Hello\n", None, None, None), + (11, "wing", "archive", None, None, None), + (11, "room", "stories", None, None, None), + (11, "source_file", "import.md", None, None, None), + (12, "chroma:document", "# Normal\n\nbody", None, None, None), + (12, "wing", "notes", None, None, None), + (12, "room", "general", None, None, None), + ] + con.executemany("insert into embedding_metadata values (?, ?, ?, ?, ?, ?)", metadata) + con.commit() + con.close() + + def test_list_groups_chunks_as_one_logical_drawer(self) -> None: + drawers = server.read_drawers() + + self.assertEqual([drawer["drawer_id"] for drawer in drawers], ["drawer_story", "drawer_normal"]) + story = drawers[0] + self.assertEqual(story["content"], "# Hello\nworld") + self.assertEqual(story["title"], "Hello") + self.assertEqual(story["wing"], "archive") + self.assertEqual(story["source_file"], "import.md") + self.assertEqual(story["etag"], server.content_etag("# Hello\nworld")) + + def test_lazy_read_search_and_delete_scope_keep_parent_identity(self) -> None: + # A hit on the second physical row must return both chunks as the + # logical parent, otherwise opening the search result loses content. + matched = server.read_drawers(light=True, ids=[10]) + self.assertEqual(len(matched), 1) + self.assertEqual(matched[0]["drawer_id"], "drawer_story") + self.assertEqual(matched[0]["content"], "# Hello\n") + self.assertTrue(matched[0]["truncated"]) + + drawer = server.read_single_drawer("drawer_story") + self.assertIsNotNone(drawer) + self.assertEqual(drawer["drawer_id"], "drawer_story") + self.assertEqual(drawer["content"], "# Hello\nworld") + self.assertFalse(drawer["truncated"]) + + label, targets = server.drawers_for_delete({"scope": "drawer", "drawer_id": "drawer_story"}) + self.assertEqual(label, "memory") + self.assertEqual([target["drawer_id"] for target in targets], ["drawer_story"]) + + with ( + mock.patch.object(server, "mempalace_update_drawer", return_value={"success": True}) as update, + mock.patch.object(server, "_mark_drawer_self_seen"), + ): + server.update_memory({ + "drawer_id": "drawer_story", + "content": "# Hello\nupdated world", + "etag": drawer["etag"], + }) + update.assert_called_once_with("drawer_story", "# Hello\nupdated world", None, None) + + with ( + mock.patch.object(server, "mempalace_delete_drawer", return_value={"success": True}) as delete, + mock.patch.object(server, "log_version"), + ): + server.delete_memories({"scope": "drawer", "drawer_id": "drawer_story", "confirm": "DELETE"}) + delete.assert_called_once_with("drawer_story") + + +if __name__ == "__main__": + unittest.main()