unique-user-memory 2026.24.0__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.
- unique_user_memory-2026.24.0/PKG-INFO +17 -0
- unique_user_memory-2026.24.0/README.md +3 -0
- unique_user_memory-2026.24.0/pyproject.toml +83 -0
- unique_user_memory-2026.24.0/unique_user_memory/__init__.py +4 -0
- unique_user_memory-2026.24.0/unique_user_memory/config.py +32 -0
- unique_user_memory-2026.24.0/unique_user_memory/tests/test_user_memory.py +434 -0
- unique_user_memory-2026.24.0/unique_user_memory/user_memory.py +548 -0
- unique_user_memory-2026.24.0/unique_user_memory/user_memory_postprocessor.py +71 -0
- unique_user_memory-2026.24.0/unique_user_memory/user_memory_prompts.py +185 -0
|
@@ -0,0 +1,17 @@
|
|
|
1
|
+
Metadata-Version: 2.3
|
|
2
|
+
Name: unique-user-memory
|
|
3
|
+
Version: 2026.24.0
|
|
4
|
+
Summary:
|
|
5
|
+
Author: Fabian Schläpfer
|
|
6
|
+
Author-email: Fabian Schläpfer <fabian@unique.ch>
|
|
7
|
+
License: Proprietary
|
|
8
|
+
Requires-Dist: jinja2>=3.1.6
|
|
9
|
+
Requires-Dist: pydantic>=2.8.2
|
|
10
|
+
Requires-Dist: unique-sdk>=2026.24.0
|
|
11
|
+
Requires-Dist: unique-toolkit>=2026.24.0
|
|
12
|
+
Requires-Python: >=3.12, <4
|
|
13
|
+
Description-Content-Type: text/markdown
|
|
14
|
+
|
|
15
|
+
# Unique User Memory
|
|
16
|
+
|
|
17
|
+
Postprocessor package for persistent per-user memory in Unique AI.
|
|
@@ -0,0 +1,83 @@
|
|
|
1
|
+
[project]
|
|
2
|
+
name = "unique_user_memory"
|
|
3
|
+
version = "2026.24.0"
|
|
4
|
+
description = ""
|
|
5
|
+
authors = [
|
|
6
|
+
{ name = "Fabian Schläpfer", email = "fabian@unique.ch" },
|
|
7
|
+
]
|
|
8
|
+
readme = "README.md"
|
|
9
|
+
license = { text = "Proprietary" }
|
|
10
|
+
requires-python = ">=3.12,<4"
|
|
11
|
+
dependencies = [
|
|
12
|
+
"jinja2>=3.1.6",
|
|
13
|
+
"pydantic>=2.8.2",
|
|
14
|
+
"unique-sdk>=2026.24.0",
|
|
15
|
+
"unique-toolkit>=2026.24.0",
|
|
16
|
+
]
|
|
17
|
+
|
|
18
|
+
[dependency-groups]
|
|
19
|
+
dev = []
|
|
20
|
+
|
|
21
|
+
[build-system]
|
|
22
|
+
requires = ["uv_build>=0.7.19,<0.8"]
|
|
23
|
+
build-backend = "uv_build"
|
|
24
|
+
|
|
25
|
+
[tool.uv.build-backend]
|
|
26
|
+
module-root = ""
|
|
27
|
+
|
|
28
|
+
[tool.uv]
|
|
29
|
+
exclude-newer = "2 weeks"
|
|
30
|
+
constraint-dependencies = ["lxml>=5.0.0"]
|
|
31
|
+
|
|
32
|
+
[tool.uv.exclude-newer-package]
|
|
33
|
+
"unique-sdk" = false
|
|
34
|
+
"unique-toolkit" = false
|
|
35
|
+
|
|
36
|
+
[tool.uv.sources]
|
|
37
|
+
unique-sdk = { workspace = true }
|
|
38
|
+
unique-toolkit = { workspace = true }
|
|
39
|
+
|
|
40
|
+
[tool.poe.tasks]
|
|
41
|
+
lint = "ruff check ."
|
|
42
|
+
lint-fix = "ruff check . --fix"
|
|
43
|
+
format = "ruff format ."
|
|
44
|
+
test = "pytest"
|
|
45
|
+
typecheck = "basedpyright"
|
|
46
|
+
depcheck = "deptry ."
|
|
47
|
+
coverage = "pytest --cov=unique_user_memory --cov-report=term-missing"
|
|
48
|
+
ci-typecheck = { shell = "bash $(git rev-parse --show-toplevel)/.github/scripts/dev.sh typecheck" }
|
|
49
|
+
ci-coverage = { shell = "bash $(git rev-parse --show-toplevel)/.github/scripts/dev.sh coverage" }
|
|
50
|
+
|
|
51
|
+
[tool.pytest.ini_options]
|
|
52
|
+
addopts = "--strict-markers --import-mode=importlib"
|
|
53
|
+
asyncio_mode = "auto"
|
|
54
|
+
markers = [
|
|
55
|
+
"ai: AI-authored or AI-generated tests",
|
|
56
|
+
"asyncio: asyncio tests",
|
|
57
|
+
"integration: integration tests that require API access or credentials",
|
|
58
|
+
"serial: tests that must run serially",
|
|
59
|
+
"unit: unit tests",
|
|
60
|
+
"verified: AI-generated tests with human verification",
|
|
61
|
+
]
|
|
62
|
+
filterwarnings = [
|
|
63
|
+
"ignore::DeprecationWarning",
|
|
64
|
+
]
|
|
65
|
+
|
|
66
|
+
[tool.deptry]
|
|
67
|
+
known_first_party = ["unique_user_memory"]
|
|
68
|
+
extend_exclude = ["unique_user_memory/tests"]
|
|
69
|
+
|
|
70
|
+
[tool.deptry.per_rule_ignores]
|
|
71
|
+
DEP002 = ["unique-sdk", "unique-toolkit"]
|
|
72
|
+
DEP003 = ["unique_sdk", "unique_toolkit"]
|
|
73
|
+
|
|
74
|
+
[tool.basedpyright]
|
|
75
|
+
typeCheckingMode = "standard"
|
|
76
|
+
include = ["unique_user_memory"]
|
|
77
|
+
exclude = ["**/tests/**", "**/test_*.py"]
|
|
78
|
+
|
|
79
|
+
[tool.ruff]
|
|
80
|
+
target-version = "py311"
|
|
81
|
+
|
|
82
|
+
[tool.ruff.lint]
|
|
83
|
+
extend-select = ["I"]
|
|
@@ -0,0 +1,32 @@
|
|
|
1
|
+
from typing import Annotated
|
|
2
|
+
|
|
3
|
+
from pydantic import BaseModel, Field
|
|
4
|
+
from unique_toolkit._common.pydantic.rjsf_tags import RJSFMetaTag
|
|
5
|
+
from unique_toolkit._common.validators import LMI
|
|
6
|
+
from unique_toolkit.agentic.tools.config import get_configuration_dict
|
|
7
|
+
from unique_toolkit.language_model.default_language_model import DEFAULT_GPT_4o
|
|
8
|
+
from unique_toolkit.language_model.infos import LanguageModelInfo
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
class UserMemoryConfig(BaseModel):
|
|
12
|
+
model_config = get_configuration_dict()
|
|
13
|
+
|
|
14
|
+
enabled: Annotated[bool, RJSFMetaTag.SpecialWidget.hidden()] = Field(
|
|
15
|
+
default=False,
|
|
16
|
+
description="Whether to enable persistent per-user memory.",
|
|
17
|
+
)
|
|
18
|
+
language_model: LMI = Field(
|
|
19
|
+
default=LanguageModelInfo.from_name(DEFAULT_GPT_4o),
|
|
20
|
+
description="The language model used for post-turn memory consolidation.",
|
|
21
|
+
)
|
|
22
|
+
max_tokens: int = Field(
|
|
23
|
+
default=2000,
|
|
24
|
+
ge=500,
|
|
25
|
+
le=8000,
|
|
26
|
+
description="Maximum size of the memory profile in tokens.",
|
|
27
|
+
)
|
|
28
|
+
root_folder: Annotated[str, RJSFMetaTag.SpecialWidget.hidden()] = Field(
|
|
29
|
+
default="user-memory",
|
|
30
|
+
min_length=1,
|
|
31
|
+
description="Root KB folder used to store per-user memory profiles.",
|
|
32
|
+
)
|
|
@@ -0,0 +1,434 @@
|
|
|
1
|
+
from unittest.mock import AsyncMock, MagicMock
|
|
2
|
+
|
|
3
|
+
import pytest
|
|
4
|
+
|
|
5
|
+
from unique_user_memory.config import UserMemoryConfig
|
|
6
|
+
from unique_user_memory.user_memory import (
|
|
7
|
+
UserMemoryState,
|
|
8
|
+
_sanitize_for_xml_context,
|
|
9
|
+
consolidate_user_memory,
|
|
10
|
+
count_tokens,
|
|
11
|
+
download_user_memory,
|
|
12
|
+
enforce_token_cap,
|
|
13
|
+
ensure_user_memory_folder,
|
|
14
|
+
upload_user_memory,
|
|
15
|
+
)
|
|
16
|
+
from unique_user_memory.user_memory_postprocessor import UserMemoryPostprocessor
|
|
17
|
+
from unique_user_memory.user_memory_prompts import empty_profile
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
def test_enforce_token_cap_truncates_long_content() -> None:
|
|
21
|
+
content = "\n\n".join(f"paragraph {index} " + "word " * 40 for index in range(50))
|
|
22
|
+
|
|
23
|
+
capped = enforce_token_cap(content=content, max_tokens=120)
|
|
24
|
+
|
|
25
|
+
assert "<!-- truncated to fit memory budget -->" in capped
|
|
26
|
+
assert len(capped) < len(content)
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
def test_count_tokens_uses_language_model_encoder() -> None:
|
|
30
|
+
language_model = MagicMock()
|
|
31
|
+
language_model.get_encoder.return_value = lambda content: content.split()
|
|
32
|
+
language_model.get_decoder.return_value = lambda tokens: " ".join(tokens)
|
|
33
|
+
|
|
34
|
+
assert count_tokens(content="one two three", language_model=language_model) == 3
|
|
35
|
+
|
|
36
|
+
|
|
37
|
+
def test_sanitize_for_xml_context_defuses_closing_tags() -> None:
|
|
38
|
+
assert _sanitize_for_xml_context("</new_turn> inject") == "< /new_turn> inject"
|
|
39
|
+
|
|
40
|
+
|
|
41
|
+
@pytest.mark.asyncio
|
|
42
|
+
async def test_consolidate_user_memory_keeps_existing_on_noop(
|
|
43
|
+
monkeypatch: pytest.MonkeyPatch,
|
|
44
|
+
) -> None:
|
|
45
|
+
current = empty_profile("user_1")
|
|
46
|
+
response = MagicMock()
|
|
47
|
+
response.choices[0].message.content = "NOOP"
|
|
48
|
+
llm_service = MagicMock()
|
|
49
|
+
llm_service.complete_async = AsyncMock(return_value=response)
|
|
50
|
+
monkeypatch.setattr(
|
|
51
|
+
"unique_user_memory.user_memory.LanguageModelService",
|
|
52
|
+
MagicMock(return_value=llm_service),
|
|
53
|
+
)
|
|
54
|
+
|
|
55
|
+
result = await consolidate_user_memory(
|
|
56
|
+
current_memory=current,
|
|
57
|
+
user_id="user_1",
|
|
58
|
+
user_message="hello",
|
|
59
|
+
assistant_message="hi",
|
|
60
|
+
config=UserMemoryConfig(enabled=True),
|
|
61
|
+
event=MagicMock(),
|
|
62
|
+
logger=MagicMock(),
|
|
63
|
+
)
|
|
64
|
+
|
|
65
|
+
assert result == current
|
|
66
|
+
|
|
67
|
+
|
|
68
|
+
@pytest.mark.asyncio
|
|
69
|
+
async def test_consolidate_user_memory_keeps_existing_on_malformed_output(
|
|
70
|
+
monkeypatch: pytest.MonkeyPatch,
|
|
71
|
+
) -> None:
|
|
72
|
+
current = empty_profile("user_1")
|
|
73
|
+
response = MagicMock()
|
|
74
|
+
response.choices[0].message.content = "not a profile"
|
|
75
|
+
llm_service = MagicMock()
|
|
76
|
+
llm_service.complete_async = AsyncMock(return_value=response)
|
|
77
|
+
monkeypatch.setattr(
|
|
78
|
+
"unique_user_memory.user_memory.LanguageModelService",
|
|
79
|
+
MagicMock(return_value=llm_service),
|
|
80
|
+
)
|
|
81
|
+
|
|
82
|
+
result = await consolidate_user_memory(
|
|
83
|
+
current_memory=current,
|
|
84
|
+
user_id="user_1",
|
|
85
|
+
user_message="remember I like concise answers",
|
|
86
|
+
assistant_message="noted",
|
|
87
|
+
config=UserMemoryConfig(enabled=True),
|
|
88
|
+
event=MagicMock(),
|
|
89
|
+
logger=MagicMock(),
|
|
90
|
+
)
|
|
91
|
+
|
|
92
|
+
assert result == current
|
|
93
|
+
|
|
94
|
+
|
|
95
|
+
@pytest.mark.asyncio
|
|
96
|
+
async def test_download_user_memory_returns_empty_when_file_missing(
|
|
97
|
+
monkeypatch: pytest.MonkeyPatch,
|
|
98
|
+
) -> None:
|
|
99
|
+
search_contents = AsyncMock(return_value=[])
|
|
100
|
+
monkeypatch.setattr(
|
|
101
|
+
"unique_user_memory.user_memory.search_contents_async",
|
|
102
|
+
search_contents,
|
|
103
|
+
)
|
|
104
|
+
|
|
105
|
+
result = await download_user_memory(
|
|
106
|
+
scope_id="scope_1",
|
|
107
|
+
user_id="user_1",
|
|
108
|
+
company_id="company_1",
|
|
109
|
+
logger=MagicMock(),
|
|
110
|
+
)
|
|
111
|
+
|
|
112
|
+
assert result == ""
|
|
113
|
+
search_contents.assert_awaited_once_with(
|
|
114
|
+
user_id="user_1",
|
|
115
|
+
company_id="company_1",
|
|
116
|
+
chat_id=None,
|
|
117
|
+
where={"ownerId": {"equals": "scope_1"}},
|
|
118
|
+
)
|
|
119
|
+
|
|
120
|
+
|
|
121
|
+
@pytest.mark.asyncio
|
|
122
|
+
async def test_download_user_memory_downloads_existing_file_to_memory(
|
|
123
|
+
monkeypatch: pytest.MonkeyPatch,
|
|
124
|
+
) -> None:
|
|
125
|
+
content = MagicMock()
|
|
126
|
+
content.id = "content_1"
|
|
127
|
+
content.key = "memory.md"
|
|
128
|
+
search_contents = AsyncMock(return_value=[content])
|
|
129
|
+
download_content = AsyncMock(return_value=b"# User Memory\n\n## Identity\n- Test")
|
|
130
|
+
monkeypatch.setattr(
|
|
131
|
+
"unique_user_memory.user_memory.search_contents_async",
|
|
132
|
+
search_contents,
|
|
133
|
+
)
|
|
134
|
+
monkeypatch.setattr(
|
|
135
|
+
"unique_user_memory.user_memory.download_content_to_bytes_async",
|
|
136
|
+
download_content,
|
|
137
|
+
)
|
|
138
|
+
|
|
139
|
+
result = await download_user_memory(
|
|
140
|
+
scope_id="scope_1",
|
|
141
|
+
user_id="user_1",
|
|
142
|
+
company_id="company_1",
|
|
143
|
+
logger=MagicMock(),
|
|
144
|
+
)
|
|
145
|
+
|
|
146
|
+
assert result == "# User Memory\n\n## Identity\n- Test"
|
|
147
|
+
search_contents.assert_awaited_once_with(
|
|
148
|
+
user_id="user_1",
|
|
149
|
+
company_id="company_1",
|
|
150
|
+
chat_id=None,
|
|
151
|
+
where={"ownerId": {"equals": "scope_1"}},
|
|
152
|
+
)
|
|
153
|
+
download_content.assert_awaited_once_with(
|
|
154
|
+
user_id="user_1",
|
|
155
|
+
company_id="company_1",
|
|
156
|
+
content_id="content_1",
|
|
157
|
+
chat_id=None,
|
|
158
|
+
)
|
|
159
|
+
|
|
160
|
+
|
|
161
|
+
@pytest.mark.asyncio
|
|
162
|
+
async def test_ensure_user_memory_folder_returns_existing_user_folder(
|
|
163
|
+
monkeypatch: pytest.MonkeyPatch,
|
|
164
|
+
) -> None:
|
|
165
|
+
get_info = AsyncMock(
|
|
166
|
+
side_effect=[
|
|
167
|
+
{"id": "scope_root"},
|
|
168
|
+
{"id": "scope_user"},
|
|
169
|
+
]
|
|
170
|
+
)
|
|
171
|
+
create_paths = AsyncMock()
|
|
172
|
+
add_access = AsyncMock()
|
|
173
|
+
get_groups = AsyncMock()
|
|
174
|
+
monkeypatch.setattr(
|
|
175
|
+
"unique_user_memory.user_memory.unique_sdk.Folder.get_info_async",
|
|
176
|
+
get_info,
|
|
177
|
+
)
|
|
178
|
+
monkeypatch.setattr(
|
|
179
|
+
"unique_user_memory.user_memory.unique_sdk.Folder.create_paths_async",
|
|
180
|
+
create_paths,
|
|
181
|
+
)
|
|
182
|
+
monkeypatch.setattr(
|
|
183
|
+
"unique_user_memory.user_memory.unique_sdk.Folder.add_access_async",
|
|
184
|
+
add_access,
|
|
185
|
+
)
|
|
186
|
+
monkeypatch.setattr(
|
|
187
|
+
"unique_user_memory.user_memory.unique_sdk.Group.get_groups_async",
|
|
188
|
+
get_groups,
|
|
189
|
+
)
|
|
190
|
+
|
|
191
|
+
result = await ensure_user_memory_folder(
|
|
192
|
+
user_id="user_1",
|
|
193
|
+
company_id="company_1",
|
|
194
|
+
root_folder="user-memory",
|
|
195
|
+
logger=MagicMock(),
|
|
196
|
+
)
|
|
197
|
+
|
|
198
|
+
assert result == "scope_user"
|
|
199
|
+
get_info.assert_any_await(
|
|
200
|
+
user_id="user_1",
|
|
201
|
+
company_id="company_1",
|
|
202
|
+
folderPath="/user-memory",
|
|
203
|
+
)
|
|
204
|
+
get_info.assert_any_await(
|
|
205
|
+
user_id="user_1",
|
|
206
|
+
company_id="company_1",
|
|
207
|
+
folderPath="/user-memory/user_1",
|
|
208
|
+
)
|
|
209
|
+
create_paths.assert_not_awaited()
|
|
210
|
+
get_groups.assert_not_awaited()
|
|
211
|
+
add_access.assert_not_awaited()
|
|
212
|
+
|
|
213
|
+
|
|
214
|
+
@pytest.mark.asyncio
|
|
215
|
+
async def test_ensure_user_memory_folder_creates_private_user_folder_under_root(
|
|
216
|
+
monkeypatch: pytest.MonkeyPatch,
|
|
217
|
+
) -> None:
|
|
218
|
+
get_info = AsyncMock(
|
|
219
|
+
side_effect=[
|
|
220
|
+
{"id": "scope_root"},
|
|
221
|
+
RuntimeError("missing user folder"),
|
|
222
|
+
]
|
|
223
|
+
)
|
|
224
|
+
create_paths = AsyncMock(return_value={"createdFolders": [{"id": "scope_user"}]})
|
|
225
|
+
add_access = AsyncMock()
|
|
226
|
+
monkeypatch.setattr(
|
|
227
|
+
"unique_user_memory.user_memory.unique_sdk.Folder.get_info_async",
|
|
228
|
+
get_info,
|
|
229
|
+
)
|
|
230
|
+
monkeypatch.setattr(
|
|
231
|
+
"unique_user_memory.user_memory.unique_sdk.Folder.create_paths_async",
|
|
232
|
+
create_paths,
|
|
233
|
+
)
|
|
234
|
+
monkeypatch.setattr(
|
|
235
|
+
"unique_user_memory.user_memory.unique_sdk.Folder.add_access_async",
|
|
236
|
+
add_access,
|
|
237
|
+
)
|
|
238
|
+
|
|
239
|
+
result = await ensure_user_memory_folder(
|
|
240
|
+
user_id="user_1",
|
|
241
|
+
company_id="company_1",
|
|
242
|
+
root_folder="user-memory",
|
|
243
|
+
logger=MagicMock(),
|
|
244
|
+
)
|
|
245
|
+
|
|
246
|
+
assert result == "scope_user"
|
|
247
|
+
create_paths.assert_awaited_once_with(
|
|
248
|
+
user_id="user_1",
|
|
249
|
+
company_id="company_1",
|
|
250
|
+
parentScopeId="scope_root",
|
|
251
|
+
relativePaths=["user_1"],
|
|
252
|
+
inheritAccess=False,
|
|
253
|
+
)
|
|
254
|
+
add_access.assert_awaited_once_with(
|
|
255
|
+
user_id="user_1",
|
|
256
|
+
company_id="company_1",
|
|
257
|
+
scopeId="scope_user",
|
|
258
|
+
scopeAccesses=[
|
|
259
|
+
{"entityId": "user_1", "type": "READ", "entityType": "USER"},
|
|
260
|
+
{"entityId": "user_1", "type": "WRITE", "entityType": "USER"},
|
|
261
|
+
],
|
|
262
|
+
applyToSubScopes=True,
|
|
263
|
+
)
|
|
264
|
+
|
|
265
|
+
|
|
266
|
+
@pytest.mark.asyncio
|
|
267
|
+
async def test_ensure_user_memory_folder_returns_none_when_access_grant_fails_after_create(
|
|
268
|
+
monkeypatch: pytest.MonkeyPatch,
|
|
269
|
+
) -> None:
|
|
270
|
+
get_info = AsyncMock(
|
|
271
|
+
side_effect=[
|
|
272
|
+
{"id": "scope_root"},
|
|
273
|
+
RuntimeError("missing user folder"),
|
|
274
|
+
]
|
|
275
|
+
)
|
|
276
|
+
create_paths = AsyncMock(return_value={"createdFolders": [{"id": "scope_user"}]})
|
|
277
|
+
grant_error = RuntimeError("grant failed")
|
|
278
|
+
add_access = AsyncMock(side_effect=grant_error)
|
|
279
|
+
logger = MagicMock()
|
|
280
|
+
monkeypatch.setattr(
|
|
281
|
+
"unique_user_memory.user_memory.unique_sdk.Folder.get_info_async",
|
|
282
|
+
get_info,
|
|
283
|
+
)
|
|
284
|
+
monkeypatch.setattr(
|
|
285
|
+
"unique_user_memory.user_memory.unique_sdk.Folder.create_paths_async",
|
|
286
|
+
create_paths,
|
|
287
|
+
)
|
|
288
|
+
monkeypatch.setattr(
|
|
289
|
+
"unique_user_memory.user_memory.unique_sdk.Folder.add_access_async",
|
|
290
|
+
add_access,
|
|
291
|
+
)
|
|
292
|
+
|
|
293
|
+
result = await ensure_user_memory_folder(
|
|
294
|
+
user_id="user_1",
|
|
295
|
+
company_id="company_1",
|
|
296
|
+
root_folder="user-memory",
|
|
297
|
+
logger=logger,
|
|
298
|
+
)
|
|
299
|
+
|
|
300
|
+
assert result is None
|
|
301
|
+
create_paths.assert_awaited_once_with(
|
|
302
|
+
user_id="user_1",
|
|
303
|
+
company_id="company_1",
|
|
304
|
+
parentScopeId="scope_root",
|
|
305
|
+
relativePaths=["user_1"],
|
|
306
|
+
inheritAccess=False,
|
|
307
|
+
)
|
|
308
|
+
add_access.assert_awaited_once_with(
|
|
309
|
+
user_id="user_1",
|
|
310
|
+
company_id="company_1",
|
|
311
|
+
scopeId="scope_user",
|
|
312
|
+
scopeAccesses=[
|
|
313
|
+
{"entityId": "user_1", "type": "READ", "entityType": "USER"},
|
|
314
|
+
{"entityId": "user_1", "type": "WRITE", "entityType": "USER"},
|
|
315
|
+
],
|
|
316
|
+
applyToSubScopes=True,
|
|
317
|
+
)
|
|
318
|
+
logger.warning.assert_called_with(
|
|
319
|
+
"[user-memory] failed to grant read/write access on scope %s "
|
|
320
|
+
"for user %s: [%s] %s",
|
|
321
|
+
"scope_user",
|
|
322
|
+
"user_1",
|
|
323
|
+
"RuntimeError",
|
|
324
|
+
grant_error,
|
|
325
|
+
)
|
|
326
|
+
|
|
327
|
+
|
|
328
|
+
@pytest.mark.asyncio
|
|
329
|
+
async def test_upload_user_memory_writes_hidden_skip_ingestion_file(
|
|
330
|
+
monkeypatch: pytest.MonkeyPatch,
|
|
331
|
+
) -> None:
|
|
332
|
+
upload_content = AsyncMock()
|
|
333
|
+
monkeypatch.setattr(
|
|
334
|
+
"unique_user_memory.user_memory.upload_content_from_bytes_async",
|
|
335
|
+
upload_content,
|
|
336
|
+
)
|
|
337
|
+
|
|
338
|
+
result = await upload_user_memory(
|
|
339
|
+
scope_id="scope_1",
|
|
340
|
+
content="# User Memory\n\n## Identity\n- Test",
|
|
341
|
+
user_id="user_1",
|
|
342
|
+
company_id="company_1",
|
|
343
|
+
logger=MagicMock(),
|
|
344
|
+
)
|
|
345
|
+
|
|
346
|
+
assert result is True
|
|
347
|
+
upload_content.assert_awaited_once()
|
|
348
|
+
assert upload_content.call_args.kwargs["content"] == (
|
|
349
|
+
b"# User Memory\n\n## Identity\n- Test"
|
|
350
|
+
)
|
|
351
|
+
assert upload_content.call_args.kwargs["content_name"] == "memory.md"
|
|
352
|
+
assert upload_content.call_args.kwargs["mime_type"] == "text/markdown"
|
|
353
|
+
assert upload_content.call_args.kwargs["scope_id"] == "scope_1"
|
|
354
|
+
assert upload_content.call_args.kwargs["ingestion_config"] == {
|
|
355
|
+
"uniqueIngestionMode": "SKIP_INGESTION",
|
|
356
|
+
"hideInChat": True,
|
|
357
|
+
}
|
|
358
|
+
|
|
359
|
+
|
|
360
|
+
@pytest.mark.asyncio
|
|
361
|
+
async def test_user_memory_postprocessor_logs_success_when_upload_succeeds(
|
|
362
|
+
monkeypatch: pytest.MonkeyPatch,
|
|
363
|
+
) -> None:
|
|
364
|
+
updated_memory = "# User Memory\n\n## Identity\n- Updated"
|
|
365
|
+
consolidate = AsyncMock(return_value=updated_memory)
|
|
366
|
+
upload = AsyncMock(return_value=True)
|
|
367
|
+
monkeypatch.setattr(
|
|
368
|
+
"unique_user_memory.user_memory_postprocessor.consolidate_user_memory",
|
|
369
|
+
consolidate,
|
|
370
|
+
)
|
|
371
|
+
monkeypatch.setattr(
|
|
372
|
+
"unique_user_memory.user_memory_postprocessor.upload_user_memory",
|
|
373
|
+
upload,
|
|
374
|
+
)
|
|
375
|
+
event = MagicMock()
|
|
376
|
+
event.user_id = "user_1"
|
|
377
|
+
event.company_id = "company_1"
|
|
378
|
+
event.payload.user_message.text = "remember this"
|
|
379
|
+
loop_response = MagicMock()
|
|
380
|
+
loop_response.message.text = "noted"
|
|
381
|
+
logger = MagicMock()
|
|
382
|
+
postprocessor = UserMemoryPostprocessor(
|
|
383
|
+
config=UserMemoryConfig(enabled=True),
|
|
384
|
+
event=event,
|
|
385
|
+
state=UserMemoryState(scope_id="scope_1", text=empty_profile("user_1")),
|
|
386
|
+
logger=logger,
|
|
387
|
+
)
|
|
388
|
+
|
|
389
|
+
await postprocessor.run(loop_response)
|
|
390
|
+
|
|
391
|
+
upload.assert_awaited_once_with(
|
|
392
|
+
scope_id="scope_1",
|
|
393
|
+
content=updated_memory,
|
|
394
|
+
user_id="user_1",
|
|
395
|
+
company_id="company_1",
|
|
396
|
+
logger=logger,
|
|
397
|
+
)
|
|
398
|
+
logger.info.assert_any_call("[user-memory] memory updated and uploaded")
|
|
399
|
+
|
|
400
|
+
|
|
401
|
+
@pytest.mark.asyncio
|
|
402
|
+
async def test_user_memory_postprocessor_does_not_log_success_when_upload_fails(
|
|
403
|
+
monkeypatch: pytest.MonkeyPatch,
|
|
404
|
+
) -> None:
|
|
405
|
+
updated_memory = "# User Memory\n\n## Identity\n- Updated"
|
|
406
|
+
monkeypatch.setattr(
|
|
407
|
+
"unique_user_memory.user_memory_postprocessor.consolidate_user_memory",
|
|
408
|
+
AsyncMock(return_value=updated_memory),
|
|
409
|
+
)
|
|
410
|
+
monkeypatch.setattr(
|
|
411
|
+
"unique_user_memory.user_memory_postprocessor.upload_user_memory",
|
|
412
|
+
AsyncMock(return_value=False),
|
|
413
|
+
)
|
|
414
|
+
event = MagicMock()
|
|
415
|
+
event.user_id = "user_1"
|
|
416
|
+
event.company_id = "company_1"
|
|
417
|
+
event.payload.user_message.text = "remember this"
|
|
418
|
+
loop_response = MagicMock()
|
|
419
|
+
loop_response.message.text = "noted"
|
|
420
|
+
logger = MagicMock()
|
|
421
|
+
postprocessor = UserMemoryPostprocessor(
|
|
422
|
+
config=UserMemoryConfig(enabled=True),
|
|
423
|
+
event=event,
|
|
424
|
+
state=UserMemoryState(scope_id="scope_1", text=empty_profile("user_1")),
|
|
425
|
+
logger=logger,
|
|
426
|
+
)
|
|
427
|
+
|
|
428
|
+
await postprocessor.run(loop_response)
|
|
429
|
+
|
|
430
|
+
assert all(
|
|
431
|
+
call.args != ("[user-memory] memory updated and uploaded",)
|
|
432
|
+
for call in logger.info.call_args_list
|
|
433
|
+
)
|
|
434
|
+
logger.warning.assert_any_call("[user-memory] memory update was not uploaded")
|
|
@@ -0,0 +1,548 @@
|
|
|
1
|
+
import re
|
|
2
|
+
from dataclasses import dataclass
|
|
3
|
+
from logging import Logger
|
|
4
|
+
|
|
5
|
+
import unique_sdk
|
|
6
|
+
from unique_toolkit.app import ChatEvent
|
|
7
|
+
from unique_toolkit.content.functions import (
|
|
8
|
+
download_content_to_bytes_async,
|
|
9
|
+
search_contents_async,
|
|
10
|
+
upload_content_from_bytes_async,
|
|
11
|
+
)
|
|
12
|
+
from unique_toolkit.language_model import (
|
|
13
|
+
DEFAULT_GPT_4o,
|
|
14
|
+
LanguageModelMessages,
|
|
15
|
+
LanguageModelService,
|
|
16
|
+
LanguageModelSystemMessage,
|
|
17
|
+
LanguageModelUserMessage,
|
|
18
|
+
TypeDecoder,
|
|
19
|
+
TypeEncoder,
|
|
20
|
+
)
|
|
21
|
+
from unique_toolkit.language_model.infos import LanguageModelInfo
|
|
22
|
+
|
|
23
|
+
from unique_user_memory.config import UserMemoryConfig
|
|
24
|
+
from unique_user_memory.user_memory_prompts import (
|
|
25
|
+
SECTION_HEADINGS,
|
|
26
|
+
consolidation_system_prompt,
|
|
27
|
+
consolidation_user_prompt,
|
|
28
|
+
empty_profile,
|
|
29
|
+
)
|
|
30
|
+
|
|
31
|
+
MEMORY_FILENAME = "memory.md"
|
|
32
|
+
MIME_TYPE = "text/markdown"
|
|
33
|
+
_LLM_OUTPUT_HEADROOM_TOKENS = 200
|
|
34
|
+
_TRUNCATION_MARKER = "\n\n<!-- truncated to fit memory budget -->"
|
|
35
|
+
_DEFAULT_LANGUAGE_MODEL = LanguageModelInfo.from_name(DEFAULT_GPT_4o)
|
|
36
|
+
_FRONTMATTER_RE = re.compile(r"^---\n.*?\n---\n", re.DOTALL)
|
|
37
|
+
|
|
38
|
+
|
|
39
|
+
@dataclass(frozen=True)
|
|
40
|
+
class UserMemoryState:
|
|
41
|
+
scope_id: str
|
|
42
|
+
text: str
|
|
43
|
+
|
|
44
|
+
|
|
45
|
+
def _get_model_tokenizer(
|
|
46
|
+
*,
|
|
47
|
+
language_model: LanguageModelInfo,
|
|
48
|
+
) -> tuple[TypeEncoder, TypeDecoder]:
|
|
49
|
+
model_info = language_model or _DEFAULT_LANGUAGE_MODEL
|
|
50
|
+
return model_info.get_encoder(), model_info.get_decoder()
|
|
51
|
+
|
|
52
|
+
|
|
53
|
+
def _count_tokens(
|
|
54
|
+
*,
|
|
55
|
+
content: str,
|
|
56
|
+
encode: TypeEncoder,
|
|
57
|
+
) -> int:
|
|
58
|
+
if not content:
|
|
59
|
+
return 0
|
|
60
|
+
return len(encode(content))
|
|
61
|
+
|
|
62
|
+
|
|
63
|
+
def count_tokens(
|
|
64
|
+
*,
|
|
65
|
+
content: str,
|
|
66
|
+
language_model: LanguageModelInfo = _DEFAULT_LANGUAGE_MODEL,
|
|
67
|
+
) -> int:
|
|
68
|
+
encode, _ = _get_model_tokenizer(language_model=language_model)
|
|
69
|
+
return _count_tokens(content=content, encode=encode)
|
|
70
|
+
|
|
71
|
+
|
|
72
|
+
def enforce_token_cap(
|
|
73
|
+
*,
|
|
74
|
+
content: str,
|
|
75
|
+
max_tokens: int,
|
|
76
|
+
language_model: LanguageModelInfo = _DEFAULT_LANGUAGE_MODEL,
|
|
77
|
+
) -> str:
|
|
78
|
+
if not content:
|
|
79
|
+
return content
|
|
80
|
+
|
|
81
|
+
encode, decode = _get_model_tokenizer(language_model=language_model)
|
|
82
|
+
|
|
83
|
+
if _count_tokens(content=content, encode=encode) <= max_tokens:
|
|
84
|
+
return content
|
|
85
|
+
|
|
86
|
+
marker_tokens = _count_tokens(content=_TRUNCATION_MARKER, encode=encode)
|
|
87
|
+
target = max(0, max_tokens - marker_tokens)
|
|
88
|
+
paragraphs = content.split("\n\n")
|
|
89
|
+
accepted: list[str] = []
|
|
90
|
+
running = 0
|
|
91
|
+
for paragraph in paragraphs:
|
|
92
|
+
paragraph_tokens = _count_tokens(content=paragraph + "\n\n", encode=encode)
|
|
93
|
+
if running + paragraph_tokens > target:
|
|
94
|
+
break
|
|
95
|
+
accepted.append(paragraph)
|
|
96
|
+
running += paragraph_tokens
|
|
97
|
+
|
|
98
|
+
if accepted:
|
|
99
|
+
truncated = "\n\n".join(accepted)
|
|
100
|
+
else:
|
|
101
|
+
truncated = decode(encode(content)[:target])
|
|
102
|
+
|
|
103
|
+
result = truncated.rstrip() + _TRUNCATION_MARKER
|
|
104
|
+
shrink = target
|
|
105
|
+
while _count_tokens(content=result, encode=encode) > max_tokens and shrink > 0:
|
|
106
|
+
shrink -= 1
|
|
107
|
+
body = decode(encode(result)[:shrink]).rstrip()
|
|
108
|
+
result = body + _TRUNCATION_MARKER
|
|
109
|
+
|
|
110
|
+
if _count_tokens(content=result, encode=encode) > max_tokens:
|
|
111
|
+
result = decode(encode(result)[:max_tokens])
|
|
112
|
+
|
|
113
|
+
return result
|
|
114
|
+
|
|
115
|
+
|
|
116
|
+
async def load_user_memory(
|
|
117
|
+
*,
|
|
118
|
+
event: ChatEvent,
|
|
119
|
+
config: UserMemoryConfig,
|
|
120
|
+
logger: Logger,
|
|
121
|
+
) -> UserMemoryState | None:
|
|
122
|
+
user_id = event.user_id
|
|
123
|
+
company_id = event.company_id
|
|
124
|
+
if not user_id or not company_id:
|
|
125
|
+
logger.warning("[user-memory] empty user_id/company_id - skipping memory load")
|
|
126
|
+
return None
|
|
127
|
+
|
|
128
|
+
scope_id = await ensure_user_memory_folder(
|
|
129
|
+
user_id=user_id,
|
|
130
|
+
company_id=company_id,
|
|
131
|
+
root_folder=config.root_folder,
|
|
132
|
+
logger=logger,
|
|
133
|
+
)
|
|
134
|
+
if scope_id is None:
|
|
135
|
+
logger.info("[user-memory] folder ensure failed - running without memory")
|
|
136
|
+
return None
|
|
137
|
+
|
|
138
|
+
text = await download_user_memory(
|
|
139
|
+
scope_id=scope_id,
|
|
140
|
+
user_id=user_id,
|
|
141
|
+
company_id=company_id,
|
|
142
|
+
logger=logger,
|
|
143
|
+
)
|
|
144
|
+
return UserMemoryState(
|
|
145
|
+
scope_id=scope_id,
|
|
146
|
+
text=enforce_token_cap(
|
|
147
|
+
content=text,
|
|
148
|
+
max_tokens=config.max_tokens,
|
|
149
|
+
language_model=config.language_model,
|
|
150
|
+
),
|
|
151
|
+
)
|
|
152
|
+
|
|
153
|
+
|
|
154
|
+
async def ensure_user_memory_folder(
|
|
155
|
+
*,
|
|
156
|
+
user_id: str,
|
|
157
|
+
company_id: str,
|
|
158
|
+
root_folder: str,
|
|
159
|
+
logger: Logger,
|
|
160
|
+
) -> str | None:
|
|
161
|
+
root_scope_id = await _resolve_root_folder(
|
|
162
|
+
user_id=user_id,
|
|
163
|
+
company_id=company_id,
|
|
164
|
+
root_folder=root_folder,
|
|
165
|
+
logger=logger,
|
|
166
|
+
)
|
|
167
|
+
if root_scope_id is None:
|
|
168
|
+
return None
|
|
169
|
+
|
|
170
|
+
user_folder_path = f"/{root_folder.strip('/')}/{user_id}"
|
|
171
|
+
scope_id: str | None = None
|
|
172
|
+
try:
|
|
173
|
+
info = await unique_sdk.Folder.get_info_async(
|
|
174
|
+
user_id=user_id,
|
|
175
|
+
company_id=company_id,
|
|
176
|
+
folderPath=user_folder_path,
|
|
177
|
+
)
|
|
178
|
+
return info.get("id")
|
|
179
|
+
except Exception:
|
|
180
|
+
logger.warning("[user-memory] user memory folder not found - creating new one")
|
|
181
|
+
|
|
182
|
+
try:
|
|
183
|
+
created = await unique_sdk.Folder.create_paths_async(
|
|
184
|
+
user_id=user_id,
|
|
185
|
+
company_id=company_id,
|
|
186
|
+
parentScopeId=root_scope_id,
|
|
187
|
+
relativePaths=[user_id],
|
|
188
|
+
inheritAccess=False,
|
|
189
|
+
)
|
|
190
|
+
except Exception as exc:
|
|
191
|
+
logger.warning(
|
|
192
|
+
"[user-memory] failed to create user folder %s: [%s] %s",
|
|
193
|
+
user_folder_path,
|
|
194
|
+
type(exc).__name__,
|
|
195
|
+
exc,
|
|
196
|
+
)
|
|
197
|
+
return None
|
|
198
|
+
|
|
199
|
+
created_folders = (created or {}).get("createdFolders", []) or []
|
|
200
|
+
if len(created_folders) > 1:
|
|
201
|
+
logger.warning(
|
|
202
|
+
"[user-memory] create_paths returned %d folders for %s, "
|
|
203
|
+
"expected exactly 1; using the first one",
|
|
204
|
+
len(created_folders),
|
|
205
|
+
user_folder_path,
|
|
206
|
+
)
|
|
207
|
+
scope_id = created_folders[0].get("id") if created_folders else None
|
|
208
|
+
if not scope_id:
|
|
209
|
+
logger.warning(
|
|
210
|
+
"[user-memory] create_paths returned no folder id for %s",
|
|
211
|
+
user_folder_path,
|
|
212
|
+
)
|
|
213
|
+
return None
|
|
214
|
+
|
|
215
|
+
try:
|
|
216
|
+
await unique_sdk.Folder.add_access_async(
|
|
217
|
+
user_id=user_id,
|
|
218
|
+
company_id=company_id,
|
|
219
|
+
scopeId=scope_id,
|
|
220
|
+
scopeAccesses=[
|
|
221
|
+
{
|
|
222
|
+
"entityId": user_id,
|
|
223
|
+
"type": "READ",
|
|
224
|
+
"entityType": "USER",
|
|
225
|
+
},
|
|
226
|
+
{
|
|
227
|
+
"entityId": user_id,
|
|
228
|
+
"type": "WRITE",
|
|
229
|
+
"entityType": "USER",
|
|
230
|
+
},
|
|
231
|
+
],
|
|
232
|
+
applyToSubScopes=True,
|
|
233
|
+
)
|
|
234
|
+
except Exception as exc:
|
|
235
|
+
logger.warning(
|
|
236
|
+
"[user-memory] failed to grant read/write access on scope %s "
|
|
237
|
+
"for user %s: [%s] %s",
|
|
238
|
+
scope_id,
|
|
239
|
+
user_id,
|
|
240
|
+
type(exc).__name__,
|
|
241
|
+
exc,
|
|
242
|
+
)
|
|
243
|
+
return None
|
|
244
|
+
|
|
245
|
+
return scope_id
|
|
246
|
+
|
|
247
|
+
|
|
248
|
+
async def _resolve_root_folder(
|
|
249
|
+
*,
|
|
250
|
+
user_id: str,
|
|
251
|
+
company_id: str,
|
|
252
|
+
root_folder: str,
|
|
253
|
+
logger: Logger,
|
|
254
|
+
) -> str | None:
|
|
255
|
+
root_path = f"/{root_folder.strip('/')}"
|
|
256
|
+
try:
|
|
257
|
+
root_info = await unique_sdk.Folder.get_info_async(
|
|
258
|
+
user_id=user_id,
|
|
259
|
+
company_id=company_id,
|
|
260
|
+
folderPath=root_path,
|
|
261
|
+
)
|
|
262
|
+
except Exception as exc:
|
|
263
|
+
logger.warning(
|
|
264
|
+
"[user-memory] failed to resolve pre-provisioned root folder %s: [%s] %s",
|
|
265
|
+
root_path,
|
|
266
|
+
type(exc).__name__,
|
|
267
|
+
exc,
|
|
268
|
+
)
|
|
269
|
+
return None
|
|
270
|
+
|
|
271
|
+
root_scope_id = root_info.get("id")
|
|
272
|
+
if not root_scope_id:
|
|
273
|
+
logger.warning(
|
|
274
|
+
"[user-memory] root folder lookup returned no id for %s",
|
|
275
|
+
root_path,
|
|
276
|
+
)
|
|
277
|
+
return None
|
|
278
|
+
|
|
279
|
+
return root_scope_id
|
|
280
|
+
|
|
281
|
+
|
|
282
|
+
async def download_user_memory(
|
|
283
|
+
*,
|
|
284
|
+
scope_id: str,
|
|
285
|
+
user_id: str,
|
|
286
|
+
company_id: str,
|
|
287
|
+
logger: Logger,
|
|
288
|
+
) -> str:
|
|
289
|
+
try:
|
|
290
|
+
contents = await search_contents_async(
|
|
291
|
+
user_id=user_id,
|
|
292
|
+
company_id=company_id,
|
|
293
|
+
chat_id=None,
|
|
294
|
+
where={"ownerId": {"equals": scope_id}},
|
|
295
|
+
)
|
|
296
|
+
except Exception as exc:
|
|
297
|
+
logger.warning(
|
|
298
|
+
"[user-memory] failed to list contents in scope %s: [%s] %s",
|
|
299
|
+
scope_id,
|
|
300
|
+
type(exc).__name__,
|
|
301
|
+
exc,
|
|
302
|
+
)
|
|
303
|
+
return ""
|
|
304
|
+
|
|
305
|
+
memory_content = next(
|
|
306
|
+
(content for content in contents if (content.key or "") == MEMORY_FILENAME),
|
|
307
|
+
None,
|
|
308
|
+
)
|
|
309
|
+
if memory_content is None:
|
|
310
|
+
logger.debug(
|
|
311
|
+
"[user-memory] no %s in scope %s - first turn for this user",
|
|
312
|
+
MEMORY_FILENAME,
|
|
313
|
+
scope_id,
|
|
314
|
+
)
|
|
315
|
+
return ""
|
|
316
|
+
|
|
317
|
+
try:
|
|
318
|
+
content_bytes = await download_content_to_bytes_async(
|
|
319
|
+
user_id=user_id,
|
|
320
|
+
company_id=company_id,
|
|
321
|
+
content_id=memory_content.id,
|
|
322
|
+
chat_id=None,
|
|
323
|
+
)
|
|
324
|
+
text = content_bytes.decode("utf-8", errors="replace")
|
|
325
|
+
logger.info(
|
|
326
|
+
"[user-memory] downloaded %s (%d bytes) from scope %s",
|
|
327
|
+
MEMORY_FILENAME,
|
|
328
|
+
len(content_bytes),
|
|
329
|
+
scope_id,
|
|
330
|
+
)
|
|
331
|
+
return text
|
|
332
|
+
except Exception as exc:
|
|
333
|
+
logger.warning(
|
|
334
|
+
"[user-memory] failed to download %s from scope %s: [%s] %s",
|
|
335
|
+
MEMORY_FILENAME,
|
|
336
|
+
scope_id,
|
|
337
|
+
type(exc).__name__,
|
|
338
|
+
exc,
|
|
339
|
+
)
|
|
340
|
+
return ""
|
|
341
|
+
|
|
342
|
+
|
|
343
|
+
async def upload_user_memory(
|
|
344
|
+
*,
|
|
345
|
+
scope_id: str,
|
|
346
|
+
content: str,
|
|
347
|
+
user_id: str,
|
|
348
|
+
company_id: str,
|
|
349
|
+
logger: Logger,
|
|
350
|
+
) -> bool:
|
|
351
|
+
if not content.strip():
|
|
352
|
+
logger.warning(
|
|
353
|
+
"[user-memory] refusing to upload empty memory file to scope %s",
|
|
354
|
+
scope_id,
|
|
355
|
+
)
|
|
356
|
+
return False
|
|
357
|
+
|
|
358
|
+
try:
|
|
359
|
+
await upload_content_from_bytes_async(
|
|
360
|
+
user_id=user_id,
|
|
361
|
+
company_id=company_id,
|
|
362
|
+
content=content.encode("utf-8"),
|
|
363
|
+
content_name=MEMORY_FILENAME,
|
|
364
|
+
mime_type=MIME_TYPE,
|
|
365
|
+
scope_id=scope_id,
|
|
366
|
+
ingestion_config={
|
|
367
|
+
"uniqueIngestionMode": "SKIP_INGESTION",
|
|
368
|
+
"hideInChat": True,
|
|
369
|
+
},
|
|
370
|
+
)
|
|
371
|
+
logger.info(
|
|
372
|
+
"[user-memory] uploaded %s (%d bytes) to scope %s",
|
|
373
|
+
MEMORY_FILENAME,
|
|
374
|
+
len(content.encode("utf-8")),
|
|
375
|
+
scope_id,
|
|
376
|
+
)
|
|
377
|
+
return True
|
|
378
|
+
except Exception as exc:
|
|
379
|
+
logger.error(
|
|
380
|
+
"[user-memory] upload failed for scope %s: [%s] %s",
|
|
381
|
+
scope_id,
|
|
382
|
+
type(exc).__name__,
|
|
383
|
+
exc,
|
|
384
|
+
exc_info=True,
|
|
385
|
+
)
|
|
386
|
+
return False
|
|
387
|
+
|
|
388
|
+
|
|
389
|
+
async def consolidate_user_memory(
|
|
390
|
+
*,
|
|
391
|
+
current_memory: str,
|
|
392
|
+
user_id: str,
|
|
393
|
+
user_message: str,
|
|
394
|
+
assistant_message: str,
|
|
395
|
+
config: UserMemoryConfig,
|
|
396
|
+
event: ChatEvent,
|
|
397
|
+
logger: Logger,
|
|
398
|
+
) -> str:
|
|
399
|
+
safe_current = enforce_token_cap(
|
|
400
|
+
content=current_memory,
|
|
401
|
+
max_tokens=config.max_tokens,
|
|
402
|
+
language_model=config.language_model,
|
|
403
|
+
)
|
|
404
|
+
logger.info(
|
|
405
|
+
"[user-memory] consolidating turn - existing_memory_tokens=%d | "
|
|
406
|
+
"user_msg_chars=%d | assistant_msg_chars=%d",
|
|
407
|
+
count_tokens(content=safe_current, language_model=config.language_model),
|
|
408
|
+
len(user_message or ""),
|
|
409
|
+
len(assistant_message or ""),
|
|
410
|
+
)
|
|
411
|
+
|
|
412
|
+
if not (user_message or "").strip() and not (assistant_message or "").strip():
|
|
413
|
+
return safe_current or enforce_token_cap(
|
|
414
|
+
content=empty_profile(user_id),
|
|
415
|
+
max_tokens=config.max_tokens,
|
|
416
|
+
language_model=config.language_model,
|
|
417
|
+
)
|
|
418
|
+
|
|
419
|
+
if not safe_current.strip():
|
|
420
|
+
safe_current = empty_profile(user_id)
|
|
421
|
+
|
|
422
|
+
try:
|
|
423
|
+
llm_service = LanguageModelService(event)
|
|
424
|
+
except Exception as exc:
|
|
425
|
+
logger.error(
|
|
426
|
+
"[user-memory] cannot construct LanguageModelService: [%s] %s",
|
|
427
|
+
type(exc).__name__,
|
|
428
|
+
exc,
|
|
429
|
+
)
|
|
430
|
+
return safe_current
|
|
431
|
+
|
|
432
|
+
messages = LanguageModelMessages(
|
|
433
|
+
[
|
|
434
|
+
LanguageModelSystemMessage(
|
|
435
|
+
content=consolidation_system_prompt(config.max_tokens)
|
|
436
|
+
),
|
|
437
|
+
LanguageModelUserMessage(
|
|
438
|
+
content=consolidation_user_prompt(
|
|
439
|
+
user_id=user_id,
|
|
440
|
+
existing_memory=safe_current,
|
|
441
|
+
user_message=_sanitize_for_xml_context(user_message or ""),
|
|
442
|
+
assistant_message=_sanitize_for_xml_context(
|
|
443
|
+
assistant_message or ""
|
|
444
|
+
),
|
|
445
|
+
)
|
|
446
|
+
),
|
|
447
|
+
]
|
|
448
|
+
)
|
|
449
|
+
|
|
450
|
+
try:
|
|
451
|
+
response = await llm_service.complete_async(
|
|
452
|
+
messages=messages,
|
|
453
|
+
model_name=config.language_model.name,
|
|
454
|
+
other_options={
|
|
455
|
+
"max_tokens": config.max_tokens + _LLM_OUTPUT_HEADROOM_TOKENS,
|
|
456
|
+
},
|
|
457
|
+
)
|
|
458
|
+
except Exception as exc:
|
|
459
|
+
logger.warning(
|
|
460
|
+
"[user-memory] consolidation LLM call failed (model=%s): [%s] %s",
|
|
461
|
+
config.language_model.name,
|
|
462
|
+
type(exc).__name__,
|
|
463
|
+
exc,
|
|
464
|
+
)
|
|
465
|
+
return safe_current
|
|
466
|
+
|
|
467
|
+
try:
|
|
468
|
+
raw = response.choices[0].message.content or ""
|
|
469
|
+
except Exception as exc:
|
|
470
|
+
logger.warning(
|
|
471
|
+
"[user-memory] could not extract content from LLM response: [%s] %s",
|
|
472
|
+
type(exc).__name__,
|
|
473
|
+
exc,
|
|
474
|
+
)
|
|
475
|
+
return safe_current
|
|
476
|
+
|
|
477
|
+
if not isinstance(raw, str):
|
|
478
|
+
logger.warning(
|
|
479
|
+
"[user-memory] LLM returned non-string content (%s)",
|
|
480
|
+
type(raw).__name__,
|
|
481
|
+
)
|
|
482
|
+
return safe_current
|
|
483
|
+
|
|
484
|
+
if raw.strip().upper() == "NOOP":
|
|
485
|
+
logger.info("[user-memory] consolidation NOOP - keeping existing memory")
|
|
486
|
+
return safe_current
|
|
487
|
+
|
|
488
|
+
candidate = _strip_code_fences(raw).strip()
|
|
489
|
+
if not _is_well_formed_profile(candidate):
|
|
490
|
+
logger.warning(
|
|
491
|
+
"[user-memory] LLM output did not look like a profile (%d chars)",
|
|
492
|
+
len(candidate),
|
|
493
|
+
)
|
|
494
|
+
return safe_current
|
|
495
|
+
|
|
496
|
+
capped = enforce_token_cap(
|
|
497
|
+
content=candidate,
|
|
498
|
+
max_tokens=config.max_tokens,
|
|
499
|
+
language_model=config.language_model,
|
|
500
|
+
)
|
|
501
|
+
if (
|
|
502
|
+
safe_current
|
|
503
|
+
and _FRONTMATTER_RE.sub("", capped).strip()
|
|
504
|
+
== _FRONTMATTER_RE.sub("", safe_current).strip()
|
|
505
|
+
):
|
|
506
|
+
logger.debug("[user-memory] memory body unchanged - skipping update")
|
|
507
|
+
return safe_current
|
|
508
|
+
|
|
509
|
+
logger.info(
|
|
510
|
+
"[user-memory] consolidation produced %d tokens (cap=%d)",
|
|
511
|
+
count_tokens(content=capped, language_model=config.language_model),
|
|
512
|
+
config.max_tokens,
|
|
513
|
+
)
|
|
514
|
+
return capped
|
|
515
|
+
|
|
516
|
+
|
|
517
|
+
def _sanitize_for_xml_context(text: str) -> str:
|
|
518
|
+
return text.replace("</", "< /")
|
|
519
|
+
|
|
520
|
+
|
|
521
|
+
def _is_well_formed_profile(content: str) -> bool:
|
|
522
|
+
if not content or len(content.strip()) < 20:
|
|
523
|
+
return False
|
|
524
|
+
if _FRONTMATTER_RE.match(content):
|
|
525
|
+
return True
|
|
526
|
+
return "## Identity" in content or "# User Memory" in content
|
|
527
|
+
|
|
528
|
+
|
|
529
|
+
def _strip_code_fences(text: str) -> str:
|
|
530
|
+
stripped = text.strip()
|
|
531
|
+
if not stripped.startswith("```"):
|
|
532
|
+
return text
|
|
533
|
+
lines = stripped.splitlines()
|
|
534
|
+
if len(lines) >= 3 and lines[-1].strip() == "```":
|
|
535
|
+
return "\n".join(lines[1:-1])
|
|
536
|
+
return text
|
|
537
|
+
|
|
538
|
+
|
|
539
|
+
def looks_empty_profile(content: str) -> bool:
|
|
540
|
+
if not content:
|
|
541
|
+
return True
|
|
542
|
+
return content.count("_(empty)_") >= len(SECTION_HEADINGS)
|
|
543
|
+
|
|
544
|
+
|
|
545
|
+
def prompt_section(memory: str | None) -> str:
|
|
546
|
+
if not memory or not memory.strip():
|
|
547
|
+
return ""
|
|
548
|
+
return memory.strip()
|
|
@@ -0,0 +1,71 @@
|
|
|
1
|
+
from logging import Logger
|
|
2
|
+
|
|
3
|
+
from unique_toolkit.agentic.postprocessor.postprocessor_manager import Postprocessor
|
|
4
|
+
from unique_toolkit.app.schemas import ChatEvent
|
|
5
|
+
from unique_toolkit.language_model.schemas import LanguageModelStreamResponse
|
|
6
|
+
|
|
7
|
+
from unique_user_memory.config import UserMemoryConfig
|
|
8
|
+
from unique_user_memory.user_memory import (
|
|
9
|
+
UserMemoryState,
|
|
10
|
+
consolidate_user_memory,
|
|
11
|
+
upload_user_memory,
|
|
12
|
+
)
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
class UserMemoryPostprocessor(Postprocessor):
|
|
16
|
+
def __init__(
|
|
17
|
+
self,
|
|
18
|
+
*,
|
|
19
|
+
config: UserMemoryConfig,
|
|
20
|
+
event: ChatEvent,
|
|
21
|
+
state: UserMemoryState,
|
|
22
|
+
logger: Logger,
|
|
23
|
+
) -> None:
|
|
24
|
+
super().__init__(name="UserMemoryPostprocessor")
|
|
25
|
+
self._config = config
|
|
26
|
+
self._event = event
|
|
27
|
+
self._state = state
|
|
28
|
+
self._logger = logger
|
|
29
|
+
self._new_memory: str | None = None
|
|
30
|
+
|
|
31
|
+
async def run(self, loop_response: LanguageModelStreamResponse) -> None:
|
|
32
|
+
self._logger.info("[user-memory] running postprocessor")
|
|
33
|
+
user_id = self._event.user_id
|
|
34
|
+
company_id = self._event.company_id
|
|
35
|
+
if not user_id or not company_id:
|
|
36
|
+
return
|
|
37
|
+
|
|
38
|
+
self._new_memory = await consolidate_user_memory(
|
|
39
|
+
current_memory=self._state.text,
|
|
40
|
+
user_id=user_id,
|
|
41
|
+
user_message=self._event.payload.user_message.text or "",
|
|
42
|
+
assistant_message=loop_response.message.text or "",
|
|
43
|
+
config=self._config,
|
|
44
|
+
event=self._event,
|
|
45
|
+
logger=self._logger,
|
|
46
|
+
)
|
|
47
|
+
|
|
48
|
+
if self._new_memory == self._state.text:
|
|
49
|
+
self._logger.debug("[user-memory] consolidation NOOP - skipping upload")
|
|
50
|
+
return
|
|
51
|
+
|
|
52
|
+
uploaded = await upload_user_memory(
|
|
53
|
+
scope_id=self._state.scope_id,
|
|
54
|
+
content=self._new_memory,
|
|
55
|
+
user_id=user_id,
|
|
56
|
+
company_id=company_id,
|
|
57
|
+
logger=self._logger,
|
|
58
|
+
)
|
|
59
|
+
if not uploaded:
|
|
60
|
+
self._logger.warning("[user-memory] memory update was not uploaded")
|
|
61
|
+
return
|
|
62
|
+
|
|
63
|
+
self._logger.info("[user-memory] memory updated and uploaded")
|
|
64
|
+
|
|
65
|
+
def apply_postprocessing_to_response(
|
|
66
|
+
self, loop_response: LanguageModelStreamResponse
|
|
67
|
+
) -> bool:
|
|
68
|
+
return False
|
|
69
|
+
|
|
70
|
+
async def remove_from_text(self, text: str) -> str:
|
|
71
|
+
return text
|
|
@@ -0,0 +1,185 @@
|
|
|
1
|
+
from datetime import datetime, timezone
|
|
2
|
+
|
|
3
|
+
from jinja2 import Template
|
|
4
|
+
|
|
5
|
+
SECTION_HEADINGS: tuple[str, ...] = (
|
|
6
|
+
"Identity",
|
|
7
|
+
"Communication Preferences",
|
|
8
|
+
"Work Context",
|
|
9
|
+
"Skills & Expertise",
|
|
10
|
+
"Recent Topics",
|
|
11
|
+
"Open Questions / Follow-ups",
|
|
12
|
+
)
|
|
13
|
+
|
|
14
|
+
_EMPTY_PROFILE_TEMPLATE = """\
|
|
15
|
+
---
|
|
16
|
+
user_id: {{ user_id }}
|
|
17
|
+
schema_version: 1
|
|
18
|
+
last_updated: {{ timestamp }}
|
|
19
|
+
turn_count: 0
|
|
20
|
+
---
|
|
21
|
+
|
|
22
|
+
# User Memory
|
|
23
|
+
|
|
24
|
+
{% for heading in section_headings -%}
|
|
25
|
+
## {{ heading }}
|
|
26
|
+
_(empty)_
|
|
27
|
+
{% if not loop.last %}
|
|
28
|
+
{% endif %}
|
|
29
|
+
{% endfor -%}
|
|
30
|
+
"""
|
|
31
|
+
|
|
32
|
+
|
|
33
|
+
def empty_profile(user_id: str) -> str:
|
|
34
|
+
timestamp = datetime.now(timezone.utc).isoformat(timespec="seconds")
|
|
35
|
+
return Template(_EMPTY_PROFILE_TEMPLATE).render(
|
|
36
|
+
user_id=user_id,
|
|
37
|
+
timestamp=timestamp,
|
|
38
|
+
section_headings=SECTION_HEADINGS,
|
|
39
|
+
)
|
|
40
|
+
|
|
41
|
+
|
|
42
|
+
_CONSOLIDATION_SYSTEM_PROMPT_TEMPLATE = """\
|
|
43
|
+
You are a memory-consolidation engine for the Unique AI platform.
|
|
44
|
+
|
|
45
|
+
Your sole job is to maintain a structured Markdown profile of an end
|
|
46
|
+
user, distilled from their conversations with AI assistants. The profile
|
|
47
|
+
is read on every future turn and shapes how the assistant addresses,
|
|
48
|
+
helps, and remembers the user. It is therefore a high-leverage artefact:
|
|
49
|
+
small mistakes in extraction compound across every future conversation.
|
|
50
|
+
|
|
51
|
+
# Inputs
|
|
52
|
+
|
|
53
|
+
You receive two XML blocks:
|
|
54
|
+
|
|
55
|
+
1. `<existing_memory>` - the current profile file (Markdown, with YAML
|
|
56
|
+
frontmatter). May be empty on the user's first turn.
|
|
57
|
+
2. `<new_turn>` - the most recent user message and the assistant's
|
|
58
|
+
reply, prefixed with `user:` and `assistant:`.
|
|
59
|
+
|
|
60
|
+
# Output
|
|
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.
|
|
65
|
+
|
|
66
|
+
The body MUST contain exactly these section headings, in this order, even
|
|
67
|
+
when a section is empty (use the literal string `_(empty)_` as a placeholder):
|
|
68
|
+
|
|
69
|
+
{{ section_list }}
|
|
70
|
+
|
|
71
|
+
# Operations
|
|
72
|
+
|
|
73
|
+
For each candidate fact in `<new_turn>`, decide one of:
|
|
74
|
+
|
|
75
|
+
- ADD - the fact is new and stable enough to remember (preferences,
|
|
76
|
+
identity attributes, ongoing projects, skills, dated topics). Add it as
|
|
77
|
+
a bullet in the most appropriate section.
|
|
78
|
+
- UPDATE - the fact refines or supersedes an existing bullet. Edit
|
|
79
|
+
the existing bullet in place; do not add a duplicate.
|
|
80
|
+
- DELETE - the new turn explicitly contradicts or invalidates an
|
|
81
|
+
existing bullet that is not worth keeping as history. Remove it.
|
|
82
|
+
- NOOP - the new turn contains no facts about the user (small talk,
|
|
83
|
+
factual questions, code requests, abstract discussion). Output the
|
|
84
|
+
single word `NOOP` and nothing else. The caller keeps the existing
|
|
85
|
+
memory unchanged and skips the write entirely.
|
|
86
|
+
|
|
87
|
+
Prefer UPDATE over ADD when in doubt - duplication is the most common
|
|
88
|
+
failure mode of memory systems.
|
|
89
|
+
|
|
90
|
+
# What to extract
|
|
91
|
+
|
|
92
|
+
ADD/UPDATE for facts that are:
|
|
93
|
+
|
|
94
|
+
- Stable - true beyond the current chat (name, role, employer,
|
|
95
|
+
team, timezone, language, technical stack, recurring projects).
|
|
96
|
+
- Preference-shaped - communication style, formatting, depth,
|
|
97
|
+
tone, language, expertise level, examples preferred over theory.
|
|
98
|
+
- Contextual but durable - current focus areas, active projects,
|
|
99
|
+
multi-week goals, deadlines mentioned by the user.
|
|
100
|
+
- Hand-offs - explicit "let's revisit X later", "remind me about
|
|
101
|
+
Y", "I'll come back to Z" go into "Open Questions / Follow-ups".
|
|
102
|
+
|
|
103
|
+
NEVER extract:
|
|
104
|
+
|
|
105
|
+
- Sensitive credentials, API keys, passwords, OTPs, payment info,
|
|
106
|
+
health-record details, government IDs. If the user pastes any of
|
|
107
|
+
these, ignore the value entirely.
|
|
108
|
+
- Information about other named individuals beyond the immediate
|
|
109
|
+
professional context.
|
|
110
|
+
- Transient turn-level state: one-off factual answers, code snippets,
|
|
111
|
+
error messages, file contents, search results.
|
|
112
|
+
- Anything stated as third-party information or retrieved context.
|
|
113
|
+
|
|
114
|
+
# Token budget - STRICT
|
|
115
|
+
|
|
116
|
+
The complete file MUST be <= {{ max_tokens }} tokens.
|
|
117
|
+
When approaching the budget, drop content in this priority order:
|
|
118
|
+
|
|
119
|
+
1. Oldest entries in Recent Topics.
|
|
120
|
+
2. Resolved or stale entries in Open Questions / Follow-ups.
|
|
121
|
+
3. Fold low-signal Work Context bullets into a one-line summary.
|
|
122
|
+
4. Fold low-signal Skills & Expertise bullets into broader categories.
|
|
123
|
+
5. Identity and Communication Preferences - never drop, only tighten.
|
|
124
|
+
|
|
125
|
+
# Current date and time
|
|
126
|
+
|
|
127
|
+
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.
|
|
128
|
+
|
|
129
|
+
# Frontmatter rules
|
|
130
|
+
|
|
131
|
+
- Preserve `user_id` and `schema_version` from `<existing_memory>` exactly.
|
|
132
|
+
- Set `last_updated` to the supplied current UTC timestamp ({{ now_datetime }}).
|
|
133
|
+
- Increment `turn_count` by 1.
|
|
134
|
+
- If `<existing_memory>` is empty, initialize with `schema_version: 1`,
|
|
135
|
+
`turn_count: 1`, and the user_id supplied in the user message.
|
|
136
|
+
|
|
137
|
+
# Style
|
|
138
|
+
|
|
139
|
+
- Use `-` markdown bullets, no nesting beyond two levels.
|
|
140
|
+
- Keep bullets short.
|
|
141
|
+
- For dated entries in Recent Topics, prefix with `YYYY-MM-DD HH:MM UTC:`
|
|
142
|
+
using the supplied current date and time ({{ now_datetime }}).
|
|
143
|
+
- No emojis in section headings.
|
|
144
|
+
"""
|
|
145
|
+
|
|
146
|
+
|
|
147
|
+
def consolidation_system_prompt(max_tokens: int) -> str:
|
|
148
|
+
section_list = "\n".join(f"- ## {heading}" for heading in SECTION_HEADINGS)
|
|
149
|
+
now = datetime.now(timezone.utc)
|
|
150
|
+
return Template(_CONSOLIDATION_SYSTEM_PROMPT_TEMPLATE).render(
|
|
151
|
+
max_tokens=max_tokens,
|
|
152
|
+
section_list=section_list,
|
|
153
|
+
now_datetime=now.strftime("%Y-%m-%d %H:%M UTC"),
|
|
154
|
+
)
|
|
155
|
+
|
|
156
|
+
|
|
157
|
+
_CONSOLIDATION_USER_PROMPT_TEMPLATE = """\
|
|
158
|
+
User ID: {{ user_id }}
|
|
159
|
+
|
|
160
|
+
<existing_memory>
|
|
161
|
+
{{ existing_memory }}
|
|
162
|
+
</existing_memory>
|
|
163
|
+
|
|
164
|
+
<new_turn>
|
|
165
|
+
user: {{ user_message }}
|
|
166
|
+
assistant: {{ assistant_message }}
|
|
167
|
+
</new_turn>
|
|
168
|
+
|
|
169
|
+
Return the complete rewritten profile file now.
|
|
170
|
+
"""
|
|
171
|
+
|
|
172
|
+
|
|
173
|
+
def consolidation_user_prompt(
|
|
174
|
+
user_id: str,
|
|
175
|
+
existing_memory: str,
|
|
176
|
+
user_message: str,
|
|
177
|
+
assistant_message: str,
|
|
178
|
+
) -> str:
|
|
179
|
+
existing = existing_memory.strip() or "(empty - this is the user's first turn)"
|
|
180
|
+
return Template(_CONSOLIDATION_USER_PROMPT_TEMPLATE).render(
|
|
181
|
+
user_id=user_id,
|
|
182
|
+
existing_memory=existing,
|
|
183
|
+
user_message=(user_message or "").strip(),
|
|
184
|
+
assistant_message=(assistant_message or "").strip(),
|
|
185
|
+
)
|