unique-user-memory 2026.30.0.dev0__tar.gz → 2026.30.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.30.0.dev0
3
+ Version: 2026.30.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.30.0.dev0,<2026.30.0rc0
11
- Requires-Dist: unique-toolkit>=2026.30.0.dev0,<2026.30.0rc0
10
+ Requires-Dist: unique-sdk>=2026.30.0.dev9,<2026.30.0rc0
11
+ Requires-Dist: unique-toolkit>=2026.30.0.dev22,<2026.30.0rc0
12
12
  Requires-Python: >=3.12, <4
13
13
  Description-Content-Type: text/markdown
14
14
 
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "unique_user_memory"
3
- version = "2026.30.0.dev0"
3
+ version = "2026.30.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.30.0.dev0,<2026.30.0rc0",
15
- "unique-toolkit>=2026.30.0.dev0,<2026.30.0rc0",
14
+ "unique-sdk>=2026.30.0.dev9,<2026.30.0rc0",
15
+ "unique-toolkit>=2026.30.0.dev22,<2026.30.0rc0",
16
16
  ]
17
17
 
18
18
  [dependency-groups]
@@ -34,6 +34,22 @@ class UserMemoryConfig(BaseModel):
34
34
  le=8000,
35
35
  description="Maximum size of the memory profile in tokens.",
36
36
  )
