unique-user-memory 2026.32.0.dev0__tar.gz → 2026.32.0.dev2__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -1,14 +1,14 @@
1
1
  Metadata-Version: 2.3
2
2
  Name: unique-user-memory
3
- Version: 2026.32.0.dev0
3
+ Version: 2026.32.0.dev2
4
4
  Summary:
5
5
  Author: Fabian Schläpfer
6
6
  Author-email: Fabian Schläpfer <fabian@unique.ch>
7
7
  License: Proprietary
8
8
  Requires-Dist: jinja2>=3.1.6
9
9
  Requires-Dist: pydantic>=2.8.2
10
- Requires-Dist: unique-sdk>=2026.32.0.dev0,<2026.32.0rc0
11
- Requires-Dist: unique-toolkit>=2026.32.0.dev0,<2026.32.0rc0
10
+ Requires-Dist: unique-sdk>=2026.32.0.dev5,<2026.32.0rc0
11
+ Requires-Dist: unique-toolkit>=2026.32.0.dev4,<2026.32.0rc0
12
12
  Requires-Python: >=3.12, <4
13
13
  Description-Content-Type: text/markdown
14
14
 
@@ -16,7 +16,7 @@ Description-Content-Type: text/markdown
16
16
 
17
17
  Persistent per-user memory for Unique AI agents.
18
18
 
19
- `unique_user_memory` stores a compact Markdown profile for each user and updates it after every agent turn. The profile is loaded before the next turn so the assistant can remember stable user context such as communication preferences, work context, expertise, recent topics, and open follow-ups.
19
+ `unique_user_memory` stores a compact Markdown profile for each user and updates it after every agent turn. The profile is loaded before the next turn so the assistant can remember stable user context such as communication preferences, work context, expertise, recent topics, and concrete future tasks.
20
20
 
21
21
  ## What It Does
22
22
 
@@ -76,7 +76,7 @@ _(empty)_
76
76
  ## Recent Topics
77
77
  _(empty)_
78
78
 
79
- ## Open Questions / Follow-ups
79
+ ## Follow-ups
80
80
  _(empty)_
81
81
  ```
82
82
 
@@ -2,7 +2,7 @@
2
2
 
3
3
  Persistent per-user memory for Unique AI agents.
4
4
 
5
- `unique_user_memory` stores a compact Markdown profile for each user and updates it after every agent turn. The profile is loaded before the next turn so the assistant can remember stable user context such as communication preferences, work context, expertise, recent topics, and open follow-ups.
5
+ `unique_user_memory` stores a compact Markdown profile for each user and updates it after every agent turn. The profile is loaded before the next turn so the assistant can remember stable user context such as communication preferences, work context, expertise, recent topics, and concrete future tasks.
6
6
 
7
7
  ## What It Does
8
8
 
@@ -62,7 +62,7 @@ _(empty)_
62
62
  ## Recent Topics
63
63
  _(empty)_
64
64
 
65
- ## Open Questions / Follow-ups
65
+ ## Follow-ups
66
66
  _(empty)_
67
67
  ```
68
68
 
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "unique_user_memory"
3
- version = "2026.32.0.dev0"
3
+ version = "2026.32.0.dev2"
4
4
  description = ""
