feat: add AAAK expand for pre-embedding semantic quality

mempalace/dialect.py

@@ -93,6 +93,12 @@
     "surprise": "surprise",
 }
 
+# Reverse map: emotion code -> readable word (for AAAK expansion)
+_REVERSE_EMOTIONS = {}
+for _word, _code in EMOTION_CODES.items():
+    if _code not in _REVERSE_EMOTIONS:
+        _REVERSE_EMOTIONS[_code] = _word
+
 # Keywords that signal emotions in plain text
 _EMOTION_SIGNALS = {
     "decided": "determ",
@@ -932,6 +938,107 @@ def decode(self, dialect_text: str) -> dict:
 
         return result
 
+    def expand(self, dialect_text: str) -> str:
+        """Expand AAAK Dialect back into natural-language fragments for embedding.
+
+        This is NOT a lossless reconstruction — the original text is gone.
+        It reassembles the structured fields (entities, topics, key sentence,
+        emotions) into a readable string that embeds well in vector space.
+
+        Args:
+            dialect_text: AAAK-formatted string (from compress() or encode_file())
+
+        Returns:
+            A natural-language string suitable for semantic embedding.
+        """
+        parsed = self.decode(dialect_text)
+
+        parts = []
+
+        # Header: wing, room, date, file title
+        header = parsed.get("header", {})
+        for key in ("entities", "title", "date"):
+            val = header.get(key, "").strip()
+            if val and val != "?":
+                parts.append(val)
+
+        # Reverse entity code map for expansion
+        code_to_name = {}
+        for name, code in self.entity_codes.items():
+            if not name.islower() and code not in code_to_name:
+                code_to_name[code] = name
+
+        # Zettels: extract entity names, topics, key quote, emotions
+        for zettel_line in parsed.get("zettels", []):
+            fields = zettel_line.split("|")
+            for field in fields:
+                field = field.strip()
+                if not field:
+                    continue
+
+                # Entity field (starts with digit + colon, e.g. "0:ALC+BOB")
+                if field and field[0].isdigit() and ":" in field:
+                    codes = field.split(":", 1)[1]
+                    for code in codes.split("+"):
+                        code = code.strip()
+                        if code and code != "???":
+                            expanded = code_to_name.get(code, code)
+                            parts.append(expanded)
+
+                # Quoted key sentence
+                elif field.startswith('"') and field.endswith('"'):
+                    parts.append(field.strip('"'))
+
+                # Emotion codes
+                elif field in _REVERSE_EMOTIONS:
+                    parts.append(_REVERSE_EMOTIONS[field])
+
+                # Combined emotions/flags (e.g. "determ+hope", "DECISION+PIVOT")
+                # Require at least one known emotion code to avoid misclassifying
+                # topics like "TCP+UDP" as emotions.
+                elif "+" in field and any(
+                    f.strip() in _REVERSE_EMOTIONS for f in field.split("+")
+                ):
+                    for f in field.split("+"):
+                        f = f.strip()
+                        if f in _REVERSE_EMOTIONS:
+                            parts.append(_REVERSE_EMOTIONS[f])
+                        elif f.isupper():
+                            parts.append(f.lower())
+
+                # Topic keywords (underscored or plain)
+                else:
+                    parts.append(field.replace("_", " "))
+
+        # Arc line
+        arc = parsed.get("arc", "")
+        if arc:
+            for emotion in arc.split("->"):
+                emotion = emotion.strip()
+                if emotion in _REVERSE_EMOTIONS:
+                    parts.append(_REVERSE_EMOTIONS[emotion])
+
+        return " ".join(parts)
+
+    @staticmethod
+    def looks_like_aaak(text: str) -> bool:
+        """Check if text looks like AAAK Dialect format.
+
+        Heuristic: AAAK lines contain pipe separators and the first
+        field of at least one line has a colon (e.g. "0:ALC+BOB|...").
+        """
+        if "|" not in text:
+            return False
+        for line in text.strip().split("\n"):
+            line = line.strip()
+            if not line or line.startswith("ARC:") or line.startswith("T:"):
+                continue
+            if "|" in line:
+                first_field = line.split("|")[0]
+                if ":" in first_field and first_field[0].isdigit():
+                    return True
+        return False
+
     # === STATS ===
 
     @staticmethod

mempalace/mcp_server.py

@@ -508,13 +508,20 @@ def tool_diary_write(agent_name: str, entry: str, topic: str = "general"):
     )
 
     try:
-        # TODO: Future versions should expand AAAK before embedding to improve
-        # semantic search quality. For now, store raw AAAK in metadata so it's
-        # preserved, and keep the document as-is for embedding (even though
-        # compressed AAAK degrades embedding quality).
+        # If entry is AAAK-compressed, expand it for better embedding quality
+        # while preserving the original compressed form in metadata.
+        from mempalace.dialect import Dialect
+
+        embed_text = entry
+        meta_extra = {}
+        if Dialect.looks_like_aaak(entry):
+            _expand_dialect = Dialect()
+            embed_text = _expand_dialect.expand(entry)
+            meta_extra["aaak_compressed"] = entry
+
         col.add(
             ids=[entry_id],
-            documents=[entry],
+            documents=[embed_text],
             metadatas=[
                 {
                     "wing": wing,
@@ -525,6 +532,7 @@ def tool_diary_write(agent_name: str, entry: str, topic: str = "general"):
                     "agent": agent_name,
                     "filed_at": now.isoformat(),
                     "date": now.strftime("%Y-%m-%d"),
+                    **meta_extra,
                 }
             ],
         )

