personal memory agent
Add --targeted flag to sol agents check for cost reduction
Automated health rechecks (from cogitate/generate agents) now use
--targeted to check only configured provider+tier pairs (~6 instead
of 18). Adds flock-based dedup so concurrent rechecks don't overlap,
and logs token usage for generate health checks.
+197
-1
+149
tests/test_agents_check.py
+149
tests/test_agents_check.py
···
39
39
tier=None,
40
40
json=False,
41
41
timeout=1,
42
+
targeted=False,
42
43
)
43
44
44
45
with pytest.raises(SystemExit) as exc_info:
···
86
87
tier=None,
87
88
json=False,
88
89
timeout=1,
90
+
targeted=False,
89
91
)
90
92
91
93
with pytest.raises(SystemExit) as exc_info:
···
130
132
tier=None,
131
133
json=False,
132
134
timeout=1,
135
+
targeted=False,
133
136
)
134
137
135
138
with pytest.raises(SystemExit) as exc_info:
···
176
179
tier=None,
177
180
json=False,
178
181
timeout=1,
182
+
targeted=False,
179
183
)
180
184
181
185
with pytest.raises(SystemExit) as exc_info:
···
200
204
assert len(reused) == 4
201
205
assert all(result["reused_from"] == "pro" for result in reused)
202
206
assert all(result["elapsed_s"] == 0.0 for result in reused)
207
+
208
+
209
+
def test_run_check_targeted_filters_to_configured_pairs(tmp_path, monkeypatch):
210
+
"""--targeted filters checks to only configured provider+tier pairs."""
211
+
import think.agents as agents
212
+
213
+
fake_registry = {"provA": object(), "provB": object(), "provC": object()}
214
+
fake_defaults = {
215
+
"provA": {1: "a-pro", 2: "a-flash", 3: "a-lite"},
216
+
"provB": {1: "b-pro", 2: "b-flash", 3: "b-lite"},
217
+
"provC": {1: "c-pro", 2: "c-flash", 3: "c-lite"},
218
+
}
219
+
fake_type_defaults = {
220
+
"generate": {"provider": "provA", "tier": 2, "backup": "provB"},
221
+
"cogitate": {"provider": "provC", "tier": 2, "backup": "provB"},
222
+
}
223
+
224
+
monkeypatch.setattr("think.providers.PROVIDER_REGISTRY", fake_registry)
225
+
monkeypatch.setattr("think.models.PROVIDER_DEFAULTS", fake_defaults)
226
+
monkeypatch.setattr("think.models.TYPE_DEFAULTS", fake_type_defaults)
227
+
monkeypatch.setattr(agents, "get_journal", lambda: str(tmp_path))
228
+
monkeypatch.setattr(agents, "_check_generate", lambda *_args: (True, "ok"))
229
+
230
+
async def mock_check_cogitate(*_args):
231
+
return True, "ok"
232
+
233
+
monkeypatch.setattr(agents, "_check_cogitate", mock_check_cogitate)
234
+
235
+
# Mock get_config to return no overrides (use TYPE_DEFAULTS)
236
+
monkeypatch.setattr("think.utils.get_config", lambda: {})
237
+
238
+
# Mock get_backup_provider to return the backup from fake_type_defaults
239
+
def fake_get_backup(agent_type):
240
+
d = fake_type_defaults[agent_type]
241
+
if d["backup"] == d["provider"]:
242
+
return None
243
+
return d["backup"]
244
+
245
+
monkeypatch.setattr("think.models.get_backup_provider", fake_get_backup)
246
+
247
+
args = argparse.Namespace(
248
+
provider=None,
249
+
interface=None,
250
+
tier=None,
251
+
json=True,
252
+
timeout=1,
253
+
targeted=True,
254
+
)
255
+
256
+
with pytest.raises(SystemExit) as exc_info:
257
+
asyncio.run(agents._run_check(args))
258
+
259
+
assert exc_info.value.code == 0
260
+
261
+
health_file = tmp_path / "health" / "agents.json"
262
+
payload = json.loads(health_file.read_text())
263
+
# Expected targeted pairs: (provA, 2), (provB, 2), (provC, 2) = 3 pairs × 2 interfaces = 6 checks
264
+
assert payload["summary"]["total"] == 6
265
+
checked_pairs = {(r["provider"], r["tier"]) for r in payload["results"]}
266
+
assert checked_pairs == {("provA", "flash"), ("provB", "flash"), ("provC", "flash")}
267
+
268
+
269
+
def test_run_check_targeted_flock_dedup(tmp_path, monkeypatch):
270
+
"""--targeted exits silently when another targeted check holds the lock."""
271
+
import fcntl
272
+
273
+
import think.agents as agents
274
+
275
+
fake_registry = {"fake": object()}
276
+
fake_defaults = {"fake": {1: "m", 2: "m", 3: "m"}}
277
+
fake_type_defaults = {
278
+
"generate": {"provider": "fake", "tier": 2, "backup": "fake"},
279
+
"cogitate": {"provider": "fake", "tier": 2, "backup": "fake"},
280
+
}
281
+
282
+
monkeypatch.setattr("think.providers.PROVIDER_REGISTRY", fake_registry)
283
+
monkeypatch.setattr("think.models.PROVIDER_DEFAULTS", fake_defaults)
284
+
monkeypatch.setattr("think.models.TYPE_DEFAULTS", fake_type_defaults)
285
+
monkeypatch.setattr(agents, "get_journal", lambda: str(tmp_path))
286
+
monkeypatch.setattr("think.utils.get_config", lambda: {})
287
+
monkeypatch.setattr("think.models.get_backup_provider", lambda _: None)
288
+
289
+
# Pre-acquire the lock to simulate a concurrent check
290
+
lock_dir = tmp_path / "health"
291
+
lock_dir.mkdir(parents=True, exist_ok=True)
292
+
lock_file = open(lock_dir / "recheck.lock", "w")
293
+
fcntl.flock(lock_file, fcntl.LOCK_EX | fcntl.LOCK_NB)
294
+
295
+
gen_mock = MagicMock(return_value=(True, "ok"))
296
+
monkeypatch.setattr(agents, "_check_generate", gen_mock)
297
+
298
+
args = argparse.Namespace(
299
+
provider=None,
300
+
interface=None,
301
+
tier=None,
302
+
json=False,
303
+
timeout=1,
304
+
targeted=True,
305
+
)
306
+
307
+
# Should return silently (no SystemExit, no checks run)
308
+
asyncio.run(agents._run_check(args))
309
+
assert gen_mock.call_count == 0
310
+
311
+
# No health file written
312
+
assert not (tmp_path / "health" / "agents.json").exists()
313
+
314
+
lock_file.close()
315
+
316
+
317
+
def test_check_generate_logs_token_usage(monkeypatch):
318
+
"""_check_generate logs token usage when result includes usage data."""
319
+
import think.agents as agents
320
+
321
+
fake_module = MagicMock()
322
+
fake_module.run_generate.return_value = {
323
+
"text": "OK",
324
+
"usage": {"input_tokens": 5, "output_tokens": 2},
325
+
}
326
+
327
+
monkeypatch.setattr(
328
+
"think.providers.get_provider_module", lambda _: fake_module
329
+
)
330
+
monkeypatch.setattr(
331
+
"think.providers.PROVIDER_METADATA",
332
+
{"fake": {"env_key": "FAKE_API_KEY"}},
333
+
)
334
+
monkeypatch.setattr(
335
+
"think.models.PROVIDER_DEFAULTS", {"fake": {2: "fake-flash"}}
336
+
)
337
+
monkeypatch.setenv("FAKE_API_KEY", "test-key")
338
+
339
+
log_mock = MagicMock()
340
+
monkeypatch.setattr("think.models.log_token_usage", log_mock)
341
+
342
+
ok, msg = agents._check_generate("fake", 2, 30)
343
+
344
+
assert ok is True
345
+
assert msg == "OK"
346
+
log_mock.assert_called_once_with(
347
+
model="fake-flash",
348
+
usage={"input_tokens": 5, "output_tokens": 2},
349
+
context="health.check.generate",
350
+
type="generate",
351
+
)
203
352
204
353
205
354
def test_cortex_start_emits_agents_check(tmp_path):
+47
think/agents.py
+47
think/agents.py
···
1205
1205
)
1206
1206
text = result.get("text", "") if isinstance(result, dict) else ""
1207
1207
if text:
1208
+
usage = result.get("usage") if isinstance(result, dict) else None
1209
+
if usage:
1210
+
from think.models import log_token_usage
1211
+
1212
+
log_token_usage(
1213
+
model=PROVIDER_DEFAULTS[provider_name][tier],
1214
+
usage=usage,
1215
+
context="health.check.generate",
1216
+
type="generate",
1217
+
)
1208
1218
return True, "OK"
1209
1219
return False, "FAIL: empty response text"
1210
1220
except Exception as exc:
···
1242
1252
1243
1253
load_dotenv()
1244
1254
1255
+
# --targeted: only check configured provider+tier pairs
1256
+
targeted_pairs = None
1257
+
if args.targeted and not args.provider and not args.tier:
1258
+
import fcntl
1259
+
1260
+
from think.models import TYPE_DEFAULTS, get_backup_provider
1261
+
from think.utils import get_config
1262
+
1263
+
targeted_pairs = set()
1264
+
config = get_config()
1265
+
providers_config = config.get("providers", {})
1266
+
for agent_type, defaults in TYPE_DEFAULTS.items():
1267
+
type_config = providers_config.get(agent_type, {})
1268
+
provider = type_config.get("provider", defaults["provider"])
1269
+
tier = type_config.get("tier", defaults["tier"])
1270
+
targeted_pairs.add((provider, tier))
1271
+
backup = get_backup_provider(agent_type)
1272
+
if backup:
1273
+
targeted_pairs.add((backup, tier))
1274
+
1275
+
# flock dedup: only one targeted check runs at a time
1276
+
lock_dir = Path(get_journal()) / "health"
1277
+
lock_dir.mkdir(parents=True, exist_ok=True)
1278
+
lock_fd = open(lock_dir / "recheck.lock", "w")
1279
+
try:
1280
+
fcntl.flock(lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
1281
+
except OSError:
1282
+
lock_fd.close()
1283
+
return
1284
+
1245
1285
if args.provider:
1246
1286
providers = args.provider
1247
1287
for name in providers:
···
1279
1319
1280
1320
for provider_name in providers:
1281
1321
for tier in tiers:
1322
+
if targeted_pairs is not None and (provider_name, tier) not in targeted_pairs:
1323
+
continue
1282
1324
model = PROVIDER_DEFAULTS[provider_name][tier]
1283
1325
for interface_name in interfaces:
1284
1326
cache_key = (provider_name, model, interface_name)
···
1415
1457
)
1416
1458
check_parser.add_argument(
1417
1459
"--json", action="store_true", help="Output results as JSON"
1460
+
)
1461
+
check_parser.add_argument(
1462
+
"--targeted",
1463
+
action="store_true",
1464
+
help="Only check configured provider+tier pairs (used by automated rechecks)",
1418
1465
)
1419
1466
1420
1467
args = setup_cli(parser)