5
5
  authors = [
6
6
  { name = "Fabian Schläpfer", email = "fabian@unique.ch" },
@@ -11,8 +11,8 @@ requires-python = ">=3.12,<4"
11
11
  dependencies = [
12
12
  "jinja2>=3.1.6",
13
13
  "pydantic>=2.8.2",
14
- "unique-sdk>=2026.32.0.dev0,<2026.32.0rc0",
15
- "unique-toolkit>=2026.32.0.dev0,<2026.32.0rc0",
14
+ "unique-sdk>=2026.32.0.dev5,<2026.32.0rc0",
15
+ "unique-toolkit>=2026.32.0.dev4,<2026.32.0rc0",
16
16
  ]
17
17
 
18
18
  [dependency-groups]
@@ -6,6 +6,8 @@ from unique_toolkit.language_model.default_language_model import (
6
6
  DEFAULT_LANGUAGE_MODEL,
7
7
  )
8
8
  from unique_toolkit.language_model.infos import LanguageModelInfo
9
+ from unique_toolkit.language_model.invocation_stats import LanguageModelInvocationStats
10
+ from unique_toolkit.language_model.schemas import LanguageModelTokenUsage
9
11
 
10
12
  from unique_user_memory.config import UserMemoryConfig
11
13
  from unique_user_memory.user_memory import (
@@ -23,11 +25,26 @@ from unique_user_memory.user_memory import (
23
25
  upload_user_memory,
24
26
  )
25
27
  from unique_user_memory.user_memory_postprocessor import UserMemoryPostprocessor
26
- from unique_user_memory.user_memory_prompts import empty_profile
28
+ from unique_user_memory.user_memory_prompts import (
29
+ consolidation_system_prompt,
30
+ empty_profile,
31
+ memory_gate_system_prompt,
32
+ )
27
33
 
28
34
  _TEST_LANGUAGE_MODEL = LanguageModelInfo.from_name(DEFAULT_LANGUAGE_MODEL)
29
35
 
30
36
 
37
+ def test_memory_profile_keeps_follow_up_tasks_but_excludes_open_questions() -> None:
38
+ profile = empty_profile("user_1")
39
+ consolidation_prompt = consolidation_system_prompt(2000)
40
+ gate_prompt = memory_gate_system_prompt()
41
+
42
+ assert "## Open Questions / Follow-ups" not in profile
43
+ assert "## Follow-ups" in profile
44
+ assert "concrete tasks the user intends to complete" in consolidation_prompt
45
+ assert "Concrete future tasks the user intends to complete" in gate_prompt
46
+
47
+
31
48
  def test_enforce_token_cap_truncates_long_content() -> None:
32
49
  content = "\n\n".join(f"paragraph {index} " + "word " * 40 for index in range(50))
33
50
 
@@ -144,6 +161,35 @@ async def test_condense_user_memory_rejects_non_profile_output(
144
161
  assert result is None
145
162
 
146
163
 
164
+ @pytest.mark.ai
165
+ @pytest.mark.asyncio
166
+ async def test_condense_user_memory_accepts_frontmatter_output(
167
+ monkeypatch: pytest.MonkeyPatch,
168
+ ) -> None:
169
+ """Accept legacy LLM output while returning only the condensed profile body."""
170
+ condensed = "# User Memory\n\n## Identity\n- concise summary"
171
+ response = MagicMock()
172
+ response.choices[0].message.content = (
173
+ "---\nuser_id: stale-user\nturn_count: 99\n---\n\n" + condensed
174
+ )
175
+ llm_service = MagicMock()
176
+ llm_service.complete_async = AsyncMock(return_value=response)
177
+ monkeypatch.setattr(
178
+ "unique_user_memory.user_memory.LanguageModelService",
179
+ MagicMock(return_value=llm_service),
180
+ )
181
+
182
+ result = await condense_user_memory(
183
+ content="# User Memory\n\n## Identity\n- lots of stuff",
184
+ max_tokens=2000,
185
+ language_model=_TEST_LANGUAGE_MODEL,
186
+ event=MagicMock(),
187
+ logger=MagicMock(),
188
+ )
189
+
190
+ assert result == condensed
191
+
192
+
147
193
  def test_count_tokens_uses_language_model_encoder() -> None:
148
194
  language_model = MagicMock()
149
195
  language_model.get_encoder.return_value = lambda content: content.split()
@@ -281,10 +327,86 @@ async def test_consolidate_user_memory_runs_full_rewrite_when_gate_update(
281
327
  logger=MagicMock(),
282
328
  )
283
329
 
284
- assert result == rewritten
330
+ assert result.endswith(f"{rewritten}\n")
331
+ assert "user_id: user_1" in result
332
+ assert "schema_version: 1" in result
333
+ assert "turn_count: 1" in result
285
334
  llm_service.complete_async.assert_awaited_once()
286
335
 
287
336
 
337
+ @pytest.mark.asyncio
338
+ async def test_consolidate_user_memory_adds_frontmatter_to_llm_body(
339
+ monkeypatch: pytest.MonkeyPatch,
340
+ ) -> None:
341
+ rewritten = "# User Memory\n\n## Identity\n- Prefers concise answers"
342
+ response = MagicMock()
343
+ response.choices[0].message.content = rewritten
344
+ llm_service = MagicMock()
345
+ llm_service.complete_async = AsyncMock(return_value=response)
346
+ monkeypatch.setattr(
347
+ "unique_user_memory.user_memory.LanguageModelService",
348
+ MagicMock(return_value=llm_service),
349
+ )
350
+ monkeypatch.setattr(
351
+ "unique_user_memory.user_memory.should_consolidate_memory",
352
+ AsyncMock(return_value=True),
353
+ )
354
+
355
+ result = await consolidate_user_memory(
356
+ current_memory="",
357
+ user_id="authenticated-user",
358
+ user_message="remember I like concise answers",
359
+ assistant_message="noted",
360
+ config=UserMemoryConfig(),
361
+ language_model=_TEST_LANGUAGE_MODEL,
362
+ event=MagicMock(),
363
+ logger=MagicMock(),
364
+ )
365
+
366
+ assert "user_id: authenticated-user" in result
367
+ assert "schema_version: 1" in result
368
+ assert "turn_count: 1" in result
369
+
370
+
371
+ @pytest.mark.ai
372
+ @pytest.mark.asyncio
373
+ async def test_consolidate_user_memory_replaces_llm_frontmatter(
374
+ monkeypatch: pytest.MonkeyPatch,
375
+ ) -> None:
376
+ """Strip untrusted legacy metadata before assembling the updated profile."""
377
+ rewritten = "# User Memory\n\n## Identity\n- Prefers concise answers"
378
+ response = MagicMock()
379
+ response.choices[0].message.content = (
380
+ "---\nuser_id: stale-user\nturn_count: 99\n---\n\n" + rewritten
381
+ )
382
+ llm_service = MagicMock()
383
+ llm_service.complete_async = AsyncMock(return_value=response)
384
+ monkeypatch.setattr(
385
+ "unique_user_memory.user_memory.LanguageModelService",
386
+ MagicMock(return_value=llm_service),
387
+ )
388
+ monkeypatch.setattr(
389
+ "unique_user_memory.user_memory.should_consolidate_memory",
390
+ AsyncMock(return_value=True),
391
+ )
392
+
393
+ result = await consolidate_user_memory(
394
+ current_memory=empty_profile("authenticated-user"),
395
+ user_id="authenticated-user",
396
+ user_message="remember I like concise answers",
397
+ assistant_message="noted",
398
+ config=UserMemoryConfig(),
399
+ language_model=_TEST_LANGUAGE_MODEL,
400
+ event=MagicMock(),
401
+ logger=MagicMock(),
402
+ )
403
+
404
+ assert result.endswith(f"{rewritten}\n")
405
+ assert "user_id: authenticated-user" in result
406
+ assert "stale-user" not in result
407
+ assert "turn_count: 1" in result
408
+
409
+
288
410
  @pytest.mark.asyncio
289
411
  async def test_consolidate_user_memory_skips_gate_when_disabled(
290
412
  monkeypatch: pytest.MonkeyPatch,
@@ -316,7 +438,9 @@ async def test_consolidate_user_memory_skips_gate_when_disabled(
316
438
  logger=MagicMock(),
317
439
  )
318
440
 
319
- assert result == rewritten
441
+ assert result.endswith(f"{rewritten}\n")
442
+ assert "user_id: user_1" in result
443
+ assert "turn_count: 1" in result
320
444
  gate.assert_not_awaited()
321
445
  llm_service.complete_async.assert_awaited_once()
322
446
 
@@ -434,7 +558,9 @@ async def test_consolidate_user_memory_invokes_update_callbacks_on_rewrite(
434
558
  on_update_end=on_end,
435
559
  )
436
560
 
437
- assert result == rewritten
561
+ assert result.endswith(f"{rewritten}\n")
562
+ assert "user_id: user_1" in result
563
+ assert "turn_count: 1" in result
438
564
  on_start.assert_awaited_once()
439
565
  on_end.assert_awaited_once()
440
566
  assert events == ["start", "end"]
@@ -964,6 +1090,106 @@ async def test_user_memory_postprocessor_logs_success_when_upload_succeeds(
964
1090
  )
965
1091
 
966
1092
 
1093
+ @pytest.mark.ai
1094
+ @pytest.mark.asyncio
1095
+ async def test_user_memory_postprocessor_run_resets_invocation_stats(
1096
+ monkeypatch: pytest.MonkeyPatch,
1097
+ ) -> None:
1098
+ """Purpose: Verify each run reports only usage attributable to that run.
1099
+ Why this matters: Reused postprocessors must not inflate token analytics.
1100
+ Setup summary: Run twice with distinct usage and assert the second excludes the first.
1101
+ """
1102
+ load_stats = LanguageModelInvocationStats.from_usage(
1103
+ _TEST_LANGUAGE_MODEL.name,
1104
+ LanguageModelTokenUsage(total_tokens=2),
1105
+ source="user_memory_load_condense",
1106
+ )
1107
+ first_run_stats = LanguageModelInvocationStats.from_usage(
1108
+ _TEST_LANGUAGE_MODEL.name,
1109
+ LanguageModelTokenUsage(total_tokens=3),
1110
+ source="user_memory_consolidate_first",
1111
+ )
1112
+ second_run_stats = LanguageModelInvocationStats.from_usage(
1113
+ _TEST_LANGUAGE_MODEL.name,
1114
+ LanguageModelTokenUsage(total_tokens=5),
1115
+ source="user_memory_consolidate_second",
1116
+ )
1117
+ run_stats = iter((first_run_stats, second_run_stats))
1118
+
1119
+ async def consolidate(*, invocation_stats, **kwargs) -> str: # type: ignore[no-untyped-def]
1120
+ invocation_stats.append(next(run_stats))
1121
+ return "# User Memory\n\n## Identity\n- unchanged"
1122
+
1123
+ monkeypatch.setattr(
1124
+ "unique_user_memory.user_memory_postprocessor.consolidate_user_memory",
1125
+ consolidate,
1126
+ )
1127
+ event = MagicMock()
1128
+ event.user_id = "user_1"
1129
+ event.company_id = "company_1"
1130
+ event.payload.user_message.text = "remember this"
1131
+ loop_response = MagicMock()
1132
+ loop_response.message.text = "noted"
1133
+ state = UserMemoryState(
1134
+ scope_id="scope_1",
1135
+ text="# User Memory\n\n## Identity\n- unchanged",
1136
+ load_invocation_stats=(load_stats,),
1137
+ )
1138
+ postprocessor = UserMemoryPostprocessor(
1139
+ config=UserMemoryConfig(),
1140
+ language_model=_TEST_LANGUAGE_MODEL,
1141
+ event=event,
1142
+ state=state,
1143
+ logger=MagicMock(),
1144
+ chat_service=MagicMock(),
1145
+ )
1146
+
1147
+ await postprocessor.run(loop_response)
1148
+ first_reported_stats = postprocessor.invocation_stats
1149
+ await postprocessor.run(loop_response)
1150
+
1151
+ assert first_reported_stats == [load_stats, first_run_stats]
1152
+ assert postprocessor.invocation_stats == [second_run_stats]
1153
+
1154
+
1155
+ @pytest.mark.ai
1156
+ def test_user_memory_postprocessor_take_pending_invocation_stats_drains_once() -> None:
1157
+ """Purpose: Verify load-time usage is reported exactly once, however it's read.
1158
+ Why this matters: A turn that exits before `run()` (cancellation, empty
1159
+ response, a control-taking tool) must still report the load-time condense
1160
+ tokens, and a turn that does reach `run()` must not double-count them.
1161
+ Setup summary: Take the pending stats directly, then run(), and assert
1162
+ run() no longer reports the already-taken load stats.
1163
+ """
1164
+ load_stats = LanguageModelInvocationStats.from_usage(
1165
+ _TEST_LANGUAGE_MODEL.name,
1166
+ LanguageModelTokenUsage(total_tokens=2),
1167
+ source="user_memory_load_condense",
1168
+ )
1169
+ event = MagicMock()
1170
+ event.user_id = "user_1"
1171
+ event.company_id = "company_1"
1172
+ event.payload.user_message.text = "remember this"
1173
+ state = UserMemoryState(
1174
+ scope_id="scope_1",
1175
+ text="# User Memory\n\n## Identity\n- unchanged",
1176
+ load_invocation_stats=(load_stats,),
1177
+ )
1178
+ postprocessor = UserMemoryPostprocessor(
1179
+ config=UserMemoryConfig(),
1180
+ language_model=_TEST_LANGUAGE_MODEL,
1181
+ event=event,
1182
+ state=state,
1183
+ logger=MagicMock(),
1184
+ chat_service=MagicMock(),
1185
+ )
1186
+
1187
+ taken = postprocessor.take_pending_invocation_stats()
1188
+
1189
+ assert taken == [load_stats]
1190
+ assert postprocessor.take_pending_invocation_stats() == []
1191
+
1192
+
967
1193
  @pytest.mark.asyncio
968
1194
  async def test_user_memory_postprocessor_does_not_log_success_when_upload_fails(
969
1195
  monkeypatch: pytest.MonkeyPatch,
@@ -1,6 +1,7 @@
1
1
  import re
2
2
  from collections.abc import Awaitable, Callable
3
3
  from dataclasses import dataclass
4
+ from datetime import datetime, timezone
4
5
  from logging import Logger
5
6
 
6
7
  import unique_sdk
@@ -20,6 +21,9 @@ from unique_toolkit.language_model import (
20
21
  TypeEncoder,
21
22
  )
22
23
  from unique_toolkit.language_model.infos import LanguageModelInfo
24
+ from unique_toolkit.language_model.invocation_stats import (
25
+ LanguageModelInvocationStats,
26
+ )
23
27
 
24
28
  from unique_user_memory.config import UserMemoryConfig
25
29
  from unique_user_memory.user_memory_prompts import (
@@ -56,12 +60,47 @@ _CONDENSE_TARGET_RATIO = 0.9
56
60
  _TRUNCATION_MARKER = "\n\n<!-- truncated to fit memory budget -->"
57
61
  _DEFAULT_LANGUAGE_MODEL = LanguageModelInfo.from_name(DEFAULT_GPT_4o)
58
62
  _FRONTMATTER_RE = re.compile(r"^---\n.*?\n---\n", re.DOTALL)
63
+ _TURN_COUNT_RE = re.compile(r"^turn_count:\s*(\d+)\s*$", re.MULTILINE)
64
+
65
+
66
+ def _profile_body(content: str) -> str:
67
+ return _FRONTMATTER_RE.sub("", content, count=1).strip()
68
+
69
+
70
+ def _turn_count(content: str) -> int:
71
+ frontmatter_match = _FRONTMATTER_RE.match(content)
72
+ if frontmatter_match is None:
73
+ return 0
74
+ match = _TURN_COUNT_RE.search(frontmatter_match.group(0))
75
+ return int(match.group(1)) if match else 0
76
+
77
+
78
+ def _assemble_profile(*, body: str, user_id: str, turn_count: int) -> str:
79
+ """Prefix an LLM-generated body with trusted application metadata."""
80
+ timestamp = datetime.now(timezone.utc).isoformat(timespec="seconds")
81
+ return (
82
+ "---\n"
83
+ f"user_id: {user_id}\n"
84
+ "schema_version: 1\n"
85
+ f"last_updated: {timestamp}\n"
86
+ f"turn_count: {turn_count}\n"
87
+ "---\n\n"
88
+ f"{body.strip()}\n"
89
+ )
90
+
91
+
92
+ def _restore_frontmatter(original: str, body: str) -> str:
93
+ match = _FRONTMATTER_RE.match(original)
94
+ if match is None:
95
+ return body
96
+ return f"{match.group(0).rstrip()}\n\n{body.strip()}\n"
59
97
 
60
98
 
61
99
  @dataclass(frozen=True)
62
100
  class UserMemoryState:
63
101
  scope_id: str
64
102
  text: str
103
+ load_invocation_stats: tuple[LanguageModelInvocationStats, ...] = ()
65
104
 
66
105
 
67
106
  def _get_model_tokenizer(
@@ -146,6 +185,8 @@ async def condense_user_memory(
146
185
  language_model: LanguageModelInfo,
147
186
  event: ChatEvent,
148
187
  logger: Logger,
188
+ invocation_stats: list[LanguageModelInvocationStats] | None = None,
189
+ invocation_source: str = "user_memory_condense",
149
190
  ) -> str | None:
150
191
  """Ask the LLM to rewrite an oversized profile into a shorter one.
151
192
 
@@ -154,7 +195,8 @@ async def condense_user_memory(
154
195
  condensed profile, or ``None`` when the call fails or the output does
155
196
  not look like a profile (the caller then falls back to a hard cut).
156
197
  """
157
- current_tokens = count_tokens(content=content, language_model=language_model)
198
+ body = _profile_body(content)
199
+ current_tokens = count_tokens(content=body, language_model=language_model)
158
200
  target_tokens = max(1, int(max_tokens * _CONDENSE_TARGET_RATIO))
159
201
 
160
202
  try:
@@ -177,7 +219,7 @@ async def condense_user_memory(
177
219
  )
178
220
  ),
179
221
  LanguageModelUserMessage(
180
- content=condensation_user_prompt(_sanitize_for_xml_context(content))
222
+ content=condensation_user_prompt(_sanitize_for_xml_context(body))
181
223
  ),
182
224
  ]
183
225
  )
@@ -197,6 +239,15 @@ async def condense_user_memory(
197
239
  )
198
240
  return None
199
241
 
242
+ if invocation_stats is not None and response.usage is not None:
243
+ invocation_stats.append(
244
+ LanguageModelInvocationStats.from_usage(
245
+ language_model.name,
246
+ response.usage,
247
+ source=invocation_source,
248
+ )
249
+ )
250
+
200
251
  try:
201
252
  raw = response.choices[0].message.content or ""
202
253
  except Exception as exc:
@@ -214,7 +265,7 @@ async def condense_user_memory(
214
265
  )
215
266
  return None
216
267
 
217
- candidate = _strip_code_fences(raw).strip()
268
+ candidate = _profile_body(_strip_code_fences(raw))
218
269
  if not _is_well_formed_profile(candidate):
219
270
  logger.warning(
220
271
  "[user-memory] condense output did not look like a profile (%d chars)",
@@ -232,6 +283,8 @@ async def fit_user_memory(
232
283
  language_model: LanguageModelInfo,
233
284
  event: ChatEvent,
234
285
  logger: Logger,
286
+ invocation_stats: list[LanguageModelInvocationStats] | None = None,
287
+ invocation_source: str = "user_memory_condense",
235
288
  ) -> str:
236
289
  """Ensure ``content`` fits ``max_tokens``, condensing before cutting.
237
290
 
@@ -257,8 +310,11 @@ async def fit_user_memory(
257
310
  language_model=language_model,
258
311
  event=event,
259
312
  logger=logger,
313
+ invocation_stats=invocation_stats,
314
+ invocation_source=invocation_source,
260
315
  )
261
316
  if condensed is not None:
317
+ condensed = _restore_frontmatter(content, condensed)
262
318
  condensed_tokens = count_tokens(
263
319
  content=condensed, language_model=language_model
264
320
  )
@@ -321,6 +377,7 @@ async def load_user_memory(
321
377
  company_id=company_id,
322
378
  logger=logger,
323
379
  )
380
+ invocation_stats: list[LanguageModelInvocationStats] = []
324
381
  return UserMemoryState(
325
382
  scope_id=scope_id,
326
383
  text=await fit_user_memory(
@@ -329,7 +386,10 @@ async def load_user_memory(
329
386
  language_model=language_model,
330
387
  event=event,
331
388
  logger=logger,
389
+ invocation_stats=invocation_stats,
390
+ invocation_source="user_memory_load_condense",
332
391
  ),
392
+ load_invocation_stats=tuple(invocation_stats),
333
393
  )
334
394
 
335
395
 
@@ -577,6 +637,7 @@ async def should_consolidate_memory(
577
637
  language_model: LanguageModelInfo,
578
638
  event: ChatEvent,
579
639
  logger: Logger,
640
+ invocation_stats: list[LanguageModelInvocationStats] | None = None,
580
641
  ) -> bool:
581
642
  """Cheaply decide whether the turn warrants a full memory rewrite.
582
643
 
@@ -627,6 +688,15 @@ async def should_consolidate_memory(
627
688
  )
628
689
  return True
629
690
 
691
+ if invocation_stats is not None and response.usage is not None:
692
+ invocation_stats.append(
693
+ LanguageModelInvocationStats.from_usage(
694
+ language_model.name,
695
+ response.usage,
696
+ source="user_memory_gate",
697
+ )
698
+ )
699
+
630
700
  try:
631
701
  raw = response.choices[0].message.content or ""
632
702
  except Exception as exc:
@@ -665,6 +735,7 @@ async def consolidate_user_memory(
665
735
  logger: Logger,
666
736
  on_update_start: Callable[[], Awaitable[None]] = noop_update_callback,
667
737
  on_update_end: Callable[[], Awaitable[None]] = noop_update_callback,
738
+ invocation_stats: list[LanguageModelInvocationStats] | None = None,
668
739
  ) -> str:
669
740
  """Consolidate the latest turn into the user's memory profile.
670
741
 
@@ -702,6 +773,7 @@ async def consolidate_user_memory(
702
773
  language_model=language_model,
703
774
  event=event,
704
775
  logger=logger,
776
+ invocation_stats=invocation_stats,
705
777
  ):
706
778
  return safe_current
707
779
 
@@ -716,6 +788,7 @@ async def consolidate_user_memory(
716
788
  language_model=language_model,
717
789
  event=event,
718
790
  logger=logger,
791
+ invocation_stats=invocation_stats,
719
792
  )
720
793
  finally:
721
794
  await on_update_end()
@@ -731,6 +804,7 @@ async def _rewrite_user_memory(
731
804
  language_model: LanguageModelInfo,
732
805
  event: ChatEvent,
733
806
  logger: Logger,
807
+ invocation_stats: list[LanguageModelInvocationStats] | None = None,
734
808
  ) -> str:
735
809
  if not safe_current.strip():
736
810
  safe_current = empty_profile(user_id)
@@ -752,8 +826,7 @@ async def _rewrite_user_memory(
752
826
  ),
753
827
  LanguageModelUserMessage(
754
828
  content=consolidation_user_prompt(
755
- user_id=user_id,
756
- existing_memory=safe_current,
829
+ existing_memory=_profile_body(safe_current),
757
830
  user_message=_sanitize_for_xml_context(user_message or ""),
758
831
  assistant_message=_sanitize_for_xml_context(
759
832
  assistant_message or ""
@@ -780,6 +853,15 @@ async def _rewrite_user_memory(
780
853
  )
781
854
  return safe_current
782
855
 
856
+ if invocation_stats is not None and response.usage is not None:
857
+ invocation_stats.append(
858
+ LanguageModelInvocationStats.from_usage(
859
+ language_model.name,
860
+ response.usage,
861
+ source="user_memory_consolidation",
862
+ )
863
+ )
864
+
783
865
  try:
784
866
  raw = response.choices[0].message.content or ""
785
867
  except Exception as exc:
@@ -801,29 +883,32 @@ async def _rewrite_user_memory(
801
883
  logger.info("[user-memory] consolidation NOOP - keeping existing memory")
802
884
  return safe_current
803
885
 
804
- candidate = _strip_code_fences(raw).strip()
805
- if not _is_well_formed_profile(candidate):
886
+ candidate_body = _profile_body(_strip_code_fences(raw))
887
+ if not _is_well_formed_profile(candidate_body):
806
888
  logger.warning(
807
889
  "[user-memory] LLM output did not look like a profile (%d chars)",
808
- len(candidate),
890
+ len(candidate_body),
809
891
  )
810
892
  return safe_current
811
893
 
894
+ if safe_current and candidate_body == _profile_body(safe_current):
895
+ logger.debug("[user-memory] memory body unchanged - skipping update")
896
+ return safe_current
897
+
898
+ candidate = _assemble_profile(
899
+ body=candidate_body,
900
+ user_id=user_id,
901
+ turn_count=_turn_count(safe_current) + 1,
902
+ )
812
903
  capped = await fit_user_memory(
813
904
  content=candidate,
814
905
  max_tokens=config.max_tokens,
815
906
  language_model=language_model,
816
907
  event=event,
817
908
  logger=logger,
909
+ invocation_stats=invocation_stats,
910
+ invocation_source="user_memory_post_consolidation_condense",
818
911
  )
819
- if (
820
- safe_current
821
- and _FRONTMATTER_RE.sub("", capped).strip()
822
- == _FRONTMATTER_RE.sub("", safe_current).strip()
823
- ):
824
- logger.debug("[user-memory] memory body unchanged - skipping update")
825
- return safe_current
826
-
827
912
  logger.info(
828
913
  "[user-memory] consolidation produced %d tokens (cap=%d)",
829
914
  count_tokens(content=capped, language_model=language_model),
@@ -839,9 +924,7 @@ def _sanitize_for_xml_context(text: str) -> str:
839
924
  def _is_well_formed_profile(content: str) -> bool:
840
925
  if not content or len(content.strip()) < 20:
841
926
  return False
842
- if _FRONTMATTER_RE.match(content):
843
- return True
844
- return "## Identity" in content or "# User Memory" in content
927
+ return content.startswith("# User Memory") and "## Identity" in content
845
928
 
846
929
 
847
930
  def _strip_code_fences(text: str) -> str:
@@ -9,6 +9,7 @@ from unique_toolkit.language_model.default_language_model import (
9
9
  DEFAULT_LANGUAGE_MODEL,
10
10
  )
11
11
  from unique_toolkit.language_model.infos import LanguageModelInfo
12
+ from unique_toolkit.language_model.invocation_stats import LanguageModelInvocationStats
12
13
  from unique_toolkit.language_model.schemas import LanguageModelStreamResponse
13
14
 
14
15
  from unique_user_memory.config import UserMemoryConfig
@@ -49,6 +50,28 @@ class UserMemoryPostprocessor(Postprocessor):
49
50
  self._logger = logger
50
51
  self._new_memory: str | None = None
51
52
  self._chat_service: ChatService = chat_service
53
+ self._pending_load_invocation_stats = list(state.load_invocation_stats)
54
+ self._invocation_stats: list[LanguageModelInvocationStats] = []
55
+
56
+ @property
57
+ def invocation_stats(self) -> list[LanguageModelInvocationStats]:
58
+ return list(self._invocation_stats)
59
+
60
+ def take_pending_invocation_stats(self) -> list[LanguageModelInvocationStats]:
61
+ """Pop load-time condense stats not yet reported.
62
+
63
+ `UniqueAI` calls this unconditionally at the start of every turn so a
64
+ turn that exits before `run()` (cancellation, empty response, a
65
+ control-taking tool) still reports the tokens spent condensing the
66
+ loaded profile. If `run()` does execute, it drains the same pending
67
+ list itself, so whichever of the two runs first "wins" and the other
68
+ sees an empty list -- the tokens are never double-counted or lost.
69
+ """
70
+ stats, self._pending_load_invocation_stats = (
71
+ self._pending_load_invocation_stats,
72
+ [],
73
+ )
74
+ return stats
52
75
 
53
76
  async def run(self, loop_response: LanguageModelStreamResponse) -> bool:
54
77
  """Consolidate and upload user memory for this turn.
@@ -56,6 +79,7 @@ class UserMemoryPostprocessor(Postprocessor):
56
79
  Returns True if the memory profile changed and was uploaded, False
57
80
  otherwise (no user/company, NOOP consolidation, or failed upload).
58
81
  """
82
+ self._invocation_stats = self.take_pending_invocation_stats()
59
83
  self._logger.info("[user-memory] running postprocessor")
60
84
  user_id = self._event.user_id
61
85
  company_id = self._event.company_id
@@ -99,6 +123,7 @@ class UserMemoryPostprocessor(Postprocessor):
99
123
  logger=self._logger,
100
124
  on_update_start=on_update_start,
101
125
  on_update_end=on_update_end,
126
+ invocation_stats=self._invocation_stats,
102
127
  )
103
128
 
104
129
  if self._new_memory == self._state.text:
@@ -8,7 +8,7 @@ SECTION_HEADINGS: tuple[str, ...] = (
8
8
  "Work Context",
9
9
  "Skills & Expertise",
10
10
  "Recent Topics",
11
- "Open Questions / Follow-ups",
11
+ "Follow-ups",
12
12
  )
13
13
 
14
14
  _EMPTY_PROFILE_TEMPLATE = """\
@@ -52,16 +52,16 @@ small mistakes in extraction compound across every future conversation.
52
52
 
53
53
  You receive two XML blocks:
54
54
 
55
- 1. `<existing_memory>` - the current profile file (Markdown, with YAML
56
- frontmatter). May be empty on the user's first turn.
55
+ 1. `<existing_memory>` - the current Markdown profile body. May be empty
56
+ on the user's first turn.
57
57
  2. `<new_turn>` - the most recent user message and the assistant's
58
58
  reply, prefixed with `user:` and `assistant:`.
59
59
 
60
60
  # Output
61
61
 
62
- Return the complete, rewritten profile file as Markdown - frontmatter
63
- followed by the body. Do NOT emit a diff. Do NOT wrap the output in
64
- ``` fences. Do NOT add commentary before or after the file.
62
+ Return the complete, rewritten Markdown profile body, starting with
63
+ `# User Memory`. Do NOT emit a diff. Do NOT wrap the output in
64
+ ``` fences. Do NOT add commentary before or after the body.
65
65
 
66
66
  The body MUST contain exactly these section headings, in this order, even
67
67
  when a section is empty (use the literal string `_(empty)_` as a placeholder):
@@ -119,8 +119,8 @@ ADD/UPDATE for facts that are:
119
119
  tone, language, expertise level, examples preferred over theory.
120
120
  - Contextual but durable - current focus areas, active projects,
121
121
  multi-week goals, deadlines mentioned by the user.
122
- - Hand-offs - explicit "let's revisit X later", "remind me about
123
- Y", "I'll come back to Z" go into "Open Questions / Follow-ups".
122
+ - Follow-ups - concrete tasks the user intends to complete in the future,
123
+ or tasks they explicitly ask to be reminded about.
124
124
 
125
125
  NEVER extract:
126
126
 
@@ -135,11 +135,11 @@ NEVER extract:
135
135
 
136
136
  # Word budget - STRICT
137
137
 
138
- The complete file MUST be <= {{ max_words }} words (corresponding to {{ max_tokens }} tokens).
138
+ The complete body MUST be <= {{ max_words }} words (corresponding to {{ max_tokens }} tokens).
139
139
  When approaching the budget, drop content in this priority order:
140
140
 
141
141
  1. Oldest entries in Recent Topics.
142
- 2. Resolved or stale entries in Open Questions / Follow-ups.
142
+ 2. Completed, cancelled, or stale entries in Follow-ups.
143
143
  3. Fold low-signal Work Context bullets into a one-line summary.
144
144
  4. Fold low-signal Skills & Expertise bullets into broader categories.
145
145
  5. Identity and Communication Preferences - never drop, only tighten.
@@ -148,14 +148,6 @@ When approaching the budget, drop content in this priority order:
148
148
 
149
149
  The current UTC date and time is **{{ now_datetime }}**. You do NOT know the date from any other source - always use this supplied value. Never guess or infer the date.
150
150
 
151
- # Frontmatter rules
152
-
153
- - Preserve `user_id` and `schema_version` from `<existing_memory>` exactly.
154
- - Set `last_updated` to the supplied current UTC timestamp ({{ now_datetime }}).
155
- - Increment `turn_count` by 1.
156
- - If `<existing_memory>` is empty, initialize with `schema_version: 1`,
157
- `turn_count: 1`, and the user_id supplied in the user message.
158
-
159
151
  # Style
160
152
 
161
153
  - Use `-` markdown bullets, no nesting beyond two levels.
@@ -180,8 +172,8 @@ def consolidation_system_prompt(max_tokens: int) -> str:
180
172
  _CONDENSATION_SYSTEM_PROMPT_TEMPLATE = """\
181
173
  You are a memory-compaction engine for the Unique AI platform.
182
174
 
183
- You are given an existing user-memory profile (Markdown with YAML
184
- frontmatter) that is OVER its size budget. Your job is to rewrite it so
175
+ You are given an existing user-memory Markdown body that is OVER its
176
+ size budget. Your job is to rewrite it so
185
177
  it becomes materially SHORTER while preserving every durable, high-signal
186
178
  fact about the user. This is lossy compression, not deletion of meaning.
187
179
 
@@ -198,8 +190,8 @@ fact about the user. This is lossy compression, not deletion of meaning.
198
190
  overlapping information into a single clear bullet. Redundancy is the
199
191
  main reason this profile is oversized - collapse it aggressively.
200
192
  2. Delete outdated, stale, resolved, or superseded entries: old
201
- "Recent Topics", answered "Open Questions / Follow-ups", and facts a
202
- later bullet already contradicts or refines.
193
+ "Recent Topics", completed or cancelled "Follow-ups", and facts a later
194
+ bullet already contradicts or refines.
203
195
  3. Tighten verbose, flowery, or repetitive prose into short factual
204
196
  bullets. Remove hedging and filler.
205
197
  4. Fold low-signal "Work Context" and "Skills & Expertise" bullets into
@@ -211,8 +203,6 @@ fact about the user. This is lossy compression, not deletion of meaning.
211
203
  # Hard rules
212
204
 
213
205
  - NEVER invent, embellish, or add facts that are not already present.
214
- - Preserve the YAML frontmatter. Keep `user_id` and `schema_version`
215
- exactly; keep `last_updated` and `turn_count` as they are.
216
206
  - Keep exactly these section headings, in this order, even if a section
217
207
  becomes empty (use the literal string `_(empty)_`):
218
208
 
@@ -224,9 +214,9 @@ fact about the user. This is lossy compression, not deletion of meaning.
224
214
 
225
215
  # Output
226
216
 
227
- Return ONLY the complete rewritten profile file - frontmatter followed by
228
- the body. Do NOT emit a diff, do NOT wrap the output in ``` fences, and
229
- do NOT add any commentary before or after the file.
217
+ Return the complete rewritten profile body, starting with
218
+ `# User Memory`. Do NOT emit a diff, or commentary,
219
+ and do NOT wrap the output in ``` fences.
230
220
  """
231
221
 
232
222
 
@@ -254,7 +244,7 @@ _CONDENSATION_USER_PROMPT_TEMPLATE = """\
254
244
  {{ profile }}
255
245
  </profile_to_condense>
256
246
 
257
- Return the complete, condensed profile file now.
247
+ Return the complete, condensed profile body now.
258
248
  """
259
249
 
260
250
 
@@ -263,8 +253,6 @@ def condensation_user_prompt(profile: str) -> str:
263
253
 
264
254
 
265
255
  _CONSOLIDATION_USER_PROMPT_TEMPLATE = """\
266
- User ID: {{ user_id }}
267
-
268
256
  <existing_memory>
269
257
  {{ existing_memory }}
270
258
  </existing_memory>
@@ -274,19 +262,17 @@ user: {{ user_message }}
274
262
  assistant: {{ assistant_message }}
275
263
  </new_turn>
276
264
 
277
- Return the complete rewritten profile file now.
265
+ Return the complete rewritten profile body now.
278
266
  """
279
267
 
280
268
 
281
269
  def consolidation_user_prompt(
282
- user_id: str,
283
270
  existing_memory: str,
284
271
  user_message: str,
285
272
  assistant_message: str,
286
273
  ) -> str:
287
274
  existing = existing_memory.strip() or "(empty - this is the user's first turn)"
288
275
  return Template(_CONSOLIDATION_USER_PROMPT_TEMPLATE).render(
289
- user_id=user_id,
290
276
  existing_memory=existing,
291
277
  user_message=(user_message or "").strip(),
292
278
  assistant_message=(assistant_message or "").strip(),
@@ -327,7 +313,8 @@ only the single word `UPDATE` or `NOOP`.
327
313
  expertise level.
328
314
  - Durable context: current focus areas, active projects, multi-week
329
315
  goals, deadlines stated by the user.
330
- - Explicit hand-offs: "remind me about X", "let's revisit Y later".
316
+ - Concrete future tasks the user intends to complete or explicitly asks
317
+ to be reminded about.
331
318
 
332
319
  # What NEVER justifies UPDATE (lean NOOP)
333
320