tests/test_dialect.py

@@ -155,3 +155,94 @@ def test_decode_roundtrip(self):
         assert decoded["header"]["file"] == "001"
         assert decoded["arc"] == "journey"
         assert len(decoded["zettels"]) == 1
+
+
+class TestExpand:
+    def test_expand_basic(self):
+        d = Dialect(entities={"Alice": "ALC", "Bob": "BOB"})
+        aaak = '001|ALC+BOB|2025-01-01|test_title\n001:ALC+BOB|memory_ai|"we decided to use GraphQL"|determ'
+        expanded = d.expand(aaak)
+        assert "Alice" in expanded or "ALC" in expanded
+        assert "memory" in expanded or "ai" in expanded
+        assert "we decided to use GraphQL" in expanded
+
+    def test_expand_preserves_entities(self):
+        d = Dialect(entities={"Alice": "ALC"})
+        aaak = '001|ALC|2025-01-01|meeting\n001:ALC|project_auth|"switched to JWT"|determ+convict'
+        expanded = d.expand(aaak)
+        assert "Alice" in expanded
+
+    def test_expand_preserves_key_sentence(self):
+        d = Dialect()
+        aaak = 'project|backend|2025-01-01|auth\n0:???|jwt_tokens|"Tokens expire after 24 hours"'
+        expanded = d.expand(aaak)
+        assert "Tokens expire after 24 hours" in expanded
+
+    def test_expand_handles_emotions(self):
+        d = Dialect()
+        aaak = '001|ALC|2025-01-01|talk\n001:ALC|life|"a turning point"|hope+joy'
+        expanded = d.expand(aaak)
+        assert "hope" in expanded
+        assert "joy" in expanded
+
+    def test_expand_empty_input(self):
+        d = Dialect()
+        expanded = d.expand("")
+        assert isinstance(expanded, str)
+
+    def test_expand_uppercase_topic_not_treated_as_emotion(self):
+        """Topics like TCP+UDP should not be misclassified as combined emotions."""
+        d = Dialect()
+        aaak = '001|SRV|2025-01-01|networking\n001:SRV|TCP+UDP|"packet routing"'
+        expanded = d.expand(aaak)
+        # TCP+UDP should be treated as topic, not emotions
+        assert "tcp" in expanded.lower() or "udp" in expanded.lower()
+
+    def test_expand_protocol_names_with_plus(self):
+        """Common tech protocol names with + should not be misclassified.
+
+        Regression test suggested by @web3guru888 — protocol names like
+        HTTPS+REST and TCP+UDP are common in tech docs and must not be
+        treated as combined emotion codes.
+        """
+        d = Dialect()
+        for protocol in ("HTTPS+REST", "TCP+UDP", "HTTP+JSON", "GRPC+PROTOBUF"):
+            aaak = f'001|API|2025-01-01|stack\n001:API|{protocol}|"design notes"'
+            expanded = d.expand(aaak)
+            # Each protocol part should appear in the expanded output
+            for part in protocol.split("+"):
+                assert part.lower() in expanded.lower(), (
+                    f"Protocol part {part!r} missing from expansion of {protocol!r}: {expanded!r}"
+                )
+
+    def test_expand_roundtrip_from_compress(self):
+        """Compress text then expand — key terms should survive."""
+        d = Dialect(entities={"Alice": "ALC"})
+        original = "Alice decided to switch from REST to GraphQL for the API layer."
+        compressed = d.compress(original)
+        expanded = d.expand(compressed)
+        # The key concept should survive the round-trip
+        assert "graphql" in expanded.lower() or "api" in expanded.lower()
+
+
+class TestLooksLikeAaak:
+    def test_positive_plain_compress(self):
+        d = Dialect()
+        compressed = d.compress(
+            "We decided to use PostgreSQL for the database.",
+            metadata={"wing": "project", "room": "backend"},
+        )
+        assert d.looks_like_aaak(compressed)
+
+    def test_positive_zettel_format(self):
+        aaak = '001|ALC|2025-01-01|title\n001:ALC|topic|"quote"|0.9|joy'
+        assert Dialect.looks_like_aaak(aaak)
+
+    def test_negative_plain_text(self):
+        assert not Dialect.looks_like_aaak("This is just plain English text.")
+
+    def test_negative_empty(self):
+        assert not Dialect.looks_like_aaak("")
+
+    def test_negative_pipe_but_no_colon_digit(self):
+        assert not Dialect.looks_like_aaak("wing|room|date|file")