37
+ consolidation_gate_enabled: bool = Field(
38
+ default=True,
39
+ description=(
40
+ "When true, a cheap single-word LLM 'gate' decides whether the turn "
41
+ "warrants a full memory rewrite before the expensive consolidation "
42
+ "runs."
43
+ ),
44
+ )
45
+ updating_notice_enabled: bool = Field(
46
+ default=True,
47
+ description=(
48
+ "When true, a transient 'updating context memory' notice is appended to the "
49
+ "assistant message while the memory rewrite runs, and removed "
50
+ "again once it completes."
51
+ ),
52
+ )
37
53
  root_folder: Annotated[str, RJSFMetaTag.SpecialWidget.hidden()] = Field(
38
54
  default="user-memory",
39
55
  min_length=1,
@@ -1,3 +1,4 @@
1
+ import asyncio
1
2
  from unittest.mock import AsyncMock, MagicMock
2
3
 
3
4
  import pytest
@@ -17,6 +18,8 @@ from unique_user_memory.user_memory import (
17
18
  enforce_token_cap,
18
19
  ensure_user_memory_folder,
19
20
  fit_user_memory,
21
+ noop_update_callback,
22
+ should_consolidate_memory,
20
23
  upload_user_memory,
21
24
  )
22
25
  from unique_user_memory.user_memory_postprocessor import UserMemoryPostprocessor
@@ -166,6 +169,10 @@ async def test_consolidate_user_memory_keeps_existing_on_noop(
166
169
  "unique_user_memory.user_memory.LanguageModelService",
167
170
  MagicMock(return_value=llm_service),
168
171
  )
172
+ monkeypatch.setattr(
173
+ "unique_user_memory.user_memory.should_consolidate_memory",
174
+ AsyncMock(return_value=True),
175
+ )
169
176
 
170
177
  result = await consolidate_user_memory(
171
178
  current_memory=current,
@@ -194,6 +201,298 @@ async def test_consolidate_user_memory_keeps_existing_on_malformed_output(
194
201
  "unique_user_memory.user_memory.LanguageModelService",
195
202
  MagicMock(return_value=llm_service),
196
203
  )
204
+ monkeypatch.setattr(
205
+ "unique_user_memory.user_memory.should_consolidate_memory",
206
+ AsyncMock(return_value=True),
207
+ )
208
+
209
+ result = await consolidate_user_memory(
210
+ current_memory=current,
211
+ user_id="user_1",
212
+ user_message="remember I like concise answers",
213
+ assistant_message="noted",
214
+ config=UserMemoryConfig(),
215
+ language_model=_TEST_LANGUAGE_MODEL,
216
+ event=MagicMock(),
217
+ logger=MagicMock(),
218
+ )
219
+
220
+ assert result == current
221
+
222
+
223
+ @pytest.mark.asyncio
224
+ async def test_consolidate_user_memory_skips_full_rewrite_when_gate_noop(
225
+ monkeypatch: pytest.MonkeyPatch,
226
+ ) -> None:
227
+ current = empty_profile("user_1")
228
+ llm_service = MagicMock()
229
+ llm_service.complete_async = AsyncMock()
230
+ monkeypatch.setattr(
231
+ "unique_user_memory.user_memory.LanguageModelService",
232
+ MagicMock(return_value=llm_service),
233
+ )
234
+ monkeypatch.setattr(
235
+ "unique_user_memory.user_memory.should_consolidate_memory",
236
+ AsyncMock(return_value=False),
237
+ )
238
+
239
+ result = await consolidate_user_memory(
240
+ current_memory=current,
241
+ user_id="user_1",
242
+ user_message="hello",
243
+ assistant_message="hi",
244
+ config=UserMemoryConfig(),
245
+ language_model=_TEST_LANGUAGE_MODEL,
246
+ event=MagicMock(),
247
+ logger=MagicMock(),
248
+ )
249
+
250
+ assert result == current
251
+ llm_service.complete_async.assert_not_awaited()
252
+
253
+
254
+ @pytest.mark.asyncio
255
+ async def test_consolidate_user_memory_runs_full_rewrite_when_gate_update(
256
+ monkeypatch: pytest.MonkeyPatch,
257
+ ) -> None:
258
+ current = empty_profile("user_1")
259
+ rewritten = "# User Memory\n\n## Identity\n- Prefers concise answers"
260
+ response = MagicMock()
261
+ response.choices[0].message.content = rewritten
262
+ llm_service = MagicMock()
263
+ llm_service.complete_async = AsyncMock(return_value=response)
264
+ monkeypatch.setattr(
265
+ "unique_user_memory.user_memory.LanguageModelService",
266
+ MagicMock(return_value=llm_service),
267
+ )
268
+ monkeypatch.setattr(
269
+ "unique_user_memory.user_memory.should_consolidate_memory",
270
+ AsyncMock(return_value=True),
271
+ )
272
+
273
+ result = await consolidate_user_memory(
274
+ current_memory=current,
275
+ user_id="user_1",
276
+ user_message="remember I like concise answers",
277
+ assistant_message="noted",
278
+ config=UserMemoryConfig(),
279
+ language_model=_TEST_LANGUAGE_MODEL,
280
+ event=MagicMock(),
281
+ logger=MagicMock(),
282
+ )
283
+
284
+ assert result == rewritten
285
+ llm_service.complete_async.assert_awaited_once()
286
+
287
+
288
+ @pytest.mark.asyncio
289
+ async def test_consolidate_user_memory_skips_gate_when_disabled(
290
+ monkeypatch: pytest.MonkeyPatch,
291
+ ) -> None:
292
+ current = empty_profile("user_1")
293
+ rewritten = "# User Memory\n\n## Identity\n- Prefers concise answers"
294
+ response = MagicMock()
295
+ response.choices[0].message.content = rewritten
296
+ llm_service = MagicMock()
297
+ llm_service.complete_async = AsyncMock(return_value=response)
298
+ monkeypatch.setattr(
299
+ "unique_user_memory.user_memory.LanguageModelService",
300
+ MagicMock(return_value=llm_service),
301
+ )
302
+ gate = AsyncMock(return_value=False)
303
+ monkeypatch.setattr(
304
+ "unique_user_memory.user_memory.should_consolidate_memory",
305
+ gate,
306
+ )
307
+
308
+ result = await consolidate_user_memory(
309
+ current_memory=current,
310
+ user_id="user_1",
311
+ user_message="hello",
312
+ assistant_message="hi",
313
+ config=UserMemoryConfig(consolidation_gate_enabled=False),
314
+ language_model=_TEST_LANGUAGE_MODEL,
315
+ event=MagicMock(),
316
+ logger=MagicMock(),
317
+ )
318
+
319
+ assert result == rewritten
320
+ gate.assert_not_awaited()
321
+ llm_service.complete_async.assert_awaited_once()
322
+
323
+
324
+ @pytest.mark.asyncio
325
+ async def test_should_consolidate_memory_returns_false_on_noop(
326
+ monkeypatch: pytest.MonkeyPatch,
327
+ ) -> None:
328
+ response = MagicMock()
329
+ response.choices[0].message.content = "NOOP"
330
+ llm_service = MagicMock()
331
+ llm_service.complete_async = AsyncMock(return_value=response)
332
+ monkeypatch.setattr(
333
+ "unique_user_memory.user_memory.LanguageModelService",
334
+ MagicMock(return_value=llm_service),
335
+ )
336
+
337
+ result = await should_consolidate_memory(
338
+ current_memory=empty_profile("user_1"),
339
+ user_id="user_1",
340
+ user_message="hello",
341
+ assistant_message="hi",
342
+ language_model=_TEST_LANGUAGE_MODEL,
343
+ event=MagicMock(),
344
+ logger=MagicMock(),
345
+ )
346
+
347
+ assert result is False
348
+ assert (
349
+ llm_service.complete_async.call_args.kwargs["other_options"]["max_tokens"] == 4
350
+ )
351
+
352
+
353
+ @pytest.mark.asyncio
354
+ async def test_should_consolidate_memory_returns_true_on_update(
355
+ monkeypatch: pytest.MonkeyPatch,
356
+ ) -> None:
357
+ response = MagicMock()
358
+ response.choices[0].message.content = "UPDATE"
359
+ llm_service = MagicMock()
360
+ llm_service.complete_async = AsyncMock(return_value=response)
361
+ monkeypatch.setattr(
362
+ "unique_user_memory.user_memory.LanguageModelService",
363
+ MagicMock(return_value=llm_service),
364
+ )
365
+
366
+ result = await should_consolidate_memory(
367
+ current_memory=empty_profile("user_1"),
368
+ user_id="user_1",
369
+ user_message="remember I like concise answers",
370
+ assistant_message="noted",
371
+ language_model=_TEST_LANGUAGE_MODEL,
372
+ event=MagicMock(),
373
+ logger=MagicMock(),
374
+ )
375
+
376
+ assert result is True
377
+
378
+
379
+ @pytest.mark.asyncio
380
+ async def test_should_consolidate_memory_falls_back_to_true_on_error(
381
+ monkeypatch: pytest.MonkeyPatch,
382
+ ) -> None:
383
+ llm_service = MagicMock()
384
+ llm_service.complete_async = AsyncMock(side_effect=RuntimeError("boom"))
385
+ monkeypatch.setattr(
386
+ "unique_user_memory.user_memory.LanguageModelService",
387
+ MagicMock(return_value=llm_service),
388
+ )
389
+
390
+ result = await should_consolidate_memory(
391
+ current_memory=empty_profile("user_1"),
392
+ user_id="user_1",
393
+ user_message="hello",
394
+ assistant_message="hi",
395
+ language_model=_TEST_LANGUAGE_MODEL,
396
+ event=MagicMock(),
397
+ logger=MagicMock(),
398
+ )
399
+
400
+ assert result is True
401
+
402
+
403
+ @pytest.mark.asyncio
404
+ async def test_consolidate_user_memory_invokes_update_callbacks_on_rewrite(
405
+ monkeypatch: pytest.MonkeyPatch,
406
+ ) -> None:
407
+ rewritten = "# User Memory\n\n## Identity\n- Prefers concise answers"
408
+ response = MagicMock()
409
+ response.choices[0].message.content = rewritten
410
+ llm_service = MagicMock()
411
+ llm_service.complete_async = AsyncMock(return_value=response)
412
+ monkeypatch.setattr(
413
+ "unique_user_memory.user_memory.LanguageModelService",
414
+ MagicMock(return_value=llm_service),
415
+ )
416
+ monkeypatch.setattr(
417
+ "unique_user_memory.user_memory.should_consolidate_memory",
418
+ AsyncMock(return_value=True),
419
+ )
420
+ events: list[str] = []
421
+ on_start = AsyncMock(side_effect=lambda: events.append("start"))
422
+ on_end = AsyncMock(side_effect=lambda: events.append("end"))
423
+
424
+ result = await consolidate_user_memory(
425
+ current_memory=empty_profile("user_1"),
426
+ user_id="user_1",
427
+ user_message="remember I like concise answers",
428
+ assistant_message="noted",
429
+ config=UserMemoryConfig(),
430
+ language_model=_TEST_LANGUAGE_MODEL,
431
+ event=MagicMock(),
432
+ logger=MagicMock(),
433
+ on_update_start=on_start,
434
+ on_update_end=on_end,
435
+ )
436
+
437
+ assert result == rewritten
438
+ on_start.assert_awaited_once()
439
+ on_end.assert_awaited_once()
440
+ assert events == ["start", "end"]
441
+
442
+
443
+ @pytest.mark.asyncio
444
+ async def test_consolidate_user_memory_skips_update_callbacks_on_noop(
445
+ monkeypatch: pytest.MonkeyPatch,
446
+ ) -> None:
447
+ current = empty_profile("user_1")
448
+ llm_service = MagicMock()
449
+ llm_service.complete_async = AsyncMock()
450
+ monkeypatch.setattr(
451
+ "unique_user_memory.user_memory.LanguageModelService",
452
+ MagicMock(return_value=llm_service),
453
+ )
454
+ monkeypatch.setattr(
455
+ "unique_user_memory.user_memory.should_consolidate_memory",
456
+ AsyncMock(return_value=False),
457
+ )
458
+ on_start = AsyncMock()
459
+ on_end = AsyncMock()
460
+
461
+ result = await consolidate_user_memory(
462
+ current_memory=current,
463
+ user_id="user_1",
464
+ user_message="hello",
465
+ assistant_message="hi",
466
+ config=UserMemoryConfig(),
467
+ language_model=_TEST_LANGUAGE_MODEL,
468
+ event=MagicMock(),
469
+ logger=MagicMock(),
470
+ on_update_start=on_start,
471
+ on_update_end=on_end,
472
+ )
473
+
474
+ assert result == current
475
+ on_start.assert_not_awaited()
476
+ on_end.assert_not_awaited()
477
+
478
+
479
+ @pytest.mark.asyncio
480
+ async def test_consolidate_user_memory_invokes_update_end_on_rewrite_error(
481
+ monkeypatch: pytest.MonkeyPatch,
482
+ ) -> None:
483
+ current = empty_profile("user_1")
484
+ llm_service = MagicMock()
485
+ llm_service.complete_async = AsyncMock(side_effect=RuntimeError("boom"))
486
+ monkeypatch.setattr(
487
+ "unique_user_memory.user_memory.LanguageModelService",
488
+ MagicMock(return_value=llm_service),
489
+ )
490
+ monkeypatch.setattr(
491
+ "unique_user_memory.user_memory.should_consolidate_memory",
492
+ AsyncMock(return_value=True),
493
+ )
494
+ on_start = AsyncMock()
495
+ on_end = AsyncMock()
197
496
 
198
497
  result = await consolidate_user_memory(
199
498
  current_memory=current,
@@ -204,9 +503,153 @@ async def test_consolidate_user_memory_keeps_existing_on_malformed_output(
204
503
  language_model=_TEST_LANGUAGE_MODEL,
205
504
  event=MagicMock(),
206
505
  logger=MagicMock(),
506
+ on_update_start=on_start,
507
+ on_update_end=on_end,
207
508
  )
208
509
 
209
510
  assert result == current
511
+ on_start.assert_awaited_once()
512
+ on_end.assert_awaited_once()
513
+
514
+
515
+ @pytest.mark.asyncio
516
+ async def test_consolidate_user_memory_invokes_update_end_when_start_cancelled(
517
+ monkeypatch: pytest.MonkeyPatch,
518
+ ) -> None:
519
+ # on_update_start may show the notice and then be cancelled. CancelledError
520
+ # is a BaseException (not caught in _set_message_content), so the cleanup
521
+ # callback must still run to remove the transient notice.
522
+ current = empty_profile("user_1")
523
+ llm_service = MagicMock()
524
+ llm_service.complete_async = AsyncMock()
525
+ monkeypatch.setattr(
526
+ "unique_user_memory.user_memory.LanguageModelService",
527
+ MagicMock(return_value=llm_service),
528
+ )
529
+ monkeypatch.setattr(
530
+ "unique_user_memory.user_memory.should_consolidate_memory",
531
+ AsyncMock(return_value=True),
532
+ )
533
+ on_start = AsyncMock(side_effect=asyncio.CancelledError)
534
+ on_end = AsyncMock()
535
+
536
+ with pytest.raises(asyncio.CancelledError):
537
+ await consolidate_user_memory(
538
+ current_memory=current,
539
+ user_id="user_1",
540
+ user_message="remember I like concise answers",
541
+ assistant_message="noted",
542
+ config=UserMemoryConfig(),
543
+ language_model=_TEST_LANGUAGE_MODEL,
544
+ event=MagicMock(),
545
+ logger=MagicMock(),
546
+ on_update_start=on_start,
547
+ on_update_end=on_end,
548
+ )
549
+
550
+ on_start.assert_awaited_once()
551
+ on_end.assert_awaited_once()
552
+ llm_service.complete_async.assert_not_awaited()
553
+
554
+
555
+ @pytest.mark.asyncio
556
+ async def test_user_memory_postprocessor_shows_and_removes_updating_notice(
557
+ monkeypatch: pytest.MonkeyPatch,
558
+ ) -> None:
559
+ updated_memory = "# User Memory\n\n## Identity\n- Updated"
560
+ original_text = "Here is your answer."
561
+
562
+ async def fake_consolidate(*, on_update_start, on_update_end, **kwargs) -> str: # type: ignore[no-untyped-def]
563
+ await on_update_start()
564
+ await on_update_end()
565
+ return updated_memory
566
+
567
+ monkeypatch.setattr(
568
+ "unique_user_memory.user_memory_postprocessor.consolidate_user_memory",
569
+ fake_consolidate,
570
+ )
571
+ monkeypatch.setattr(
572
+ "unique_user_memory.user_memory_postprocessor.upload_user_memory",
573
+ AsyncMock(return_value=True),
574
+ )
575
+ chat_service = MagicMock()
576
+ chat_service.modify_assistant_message_async = AsyncMock()
577
+ monkeypatch.setattr(
578
+ "unique_user_memory.user_memory_postprocessor.ChatService",
579
+ MagicMock(return_value=chat_service),
580
+ )
581
+ event = MagicMock()
582
+ event.user_id = "user_1"
583
+ event.company_id = "company_1"
584
+ event.payload.user_message.text = "remember this"
585
+ loop_response = MagicMock()
586
+ loop_response.message.text = original_text
587
+ loop_response.message.id = "msg_1"
588
+ loop_response.message.references = []
589
+ postprocessor = UserMemoryPostprocessor(
590
+ config=UserMemoryConfig(),
591
+ language_model=_TEST_LANGUAGE_MODEL,
592
+ event=event,
593
+ state=UserMemoryState(scope_id="scope_1", text=empty_profile("user_1")),
594
+ logger=MagicMock(),
595
+ chat_service=chat_service,
596
+ )
597
+
598
+ await postprocessor.run(loop_response)
599
+
600
+ calls = chat_service.modify_assistant_message_async.await_args_list
601
+ assert len(calls) == 2
602
+ assert calls[0].kwargs["content"].startswith(original_text)
603
+ assert calls[0].kwargs["content"] != original_text
604
+ assert calls[0].kwargs["message_id"] == "msg_1"
605
+ assert calls[1].kwargs["content"] == original_text
606
+
607
+
608
+ @pytest.mark.asyncio
609
+ async def test_user_memory_postprocessor_skips_notice_when_disabled(
610
+ monkeypatch: pytest.MonkeyPatch,
611
+ ) -> None:
612
+ updated_memory = "# User Memory\n\n## Identity\n- Updated"
613
+
614
+ async def fake_consolidate(*, on_update_start, on_update_end, **kwargs) -> str: # type: ignore[no-untyped-def]
615
+ assert on_update_start is noop_update_callback
616
+ assert on_update_end is noop_update_callback
617
+ return updated_memory
618
+
619
+ monkeypatch.setattr(
620
+ "unique_user_memory.user_memory_postprocessor.consolidate_user_memory",
621
+ fake_consolidate,
622
+ )
623
+ monkeypatch.setattr(
624
+ "unique_user_memory.user_memory_postprocessor.upload_user_memory",
625
+ AsyncMock(return_value=True),
626
+ )
627
+ chat_service = MagicMock()
628
+ chat_service.modify_assistant_message_async = AsyncMock()
629
+ monkeypatch.setattr(
630
+ "unique_user_memory.user_memory_postprocessor.ChatService",
631
+ MagicMock(return_value=chat_service),
632
+ )
633
+ event = MagicMock()
634
+ event.user_id = "user_1"
635
+ event.company_id = "company_1"
636
+ event.payload.user_message.text = "remember this"
637
+ loop_response = MagicMock()
638
+ loop_response.message.text = "answer"
639
+ loop_response.message.id = "msg_1"
640
+ loop_response.message.references = []
641
+ postprocessor = UserMemoryPostprocessor(
642
+ config=UserMemoryConfig(updating_notice_enabled=False),
643
+ language_model=_TEST_LANGUAGE_MODEL,
644
+ event=event,
645
+ state=UserMemoryState(scope_id="scope_1", text=empty_profile("user_1")),
646
+ logger=MagicMock(),
647
+ chat_service=chat_service,
648
+ )
649
+
650
+ await postprocessor.run(loop_response)
651
+
652
+ chat_service.modify_assistant_message_async.assert_not_awaited()
210
653
 
211
654
 
212
655
  @pytest.mark.asyncio
@@ -496,16 +939,19 @@ async def test_user_memory_postprocessor_logs_success_when_upload_succeeds(
496
939
  loop_response = MagicMock()
497
940
  loop_response.message.text = "noted"
498
941
  logger = MagicMock()
942
+ chat_service = MagicMock()
499
943
  postprocessor = UserMemoryPostprocessor(
500
944
  config=UserMemoryConfig(),
501
945
  language_model=_TEST_LANGUAGE_MODEL,
502
946
  event=event,
503
947
  state=UserMemoryState(scope_id="scope_1", text=empty_profile("user_1")),
504
948
  logger=logger,
949
+ chat_service=chat_service,
505
950
  )
506
951
 
507
- await postprocessor.run(loop_response)
952
+ updated = await postprocessor.run(loop_response)
508
953
 
954
+ assert updated is True
509
955
  upload.assert_awaited_once_with(
510
956
  scope_id="scope_1",
511
957
  content=updated_memory,
@@ -538,16 +984,19 @@ async def test_user_memory_postprocessor_does_not_log_success_when_upload_fails(
538
984
  loop_response = MagicMock()
539
985
  loop_response.message.text = "noted"
540
986
  logger = MagicMock()
987
+ chat_service = MagicMock()
541
988
  postprocessor = UserMemoryPostprocessor(
542
989
  config=UserMemoryConfig(),
543
990
  language_model=_TEST_LANGUAGE_MODEL,
544
991
  event=event,
545
992
  state=UserMemoryState(scope_id="scope_1", text=empty_profile("user_1")),
546
993
  logger=logger,
994
+ chat_service=chat_service,
547
995
  )
548
996
 
549
- await postprocessor.run(loop_response)
997
+ updated = await postprocessor.run(loop_response)
550
998
 
999
+ assert updated is False
551
1000
  assert all(
552
1001
  call.args != ("[user-memory] memory updated and uploaded successfully",)
553
1002
  for call in logger.info.call_args_list
@@ -1,4 +1,5 @@
1
1
  import re
2
+ from collections.abc import Awaitable, Callable
2
3
  from dataclasses import dataclass
3
4
  from logging import Logger
4
5
 
@@ -28,11 +29,27 @@ from unique_user_memory.user_memory_prompts import (
28
29
  consolidation_system_prompt,
29
30
  consolidation_user_prompt,
30
31
  empty_profile,
32
+ memory_gate_system_prompt,
33
+ memory_gate_user_prompt,
31
34
  )
32
35
 
33
36
  MEMORY_FILENAME = "memory.md"
34
37
  MIME_TYPE = "text/markdown"
35
38
  _LLM_OUTPUT_HEADROOM_TOKENS = 200
39
+
40
+
41
+ async def noop_update_callback() -> None:
42
+ """Default update hook that does nothing.
43
+
44
+ Used as the default for ``on_update_start`` / ``on_update_end`` so callers
45
+ that do not need update notifications can be awaited unconditionally.
46
+ """
47
+ return None
48
+
49
+
50
+ # The gate only ever replies with the single word UPDATE or NOOP; a tiny
51
+ # output budget keeps the common (NOOP) path cheap and fast.
52
+ _GATE_MAX_TOKENS = 4
36
53
  # When condensing an oversized profile, aim below the hard cap so the LLM
37
54
  # output leaves headroom and the hard-cut safety net rarely has to fire.
38
55
  _CONDENSE_TARGET_RATIO = 0.9
@@ -551,6 +568,91 @@ async def upload_user_memory(
551
568
  return False
552
569
 
553
570
 
571
+ async def should_consolidate_memory(
572
+ *,
573
+ current_memory: str,
574
+ user_id: str,
575
+ user_message: str,
576
+ assistant_message: str,
577
+ language_model: LanguageModelInfo,
578
+ event: ChatEvent,
579
+ logger: Logger,
580
+ ) -> bool:
581
+ """Cheaply decide whether the turn warrants a full memory rewrite.
582
+
583
+ Runs a focused LLM call capped at a few output tokens that answers
584
+ ``UPDATE`` or ``NOOP``. Returns ``False`` only on an explicit ``NOOP``
585
+ so the caller can skip the expensive full-profile regeneration. Any
586
+ error (or ambiguous output) falls back to ``True`` so behaviour stays
587
+ identical to the pre-gate path.
588
+ """
589
+ try:
590
+ llm_service = LanguageModelService(event)
591
+ except Exception as exc:
592
+ logger.warning(
593
+ "[user-memory] cannot construct LanguageModelService for gate: [%s] %s",
594
+ type(exc).__name__,
595
+ exc,
596
+ )
597
+ return True
598
+
599
+ messages = LanguageModelMessages(
600
+ [
601
+ LanguageModelSystemMessage(content=memory_gate_system_prompt()),
602
+ LanguageModelUserMessage(
603
+ content=memory_gate_user_prompt(
604
+ user_id=user_id,
605
+ existing_memory=current_memory,
606
+ user_message=_sanitize_for_xml_context(user_message or ""),
607
+ assistant_message=_sanitize_for_xml_context(
608
+ assistant_message or ""
609
+ ),
610
+ )
611
+ ),
612
+ ]
613
+ )
614
+
615
+ try:
616
+ response = await llm_service.complete_async(
617
+ messages=messages,
618
+ model_name=language_model.name,
619
+ other_options={"max_tokens": _GATE_MAX_TOKENS},
620
+ )
621
+ except Exception as exc:
622
+ logger.warning(
623
+ "[user-memory] gate LLM call failed (model=%s): [%s] %s",
624
+ language_model.name,
625
+ type(exc).__name__,
626
+ exc,
627
+ )
628
+ return True
629
+
630
+ try:
631
+ raw = response.choices[0].message.content or ""
632
+ except Exception as exc:
633
+ logger.warning(
634
+ "[user-memory] could not extract content from gate response: [%s] %s",
635
+ type(exc).__name__,
636
+ exc,
637
+ )
638
+ return True
639
+
640
+ if not isinstance(raw, str):
641
+ logger.warning(
642
+ "[user-memory] gate returned non-string content (%s)",
643
+ type(raw).__name__,
644
+ )
645
+ return True
646
+
647
+ decision = raw.strip().upper()
648
+ if decision.startswith("NOOP"):
649
+ logger.info("[user-memory] gate decided NOOP - skipping consolidation")
650
+ return False
651
+
652
+ logger.info("[user-memory] gate decided UPDATE - consolidating")
653
+ return True
654
+
655
+
554
656
  async def consolidate_user_memory(
555
657
  *,
556
658
  current_memory: str,
@@ -561,7 +663,17 @@ async def consolidate_user_memory(
561
663
  language_model: LanguageModelInfo,
562
664
  event: ChatEvent,
563
665
  logger: Logger,
666
+ on_update_start: Callable[[], Awaitable[None]] = noop_update_callback,
667
+ on_update_end: Callable[[], Awaitable[None]] = noop_update_callback,
564
668
  ) -> str:
669
+ """Consolidate the latest turn into the user's memory profile.
670
+
671
+ When the (optional) gate decides the turn is worth remembering, the
672
+ expensive full-profile rewrite runs. ``on_update_start`` is awaited
673
+ just before that slow rewrite begins and ``on_update_end`` right after
674
+ it finishes (even on failure), so a caller can surface a transient
675
+ "updating memory" notice to the user only while real work happens.
676
+ """
565
677
  safe_current = enforce_token_cap(
566
678
  content=current_memory,
567
679
  max_tokens=config.max_tokens,
@@ -582,6 +694,44 @@ async def consolidate_user_memory(
582
694
  language_model=language_model,
583
695
  )
584
696
 
697
+ if config.consolidation_gate_enabled and not await should_consolidate_memory(
698
+ current_memory=safe_current,
699
+ user_id=user_id,
700
+ user_message=user_message,
701
+ assistant_message=assistant_message,
702
+ language_model=language_model,
703
+ event=event,
704
+ logger=logger,
705
+ ):
706
+ return safe_current
707
+
708
+ try:
709
+ await on_update_start()
710
+ return await _rewrite_user_memory(
711
+ safe_current=safe_current,
712
+ user_id=user_id,
713
+ user_message=user_message,
714
+ assistant_message=assistant_message,
715
+ config=config,
716
+ language_model=language_model,
717
+ event=event,
718
+ logger=logger,
719
+ )
720
+ finally:
721
+ await on_update_end()
722
+
723
+
724
+ async def _rewrite_user_memory(
725
+ *,
726
+ safe_current: str,
727
+ user_id: str,
728
+ user_message: str,
729
+ assistant_message: str,
730
+ config: UserMemoryConfig,
731
+ language_model: LanguageModelInfo,
732
+ event: ChatEvent,
733
+ logger: Logger,
734
+ ) -> str:
585
735
  if not safe_current.strip():
586
736
  safe_current = empty_profile(user_id)
587
737
 
@@ -1,7 +1,10 @@
1
+ from collections.abc import Awaitable, Callable
1
2
  from logging import Logger
2
3
 
3
4
  from unique_toolkit.agentic.postprocessor.postprocessor_manager import Postprocessor
4
5
  from unique_toolkit.app.schemas import ChatEvent
6
+ from unique_toolkit.chat.service import ChatService
7
+ from unique_toolkit.content.schemas import ContentReference
5
8
  from unique_toolkit.language_model.default_language_model import (
6
9
  DEFAULT_LANGUAGE_MODEL,
7
10
  )
@@ -12,9 +15,14 @@ from unique_user_memory.config import UserMemoryConfig
12
15
  from unique_user_memory.user_memory import (
13
16
  UserMemoryState,
14
17
  consolidate_user_memory,
18
+ noop_update_callback,
15
19
  upload_user_memory,
16
20
  )
17
21
 
22
+ # Transient marker appended to the assistant message while the (slow) memory
23
+ # rewrite runs; removed again once consolidation finishes.
24
+ _UPDATING_NOTICE = "\n\n---\n\n🧠 _Updating context memory…_"
25
+
18
26
 
19
27
  class UserMemoryPostprocessor(Postprocessor):
20
28
  def __init__(
@@ -27,6 +35,7 @@ class UserMemoryPostprocessor(Postprocessor):
27
35
  event: ChatEvent,
28
36
  state: UserMemoryState,
29
37
  logger: Logger,
38
+ chat_service: ChatService,
30
39
  ) -> None:
31
40
  super().__init__(name="UserMemoryPostprocessor")
32
41
  self._config = config
@@ -39,13 +48,45 @@ class UserMemoryPostprocessor(Postprocessor):
39
48
  self._state = state
40
49
  self._logger = logger
41
50
  self._new_memory: str | None = None
51
+ self._chat_service: ChatService = chat_service
52
+
53
+ async def run(self, loop_response: LanguageModelStreamResponse) -> bool:
54
+ """Consolidate and upload user memory for this turn.
42
55
 
43
- async def run(self, loop_response: LanguageModelStreamResponse) -> None:
56
+ Returns True if the memory profile changed and was uploaded, False
57
+ otherwise (no user/company, NOOP consolidation, or failed upload).
58
+ """
44
59
  self._logger.info("[user-memory] running postprocessor")
45
60
  user_id = self._event.user_id
46
61
  company_id = self._event.company_id
47
62
  if not user_id or not company_id:
48
- return
63
+ return False
64
+
65
+ on_update_start: Callable[[], Awaitable[None]] = noop_update_callback
66
+ on_update_end: Callable[[], Awaitable[None]] = noop_update_callback
67
+ if self._config.updating_notice_enabled:
68
+ original_text = loop_response.message.text or ""
69
+ message_id = loop_response.message.id
70
+ references = loop_response.message.references
71
+
72
+ async def _on_update_start() -> None:
73
+ await self._set_message_content(
74
+ content=original_text + _UPDATING_NOTICE,
75
+ message_id=message_id,
76
+ references=references,
77
+ action="show updating notice",
78
+ )
79
+
80
+ async def _on_update_end() -> None:
81
+ await self._set_message_content(
82
+ content=original_text,
83
+ message_id=message_id,
84
+ references=references,
85
+ action="remove updating notice",
86
+ )
87
+
88
+ on_update_start = _on_update_start
89
+ on_update_end = _on_update_end
49
90
 
50
91
  self._new_memory = await consolidate_user_memory(
51
92
  current_memory=self._state.text,
@@ -56,11 +97,13 @@ class UserMemoryPostprocessor(Postprocessor):
56
97
  language_model=self._language_model,
57
98
  event=self._event,
58
99
  logger=self._logger,
100
+ on_update_start=on_update_start,
101
+ on_update_end=on_update_end,
59
102
  )
60
103
 
61
104
  if self._new_memory == self._state.text:
62
105
  self._logger.debug("[user-memory] consolidation NOOP - skipping upload")
63
- return
106
+ return False
64
107
 
65
108
  uploaded = await upload_user_memory(
66
109
  scope_id=self._state.scope_id,
@@ -71,9 +114,32 @@ class UserMemoryPostprocessor(Postprocessor):
71
114
  )
72
115
  if not uploaded:
73
116
  self._logger.warning("[user-memory] memory update was not uploaded")
74
- return
117
+ return False
75
118
 
76
119
  self._logger.info("[user-memory] memory updated and uploaded successfully")
120
+ return True
121
+
122
+ async def _set_message_content(
123
+ self,
124
+ *,
125
+ content: str,
126
+ message_id: str | None,
127
+ references: list[ContentReference] | None,
128
+ action: str,
129
+ ) -> None:
130
+ try:
131
+ await self._chat_service.modify_assistant_message_async(
132
+ content=content,
133
+ message_id=message_id,
134
+ references=references,
135
+ )
136
+ except Exception as exc:
137
+ self._logger.warning(
138
+ "[user-memory] failed to %s: [%s] %s",
139
+ action,
140
+ type(exc).__name__,
141
+ exc,
142
+ )
77
143
 
78
144
  def apply_postprocessing_to_response(
79
145
  self, loop_response: LanguageModelStreamResponse
@@ -291,3 +291,85 @@ def consolidation_user_prompt(
291
291
  user_message=(user_message or "").strip(),
292
292
  assistant_message=(assistant_message or "").strip(),
293
293
  )
294
+
295
+
296
+ _GATE_SYSTEM_PROMPT_TEMPLATE = """\
297
+ You are the decision gate for a user-memory system on the Unique AI
298
+ platform. A structured Markdown profile of the user is maintained across
299
+ conversations. Rewriting that profile is expensive, so it must happen
300
+ only when the latest turn actually adds new, durable knowledge about the
301
+ user.
302
+
303
+ # Inputs
304
+
305
+ 1. `<existing_memory>` - the current profile (may be empty on the first turn).
306
+ 2. `<new_turn>` - the most recent user message and the assistant's reply.
307
+
308
+ # Your task
309
+
310
+ Reply with EXACTLY ONE uppercase word and nothing else:
311
+
312
+ - `UPDATE` - the new turn contains at least one durable fact about the
313
+ user that is NOT already captured in `<existing_memory>`, or that
314
+ changes/contradicts something already stored.
315
+ - `NOOP` - otherwise. Choose `NOOP` for small talk, greetings, factual
316
+ questions, code or writing requests, abstract discussion, and for any
317
+ fact that is already present in `<existing_memory>`.
318
+
319
+ Do NOT output the profile, an explanation, punctuation, or code fences -
320
+ only the single word `UPDATE` or `NOOP`.
321
+
322
+ # What counts as a durable fact (lean UPDATE)
323
+
324
+ - Stable attributes: name, role, employer, team, timezone, language,
325
+ technical stack, recurring projects.
326
+ - Preferences: communication style, formatting, depth, tone, language,
327
+ expertise level.
328
+ - Durable context: current focus areas, active projects, multi-week
329
+ goals, deadlines stated by the user.
330
+ - Explicit hand-offs: "remind me about X", "let's revisit Y later".
331
+
332
+ # What NEVER justifies UPDATE (lean NOOP)
333
+
334
+ - Credentials, API keys, passwords, payment or health details, IDs.
335
+ - Facts about other people beyond the immediate professional context.
336
+ - Transient turn state: one-off answers, code snippets, error messages,
337
+ file contents, search results.
338
+ - Anything stated as third-party or retrieved context.
339
+ - Facts already captured in `<existing_memory>`.
340
+ """
341
+
342
+
343
+ def memory_gate_system_prompt() -> str:
344
+ return Template(_GATE_SYSTEM_PROMPT_TEMPLATE).render()
345
+
346
+
347
+ _GATE_USER_PROMPT_TEMPLATE = """\
348
+ User ID: {{ user_id }}
349
+
350
+ <existing_memory>
351
+ {{ existing_memory }}
352
+ </existing_memory>
353
+
354
+ <new_turn>
355
+ user: {{ user_message }}
356
+ assistant: {{ assistant_message }}
357
+ </new_turn>
358
+
359
+ Reply with the single word UPDATE or NOOP now.
360
+ """
361
+
362
+
363
+ def memory_gate_user_prompt(
364
+ user_id: str,
365
+ existing_memory: str,
366
+ user_message: str,
367
+ assistant_message: str,
368
+ ) -> str:
369
+ existing = existing_memory.strip() or "(empty - this is the user's first turn)"
370
+ return Template(_GATE_USER_PROMPT_TEMPLATE).render(
371
+ user_id=user_id,
372
+ existing_memory=existing,
373
+ user_message=(user_message or "").strip(),
374
+ assistant_message=(assistant_message or "").strip(),
375
+ )