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

{unique_user_memory-2026.32.0.dev1 → unique_user_memory-2026.32.0.dev2}/PKG-INFO

@@ -1,14 +1,14 @@
 Metadata-Version: 2.3
 Name: unique-user-memory
-Version: 2026.32.0.dev1
+Version: 2026.32.0.dev2
 Summary: 
 Author: Fabian Schläpfer
 Author-email: Fabian Schläpfer <fabian@unique.ch>
 License: Proprietary
 Requires-Dist: jinja2>=3.1.6
 Requires-Dist: pydantic>=2.8.2
-Requires-Dist: unique-sdk>=2026.32.0.dev0,<2026.32.0rc0
-Requires-Dist: unique-toolkit>=2026.32.0.dev0,<2026.32.0rc0
+Requires-Dist: unique-sdk>=2026.32.0.dev5,<2026.32.0rc0
+Requires-Dist: unique-toolkit>=2026.32.0.dev4,<2026.32.0rc0
 Requires-Python: >=3.12, <4
 Description-Content-Type: text/markdown
 

{unique_user_memory-2026.32.0.dev1 → unique_user_memory-2026.32.0.dev2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "unique_user_memory"
-version = "2026.32.0.dev1"
+version = "2026.32.0.dev2"
 description = ""
 authors = [
     { name = "Fabian Schläpfer", email = "fabian@unique.ch" },
@@ -11,8 +11,8 @@ requires-python = ">=3.12,<4"
 dependencies = [
     "jinja2>=3.1.6",
     "pydantic>=2.8.2",
-    "unique-sdk>=2026.32.0.dev0,<2026.32.0rc0",
-    "unique-toolkit>=2026.32.0.dev0,<2026.32.0rc0",
+    "unique-sdk>=2026.32.0.dev5,<2026.32.0rc0",
+    "unique-toolkit>=2026.32.0.dev4,<2026.32.0rc0",
 ]
 
 [dependency-groups]

{unique_user_memory-2026.32.0.dev1 → unique_user_memory-2026.32.0.dev2}/unique_user_memory/tests/test_user_memory.py

@@ -6,6 +6,8 @@ from unique_toolkit.language_model.default_language_model import (
     DEFAULT_LANGUAGE_MODEL,
 )
 from unique_toolkit.language_model.infos import LanguageModelInfo
+from unique_toolkit.language_model.invocation_stats import LanguageModelInvocationStats
+from unique_toolkit.language_model.schemas import LanguageModelTokenUsage
 
 from unique_user_memory.config import UserMemoryConfig
 from unique_user_memory.user_memory import (
@@ -1088,6 +1090,106 @@ async def test_user_memory_postprocessor_logs_success_when_upload_succeeds(
     )
 
 
+@pytest.mark.ai
+@pytest.mark.asyncio
+async def test_user_memory_postprocessor_run_resets_invocation_stats(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Purpose: Verify each run reports only usage attributable to that run.
+    Why this matters: Reused postprocessors must not inflate token analytics.
+    Setup summary: Run twice with distinct usage and assert the second excludes the first.
+    """
+    load_stats = LanguageModelInvocationStats.from_usage(
+        _TEST_LANGUAGE_MODEL.name,
+        LanguageModelTokenUsage(total_tokens=2),
+        source="user_memory_load_condense",
+    )
+    first_run_stats = LanguageModelInvocationStats.from_usage(
+        _TEST_LANGUAGE_MODEL.name,
+        LanguageModelTokenUsage(total_tokens=3),
+        source="user_memory_consolidate_first",
+    )
+    second_run_stats = LanguageModelInvocationStats.from_usage(
+        _TEST_LANGUAGE_MODEL.name,
+        LanguageModelTokenUsage(total_tokens=5),
+        source="user_memory_consolidate_second",
+    )
+    run_stats = iter((first_run_stats, second_run_stats))
+
+    async def consolidate(*, invocation_stats, **kwargs) -> str:  # type: ignore[no-untyped-def]
+        invocation_stats.append(next(run_stats))
+        return "# User Memory\n\n## Identity\n- unchanged"
+
+    monkeypatch.setattr(
+        "unique_user_memory.user_memory_postprocessor.consolidate_user_memory",
+        consolidate,
+    )
+    event = MagicMock()
+    event.user_id = "user_1"
+    event.company_id = "company_1"
+    event.payload.user_message.text = "remember this"
+    loop_response = MagicMock()
+    loop_response.message.text = "noted"
+    state = UserMemoryState(
+        scope_id="scope_1",
+        text="# User Memory\n\n## Identity\n- unchanged",
+        load_invocation_stats=(load_stats,),
+    )
+    postprocessor = UserMemoryPostprocessor(
+        config=UserMemoryConfig(),
+        language_model=_TEST_LANGUAGE_MODEL,
+        event=event,
+        state=state,
+        logger=MagicMock(),
+        chat_service=MagicMock(),
+    )
+
+    await postprocessor.run(loop_response)
+    first_reported_stats = postprocessor.invocation_stats
+    await postprocessor.run(loop_response)
+
+    assert first_reported_stats == [load_stats, first_run_stats]
+    assert postprocessor.invocation_stats == [second_run_stats]
+
+
+@pytest.mark.ai
+def test_user_memory_postprocessor_take_pending_invocation_stats_drains_once() -> None:
+    """Purpose: Verify load-time usage is reported exactly once, however it's read.
+    Why this matters: A turn that exits before `run()` (cancellation, empty
+    response, a control-taking tool) must still report the load-time condense
+    tokens, and a turn that does reach `run()` must not double-count them.
+    Setup summary: Take the pending stats directly, then run(), and assert
+    run() no longer reports the already-taken load stats.
+    """
+    load_stats = LanguageModelInvocationStats.from_usage(
+        _TEST_LANGUAGE_MODEL.name,
+        LanguageModelTokenUsage(total_tokens=2),
+        source="user_memory_load_condense",
+    )
+    event = MagicMock()
+    event.user_id = "user_1"
+    event.company_id = "company_1"
+    event.payload.user_message.text = "remember this"
+    state = UserMemoryState(
+        scope_id="scope_1",
+        text="# User Memory\n\n## Identity\n- unchanged",
+        load_invocation_stats=(load_stats,),
+    )
+    postprocessor = UserMemoryPostprocessor(
+        config=UserMemoryConfig(),
+        language_model=_TEST_LANGUAGE_MODEL,
+        event=event,
+        state=state,
+        logger=MagicMock(),
+        chat_service=MagicMock(),
+    )
+
+    taken = postprocessor.take_pending_invocation_stats()
+
+    assert taken == [load_stats]
+    assert postprocessor.take_pending_invocation_stats() == []
+
+
 @pytest.mark.asyncio
 async def test_user_memory_postprocessor_does_not_log_success_when_upload_fails(
     monkeypatch: pytest.MonkeyPatch,

{unique_user_memory-2026.32.0.dev1 → unique_user_memory-2026.32.0.dev2}/unique_user_memory/user_memory.py

@@ -21,6 +21,9 @@ from unique_toolkit.language_model import (
     TypeEncoder,
 )
 from unique_toolkit.language_model.infos import LanguageModelInfo
+from unique_toolkit.language_model.invocation_stats import (
+    LanguageModelInvocationStats,
+)
 
 from unique_user_memory.config import UserMemoryConfig
 from unique_user_memory.user_memory_prompts import (
@@ -97,6 +100,7 @@ def _restore_frontmatter(original: str, body: str) -> str:
 class UserMemoryState:
     scope_id: str
     text: str
+    load_invocation_stats: tuple[LanguageModelInvocationStats, ...] = ()
 
 
 def _get_model_tokenizer(
@@ -181,6 +185,8 @@ async def condense_user_memory(
     language_model: LanguageModelInfo,
     event: ChatEvent,
     logger: Logger,
+    invocation_stats: list[LanguageModelInvocationStats] | None = None,
+    invocation_source: str = "user_memory_condense",
 ) -> str | None:
     """Ask the LLM to rewrite an oversized profile into a shorter one.
 
@@ -233,6 +239,15 @@ async def condense_user_memory(
         )
         return None
 
+    if invocation_stats is not None and response.usage is not None:
+        invocation_stats.append(
+            LanguageModelInvocationStats.from_usage(
+                language_model.name,
+                response.usage,
+                source=invocation_source,
+            )
+        )
+
     try:
         raw = response.choices[0].message.content or ""
     except Exception as exc:
@@ -268,6 +283,8 @@ async def fit_user_memory(
     language_model: LanguageModelInfo,
     event: ChatEvent,
     logger: Logger,
+    invocation_stats: list[LanguageModelInvocationStats] | None = None,
+    invocation_source: str = "user_memory_condense",
 ) -> str:
     """Ensure ``content`` fits ``max_tokens``, condensing before cutting.
 
@@ -293,6 +310,8 @@ async def fit_user_memory(
         language_model=language_model,
         event=event,
         logger=logger,
+        invocation_stats=invocation_stats,
+        invocation_source=invocation_source,
     )
     if condensed is not None:
         condensed = _restore_frontmatter(content, condensed)
@@ -358,6 +377,7 @@ async def load_user_memory(
         company_id=company_id,
         logger=logger,
     )
+    invocation_stats: list[LanguageModelInvocationStats] = []
     return UserMemoryState(
         scope_id=scope_id,
         text=await fit_user_memory(
@@ -366,7 +386,10 @@ async def load_user_memory(
             language_model=language_model,
             event=event,
             logger=logger,
+            invocation_stats=invocation_stats,
+            invocation_source="user_memory_load_condense",
         ),
+        load_invocation_stats=tuple(invocation_stats),
     )
 
 
@@ -614,6 +637,7 @@ async def should_consolidate_memory(
     language_model: LanguageModelInfo,
     event: ChatEvent,
     logger: Logger,
+    invocation_stats: list[LanguageModelInvocationStats] | None = None,
 ) -> bool:
     """Cheaply decide whether the turn warrants a full memory rewrite.
 
@@ -664,6 +688,15 @@ async def should_consolidate_memory(
         )
         return True
 
+    if invocation_stats is not None and response.usage is not None:
+        invocation_stats.append(
+            LanguageModelInvocationStats.from_usage(
+                language_model.name,
+                response.usage,
+                source="user_memory_gate",
+            )
+        )
+
     try:
         raw = response.choices[0].message.content or ""
     except Exception as exc:
@@ -702,6 +735,7 @@ async def consolidate_user_memory(
     logger: Logger,
     on_update_start: Callable[[], Awaitable[None]] = noop_update_callback,
     on_update_end: Callable[[], Awaitable[None]] = noop_update_callback,
+    invocation_stats: list[LanguageModelInvocationStats] | None = None,
 ) -> str:
     """Consolidate the latest turn into the user's memory profile.
 
@@ -739,6 +773,7 @@ async def consolidate_user_memory(
         language_model=language_model,
         event=event,
         logger=logger,
+        invocation_stats=invocation_stats,
     ):
         return safe_current
 
@@ -753,6 +788,7 @@ async def consolidate_user_memory(
             language_model=language_model,
             event=event,
             logger=logger,
+            invocation_stats=invocation_stats,
         )
     finally:
         await on_update_end()
@@ -768,6 +804,7 @@ async def _rewrite_user_memory(
     language_model: LanguageModelInfo,
     event: ChatEvent,
     logger: Logger,
+    invocation_stats: list[LanguageModelInvocationStats] | None = None,
 ) -> str:
     if not safe_current.strip():
         safe_current = empty_profile(user_id)
@@ -816,6 +853,15 @@ async def _rewrite_user_memory(
         )
         return safe_current
 
+    if invocation_stats is not None and response.usage is not None:
+        invocation_stats.append(
+            LanguageModelInvocationStats.from_usage(
+                language_model.name,
+                response.usage,
+                source="user_memory_consolidation",
+            )
+        )
+
     try:
         raw = response.choices[0].message.content or ""
     except Exception as exc:
@@ -860,6 +906,8 @@ async def _rewrite_user_memory(
         language_model=language_model,
         event=event,
         logger=logger,
+        invocation_stats=invocation_stats,
+        invocation_source="user_memory_post_consolidation_condense",
     )
     logger.info(
         "[user-memory] consolidation produced %d tokens (cap=%d)",

{unique_user_memory-2026.32.0.dev1 → unique_user_memory-2026.32.0.dev2}/unique_user_memory/user_memory_postprocessor.py

@@ -9,6 +9,7 @@ from unique_toolkit.language_model.default_language_model import (
     DEFAULT_LANGUAGE_MODEL,
 )
 from unique_toolkit.language_model.infos import LanguageModelInfo
+from unique_toolkit.language_model.invocation_stats import LanguageModelInvocationStats
 from unique_toolkit.language_model.schemas import LanguageModelStreamResponse
 
 from unique_user_memory.config import UserMemoryConfig
@@ -49,6 +50,28 @@ class UserMemoryPostprocessor(Postprocessor):
         self._logger = logger
         self._new_memory: str | None = None
         self._chat_service: ChatService = chat_service
+        self._pending_load_invocation_stats = list(state.load_invocation_stats)
+        self._invocation_stats: list[LanguageModelInvocationStats] = []
+
+    @property
+    def invocation_stats(self) -> list[LanguageModelInvocationStats]:
+        return list(self._invocation_stats)
+
+    def take_pending_invocation_stats(self) -> list[LanguageModelInvocationStats]:
+        """Pop load-time condense stats not yet reported.
+
+        `UniqueAI` calls this unconditionally at the start of every turn so a
+        turn that exits before `run()` (cancellation, empty response, a
+        control-taking tool) still reports the tokens spent condensing the
+        loaded profile. If `run()` does execute, it drains the same pending
+        list itself, so whichever of the two runs first "wins" and the other
+        sees an empty list -- the tokens are never double-counted or lost.
+        """
+        stats, self._pending_load_invocation_stats = (
+            self._pending_load_invocation_stats,
+            [],
+        )
+        return stats
 
     async def run(self, loop_response: LanguageModelStreamResponse) -> bool:
         """Consolidate and upload user memory for this turn.
@@ -56,6 +79,7 @@ class UserMemoryPostprocessor(Postprocessor):
         Returns True if the memory profile changed and was uploaded, False
         otherwise (no user/company, NOOP consolidation, or failed upload).
         """
+        self._invocation_stats = self.take_pending_invocation_stats()
         self._logger.info("[user-memory] running postprocessor")
         user_id = self._event.user_id
         company_id = self._event.company_id
@@ -99,6 +123,7 @@ class UserMemoryPostprocessor(Postprocessor):
             logger=self._logger,
             on_update_start=on_update_start,
             on_update_end=on_update_end,
+            invocation_stats=self._invocation_stats,
         )
 
         if self._new_memory == self._state.text: