agent-memory-git 0.1.0__py3-none-any.whl
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.
- agent_memory/PROTOCOL.md +80 -0
- agent_memory/__init__.py +1 -0
- agent_memory/__main__.py +7 -0
- agent_memory/capture.py +207 -0
- agent_memory/cli.py +317 -0
- agent_memory/config.py +53 -0
- agent_memory/context.py +73 -0
- agent_memory/embed_worker.py +90 -0
- agent_memory/frontmatter.py +146 -0
- agent_memory/gitsync.py +119 -0
- agent_memory/hooks.py +95 -0
- agent_memory/importers.py +231 -0
- agent_memory/install.py +483 -0
- agent_memory/lint.py +45 -0
- agent_memory/select.py +105 -0
- agent_memory/sessions.py +64 -0
- agent_memory/store.py +97 -0
- agent_memory/ui.py +410 -0
- agent_memory/ui_server.py +371 -0
- agent_memory/util.py +101 -0
- agent_memory/webui/app.css +192 -0
- agent_memory/webui/app.js +635 -0
- agent_memory/webui/index.html +31 -0
- agent_memory/workdir.py +135 -0
- agent_memory_git-0.1.0.dist-info/METADATA +183 -0
- agent_memory_git-0.1.0.dist-info/RECORD +29 -0
- agent_memory_git-0.1.0.dist-info/WHEEL +4 -0
- agent_memory_git-0.1.0.dist-info/entry_points.txt +3 -0
- agent_memory_git-0.1.0.dist-info/licenses/LICENSE +21 -0
agent_memory/sessions.py
ADDED
|
@@ -0,0 +1,64 @@
|
|
|
1
|
+
"""Best-effort lookup of a session's human-readable name."""
|
|
2
|
+
|
|
3
|
+
import json
|
|
4
|
+
import os
|
|
5
|
+
import sqlite3
|
|
6
|
+
from pathlib import Path
|
|
7
|
+
|
|
8
|
+
|
|
9
|
+
def _claude_title(transcript):
|
|
10
|
+
title = ""
|
|
11
|
+
try:
|
|
12
|
+
with open(transcript, encoding="utf-8", errors="replace") as f:
|
|
13
|
+
for raw in f:
|
|
14
|
+
if '"custom-title"' not in raw:
|
|
15
|
+
continue
|
|
16
|
+
try:
|
|
17
|
+
title = json.loads(raw).get("customTitle") or title
|
|
18
|
+
except ValueError:
|
|
19
|
+
pass
|
|
20
|
+
except OSError:
|
|
21
|
+
pass
|
|
22
|
+
return title
|
|
23
|
+
|
|
24
|
+
|
|
25
|
+
def _codex_thread(session_id, column):
|
|
26
|
+
home = Path(os.environ.get("CODEX_HOME", "~/.codex")).expanduser()
|
|
27
|
+
dbs = sorted(home.glob("state_*.sqlite"))
|
|
28
|
+
if not dbs or not session_id:
|
|
29
|
+
return ""
|
|
30
|
+
try:
|
|
31
|
+
con = sqlite3.connect(f"file:{dbs[-1]}?mode=ro", uri=True, timeout=1)
|
|
32
|
+
try:
|
|
33
|
+
row = con.execute(f"select {column} from threads where id = ?",
|
|
34
|
+
(session_id,)).fetchone()
|
|
35
|
+
finally:
|
|
36
|
+
con.close()
|
|
37
|
+
except sqlite3.Error:
|
|
38
|
+
return ""
|
|
39
|
+
return (row[0] or "") if row else ""
|
|
40
|
+
|
|
41
|
+
|
|
42
|
+
def _codex_title(session_id):
|
|
43
|
+
return _codex_thread(session_id, "coalesce(nullif(name, ''), title)")
|
|
44
|
+
|
|
45
|
+
|
|
46
|
+
def transcript_path(tool, session_id):
|
|
47
|
+
"""Locate a session transcript when the hook payload did not carry one."""
|
|
48
|
+
if not session_id:
|
|
49
|
+
return ""
|
|
50
|
+
if tool == "codex":
|
|
51
|
+
return _codex_thread(session_id, "rollout_path")
|
|
52
|
+
claude = Path(os.environ.get("CLAUDE_CONFIG_DIR", "~/.claude")).expanduser()
|
|
53
|
+
hits = list(claude.glob(f"projects/*/{session_id}.jsonl"))
|
|
54
|
+
return str(hits[0]) if hits else ""
|
|
55
|
+
|
|
56
|
+
|
|
57
|
+
def title(tool, session_id, transcript=None):
|
|
58
|
+
if tool == "codex":
|
|
59
|
+
name = _codex_title(session_id)
|
|
60
|
+
elif transcript:
|
|
61
|
+
name = _claude_title(transcript)
|
|
62
|
+
else:
|
|
63
|
+
name = ""
|
|
64
|
+
return " ".join(str(name).split())[:120]
|
agent_memory/store.py
ADDED
|
@@ -0,0 +1,97 @@
|
|
|
1
|
+
"""Reading memory files from the data repository."""
|
|
2
|
+
|
|
3
|
+
import json
|
|
4
|
+
from pathlib import Path
|
|
5
|
+
|
|
6
|
+
from . import frontmatter
|
|
7
|
+
from .util import SCOPES, TYPES, atomic_write
|
|
8
|
+
|
|
9
|
+
MEMORIES = "memories"
|
|
10
|
+
STAGING = "staging"
|
|
11
|
+
CACHE = ".cache"
|
|
12
|
+
|
|
13
|
+
|
|
14
|
+
def memories_dir(data_dir):
|
|
15
|
+
return Path(data_dir) / MEMORIES
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
def staging_dir(data_dir):
|
|
19
|
+
return Path(data_dir) / STAGING
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
def first_line(body):
|
|
23
|
+
for line in body.splitlines():
|
|
24
|
+
line = line.strip().lstrip("#").strip()
|
|
25
|
+
if line:
|
|
26
|
+
return line
|
|
27
|
+
return ""
|
|
28
|
+
|
|
29
|
+
|
|
30
|
+
def normalize(meta, body, fallback_id):
|
|
31
|
+
"""Fill defaults so callers can rely on the core fields."""
|
|
32
|
+
legacy = meta.get("metadata") if isinstance(meta.get("metadata"), dict) else {}
|
|
33
|
+
typ = str(meta.get("type") or legacy.get("type") or "project").lower()
|
|
34
|
+
scope = str(meta.get("scope") or "here").lower()
|
|
35
|
+
source = meta.get("source") if isinstance(meta.get("source"), dict) else {}
|
|
36
|
+
supersedes = meta.get("supersedes") or []
|
|
37
|
+
if isinstance(supersedes, str):
|
|
38
|
+
supersedes = [s.strip() for s in supersedes.split(",") if s.strip()]
|
|
39
|
+
return {
|
|
40
|
+
"id": str(meta.get("id") or fallback_id),
|
|
41
|
+
"name": str(meta.get("name") or ""),
|
|
42
|
+
"type": typ if typ in TYPES else "project",
|
|
43
|
+
"scope": scope if scope in SCOPES else "here",
|
|
44
|
+
"description": str(meta.get("description") or first_line(body))[:300],
|
|
45
|
+
"created": str(meta.get("created") or ""),
|
|
46
|
+
"updated": str(meta.get("updated") or ""),
|
|
47
|
+
"supersedes": [str(s) for s in supersedes],
|
|
48
|
+
"source": {k: str(v) for k, v in source.items() if v not in (None, "")},
|
|
49
|
+
}
|
|
50
|
+
|
|
51
|
+
|
|
52
|
+
def read(path):
|
|
53
|
+
text = Path(path).read_text("utf-8", errors="replace")
|
|
54
|
+
meta, body = frontmatter.split(text)
|
|
55
|
+
return meta, body
|
|
56
|
+
|
|
57
|
+
|
|
58
|
+
def load_all(data_dir):
|
|
59
|
+
"""All memories as normalized dicts with a `rel` path, using an mtime cache."""
|
|
60
|
+
mdir = memories_dir(data_dir)
|
|
61
|
+
cache_file = Path(data_dir) / CACHE / "index.json"
|
|
62
|
+
try:
|
|
63
|
+
cache = json.loads(cache_file.read_text("utf-8"))
|
|
64
|
+
except (OSError, ValueError):
|
|
65
|
+
cache = {}
|
|
66
|
+
fresh, out, dirty = {}, [], False
|
|
67
|
+
for path in sorted(mdir.rglob("*.md")) if mdir.is_dir() else []:
|
|
68
|
+
rel = path.relative_to(mdir).as_posix()
|
|
69
|
+
try:
|
|
70
|
+
st = path.stat()
|
|
71
|
+
except OSError:
|
|
72
|
+
continue
|
|
73
|
+
key = [st.st_mtime_ns, st.st_size]
|
|
74
|
+
entry = cache.get(rel)
|
|
75
|
+
if not entry or entry.get("k") != key:
|
|
76
|
+
try:
|
|
77
|
+
meta, body = read(path)
|
|
78
|
+
except OSError:
|
|
79
|
+
continue
|
|
80
|
+
entry = {"k": key, "m": normalize(meta, body, path.stem)}
|
|
81
|
+
dirty = True
|
|
82
|
+
fresh[rel] = entry
|
|
83
|
+
out.append({**entry["m"], "rel": rel})
|
|
84
|
+
if dirty or len(fresh) != len(cache):
|
|
85
|
+
try:
|
|
86
|
+
atomic_write(cache_file, json.dumps(fresh, ensure_ascii=False))
|
|
87
|
+
except OSError:
|
|
88
|
+
pass
|
|
89
|
+
return out
|
|
90
|
+
|
|
91
|
+
|
|
92
|
+
def find(data_dir, key):
|
|
93
|
+
"""Look up a memory by id, name or relative path."""
|
|
94
|
+
for m in load_all(data_dir):
|
|
95
|
+
if key in (m["id"], m["name"], m["rel"]) or m["rel"].endswith("/" + key):
|
|
96
|
+
return m
|
|
97
|
+
return None
|
agent_memory/ui.py
ADDED
|
@@ -0,0 +1,410 @@
|
|
|
1
|
+
"""`mem ui`: the data behind the local page for searching, reading and editing memories.
|
|
2
|
+
|
|
3
|
+
Similar memories and search by meaning come from a local embedding model
|
|
4
|
+
(Qwen3-Embedding-0.6B, int8 ONNX, run in an isolated `uv` environment) when
|
|
5
|
+
it is present, and from keyword vectors (TF-IDF over English words,
|
|
6
|
+
identifiers and Chinese words) otherwise.
|
|
7
|
+
"""
|
|
8
|
+
|
|
9
|
+
import datetime as _dt
|
|
10
|
+
import hashlib
|
|
11
|
+
import json
|
|
12
|
+
import math
|
|
13
|
+
import os
|
|
14
|
+
import re
|
|
15
|
+
import shutil
|
|
16
|
+
import subprocess
|
|
17
|
+
import sys
|
|
18
|
+
import threading
|
|
19
|
+
import urllib.request
|
|
20
|
+
from collections import Counter
|
|
21
|
+
from pathlib import Path
|
|
22
|
+
|
|
23
|
+
import numpy as np
|
|
24
|
+
|
|
25
|
+
from . import config, lint, store
|
|
26
|
+
from .util import HOME
|
|
27
|
+
|
|
28
|
+
WORKER = Path(__file__).with_name("embed_worker.py")
|
|
29
|
+
|
|
30
|
+
# The embedding model, pinned to one revision and checked by hash.
|
|
31
|
+
MODEL = {
|
|
32
|
+
"name": "Qwen3-Embedding-0.6B (int8 ONNX)",
|
|
33
|
+
"dir": "qwen3-embedding-0.6b-int8",
|
|
34
|
+
"repo": "onnx-community/Qwen3-Embedding-0.6B-ONNX",
|
|
35
|
+
"revision": "c25a394dd583836952667c12f008335071b3f43d",
|
|
36
|
+
"files": {"onnx/model_quantized.onnx": "87cd124e0ef1fd1f223ebc283efccbaeac386d0b08344701c46975d0657b591f",
|
|
37
|
+
"config.json": None, "tokenizer.json": None, "tokenizer_config.json": None,
|
|
38
|
+
"special_tokens_map.json": None},
|
|
39
|
+
# Plain text, no instruction prefix: on known same-topic pairs (including Chinese/English
|
|
40
|
+
# ones) this ranked partners closer than the topic or similarity instructions did.
|
|
41
|
+
"instruction": "",
|
|
42
|
+
}
|
|
43
|
+
WORKER_DEPS = ["onnxruntime==1.31.0", "tokenizers==0.23.3", "numpy"]
|
|
44
|
+
EMBED_CHARS = 2000 # characters of each memory given to the model
|
|
45
|
+
TITLE_WEIGHT = 3 # description and name count this many times in keyword vectors
|
|
46
|
+
NEIGHBORS = 6 # similar memories listed for each memory
|
|
47
|
+
KEYWORDS = 4 # keywords shown for each memory
|
|
48
|
+
SEARCH_INSTRUCTION = "Instruct: Given a search query, retrieve memory notes about it\nQuery: "
|
|
49
|
+
|
|
50
|
+
_ASCII = re.compile(r"[A-Za-z][A-Za-z0-9_+#-]*[A-Za-z0-9+#]|[A-Za-z]")
|
|
51
|
+
_CJK = re.compile(r"[㐀-鿿]+")
|
|
52
|
+
_STOP = set("""
|
|
53
|
+
a about above after again against all also am an and any are as at be because been before being
|
|
54
|
+
below between both but by can could did do does doing done down during each either else etc even
|
|
55
|
+
ever every few for from further get gets got had has have having he her here hers him his how i if
|
|
56
|
+
in into is it its itself just keep kept let like made make makes many may me might more most much
|
|
57
|
+
must my no nor not now of off often on once one only or other our out over own per put rather re
|
|
58
|
+
same see seen she should since so some such than that the their them then there these they this
|
|
59
|
+
those though through thus to too under until up upon us use used uses using very via was way we
|
|
60
|
+
well were what when where whether which while who whom whose why will with within without would
|
|
61
|
+
yes yet you your also first new next last still back via vs eg ie non
|
|
62
|
+
md json txt yaml yml py js ts html css sh com org www http https github io dev
|
|
63
|
+
task tasks memory memories user users agent agents codex claude session sessions file files
|
|
64
|
+
when what thing things work works working worked step steps run runs running ran note notes
|
|
65
|
+
group scope cwd thread_id updated_at rollout rollout_path rollout_summary
|
|
66
|
+
rollout_summary_files rollout_summaries applies_to reuse_rule keywords reusable knowledge failures
|
|
67
|
+
differently preferences apply how why outcome success partial history summary
|
|
68
|
+
asked plus real never check checks verify verified evidence explicit said complete final current
|
|
69
|
+
live state local source remote text path test symptom instead before after already actual ad-hoc
|
|
70
|
+
""".split())
|
|
71
|
+
# Paths, URLs, hashes and ids say where something lives, not what it is about.
|
|
72
|
+
_NOISE = re.compile(r"(?:https?://\S+|(?:~|/|\.\.?/)[\w.@%+-]+(?:/[\w.@%+-]*)+"
|
|
73
|
+
r"|\b[\w.-]+\.(?:md|json|jsonl|py|ts|js|html|yaml|yml|toml|txt|sh)\b"
|
|
74
|
+
r"|\b[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}\b|\b[0-9a-f]{7,40}\b"
|
|
75
|
+
r"|\b\d{4}-\d\d-\d\d(?:[T ][\d:.+-]+)?\b)")
|
|
76
|
+
CJK_MAX = 4
|
|
77
|
+
_STOP_CJK = set("""
|
|
78
|
+
的是 是的 一个 这个 那个 可以 需要 如果 因为 所以 但是 而且 或者 以及 我们 你们 他们 用户 时候 没有 不是 就是
|
|
79
|
+
已经 进行 通过 使用 这样 这些 那些 什么 怎么 为了 之后 之前 然后 还是 一下 一些 这里 那里 其中 以后 以前 不要
|
|
80
|
+
必须 应该 能够 只是 也是 都是 就会 会在 在这 这次 每次 一次 当前 目前 现在 之类 等等 方面 问题 内容 情况
|
|
81
|
+
相关 进行 这种 时候 需要 不能 不会 只有 直接 所有 其他 自己 一样 部分 以上 以下 下来 起来 出来 每个 各个 任何 一定 先给
|
|
82
|
+
""".split())
|
|
83
|
+
# Function characters that start no word: 的了和与… A string beginning with one is a fragment.
|
|
84
|
+
_CJK_BAD_START = set("的了和与及或就都也把被让给并而时在是对向从个件等上不每一这那该各再又还才只很最更已将会能要可使于以之其此")
|
|
85
|
+
_CJK_KEEP = {"完成", "生成", "集成", "构成", "组成", "前后", "最后", "然后", "之后", "背后", "上一", "中文"}
|
|
86
|
+
# ...and characters that end no word: a string ending in one runs into the next word.
|
|
87
|
+
_CJK_BAD_END = set("的了着过是在和与及等上中里成到给把被让后不也都就又再才很一个这那各每将会能要可于以之其此任")
|
|
88
|
+
|
|
89
|
+
|
|
90
|
+
# ------------------------------------------------------------ keyword vectors
|
|
91
|
+
|
|
92
|
+
def clean(text):
|
|
93
|
+
return _NOISE.sub(" ", text)
|
|
94
|
+
|
|
95
|
+
|
|
96
|
+
def chinese_words(texts):
|
|
97
|
+
"""Recurring Chinese words, found without a dictionary.
|
|
98
|
+
|
|
99
|
+
Takes every 2-4 character string that appears in at least two texts, then
|
|
100
|
+
drops fragments: a string that nearly always occurs inside one longer
|
|
101
|
+
string (e.g. 京时 inside 北京时间) is part of that word, not a word.
|
|
102
|
+
"""
|
|
103
|
+
df = Counter()
|
|
104
|
+
for text in texts:
|
|
105
|
+
grams = set()
|
|
106
|
+
for run in _CJK.findall(text):
|
|
107
|
+
for n in range(2, CJK_MAX + 1):
|
|
108
|
+
grams.update(run[i:i + n] for i in range(len(run) - n + 1))
|
|
109
|
+
df.update(grams)
|
|
110
|
+
frequent = {g for g, d in df.items() if d >= 2}
|
|
111
|
+
words = set()
|
|
112
|
+
for g in frequent:
|
|
113
|
+
longer = [h for h in frequent if len(h) == len(g) + 1 and g in h]
|
|
114
|
+
if not any(df[h] >= 0.8 * df[g] for h in longer):
|
|
115
|
+
words.add(g)
|
|
116
|
+
return {w for w in words - _STOP_CJK
|
|
117
|
+
if (w[0] not in _CJK_BAD_START or w == "时间") and (w[-1] not in _CJK_BAD_END or w in _CJK_KEEP)}
|
|
118
|
+
|
|
119
|
+
|
|
120
|
+
def tokens(text, words=frozenset()):
|
|
121
|
+
"""Lowercased English words and identifiers, plus Chinese words.
|
|
122
|
+
|
|
123
|
+
Chinese runs are split by longest match against `words`; characters that
|
|
124
|
+
start no known word are skipped.
|
|
125
|
+
"""
|
|
126
|
+
text = clean(text)
|
|
127
|
+
out = []
|
|
128
|
+
for w in _ASCII.findall(text):
|
|
129
|
+
low = w.lower().strip("-_")
|
|
130
|
+
if len(low) < 2 or low in _STOP or low.isdigit():
|
|
131
|
+
continue
|
|
132
|
+
if len(low) == 2 and not w.isupper(): # keep acronyms like S3, UI
|
|
133
|
+
continue
|
|
134
|
+
out.append(low)
|
|
135
|
+
for run in _CJK.findall(text):
|
|
136
|
+
i = 0
|
|
137
|
+
while i < len(run) - 1:
|
|
138
|
+
for n in range(min(CJK_MAX, len(run) - i), 1, -1):
|
|
139
|
+
if run[i:i + n] in words:
|
|
140
|
+
out.append(run[i:i + n])
|
|
141
|
+
i += n
|
|
142
|
+
break
|
|
143
|
+
else:
|
|
144
|
+
i += 1
|
|
145
|
+
return out
|
|
146
|
+
|
|
147
|
+
|
|
148
|
+
def _display_forms(texts):
|
|
149
|
+
"""Most common original spelling of each token (GPU, not gpu), for showing keywords."""
|
|
150
|
+
seen = {}
|
|
151
|
+
for text in texts:
|
|
152
|
+
for w in _ASCII.findall(text):
|
|
153
|
+
seen.setdefault(w.lower().strip("-_"), Counter())[w] += 1
|
|
154
|
+
return {k: c.most_common(1)[0][0] for k, c in seen.items()}
|
|
155
|
+
|
|
156
|
+
|
|
157
|
+
# Plain English words that score high in TF-IDF but say little about what a memory is about.
|
|
158
|
+
_GENERIC = set("""
|
|
159
|
+
actually specified large small nontrivial declaring declared truth human builds build built pin pinned
|
|
160
|
+
following follow follows needed need needs means mean simple simply clear clearly important exact exactly
|
|
161
|
+
whole entire full fully main mainly certain directly direct proper properly various several general
|
|
162
|
+
correct correctly wrong instead later earlier often usually rather quite really already able unable
|
|
163
|
+
item items part parts kind way ways case cases point points detail details thing result results
|
|
164
|
+
come comes edit edits edited anything something everything separate change changes changed
|
|
165
|
+
fail fails failed code example examples
|
|
166
|
+
""".split())
|
|
167
|
+
|
|
168
|
+
|
|
169
|
+
def _term_like(display):
|
|
170
|
+
"""Acronyms, identifiers and Chinese words make better keywords than plain English.
|
|
171
|
+
|
|
172
|
+
A capital only at the start (Choose, Pack) is usually the start of a sentence, so it
|
|
173
|
+
does not count; one further in (SigLIP, GPU, OpenReview) does.
|
|
174
|
+
"""
|
|
175
|
+
return bool(re.search(r".[A-Z]|[0-9_.+#-]|[\u3400-\u9fff]", display))
|
|
176
|
+
|
|
177
|
+
|
|
178
|
+
def keyword_picks(x, vocab, forms, k=KEYWORDS):
|
|
179
|
+
"""Column indices of each memory's keywords, best first."""
|
|
180
|
+
skip = _GENERIC | {w.lower() for w in config.load().get("keyword_stopwords", [])}
|
|
181
|
+
weight = np.array([0 if t in skip else (1.6 if _term_like(forms.get(t, t)) else 1.0) for t in vocab])
|
|
182
|
+
scored = x * weight
|
|
183
|
+
return [[int(j) for j in np.argsort(-row)[:k] if row[j] > 0] for row in scored]
|
|
184
|
+
|
|
185
|
+
|
|
186
|
+
def vectorize(docs):
|
|
187
|
+
"""TF-IDF rows (L2-normalized) and the vocabulary.
|
|
188
|
+
|
|
189
|
+
docs: list of (title, body). Title tokens count TITLE_WEIGHT times.
|
|
190
|
+
Terms in one document only carry no similarity and are dropped, as are
|
|
191
|
+
terms in more than half of the documents.
|
|
192
|
+
"""
|
|
193
|
+
words = chinese_words(clean(t + " " + b) for t, b in docs)
|
|
194
|
+
counts = [Counter(tokens(t, words) * TITLE_WEIGHT + tokens(b, words)) for t, b in docs]
|
|
195
|
+
df = Counter(term for c in counts for term in c)
|
|
196
|
+
n = len(docs)
|
|
197
|
+
vocab = sorted(t for t, d in df.items() if 2 <= d <= max(2, n // 2))
|
|
198
|
+
index = {t: i for i, t in enumerate(vocab)}
|
|
199
|
+
x = np.zeros((n, len(vocab)))
|
|
200
|
+
for row, c in enumerate(counts):
|
|
201
|
+
for term, tf in c.items():
|
|
202
|
+
j = index.get(term)
|
|
203
|
+
if j is not None:
|
|
204
|
+
x[row, j] = (1 + math.log(tf)) * math.log((1 + n) / (1 + df[term]))
|
|
205
|
+
norms = np.linalg.norm(x, axis=1, keepdims=True)
|
|
206
|
+
return x / np.where(norms == 0, 1, norms), vocab
|
|
207
|
+
|
|
208
|
+
|
|
209
|
+
# ------------------------------------------------------------ meaning vectors
|
|
210
|
+
|
|
211
|
+
_TEMPLATE_LINE = re.compile(r"^\s*(?:-\s*)?rollout_summaries/|^\s*#+\s*(?:rollout_summary_files|keywords)\s*$", re.I)
|
|
212
|
+
|
|
213
|
+
|
|
214
|
+
def prose(title, body):
|
|
215
|
+
"""What the model reads: the description and body without bookkeeping lines."""
|
|
216
|
+
lines = [line for line in body.splitlines() if not _TEMPLATE_LINE.search(line)]
|
|
217
|
+
text = clean(f"{title}\n" + "\n".join(lines))
|
|
218
|
+
text = re.sub(r"[#*`>|]+", " ", text)
|
|
219
|
+
return re.sub(r"[ \t]+", " ", text).strip()[:EMBED_CHARS]
|
|
220
|
+
|
|
221
|
+
|
|
222
|
+
def model_dir():
|
|
223
|
+
custom = config.load().get("embed_model_dir")
|
|
224
|
+
return Path(custom).expanduser() if custom else Path("~/.cache/agent-memory/models").expanduser() / MODEL["dir"]
|
|
225
|
+
|
|
226
|
+
|
|
227
|
+
def model_ready():
|
|
228
|
+
d = model_dir()
|
|
229
|
+
return (d / "model_quantized.onnx").exists() and (d / "tokenizer.json").exists() and bool(shutil.which("uv"))
|
|
230
|
+
|
|
231
|
+
|
|
232
|
+
def fetch_model():
|
|
233
|
+
"""Download the pinned model files (about 625 MB) and check the ONNX hash."""
|
|
234
|
+
d = model_dir()
|
|
235
|
+
d.mkdir(parents=True, exist_ok=True)
|
|
236
|
+
base = f"https://huggingface.co/{MODEL['repo']}/resolve/{MODEL['revision']}"
|
|
237
|
+
for rel, sha in MODEL["files"].items():
|
|
238
|
+
target = d / Path(rel).name
|
|
239
|
+
if target.exists() and (sha is None or _sha256(target) == sha):
|
|
240
|
+
continue
|
|
241
|
+
print(f"downloading {rel}…", flush=True)
|
|
242
|
+
part = target.with_suffix(target.suffix + ".part")
|
|
243
|
+
urllib.request.urlretrieve(f"{base}/{rel}", part)
|
|
244
|
+
if sha and _sha256(part) != sha:
|
|
245
|
+
part.unlink()
|
|
246
|
+
raise SystemExit(f"{rel}: checksum mismatch, download discarded")
|
|
247
|
+
part.replace(target)
|
|
248
|
+
(d / "REVISION").write_text(MODEL["revision"] + "\n")
|
|
249
|
+
print(f"model ready in {d}")
|
|
250
|
+
|
|
251
|
+
|
|
252
|
+
def _sha256(path):
|
|
253
|
+
h = hashlib.sha256()
|
|
254
|
+
with open(path, "rb") as f:
|
|
255
|
+
for chunk in iter(lambda: f.read(1 << 20), b""):
|
|
256
|
+
h.update(chunk)
|
|
257
|
+
return h.hexdigest()
|
|
258
|
+
|
|
259
|
+
|
|
260
|
+
def embed(data_dir, texts):
|
|
261
|
+
"""Unit vectors for texts, reusing cached ones; None if the model is unavailable."""
|
|
262
|
+
if not model_ready():
|
|
263
|
+
return None
|
|
264
|
+
cache_file = Path(data_dir) / ".cache" / "embeddings" / f"{MODEL['dir']}.npz"
|
|
265
|
+
keys = [hashlib.sha1((MODEL["revision"] + MODEL["instruction"] + t).encode()).hexdigest() for t in texts]
|
|
266
|
+
known = {}
|
|
267
|
+
if cache_file.exists():
|
|
268
|
+
with np.load(cache_file) as z:
|
|
269
|
+
known = dict(zip(z["keys"].tolist(), z["vectors"]))
|
|
270
|
+
todo = [i for i, k in enumerate(keys) if k not in known]
|
|
271
|
+
if todo:
|
|
272
|
+
out = cache_file.with_name("new-vectors.npy")
|
|
273
|
+
out.parent.mkdir(parents=True, exist_ok=True)
|
|
274
|
+
job = {"model_dir": str(model_dir()), "out": str(out), "max_tokens": 512,
|
|
275
|
+
"texts": [MODEL["instruction"] + texts[i] for i in todo]}
|
|
276
|
+
print(f"embedding {len(todo)} memories with {MODEL['name']}…", flush=True)
|
|
277
|
+
cmd = ["uv", "run", "--no-project", "--quiet"] + [a for dep in WORKER_DEPS for a in ("--with", dep)] \
|
|
278
|
+
+ ["python", "-I", str(WORKER)]
|
|
279
|
+
r = subprocess.run(cmd, input=json.dumps(job), text=True, cwd=str(out.parent),
|
|
280
|
+
env={**os.environ, "PYTHONWARNINGS": "ignore"})
|
|
281
|
+
if r.returncode != 0:
|
|
282
|
+
print("embedding failed; falling back to keywords", file=sys.stderr)
|
|
283
|
+
return None
|
|
284
|
+
for i, vec in zip(todo, np.load(out)):
|
|
285
|
+
known[keys[i]] = vec
|
|
286
|
+
out.unlink()
|
|
287
|
+
live = {k: known[k] for k in keys} # drop vectors of edited or deleted memories
|
|
288
|
+
np.savez(cache_file, keys=np.array(list(live)), vectors=np.array(list(live.values())))
|
|
289
|
+
return np.array([known[k] for k in keys])
|
|
290
|
+
|
|
291
|
+
|
|
292
|
+
def closeness(vectors):
|
|
293
|
+
"""Cosine similarity rescaled so a typical pair is 0 and the closest pair is 1.
|
|
294
|
+
|
|
295
|
+
Embedding cosines bunch together (most unrelated pairs still score 0.3-0.5),
|
|
296
|
+
so the raw number says little; the rescaled one reads like a percentage.
|
|
297
|
+
"""
|
|
298
|
+
raw = vectors @ vectors.T
|
|
299
|
+
off = raw[~np.eye(len(raw), dtype=bool)]
|
|
300
|
+
if not len(off):
|
|
301
|
+
return np.ones_like(raw)
|
|
302
|
+
lo, hi = np.percentile(off, 50), off.max()
|
|
303
|
+
sim = np.clip((raw - lo) / max(hi - lo, 1e-6), 0, 1)
|
|
304
|
+
np.fill_diagonal(sim, 1)
|
|
305
|
+
return sim
|
|
306
|
+
|
|
307
|
+
|
|
308
|
+
# ------------------------------------------------------------ page data
|
|
309
|
+
|
|
310
|
+
def _short(path):
|
|
311
|
+
home = str(HOME)
|
|
312
|
+
return "~" + path[len(home):] if path and (path == home or path.startswith(home + "/")) else path or ""
|
|
313
|
+
|
|
314
|
+
|
|
315
|
+
def _place(src):
|
|
316
|
+
if src.get("git_remote"):
|
|
317
|
+
return src["git_remote"]
|
|
318
|
+
return _short(src.get("git_root") or src.get("cwd") or "")
|
|
319
|
+
|
|
320
|
+
|
|
321
|
+
def build(data_dir, semantic=True, vectors_out=None):
|
|
322
|
+
"""Every memory with its whole file, and the ones most similar to it."""
|
|
323
|
+
now = _dt.datetime.now().astimezone()
|
|
324
|
+
mdir = store.memories_dir(data_dir)
|
|
325
|
+
items, docs = [], []
|
|
326
|
+
for m in store.load_all(data_dir):
|
|
327
|
+
try:
|
|
328
|
+
raw = (mdir / m["rel"]).read_text("utf-8", errors="replace")
|
|
329
|
+
except OSError:
|
|
330
|
+
continue
|
|
331
|
+
meta, body = store.read(mdir / m["rel"])
|
|
332
|
+
items.append((m, raw, meta, body))
|
|
333
|
+
docs.append((f"{m['description']} {m['name'].replace('-', ' ')}", body))
|
|
334
|
+
base = {"generated": now.isoformat(timespec="minutes"), "dataDir": _short(str(data_dir))}
|
|
335
|
+
if not items:
|
|
336
|
+
return {**base, "memories": [], "places": [], "method": "", "semantic": False}
|
|
337
|
+
|
|
338
|
+
# keywords: the words that most set each memory apart (TF-IDF), shared with at least one other
|
|
339
|
+
x, vocab = vectorize(docs)
|
|
340
|
+
forms = _display_forms(clean(t + " " + b) for t, b in docs)
|
|
341
|
+
vectors = embed(data_dir, [prose(t, b) for t, b in docs]) if semantic else None
|
|
342
|
+
if vectors is not None:
|
|
343
|
+
sim, space = closeness(vectors), vectors
|
|
344
|
+
method = f"{MODEL['name']}, local"
|
|
345
|
+
else:
|
|
346
|
+
sim, space = np.clip(x @ x.T, 0, 1), x
|
|
347
|
+
np.fill_diagonal(sim, 1)
|
|
348
|
+
method = "keywords only (no local model)"
|
|
349
|
+
if vectors_out is not None:
|
|
350
|
+
vectors_out[:] = [space]
|
|
351
|
+
|
|
352
|
+
picks = keyword_picks(x, vocab, forms)
|
|
353
|
+
names = {m["name"] for m, _, _, _ in items if m["name"]}
|
|
354
|
+
rows, places = [], {}
|
|
355
|
+
for i, (m, raw, meta, body) in enumerate(items):
|
|
356
|
+
order = [int(j) for j in np.argsort(-sim[i]) if j != i][:NEIGHBORS]
|
|
357
|
+
src = m["source"]
|
|
358
|
+
rows.append({
|
|
359
|
+
"id": m["id"], "rel": m["rel"], "path": str(mdir / m["rel"]),
|
|
360
|
+
"hash": hashlib.sha1(raw.encode()).hexdigest()[:16], "raw": raw,
|
|
361
|
+
"meta": json.loads(json.dumps(meta, default=str)), "body": body,
|
|
362
|
+
"description": m["description"], "type": m["type"], "scope": m["scope"],
|
|
363
|
+
"place": _place(src) if m["scope"] == "here" else "",
|
|
364
|
+
"tool": src.get("tool") or "", "host": src.get("host") or "",
|
|
365
|
+
"created": m["created"], "updated": m["updated"],
|
|
366
|
+
"similar": [{"i": j, "s": round(float(sim[i, j]), 3)} for j in order if sim[i, j] > 0],
|
|
367
|
+
"keywords": [{"k": vocab[j], "t": forms.get(vocab[j], vocab[j])} for j in picks[i]],
|
|
368
|
+
"terms": [vocab[j] for j in np.flatnonzero(x[i])],
|
|
369
|
+
"issues": [message for _, message in lint.check(m, body, names)],
|
|
370
|
+
})
|
|
371
|
+
if m["scope"] == "here" and (src.get("git_remote") or src.get("cwd")):
|
|
372
|
+
key = _place(src)
|
|
373
|
+
places.setdefault(key, {"key": key, "count": 0,
|
|
374
|
+
"ctx": {k: src[k] for k in ("cwd", "git_root", "git_remote", "git_branch", "host")
|
|
375
|
+
if src.get(k)}})
|
|
376
|
+
places[key]["count"] += 1
|
|
377
|
+
return {**base, "memories": rows, "method": method, "semantic": vectors is not None,
|
|
378
|
+
"places": sorted(places.values(), key=lambda p: (-p["count"], p["key"]))}
|
|
379
|
+
|
|
380
|
+
|
|
381
|
+
# ------------------------------------------------------------ query embedding
|
|
382
|
+
|
|
383
|
+
class QueryEncoder:
|
|
384
|
+
"""A long-lived embedding worker for search queries (model loads once)."""
|
|
385
|
+
|
|
386
|
+
def __init__(self):
|
|
387
|
+
self.proc = None
|
|
388
|
+
self.lock = threading.Lock()
|
|
389
|
+
|
|
390
|
+
def encode(self, text):
|
|
391
|
+
if not model_ready():
|
|
392
|
+
return None
|
|
393
|
+
with self.lock:
|
|
394
|
+
if self.proc is None or self.proc.poll() is not None:
|
|
395
|
+
cmd = ["uv", "run", "--no-project", "--quiet"] + \
|
|
396
|
+
[a for dep in WORKER_DEPS for a in ("--with", dep)] + \
|
|
397
|
+
["python", "-I", str(WORKER), "--serve", str(model_dir())]
|
|
398
|
+
self.proc = subprocess.Popen(cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE,
|
|
399
|
+
stderr=subprocess.DEVNULL, text=True, bufsize=1,
|
|
400
|
+
cwd=str(model_dir()))
|
|
401
|
+
self.proc.stdin.write(json.dumps({"text": SEARCH_INSTRUCTION + text}) + "\n")
|
|
402
|
+
self.proc.stdin.flush()
|
|
403
|
+
line = self.proc.stdout.readline()
|
|
404
|
+
if not line:
|
|
405
|
+
return None
|
|
406
|
+
return np.array(json.loads(line)["vector"], dtype=np.float32)
|
|
407
|
+
|
|
408
|
+
def close(self):
|
|
409
|
+
if self.proc and self.proc.poll() is None:
|
|
410
|
+
self.proc.terminate()
|