Compare commits
409
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
b323df75b6 | ||
|
|
f8e588c042 | ||
|
|
e852acc7c0 | ||
|
|
e9a14da690 | ||
|
|
3e3e3f372c | ||
|
|
ea0d5ead0a | ||
|
|
e9d1676ef3 | ||
|
|
2cf31d7fbe | ||
|
|
aefc5e8e7d | ||
|
|
18bff79aea | ||
|
|
4c3a0bf295 | ||
|
|
76a74b3181 | ||
|
|
6d61772190 | ||
|
|
d733772e39 | ||
|
|
5f1237279b | ||
|
|
3a48b6e5bb | ||
|
|
f1c825d884 | ||
|
|
4b2041eb3d | ||
|
|
cbbc864694 | ||
|
|
5b2f5d82c2 | ||
|
|
13d9f2df95 | ||
|
|
eb47374fa3 | ||
|
|
b12646f49e | ||
|
|
404fe467b6 | ||
|
|
fecacc4ea8 | ||
|
|
4d52d2974c | ||
|
|
cb4c4c06b3 | ||
|
|
79fe411101 | ||
|
|
d3946f17cd | ||
|
|
a3298cdf03 | ||
|
|
475e04d244 | ||
|
|
18a58ef3aa | ||
|
|
ad85affb13 | ||
|
|
53948468aa | ||
|
|
c15b565c80 | ||
|
|
3211952c6e | ||
|
|
fd6ed94060 | ||
|
|
8d575a9b30 | ||
|
|
b5aec64a49 | ||
|
|
f9fb3e934a | ||
|
|
afdea53a96 | ||
|
|
25a7237b1f | ||
|
|
5e557dbfc5 | ||
|
|
d637008fbe | ||
|
|
e1b3d438ef | ||
|
|
8de576b4b7 | ||
|
|
797faf7981 | ||
|
|
436bc7c156 | ||
|
|
59d3f078cd | ||
|
|
359e619a42 | ||
|
|
f6b3ce3e33 | ||
|
|
9d828cc7b1 | ||
|
|
42e8c53da9 | ||
|
|
ba5a4813b3 | ||
|
|
f2ae61eda3 | ||
|
|
468cc4b7d7 | ||
|
|
dc9d033d16 | ||
|
|
bb895faaf9 | ||
|
|
926f752912 | ||
|
|
94d65a591a | ||
|
|
aad36e5eee | ||
|
|
93fa0b016b | ||
|
|
559ab9dbff | ||
|
|
76d8ba8a51 | ||
|
|
362c719f1e | ||
|
|
c335192cd3 | ||
|
|
11ecfe54c2 | ||
|
|
751deb47be | ||
|
|
5f8a030615 | ||
|
|
f572d8647d | ||
|
|
0ca0e87a08 | ||
|
|
04b7c9a188 | ||
|
|
ebae35670e | ||
|
|
333812c85a | ||
|
|
417efb35d8 | ||
|
|
8e953c2300 | ||
|
|
cac55fedad | ||
|
|
b5548ac25c | ||
|
|
5c15f28afe | ||
|
|
f77f2188d6 | ||
|
|
efcf36aaac | ||
|
|
84fd3b1767 | ||
|
|
4b76d8be22 | ||
|
|
d238d2f7d7 | ||
|
|
4278f0989d | ||
|
|
17b7f46ae7 | ||
|
|
06e9c7054e | ||
|
|
5061ff6643 | ||
|
|
dbca379410 | ||
|
|
13cacf21a7 | ||
|
|
bda7ffd289 | ||
|
|
d6bdb62a02 | ||
|
|
0aa3480e8c | ||
|
|
05b55d20f5 | ||
|
|
4d22a882f5 | ||
|
|
2b1c4989f3 | ||
|
|
55883dc517 | ||
|
|
06eaf09493 | ||
|
|
736d1c2f7b | ||
|
|
bb26f49e93 | ||
|
|
47b679ec24 | ||
|
|
a1ec656424 | ||
|
|
cb0e1dea40 | ||
|
|
bf6c12d550 | ||
|
|
237a45fdb5 | ||
|
|
8396b51d92 | ||
|
|
8e5fdf28dd | ||
|
|
abb5ba3498 | ||
|
|
b5d8439c8f | ||
|
|
500a9e637a | ||
|
|
486a8c41d6 | ||
|
|
1ce308f359 | ||
|
|
d571199cd6 | ||
|
|
baa923debc | ||
|
|
d19f54c770 | ||
|
|
24825200b0 | ||
|
|
cd40649393 | ||
|
|
79c9f4afeb | ||
|
|
3868c8d055 | ||
|
|
c3b98998b7 | ||
|
|
a90f922376 | ||
|
|
53325898e6 | ||
|
|
b651f43f22 | ||
|
|
d105038323 | ||
|
|
bc604ab91b | ||
|
|
aa38790dd6 | ||
|
|
f61d383acf | ||
|
|
29aea56281 | ||
|
|
218e6d34b1 | ||
|
|
0f9dc55084 | ||
|
|
b769045b64 | ||
|
|
74cff93098 | ||
|
|
556c2c76c6 | ||
|
|
cc1eaeeeba | ||
|
|
6268654bf6 | ||
|
|
08995e3013 | ||
|
|
3a1841c421 | ||
|
|
94f4adfbed | ||
|
|
70c09adcf0 | ||
|
|
c5643fdf5c | ||
|
|
6a460d2c9a | ||
|
|
97b00ca75a | ||
|
|
9452ac29da | ||
|
|
40d2b7f6b8 | ||
|
|
b1a0ef5f7d | ||
|
|
4f4c2988e9 | ||
|
|
feac397324 | ||
|
|
b8e1524a17 | ||
|
|
979999651a | ||
|
|
91160f3bab | ||
|
|
2dc40f3bc7 | ||
|
|
ff7a087a37 | ||
|
|
6500944c74 | ||
|
|
8133c5ab7e | ||
|
|
9a1ba951fa | ||
|
|
20caf8aa5c | ||
|
|
ac4df7eb8e | ||
|
|
af196287e4 | ||
|
|
678ca0e908 | ||
|
|
6fe0dd690f | ||
|
|
2620a2a3fa | ||
|
|
ca755f8ca2 | ||
|
|
8f19087c2b | ||
|
|
4708a3661b | ||
|
|
e57765e012 | ||
|
|
2d0cd46084 | ||
|
|
3feb111c86 | ||
|
|
581b7c48cc | ||
|
|
ed248447e2 | ||
|
|
5792b2b864 | ||
|
|
e5b4c52d7e | ||
|
|
b475f5cca0 | ||
|
|
4724f26d33 | ||
|
|
c908fade19 | ||
|
|
c65bf5c9eb | ||
|
|
2acd66df44 | ||
|
|
f7ff5341f7 | ||
|
|
dcd3ba57e4 | ||
|
|
e5cd239401 | ||
|
|
1fa2de3327 | ||
|
|
ed120a256d | ||
|
|
73b575c7a3 | ||
|
|
6d5157575c | ||
|
|
1a4388ae49 | ||
|
|
0c6d54dc8a | ||
|
|
370d930341 | ||
|
|
0e5aa8896e | ||
|
|
0f47c7a8dc | ||
|
|
7f05187325 | ||
|
|
ada3329bb9 | ||
|
|
a514d39624 | ||
|
|
d06fdd78cc | ||
|
|
7a9ea70580 | ||
|
|
64fe5e81f2 | ||
|
|
d29f4703e9 | ||
|
|
c41ad9bd75 | ||
|
|
af8cf142d7 | ||
|
|
dfac776dd5 | ||
|
|
31218127e0 | ||
|
|
57c18bc298 | ||
|
|
6a0b85f108 | ||
|
|
1ff09ccf9c | ||
|
|
ff4dc116c3 | ||
|
|
a6c875156b | ||
|
|
029e5d47d6 | ||
|
|
489d55fa62 | ||
|
|
47a7d43809 | ||
|
|
21928d7c95 | ||
|
|
552feb24b2 | ||
|
|
3cc3713829 | ||
|
|
03a1fd08c0 | ||
|
|
0db0d3ec93 | ||
|
|
fa69b5b73b | ||
|
|
dbf3b9d9bc | ||
|
|
1942cf2cd8 | ||
|
|
441cf2272e | ||
|
|
4dc8348348 | ||
|
|
1bb7e03429 | ||
|
|
7b161740d0 | ||
|
|
6428a83713 | ||
|
|
705757f362 | ||
|
|
dd6766c37d | ||
|
|
8ca1f20f93 | ||
|
|
a23187a456 | ||
|
|
18650712fa | ||
|
|
bd853be356 | ||
|
|
17d0a0c068 | ||
|
|
be6caf9dcf | ||
|
|
91ee2537e2 | ||
|
|
c1fae2ae1b | ||
|
|
434dbee64c | ||
|
|
c3dfaf3dd9 | ||
|
|
234f5a0621 | ||
|
|
d7c32f9633 | ||
|
|
dcef72480b | ||
|
|
3a65d5ed29 | ||
|
|
1265563fc8 | ||
|
|
036ba19b65 | ||
|
|
6c8a92c318 | ||
|
|
912f8e22d1 | ||
|
|
5126e0bb08 | ||
|
|
41d71a9818 | ||
|
|
9fe339dfb1 | ||
|
|
7d1aab8b8d | ||
|
|
2ea5df3db0 | ||
|
|
820019f3c6 | ||
|
|
0679d38e8e | ||
|
|
0d2dbe756d | ||
|
|
ea460c062d | ||
|
|
6ba98c040e | ||
|
|
b82fb603c2 | ||
|
|
c1fadc008a | ||
|
|
4754efc419 | ||
|
|
a404071d3b | ||
|
|
e07ca1bd0a | ||
|
|
d61e8ffff6 | ||
|
|
c20e08fecc | ||
|
|
2142d43f6f | ||
|
|
0c38d46ee9 | ||
|
|
375ec091f3 | ||
|
|
eb5b29f067 | ||
|
|
000fb9ddbe | ||
|
|
dfa02c8b61 | ||
|
|
d28b852732 | ||
|
|
6e03dd2d4c | ||
|
|
b11e053323 | ||
|
|
11154d48b7 | ||
|
|
a483682da6 | ||
|
|
8a7a70b828 | ||
|
|
dddd571a99 | ||
|
|
8bd9ce194b | ||
|
|
0e0fd14ed4 | ||
|
|
946a80bfb8 | ||
|
|
07af5b4a37 | ||
|
|
59b008a461 | ||
|
|
188f8fcb32 | ||
|
|
d7059d2840 | ||
|
|
5adaf60a9f | ||
|
|
3fb33b9873 | ||
|
|
ca97f947d9 | ||
|
|
347b9c23c4 | ||
|
|
ebe438b25b | ||
|
|
c05ca6529f | ||
|
|
6a6d4f2261 | ||
|
|
cf7aece729 | ||
|
|
36e94454e6 | ||
|
|
5bfef3caa4 | ||
|
|
c1908a0205 | ||
|
|
54354e4735 | ||
|
|
4df4b398f5 | ||
|
|
81aa4979b3 | ||
|
|
0108cd7019 | ||
|
|
aad0af9756 | ||
|
|
bd49f6a7c7 | ||
|
|
263eba1342 | ||
|
|
418524051d | ||
|
|
327aa05e80 | ||
|
|
9bf0023163 | ||
|
|
d64921221e | ||
|
|
7bb3d1925b | ||
|
|
eca0fd5a29 | ||
|
|
44398633bb | ||
|
|
52b893b93b | ||
|
|
52c216c1fb | ||
|
|
d9bc612a3c | ||
|
|
1fe43ec3fb | ||
|
|
ca87e29891 | ||
|
|
d2b14e51ee | ||
|
|
9676fc1699 | ||
|
|
73a5b576c9 | ||
|
|
685e50b9af | ||
|
|
c27fafb298 | ||
|
|
37fa0adf93 | ||
|
|
b95055ba28 | ||
|
|
4b78761d20 | ||
|
|
1549987015 | ||
|
|
b0e5d103d8 | ||
|
|
5ab6bdc9b6 | ||
|
|
ed1083803b | ||
|
|
ec3b415c42 | ||
|
|
3a188971b9 | ||
|
|
b3e32ce5b6 | ||
|
|
d0bbfa1015 | ||
|
|
043a68fa44 | ||
|
|
7542035e44 | ||
|
|
bee6f5d114 | ||
|
|
e20b1815fc | ||
|
|
395823f7b6 | ||
|
|
a910fd8a0b | ||
|
|
64ee029a18 | ||
|
|
25df91ca53 | ||
|
|
8987fb8267 | ||
|
|
86ff344c93 | ||
|
|
b6c7b2a2e9 | ||
|
|
9efce6a470 | ||
|
|
6cc7484b78 | ||
|
|
920afcce20 | ||
|
|
6d82c8a554 | ||
|
|
4b52b10e2e | ||
|
|
ac06df1ade | ||
|
|
5f1a867650 | ||
|
|
d284119246 | ||
|
|
84b9aa56ce | ||
|
|
f58ecee0b9 | ||
|
|
b9d16fe86f | ||
|
|
2ccde7a5cd | ||
|
|
d2ca26afaf | ||
|
|
b52feb305e | ||
|
|
558b2f8b67 | ||
|
|
383f0caa16 | ||
|
|
16e8c4e216 | ||
|
|
c6a1b507ba | ||
|
|
cb4fe70b63 | ||
|
|
408d7c34c8 | ||
|
|
7f2df54e01 | ||
|
|
1758d8510b | ||
|
|
7c3a5619f9 | ||
|
|
2b5b47f97d | ||
|
|
b2508bf2d6 | ||
|
|
b0fb1111ec | ||
|
|
4bf126bf52 | ||
|
|
e4449326e6 | ||
|
|
7bab4db28d | ||
|
|
e29ee58603 | ||
|
|
82e67315a4 | ||
|
|
2e82edee14 | ||
|
|
0338ab4d73 | ||
|
|
f4106fff55 | ||
|
|
d4f700ed22 | ||
|
|
041f0f13f4 | ||
|
|
9dee1c594b | ||
|
|
a1ebb2d9c6 | ||
|
|
06ddf041e5 | ||
|
|
d251fcb7d2 | ||
|
|
e839c65537 | ||
|
|
0cde79b831 | ||
|
|
8767a518db | ||
|
|
10ed288d80 | ||
|
|
7143684a81 | ||
|
|
11da432db2 | ||
|
|
73d3231bbd | ||
|
|
59d825dfca | ||
|
|
ae7099fd02 | ||
|
|
56db6d7cf6 | ||
|
|
0eb52762ae | ||
|
|
e93c560288 | ||
|
|
1f213d00f7 | ||
|
|
8f51f99dde | ||
|
|
5cc1482a72 | ||
|
|
639d84ad32 | ||
|
|
1c74f795a6 | ||
|
|
b4f9fbe1b5 | ||
|
|
b992ba996d | ||
|
|
f00d3c7f66 | ||
|
|
e97b615547 | ||
|
|
29cc1d7fdc | ||
|
|
016b5f0363 | ||
|
|
dd7e252452 | ||
|
|
fda1a77f70 | ||
|
|
0accef8e98 | ||
|
|
c77e2368de | ||
|
|
38ef0247c2 | ||
|
|
a158b819f3 | ||
|
|
381963c28a | ||
|
|
ba158c9cdb | ||
|
|
767a2c0061 | ||
|
|
36334f27a1 | ||
|
|
6a479dddb9 | ||
|
|
7058d1aad7 |
@@ -1,7 +1,7 @@
|
||||
{
|
||||
"$schema": "https://anthropic.com/claude-code/marketplace.schema.json",
|
||||
"name": "hindsight",
|
||||
"version": "0.7.2",
|
||||
"version": "0.7.5",
|
||||
"description": "Official Hindsight integrations for Claude Code",
|
||||
"owner": {
|
||||
"name": "vectorize-io"
|
||||
@@ -11,6 +11,11 @@
|
||||
"name": "hindsight-memory",
|
||||
"description": "Automatic long-term memory for Claude Code via Hindsight",
|
||||
"source": "./hindsight-integrations/claude-code"
|
||||
},
|
||||
{
|
||||
"name": "hindsight-zcode",
|
||||
"description": "No-MCP long-term memory for ZCode via Hindsight hooks",
|
||||
"source": "./hindsight-integrations/zcode"
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
@@ -78,6 +78,11 @@ results = await asyncio.gather(*tasks, return_exceptions=True)
|
||||
- **Authentication/tenancy is enforced inside each engine method, not assumed by the handler.** Every engine method that touches bank-scoped data must authenticate via `request_context` — typically `await self._authenticate_tenant(request_context)` (often indirectly through `get_bank_profile(...)`) — so the correct tenant schema is resolved before any query runs. Handlers must thread `request_context` through to the engine method; never query a tenant-scoped table assuming the schema is already set.
|
||||
- Engine methods return typed models (Pydantic/dataclass), not raw dicts (see Type Safety).
|
||||
|
||||
### Database Locking
|
||||
- **Never use PostgreSQL advisory locks** (`pg_advisory_lock`, `pg_try_advisory_lock`, `pg_advisory_xact_lock`, `pg_advisory_unlock`, …) in migrations, engine code, or anything else. Hindsight runs against connection poolers and managed/PG-compatible services where advisory locks are unreliable or unsupported: session-level locks silently leak or vanish when a pooler hands the session to another client, and callers can block forever on a lock the server never grants. Reject any new occurrence, including ones that look "safe" because they are transaction-scoped.
|
||||
- The pre-existing usage in `hindsight_api/migrations.py` is grandfathered, not a precedent — it is tracked for removal. Don't copy it.
|
||||
- Design the concurrency out instead of locking around it: give each process its own object to write (e.g. per-schema DDL rather than a shared `public.` object), make the operation idempotent, or use a real row/table constraint (`INSERT ... ON CONFLICT`, `SELECT ... FOR UPDATE` in a fixed order). See #2690 for a migration that reached for `pg_advisory_xact_lock` and had to be reverted.
|
||||
|
||||
### Branch Hygiene
|
||||
- **Always start new feature branches from `origin/main`** — rebase to ensure a clean base.
|
||||
- **Only include commits relevant to the PR/branch/feature** — no unrelated changes. If the branch contains commits that don't belong, they must be removed before merging.
|
||||
@@ -154,6 +159,18 @@ If any files in `hindsight-api-slim/hindsight_api/api/` were changed:
|
||||
- Were the client SDKs regenerated? (`./scripts/generate-clients.sh`)
|
||||
- Were the control plane proxy routes updated? (`hindsight-control-plane/src/app/api/`)
|
||||
|
||||
### 7a. Check TS/Python wrapper-client parity
|
||||
|
||||
Two of the generated SDKs ship a **hand-written, maintained convenience wrapper** on top of the auto-generated low-level client — and *only* these two:
|
||||
- **TypeScript**: `hindsight-clients/typescript/src/index.ts` (`HindsightClient`)
|
||||
- **Python**: `hindsight-clients/python/hindsight_client/hindsight_client.py` (`Hindsight`)
|
||||
|
||||
(The Rust/Go/etc. clients are generated-only — no wrapper to keep in sync.)
|
||||
|
||||
These wrappers are what most third-party consumers actually call, and they must expose the same surface. **If a change touches one wrapper's method — adds/removes a parameter, changes a default, forwards a new query/body field — the equivalent method in the *other* wrapper must get the same change in the same (or an immediately-following) PR.** A parameter that exists in the generated SDK but is dropped by one wrapper silently strips it for every consumer of that language (this is exactly what #2975 / #3042 fixed for `detail`/`tags_match`/`limit`/`offset` on `listMentalModels`/`getMentalModel`). **Should fix** — flag any wrapper method that gains capabilities in one language but not the other, and add a matching mapping regression test on both sides.
|
||||
|
||||
Note: the `client-coverage-check` CI tool only validates **request-body** fields, not GET **query** parameters — so query-param parity gaps are *not* caught automatically and must be checked by hand here.
|
||||
|
||||
### 7b. Check API-layer data-access boundary
|
||||
|
||||
For each changed handler in `hindsight-api-slim/hindsight_api/api/` (e.g. `http.py`, `mcp.py`):
|
||||
@@ -204,6 +221,14 @@ in `hindsight-api-slim/hindsight_api/config.py`):
|
||||
The `test_bundled_template_matches_repo_root` sync test fails on drift; if the
|
||||
root file changed without re-copying, flag it as a **must fix**.
|
||||
|
||||
### 11c. Check for advisory locks
|
||||
|
||||
Grep the diff for `advisory` (`git diff main...HEAD | grep -in advisory`). Any new
|
||||
`pg_advisory_lock` / `pg_try_advisory_lock` / `pg_advisory_xact_lock` /
|
||||
`pg_advisory_unlock` call is a **must fix** — see Database Locking above. Point the
|
||||
author at the alternatives (per-process objects, idempotent DDL, row-level
|
||||
constraints) rather than just asking them to drop the lock.
|
||||
|
||||
### 12. Review against other coding standards
|
||||
|
||||
Check the diff for violations of the standards listed above:
|
||||
|
||||
+140
-1
@@ -2,7 +2,7 @@
|
||||
# Copy this file to .env and fill in your values
|
||||
|
||||
# LLM Configuration (Required)
|
||||
# Supported providers: openai, groq, ollama, gemini, anthropic, lmstudio, vertexai, minimax, deepseek, zai, atlas, volcano
|
||||
# Supported providers: openai, openai-responses, groq, ollama, gemini, anthropic, lmstudio, vertexai, minimax, deepseek, zai, atlas, volcano
|
||||
HINDSIGHT_API_LLM_PROVIDER=openai
|
||||
HINDSIGHT_API_LLM_API_KEY=your-api-key-here
|
||||
HINDSIGHT_API_LLM_MODEL=gpt-4o-mini
|
||||
@@ -21,6 +21,25 @@ HINDSIGHT_API_LLM_BASE_URL=https://api.openai.com/v1
|
||||
# HINDSIGHT_API_LLM_TEMPERATURE_REFLECT=0.9
|
||||
# HINDSIGHT_API_LLM_TEMPERATURE_CONSOLIDATION=0.0
|
||||
|
||||
# Grammar-enforce structured output (json_schema strict) instead of the soft
|
||||
# schema-in-prompt path. Helps weaker self-hosted models that emit prose preambles
|
||||
# or invalid JSON. The global override below applies to every operation;
|
||||
# per-operation overrides take precedence, in both directions -- set one to false
|
||||
# to opt that operation out while the global flag is on.
|
||||
# HINDSIGHT_API_LLM_STRICT_SCHEMA=false
|
||||
# HINDSIGHT_API_LLM_STRICT_SCHEMA_RETAIN=true
|
||||
# HINDSIGHT_API_LLM_STRICT_SCHEMA_REFLECT=true
|
||||
# HINDSIGHT_API_LLM_STRICT_SCHEMA_CONSOLIDATION=true
|
||||
|
||||
# Some backends, including Bedrock Converse, reject JSON Schema maxItems.
|
||||
# Disable it only for those backends; consolidation still enforces the cap.
|
||||
# HINDSIGHT_API_LLM_SUPPORTS_MAX_ITEMS=true
|
||||
|
||||
# Diagnostic: on any LLM 4xx, log the exact assembled request ([LLM_4XX_DUMP]) --
|
||||
# serialized request config (message bodies stripped) + capped per-message previews.
|
||||
# For debugging otherwise-unreproducible rejected calls. Off by default.
|
||||
# HINDSIGHT_API_LLM_DEBUG_DUMP_4XX=false
|
||||
|
||||
# Example: Anthropic Claude configuration
|
||||
# HINDSIGHT_API_LLM_PROVIDER=anthropic
|
||||
# HINDSIGHT_API_LLM_API_KEY=your-anthropic-api-key
|
||||
@@ -38,6 +57,12 @@ HINDSIGHT_API_LLM_BASE_URL=https://api.openai.com/v1
|
||||
# HINDSIGHT_API_LLM_API_KEY=your-minimax-api-key
|
||||
# HINDSIGHT_API_LLM_MODEL=MiniMax-M3 # or MiniMax-M2.7 for the previous generation
|
||||
|
||||
# Example: OpenAI Responses API (/v1/responses) — reasoning + function tools together
|
||||
# HINDSIGHT_API_LLM_PROVIDER=openai-responses
|
||||
# HINDSIGHT_API_LLM_API_KEY=your-openai-api-key
|
||||
# HINDSIGHT_API_LLM_MODEL=gpt-5.6 # reasoning model (gpt-5.x / o-series); e.g. gpt-5.6-terra
|
||||
# HINDSIGHT_API_LLM_REASONING_EFFORT=high # sent alongside tools, unlike chat/completions
|
||||
|
||||
# Example: DeepSeek configuration (https://api.deepseek.com)
|
||||
# HINDSIGHT_API_LLM_PROVIDER=deepseek
|
||||
# HINDSIGHT_API_LLM_API_KEY=your-deepseek-api-key
|
||||
@@ -59,6 +84,15 @@ HINDSIGHT_API_LLM_BASE_URL=https://api.openai.com/v1
|
||||
# HINDSIGHT_API_LLM_BASE_URL=http://localhost:1234/v1
|
||||
# HINDSIGHT_API_LLM_MODEL=qwen2.5-32b-instruct
|
||||
|
||||
# Example: Ollama local configuration (native provider)
|
||||
# HINDSIGHT_API_LLM_PROVIDER=ollama
|
||||
# HINDSIGHT_API_LLM_BASE_URL=http://localhost:11434/v1
|
||||
# HINDSIGHT_API_LLM_MODEL=gemma3:12b
|
||||
# Native Ollama context-window override (num_ctx). Leave unset to let Ollama use
|
||||
# the model Modelfile / server default; set a positive integer only to force a
|
||||
# specific context size (e.g. 16384 to keep the previous request behavior).
|
||||
# HINDSIGHT_API_LLM_OLLAMA_NUM_CTX=16384
|
||||
|
||||
# Multi-LLM strategies: configure extra LLMs by index alongside the primary above,
|
||||
# then pick a routing strategy. Unset = single primary LLM (default). Members are
|
||||
# numbered from 1; indices must be contiguous. Each operation can override with a
|
||||
@@ -80,6 +114,16 @@ HINDSIGHT_API_LOG_LEVEL=info
|
||||
# Unset uses HINDSIGHT_API_RETAIN_CHUNK_SIZE as the structured-chunk limit.
|
||||
# HINDSIGHT_API_RETAIN_STRUCTURED_CHUNK_SIZE=
|
||||
|
||||
# When true, a retain operation that hit any fact-extraction errors is marked
|
||||
# 'failed' (not 'completed'), surfacing silently-dropped facts. Default false.
|
||||
# HINDSIGHT_API_FAIL_ON_EXTRACTION_ERRORS=false
|
||||
|
||||
# Wall-clock ceiling (seconds) for one retain task in the worker. A retain that
|
||||
# blocks indefinitely is cancelled and marked 'failed' — and so becomes
|
||||
# retryable — instead of holding its worker slot until the process restarts.
|
||||
# Set well above your slowest healthy retain; 0 disables. Default 3600.
|
||||
# HINDSIGHT_API_RETAIN_WALL_TIMEOUT=3600
|
||||
|
||||
# Dry-run extraction preview endpoint (POST /memories/dry-run-extract). Enabled by default; it makes
|
||||
# a real LLM call but stores nothing. Set to false to remove the endpoint (returns 404).
|
||||
# HINDSIGHT_API_ENABLE_DRY_RUN_EXTRACT=true
|
||||
@@ -95,7 +139,13 @@ HINDSIGHT_API_LOG_LEVEL=info
|
||||
# HINDSIGHT_API_READ_DATABASE_URL= # Optional read-replica URL. When set, recall queries (semantic, BM25, graph, temporal) flow through a separate pool against this URL, offloading the primary. Typically points to a read-only endpoint (CNPG's <cluster>-ro service or Aurora reader endpoint).
|
||||
# HINDSIGHT_API_MIGRATION_DATABASE_URL= # Direct PostgreSQL URL for migrations (bypasses PgBouncer). Falls back to DATABASE_URL.
|
||||
# HINDSIGHT_API_DATABASE_SCHEMA=public # PostgreSQL schema name (default: public)
|
||||
# HINDSIGHT_API_DB_MAX_PARALLEL_WORKERS_PER_GATHER= # Optional cap on Postgres planner parallelism for this process's pool connections. Unset leaves the server default; 0 makes background/bulk queries run serially (useful on worker processes sharing a primary with latency-sensitive traffic).
|
||||
# HINDSIGHT_API_ENTITY_TRGM_SIMILARITY_THRESHOLD=0.15 # Postgres pg_trgm.similarity_threshold applied on every pool connection, used by entity resolution's % trigram match. Must be in (0, 1]. Lower catches more substring-ish matches at higher CPU cost on large entity sets; higher is stricter and cheaper.
|
||||
# HINDSIGHT_API_ENTITY_INTRABATCH_MERGE_SIMILARITY=0.5 # Trigram similarity (pg_trgm-equivalent, computed in-memory) at/above which two new names created by the SAME retain are merged into one entity (in-batch dedup of surface-form variants). Must be in (0, 1]. A merge cutoff, stricter than the recall threshold above; raise toward 1.0 to merge only near-identical forms.
|
||||
# HINDSIGHT_API_RETAIN_ENTITY_RESOLUTION_MAX_CANDIDATES=200 # Max candidates scored per entity mention during retain. The fuzzy lookup keeps only this many best matches per name (ranked by trigram/Jaro-Winkler similarity) before scoring them one by one. On banks holding thousands of near-identical names an uncapped set turns one retain into minutes of CPU that stall the worker's health checks. Raise only if entities that should merge are being duplicated.
|
||||
# HINDSIGHT_API_MIGRATION_CONCURRENCY=1 # Tenant schemas to migrate concurrently (PG only, each in its own process; per-schema work stays sequential). Each worker has ~1-2s startup cost + uses ~3 DB connections, so it only pays off with many schemas (tens+) or slow migrations; keep concurrency*3 <= spare max_connections. Default: 1 (sequential).
|
||||
# HINDSIGHT_API_OPERATION_RETENTION_DAYS=30 # Prune terminal operation rows, payloads, and metadata after this many days; 0 (the default) keeps them forever.
|
||||
# HINDSIGHT_API_OPERATION_CLEANUP_BATCH_SIZE=1000 # Maximum expired terminal rows deleted per tenant schema in each cleanup cycle; must be positive.
|
||||
|
||||
# Vector Extension (Optional - uses pgvector by default)
|
||||
# Options: "pgvector" (default), "vchord", "pgvectorscale" (DiskANN)
|
||||
@@ -115,6 +165,11 @@ HINDSIGHT_API_LOG_LEVEL=info
|
||||
# chinese_lindera/lindera(chinese), japanese_lindera/lindera(japanese),
|
||||
# korean_lindera/lindera(korean), ngram(min,max), edge_ngram(min,max)
|
||||
# HINDSIGHT_API_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER=
|
||||
# Optional cap on the number of terms in the native PostgreSQL BM25 tsquery.
|
||||
# Long queries OR-join every normalized token, which can match too much of a
|
||||
# large bank. 0 (default) keeps the historical uncapped behavior; a positive
|
||||
# value bounds only the native backend (other BM25 backends get the raw query).
|
||||
# HINDSIGHT_API_BM25_MAX_QUERY_TERMS=0
|
||||
|
||||
# File Parser (Optional - uses markitdown by default)
|
||||
# HINDSIGHT_API_FILE_PARSER=markitdown
|
||||
@@ -127,12 +182,19 @@ HINDSIGHT_API_LOG_LEVEL=info
|
||||
# HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_BASE_URL=
|
||||
# HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_MODEL=
|
||||
# HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_PROMPT=
|
||||
# Optional JSON dict of custom headers for the OCR OpenAI client (e.g. proxies / request tracing).
|
||||
# HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_DEFAULT_HEADERS=
|
||||
|
||||
# Embeddings Configuration (Optional - uses local by default)
|
||||
# Provider: "local" (default), "onnx", "tei", "openai", "cohere", "google", "openrouter", "zeroentropy", "litellm", or "litellm-sdk"
|
||||
# HINDSIGHT_API_EMBEDDINGS_PROVIDER=local
|
||||
# For local provider:
|
||||
# HINDSIGHT_API_EMBEDDINGS_LOCAL_MODEL=BAAI/bge-small-en-v1.5
|
||||
# Force CPU if local embeddings hit MPS/XPC instability on macOS:
|
||||
# HINDSIGHT_API_EMBEDDINGS_LOCAL_FORCE_CPU=false
|
||||
# Opt in to the Apple Silicon MPS GPU (off by default: MPS leaks memory under
|
||||
# variable-length workloads). CUDA/XPU still auto-select regardless:
|
||||
# HINDSIGHT_API_EMBEDDINGS_LOCAL_ALLOW_MPS=false
|
||||
# For ONNX provider (local CPU embeddings without an Ollama/TEI sidecar):
|
||||
# HINDSIGHT_API_EMBEDDINGS_PROVIDER=onnx
|
||||
# HINDSIGHT_API_EMBEDDINGS_ONNX_MODEL_ID=intfloat/multilingual-e5-small
|
||||
@@ -148,6 +210,11 @@ HINDSIGHT_API_LOG_LEVEL=info
|
||||
# HINDSIGHT_API_EMBEDDINGS_ONNX_TOKENIZER_NAME_OR_PATH=/models/multilingual-e5-small
|
||||
# Optional for China network / restricted HF access:
|
||||
# HF_ENDPOINT=https://hf-mirror.com
|
||||
# Applies to any provider: cap each input at this many tiktoken tokens before
|
||||
# embedding, so oversized content is truncated instead of failing the embed call
|
||||
# permanently (e.g. Bedrock Titan V2's 8192, or a llama.cpp server's context). Off
|
||||
# by default. (Deprecated alias: HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_MAX_INPUT_TOKENS)
|
||||
# HINDSIGHT_API_EMBEDDINGS_MAX_INPUT_TOKENS=8192
|
||||
# For TEI provider:
|
||||
# HINDSIGHT_API_EMBEDDINGS_TEI_URL=http://localhost:8080
|
||||
# For OpenAI-compatible embeddings:
|
||||
@@ -169,13 +236,58 @@ HINDSIGHT_API_LOG_LEVEL=info
|
||||
# DeepSeek note: DeepSeek is supported for LLM calls, but not for embeddings.
|
||||
# If using DeepSeek as LLM provider, keep embeddings on local/openai/cohere/google/etc.
|
||||
|
||||
# Embedding similarity thresholds. These defaults preserve the behavior calibrated
|
||||
# for BAAI/bge-small-en-v1.5. Recalibrate each threshold independently when changing
|
||||
# embedding models because cosine-similarity distributions are model-dependent.
|
||||
# HINDSIGHT_API_SEMANTIC_MIN_SIMILARITY=0.3
|
||||
# HINDSIGHT_API_GRAPH_SEED_MIN_SIMILARITY=0.3
|
||||
# HINDSIGHT_API_TEMPORAL_SEMANTIC_MIN_SIMILARITY=0.1
|
||||
# HINDSIGHT_API_SEMANTIC_LINK_MIN_SIMILARITY=0.7
|
||||
# HINDSIGHT_API_CONSOLIDATION_DEDUP_THRESHOLD=0.97
|
||||
|
||||
# Recall pipeline stages (all on by default). Each is hierarchical, so a single
|
||||
# bank can switch a stage off via the config API without changing the server
|
||||
# default. Turning all three off leaves semantic + BM25 fused by RRF, the
|
||||
# lowest-latency recall path.
|
||||
# Temporal retrieval arm, plus the date-aware query analysis that feeds it:
|
||||
# HINDSIGHT_API_ENABLE_TEMPORAL_RETRIEVAL=true
|
||||
# Entity/link graph traversal arm:
|
||||
# HINDSIGHT_API_ENABLE_GRAPH_RETRIEVAL=true
|
||||
# Cross-encoder rerank of the fused candidates (false = use the RRF order):
|
||||
# HINDSIGHT_API_ENABLE_RERANKING=true
|
||||
|
||||
# Reranker Configuration (Optional - uses local by default)
|
||||
# Provider: "local" (default) or "tei" (HuggingFace Text Embeddings Inference)
|
||||
# HINDSIGHT_API_RERANKER_PROVIDER=local
|
||||
# Trusted gateway attribution (disabled by default). When enabled, remote
|
||||
# reranker requests include X-Hindsight-Bank-Id with the current bank ID.
|
||||
# HINDSIGHT_API_RERANKER_SEND_BANK_AS_HEADER=false
|
||||
# For local provider:
|
||||
# HINDSIGHT_API_RERANKER_LOCAL_MODEL=cross-encoder/ms-marco-MiniLM-L-6-v2
|
||||
# Force CPU if the local reranker hits MPS/XPC instability on macOS:
|
||||
# HINDSIGHT_API_RERANKER_LOCAL_FORCE_CPU=false
|
||||
# Opt in to the Apple Silicon MPS GPU (off by default: MPS leaks memory under
|
||||
# variable-length workloads). CUDA/XPU still auto-select regardless:
|
||||
# HINDSIGHT_API_RERANKER_LOCAL_ALLOW_MPS=false
|
||||
# For TEI provider:
|
||||
# HINDSIGHT_API_RERANKER_TEI_URL=http://localhost:8081
|
||||
# Max candidates the cross-encoder reranks per recall (RRF pre-filters the rest):
|
||||
# HINDSIGHT_API_RERANKER_MAX_CANDIDATES=300
|
||||
# Optionally scale that cap by the recall budget level (the cross-encoder dominates
|
||||
# a large recall's latency). 0 = fall back to the flat cap above; fully backwards-compatible.
|
||||
# HINDSIGHT_API_RERANKER_MAX_CANDIDATES_LOW=0
|
||||
# HINDSIGHT_API_RERANKER_MAX_CANDIDATES_MID=0
|
||||
# HINDSIGHT_API_RERANKER_MAX_CANDIDATES_HIGH=0
|
||||
|
||||
# Reranker failover chain: extra rerankers tried, in order, when the one above
|
||||
# fails. Members are numbered from 1 (indices must be contiguous) and every
|
||||
# setting of member n carries the same index. A member inherits nothing from the
|
||||
# primary, so spell out everything it needs. Unset = no fallback (default): a
|
||||
# failing reranker fails the recall. End the chain with "rrf" to fail open and
|
||||
# keep the retrieval order instead.
|
||||
# HINDSIGHT_API_RERANKER_1_PROVIDER=cohere
|
||||
# HINDSIGHT_API_RERANKER_1_COHERE_API_KEY=your-cohere-api-key
|
||||
# HINDSIGHT_API_RERANKER_2_PROVIDER=rrf
|
||||
|
||||
# Observability & Tracing (Optional - disabled by default)
|
||||
# Enable OpenTelemetry tracing for LLM calls (GenAI semantic conventions)
|
||||
@@ -195,6 +307,33 @@ HINDSIGHT_API_LOG_LEVEL=info
|
||||
# Expose async-operation queue + consolidation-backlog gauges on /metrics.
|
||||
# Runs periodic per-schema COUNT queries on a background task (disabled by default).
|
||||
# HINDSIGHT_API_METRICS_BACKLOG_ENABLED=true
|
||||
#
|
||||
# Runtime-stall observability (enabled by default). When a liveness probe fails,
|
||||
# these tell you WHY: a blocked event loop vs DB connection-pool exhaustion.
|
||||
# The loop watchdog logs the offending stack when the loop is unresponsive; the
|
||||
# DB-pool acquire timing logs (and exposes hindsight.db.pool.waiting) when
|
||||
# callers queue for a connection. Both are cheap; tune or disable if needed.
|
||||
# HINDSIGHT_API_LOOP_WATCHDOG_ENABLED=false
|
||||
# HINDSIGHT_API_LOOP_WATCHDOG_STALL_THRESHOLD_MS=1000
|
||||
# HINDSIGHT_API_LOOP_WATCHDOG_POLL_INTERVAL_MS=250
|
||||
# HINDSIGHT_API_DB_ACQUIRE_WARN_THRESHOLD_MS=1000
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Webhooks (Optional)
|
||||
# -----------------------------------------------------------------------------
|
||||
|
||||
# Outbound webhook delivery targets caller-supplied URLs. To prevent SSRF, the
|
||||
# delivery worker blocks private, loopback, and link-local destinations
|
||||
# (including the cloud metadata address 169.254.169.254) by default. List hosts
|
||||
# or IP/CIDR ranges here (comma-separated) to re-permit specific internal
|
||||
# destinations — e.g. 127.0.0.1 for local testing, or an internal receiver.
|
||||
# HINDSIGHT_API_WEBHOOK_ALLOWED_HOSTS=127.0.0.1,internal-receiver.svc,10.0.0.0/8
|
||||
|
||||
# Whether the webhook delivery-history API returns the raw upstream response
|
||||
# body. Off by default: returning arbitrary response bodies to callers is an
|
||||
# information-exfiltration primitive. The delivery status code is always
|
||||
# returned regardless. Enable only if you trust your webhook destinations.
|
||||
# HINDSIGHT_API_WEBHOOK_EXPOSE_RESPONSE_BODY=false
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Control Plane (Optional)
|
||||
|
||||
File diff suppressed because one or more lines are too long
|
After Width: | Height: | Size: 40 KiB |
File diff suppressed because it is too large
Load Diff
@@ -25,6 +25,20 @@ jobs:
|
||||
with:
|
||||
python-version-file: ".python-version"
|
||||
|
||||
# Each package is built from its own directory, so stage the repository's
|
||||
# canonical license inside each isolated build context.
|
||||
- name: Stage Python package licenses
|
||||
run: |
|
||||
for package in \
|
||||
hindsight-clients/python \
|
||||
hindsight-api-slim \
|
||||
hindsight-api \
|
||||
hindsight-all \
|
||||
hindsight-all-slim \
|
||||
hindsight-embed; do
|
||||
cp LICENSE "$package/LICENSE"
|
||||
done
|
||||
|
||||
# Build all packages
|
||||
- name: Build hindsight-client
|
||||
working-directory: ./hindsight-clients/python
|
||||
@@ -50,6 +64,24 @@ jobs:
|
||||
working-directory: ./hindsight-embed
|
||||
run: uv build --out-dir dist
|
||||
|
||||
- name: Verify Python package licenses
|
||||
run: |
|
||||
for package in \
|
||||
hindsight-clients/python \
|
||||
hindsight-api-slim \
|
||||
hindsight-api \
|
||||
hindsight-all \
|
||||
hindsight-all-slim \
|
||||
hindsight-embed; do
|
||||
wheel=$(find "$package/dist" -maxdepth 1 -name '*.whl' -print -quit)
|
||||
sdist=$(find "$package/dist" -maxdepth 1 -name '*.tar.gz' -print -quit)
|
||||
|
||||
unzip -Z1 "$wheel" | grep -Eq '\.dist-info/licenses/LICENSE$'
|
||||
unzip -p "$wheel" '*/METADATA' | grep -Fxq 'License-Expression: MIT'
|
||||
unzip -p "$wheel" '*/METADATA' | grep -Fxq 'License-File: LICENSE'
|
||||
tar -tzf "$sdist" | grep -Eq '/LICENSE$'
|
||||
done
|
||||
|
||||
# Publish in order (client and api-slim first, then api/all wrappers which depend on them)
|
||||
- name: Publish hindsight-client to PyPI
|
||||
uses: pypa/gh-action-pypi-publish@release/v1
|
||||
|
||||
@@ -0,0 +1,29 @@
|
||||
name: Update star history
|
||||
|
||||
on:
|
||||
schedule:
|
||||
- cron: '17 3 * * *'
|
||||
workflow_dispatch:
|
||||
|
||||
permissions:
|
||||
contents: write
|
||||
|
||||
jobs:
|
||||
update:
|
||||
concurrency:
|
||||
group: star-history
|
||||
cancel-in-progress: false
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: nicoloboschi/gh-stars@v1
|
||||
with:
|
||||
github-token: ${{ secrets.GITHUB_TOKEN }}
|
||||
line-color: '#14b8a6'
|
||||
- name: Commit chart
|
||||
run: |
|
||||
git config user.name 'github-actions[bot]'
|
||||
git config user.email '41898282+github-actions[bot]@users.noreply.github.com'
|
||||
git add .github/star-history/data.json .github/star-history/chart.svg
|
||||
git diff --cached --quiet || git commit -m 'chore: update star history'
|
||||
git push
|
||||
+272
-17
@@ -36,11 +36,14 @@ jobs:
|
||||
integrations-composio: ${{ steps.filter.outputs.integrations-composio }}
|
||||
integrations-chat: ${{ steps.filter.outputs.integrations-chat }}
|
||||
integrations-claude-code: ${{ steps.filter.outputs.integrations-claude-code }}
|
||||
integrations-coding-agents: ${{ steps.filter.outputs.integrations-coding-agents }}
|
||||
integrations-cline: ${{ steps.filter.outputs.integrations-cline }}
|
||||
integrations-codex: ${{ steps.filter.outputs.integrations-codex }}
|
||||
integrations-github-copilot: ${{ steps.filter.outputs.integrations-github-copilot }}
|
||||
integrations-continue: ${{ steps.filter.outputs.integrations-continue }}
|
||||
integrations-cursor-cli: ${{ steps.filter.outputs.integrations-cursor-cli }}
|
||||
integrations-zcode: ${{ steps.filter.outputs.integrations-zcode }}
|
||||
integrations-copilot-cli: ${{ steps.filter.outputs.integrations-copilot-cli }}
|
||||
integrations-crewai: ${{ steps.filter.outputs.integrations-crewai }}
|
||||
integrations-litellm: ${{ steps.filter.outputs.integrations-litellm }}
|
||||
integrations-pydantic-ai: ${{ steps.filter.outputs.integrations-pydantic-ai }}
|
||||
@@ -142,6 +145,8 @@ jobs:
|
||||
- 'hindsight-integrations/chat/**'
|
||||
integrations-claude-code:
|
||||
- 'hindsight-integrations/claude-code/**'
|
||||
integrations-coding-agents:
|
||||
- 'hindsight-integrations/coding-agents/**'
|
||||
integrations-cline:
|
||||
- 'hindsight-integrations/cline/**'
|
||||
integrations-codex:
|
||||
@@ -152,6 +157,8 @@ jobs:
|
||||
- 'hindsight-integrations/continue/**'
|
||||
integrations-cursor-cli:
|
||||
- 'hindsight-integrations/cursor-cli/**'
|
||||
integrations-copilot-cli:
|
||||
- 'hindsight-integrations/copilot-cli/**'
|
||||
integrations-crewai:
|
||||
- 'hindsight-integrations/crewai/**'
|
||||
integrations-litellm:
|
||||
@@ -180,6 +187,8 @@ jobs:
|
||||
- 'hindsight-integrations/cursor/**'
|
||||
integrations-zed:
|
||||
- 'hindsight-integrations/zed/**'
|
||||
integrations-zcode:
|
||||
- 'hindsight-integrations/zcode/**'
|
||||
integrations-n8n:
|
||||
- 'hindsight-integrations/n8n/**'
|
||||
integrations-zapier:
|
||||
@@ -283,6 +292,18 @@ jobs:
|
||||
working-directory: ./hindsight-api-slim
|
||||
run: uv build
|
||||
|
||||
# `uv build` only packages the source; it does not prove the dependency set
|
||||
# resolves or that the code imports on this interpreter. Install into a fresh
|
||||
# env and run a byte-compile + import smoke test so the matrix actually
|
||||
# exercises each Python version (notably 3.14).
|
||||
- name: Install and smoke-test on Python ${{ matrix.python-version }}
|
||||
working-directory: ./hindsight-api-slim
|
||||
run: |
|
||||
uv venv --python ${{ matrix.python-version }} .venv-smoke
|
||||
VIRTUAL_ENV=.venv-smoke uv pip install .
|
||||
.venv-smoke/bin/python -m compileall -q hindsight_api
|
||||
.venv-smoke/bin/python -c "import hindsight_api, hindsight_api.main, hindsight_api.config; from hindsight_api.engine import memory_engine, llm_wrapper; print('import OK')"
|
||||
|
||||
build-typescript-client:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
@@ -454,6 +475,42 @@ jobs:
|
||||
working-directory: ./hindsight-integrations/openclaw
|
||||
run: ./scripts/smoke-test.sh
|
||||
|
||||
test-coding-agents:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
(github.event_name == 'workflow_dispatch' ||
|
||||
needs.detect-changes.outputs.integrations-coding-agents == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 15
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.head.sha || '' }}
|
||||
|
||||
- uses: actions/setup-node@v6
|
||||
with:
|
||||
node-version: '22'
|
||||
cache: npm
|
||||
cache-dependency-path: hindsight-integrations/coding-agents/package-lock.json
|
||||
|
||||
- name: Install dependencies
|
||||
working-directory: ./hindsight-integrations/coding-agents
|
||||
run: npm ci
|
||||
|
||||
- name: Typecheck
|
||||
working-directory: ./hindsight-integrations/coding-agents
|
||||
run: npx tsc --noEmit
|
||||
|
||||
- name: Unit tests
|
||||
working-directory: ./hindsight-integrations/coding-agents
|
||||
run: npm test
|
||||
|
||||
- name: Build
|
||||
working-directory: ./hindsight-integrations/coding-agents
|
||||
run: npm run build
|
||||
|
||||
test-claude-code-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
@@ -520,22 +577,17 @@ jobs:
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.head.sha || '' }}
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v6
|
||||
- name: Set up Node.js
|
||||
uses: actions/setup-node@v6
|
||||
with:
|
||||
python-version: '3.11'
|
||||
|
||||
- name: Install package and pytest
|
||||
working-directory: ./hindsight-integrations/zed
|
||||
# Installs the package (incl. the zstandard runtime dep) so the threads.db
|
||||
# reader tests can decompress Zed's zstd blobs.
|
||||
run: pip install -e . pytest
|
||||
node-version: '22'
|
||||
|
||||
- name: Run tests
|
||||
working-directory: ./hindsight-integrations/zed
|
||||
# PR CI runs only the deterministic bucket; the real-LLM E2E bucket
|
||||
# (requires_real_llm) needs a live Hindsight server and runs separately.
|
||||
run: python -m pytest tests/ -v -m "not requires_real_llm"
|
||||
# Config-only integration with no dependencies — it uses Node's built-in
|
||||
# test runner. The runtime MCP bridge is `npx mcp-remote` (Node), so this
|
||||
# integration requires only Node.js (no Python).
|
||||
run: npm test
|
||||
|
||||
test-omo-integration:
|
||||
needs: [detect-changes]
|
||||
@@ -702,6 +754,43 @@ jobs:
|
||||
working-directory: ./hindsight-integrations/cursor-cli
|
||||
run: uv run pytest tests -v
|
||||
|
||||
test-zcode-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
(github.event_name == 'workflow_dispatch' ||
|
||||
needs.detect-changes.outputs.integrations-zcode == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.head.sha || '' }}
|
||||
|
||||
- name: Install uv
|
||||
uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
enable-cache: true
|
||||
prune-cache: false
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v6
|
||||
with:
|
||||
python-version-file: ".python-version"
|
||||
|
||||
- name: Build zcode integration
|
||||
working-directory: ./hindsight-integrations/zcode
|
||||
run: uv build
|
||||
|
||||
- name: Install dependencies
|
||||
working-directory: ./hindsight-integrations/zcode
|
||||
run: uv sync --frozen
|
||||
|
||||
- name: Run tests
|
||||
working-directory: ./hindsight-integrations/zcode
|
||||
run: uv run pytest tests -v
|
||||
|
||||
build-ai-sdk-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
@@ -1229,10 +1318,10 @@ jobs:
|
||||
|
||||
build-docs:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
(github.event_name == 'workflow_dispatch' ||
|
||||
needs.detect-changes.outputs.docs == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
# Keep the production docs build as an unconditional PR check. OpenAPI
|
||||
# generation used to build the site again inside verify-generated-files;
|
||||
# running the existing job for every PR preserves that coverage without
|
||||
# serializing two full Docusaurus builds in the generated-files check.
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
|
||||
@@ -1836,6 +1925,27 @@ jobs:
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.head.sha || '' }}
|
||||
|
||||
# The Oracle 23ai `free` service image is large; combined with the Python ML
|
||||
# deps (torch) it exhausts the runner's ~14 GB root disk, so uv fails to
|
||||
# extract a wheel with "No space left on device (os error 28)".
|
||||
# Only the cheap, high-yield reclaims are enabled — the Android/.NET/Haskell
|
||||
# dirs plus swap are a few `rm -rf`s worth ~16-21 GB, which is ample headroom:
|
||||
# - large-packages runs apt-get remove and costs minutes for little gain;
|
||||
# - tool-cache would delete the preinstalled Python that actions/setup-python
|
||||
# then has to re-download, making the job slower, not faster;
|
||||
# - docker-images is pointless here (the Oracle service container is already
|
||||
# running, so its image is in use and cannot be pruned anyway).
|
||||
- name: Free Disk Space
|
||||
uses: jlumbroso/free-disk-space@main
|
||||
with:
|
||||
tool-cache: false
|
||||
android: true
|
||||
dotnet: true
|
||||
haskell: true
|
||||
large-packages: false
|
||||
docker-images: false
|
||||
swap-storage: true
|
||||
|
||||
- name: Setup Oracle test user
|
||||
# The SYSTEM tablespace uses manual segment space management which
|
||||
# doesn't support VECTOR types. Create an ASSM tablespace and a
|
||||
@@ -2196,6 +2306,27 @@ jobs:
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.head.sha || '' }}
|
||||
|
||||
# The Oracle 23ai `free` service image is large; combined with the Python ML
|
||||
# deps (torch) it exhausts the runner's ~14 GB root disk, so uv fails to
|
||||
# extract a wheel with "No space left on device (os error 28)".
|
||||
# Only the cheap, high-yield reclaims are enabled — the Android/.NET/Haskell
|
||||
# dirs plus swap are a few `rm -rf`s worth ~16-21 GB, which is ample headroom:
|
||||
# - large-packages runs apt-get remove and costs minutes for little gain;
|
||||
# - tool-cache would delete the preinstalled Python that actions/setup-python
|
||||
# then has to re-download, making the job slower, not faster;
|
||||
# - docker-images is pointless here (the Oracle service container is already
|
||||
# running, so its image is in use and cannot be pruned anyway).
|
||||
- name: Free Disk Space
|
||||
uses: jlumbroso/free-disk-space@main
|
||||
with:
|
||||
tool-cache: false
|
||||
android: true
|
||||
dotnet: true
|
||||
haskell: true
|
||||
large-packages: false
|
||||
docker-images: false
|
||||
swap-storage: true
|
||||
|
||||
- name: Setup Oracle test user
|
||||
run: |
|
||||
pip install oracledb
|
||||
@@ -2356,6 +2487,27 @@ jobs:
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.head.sha || '' }}
|
||||
|
||||
# The Oracle 23ai `free` service image is large; combined with the Python ML
|
||||
# deps (torch) it exhausts the runner's ~14 GB root disk, so uv fails to
|
||||
# extract a wheel with "No space left on device (os error 28)".
|
||||
# Only the cheap, high-yield reclaims are enabled — the Android/.NET/Haskell
|
||||
# dirs plus swap are a few `rm -rf`s worth ~16-21 GB, which is ample headroom:
|
||||
# - large-packages runs apt-get remove and costs minutes for little gain;
|
||||
# - tool-cache would delete the preinstalled Python that actions/setup-python
|
||||
# then has to re-download, making the job slower, not faster;
|
||||
# - docker-images is pointless here (the Oracle service container is already
|
||||
# running, so its image is in use and cannot be pruned anyway).
|
||||
- name: Free Disk Space
|
||||
uses: jlumbroso/free-disk-space@main
|
||||
with:
|
||||
tool-cache: false
|
||||
android: true
|
||||
dotnet: true
|
||||
haskell: true
|
||||
large-packages: false
|
||||
docker-images: false
|
||||
swap-storage: true
|
||||
|
||||
- name: Setup Oracle test user
|
||||
run: |
|
||||
pip install oracledb
|
||||
@@ -3481,6 +3633,43 @@ jobs:
|
||||
# (requires_real_llm) needs a live Hindsight server and runs separately.
|
||||
run: uv run pytest tests -v -m "not requires_real_llm"
|
||||
|
||||
test-copilot-cli-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
(github.event_name == 'workflow_dispatch' ||
|
||||
needs.detect-changes.outputs.integrations-copilot-cli == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.head.sha || '' }}
|
||||
|
||||
- name: Install uv
|
||||
uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
enable-cache: true
|
||||
prune-cache: false
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v6
|
||||
with:
|
||||
python-version-file: ".python-version"
|
||||
|
||||
- name: Build copilot-cli integration
|
||||
working-directory: ./hindsight-integrations/copilot-cli
|
||||
run: uv build
|
||||
|
||||
- name: Install dependencies
|
||||
working-directory: ./hindsight-integrations/copilot-cli
|
||||
run: uv sync --frozen
|
||||
|
||||
- name: Run tests
|
||||
working-directory: ./hindsight-integrations/copilot-cli
|
||||
run: uv run pytest tests -v
|
||||
|
||||
test-crewai-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
@@ -4339,6 +4528,69 @@ jobs:
|
||||
fi
|
||||
done || true
|
||||
|
||||
# Compatibility gate against hermes-agent's *main* branch.
|
||||
#
|
||||
# `hermes memory setup` installs `hindsight-all` into Hermes' own venv to run
|
||||
# memory in local_embedded mode, and Hermes exact-pins every direct dependency
|
||||
# (`==X.Y.Z`) as a deliberate supply-chain policy — they will not loosen a pin
|
||||
# for us. So any version range Hindsight declares that excludes one of their
|
||||
# pins makes the two impossible to co-install for every Hermes user on
|
||||
# embedded memory. That was #3251: our `cryptography>=48.0.1` / `pillow>=12.3.0`
|
||||
# against their `==46.0.7` / `==12.2.0`, which left `pip check` permanently
|
||||
# broken. Both sides bump on their own schedule, so this needs a standing gate
|
||||
# rather than a one-off fix; tracking main surfaces the next collision while
|
||||
# it is still cheap to fix on either side.
|
||||
#
|
||||
# Deliberately not gated on has_secrets — the resolution, wiring, runtime and
|
||||
# daemon-boot checks need no credentials, so this runs on fork PRs too. Only
|
||||
# retain/recall need an LLM and the script skips them when no key is present.
|
||||
test-hermes-compat:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
(github.event_name == 'workflow_dispatch' ||
|
||||
needs.detect-changes.outputs.core == 'true' ||
|
||||
needs.detect-changes.outputs.embed == 'true' ||
|
||||
needs.detect-changes.outputs.hindsight-all == 'true' ||
|
||||
needs.detect-changes.outputs.clients-python == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 45
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.head.sha || '' }}
|
||||
|
||||
- name: Install uv
|
||||
uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
enable-cache: true
|
||||
prune-cache: false
|
||||
|
||||
# hindsight-all pulls the local-ml extra, so the embedded daemon can load
|
||||
# sentence-transformers models. Cache them like the test-embed job does.
|
||||
- name: Cache HuggingFace models
|
||||
uses: actions/cache@v5
|
||||
with:
|
||||
path: ~/.cache/huggingface
|
||||
key: ${{ runner.os }}-huggingface-hermes-compat-${{ hashFiles('hindsight-api-slim/pyproject.toml') }}
|
||||
restore-keys: |
|
||||
${{ runner.os }}-huggingface-hermes-compat-
|
||||
${{ runner.os }}-huggingface-
|
||||
|
||||
- name: Run Hermes compatibility test
|
||||
run: ./scripts/test-hermes-compat.sh
|
||||
|
||||
- name: Collect embedded daemon logs on failure
|
||||
if: failure()
|
||||
run: |
|
||||
for f in ~/.hindsight/profiles/hermes-ci*.log ~/.hindsight/profiles/hermes-ci*.stderr.log; do
|
||||
if [ -f "$f" ]; then
|
||||
echo "=== $f ==="
|
||||
cat "$f"
|
||||
fi
|
||||
done || true
|
||||
|
||||
test-hindsight-all:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
@@ -4771,7 +5023,8 @@ jobs:
|
||||
cd ../hindsight-embed && uv sync --frozen --index-strategy unsafe-best-match
|
||||
|
||||
- name: Run generate-openapi
|
||||
run: ./scripts/generate-openapi.sh
|
||||
working-directory: hindsight-dev
|
||||
run: uv run generate-openapi
|
||||
|
||||
- name: Run generate-bank-template-schema
|
||||
run: ./scripts/generate-bank-template-schema.sh
|
||||
@@ -4909,6 +5162,7 @@ jobs:
|
||||
- test-github-copilot-integration
|
||||
- test-codex-integration
|
||||
- test-cursor-cli-integration
|
||||
- test-zcode-integration
|
||||
- build-ai-sdk-integration
|
||||
- test-ai-sdk-integration-deno
|
||||
- test-opencode-integration
|
||||
@@ -4964,6 +5218,7 @@ jobs:
|
||||
- test-embed
|
||||
- test-embed-windows
|
||||
- verify-embed-control-center-bundle
|
||||
- test-hermes-compat
|
||||
- test-hindsight-all
|
||||
- test-hindsight-agent-sdk
|
||||
- test-claude-agent-sdk-integration
|
||||
|
||||
@@ -5,6 +5,13 @@ build/
|
||||
dist/
|
||||
wheels/
|
||||
*.egg-info
|
||||
# Release builds stage the canonical root license in each package context.
|
||||
/hindsight-clients/python/LICENSE
|
||||
/hindsight-api-slim/LICENSE
|
||||
/hindsight-api/LICENSE
|
||||
/hindsight-all/LICENSE
|
||||
/hindsight-all-slim/LICENSE
|
||||
/hindsight-embed/LICENSE
|
||||
.mcp.json
|
||||
.playwright-mcp/
|
||||
.osgrep
|
||||
@@ -41,6 +48,7 @@ nltk_data/
|
||||
logs/
|
||||
|
||||
.DS_Store
|
||||
.sesskey
|
||||
|
||||
# Generated docs files
|
||||
hindsight-docs/static/llms-full.txt
|
||||
|
||||
@@ -286,6 +286,35 @@ When adding or modifying parameters in the dataplane API (hindsight-api), you mu
|
||||
- Update the client type definition in `lib/api.ts`
|
||||
- Update any UI components that need to use the new parameter
|
||||
|
||||
### Harness Attribution (which coding agent wrote a document)
|
||||
|
||||
`hindsight-integrations/hindsight-coding-agents/` stamps the coding agent on every
|
||||
document it retains, so the control plane can show its logo instead of another
|
||||
`key=value` chip:
|
||||
|
||||
- `metadata.harness = "<id>"` — the authoritative field
|
||||
- tag `harness:<id>` — the same value, so the documents list can filter on it
|
||||
|
||||
The ids are defined by that integration's HookSpecs
|
||||
(`src/harness/hook-lifecycle.ts`) plus the persistent-plugin entrypoints
|
||||
registered in `src/harness/registry.ts`, whose id is their
|
||||
`createPluginEntry(...)` argument — currently `antigravity-cli`, `claude-code`,
|
||||
`cline-cli`, `codex`, `copilot-cli`, `cursor-cli`, `devin-cli`, `grok-build`,
|
||||
`kilo`, `opencode`.
|
||||
|
||||
The control plane resolves the value in
|
||||
`hindsight-control-plane/src/lib/harness-logo.ts` (metadata wins over the tag) and
|
||||
renders it with `components/ui/harness-logo.tsx` in the documents table and the
|
||||
document detail dialog. **Adding a harness to the integration means adding it to
|
||||
that registry in the same change**: copy its icon from
|
||||
`hindsight-docs/static/img/icons/` (or take it from the agent's own brand assets
|
||||
when the docs site carries none) into
|
||||
`hindsight-control-plane/public/img/harness/` and add one entry. Don't register
|
||||
ids nothing writes — a test asserts the registry matches the emitted set, plus an
|
||||
explicit list of retired ids kept so already-retained documents keep their logo.
|
||||
An unregistered harness is not an error: it renders no logo and still shows as
|
||||
ordinary metadata.
|
||||
|
||||
### Adding New Integrations
|
||||
|
||||
Every new integration in `hindsight-integrations/` must satisfy all of the following before it can be merged:
|
||||
|
||||
@@ -298,7 +298,7 @@ client.reflect(bank_id="my-bank", query="What should I know about Alice?")
|
||||
---
|
||||
## Star History
|
||||
|
||||
[](https://www.star-history.com/#vectorize-io/hindsight&type=date&legend=top-left)
|
||||
[](https://github.com/vectorize-io/hindsight/stargazers)
|
||||
---
|
||||
|
||||
## Supported Platforms
|
||||
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 303 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 152 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 153 KiB |
@@ -17,7 +17,7 @@ FROM ghcr.io/vectorize-io/hindsight:latest-slim
|
||||
# `pip install` would fall back to user site-packages and not be visible
|
||||
# to the runtime python.
|
||||
RUN uv pip install --python /app/api/.venv/bin/python --no-cache \
|
||||
'sentence-transformers>=3.3.0' \
|
||||
'sentence-transformers>=5.0.0' \
|
||||
'transformers>=4.53.0' \
|
||||
'torch>=2.6.0'
|
||||
|
||||
|
||||
@@ -3,21 +3,11 @@
|
||||
# pgroonga is a multilingual full-text search extension built on Groonga.
|
||||
# It works out of the box for CJK (Chinese, Japanese, Korean) and other
|
||||
# non-whitespace-segmented languages via the TokenBigram tokenizer.
|
||||
FROM groonga/pgroonga:latest-debian-pg17
|
||||
FROM groonga/pgroonga:4.0.8-debian-17
|
||||
|
||||
# Install pgvector on top of the pgroonga base image (which already provides
|
||||
# pgroonga and the Groonga library).
|
||||
# pgroonga, the Groonga library, and the PostgreSQL PGDG package repository).
|
||||
RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||
build-essential \
|
||||
git \
|
||||
postgresql-server-dev-17 \
|
||||
postgresql-17-pgvector=0.8.6-1.pgdg13+1 \
|
||||
&& apt-get clean \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
RUN cd /tmp && \
|
||||
git clone --branch v0.8.0 https://github.com/pgvector/pgvector.git && \
|
||||
cd pgvector && \
|
||||
make && \
|
||||
make install
|
||||
|
||||
RUN rm -rf /tmp/pgvector && \
|
||||
apt-get purge -y --auto-remove build-essential git postgresql-server-dev-17
|
||||
|
||||
@@ -41,28 +41,43 @@ RUN apt-get update && apt-get install -y \
|
||||
&& rm -rf /var/lib/apt/lists/* \
|
||||
&& pip install --no-cache-dir uv
|
||||
|
||||
# Copy dependency files and README (required by pyproject.toml)
|
||||
# Copy the workspace lock and member metadata before source code so dependency
|
||||
# installation stays cacheable while matching the versions tested in CI.
|
||||
COPY pyproject.toml uv.lock ./
|
||||
COPY hindsight-all/pyproject.toml ./hindsight-all/
|
||||
COPY hindsight-api/pyproject.toml ./hindsight-api/
|
||||
COPY hindsight-api-slim/pyproject.toml ./api/
|
||||
COPY hindsight-api-slim/README.md ./api/
|
||||
|
||||
WORKDIR /app/api
|
||||
COPY hindsight-all-slim/pyproject.toml ./hindsight-all-slim/
|
||||
COPY hindsight-dev/pyproject.toml ./hindsight-dev/
|
||||
COPY hindsight-clients/python/pyproject.toml ./hindsight-clients/python/
|
||||
COPY hindsight-embed/pyproject.toml ./hindsight-embed/
|
||||
RUN ln -s api hindsight-api-slim
|
||||
|
||||
# Sync dependencies using appropriate extras based on INCLUDE_LOCAL_MODELS
|
||||
# local-ml: torch, sentence-transformers, transformers, einops, flashrank, mlx (optional)
|
||||
# embedded-db: pg0-embedded (always included for embedded PostgreSQL support)
|
||||
# ONNX Runtime embeddings are intentionally not bundled into the official
|
||||
# standalone image; install the local-onnx extra in custom images when needed.
|
||||
ENV UV_PROJECT_ENVIRONMENT=/app/api/.venv
|
||||
RUN if [ "$INCLUDE_LOCAL_MODELS" = "true" ]; then \
|
||||
uv sync --extra local-ml --extra embedded-db; \
|
||||
uv sync --locked --package hindsight-api-slim --no-install-package hindsight-api-slim --extra local-ml --extra embedded-db; \
|
||||
else \
|
||||
uv sync --extra embedded-db; \
|
||||
uv sync --locked --package hindsight-api-slim --no-install-package hindsight-api-slim --extra embedded-db; \
|
||||
fi
|
||||
|
||||
# Copy source code (alembic migrations are inside hindsight_api/)
|
||||
WORKDIR /app/api
|
||||
COPY hindsight-api-slim/hindsight_api ./hindsight_api
|
||||
|
||||
# Install the local package (uv sync only installed dependencies, not the package itself)
|
||||
RUN uv pip install -e .
|
||||
# Install the local package from the same validated lock after source is present.
|
||||
WORKDIR /app
|
||||
RUN if [ "$INCLUDE_LOCAL_MODELS" = "true" ]; then \
|
||||
uv sync --locked --package hindsight-api-slim --extra local-ml --extra embedded-db; \
|
||||
else \
|
||||
uv sync --locked --package hindsight-api-slim --extra embedded-db; \
|
||||
fi \
|
||||
&& uv pip check --python /app/api/.venv/bin/python
|
||||
|
||||
# =============================================================================
|
||||
# Stage: SDK Builder (needed for Control Plane)
|
||||
@@ -145,6 +160,8 @@ FROM python:3.11-slim AS api-only
|
||||
WORKDIR /app
|
||||
|
||||
# Note: libicu version varies by Debian version - try common versions in order
|
||||
# Runtime images use uv directly; remove pip build tooling after installation so
|
||||
# vulnerable setuptools-vendored packages and wheel are not shipped in production.
|
||||
RUN apt-get update && apt-get install -y \
|
||||
curl \
|
||||
procps \
|
||||
@@ -154,7 +171,8 @@ RUN apt-get update && apt-get install -y \
|
||||
libossp-uuid16 \
|
||||
&& (apt-get install -y libicu72 2>/dev/null || apt-get install -y libicu74 2>/dev/null || apt-get install -y libicu76 2>/dev/null || true) \
|
||||
&& rm -rf /var/lib/apt/lists/* \
|
||||
&& pip install --no-cache-dir uv
|
||||
&& pip install --no-cache-dir uv \
|
||||
&& pip uninstall --yes setuptools wheel
|
||||
|
||||
RUN useradd -m -s /bin/bash hindsight
|
||||
|
||||
@@ -292,6 +310,8 @@ WORKDIR /app
|
||||
|
||||
# Install Node.js, curl, uv, and system dependencies
|
||||
# Note: libicu version varies by Debian version - try common versions in order
|
||||
# Runtime images use uv directly; remove pip build tooling after installation so
|
||||
# vulnerable setuptools-vendored packages and wheel are not shipped in production.
|
||||
RUN apt-get update && apt-get install -y \
|
||||
curl \
|
||||
procps \
|
||||
@@ -303,7 +323,8 @@ RUN apt-get update && apt-get install -y \
|
||||
&& curl -fsSL https://deb.nodesource.com/setup_20.x | bash - \
|
||||
&& apt-get install -y nodejs \
|
||||
&& rm -rf /var/lib/apt/lists/* \
|
||||
&& pip install --no-cache-dir uv
|
||||
&& pip install --no-cache-dir uv \
|
||||
&& pip uninstall --yes setuptools wheel
|
||||
|
||||
RUN useradd -m -s /bin/bash hindsight
|
||||
|
||||
|
||||
@@ -0,0 +1,341 @@
|
||||
# v2 Knowledge Pages — Implementation Plan
|
||||
|
||||
> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking.
|
||||
|
||||
**Goal:** Make v2 knowledge pages a reliable, cleanly-tiered "wiki" surface: passive `entity_labels` tier-tagging + tag-scoped seeded pages, a `hindsight_*` MCP surface with one active `capture_initiative` verb that creates per-initiative pages linked by tag, and SessionStart/UserPromptSubmit page-roster injection.
|
||||
|
||||
**Architecture:** Shared TS core (`hindsight-integrations/hindsight-coding-agents`). Extraction stays blind to "pages"; classification is intrinsic (`knowledge:<tier>` tags), pages are tag-scoped saved views. Per-initiative navigation via a `relatedPageId:<id>` tag the synthesizer renders into `[[page:<id>]]`, with the Initiatives folder/roster as the guaranteed fallback.
|
||||
|
||||
**Tech Stack:** TypeScript, vitest, tsup bundling. Hindsight REST API (`/knowledge-base/*`, `/mental-models`, `/memories`, bank `/config`).
|
||||
|
||||
**Spec:** `docs/superpowers/specs/2026-07-25-v2-knowledge-pages-design.md`
|
||||
|
||||
**Working dir for all commands:** `hindsight-integrations/hindsight-coding-agents`
|
||||
**Test command:** `npx vitest run <file>` (fast suite; excludes `*.live.test.ts`). Full check: `npx vitest run && npx tsc --noEmit`.
|
||||
|
||||
**Conventions to follow (existing patterns):**
|
||||
- `HindsightClient` HTTP via `this.req("METHOD", this.bankUrl(path), body?)`; JSON via `await r.json()`.
|
||||
- MCP tools are SDK-free `ToolSpec { name, description, inputSchema (ZodRawShape), handler }`; wrap handler bodies in `guarded(...)`; `ok(value)` / `err(e)` result helpers.
|
||||
- Fail-open everywhere in hooks; pure logic separated from stdin/stdout plumbing.
|
||||
- Do NOT add a Claude co-author trailer to any commit.
|
||||
|
||||
---
|
||||
|
||||
## Task 1: Config field `pageRefreshEveryTurns`
|
||||
|
||||
**Files:**
|
||||
- Modify: `src/core/config.ts`
|
||||
- Test: `src/core/config.test.ts`
|
||||
|
||||
- [ ] **Step 1: Write failing test** — assert the default resolves to 10 and an override wins.
|
||||
|
||||
```ts
|
||||
it("pageRefreshEveryTurns defaults to 10 and is overridable", () => {
|
||||
expect(loadConfig({ harness: "claude-code", projectDir: process.cwd() }).pageRefreshEveryTurns).toBe(10);
|
||||
});
|
||||
```
|
||||
(Add an override case mirroring the existing override tests in this file.)
|
||||
|
||||
- [ ] **Step 2: Run** `npx vitest run src/core/config.test.ts` → FAIL (property missing).
|
||||
- [ ] **Step 3: Implement** — add `pageRefreshEveryTurns: number` to the `Config` type and default `10` in the same place `recallMaxTokens`/`reflectTimeoutMs` are defined/merged. Follow the exact merge/layering pattern already used for numeric fields.
|
||||
- [ ] **Step 4: Run** the test → PASS.
|
||||
- [ ] **Step 5: Commit** `git add src/core/config.ts src/core/config.test.ts && git commit -m "feat(core): add pageRefreshEveryTurns config (default 10)"`
|
||||
|
||||
---
|
||||
|
||||
## Task 2: `knowledge-injection.ts` — roster/preamble formatting (pure, new)
|
||||
|
||||
**Files:**
|
||||
- Create: `src/core/knowledge-injection.ts`
|
||||
- Test: `src/core/knowledge-injection.test.ts`
|
||||
|
||||
Pure, SDK-free, no network. Parses the `listPages()` payload and formats the two injections.
|
||||
|
||||
- [ ] **Step 1: Write failing tests**
|
||||
|
||||
```ts
|
||||
import { describe, expect, it } from "vitest";
|
||||
import { parsePageList, buildKnowledgePreamble, buildRosterRefresh } from "./knowledge-injection";
|
||||
|
||||
describe("parsePageList", () => {
|
||||
it("extracts {id,title} from the mental-model list shape, tolerating junk", () => {
|
||||
const raw = { items: [{ id: "p1", name: "Component map" }, { id: "p2", name: "Core concepts" }, { nope: 1 }] };
|
||||
expect(parsePageList(raw)).toEqual([{ id: "p1", title: "Component map" }, { id: "p2", title: "Core concepts" }]);
|
||||
});
|
||||
it("returns [] for null/garbage", () => {
|
||||
expect(parsePageList(null)).toEqual([]);
|
||||
expect(parsePageList(42 as unknown)).toEqual([]);
|
||||
});
|
||||
});
|
||||
|
||||
describe("buildKnowledgePreamble", () => {
|
||||
it("includes guidance, a roster of pages, and a refresh note", () => {
|
||||
const out = buildKnowledgePreamble([{ id: "p1", title: "Component map" }]);
|
||||
expect(out).toContain("<hindsight_knowledge>");
|
||||
expect(out).toContain("Component map");
|
||||
expect(out).toContain("p1");
|
||||
expect(out).toMatch(/hindsight_read_knowledge_page/);
|
||||
});
|
||||
it("has an empty-state line when there are no pages", () => {
|
||||
const out = buildKnowledgePreamble([]);
|
||||
expect(out).toMatch(/no knowledge pages yet|still learning/i);
|
||||
});
|
||||
});
|
||||
|
||||
describe("buildRosterRefresh", () => {
|
||||
it("is a compact 'current pages' block listing ids+titles", () => {
|
||||
const out = buildRosterRefresh([{ id: "p1", title: "Component map" }]);
|
||||
expect(out).toContain("Component map");
|
||||
expect(out).toContain("p1");
|
||||
});
|
||||
it("returns undefined when there are no pages (nothing to refresh)", () => {
|
||||
expect(buildRosterRefresh([])).toBeUndefined();
|
||||
});
|
||||
});
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Run** `npx vitest run src/core/knowledge-injection.test.ts` → FAIL.
|
||||
- [ ] **Step 3: Implement**
|
||||
|
||||
```ts
|
||||
export interface PageRef { id: string; title: string; }
|
||||
|
||||
/** Defensive parse of HindsightClient.listPages() (GET /mental-models?detail=metadata → {items:[{id,name}]}). */
|
||||
export function parsePageList(raw: unknown): PageRef[] {
|
||||
const items = (raw as { items?: unknown })?.items;
|
||||
if (!Array.isArray(items)) return [];
|
||||
const out: PageRef[] = [];
|
||||
for (const it of items) {
|
||||
const id = (it as { id?: unknown })?.id;
|
||||
const name = (it as { name?: unknown })?.name;
|
||||
if (typeof id === "string" && typeof name === "string") out.push({ id, title: name });
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
function roster(pages: PageRef[]): string {
|
||||
return pages.map((p) => `- ${p.title} (${p.id})`).join("\n");
|
||||
}
|
||||
|
||||
/** SessionStart: teach when/why to use pages + list what exists. Empty-state aware. */
|
||||
export function buildKnowledgePreamble(pages: PageRef[]): string {
|
||||
const body = pages.length
|
||||
? `Knowledge pages available in this repository:\n${roster(pages)}`
|
||||
: "No knowledge pages yet — Hindsight is still learning this repo; they'll appear as it processes.";
|
||||
return (
|
||||
"<hindsight_knowledge>\n" +
|
||||
"This repository has a Hindsight knowledge base: curated, continuously-updated pages summarizing its " +
|
||||
"durable engineering knowledge (architecture, components, conventions, key decisions, and in-flight initiatives).\n" +
|
||||
"Before substantial work, consult the relevant pages instead of re-deriving understanding from the code: read " +
|
||||
"Conventions before writing new code, the Component map before changing a subsystem, and an initiative's page " +
|
||||
"before continuing that feature.\n" +
|
||||
`${body}\n` +
|
||||
"Read one with hindsight_read_knowledge_page(page_id). Follow any [[page:<id>]] links you see. The list is " +
|
||||
"re-injected for you periodically as it changes.\n" +
|
||||
"</hindsight_knowledge>"
|
||||
);
|
||||
}
|
||||
|
||||
/** Periodic UserPromptSubmit refresh — compact, or undefined when there's nothing to show. */
|
||||
export function buildRosterRefresh(pages: PageRef[]): string | undefined {
|
||||
if (!pages.length) return undefined;
|
||||
return (
|
||||
"<hindsight_knowledge_refresh>\n" +
|
||||
`Current Hindsight knowledge pages (may have changed):\n${roster(pages)}\n` +
|
||||
"Read any with hindsight_read_knowledge_page(page_id).\n" +
|
||||
"</hindsight_knowledge_refresh>"
|
||||
);
|
||||
}
|
||||
```
|
||||
|
||||
- [ ] **Step 4: Run** the test → PASS.
|
||||
- [ ] **Step 5: Commit** `git add src/core/knowledge-injection.ts src/core/knowledge-injection.test.ts && git commit -m "feat(core): knowledge-injection roster/preamble formatting"`
|
||||
|
||||
---
|
||||
|
||||
## Task 3: `entity_labels` tier vocabulary + configureBank wiring
|
||||
|
||||
**Files:**
|
||||
- Modify: `src/core/missions.ts` (add `KNOWLEDGE_LABELS`)
|
||||
- Modify: `src/core/hindsight.ts` (`configureBank` PATCH sets `entity_labels`)
|
||||
- Test: `src/core/hindsight.*.test.ts` (add/extend a config test with a mock client)
|
||||
|
||||
- [ ] **Step 1: Write failing test** — assert `configureBank` PATCHes `/config` with `entity_labels` containing the `knowledge` group and its five values, and `entities_allow_free_form: true`. Use the existing fetch/req mock pattern from `hindsight.*.test.ts`; capture the PATCH body to `/config` and assert on it.
|
||||
|
||||
- [ ] **Step 2: Run** → FAIL.
|
||||
- [ ] **Step 3: Implement**
|
||||
- In `missions.ts`, export `KNOWLEDGE_LABELS` — the exact object from the spec §4 (`key:"knowledge"`, `type:"multi-values"`, `optional:true`, `tag:true`, the verbose group `description`, and the five value `{value,description}` entries: feature-work, decision, convention, component, concept).
|
||||
- In `hindsight.ts::configureBank`, extend the existing `PATCH .../config` `updates` object to include `entity_labels: [KNOWLEDGE_LABELS]` and `entities_allow_free_form: true`. Import `KNOWLEDGE_LABELS`.
|
||||
- Update the `[bank] configured …` log to mention `entity_labels`.
|
||||
- [ ] **Step 4: Run** → PASS.
|
||||
- [ ] **Step 5: Commit** `git add src/core/missions.ts src/core/hindsight.ts src/core/hindsight.*.test.ts && git commit -m "feat(core): passive knowledge entity_labels tier vocabulary + configureBank wiring"`
|
||||
|
||||
---
|
||||
|
||||
## Task 4: Tag-scoped seeded pages + Initiatives folder + link source_query
|
||||
|
||||
**Files:**
|
||||
- Modify: `src/core/missions.ts` (`PAGES` gain `tags`; Initiatives `source_query` link instruction)
|
||||
- Modify: `src/core/hindsight.ts` (`ensureFolder`, `createPages` sets page `tags` + parents Initiatives under the folder)
|
||||
- Test: `src/core/hindsight.pages.test.ts`
|
||||
|
||||
- [ ] **Step 1: Write failing tests** (mock client `req`):
|
||||
- Each seeded page POST to `/knowledge-base/pages` includes `tags: ["knowledge:<tier>"]` mapped per the spec §5 table.
|
||||
- The Initiatives page is created with `parent_id` equal to the id returned by an Initiatives folder POST to `/knowledge-base/folders`.
|
||||
- `ensureFolder("Initiatives")` returns an existing root folder's id when the tree already contains it (GET `/knowledge-base/tree`) and does NOT POST a duplicate.
|
||||
|
||||
- [ ] **Step 2: Run** → FAIL.
|
||||
- [ ] **Step 3: Implement**
|
||||
- `missions.ts`: add `tags: string[]` to each `PAGES` entry (feature-work/decision/convention/component/concept mapping). Append to the Initiatives `source_query`: *"When a source memory carries a tag of the form `relatedPageId:<id>`, include a Markdown link `[[page:<id>]]` to that page in the summary, so each initiative links to its detailed page."*
|
||||
- `hindsight.ts`: add
|
||||
```ts
|
||||
/** Find a root folder by name (case-insensitive) or create it; returns its id. */
|
||||
async ensureFolder(name: string): Promise<string | undefined> {
|
||||
try {
|
||||
const tree = (await (await this.req("GET", this.bankUrl("/knowledge-base/tree"))).json()) as
|
||||
{ roots?: { id?: string; kind?: string; name?: string }[] };
|
||||
const hit = (tree.roots || []).find((n) => n.kind === "folder" && (n.name || "").toLowerCase() === name.toLowerCase());
|
||||
if (hit?.id) return hit.id;
|
||||
} catch { /* fall through to create */ }
|
||||
try {
|
||||
const r = await this.req("POST", this.bankUrl("/knowledge-base/folders"), { name });
|
||||
return ((await r.json()) as { id?: string }).id;
|
||||
} catch { return undefined; }
|
||||
}
|
||||
```
|
||||
- In `createPages()`: before the loop, `const initiativesFolderId = await this.ensureFolder("Initiatives");`. For each page, build body `{ name, source_query, tags: p.tags, parent_id: <initiativesFolderId if this is the Initiatives page else undefined>, trigger: { fact_types:[...], refresh_after_consolidation:true } }`. (Page-level `tags` drives synthesis scoping via `RefreshTagFiltering`; `tags_match` defaults to `all_strict` when tags present.)
|
||||
- [ ] **Step 4: Run** → PASS.
|
||||
- [ ] **Step 5: Commit** `git add src/core/missions.ts src/core/hindsight.ts src/core/hindsight.pages.test.ts && git commit -m "feat(core): tag-scope seeded pages, Initiatives folder, relatedPageId link source_query"`
|
||||
|
||||
---
|
||||
|
||||
## Task 5: Client helpers — per-initiative page + marker retain
|
||||
|
||||
**Files:**
|
||||
- Modify: `src/core/hindsight.ts` (`captureInitiative`)
|
||||
- Test: `src/core/hindsight.pages.test.ts`
|
||||
|
||||
- [ ] **Step 1: Write failing tests** (mock `req`):
|
||||
- `captureInitiative({title:"Retry backoff for the uploader", summary:"…"})` → derives slug `retry-backoff-for-the-uploader`, POSTs a page id `initiative-<slug>` to `/knowledge-base/pages` with `parent_id` = the Initiatives folder and `tags: ["knowledge:feature-work"]`, AND POSTs a marker to `/memories` (via `retain`) tagged `["knowledge:feature-work","relatedPageId:initiative-<slug>"]`, strategy `session` or `document` (pick `document`), `async:true`. Returns `{ page_id: "initiative-<slug>" }`.
|
||||
- Slug is deterministic and identical between the page id and the `relatedPageId:` tag value.
|
||||
- Enhancement path: `captureInitiative({title, summary, relatesToPageId:"initiative-x"})` POSTs NO new page; marker tagged `relatedPageId:initiative-x`; returns `{ page_id: "initiative-x" }`.
|
||||
|
||||
- [ ] **Step 2: Run** → FAIL.
|
||||
- [ ] **Step 3: Implement**
|
||||
|
||||
```ts
|
||||
private slugify(s: string): string {
|
||||
return s.toLowerCase().replace(/[^a-z0-9]+/g, "-").replace(/^-+|-+$/g, "").slice(0, 60) || "initiative";
|
||||
}
|
||||
|
||||
/** Active-path capture: register a major feature as a per-initiative page + a tagged marker memory. */
|
||||
async captureInitiative(args: { title: string; summary: string; relatesToPageId?: string }): Promise<{ page_id: string }> {
|
||||
const pageId = args.relatesToPageId ?? `initiative-${this.slugify(args.title)}`;
|
||||
if (!args.relatesToPageId) {
|
||||
const folderId = await this.ensureFolder("Initiatives");
|
||||
await this.req("POST", this.bankUrl("/knowledge-base/pages"), {
|
||||
name: args.title,
|
||||
source_query: `Summarize the "${args.title}" initiative: what is being built or changed and why, and its current state — drawn from the project's memory.`,
|
||||
parent_id: folderId,
|
||||
tags: ["knowledge:feature-work", `relatedPageId:${pageId}`],
|
||||
trigger: { fact_types: ["world", "experience", "observation"], refresh_after_consolidation: true },
|
||||
});
|
||||
}
|
||||
const verb = args.relatesToPageId ? "Enhancement to an existing initiative" : "New initiative";
|
||||
const content = `${verb}: ${args.title}. ${args.summary}`;
|
||||
await this.retain(content, "initiative marker", pageId /* not a stable doc id requirement; see note */,
|
||||
["knowledge:feature-work", `relatedPageId:${pageId}`], "document", { async: true });
|
||||
return { page_id: pageId };
|
||||
}
|
||||
```
|
||||
- NOTE: use a UNIQUE document id per marker (e.g. `initiative-marker-<slug>-<n>`), NOT `pageId`, so repeated enhancement captures accrue instead of replacing. Since `Date.now()` is fine here (runtime, not a workflow script), suffix with a timestamp: `initiative-marker-${this.slugify(args.title)}-${Date.now()}`. Keep the `relatedPageId` tag equal to `pageId`.
|
||||
- Confirm `retain(content, context, documentId, tags, strategy, opts)` signature matches current `HindsightClient.retain`.
|
||||
|
||||
- [ ] **Step 4: Run** → PASS.
|
||||
- [ ] **Step 5: Commit** `git add src/core/hindsight.ts src/core/hindsight.pages.test.ts && git commit -m "feat(core): captureInitiative — per-initiative page + relatedPageId marker"`
|
||||
|
||||
---
|
||||
|
||||
## Task 6: MCP surface — `hindsight_*` grounding + `capture_initiative`; drop page CRUD
|
||||
|
||||
**Files:**
|
||||
- Modify: `src/core/knowledge-tools.ts`
|
||||
- Modify: `src/mcp-server.ts` (only if it references removed tool names)
|
||||
- Test: `src/core/knowledge-tools.test.ts`, `src/mcp-server.test.ts` (tool-count assertions)
|
||||
|
||||
- [ ] **Step 1: Write failing tests**
|
||||
- `buildKnowledgeTools(client, bankId)` returns exactly these tool names: `hindsight_get_current_bank`, `hindsight_list_knowledge_pages`, `hindsight_read_knowledge_page`, `hindsight_search_memory`, `hindsight_capture_initiative`, `hindsight_ingest_document`. (Assert the set; update any count assertion.)
|
||||
- `hindsight_capture_initiative` handler calls `client.captureInitiative` with `{title, summary, relatesToPageId?}` and returns the page id (mock client).
|
||||
- No `create_page` / `update_page` / `delete_page` tools are present.
|
||||
- Each tool still fails closed via `guarded` (a thrown client error → `isError:true`, no throw).
|
||||
|
||||
- [ ] **Step 2: Run** → FAIL.
|
||||
- [ ] **Step 3: Implement**
|
||||
- Rebuild the `buildKnowledgeTools` list: rename read/recall/ingest/bank tools to the `hindsight_*` names; drop `create_page`/`update_page`/`delete_page`; add `hindsight_capture_initiative` with `inputSchema { title: z.string(), summary: z.string(), relates_to_page_id: z.string().optional() }` calling `client.captureInitiative({ title, summary, relatesToPageId: relates_to_page_id })`.
|
||||
- Use the **verbatim agent-facing `description` strings** from the spec §6 / the brainstorm (grounding tools + the explicit WHEN/WHEN-NOT `capture_initiative` description).
|
||||
- Update `mcp-server.ts` only if it enumerates tool names; otherwise it consumes `buildKnowledgeTools` generically and needs no change.
|
||||
- [ ] **Step 4: Run** `npx vitest run src/core/knowledge-tools.test.ts src/mcp-server.test.ts` → PASS.
|
||||
- [ ] **Step 5: Commit** `git add src/core/knowledge-tools.ts src/mcp-server.ts src/core/knowledge-tools.test.ts src/mcp-server.test.ts && git commit -m "feat(mcp): hindsight_* grounding tools + capture_initiative; remove raw page CRUD from agent"`
|
||||
|
||||
---
|
||||
|
||||
## Task 7: SessionStart — preamble + roster
|
||||
|
||||
**Files:**
|
||||
- Modify: `src/core/session-start.ts`
|
||||
- Test: `src/core/session-start.test.ts`
|
||||
|
||||
- [ ] **Step 1: Write failing tests**
|
||||
- `buildSessionStartContext` now fetches pages via the client and injects `buildKnowledgePreamble(...)` instead of the static `KNOWLEDGE_MISSION`. Extend the `SeedContextClient` interface with `listPages(): Promise<unknown>`; the mock returns `{items:[{id:"p1",name:"Component map"}]}` and the output contains "Component map".
|
||||
- listPages failure is fail-open: the preamble still renders (empty-state) and the seed logic is unaffected.
|
||||
|
||||
- [ ] **Step 2: Run** → FAIL.
|
||||
- [ ] **Step 3: Implement**
|
||||
- Add `listPages` to `SeedContextClient`.
|
||||
- Replace the `parts.push(KNOWLEDGE_MISSION)` line with: fetch `const pages = parsePageList(await client.listPages().catch(() => null));` then `parts.push(buildKnowledgePreamble(pages));`. Import from `./knowledge-injection`.
|
||||
- Remove the now-unused `KNOWLEDGE_MISSION` export if nothing else references it (grep first; keep if referenced).
|
||||
- [ ] **Step 4: Run** → PASS.
|
||||
- [ ] **Step 5: Commit** `git add src/core/session-start.ts src/core/session-start.test.ts && git commit -m "feat(core): SessionStart injects page roster + guidance preamble"`
|
||||
|
||||
---
|
||||
|
||||
## Task 8: UserPromptSubmit — hook-counted periodic roster refresh
|
||||
|
||||
**Files:**
|
||||
- Modify: `src/core/hook.ts`
|
||||
- Test: `src/core/hook.test.ts`
|
||||
|
||||
- [ ] **Step 1: Write failing tests**
|
||||
- The session cache round-trips `{answer, turns}`; each `buildHookOutput` call increments `turns`.
|
||||
- Add `listPages` to the `HookClient` interface. On a turn where `turns % cfg.pageRefreshEveryTurns === 0`, the output includes `buildRosterRefresh(...)` content (assert "Component map" appears); on other turns it does not.
|
||||
- Refresh is fail-open (a `listPages` rejection doesn't break recall/injection).
|
||||
- First-turn behavior (reflect) unchanged.
|
||||
|
||||
- [ ] **Step 2: Run** → FAIL.
|
||||
- [ ] **Step 3: Implement**
|
||||
- Extend the cache read/write to `{ answer?: string; turns?: number }`. Compute `const turns = (cached.turns ?? 0) + 1;` and persist it alongside `answer`.
|
||||
- Add `listPages(): Promise<unknown>` to `HookClient`.
|
||||
- After computing `memBlock`, if `cfg.pageRefreshEveryTurns > 0 && turns % cfg.pageRefreshEveryTurns === 0`, `try { const refresh = buildRosterRefresh(parsePageList(await client.listPages())); if (refresh) blocks.push(refresh); } catch { /* fail-open */ }`. Kick the `listPages` call off concurrently with recall to avoid added latency.
|
||||
- Import from `./knowledge-injection`.
|
||||
- [ ] **Step 4: Run** `npx vitest run src/core/hook.test.ts` → PASS.
|
||||
- [ ] **Step 5: Commit** `git add src/core/hook.ts src/core/hook.test.ts && git commit -m "feat(core): UserPromptSubmit hook-counted periodic page-roster refresh"`
|
||||
|
||||
---
|
||||
|
||||
## Task 9: Full check + LLM behavior (live) verification
|
||||
|
||||
**Files:**
|
||||
- Modify: `src/system.live.test.ts` (add coverage; runs only under `HINDSIGHT_LIVE_E2E=1`)
|
||||
|
||||
- [ ] **Step 1: Full fast suite + types** — `npx vitest run && npx tsc --noEmit` → all green.
|
||||
- [ ] **Step 2: Add a live assertion** (guarded by the existing live env flag) that after seeding a small repo + one `captureInitiative`, the Initiatives page content contains a `[[page:initiative-…]]` link (verifies the `relatedPageId` → link rendering end-to-end). Keep it in the live suite; do not run in the fast job.
|
||||
- [ ] **Step 3: Manual/live run** (optional, operator): `HINDSIGHT_API_URL=http://localhost:8888 npm run test:live`.
|
||||
- [ ] **Step 4: Commit** `git add src/system.live.test.ts && git commit -m "test(live): initiative page renders relatedPageId link end-to-end"`
|
||||
|
||||
---
|
||||
|
||||
## Final review
|
||||
|
||||
- [ ] Dispatch a final code-reviewer over the whole change set against the spec (`docs/superpowers/specs/2026-07-25-v2-knowledge-pages-design.md`).
|
||||
- [ ] Rebuild + dev-install the `claude-code-v2` bundle so the running plugin picks up the new hooks/MCP (`bash scripts/dev-install.sh`); do not push/PR without explicit consent.
|
||||
- [ ] Note deferred follow-ups: session drill-down tag, `capture_decision`, `gotcha` tier, older-bank reseed requirement.
|
||||
@@ -0,0 +1,135 @@
|
||||
# v2 Knowledge Pages — Design Spec
|
||||
|
||||
**Status:** approved in brainstorm (2026-07-25), pending implementation plan
|
||||
**Scope:** `hindsight-integrations/hindsight-coding-agents` (shared TS core) + `claude-code-v2` wrapper
|
||||
**Motivation:** make knowledge pages a real, trustworthy "wiki" surface for the vectorize-crm demo (the `knowledge-pages-as-trust-surface` principle) — the agent reliably knows what pages exist, pages are cleanly tiered instead of blended, and major initiatives become first-class, linkable pages.
|
||||
|
||||
---
|
||||
|
||||
## 1. Problem
|
||||
|
||||
Three gaps in the current v2 branch:
|
||||
|
||||
1. **Page discovery is a blind fetch.** SessionStart injects a static `KNOWLEDGE_MISSION` telling the agent to call `agent_knowledge_list_pages`, but hands it **no roster** — the agent never learns a page exists unless it independently decides to call the tool. Per-turn recall injects facts, not pages.
|
||||
2. **Pages are blended.** Neither the seeded `PAGES` nor the agent's `create_page` tool scope synthesis by tag, so every page synthesizes from the whole bank filtered only by `fact_type`. Git-log, session, and survey memories all bleed into every page.
|
||||
3. **No first-class initiative tracking / linking.** Hindsight has no native page-to-page links. A "major feature" leaves no durable, navigable page a future session can pick up.
|
||||
|
||||
## 2. Principles applied
|
||||
|
||||
- Automatic/visible value; zero out-of-band CLI; memory beats code search; knowledge pages as a trust surface; minimal post-setup burden.
|
||||
- Modular units, small files, follow existing patterns (per-hook specs, fail-open, unit-testable pure cores).
|
||||
- The **memory extractor never knows what a "page" is.** Classification is by the fact's *intrinsic* nature; pages are application-side saved views. No abstraction leak into extraction.
|
||||
|
||||
## 3. Architecture overview
|
||||
|
||||
Two complementary curation paths + a discovery layer:
|
||||
|
||||
- **Passive (automatic):** `entity_labels` schema-forces the extractor to tag qualifying facts `knowledge:<tier>`. Seeded **tier pages** each filter on one tier tag. No agent effort.
|
||||
- **Active (high-signal):** one intent-named MCP verb, `hindsight_capture_initiative`, lets the agent register a major feature as a **per-initiative page** with a tag-based link back from the aggregate Initiatives page.
|
||||
- **Discovery:** SessionStart injects guidance + the page roster; the UserPromptSubmit hook re-injects a fresh roster on a fixed cadence (hook-counted, not model-counted).
|
||||
|
||||
## 4. `entity_labels` — passive tier tagging
|
||||
|
||||
One hierarchical bank config group, set by `configureBank` at seed time:
|
||||
|
||||
```jsonc
|
||||
{
|
||||
"key": "knowledge",
|
||||
"type": "multi-values", // 0, 1, or several — empty is normal
|
||||
"optional": true,
|
||||
"tag": true, // emits knowledge:<value> onto the fact's tags
|
||||
"description": "Routing labels for this project's Hindsight KNOWLEDGE PAGES — curated, human-readable summaries of the repo's DURABLE engineering knowledge (architecture, key decisions, conventions, ongoing initiatives), each page rebuilt automatically from the facts labeled for it. Mark a fact only when it is durable, reusable knowledge a developer would still want surfaced in future sessions. IMPORTANT: leave this EMPTY for routine, transient, or operational facts — a passing test, a one-off command, a status update, a debugging dead-end. MOST facts should get no label here. Assign more than one value only when the fact genuinely fits several.",
|
||||
"values": [
|
||||
{ "value": "feature-work", "description": "A new feature, initiative, or enhancement being planned or built — the capability being added and the intent behind it. Not routine bug-fixes or chores." },
|
||||
{ "value": "decision", "description": "A technical decision that will constrain future work, with its rationale — why this approach was chosen over alternatives, or a rule deliberately adopted." },
|
||||
{ "value": "convention", "description": "An established way this project does things — naming, structure, testing, error handling, or another recurring pattern a contributor is expected to follow." },
|
||||
{ "value": "component", "description": "What a specific module, file, service, or subsystem is responsible for, or how components depend on and connect to one another." },
|
||||
{ "value": "concept", "description": "A domain concept, key abstraction, or piece of project vocabulary a new contributor must understand to work effectively." }
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
Notes:
|
||||
- `tag: true` → `_inject_label_tags` copies each `knowledge:<value>` onto the fact's `tags` (no extra query infra).
|
||||
- Selectivity (multi-values + "mostly empty" instruction) prevents force-fitting routine facts into a tier.
|
||||
|
||||
## 5. Seeded tier pages (tag-scoped)
|
||||
|
||||
Created via `/knowledge-base/pages` (supports `tags`, `trigger`, `parent_id`) — **not** `/mental-models`. Each `PAGES` entry gains a `trigger.tags` pin:
|
||||
|
||||
| Page | `trigger.tags` |
|
||||
| --- | --- |
|
||||
| Initiatives and enhancements | `["knowledge:feature-work"]` |
|
||||
| Key decisions and rationale | `["knowledge:decision"]` |
|
||||
| Conventions and patterns | `["knowledge:convention"]` |
|
||||
| Component map | `["knowledge:component"]` |
|
||||
| Core concepts | `["knowledge:concept"]` |
|
||||
|
||||
`tags_match` strict enough to exclude untagged facts (`all_strict`/`any_strict`). Tag matching is exact set-ops (no wildcards) — this is *why* the vocabulary is fixed, not per-feature.
|
||||
|
||||
## 6. MCP surface
|
||||
|
||||
Raw page CRUD (`create_page`/`update_page`/`delete_page`) is **removed** from the agent. The agent sees grounding tools + one capture verb. Naming convention: `hindsight_*`.
|
||||
|
||||
**Grounding**
|
||||
- `hindsight_list_knowledge_pages` `{}` — roster: id, title, one-line coverage. (agent-facing description as drafted in brainstorm)
|
||||
- `hindsight_read_knowledge_page` `{ page_id }` — full page content; follow `[[page:<id>]]` links by re-calling.
|
||||
- `hindsight_search_memory` `{ query, max_tokens? }` — raw fact recall for specifics pages don't cover.
|
||||
- `hindsight_get_current_bank` `{}` — minor introspection (kept).
|
||||
|
||||
**Capture**
|
||||
- `hindsight_capture_initiative` `{ title, summary, relates_to_page_id? }` — the one active verb. Explicit WHEN / WHEN-NOT description (as drafted). Returns the initiative page id.
|
||||
- `hindsight_ingest_document` `{ title, content }` — existing `agent_knowledge_ingest`, reframed.
|
||||
|
||||
(Full agent-facing descriptions are captured verbatim in the brainstorm thread and will be reproduced in the implementation plan.)
|
||||
|
||||
## 7. `hindsight_capture_initiative` mechanism
|
||||
|
||||
- Derive one slug `S` from `title`. Page id = `initiative-<S>`. **The slug in the tag and the page id are the same token, derived once** (cannot drift).
|
||||
- **New initiative** (`relates_to_page_id` omitted):
|
||||
1. Create page `initiative-<S>` (title from `title`, `source_query` about that initiative) under an **"Initiatives" folder** (tag-scoped).
|
||||
2. Retain a marker memory (text = title + summary) tagged `["knowledge:feature-work", "relatedPageId:initiative-<S>"]`. **No session tag** (decided — the MCP server has no Claude session id; faking one wouldn't link to the Stop write-back's `conversation:<sessionId>` doc anyway).
|
||||
- **Enhancement** (`relates_to_page_id` given): marker only, `relatedPageId = relates_to_page_id`; no new page. Re-invoking for the same initiative accrues markers → the page re-synthesizes with progress.
|
||||
|
||||
### Link survival (why `relatedPageId` as a tag, not in prose)
|
||||
|
||||
A tag is set directly via the retain `tags` param — it **bypasses LLM extraction entirely**, so it's guaranteed present verbatim (no REF-ID-style preservation needed at extraction). Verified: the reflect/synthesis path SELECTs `tags` and serializes facts via `_prune_nulls(model_dump())`, which keeps non-empty tags → **the synthesis LLM sees the tag.** The **Initiatives page `source_query`** instructs: *"when a memory carries a `relatedPageId:<id>` tag, emit a `[[page:<id>]]` link to it."* The link id is generated from the tag value at synthesis time, so it always matches the created page id.
|
||||
|
||||
- Only **Stage 2 (synthesis)** is probabilistic now (bounded token budget may omit some entries when there are many).
|
||||
- **Guaranteed fallback:** the per-initiative page always exists (created via API, independent of any LLM stage) and appears in the **Initiatives folder / injected roster**, so navigation works even if a synthesized inline link drops.
|
||||
|
||||
## 8. Page-access injection
|
||||
|
||||
- **SessionStart** (`session-start.ts`): replace static `KNOWLEDGE_MISSION` with a preamble = (a) guidance on *when/why* to consult pages, (b) the roster fetched via `client.listPages()` (`- <title> (<id>)`, empty-state aware), (c) a note that the list refreshes periodically. Cold repo → empty roster line; roster comes alive mid-session as seeding/survey complete.
|
||||
- **UserPromptSubmit** (`hook.ts`): extend the per-session cache (`{answer}` → `{answer, turns}`); the **hook** counts user turns and, roughly every `pageRefreshEveryTurns` (default 10, approximate), calls `listPages()` and injects a compact roster refresh. Runs concurrently with recall; **fail-open** (a refresh error never blocks the turn).
|
||||
- **Shared formatting** (new `core/knowledge-injection.ts`, SDK-free/unit-testable): `parsePageList(raw) -> {id,title}[]`, `buildKnowledgePreamble(pages)`, `buildRosterRefresh(pages)`.
|
||||
- **Config:** `pageRefreshEveryTurns` (default 10).
|
||||
|
||||
## 9. Non-goals / deferred
|
||||
|
||||
- Session drill-down tag on captured markers (dropped — see §7).
|
||||
- `hindsight_capture_decision` and other capture verbs (passive path covers those tiers; revisit if the aggregate pages aren't sharp enough).
|
||||
- A `gotcha`/`pitfall` tier (five tiers for now).
|
||||
- Native page-to-page links / backlinks (Hindsight has none; we approximate via folder tree + `relatedPageId`-driven `[[page:<id>]]`).
|
||||
|
||||
## 10. Risks / migration
|
||||
|
||||
- **Older banks** need re-seeding to pick up the new `entity_labels`, the `session` retain strategy, and the tag-scoped page triggers (`configureBank` sets them). User is starting fresh with v2 banks, so acceptable; live retain fails open otherwise.
|
||||
- **Stage-2 synthesis omission** for large initiative counts — mitigated by the folder/roster fallback.
|
||||
- **Instruction adherence** for the `source_query` link-rendering and the label selectivity — both are LLM-following behaviors; cover with an `hs_llm_core` judge test, and the deterministic mechanics (tag injection, roster formatting, slug/id equality, hook turn-counting) with fast unit tests.
|
||||
|
||||
## 11. Testing
|
||||
|
||||
- **Deterministic unit tests:** `knowledge-injection` formatting + empty-state; hook turn-counter + cadence; `capture_initiative` slug→id→tag equality and request shape (mock client); tag-scoped page request bodies; entity_labels config emitted by `configureBank`.
|
||||
- **LLM judge test (`hs_llm_core`):** label selectivity (routine facts get no `knowledge:*`), and `relatedPageId` → `[[page:<id>]]` rendering in a synthesized Initiatives page.
|
||||
|
||||
## 12. File map (anticipated)
|
||||
|
||||
- `src/core/knowledge-injection.ts` (new) — roster/preamble formatting.
|
||||
- `src/core/session-start.ts` — preamble + roster.
|
||||
- `src/core/hook.ts` — cache `{answer,turns}` + periodic roster refresh.
|
||||
- `src/core/config.ts` — `pageRefreshEveryTurns`.
|
||||
- `src/core/missions.ts` — `entity_labels` group; tag-scoped `PAGES`; Initiatives `source_query` link instruction.
|
||||
- `src/core/hindsight.ts` — `configureBank` sets `entity_labels`; `createPages` pins `trigger.tags` + Initiatives folder; new `createInitiativePage`/marker retain helpers.
|
||||
- `src/core/knowledge-tools.ts` — new `hindsight_*` grounding + `capture_initiative` tools; remove raw page CRUD from agent surface.
|
||||
- Tests alongside each.
|
||||
@@ -0,0 +1,139 @@
|
||||
# Reflect + Pages Runtime — Design Spec
|
||||
|
||||
**Status:** decided (2026-07-27), reconciles the earlier reflect-based runtime with the recall-based v2 into one opinionated path
|
||||
**Scope:** `hindsight-integrations/hindsight-coding-agents` (shared TS core) + `claude-code-v2` wrapper
|
||||
**Motivation:** the 33-task coding benchmark showed the v2 recall-per-prompt runtime *underperforms no memory* (35.0 mean corrections vs 32.0 baseline), while the earlier reflect-injection runtime beats baseline by 22% (25.0). This spec restores reflect as the only deep-memory path and replaces raw per-turn recall with lightweight injection from knowledge pages — "fast like recall, organized like reflect" — keeping v2's page/curation machinery where it earned its place and deleting it where it didn't.
|
||||
|
||||
---
|
||||
|
||||
## 1. Problem
|
||||
|
||||
Two prior iterations, each half right:
|
||||
|
||||
1. **Reflect runtime (v1):** one agentic REFLECT over the bank at session start, cached and re-injected every turn. Benchmark-proven (25.0 mean corrections) — but nothing surfaced mid-session; a task that drifted away from the first message got stale context.
|
||||
2. **Recall runtime (v2):** per-prompt recall injection for turn-by-turn visibility, plus knowledge pages as a trust surface. But raw recall injects unsynthesized fact fragments — noise that *hurt*: 35.0 mean corrections, worse than running with no memory at all.
|
||||
|
||||
| Runtime | Mean corrections (33-task benchmark) | vs no-memory (32.0) |
|
||||
| --- | --- | --- |
|
||||
| Reflect-injection (v1) | **25.0** | **−22%** |
|
||||
| Recall-per-prompt (v2) | 35.0 | +9% (regression) |
|
||||
| No memory | 32.0 | baseline |
|
||||
|
||||
The reconciliation: keep reflect's synthesis quality as the deep path, keep v2's per-turn visibility principle, but source the per-turn material from the already-synthesized knowledge pages instead of raw recall.
|
||||
|
||||
## 2. Decisions
|
||||
|
||||
Explicit, decided — not options:
|
||||
|
||||
1. **Reflect restored** as the only deep-memory path (session-start, agentic synthesis, cached + re-injected every turn).
|
||||
2. **Recall removed from the runtime** entirely. No per-prompt `recall` call.
|
||||
3. **No `memoryMode` flag.** One opinionated path; config is for environment, naming, and harness wiring only — never behavior selection.
|
||||
4. **Sections, not pages, are the per-turn injection unit** — locally matched, budget-trimmed, provenance-labeled.
|
||||
5. **JSON turn transcripts** replace the markdown tool-call transcript in the Stop-hook write-back, with compact action entries.
|
||||
6. **No tags / no `entity_labels`.** The server re-synthesizes pages after consolidation; "living pages" needs no client-side tagging machinery.
|
||||
|
||||
## 3. Runtime path — session start
|
||||
|
||||
Three steps, in order, all inside existing hooks (no out-of-band CLI):
|
||||
|
||||
### 3a. Cold-repo bootstrap (kept from v2)
|
||||
|
||||
On a bank with no prior memories: automatic shallow gitlog seed + codebase survey, exactly as v2 does it. The user never runs a setup command; the first session self-seeds. (Deep ingestion of that history is §7 — the seed here stays instant.)
|
||||
|
||||
### 3b. REFLECT once, on the first task message
|
||||
|
||||
The benchmark-proven core:
|
||||
|
||||
- On the first user prompt of the session, run one **REFLECT** — agentic synthesis over the whole bank, prompted to return the *root-cause decision with exact values* (concrete file paths, config values, version numbers — not summaries of summaries).
|
||||
- Cache the result per session; **re-inject it every turn**. It is the session's durable deep context.
|
||||
- One LLM-backed call per session, on the message that actually states the task — not on session-open, where there is nothing to reflect about.
|
||||
|
||||
### 3c. Page index build
|
||||
|
||||
Fetch all knowledge pages once (existing `listPages` + page reads), split each page at headings into **sections**, and build a **local section index** in the hook process. This index is what every subsequent turn matches against (§4) — no further server calls on the hot path.
|
||||
|
||||
## 4. Runtime path — every turn
|
||||
|
||||
Per-turn visibility, satisfied at ~zero latency and ~zero cost. Injection sources from **knowledge pages, not raw recall** — the material is already synthesized and organized; the turn hook only *selects* from it.
|
||||
|
||||
Mechanism (local, deterministic — no server call, no LLM call):
|
||||
|
||||
| Aspect | Design |
|
||||
| --- | --- |
|
||||
| Unit | Page **sections** (pages split at headings at index-build time) |
|
||||
| Matching | Lexical: prompt scored against each section by weighted term overlap; **heading hits weighted higher** than body hits |
|
||||
| Selection | Top 2–3 sections |
|
||||
| Budget | Trimmed to a **~700-token total** |
|
||||
| Provenance | Each snippet labeled `From <page> › <section>` + a tool pointer to read the full page |
|
||||
| Floor | A minimum-score threshold below which **nothing is injected** — silence over noise |
|
||||
| Refresh | Section index rebuilt on the existing 10-turn roster cadence (`pageRefreshEveryTurns`) |
|
||||
|
||||
The score floor is load-bearing: the benchmark showed that injecting weak matches is worse than injecting nothing (v2's regression). An empty injection is a correct outcome, not a failure mode.
|
||||
|
||||
## 5. Write-back
|
||||
|
||||
The Stop-hook session retain is **kept** — same trigger, same fail-open behavior. What changes is the transcript format handed to extraction:
|
||||
|
||||
- **JSON turns**, not markdown: an array of `{ "role": "user" | "assistant", "text": ... }` entries for the conversational content.
|
||||
- Tool calls collapse to **compact one-line action entries**: `{ "role": "action", "text": "Edit boltons/strutils.py" }` — tool name + primary target only, **no arguments, no outputs**.
|
||||
|
||||
Rationale: extraction keeps the concrete artifacts (which files were touched, what actions occurred) without the transcript noise of full tool payloads — the markdown tool-call dumps were volume without signal.
|
||||
|
||||
## 6. Knowledge pages
|
||||
|
||||
Simplified from the v2 spec:
|
||||
|
||||
- **Dropped: tags and `entity_labels`** (v2 spec §4–5). The server already re-synthesizes pages after consolidation, so pages stay "living" with no client-side routing machinery. The extractor-never-knows-about-pages principle now holds trivially — there is nothing to route.
|
||||
- **Creation paths:**
|
||||
1. **Seeded taxonomy** at bank creation (the fixed page set, as today, minus tag triggers).
|
||||
2. **Agent-driven `capture_initiative`** at plan approval — the one active capture verb survives from v2.
|
||||
3. **Organic splitting** of pages that outgrow their scope is a **server/curator concern**, not a client feature.
|
||||
|
||||
## 7. Ingestion — progressive background deepening
|
||||
|
||||
*Status: design accepted, implementation phased separately.*
|
||||
|
||||
Replaces the manual backfill CLI as the user-facing path (the CLI was out-of-band burden; nobody runs it). The principle: converge to full-depth history through normal usage, with zero user action.
|
||||
|
||||
1. **Instant shallow seed** — the gitlog seed from §3a; the session is useful immediately.
|
||||
2. **Background deepening** — a background worker deep-ingests **per-commit-with-diffs, incrementally**, never blocking a turn.
|
||||
3. **Working-set prioritization** — commits are ingested in order of relevance to what the agent is actually doing: files the agent reads/edits get their commit histories ingested **first**. Depth arrives where it pays off.
|
||||
4. **Checkpointing** — progress persists across sessions; each session resumes deepening where the last left off, converging to full depth over normal usage.
|
||||
|
||||
The **backfill CLI survives as an internal tool** (benchmark setup, CI bank preparation) — it is no longer a documented user path.
|
||||
|
||||
## 8. Gap analysis — v2 principles under this design
|
||||
|
||||
| v2 principle | How this design satisfies it |
|
||||
| --- | --- |
|
||||
| See-it-working (automatic, visible value) | Reflect answer visible from turn 1; page-section snippets appear with explicit `From <page> › <section>` provenance, so the user sees memory working — and the score floor keeps it from visibly misfiring. |
|
||||
| No out-of-band CLI | Cold-repo auto-seed kept (§3a); backfill CLI demoted to internal-only, replaced by background deepening (§7). Nothing requires a terminal command. |
|
||||
| Reuse-over-reinvent | Reflect, `listPages`, Stop-hook retain, `capture_initiative`, and the 10-turn refresh cadence are all existing machinery recombined; the only new code is the local section index and matcher — deliberately dumb (lexical, no LLM). |
|
||||
| Preserve-intent | Reflect is prompted for root-cause decisions with exact values; JSON transcripts keep concrete action artifacts; per-commit-with-diffs deepening captures *why* the code changed, not just that it did. |
|
||||
| Near-zero-burden | No config flags to choose, no CLI to run, no tags to maintain; one LLM call per session start, everything else local. |
|
||||
|
||||
## 9. Verification gates
|
||||
|
||||
Ship gates, in order:
|
||||
|
||||
1. **Reflect-restored benchmark:** the restored runtime must recover **~25 mean corrections at n=2 on identical banks** to the original reflect run. This proves the restoration is faithful before anything is layered on.
|
||||
2. **Reflect+pages benchmark:** with per-turn section injection enabled, the score **must not regress** vs reflect-alone. Section injection earns its place by not hurting; any regression points at the floor/budget tuning.
|
||||
3. **Live system suite:** existing hook/integration suite updated for the new path — reflect caching + per-turn re-injection, section index build/refresh, score-floor silence, JSON transcript shape, action-entry compaction. Deterministic pieces (matcher scoring, budget trim, provenance formatting, transcript serialization) as fast unit tests.
|
||||
|
||||
## 10. Non-goals / deferred
|
||||
|
||||
- Any per-turn LLM or server call for injection (explicitly excluded — the local matcher is the whole point).
|
||||
- Semantic/embedding-based section matching (revisit only if lexical matching demonstrably misses; start dumb).
|
||||
- Client-side page splitting or curation (server/curator concern, §6).
|
||||
- Progressive-deepening implementation details (worker scheduling, checkpoint format) — phased separately per §7.
|
||||
|
||||
## 11. File map (anticipated)
|
||||
|
||||
- `src/core/reflect.ts` (restored) — session reflect call + per-session cache.
|
||||
- `src/core/section-index.ts` (new) — page → sections split, lexical scorer, budget trim, provenance formatting; pure/unit-testable.
|
||||
- `src/core/hook.ts` — drop recall; inject cached reflect + matched sections; index refresh on roster cadence.
|
||||
- `src/core/session-start.ts` — cold-repo seed (unchanged) + reflect trigger wiring + initial index build.
|
||||
- `src/core/transcript.ts` (new or reworked) — JSON turn serialization + action-entry compaction for the Stop hook.
|
||||
- `src/core/missions.ts` / `src/core/hindsight.ts` — remove `entity_labels` and tag-scoped page triggers; keep seeded taxonomy + `capture_initiative`.
|
||||
- `src/core/config.ts` — remove any behavior flags; keep env/naming/harness + `pageRefreshEveryTurns`.
|
||||
- Tests alongside each.
|
||||
@@ -2,8 +2,8 @@ apiVersion: v2
|
||||
name: hindsight
|
||||
description: Hindsight helm chart
|
||||
type: application
|
||||
version: 0.8.4
|
||||
appVersion: "0.8.4"
|
||||
version: 0.9.0
|
||||
appVersion: "0.9.0"
|
||||
keywords:
|
||||
- ai
|
||||
- memory
|
||||
|
||||
@@ -60,13 +60,13 @@ spec:
|
||||
valueFrom:
|
||||
fieldRef:
|
||||
fieldPath: metadata.name
|
||||
{{- /* Inherit LLM config from api.env */}}
|
||||
{{- range $key, $value := .Values.api.env }}
|
||||
- name: {{ $key }}
|
||||
value: {{ $value | quote }}
|
||||
{{- end }}
|
||||
{{- /* Worker-specific env vars */}}
|
||||
{{- range $key, $value := .Values.worker.env }}
|
||||
{{- /* Explicitly set port to override K8s service discovery env var (HINDSIGHT_API_PORT) */}}
|
||||
- name: HINDSIGHT_API_PORT
|
||||
value: {{ .Values.worker.service.targetPort | quote }}
|
||||
{{- /* Inherit LLM config from api.env, then apply worker-specific env.
|
||||
Merge (worker.env wins) so a key set in both does not emit a
|
||||
duplicate env entry, which server-side apply rejects. */}}
|
||||
{{- range $key, $value := merge (deepCopy (.Values.worker.env | default dict)) (.Values.api.env | default dict) }}
|
||||
- name: {{ $key }}
|
||||
value: {{ $value | quote }}
|
||||
{{- end }}
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@vectorize-io/hindsight-all",
|
||||
"version": "0.8.4",
|
||||
"version": "0.9.0",
|
||||
"description": "Node.js programmatic lifecycle manager for Hindsight — embeds a local hindsight daemon in a Node application. Pair with @vectorize-io/hindsight-client for memory operations.",
|
||||
"main": "dist/index.js",
|
||||
"types": "dist/index.d.ts",
|
||||
|
||||
@@ -1,15 +1,16 @@
|
||||
[build-system]
|
||||
requires = ["setuptools>=61"]
|
||||
requires = ["setuptools>=77"]
|
||||
build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "hindsight-all-slim"
|
||||
version = "0.8.4"
|
||||
version = "0.9.0"
|
||||
description = "Hindsight: Agent Memory That Works Like Human Memory - Slim All-in-One Bundle"
|
||||
license = "MIT"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.11"
|
||||
dependencies = [
|
||||
"hindsight-api-slim==0.8.4",
|
||||
"hindsight-api-slim==0.9.0",
|
||||
"hindsight-client>=0.0.7",
|
||||
"hindsight-embed>=0.1.0",
|
||||
]
|
||||
|
||||
@@ -1,15 +1,16 @@
|
||||
[build-system]
|
||||
requires = ["hatchling"]
|
||||
requires = ["hatchling>=1.27"]
|
||||
build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "hindsight-all"
|
||||
version = "0.8.4"
|
||||
version = "0.9.0"
|
||||
description = "Hindsight: Agent Memory That Works Like Human Memory - All-in-One Bundle"
|
||||
license = "MIT"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.11"
|
||||
dependencies = [
|
||||
"hindsight-api-slim[all]==0.8.4",
|
||||
"hindsight-api-slim[all]==0.9.0",
|
||||
"hindsight-client>=0.0.7",
|
||||
"hindsight-embed>=0.1.0",
|
||||
]
|
||||
@@ -21,7 +22,7 @@ hindsight-embed = { workspace = true }
|
||||
|
||||
[project.optional-dependencies]
|
||||
local-llm = [
|
||||
"hindsight-api-slim[local-llm]==0.8.4",
|
||||
"hindsight-api-slim[local-llm]==0.9.0",
|
||||
]
|
||||
test = [
|
||||
"pytest>=7.0.0",
|
||||
|
||||
@@ -53,4 +53,4 @@ __all__ = [
|
||||
"RemoteTEICrossEncoder",
|
||||
"LLMConfig",
|
||||
]
|
||||
__version__ = "0.8.4"
|
||||
__version__ = "0.9.0"
|
||||
|
||||
@@ -8,7 +8,9 @@ import asyncio
|
||||
import io
|
||||
import json
|
||||
import logging
|
||||
import struct
|
||||
import zipfile
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
@@ -16,10 +18,12 @@ from typing import Any
|
||||
import asyncpg
|
||||
import typer
|
||||
|
||||
from ..config import DEFAULT_DATABASE_SCHEMA, HindsightConfig
|
||||
from ..config import DEFAULT_DATABASE_SCHEMA, HindsightConfig, load_dotenv_for_entrypoint
|
||||
from ..engine.memory_engine import _current_schema
|
||||
from ..engine.retain.bank_utils import _vector_index_clause
|
||||
from ..engine.schema import fq_table_explicit as _fq_table
|
||||
from ..engine.transfer import export_bank
|
||||
from ..engine.vector_index_health import SchemaVectorIndexResult, repair_vector_indexes
|
||||
from ..extensions import TenantExtension, load_extension
|
||||
from ..pg0 import parse_pg0_url, resolve_database_url
|
||||
|
||||
@@ -56,6 +60,7 @@ BACKUP_TABLES = [
|
||||
"observation_history",
|
||||
"mental_models",
|
||||
"mental_model_history",
|
||||
"knowledge_pages",
|
||||
"directives",
|
||||
"async_operations",
|
||||
"webhooks",
|
||||
@@ -65,7 +70,173 @@ BACKUP_TABLES = [
|
||||
"graph_maintenance_queue",
|
||||
]
|
||||
|
||||
MANIFEST_VERSION = "1"
|
||||
MANIFEST_VERSION = "2"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class BackupColumn:
|
||||
"""A PostgreSQL column shape required to decode a binary COPY stream."""
|
||||
|
||||
name: str
|
||||
type_name: str
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class TableRestorePlan:
|
||||
"""How one table's backed-up binary COPY stream is replayed onto the target.
|
||||
|
||||
``columns`` is the target column list handed to ``copy_to_table``, in stream
|
||||
order. When the target no longer has a backed-up column, its field is stripped
|
||||
from every tuple (``dropped_field_indices``) before the stream is replayed —
|
||||
binary COPY is positional, so the column list and the tuple fields must agree.
|
||||
"""
|
||||
|
||||
columns: list[str]
|
||||
dropped_field_indices: tuple[int, ...]
|
||||
source_field_count: int
|
||||
|
||||
|
||||
# Header of a PostgreSQL binary COPY stream: an 11-byte signature, an int32 flags
|
||||
# field, and an int32 header-extension length followed by that many bytes.
|
||||
_COPY_BINARY_SIGNATURE = b"PGCOPY\n\xff\r\n\x00"
|
||||
_COPY_BINARY_HEADER_LEN = len(_COPY_BINARY_SIGNATURE) + 8
|
||||
|
||||
|
||||
def _strip_binary_copy_fields(data: bytes, plan: TableRestorePlan) -> bytes:
|
||||
"""Drop `plan.dropped_field_indices` from every tuple of a binary COPY stream.
|
||||
|
||||
Restore used to reject a backup whose columns the target no longer had — the
|
||||
preflight raised "target is missing backup columns …", which made any backup
|
||||
taken before a column-dropping migration unrestorable afterwards. Those columns
|
||||
are now ignored instead, but they cannot simply be left out of the
|
||||
``copy_to_table`` column list: binary COPY carries no column identities, so each
|
||||
tuple's fields are matched to the column list purely by position and an unedited
|
||||
stream would desynchronise (or, worse, land values in the wrong columns). So the
|
||||
stream itself is rewritten here.
|
||||
|
||||
Tuple format: int16 field count, then per field an int32 length (-1 for NULL)
|
||||
followed by that many bytes. An int16 of -1 is the end-of-data trailer.
|
||||
"""
|
||||
if not plan.dropped_field_indices:
|
||||
return data
|
||||
if not data.startswith(_COPY_BINARY_SIGNATURE):
|
||||
raise ValueError("Backup stream is not in PostgreSQL binary COPY format")
|
||||
|
||||
(extension_len,) = struct.unpack_from("!i", data, len(_COPY_BINARY_SIGNATURE) + 4)
|
||||
pos = _COPY_BINARY_HEADER_LEN + extension_len
|
||||
out = bytearray(data[:pos])
|
||||
|
||||
dropped = set(plan.dropped_field_indices)
|
||||
kept_count = plan.source_field_count - len(dropped)
|
||||
while True:
|
||||
(field_count,) = struct.unpack_from("!h", data, pos)
|
||||
pos += 2
|
||||
if field_count == -1: # end-of-data trailer
|
||||
out += struct.pack("!h", -1)
|
||||
break
|
||||
if field_count != plan.source_field_count:
|
||||
raise ValueError(
|
||||
f"Backup stream tuple has {field_count} fields, manifest declares {plan.source_field_count}"
|
||||
)
|
||||
out += struct.pack("!h", kept_count)
|
||||
for index in range(field_count):
|
||||
(length,) = struct.unpack_from("!i", data, pos)
|
||||
pos += 4
|
||||
payload = b"" if length == -1 else data[pos : pos + length]
|
||||
pos += max(length, 0)
|
||||
if index in dropped:
|
||||
continue
|
||||
out += struct.pack("!i", length)
|
||||
out += payload
|
||||
return bytes(out)
|
||||
|
||||
|
||||
async def _table_columns(conn: asyncpg.Connection, schema: str, table: str) -> list[BackupColumn]:
|
||||
rows = await conn.fetch(
|
||||
"""
|
||||
SELECT a.attname AS name, pg_catalog.format_type(a.atttypid, a.atttypmod) AS type_name
|
||||
FROM pg_catalog.pg_attribute AS a
|
||||
JOIN pg_catalog.pg_class AS c ON c.oid = a.attrelid
|
||||
JOIN pg_catalog.pg_namespace AS n ON n.oid = c.relnamespace
|
||||
WHERE n.nspname = $1 AND c.relname = $2 AND a.attnum > 0 AND NOT a.attisdropped
|
||||
AND a.attgenerated = ''
|
||||
ORDER BY a.attnum
|
||||
""",
|
||||
schema,
|
||||
table,
|
||||
)
|
||||
return [BackupColumn(name=row["name"], type_name=row["type_name"]) for row in rows]
|
||||
|
||||
|
||||
async def _validate_restore_schema(
|
||||
conn: asyncpg.Connection, manifest: dict[str, Any], schema: str
|
||||
) -> dict[str, TableRestorePlan]:
|
||||
"""Validate every COPY stream against the target before destructive work starts.
|
||||
|
||||
A column the target no longer has is **not** an error: a migration that drops a
|
||||
column would otherwise make every backup taken before it permanently
|
||||
unrestorable. Such columns are skipped (their fields are stripped from the
|
||||
stream by ``_strip_binary_copy_fields``) and reported, so the operator sees what
|
||||
was discarded instead of the restore failing outright.
|
||||
|
||||
Type mismatches remain fatal. Type equality is an exact ``format_type`` string
|
||||
match. This is deliberately stricter than binary-COPY wire compatibility (e.g.
|
||||
``varchar`` and ``text`` share a binary format yet compare unequal here): we
|
||||
would rather fail a genuinely-restorable backup with a clear, actionable error
|
||||
than silently risk a subtle binary mismatch. Restores blocked this way can be
|
||||
recovered by aligning the target schema.
|
||||
"""
|
||||
plans: dict[str, TableRestorePlan] = {}
|
||||
errors: list[str] = []
|
||||
for table, table_manifest in manifest["tables"].items():
|
||||
source_columns = [BackupColumn(**column) for column in table_manifest["columns"]]
|
||||
target_by_name = {column.name: column for column in await _table_columns(conn, schema, table)}
|
||||
unknown = [
|
||||
(index, column.name) for index, column in enumerate(source_columns) if column.name not in target_by_name
|
||||
]
|
||||
mismatched = [
|
||||
f"{column.name} ({column.type_name} in backup, {target_by_name[column.name].type_name} in target)"
|
||||
for column in source_columns
|
||||
if column.name in target_by_name and target_by_name[column.name].type_name != column.type_name
|
||||
]
|
||||
if mismatched:
|
||||
errors.append(f"{table}: incompatible column types: {', '.join(mismatched)}")
|
||||
if unknown:
|
||||
typer.echo(
|
||||
f" {table}: ignoring {len(unknown)} backup column(s) absent from the target schema: "
|
||||
f"{', '.join(name for _, name in unknown)}"
|
||||
)
|
||||
plans[table] = TableRestorePlan(
|
||||
columns=[column.name for column in source_columns if column.name in target_by_name],
|
||||
dropped_field_indices=tuple(index for index, _ in unknown),
|
||||
source_field_count=len(source_columns),
|
||||
)
|
||||
|
||||
if errors:
|
||||
details = "; ".join(errors)
|
||||
raise ValueError(f"Backup schema is incompatible with target schema '{schema}': {details}")
|
||||
return plans
|
||||
|
||||
|
||||
def _effective_backup_tables() -> list[str]:
|
||||
"""Core backup tables plus any bank-scoped tables a loaded extension declares.
|
||||
|
||||
``BACKUP_TABLES`` covers only the tables core owns. An extension that
|
||||
provisions its own bank-scoped tables (via ``TenantExtension``) declares
|
||||
them through ``extra_bank_tables()`` so they aren't dropped on restore.
|
||||
Extension tables are appended *after* the core set so restore's forward
|
||||
COPY inserts them after their FK parents (e.g. ``banks``) and the reversed
|
||||
TRUNCATE clears them before those parents.
|
||||
"""
|
||||
tables = list(BACKUP_TABLES)
|
||||
tenant_extension = load_extension("TENANT", TenantExtension)
|
||||
if tenant_extension is not None:
|
||||
seen = set(tables)
|
||||
for spec in tenant_extension.extra_bank_tables():
|
||||
if spec.include_in_backup and spec.name not in seen:
|
||||
tables.append(spec.name)
|
||||
seen.add(spec.name)
|
||||
return tables
|
||||
|
||||
|
||||
async def _admin_connect(db_url: str) -> asyncpg.Connection:
|
||||
@@ -76,7 +247,8 @@ async def _admin_connect(db_url: str) -> asyncpg.Connection:
|
||||
is the only step needed to connect. JSON codecs are registered so ``jsonb``
|
||||
columns decode to Python objects (used by the export row dumps).
|
||||
"""
|
||||
is_pg0, instance_name, _ = parse_pg0_url(db_url)
|
||||
_pg0 = parse_pg0_url(db_url)
|
||||
is_pg0, instance_name = _pg0.is_pg0, _pg0.instance_name
|
||||
if is_pg0:
|
||||
typer.echo(f"Starting embedded PostgreSQL (instance: {instance_name})...")
|
||||
conn = await asyncpg.connect(await resolve_database_url(db_url))
|
||||
@@ -85,8 +257,18 @@ async def _admin_connect(db_url: str) -> asyncpg.Connection:
|
||||
return conn
|
||||
|
||||
|
||||
async def _backup(database_url: str, output_path: Path, schema: str = "public") -> dict[str, Any]:
|
||||
"""Backup all tables to a zip file using binary COPY protocol."""
|
||||
async def _backup(
|
||||
database_url: str,
|
||||
output_path: Path,
|
||||
schema: str = "public",
|
||||
backup_tables: list[str] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Backup all tables to a zip file using binary COPY protocol.
|
||||
|
||||
``backup_tables`` defaults to the core ``BACKUP_TABLES``; callers pass the
|
||||
extension-augmented list from ``_effective_backup_tables()``.
|
||||
"""
|
||||
backup_tables = backup_tables if backup_tables is not None else BACKUP_TABLES
|
||||
conn = await asyncpg.connect(database_url)
|
||||
try:
|
||||
tables: dict[str, Any] = {}
|
||||
@@ -103,14 +285,24 @@ async def _backup(database_url: str, output_path: Path, schema: str = "public")
|
||||
# entities table was backed up.
|
||||
async with conn.transaction(isolation="repeatable_read"):
|
||||
with zipfile.ZipFile(output_path, "w", zipfile.ZIP_DEFLATED) as zf:
|
||||
for i, table in enumerate(BACKUP_TABLES, 1):
|
||||
typer.echo(f" [{i}/{len(BACKUP_TABLES)}] Backing up {table}...", nl=False)
|
||||
for i, table in enumerate(backup_tables, 1):
|
||||
typer.echo(f" [{i}/{len(backup_tables)}] Backing up {table}...", nl=False)
|
||||
|
||||
buffer = io.BytesIO()
|
||||
|
||||
# Use binary COPY for exact type preservation
|
||||
columns = await _table_columns(conn, schema, table)
|
||||
|
||||
# Pin the ordered columns into both the stream and manifest.
|
||||
# PostgreSQL binary COPY does not encode column identities, so
|
||||
# restore must validate this shape before truncating any data.
|
||||
# asyncpg requires schema_name as separate parameter
|
||||
await conn.copy_from_table(table, schema_name=schema, output=buffer, format="binary")
|
||||
await conn.copy_from_table(
|
||||
table,
|
||||
schema_name=schema,
|
||||
columns=[column.name for column in columns],
|
||||
output=buffer,
|
||||
format="binary",
|
||||
)
|
||||
|
||||
data = buffer.getvalue()
|
||||
zf.writestr(f"{table}.bin", data)
|
||||
@@ -121,6 +313,7 @@ async def _backup(database_url: str, output_path: Path, schema: str = "public")
|
||||
tables[table] = {
|
||||
"rows": row_count,
|
||||
"size_bytes": len(data),
|
||||
"columns": [{"name": column.name, "type_name": column.type_name} for column in columns],
|
||||
}
|
||||
|
||||
typer.echo(f" {row_count} rows")
|
||||
@@ -132,8 +325,20 @@ async def _backup(database_url: str, output_path: Path, schema: str = "public")
|
||||
await conn.close()
|
||||
|
||||
|
||||
async def _restore(database_url: str, input_path: Path, schema: str = "public") -> dict[str, Any]:
|
||||
"""Restore all tables from a zip file using binary COPY protocol."""
|
||||
async def _restore(
|
||||
database_url: str,
|
||||
input_path: Path,
|
||||
schema: str = "public",
|
||||
backup_tables: list[str] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Restore all tables from a zip file using binary COPY protocol.
|
||||
|
||||
``backup_tables`` defaults to the core ``BACKUP_TABLES``; callers pass the
|
||||
extension-augmented list from ``_effective_backup_tables()``. Tables named
|
||||
here but absent from the archive are truncated then skipped for restore, so
|
||||
a stale extension registration never leaves pre-restore rows behind.
|
||||
"""
|
||||
backup_tables = backup_tables if backup_tables is not None else BACKUP_TABLES
|
||||
conn = await asyncpg.connect(database_url)
|
||||
try:
|
||||
with zipfile.ZipFile(input_path, "r") as zf:
|
||||
@@ -142,29 +347,42 @@ async def _restore(database_url: str, input_path: Path, schema: str = "public")
|
||||
if manifest.get("version") != MANIFEST_VERSION:
|
||||
raise ValueError(f"Unsupported backup version: {manifest.get('version')}")
|
||||
|
||||
# Complete the compatibility check before entering the transaction
|
||||
# that truncates tables. This turns historical schema drift into an
|
||||
# actionable error without risking the target's existing data.
|
||||
restore_plans = await _validate_restore_schema(conn, manifest, schema)
|
||||
|
||||
# Use a transaction for atomic restore - either all tables are
|
||||
# restored or none are, preventing partial/inconsistent state.
|
||||
async with conn.transaction():
|
||||
typer.echo(" Clearing existing data...")
|
||||
# Truncate tables in reverse order (respects FK constraints)
|
||||
for table in reversed(BACKUP_TABLES):
|
||||
for table in reversed(backup_tables):
|
||||
qualified_table = _fq_table(table, schema)
|
||||
await conn.execute(f"TRUNCATE TABLE {qualified_table} CASCADE")
|
||||
|
||||
# Restore tables in forward order
|
||||
for i, table in enumerate(BACKUP_TABLES, 1):
|
||||
for i, table in enumerate(backup_tables, 1):
|
||||
filename = f"{table}.bin"
|
||||
if filename not in zf.namelist():
|
||||
typer.echo(f" [{i}/{len(BACKUP_TABLES)}] {table}: skipped (not in backup)")
|
||||
typer.echo(f" [{i}/{len(backup_tables)}] {table}: skipped (not in backup)")
|
||||
continue
|
||||
|
||||
expected_rows = manifest["tables"].get(table, {}).get("rows", "?")
|
||||
typer.echo(f" [{i}/{len(BACKUP_TABLES)}] Restoring {table}... {expected_rows} rows")
|
||||
typer.echo(f" [{i}/{len(backup_tables)}] Restoring {table}... {expected_rows} rows")
|
||||
|
||||
data = zf.read(filename)
|
||||
buffer = io.BytesIO(data)
|
||||
plan = restore_plans[table]
|
||||
# Strips the fields of any column the target no longer has;
|
||||
# a no-op when the schemas still line up.
|
||||
buffer = io.BytesIO(_strip_binary_copy_fields(zf.read(filename), plan))
|
||||
# asyncpg requires schema_name as separate parameter
|
||||
await conn.copy_to_table(table, schema_name=schema, source=buffer, format="binary")
|
||||
await conn.copy_to_table(
|
||||
table,
|
||||
schema_name=schema,
|
||||
columns=plan.columns,
|
||||
source=buffer,
|
||||
format="binary",
|
||||
)
|
||||
|
||||
# Refresh materialized view
|
||||
typer.echo(" Refreshing materialized views...")
|
||||
@@ -177,20 +395,22 @@ async def _restore(database_url: str, input_path: Path, schema: str = "public")
|
||||
|
||||
async def _run_backup(db_url: str, output: Path, schema: str = "public") -> dict[str, Any]:
|
||||
"""Resolve database URL and run backup."""
|
||||
is_pg0, instance_name, _ = parse_pg0_url(db_url)
|
||||
_pg0 = parse_pg0_url(db_url)
|
||||
is_pg0, instance_name = _pg0.is_pg0, _pg0.instance_name
|
||||
if is_pg0:
|
||||
typer.echo(f"Starting embedded PostgreSQL (instance: {instance_name})...")
|
||||
resolved_url = await resolve_database_url(db_url)
|
||||
return await _backup(resolved_url, output, schema)
|
||||
return await _backup(resolved_url, output, schema, backup_tables=_effective_backup_tables())
|
||||
|
||||
|
||||
async def _run_restore(db_url: str, input_file: Path, schema: str = "public") -> dict[str, Any]:
|
||||
"""Resolve database URL and run restore."""
|
||||
is_pg0, instance_name, _ = parse_pg0_url(db_url)
|
||||
_pg0 = parse_pg0_url(db_url)
|
||||
is_pg0, instance_name = _pg0.is_pg0, _pg0.instance_name
|
||||
if is_pg0:
|
||||
typer.echo(f"Starting embedded PostgreSQL (instance: {instance_name})...")
|
||||
resolved_url = await resolve_database_url(db_url)
|
||||
return await _restore(resolved_url, input_file, schema)
|
||||
return await _restore(resolved_url, input_file, schema, backup_tables=_effective_backup_tables())
|
||||
|
||||
|
||||
@app.command()
|
||||
@@ -214,7 +434,7 @@ def backup(
|
||||
manifest = asyncio.run(_run_backup(config.database_url, output, schema))
|
||||
|
||||
total_rows = sum(t["rows"] for t in manifest["tables"].values())
|
||||
typer.echo(f"Backed up {total_rows} rows across {len(BACKUP_TABLES)} tables")
|
||||
typer.echo(f"Backed up {total_rows} rows across {len(manifest['tables'])} tables")
|
||||
typer.echo(f"Backup saved to {output}")
|
||||
|
||||
|
||||
@@ -247,7 +467,7 @@ def restore(
|
||||
manifest = asyncio.run(_run_restore(config.database_url, input_file, schema))
|
||||
|
||||
total_rows = sum(t["rows"] for t in manifest["tables"].values())
|
||||
typer.echo(f"Restored {total_rows} rows across {len(BACKUP_TABLES)} tables")
|
||||
typer.echo(f"Restored {total_rows} rows across {len(manifest['tables'])} tables")
|
||||
typer.echo("Restore complete")
|
||||
|
||||
|
||||
@@ -261,17 +481,17 @@ async def _run_migration(
|
||||
"""Resolve database URL and run migrations for one schema or all discovered schemas."""
|
||||
from ..migrations import run_migrations_for_schemas
|
||||
|
||||
is_pg0, instance_name, _ = parse_pg0_url(db_url)
|
||||
_pg0 = parse_pg0_url(db_url)
|
||||
is_pg0, instance_name = _pg0.is_pg0, _pg0.instance_name
|
||||
if is_pg0:
|
||||
typer.echo(f"Starting embedded PostgreSQL (instance: {instance_name})...")
|
||||
resolved_url = await resolve_database_url(db_url)
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
tenant_extension = load_extension("TENANT", TenantExtension)
|
||||
if schema:
|
||||
schemas = [schema]
|
||||
else:
|
||||
tenant_extension = load_extension("TENANT", TenantExtension)
|
||||
|
||||
schemas = [base_schema or DEFAULT_DATABASE_SCHEMA]
|
||||
if tenant_extension:
|
||||
tenants = await tenant_extension.list_tenants()
|
||||
@@ -296,9 +516,36 @@ async def _run_migration(
|
||||
ensure_extensions=ensure_extensions,
|
||||
)
|
||||
|
||||
# After core migrations, provision any extension-owned bank-scoped tables
|
||||
# per schema so extension schema evolves on the same lifecycle as core
|
||||
# schema (rather than via a lazy first-request path).
|
||||
if tenant_extension is not None:
|
||||
await _provision_extra_bank_tables(resolved_url, schemas, tenant_extension)
|
||||
|
||||
return schemas
|
||||
|
||||
|
||||
async def _provision_extra_bank_tables(
|
||||
resolved_url: str, schemas: list[str], tenant_extension: TenantExtension
|
||||
) -> None:
|
||||
"""Run the tenant extension's table provisioner for each migrated schema.
|
||||
|
||||
Fires after core migrations complete so extension-owned bank tables are
|
||||
created/evolved on the same lifecycle as core schema. A failure aborts the
|
||||
migration command (and names the offending schema) rather than being
|
||||
swallowed — provisioning is idempotent, so the operator can fix and re-run.
|
||||
"""
|
||||
for schema in schemas:
|
||||
conn = await asyncpg.connect(resolved_url)
|
||||
try:
|
||||
await tenant_extension.provision_bank_tables(conn, schema)
|
||||
except Exception as e:
|
||||
typer.echo(f" Failed to provision extension tables for schema '{schema}': {e}", err=True)
|
||||
raise
|
||||
finally:
|
||||
await conn.close()
|
||||
|
||||
|
||||
@app.command(name="run-db-migration")
|
||||
def run_db_migration(
|
||||
schema: str | None = typer.Option(
|
||||
@@ -353,6 +600,134 @@ def run_db_migration(
|
||||
typer.echo(f"Database migrations completed successfully for {len(schemas)} schema(s)")
|
||||
|
||||
|
||||
async def _resolve_schemas(base_schema: str | None) -> list[str]:
|
||||
"""Base schema plus every discovered tenant schema, de-duplicated in order."""
|
||||
schemas = [base_schema or DEFAULT_DATABASE_SCHEMA]
|
||||
tenant_extension = load_extension("TENANT", TenantExtension)
|
||||
if tenant_extension:
|
||||
tenants = await tenant_extension.list_tenants()
|
||||
schemas.extend(tenant.schema for tenant in tenants if tenant.schema)
|
||||
return list(dict.fromkeys(schemas))
|
||||
|
||||
|
||||
async def _run_repair_bank(
|
||||
db_url: str,
|
||||
*,
|
||||
base_schema: str,
|
||||
schema: str | None,
|
||||
bank_id: str | None,
|
||||
dry_run: bool,
|
||||
) -> list[SchemaVectorIndexResult]:
|
||||
"""Reconcile per-(bank, fact_type) vector index coverage over a raw connection.
|
||||
|
||||
A single autocommit connection is used because ``CREATE INDEX CONCURRENTLY``
|
||||
(used by ``repair_vector_indexes``) cannot run inside a transaction block.
|
||||
"""
|
||||
schemas = [schema] if schema else await _resolve_schemas(base_schema)
|
||||
index_clause = _vector_index_clause()
|
||||
# Guarded by the command, but assert so this helper is never called for a
|
||||
# backend without per-bank indexes.
|
||||
assert index_clause is not None
|
||||
|
||||
conn = await _admin_connect(db_url)
|
||||
try:
|
||||
results = await repair_vector_indexes(conn, schemas, index_clause, dry_run=dry_run, bank_id=bank_id)
|
||||
for result in results:
|
||||
typer.echo(
|
||||
f" schema '{result.schema}': {result.banks_scanned} bank(s) scanned, "
|
||||
f"{result.already_present} present, {result.created} created, "
|
||||
f"{result.skipped} to-create (dry-run), {result.failed} failed"
|
||||
)
|
||||
return results
|
||||
finally:
|
||||
await conn.close()
|
||||
|
||||
|
||||
@app.command(name="repair-bank")
|
||||
def repair_bank(
|
||||
bank_id: str | None = typer.Option(
|
||||
None,
|
||||
"--bank",
|
||||
"-b",
|
||||
help="Bank id to repair. Mutually exclusive with --all.",
|
||||
),
|
||||
all_banks: bool = typer.Option(
|
||||
False,
|
||||
"--all",
|
||||
help="Repair every bank in the base schema and all discovered tenant schemas.",
|
||||
),
|
||||
schema: str | None = typer.Option(
|
||||
None,
|
||||
"--schema",
|
||||
"-s",
|
||||
help="Limit to a single schema. Defaults to the base schema plus discovered tenant schemas.",
|
||||
),
|
||||
dry_run: bool = typer.Option(
|
||||
False,
|
||||
"--dry-run",
|
||||
help="Report what would be repaired without creating or dropping any index.",
|
||||
),
|
||||
):
|
||||
"""Verify and repair a bank's per-(bank, fact_type) vector index coverage.
|
||||
|
||||
Per-bank partial vector indexes are created when a bank is first created
|
||||
(instant on an empty bank). Banks that arrive populated — via logical
|
||||
restore, a cross-version upgrade, or a vector-extension switch — never hit
|
||||
that path, so their recall silently falls back to a global index +
|
||||
post-filter (slower, under-returning). This command detects missing OR
|
||||
invalid coverage (an INVALID leftover or an index whose access method
|
||||
drifted counts as missing) and rebuilds it with CREATE INDEX CONCURRENTLY,
|
||||
so it never blocks the live fleet. Idempotent and safe to re-run — the
|
||||
escape hatch after a restore, upgrade, or backend switch.
|
||||
"""
|
||||
if bool(bank_id) == all_banks:
|
||||
typer.echo("Error: pass exactly one of --bank <id> or --all.", err=True)
|
||||
raise typer.Exit(2)
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
if not config.database_url:
|
||||
typer.echo("Error: Database URL not configured.", err=True)
|
||||
typer.echo("Set HINDSIGHT_API_DATABASE_URL environment variable.", err=True)
|
||||
raise typer.Exit(1)
|
||||
|
||||
# Backend guard: backends with a single global vector index (AlloyDB ScaNN,
|
||||
# Oracle) have no per-bank indexes to repair.
|
||||
if _vector_index_clause() is None:
|
||||
typer.echo("Configured vector backend does not use per-bank vector indexes — nothing to repair.")
|
||||
return
|
||||
|
||||
target = f"bank '{bank_id}'" if bank_id else "all banks"
|
||||
scope = f"schema '{schema}'" if schema else "base schema and all discovered tenant schemas"
|
||||
typer.echo(f"Repairing per-bank vector indexes for {target} across {scope}...")
|
||||
if dry_run:
|
||||
typer.echo("Dry run: no indexes will be created or dropped.")
|
||||
|
||||
results = asyncio.run(
|
||||
_run_repair_bank(
|
||||
config.database_url,
|
||||
base_schema=config.database_schema,
|
||||
schema=schema,
|
||||
bank_id=bank_id,
|
||||
dry_run=dry_run,
|
||||
)
|
||||
)
|
||||
|
||||
total_banks = sum(r.banks_scanned for r in results)
|
||||
total_present = sum(r.already_present for r in results)
|
||||
total_created = sum(r.created for r in results)
|
||||
total_skipped = sum(r.skipped for r in results)
|
||||
total_failed = sum(r.failed for r in results)
|
||||
typer.echo(
|
||||
f"Done: {len(results)} schema(s), {total_banks} bank(s) scanned, "
|
||||
f"{total_present} already present, {total_created} created, "
|
||||
f"{total_skipped} to-create (dry-run), {total_failed} failed"
|
||||
)
|
||||
if total_failed:
|
||||
failed_names = [name for r in results for name in r.failed_indexes]
|
||||
typer.echo(f"Failed indexes (dropped, retry with a re-run): {', '.join(failed_names)}", err=True)
|
||||
raise typer.Exit(1)
|
||||
|
||||
|
||||
async def _run_export_bank(db_url: str, bank_id: str, output: Path, schema: str, include_history: bool) -> int:
|
||||
"""Export a whole bank to a ZIP archive."""
|
||||
conn = await _admin_connect(db_url)
|
||||
@@ -360,7 +735,14 @@ async def _run_export_bank(db_url: str, bank_id: str, output: Path, schema: str,
|
||||
# export_bank resolves table names via fq_table (the _current_schema
|
||||
# contextvar); set it so the raw connection targets the right schema.
|
||||
_current_schema.set(schema)
|
||||
data = await export_bank(conn, bank_id, include_history=include_history)
|
||||
# _admin_connect registers JSON codecs, so row dumps already contain
|
||||
# decoded Python values (including JSON scalar strings).
|
||||
data = await export_bank(
|
||||
conn,
|
||||
bank_id,
|
||||
include_history=include_history,
|
||||
bank_rows_json_encoding="decoded",
|
||||
)
|
||||
finally:
|
||||
await conn.close()
|
||||
|
||||
@@ -472,7 +854,8 @@ def import_bank_command(
|
||||
|
||||
async def _decommission_worker(db_url: str, worker_id: str, schema: str = "public") -> int:
|
||||
"""Release all tasks owned by a worker, setting them back to pending status."""
|
||||
is_pg0, instance_name, _ = parse_pg0_url(db_url)
|
||||
_pg0 = parse_pg0_url(db_url)
|
||||
is_pg0, instance_name = _pg0.is_pg0, _pg0.instance_name
|
||||
if is_pg0:
|
||||
typer.echo(f"Starting embedded PostgreSQL (instance: {instance_name})...")
|
||||
resolved_url = await resolve_database_url(db_url)
|
||||
@@ -531,7 +914,8 @@ def decommission_worker(
|
||||
|
||||
async def _decommission_all_workers(db_url: str, schema: str = "public") -> list[dict[str, Any]]:
|
||||
"""Release all processing tasks from all workers, setting them back to pending status."""
|
||||
is_pg0, instance_name, _ = parse_pg0_url(db_url)
|
||||
_pg0 = parse_pg0_url(db_url)
|
||||
is_pg0, instance_name = _pg0.is_pg0, _pg0.instance_name
|
||||
if is_pg0:
|
||||
typer.echo(f"Starting embedded PostgreSQL (instance: {instance_name})...")
|
||||
resolved_url = await resolve_database_url(db_url)
|
||||
@@ -596,7 +980,8 @@ def decommission_workers(
|
||||
|
||||
async def _worker_status(db_url: str, schema: str = "public") -> list[dict[str, Any]]:
|
||||
"""Get all processing tasks grouped by worker with their last update time."""
|
||||
is_pg0, instance_name, _ = parse_pg0_url(db_url)
|
||||
_pg0 = parse_pg0_url(db_url)
|
||||
is_pg0, instance_name = _pg0.is_pg0, _pg0.instance_name
|
||||
if is_pg0:
|
||||
typer.echo(f"Starting embedded PostgreSQL (instance: {instance_name})...")
|
||||
resolved_url = await resolve_database_url(db_url)
|
||||
@@ -664,6 +1049,7 @@ def worker_status(
|
||||
|
||||
|
||||
def main():
|
||||
load_dotenv_for_entrypoint()
|
||||
app()
|
||||
|
||||
|
||||
|
||||
@@ -96,7 +96,8 @@ def get_database_url() -> str:
|
||||
# for the sync engine used during migrations.
|
||||
database_url = to_libpq_url(database_url)
|
||||
|
||||
config.set_main_option("sqlalchemy.url", database_url)
|
||||
# Alembic stores options through ConfigParser, where '%' is interpolation.
|
||||
config.set_main_option("sqlalchemy.url", database_url.replace("%", "%%"))
|
||||
return database_url
|
||||
|
||||
|
||||
|
||||
+67
@@ -0,0 +1,67 @@
|
||||
"""Drop observation_history's FK to memory_units.
|
||||
|
||||
The history table records one snapshot per observation change, keyed by
|
||||
``(bank_id, observation_id)``. Its foreign key to ``memory_units`` existed only to
|
||||
cascade-delete history when the observation row went away.
|
||||
|
||||
That assumes every observation *is* a ``memory_units`` row, which is true only
|
||||
while Postgres is the memories store. When another store owns the memories the
|
||||
observation lives there and Postgres holds no row for it, so every history insert
|
||||
raises a foreign-key violation — swallowed by the writer as "a race with parallel
|
||||
consolidation" and logged at warning level. The audit trail goes silently empty.
|
||||
|
||||
Dropping the constraint lets history be recorded wherever the observation is
|
||||
stored. The cleanup the cascade used to do is now explicit, in the paths that
|
||||
delete observations (``_execute_delete_action``, ``clear_observations``,
|
||||
``delete_bank``). Rows orphaned by a path that misses — a document delete
|
||||
cascading through ``memory_units``, for instance — are invisible to readers,
|
||||
which always filter by ``(bank_id, observation_id)``, and are reclaimed when the
|
||||
bank is deleted.
|
||||
|
||||
Oracle builds this schema through its own DDL runner and never had the
|
||||
constraint, so the Oracle slot is a deliberate no-op.
|
||||
|
||||
Revision ID: a1c9e7f3b2d8
|
||||
Revises: c7d1e9a4b3f2
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "a1c9e7f3b2d8"
|
||||
down_revision: str | Sequence[str] | None = "c7d1e9a4b3f2"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
_CONSTRAINT = "observation_history_observation_id_fkey"
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
op.execute(f"ALTER TABLE observation_history DROP CONSTRAINT IF EXISTS {_CONSTRAINT}")
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
# Re-adding the FK requires every row to reference a live memory_unit, so
|
||||
# clear any history whose observation is not a Postgres row first — those are
|
||||
# exactly the rows this migration made possible.
|
||||
op.execute(
|
||||
"DELETE FROM observation_history h "
|
||||
"WHERE NOT EXISTS (SELECT 1 FROM memory_units m WHERE m.id = h.observation_id)"
|
||||
)
|
||||
op.execute(
|
||||
f"ALTER TABLE observation_history ADD CONSTRAINT {_CONSTRAINT} "
|
||||
"FOREIGN KEY (observation_id) REFERENCES memory_units(id) ON DELETE CASCADE"
|
||||
)
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
# Oracle never had the constraint (its schema is built by a separate DDL
|
||||
# runner), so only Postgres has anything to drop.
|
||||
run_for_dialect(pg=_pg_upgrade, oracle=None)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade, oracle=None)
|
||||
+82
@@ -0,0 +1,82 @@
|
||||
"""Add indexes for terminal cleanup and newest-first operation listing.
|
||||
|
||||
Revision ID: a8c1e4f7b0d3
|
||||
Revises: e7c3a9f1b2d5
|
||||
Create Date: 2026-07-14
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "a8c1e4f7b0d3"
|
||||
down_revision: str | Sequence[str] | None = "e7c3a9f1b2d5"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _pg_schema_prefix() -> str:
|
||||
"""Schema-qualifier for PostgreSQL multi-tenant migration runs."""
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
schema = _pg_schema_prefix()
|
||||
# These can be large tables in long-running installations. Concurrent DDL
|
||||
# keeps operation submission, polling, and status reads available.
|
||||
with op.get_context().autocommit_block():
|
||||
op.execute(
|
||||
"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_async_operations_terminal_cleanup "
|
||||
f"ON {schema}async_operations (updated_at, operation_id) "
|
||||
"WHERE status IN ('completed', 'failed', 'cancelled')"
|
||||
)
|
||||
op.execute(
|
||||
"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_async_operations_bank_created_desc "
|
||||
f"ON {schema}async_operations (bank_id, created_at DESC)"
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema = _pg_schema_prefix()
|
||||
with op.get_context().autocommit_block():
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_async_operations_bank_created_desc")
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_async_operations_terminal_cleanup")
|
||||
|
||||
|
||||
def _oracle_create_index(sql: str) -> None:
|
||||
"""Create an index idempotently for rerun-safe Oracle migrations."""
|
||||
block = (
|
||||
"BEGIN "
|
||||
"EXECUTE IMMEDIATE :stmt; "
|
||||
"EXCEPTION WHEN OTHERS THEN "
|
||||
"IF SQLCODE = -955 THEN NULL; ELSE RAISE; END IF; "
|
||||
"END;"
|
||||
)
|
||||
op.get_bind().exec_driver_sql(block, {"stmt": sql})
|
||||
|
||||
|
||||
def _oracle_upgrade() -> None:
|
||||
# Oracle migrations run with CURRENT_SCHEMA set to each tenant, so table
|
||||
# and index names intentionally remain unqualified here.
|
||||
_oracle_create_index(
|
||||
"CREATE INDEX idx_async_operations_terminal_cleanup ON async_operations (updated_at, operation_id, status)"
|
||||
)
|
||||
_oracle_create_index(
|
||||
"CREATE INDEX idx_async_operations_bank_created_desc ON async_operations (bank_id, created_at DESC)"
|
||||
)
|
||||
|
||||
|
||||
def _oracle_downgrade() -> None:
|
||||
op.execute("DROP INDEX idx_async_operations_bank_created_desc")
|
||||
op.execute("DROP INDEX idx_async_operations_terminal_cleanup")
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade, oracle=_oracle_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade, oracle=_oracle_downgrade)
|
||||
+126
@@ -0,0 +1,126 @@
|
||||
"""Add knowledge_pages table (knowledge-base hierarchy).
|
||||
|
||||
The knowledge base organizes synthesized mental models into a navigable tree of
|
||||
**folders** and **pages**. A page references the mental model that holds its
|
||||
content (``mental_model_id``); a folder is a pure container (``mental_model_id``
|
||||
NULL). Hierarchy is a single self-referential ``parent_id`` so folders can nest
|
||||
arbitrarily. Content stays in ``mental_models`` — this table is metadata + tree
|
||||
structure only.
|
||||
|
||||
``managed`` lets a client tag a node as system-owned vs. hand-authored; it
|
||||
carries no server-side behaviour. A partial unique index keeps page names unique
|
||||
within a folder (case-insensitive; root pages compared under an empty parent).
|
||||
|
||||
Revision ID: a9b8c7d6e5f4
|
||||
Revises: a1c9e7f3b2d8
|
||||
Create Date: 2026-06-25
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "a9b8c7d6e5f4"
|
||||
down_revision: str | Sequence[str] | None = "a1c9e7f3b2d8"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _pg_schema_prefix() -> str:
|
||||
"""Schema-qualifier for raw SQL on PG (multi-tenant search_path)."""
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
schema = _pg_schema_prefix()
|
||||
# parent_id self-FK cascades so deleting a folder row removes its whole
|
||||
# subtree of rows in one shot. The mental_model FK is composite (matches the
|
||||
# mental_models (id, bank_id) PK) and cascades too, so deleting a page's
|
||||
# mental model removes the page row — folders skip the FK because a NULL
|
||||
# column in a composite FK is not enforced (MATCH SIMPLE).
|
||||
op.execute(
|
||||
f"""
|
||||
CREATE TABLE IF NOT EXISTS {schema}knowledge_pages (
|
||||
id VARCHAR(64) NOT NULL,
|
||||
bank_id TEXT NOT NULL,
|
||||
parent_id VARCHAR(64),
|
||||
kind VARCHAR(16) NOT NULL,
|
||||
name TEXT NOT NULL,
|
||||
mental_model_id VARCHAR(64),
|
||||
sort_order INTEGER NOT NULL DEFAULT 0,
|
||||
managed BOOLEAN NOT NULL DEFAULT false,
|
||||
created_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT now(),
|
||||
updated_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT now(),
|
||||
CONSTRAINT pk_knowledge_pages PRIMARY KEY (id),
|
||||
CONSTRAINT ck_knowledge_pages_kind CHECK (kind IN ('folder', 'page')),
|
||||
CONSTRAINT fk_kp_bank FOREIGN KEY (bank_id)
|
||||
REFERENCES {schema}banks(bank_id) ON DELETE CASCADE,
|
||||
CONSTRAINT fk_kp_parent FOREIGN KEY (parent_id)
|
||||
REFERENCES {schema}knowledge_pages(id) ON DELETE CASCADE,
|
||||
CONSTRAINT fk_kp_mm FOREIGN KEY (mental_model_id, bank_id)
|
||||
REFERENCES {schema}mental_models(id, bank_id) ON DELETE CASCADE
|
||||
)
|
||||
"""
|
||||
)
|
||||
op.execute(
|
||||
f"CREATE INDEX IF NOT EXISTS idx_kp_bank_parent ON {schema}knowledge_pages (bank_id, parent_id, sort_order)"
|
||||
)
|
||||
# COALESCE(parent_id, '') so root-level pages (NULL parent) are also unique by
|
||||
# name — NULLs would otherwise compare distinct and allow duplicates.
|
||||
op.execute(
|
||||
"CREATE UNIQUE INDEX IF NOT EXISTS uq_kp_folder_pagename "
|
||||
f"ON {schema}knowledge_pages (bank_id, COALESCE(parent_id, ''), lower(name)) "
|
||||
"WHERE kind = 'page'"
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema = _pg_schema_prefix()
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}uq_kp_folder_pagename")
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_kp_bank_parent")
|
||||
op.execute(f"DROP TABLE IF EXISTS {schema}knowledge_pages")
|
||||
|
||||
|
||||
def _oracle_upgrade() -> None:
|
||||
# No case-insensitive unique index on Oracle: `name` is a CLOB and cannot be
|
||||
# indexed with lower(); page-name uniqueness is enforced on PG only.
|
||||
op.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS knowledge_pages (
|
||||
id VARCHAR2(64) NOT NULL,
|
||||
bank_id VARCHAR2(256) NOT NULL,
|
||||
parent_id VARCHAR2(64),
|
||||
kind VARCHAR2(16) NOT NULL,
|
||||
name CLOB NOT NULL,
|
||||
mental_model_id VARCHAR2(64),
|
||||
sort_order NUMBER DEFAULT 0 NOT NULL,
|
||||
managed NUMBER(1) DEFAULT 0 NOT NULL,
|
||||
created_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
|
||||
updated_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
|
||||
CONSTRAINT pk_knowledge_pages PRIMARY KEY (id),
|
||||
CONSTRAINT ck_knowledge_pages_kind CHECK (kind IN ('folder', 'page')),
|
||||
CONSTRAINT fk_kp_bank FOREIGN KEY (bank_id)
|
||||
REFERENCES banks(bank_id) ON DELETE CASCADE,
|
||||
CONSTRAINT fk_kp_parent FOREIGN KEY (parent_id)
|
||||
REFERENCES knowledge_pages(id) ON DELETE CASCADE,
|
||||
CONSTRAINT fk_kp_mm FOREIGN KEY (mental_model_id, bank_id)
|
||||
REFERENCES mental_models(id, bank_id) ON DELETE CASCADE
|
||||
)
|
||||
"""
|
||||
)
|
||||
op.execute("CREATE INDEX idx_kp_bank_parent ON knowledge_pages (bank_id, parent_id, sort_order)")
|
||||
|
||||
|
||||
def _oracle_downgrade() -> None:
|
||||
op.execute("DROP TABLE knowledge_pages CASCADE CONSTRAINTS")
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade, oracle=_oracle_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade, oracle=_oracle_downgrade)
|
||||
+213
@@ -0,0 +1,213 @@
|
||||
"""Add entities.entity_kind and exclude label entities from the trigram index.
|
||||
|
||||
Label entities (values of ``entity_labels`` config groups, stored as
|
||||
``key:value`` canonical names) resolve by exact match only — fuzzy resolution
|
||||
must never merge distinct label values (#1558), and since #3187 they are looked
|
||||
up via the exact-match unique index rather than probed through pg_trgm. Their
|
||||
rows were still covered by the shared trigram index, so every fuzzy probe for a
|
||||
*regular* entity name pulled them into its candidate set only to discard them
|
||||
in the bitmap recheck. On banks where a free-text label group accumulated tens
|
||||
of thousands of mutually-similar values this recheck-discard overhead dominated
|
||||
database CPU under ingest bursts (#3208).
|
||||
|
||||
"Is this row a label" was previously derived at runtime from the bank's
|
||||
``entity_labels`` config, which an index predicate cannot reference — so the
|
||||
classification is now materialised on the row:
|
||||
|
||||
1. Add ``entity_kind`` ("regular"/"label", CHECK-constrained) on both dialects.
|
||||
A kind column rather than a boolean so future entity kinds don't need
|
||||
another column.
|
||||
2. Backfill per bank by classifying ``canonical_name`` against the bank's
|
||||
``entity_labels`` config with the same ``is_label_entity()`` the resolver
|
||||
uses at insert time — a SQL reimplementation would be a second source of
|
||||
truth (and the map-group recursion doesn't translate). Banks hold at most
|
||||
tens of thousands of entities, so the synchronous per-bank backfill is fine.
|
||||
Label configs supplied only by a tenant extension (not stored in
|
||||
``banks.config``) can't be seen here; their rows stay "regular", which
|
||||
costs index size but never correctness — label *texts* still resolve via
|
||||
the exact-match unique index.
|
||||
3. Rebuild the PG trigram index as a partial index excluding label rows.
|
||||
Built CONCURRENTLY (autocommit block, invalid-leftover sweep, IF NOT
|
||||
EXISTS — same shape as 2071c7518f88) and only then drop the old full
|
||||
index, so fuzzy probes never lose index coverage. Skipped entirely when
|
||||
pg_trgm is absent (the resolver falls back to the "full" strategy, #626).
|
||||
|
||||
Oracle has no trigram index — it fuzzy-matches with a UTL_MATCH scan — so it
|
||||
only gets the column + backfill; the resolver adds the matching
|
||||
``entity_kind != 'label'`` filter to that scan.
|
||||
|
||||
Revision ID: b3e8d1c6f4a9
|
||||
Revises: f2a6d8c4b1e9
|
||||
Create Date: 2026-08-06
|
||||
"""
|
||||
|
||||
import json
|
||||
from collections.abc import Sequence
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "b3e8d1c6f4a9"
|
||||
down_revision: str | Sequence[str] | None = "f2a6d8c4b1e9"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
_OLD_INDEX = "entities_canonical_name_lower_trgm_idx"
|
||||
_NEW_INDEX = "entities_canonical_name_lower_trgm_nonlabel_idx"
|
||||
|
||||
|
||||
def _pg_schema_prefix() -> str:
|
||||
"""Schema-qualifier for raw SQL on PG (multi-tenant search_path)."""
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def _backfill_entity_kind(schema: str) -> None:
|
||||
"""Set entity_kind='label' on rows matching their bank's entity_labels config.
|
||||
|
||||
Runs the resolver's own classification (``is_label_entity``) per bank in
|
||||
Python rather than reimplementing the enum/text/map prefix rules in SQL.
|
||||
Shared by both dialects: plain SELECT/UPDATE with expanding IN binds.
|
||||
"""
|
||||
from hindsight_api.engine.retain.entity_labels import (
|
||||
build_labels_lookup,
|
||||
is_label_entity,
|
||||
parse_entity_labels,
|
||||
)
|
||||
|
||||
bind = op.get_bind()
|
||||
banks = bind.execute(sa.text(f"SELECT bank_id, config FROM {schema}banks")).fetchall()
|
||||
for bank_id, raw_config in banks:
|
||||
# PG JSONB arrives as a dict; Oracle CLOB arrives as a LOB object on
|
||||
# raw text() fetches (oracledb's fetch_lobs default) — read it into a
|
||||
# JSON string first.
|
||||
if raw_config is not None and not isinstance(raw_config, (str, dict)):
|
||||
raw_config = raw_config.read()
|
||||
config = json.loads(raw_config) if isinstance(raw_config, str) else (raw_config or {})
|
||||
labels_cfg = parse_entity_labels(config.get("entity_labels"))
|
||||
if labels_cfg is None:
|
||||
continue
|
||||
lookup = build_labels_lookup(labels_cfg)
|
||||
rows = bind.execute(
|
||||
sa.text(f"SELECT id, canonical_name FROM {schema}entities WHERE bank_id = :bank_id"),
|
||||
{"bank_id": bank_id},
|
||||
).fetchall()
|
||||
label_ids = [entity_id for entity_id, name in rows if is_label_entity(name, labels_cfg, lookup)]
|
||||
# Chunked to stay under Oracle's 1000-element IN limit; also keeps PG
|
||||
# bind arrays bounded.
|
||||
for start in range(0, len(label_ids), 500):
|
||||
chunk = label_ids[start : start + 500]
|
||||
stmt = sa.text(f"UPDATE {schema}entities SET entity_kind = 'label' WHERE id IN :ids").bindparams(
|
||||
sa.bindparam("ids", expanding=True)
|
||||
)
|
||||
bind.execute(stmt, {"ids": chunk})
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
bind = op.get_bind()
|
||||
schema = _pg_schema_prefix()
|
||||
# `or None` collapses an unset option and an explicit empty string into NULL
|
||||
# so the COALESCE below falls back to current_schema() in both cases.
|
||||
target_schema = context.config.get_main_option("target_schema") or None
|
||||
|
||||
# IF NOT EXISTS: the transactional part below commits when the autocommit
|
||||
# block is entered, so a failure during the CONCURRENTLY build leaves the
|
||||
# revision unstamped with the column already added — the retry must not
|
||||
# trip over it. The constant default is a metadata-only change on PG 11+.
|
||||
op.execute(
|
||||
f"ALTER TABLE {schema}entities ADD COLUMN IF NOT EXISTS entity_kind TEXT DEFAULT 'regular' NOT NULL "
|
||||
f"CONSTRAINT chk_entities_entity_kind CHECK (entity_kind IN ('regular', 'label'))"
|
||||
)
|
||||
_backfill_entity_kind(schema)
|
||||
|
||||
# Without pg_trgm neither the old index nor the extension's opclass exists;
|
||||
# the resolver already runs the "full" strategy there (#626).
|
||||
has_trgm = bind.execute(sa.text("SELECT EXISTS(SELECT 1 FROM pg_extension WHERE extname = 'pg_trgm')")).scalar()
|
||||
if not has_trgm:
|
||||
return
|
||||
|
||||
# CREATE INDEX CONCURRENTLY cannot run inside a transaction block; the
|
||||
# autocommit_block runs each statement outside Alembic's migration
|
||||
# transaction. Build the partial index first and drop the old full index
|
||||
# only afterwards, so fuzzy probes never lose index coverage.
|
||||
with op.get_context().autocommit_block():
|
||||
# A CONCURRENTLY build that errored on a previous run leaves an INVALID
|
||||
# index of this name behind, which IF NOT EXISTS would skip forever.
|
||||
leftover_invalid = bind.execute(
|
||||
sa.text(
|
||||
"SELECT NOT i.indisvalid "
|
||||
"FROM pg_class c "
|
||||
"JOIN pg_index i ON c.oid = i.indexrelid "
|
||||
"JOIN pg_namespace n ON c.relnamespace = n.oid "
|
||||
"WHERE c.relname = :index_name "
|
||||
" AND n.nspname = COALESCE(:target_schema, current_schema())"
|
||||
),
|
||||
{"index_name": _NEW_INDEX, "target_schema": target_schema},
|
||||
).scalar()
|
||||
if leftover_invalid:
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}{_NEW_INDEX}")
|
||||
|
||||
# The predicate must textually match the resolver's candidate query
|
||||
# (`entity_kind != 'label'`) for the planner to choose this index.
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS {_NEW_INDEX} "
|
||||
f"ON {schema}entities USING GIN (LOWER(canonical_name) gin_trgm_ops) "
|
||||
f"WHERE entity_kind != 'label'"
|
||||
)
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}{_OLD_INDEX}")
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
bind = op.get_bind()
|
||||
schema = _pg_schema_prefix()
|
||||
|
||||
has_trgm = bind.execute(sa.text("SELECT EXISTS(SELECT 1 FROM pg_extension WHERE extname = 'pg_trgm')")).scalar()
|
||||
if has_trgm:
|
||||
# Restore the full index before dropping the partial one so fuzzy
|
||||
# probes keep index coverage throughout.
|
||||
with op.get_context().autocommit_block():
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS {_OLD_INDEX} "
|
||||
f"ON {schema}entities USING GIN (LOWER(canonical_name) gin_trgm_ops)"
|
||||
)
|
||||
# Dropping the column also drops the partial index and CHECK constraint.
|
||||
op.execute(f"ALTER TABLE {schema}entities DROP COLUMN IF EXISTS entity_kind")
|
||||
|
||||
|
||||
def _oracle_upgrade() -> None:
|
||||
# Swallow ORA-01430 (column already exists) so a retry after a mid-run
|
||||
# failure is idempotent — Oracle DDL auto-commits statement by statement.
|
||||
op.execute(
|
||||
"""
|
||||
BEGIN
|
||||
EXECUTE IMMEDIATE 'ALTER TABLE entities ADD (entity_kind VARCHAR2(16) DEFAULT ''regular'' NOT NULL
|
||||
CONSTRAINT chk_entities_entity_kind CHECK (entity_kind IN (''regular'', ''label'')))';
|
||||
EXCEPTION WHEN OTHERS THEN
|
||||
IF SQLCODE != -1430 THEN RAISE; END IF;
|
||||
END;
|
||||
"""
|
||||
)
|
||||
_backfill_entity_kind("")
|
||||
|
||||
|
||||
def _oracle_downgrade() -> None:
|
||||
# Swallow ORA-00904 (column does not exist).
|
||||
op.execute(
|
||||
"""
|
||||
BEGIN
|
||||
EXECUTE IMMEDIATE 'ALTER TABLE entities DROP COLUMN entity_kind';
|
||||
EXCEPTION WHEN OTHERS THEN
|
||||
IF SQLCODE != -904 THEN RAISE; END IF;
|
||||
END;
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade, oracle=_oracle_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade, oracle=_oracle_downgrade)
|
||||
+259
@@ -0,0 +1,259 @@
|
||||
"""Install the maintenance discovery routines into the configured schema.
|
||||
|
||||
The three discovery routines driving the background maintenance loop —
|
||||
``banks_needing_consolidation()``, ``schemas_with_expired_rows(...)`` and
|
||||
``mental_models_with_cron()`` — were installed into ``public`` and gated on the
|
||||
run being the base run (no ``target_schema``) or an explicit
|
||||
``target_schema='public'`` run (``e5f6a7b8c9d0`` → ``b2d4f6a8c1e3`` →
|
||||
``c7e9f1a3b5d2``, ``f4d1c2b3a5e6``).
|
||||
|
||||
That leaves a **single-tenant deployment migrated into a dedicated, non-**
|
||||
``public`` **schema** (``HINDSIGHT_API_DATABASE_SCHEMA=<non-public>``) with no
|
||||
routines at all: the runtime migrates only that one schema, so ``target_schema``
|
||||
is never falsy or ``public``, the gate never opens, and the maintenance loop
|
||||
logs, forever::
|
||||
|
||||
function public.banks_needing_consolidation() does not exist
|
||||
function public.schemas_with_expired_rows(...) does not exist
|
||||
|
||||
The revision is stamped applied, so redeploying the same version does not help
|
||||
(issue #2638; #2056 only fixed the ``public``/base-run case).
|
||||
|
||||
**The bug was the hardcoded literal, not the gating.** These routines are
|
||||
database-global — each enumerates ``pg_class`` across every schema and dispatches
|
||||
per schema — so exactly one copy should exist, and the maintenance loop calls the
|
||||
one in ``get_config().database_schema`` (see ``fq_routine``). The old gate
|
||||
installed into whichever schema was named ``public`` instead of whichever schema
|
||||
the deployment is actually configured to use. Comparing ``target_schema`` against
|
||||
the configured schema instead of the literal fixes #2638 at the source.
|
||||
|
||||
That also keeps the property the gate existed for: exactly one migration run
|
||||
satisfies the predicate, so concurrent per-schema runs never issue competing
|
||||
``CREATE OR REPLACE`` against the same ``pg_proc`` row and cannot hit
|
||||
``tuple concurrently updated``. No cross-process coordination is required — in
|
||||
particular no advisory lock, which is unusable here because Hindsight runs behind
|
||||
connection poolers and managed PG services (see #2817).
|
||||
|
||||
Runs targeting any *other* schema drop the routines from that schema rather than
|
||||
merely skipping. An earlier revision of this migration installed a copy into
|
||||
every schema it touched, which left one dead duplicate per tenant on any database
|
||||
that ran it; the drop makes the next migration pass clean those up instead of
|
||||
leaving them behind forever.
|
||||
|
||||
PostgreSQL only: the maintenance loop and worker poller are PG-only, so the
|
||||
Oracle slot is intentionally absent (mirrors ``e5f6a7b8c9d0``).
|
||||
|
||||
Revision ID: b6d2f8a4c1e7
|
||||
Revises: a8c1e4f7b0d3
|
||||
Create Date: 2026-07-20
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
from hindsight_api.config import get_config
|
||||
|
||||
revision: str = "b6d2f8a4c1e7"
|
||||
down_revision: str | Sequence[str] | None = "a8c1e4f7b0d3"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _configured_schema() -> str:
|
||||
"""The one schema this deployment's routines live in and are called from."""
|
||||
return get_config().database_schema or "public"
|
||||
|
||||
|
||||
def _target_schema() -> str | None:
|
||||
return context.config.get_main_option("target_schema")
|
||||
|
||||
|
||||
def _is_install_run() -> bool:
|
||||
"""True for the single run that owns the routines.
|
||||
|
||||
The base run (no ``target_schema``) and the run targeting the configured
|
||||
schema are the same deployment-level run; every other target is a tenant
|
||||
schema that must not carry its own copy.
|
||||
"""
|
||||
target = _target_schema()
|
||||
return not target or target == _configured_schema()
|
||||
|
||||
|
||||
def _prefix(schema: str | None) -> str:
|
||||
"""Qualifier for ``schema``, or ``""`` to fall back to ``search_path``."""
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
if not _is_install_run():
|
||||
_drop_stray_copies()
|
||||
return
|
||||
schema = _prefix(_target_schema())
|
||||
|
||||
op.execute(
|
||||
f"""
|
||||
CREATE OR REPLACE FUNCTION {schema}banks_needing_consolidation()
|
||||
RETURNS TABLE(schema_name text, bank_id text)
|
||||
LANGUAGE plpgsql STABLE
|
||||
AS $fn$
|
||||
DECLARE
|
||||
sch text;
|
||||
BEGIN
|
||||
FOR sch IN
|
||||
SELECT n.nspname
|
||||
FROM pg_class c
|
||||
JOIN pg_namespace n ON n.oid = c.relnamespace
|
||||
WHERE c.relname = 'memory_units' AND c.relkind = 'r'
|
||||
LOOP
|
||||
BEGIN
|
||||
RETURN QUERY EXECUTE format($q$
|
||||
SELECT %1$L::text, m.bank_id
|
||||
FROM %1$I.memory_units m
|
||||
JOIN %1$I.banks b ON b.bank_id = m.bank_id
|
||||
WHERE m.consolidated_at IS NULL
|
||||
AND m.consolidation_failed_at IS NULL
|
||||
AND m.fact_type IN ('experience', 'world')
|
||||
AND COALESCE(b.config -> 'enable_auto_consolidation', 'true'::jsonb) <> 'false'::jsonb
|
||||
AND NOT EXISTS (
|
||||
SELECT 1 FROM %1$I.async_operations o
|
||||
WHERE o.bank_id = m.bank_id
|
||||
AND o.operation_type = 'consolidation'
|
||||
AND o.status IN ('pending', 'processing')
|
||||
)
|
||||
GROUP BY m.bank_id
|
||||
$q$, sch);
|
||||
EXCEPTION
|
||||
-- Schema or its tables vanished between the pg_class
|
||||
-- snapshot and this query (tenant dropped or migrating).
|
||||
WHEN undefined_table OR invalid_schema_name OR undefined_column THEN
|
||||
CONTINUE;
|
||||
END;
|
||||
END LOOP;
|
||||
END;
|
||||
$fn$;
|
||||
"""
|
||||
)
|
||||
|
||||
op.execute(
|
||||
f"""
|
||||
CREATE OR REPLACE FUNCTION {schema}schemas_with_expired_rows(
|
||||
p_table text, p_ts_col text, p_days int
|
||||
)
|
||||
RETURNS SETOF text
|
||||
LANGUAGE plpgsql STABLE
|
||||
AS $fn$
|
||||
DECLARE
|
||||
sch text;
|
||||
has_expired boolean;
|
||||
BEGIN
|
||||
IF p_days IS NULL OR p_days <= 0 THEN
|
||||
RETURN;
|
||||
END IF;
|
||||
FOR sch IN
|
||||
SELECT n.nspname
|
||||
FROM pg_class c
|
||||
JOIN pg_namespace n ON n.oid = c.relnamespace
|
||||
WHERE c.relname = p_table AND c.relkind = 'r'
|
||||
LOOP
|
||||
BEGIN
|
||||
EXECUTE format(
|
||||
'SELECT EXISTS (SELECT 1 FROM %I.%I WHERE %I < NOW() - make_interval(days => $1))',
|
||||
sch, p_table, p_ts_col
|
||||
) INTO has_expired USING p_days;
|
||||
EXCEPTION
|
||||
-- Schema or its table vanished mid-scan; skip it.
|
||||
WHEN undefined_table OR invalid_schema_name OR undefined_column THEN
|
||||
CONTINUE;
|
||||
END;
|
||||
IF has_expired THEN
|
||||
RETURN NEXT sch;
|
||||
END IF;
|
||||
END LOOP;
|
||||
END;
|
||||
$fn$;
|
||||
"""
|
||||
)
|
||||
|
||||
op.execute(
|
||||
f"""
|
||||
CREATE OR REPLACE FUNCTION {schema}mental_models_with_cron()
|
||||
RETURNS TABLE(schema_name text, bank_id text, mental_model_id text,
|
||||
refresh_cron text, last_refreshed_at timestamptz)
|
||||
LANGUAGE plpgsql STABLE
|
||||
AS $fn$
|
||||
DECLARE
|
||||
sch text;
|
||||
BEGIN
|
||||
FOR sch IN
|
||||
SELECT n.nspname
|
||||
FROM pg_class c
|
||||
JOIN pg_namespace n ON n.oid = c.relnamespace
|
||||
WHERE c.relname = 'mental_models' AND c.relkind = 'r'
|
||||
LOOP
|
||||
BEGIN
|
||||
RETURN QUERY EXECUTE format($q$
|
||||
SELECT %1$L::text, mm.bank_id::text, mm.id::text,
|
||||
mm.trigger->>'refresh_cron', mm.last_refreshed_at
|
||||
FROM %1$I.mental_models mm
|
||||
WHERE COALESCE(mm.trigger->>'refresh_cron', '') <> ''
|
||||
AND NOT EXISTS (
|
||||
SELECT 1 FROM %1$I.async_operations o
|
||||
WHERE o.bank_id = mm.bank_id
|
||||
AND o.operation_type = 'refresh_mental_model'
|
||||
AND o.status IN ('pending', 'processing')
|
||||
AND o.task_payload->>'mental_model_id' = mm.id::text
|
||||
)
|
||||
$q$, sch);
|
||||
EXCEPTION
|
||||
-- Schema or its tables vanished between the pg_class
|
||||
-- snapshot and this query (tenant dropped or migrating).
|
||||
WHEN undefined_table OR invalid_schema_name OR undefined_column THEN
|
||||
CONTINUE;
|
||||
END;
|
||||
END LOOP;
|
||||
END;
|
||||
$fn$;
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def _drop_routines(schema: str | None) -> None:
|
||||
prefix = _prefix(schema)
|
||||
op.execute(f"DROP FUNCTION IF EXISTS {prefix}mental_models_with_cron()")
|
||||
op.execute(f"DROP FUNCTION IF EXISTS {prefix}schemas_with_expired_rows(text, text, int)")
|
||||
op.execute(f"DROP FUNCTION IF EXISTS {prefix}banks_needing_consolidation()")
|
||||
|
||||
|
||||
def _drop_stray_copies() -> None:
|
||||
"""Remove per-tenant duplicates left by the first cut of this migration.
|
||||
|
||||
That version installed a copy into every schema it touched, so a database
|
||||
that ran it carries one dead duplicate per tenant — only the copy in the
|
||||
configured schema is ever called. Dropping here means the next migration pass
|
||||
cleans them up; without it they would persist for the life of the database.
|
||||
|
||||
Safe on a database that never had them: ``DROP FUNCTION IF EXISTS`` is a
|
||||
no-op, and this branch never runs for the configured schema.
|
||||
"""
|
||||
_drop_routines(_target_schema())
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
# Only drop what this migration uniquely owns. When the configured schema is
|
||||
# ``public`` the copies there belong to e5f6a7b8c9d0 / f4d1c2b3a5e6, which are
|
||||
# still applied at this point and drop them on their own downgrade — removing
|
||||
# them here would strand those migrations without the functions they claim to
|
||||
# have installed.
|
||||
if not _is_install_run() or _configured_schema() == "public":
|
||||
return
|
||||
_drop_routines(_target_schema())
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade)
|
||||
+90
@@ -0,0 +1,90 @@
|
||||
"""Add ``causal_links`` to the curation archive (invalidated_memory_units).
|
||||
|
||||
Causal edges (``caused_by`` and the historical ``causes``/``enables``/
|
||||
``prevents``) are retain-time extraction output: unlike temporal and semantic
|
||||
links they cannot be recomputed from dates or embeddings, and graph maintenance
|
||||
never rebuilds them. Invalidation MOVES a fact out of ``memory_units``, so the
|
||||
``memory_links → memory_units`` FK cascade deletes every incident edge — and
|
||||
revert had no way to bring the causal ones back (#2864).
|
||||
|
||||
This column parks the descriptors of the causal edges incident to an archived
|
||||
fact — ``[{"from_unit_id", "to_unit_id", "link_type", "weight"}, ...]`` — so
|
||||
revert can rematerialize them. It is deliberately unindexed and lives only on
|
||||
the archive: live facts keep their causal edges in ``memory_links`` (curation
|
||||
edits no longer delete them), and the archive is small, cold, and only read by
|
||||
low-frequency curation operations.
|
||||
|
||||
Revision ID: c7d1e9a4b3f2
|
||||
Revises: d7b2f8a1c934
|
||||
Create Date: 2026-07-24
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "c7d1e9a4b3f2"
|
||||
down_revision: str | Sequence[str] | None = "d7b2f8a1c934"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _pg_schema_prefix() -> str:
|
||||
"""Schema-qualifier for raw SQL on PG (multi-tenant search_path)."""
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
schema = _pg_schema_prefix()
|
||||
# NOT NULL DEFAULT is metadata-only on PG 11+, so this is cheap even on a
|
||||
# large archive. Existing rows read as "no causal edges captured" — edges
|
||||
# lost before this migration cannot be reconstructed and are not guessed.
|
||||
op.execute(
|
||||
f"ALTER TABLE {schema}invalidated_memory_units "
|
||||
f"ADD COLUMN IF NOT EXISTS causal_links JSONB NOT NULL DEFAULT '[]'::jsonb"
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema = _pg_schema_prefix()
|
||||
op.execute(f"ALTER TABLE {schema}invalidated_memory_units DROP COLUMN IF EXISTS causal_links")
|
||||
|
||||
|
||||
def _oracle_upgrade() -> None:
|
||||
# Kept in sync with PG for schema parity (curation itself is PostgreSQL-only
|
||||
# today — it introspects pg_attribute to move rows between the two tables).
|
||||
# Swallow ORA-01430 (column already exists) so the migration is idempotent.
|
||||
op.execute(
|
||||
"""
|
||||
BEGIN
|
||||
EXECUTE IMMEDIATE 'ALTER TABLE invalidated_memory_units ADD (causal_links CLOB DEFAULT ''[]''
|
||||
CONSTRAINT imu_causal_links_json CHECK (causal_links IS JSON))';
|
||||
EXCEPTION WHEN OTHERS THEN
|
||||
IF SQLCODE != -1430 THEN RAISE; END IF;
|
||||
END;
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def _oracle_downgrade() -> None:
|
||||
# Swallow ORA-00904 (column does not exist).
|
||||
op.execute(
|
||||
"""
|
||||
BEGIN
|
||||
EXECUTE IMMEDIATE 'ALTER TABLE invalidated_memory_units DROP COLUMN causal_links';
|
||||
EXCEPTION WHEN OTHERS THEN
|
||||
IF SQLCODE != -904 THEN RAISE; END IF;
|
||||
END;
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade, oracle=_oracle_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade, oracle=_oracle_downgrade)
|
||||
+150
@@ -0,0 +1,150 @@
|
||||
"""Add the ``schemas_with_expired_operations`` cross-tenant discovery routine.
|
||||
|
||||
The worker's terminal-operation cleanup (``a8c1e4f7b0d3``) opens a connection
|
||||
and a prune transaction against *every* tenant schema on every cleanup cycle,
|
||||
whether or not that tenant has anything to prune. At thousands of tenants that
|
||||
is a per-cycle query storm whose cost is paid entirely by idle schemas.
|
||||
|
||||
This is the same problem ``public.schemas_with_expired_rows`` already solves for
|
||||
the ``audit_log`` / ``llm_requests`` retention sweeps (``e5f6a7b8c9d0``): one
|
||||
round-trip returns just the schemas that actually hold expired rows, and the
|
||||
caller then does real work only there. ``async_operations`` needs its own
|
||||
routine rather than reusing that one because eligibility is not "row older than
|
||||
N days" — pending and processing rows are never prunable, so the status filter
|
||||
has to be part of the predicate.
|
||||
|
||||
Install policy mirrors ``b6d2f8a4c1e7`` (#2638/#2824), the current behaviour for
|
||||
the sibling routines: the routine is database-global — it enumerates ``pg_class``
|
||||
across every schema and dispatches per schema — so exactly one copy should exist,
|
||||
installed into the schema this deployment is *configured* to use and called from
|
||||
there via ``fq_routine``. Gating on the literal ``"public"`` instead of the
|
||||
configured schema is what left single-tenant deployments in a dedicated
|
||||
non-``public`` schema without the routine (#2638).
|
||||
|
||||
Exactly one migration run satisfies that predicate, so concurrent per-schema runs
|
||||
never issue competing ``CREATE OR REPLACE`` against the same ``pg_proc`` row and
|
||||
cannot hit ``tuple concurrently updated``. No cross-process coordination is
|
||||
required — in particular no advisory lock, which is unusable here because
|
||||
Hindsight runs behind connection poolers and managed PG services (see #2817).
|
||||
|
||||
Each per-schema probe runs in its own ``BEGIN ... EXCEPTION`` block so a tenant
|
||||
dropped mid-scan is skipped instead of aborting the sweep (see ``c7e9f1a3b5d2``).
|
||||
|
||||
Revision ID: d7b2f8a1c934
|
||||
Revises: b6d2f8a4c1e7
|
||||
Create Date: 2026-07-20
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
from hindsight_api.config import get_config
|
||||
|
||||
revision: str = "d7b2f8a1c934"
|
||||
down_revision: str | Sequence[str] | None = "b6d2f8a4c1e7"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _configured_schema() -> str:
|
||||
"""The one schema this deployment's routines live in and are called from."""
|
||||
return get_config().database_schema or "public"
|
||||
|
||||
|
||||
def _target_schema() -> str | None:
|
||||
return context.config.get_main_option("target_schema")
|
||||
|
||||
|
||||
def _is_install_run() -> bool:
|
||||
"""True for the single run that owns the routine (mirrors b6d2f8a4c1e7)."""
|
||||
target = _target_schema()
|
||||
return not target or target == _configured_schema()
|
||||
|
||||
|
||||
def _prefix(schema: str | None) -> str:
|
||||
"""Qualifier for ``schema``, or ``""`` to fall back to ``search_path``."""
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def _drop_routine(schema: str | None) -> None:
|
||||
op.execute(f"DROP FUNCTION IF EXISTS {_prefix(schema)}schemas_with_expired_operations(int)")
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
if not _is_install_run():
|
||||
# Tenant schemas must not carry their own copy: the routine is
|
||||
# database-global and only the configured schema's copy is ever called.
|
||||
# Dropping (rather than skipping) also cleans up after any interim build
|
||||
# of this branch that installed per-schema copies.
|
||||
_drop_routine(_target_schema())
|
||||
return
|
||||
schema = _prefix(_target_schema())
|
||||
op.execute(
|
||||
f"""
|
||||
CREATE OR REPLACE FUNCTION {schema}schemas_with_expired_operations(p_days int)
|
||||
RETURNS SETOF text
|
||||
LANGUAGE plpgsql STABLE
|
||||
AS $fn$
|
||||
DECLARE
|
||||
sch text;
|
||||
has_expired boolean;
|
||||
BEGIN
|
||||
-- Zero (or negative) retention means "keep forever": report nothing
|
||||
-- so the caller skips the sweep entirely.
|
||||
IF p_days IS NULL OR p_days <= 0 THEN
|
||||
RETURN;
|
||||
END IF;
|
||||
FOR sch IN
|
||||
SELECT n.nspname
|
||||
FROM pg_class c
|
||||
JOIN pg_namespace n ON n.oid = c.relnamespace
|
||||
WHERE c.relname = 'async_operations' AND c.relkind = 'r'
|
||||
LOOP
|
||||
BEGIN
|
||||
-- Matches the worker's prune predicate: only terminal rows
|
||||
-- are eligible, so a schema holding nothing but pending or
|
||||
-- processing work is correctly reported as having nothing
|
||||
-- to prune. Uses idx_async_operations_terminal_cleanup.
|
||||
EXECUTE format(
|
||||
'SELECT EXISTS ('
|
||||
' SELECT 1 FROM %I.async_operations'
|
||||
' WHERE status IN (''completed'', ''failed'', ''cancelled'')'
|
||||
' AND updated_at < NOW() - make_interval(days => $1)'
|
||||
')',
|
||||
sch
|
||||
) INTO has_expired USING p_days;
|
||||
EXCEPTION
|
||||
-- Schema or its table vanished between the pg_class
|
||||
-- snapshot and this probe (tenant dropped or migrating).
|
||||
WHEN undefined_table OR invalid_schema_name OR undefined_column THEN
|
||||
CONTINUE;
|
||||
END;
|
||||
IF has_expired THEN
|
||||
RETURN NEXT sch;
|
||||
END IF;
|
||||
END LOOP;
|
||||
END;
|
||||
$fn$;
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
# This migration is the sole creator of this routine — no older migration
|
||||
# owns a copy the way e5f6a7b8c9d0 owns the public sibling routines — so the
|
||||
# install run's own copy is always ours to drop.
|
||||
if not _is_install_run():
|
||||
return
|
||||
_drop_routine(_target_schema())
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
# Oracle slot intentionally absent: this mirrors the PostgreSQL-only
|
||||
# maintenance routines, and the Oracle worker keeps its per-schema sweep.
|
||||
run_for_dialect(pg=_pg_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade)
|
||||
+114
@@ -0,0 +1,114 @@
|
||||
"""Drop the never-written `access_count` column from memory_units (and its archive).
|
||||
|
||||
``memory_units.access_count`` has been dead since the initial schema
|
||||
(5a366d414dce): no code path anywhere in the repo ever writes it, and — despite
|
||||
the ``access_count DESC`` index created alongside it — no query ever reads or
|
||||
orders by it either. It is 0 on every row of every install. The lone remaining
|
||||
mentions were an index, a stale comment naming an ``access_count_update`` task
|
||||
type that was never implemented, and the column's name in the Oracle backend's
|
||||
numeric-RETURNING list; all three go away with this change.
|
||||
|
||||
The column is dropped from the curation archive too. ``invalidated_memory_units``
|
||||
was cloned ``LIKE memory_units`` (c9a1b2d3e4f5), so it inherited the column, and
|
||||
curation's INSERT…SELECT round-trip builds its column list from the catalog
|
||||
(``writes.py::_memory_unit_columns``) — the two tables must stay in lockstep or
|
||||
the round-trip breaks on a column-count mismatch.
|
||||
|
||||
Dropping the column implicitly drops its index on both dialects
|
||||
(``idx_memory_units_access_count`` on PG, ``idx_mu_access_count`` on Oracle), so
|
||||
PostgreSQL also stops maintaining a btree that nothing ever probed.
|
||||
|
||||
Cost: on PostgreSQL ``DROP COLUMN`` is metadata-only (the attribute is marked
|
||||
dropped, no table rewrite). On Oracle it does delete the column data row by row,
|
||||
so on a large ``memory_units`` this migration is not free — it is still bounded
|
||||
work on a single small integer column, and Oracle installs of that size can run
|
||||
it during a maintenance window ahead of the upgrade if they prefer.
|
||||
|
||||
Revision ID: e4a7c1b9d2f6
|
||||
Revises: a9b8c7d6e5f4
|
||||
Create Date: 2026-08-03
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "e4a7c1b9d2f6"
|
||||
down_revision: str | Sequence[str] | None = "a9b8c7d6e5f4"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
_TABLES = ("memory_units", "invalidated_memory_units")
|
||||
|
||||
|
||||
def _pg_schema_prefix() -> str:
|
||||
"""Schema-qualifier for raw SQL on PG (multi-tenant search_path)."""
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
schema = _pg_schema_prefix()
|
||||
for table in _TABLES:
|
||||
# Drops idx_memory_units_access_count along with the column.
|
||||
op.execute(f"ALTER TABLE {schema}{table} DROP COLUMN IF EXISTS access_count")
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema = _pg_schema_prefix()
|
||||
for table in _TABLES:
|
||||
op.execute(f"ALTER TABLE {schema}{table} ADD COLUMN IF NOT EXISTS access_count integer NOT NULL DEFAULT 0")
|
||||
# The archive was cloned without indexes; only the live table carried one.
|
||||
op.execute(f"CREATE INDEX IF NOT EXISTS idx_memory_units_access_count ON {schema}memory_units (access_count DESC)")
|
||||
|
||||
|
||||
def _oracle_upgrade() -> None:
|
||||
# Oracle has no `DROP COLUMN IF EXISTS`; swallow ORA-00904 (column does not
|
||||
# exist) so the migration is idempotent and safe on a schema that already
|
||||
# lacks the column. Dropping the column also drops idx_mu_access_count.
|
||||
for table in _TABLES:
|
||||
op.execute(
|
||||
f"""
|
||||
BEGIN
|
||||
EXECUTE IMMEDIATE 'ALTER TABLE {table} DROP COLUMN access_count';
|
||||
EXCEPTION WHEN OTHERS THEN
|
||||
IF SQLCODE != -904 THEN RAISE; END IF;
|
||||
END;
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def _oracle_downgrade() -> None:
|
||||
# Swallow ORA-01430 (column already exists) for idempotency. Matches the
|
||||
# Oracle baseline's declaration: NUMBER(10) DEFAULT 0 NOT NULL.
|
||||
for table in _TABLES:
|
||||
op.execute(
|
||||
f"""
|
||||
BEGIN
|
||||
EXECUTE IMMEDIATE
|
||||
'ALTER TABLE {table} ADD (access_count NUMBER(10) DEFAULT 0 NOT NULL)';
|
||||
EXCEPTION WHEN OTHERS THEN
|
||||
IF SQLCODE != -1430 THEN RAISE; END IF;
|
||||
END;
|
||||
"""
|
||||
)
|
||||
# ORA-00955: index name already in use.
|
||||
op.execute(
|
||||
"""
|
||||
BEGIN
|
||||
EXECUTE IMMEDIATE 'CREATE INDEX idx_mu_access_count ON memory_units(access_count DESC)';
|
||||
EXCEPTION WHEN OTHERS THEN
|
||||
IF SQLCODE != -955 THEN RAISE; END IF;
|
||||
END;
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade, oracle=_oracle_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade, oracle=_oracle_downgrade)
|
||||
+96
@@ -0,0 +1,96 @@
|
||||
"""Drop the search_vector column from the curation archive (invalidated_memory_units).
|
||||
|
||||
The archive is cold storage, never a recall surface, and carries no text-search
|
||||
index. Like ``embedding`` (dropped in d4f6a8c2e1b3), ``search_vector`` is a
|
||||
recall-surface column whose type follows the configured text-search backend, so
|
||||
it has no business living on the archive. Earlier curation code copied the live
|
||||
row's ``search_vector`` into ``invalidated_memory_units`` on invalidate; the
|
||||
engine now leaves it out on invalidate and recomputes it on revert, so the
|
||||
column is dead weight.
|
||||
|
||||
Dropping it removes a latent failure mode (#2503): under a non-native backend
|
||||
(pgroonga / pg_textsearch / pg_search / vchord) ``ensure_text_search_extension``
|
||||
reconciles ``memory_units.search_vector`` to ``text`` / ``bm25vector`` but never
|
||||
touched the archive, which the ``LIKE memory_units`` clone (c9a1b2d3e4f5) created
|
||||
as ``tsvector``. The type mismatch then broke the curation INSERT … SELECT
|
||||
round-trip:
|
||||
|
||||
column "search_vector" is of type tsvector but expression is of type text
|
||||
|
||||
With no column at all, there is nothing to mismatch. Unlike ``embedding`` (whose
|
||||
creation sites already omit it), the ``LIKE`` clone still adds ``search_vector``,
|
||||
so this migration does real work on both fresh and existing PostgreSQL databases.
|
||||
|
||||
DROP COLUMN is a metadata-only operation on both PostgreSQL and Oracle 23ai (no
|
||||
table rewrite), so it is cheap even across many tenant schemas. The downgrade
|
||||
re-adds an empty ``tsvector`` column (its original creation type).
|
||||
|
||||
Revision ID: e7c3a9f1b2d5
|
||||
Revises: b57a7c9e0d13
|
||||
Create Date: 2026-07-02
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "e7c3a9f1b2d5"
|
||||
down_revision: str | Sequence[str] | None = "b57a7c9e0d13"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _pg_schema_prefix() -> str:
|
||||
"""Schema-qualifier for raw SQL on PG (multi-tenant search_path)."""
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
schema = _pg_schema_prefix()
|
||||
op.execute(f"ALTER TABLE {schema}invalidated_memory_units DROP COLUMN IF EXISTS search_vector")
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema = _pg_schema_prefix()
|
||||
# Re-add as the original tsvector creation type; comes back empty regardless.
|
||||
op.execute(f"ALTER TABLE {schema}invalidated_memory_units ADD COLUMN IF NOT EXISTS search_vector tsvector")
|
||||
|
||||
|
||||
def _oracle_upgrade() -> None:
|
||||
# Oracle has no `DROP COLUMN IF EXISTS`; swallow ORA-00904 (column does not
|
||||
# exist) so the migration is idempotent and safe on a schema whose baseline
|
||||
# may already omit the column.
|
||||
op.execute(
|
||||
"""
|
||||
BEGIN
|
||||
EXECUTE IMMEDIATE 'ALTER TABLE invalidated_memory_units DROP COLUMN search_vector';
|
||||
EXCEPTION WHEN OTHERS THEN
|
||||
IF SQLCODE != -904 THEN RAISE; END IF;
|
||||
END;
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def _oracle_downgrade() -> None:
|
||||
# Swallow ORA-01430 (column already exists) for idempotency. Oracle stores
|
||||
# search_vector as CLOB (see the Oracle baseline), so re-add it as CLOB.
|
||||
op.execute(
|
||||
"""
|
||||
BEGIN
|
||||
EXECUTE IMMEDIATE 'ALTER TABLE invalidated_memory_units ADD (search_vector CLOB)';
|
||||
EXCEPTION WHEN OTHERS THEN
|
||||
IF SQLCODE != -1430 THEN RAISE; END IF;
|
||||
END;
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade, oracle=_oracle_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade, oracle=_oracle_downgrade)
|
||||
+85
@@ -0,0 +1,85 @@
|
||||
"""Repair: drop the stale global memory_units vector index on per-bank backends.
|
||||
|
||||
Revision ID: f2a6d8c4b1e9
|
||||
Revises: e4a7c1b9d2f6
|
||||
Create Date: 2026-08-06
|
||||
|
||||
Migration d5e6f7a8b9c0 dropped the global ``idx_memory_units_embedding`` for
|
||||
per-bank backends (every vector search is bank + fact_type scoped and served
|
||||
by the ``idx_mu_emb_*`` partial indexes; the global index is never chosen by
|
||||
the planner). However, older versions of the post-migration reconcile
|
||||
(``ensure_vector_extension``) recreated the index when they found none, so
|
||||
schemas that were provisioned or reconciled in that window carry it to this
|
||||
day — paying a second vector graph insertion on every ``memory_units`` write
|
||||
for an index no query uses.
|
||||
|
||||
This repair drops the leftover index. It is intentionally a migration, not
|
||||
runtime reconcile behavior: ``DROP INDEX`` takes an ACCESS EXCLUSIVE lock on
|
||||
``memory_units``, which belongs in the versioned, once-per-schema migration
|
||||
path — not in code that runs at unpredictable times during startup or tenant
|
||||
provisioning. The reconcile now leaves memory_units vector-index DDL to
|
||||
migrations entirely on per-bank backends.
|
||||
|
||||
ScaNN deployments keep the global index by design (filtered vector search over
|
||||
a global index; per-bank partial indexes cannot be built safely there), so the
|
||||
migration is a no-op for them.
|
||||
"""
|
||||
|
||||
import os
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "f2a6d8c4b1e9"
|
||||
down_revision: str | Sequence[str] | None = "e4a7c1b9d2f6"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _pg_schema_prefix() -> str:
|
||||
"""Schema-qualifier for raw SQL on PG (multi-tenant search_path)."""
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def _configured_vector_extension() -> str:
|
||||
ext = os.getenv("HINDSIGHT_API_VECTOR_EXTENSION", "pgvector").lower()
|
||||
if ext not in {"pgvector", "pgvectorscale", "vchord", "scann"}:
|
||||
raise ValueError(
|
||||
f"Invalid HINDSIGHT_API_VECTOR_EXTENSION: {ext}. Must be 'pgvector', 'vchord', 'pgvectorscale', or 'scann'"
|
||||
)
|
||||
return ext
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
# ScaNN uses a global vector index by design — nothing stale to repair.
|
||||
if _configured_vector_extension() == "scann":
|
||||
return
|
||||
schema = _pg_schema_prefix()
|
||||
# DROP INDEX needs ACCESS EXCLUSIVE on memory_units. While it waits for
|
||||
# in-flight transactions, every new query on the table queues behind it,
|
||||
# so on a write-busy schema an unbounded wait can pile up traffic. Fail
|
||||
# fast instead: the migration errors, the schema stays below head, and
|
||||
# the next migration pass retries — preferable to freezing the table.
|
||||
# SET LOCAL scopes the timeout to this migration's transaction.
|
||||
op.execute("SET LOCAL lock_timeout = '10s'")
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_memory_units_embedding")
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
# Intentional no-op: recreating a potentially multi-GB vector index that no
|
||||
# query uses is not a safe downgrade action. Downgrading past d5e6f7a8b9c0
|
||||
# restores the global index for deployments that genuinely need it.
|
||||
pass
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
# PG-only repair: the stale index is a PostgreSQL artifact of the old
|
||||
# reconcile; Oracle deployments never had a reconcile that created it.
|
||||
run_for_dialect(pg=_pg_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,171 @@
|
||||
"""Markdown rendering for knowledge pages.
|
||||
|
||||
Knowledge pages render as *read-only* markdown documents over the existing mental
|
||||
models: each mental model becomes a markdown body with a YAML frontmatter block
|
||||
(``type`` required; ``title``/``description``/``tags``/``timestamp`` optional).
|
||||
|
||||
This module is intentionally pure: every function transforms the mental-model
|
||||
dicts returned by ``MemoryEngine.list_mental_models`` / ``get_mental_model`` and
|
||||
never touches the database. That keeps rendering unit-testable without a DB or
|
||||
LLM and lets the HTTP layer stay a thin wrapper.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
# Every page carries exactly one ``type`` frontmatter field. We default to this
|
||||
# when a page does not declare one via a ``type:<x>`` tag.
|
||||
DEFAULT_PAGE_TYPE = "knowledge-page"
|
||||
|
||||
# A page declares its ``type`` through a tag of the form ``type:runbook``.
|
||||
# This keeps rendering schema-free (no new mental_models column): the type is
|
||||
# lifted from the existing tags array.
|
||||
TYPE_TAG_PREFIX = "type:"
|
||||
|
||||
INDEX_FILENAME = "index.md"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class PageType:
|
||||
"""A page's ``type`` and the tags that remain after the type tag is split off."""
|
||||
|
||||
type: str
|
||||
display_tags: list[str]
|
||||
|
||||
|
||||
def _scalar(value: Any) -> str:
|
||||
"""Emit a YAML-safe double-quoted scalar.
|
||||
|
||||
We always double-quote so arbitrary page names / source queries can't be
|
||||
misread as YAML special forms (``true``, ``2026-01-01``, ``- x``, etc.).
|
||||
"""
|
||||
text = str(value)
|
||||
escaped = text.replace("\\", "\\\\").replace('"', '\\"').replace("\n", "\\n").replace("\r", "")
|
||||
return f'"{escaped}"'
|
||||
|
||||
|
||||
def page_type(tags: list[str] | None) -> PageType:
|
||||
"""Split a ``type`` out of the tag list.
|
||||
|
||||
The first ``type:<x>`` tag wins; all ``type:`` tags are removed from the
|
||||
returned ``display_tags`` so they don't leak into the page's displayed tags.
|
||||
Falls back to :data:`DEFAULT_PAGE_TYPE`.
|
||||
"""
|
||||
resolved = DEFAULT_PAGE_TYPE
|
||||
display: list[str] = []
|
||||
for tag in tags or []:
|
||||
if tag.startswith(TYPE_TAG_PREFIX):
|
||||
suffix = tag[len(TYPE_TAG_PREFIX) :].strip()
|
||||
if suffix and resolved == DEFAULT_PAGE_TYPE:
|
||||
resolved = suffix
|
||||
continue
|
||||
display.append(tag)
|
||||
return PageType(type=resolved, display_tags=display)
|
||||
|
||||
|
||||
def _timestamp(mm: dict[str, Any]) -> str | None:
|
||||
return mm.get("last_refreshed_at") or mm.get("created_at")
|
||||
|
||||
|
||||
def frontmatter(mm: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Build the ordered frontmatter mapping for a mental model.
|
||||
|
||||
``None``/empty values are dropped by :func:`render_frontmatter`.
|
||||
"""
|
||||
pt = page_type(mm.get("tags"))
|
||||
return {
|
||||
"id": mm.get("id"),
|
||||
"type": pt.type,
|
||||
"title": mm.get("name"),
|
||||
"description": mm.get("source_query"),
|
||||
"tags": pt.display_tags,
|
||||
"timestamp": _timestamp(mm),
|
||||
}
|
||||
|
||||
|
||||
def render_frontmatter(fm: dict[str, Any]) -> str:
|
||||
"""Render a frontmatter mapping into a ``---`` fenced YAML block."""
|
||||
lines = ["---"]
|
||||
for key, value in fm.items():
|
||||
if value is None:
|
||||
continue
|
||||
if isinstance(value, list):
|
||||
if not value:
|
||||
continue
|
||||
lines.append(f"{key}:")
|
||||
lines.extend(f" - {_scalar(item)}" for item in value)
|
||||
else:
|
||||
lines.append(f"{key}: {_scalar(value)}")
|
||||
lines.append("---")
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def render_document(mm: dict[str, Any]) -> str:
|
||||
"""Render a full markdown document: frontmatter block + markdown body."""
|
||||
body = (mm.get("content") or "").strip()
|
||||
return f"{render_frontmatter(frontmatter(mm))}\n\n{body}\n" if body else f"{render_frontmatter(frontmatter(mm))}\n"
|
||||
|
||||
|
||||
def page_filename(page_id: str) -> str:
|
||||
"""Bundle filename for a page id."""
|
||||
return f"{page_id}.md"
|
||||
|
||||
|
||||
def log_filename(page_id: str) -> str:
|
||||
"""Reserved per-page history filename."""
|
||||
return f"{page_id}.log.md"
|
||||
|
||||
|
||||
def render_index(nodes: list[dict[str, Any]]) -> str:
|
||||
"""Render the reserved ``index.md`` — nested markdown navigation over the tree.
|
||||
|
||||
``nodes`` is the flat folder/page list (each with ``id``, ``kind``, ``name``,
|
||||
``parent_id``); folders nest their children, pages link to their ``.md``.
|
||||
"""
|
||||
fm = render_frontmatter({"type": "index", "title": "Knowledge base"})
|
||||
lines = [fm, "", "# Knowledge base", ""]
|
||||
|
||||
children: dict[Any, list[dict[str, Any]]] = {}
|
||||
for node in nodes:
|
||||
children.setdefault(node.get("parent_id"), []).append(node)
|
||||
|
||||
def walk(parent: Any, depth: int) -> None:
|
||||
ordered = sorted(children.get(parent, []), key=lambda n: (n.get("sort_order", 0), n.get("name") or ""))
|
||||
for node in ordered:
|
||||
indent = " " * depth
|
||||
if node.get("kind") == "folder":
|
||||
lines.append(f"{indent}- **{node['name']}/**")
|
||||
walk(node["id"], depth + 1)
|
||||
else:
|
||||
description = node.get("source_query") or node.get("description")
|
||||
link = f"{indent}- [{node['name']}](./{page_filename(node['id'])})"
|
||||
lines.append(f"{link} — {description}" if description else link)
|
||||
|
||||
walk(None, 0)
|
||||
if len(lines) == 4:
|
||||
lines.append("_No knowledge pages yet._")
|
||||
return "\n".join(lines) + "\n"
|
||||
|
||||
|
||||
def render_log(mm: dict[str, Any], history: list[dict[str, Any]]) -> str:
|
||||
"""Render the reserved per-page ``log.md`` from refresh history.
|
||||
|
||||
Each history entry is ``{previous_content, previous_reflect_response,
|
||||
changed_at}`` (newest first), capturing the content *before* a refresh.
|
||||
"""
|
||||
name = mm.get("name") or mm.get("id")
|
||||
fm = render_frontmatter({"type": "log", "title": f"{name} — history"})
|
||||
lines = [fm, "", f"# {name} — history", ""]
|
||||
if not history:
|
||||
lines.append("_No refresh history._")
|
||||
return "\n".join(lines) + "\n"
|
||||
for entry in history:
|
||||
changed_at = entry.get("changed_at") or "unknown"
|
||||
previous = (entry.get("previous_content") or "").strip()
|
||||
lines.append(f"## {changed_at}")
|
||||
lines.append("")
|
||||
lines.append(previous if previous else "_(empty)_")
|
||||
lines.append("")
|
||||
return "\n".join(lines).rstrip() + "\n"
|
||||
@@ -66,11 +66,6 @@ def color_end(text: str) -> str:
|
||||
return color(text, 1.0)
|
||||
|
||||
|
||||
def color_mid(text: str) -> str:
|
||||
"""Color text with gradient middle color."""
|
||||
return color(text, 0.5)
|
||||
|
||||
|
||||
def dim(text: str) -> str:
|
||||
"""Dim/gray text."""
|
||||
return f"\033[38;2;128;128;128m{text}\033[0m"
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -11,8 +11,10 @@ multiple API servers.
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
from dataclasses import asdict, replace
|
||||
from typing import TYPE_CHECKING, Any
|
||||
from dataclasses import asdict, fields, replace
|
||||
from functools import lru_cache
|
||||
from types import UnionType
|
||||
from typing import TYPE_CHECKING, Any, Union, get_args, get_origin
|
||||
|
||||
from hindsight_api.config import (
|
||||
RECALL_BUDGET_FUNCTIONS,
|
||||
@@ -32,6 +34,14 @@ if TYPE_CHECKING:
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class BankConfigPersistenceConflictError(ValueError):
|
||||
"""Raised when a validated bank config update can no longer be persisted."""
|
||||
|
||||
def __init__(self, bank_id: str):
|
||||
self.bank_id = bank_id
|
||||
super().__init__(f"Cannot update config for bank '{bank_id}': the bank does not exist")
|
||||
|
||||
|
||||
def _validate_retain_strategy_chunking(base_config: HindsightConfig, strategies: Any) -> None:
|
||||
"""Validate retain strategy chunking with the same semantics as apply_strategy()."""
|
||||
if not isinstance(strategies, dict):
|
||||
@@ -128,12 +138,13 @@ class ConfigResolver:
|
||||
# Return full config object (dataclass doesn't have __init__ that accepts kwargs, so we update the object)
|
||||
# Create a new config instance by copying the global config and updating fields
|
||||
resolved_config = HindsightConfig(**config_dict)
|
||||
# Multi-LLM chains are static credential fields (never tenant/bank-overridable),
|
||||
# but asdict() above flattened their member dataclasses into plain dicts. Restore
|
||||
# the original typed objects from the global config so the resolved object stays
|
||||
# well-typed for any consumer that reads them.
|
||||
# Multi-LLM chains and the reranker failover chain are static credential fields
|
||||
# (never tenant/bank-overridable), but asdict() above flattened their member
|
||||
# dataclasses into plain dicts. Restore the original typed objects from the global
|
||||
# config so the resolved object stays well-typed for any consumer that reads them.
|
||||
resolved_config = replace(
|
||||
resolved_config,
|
||||
reranker_members=self._global_config.reranker_members,
|
||||
llm_members=self._global_config.llm_members,
|
||||
llm_strategy=self._global_config.llm_strategy,
|
||||
retain_llm_members=self._global_config.retain_llm_members,
|
||||
@@ -286,7 +297,8 @@ class ConfigResolver:
|
||||
|
||||
# Only return active overrides for configurable fields. JSON null is a tombstone
|
||||
# for "Server Default" in the bank-config UI and should not override defaults.
|
||||
return {k: v for k, v in normalized.items() if k in self._configurable_fields and v is not None}
|
||||
active = {k: v for k, v in normalized.items() if k in self._configurable_fields and v is not None}
|
||||
return _coerce_stored_bank_overrides(bank_id, active)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to load bank config for {bank_id}: {e}")
|
||||
|
||||
@@ -326,16 +338,22 @@ class ConfigResolver:
|
||||
k: v for k, v in normalized.items() if k in self._configurable_fields and v is not None
|
||||
}
|
||||
if overrides:
|
||||
result[row["bank_id"]] = overrides
|
||||
result[row["bank_id"]] = _coerce_stored_bank_overrides(row["bank_id"], overrides)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to bulk-load bank configs: {e}")
|
||||
return result
|
||||
|
||||
async def update_bank_config(
|
||||
self, bank_id: str, updates: dict[str, Any], context: RequestContext | None = None
|
||||
) -> None:
|
||||
async def validate_bank_config_updates(
|
||||
self,
|
||||
bank_id: str,
|
||||
updates: dict[str, Any],
|
||||
context: RequestContext | None = None,
|
||||
*,
|
||||
projected_bank_overrides: dict[str, Any] | None = None,
|
||||
check_permissions: bool = True,
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Update bank configuration overrides (with permission checking).
|
||||
Normalize and validate bank configuration overrides.
|
||||
|
||||
Args:
|
||||
bank_id: Bank identifier
|
||||
@@ -344,9 +362,16 @@ class ConfigResolver:
|
||||
or Python field format (llm_provider).
|
||||
Only configurable fields are allowed.
|
||||
context: Request context for permission checking
|
||||
projected_bank_overrides: Bank overrides to use as the validation
|
||||
base instead of loading the current bank row.
|
||||
check_permissions: Whether client field permissions apply to these
|
||||
updates. Server-owned projected values set this to false.
|
||||
|
||||
Returns:
|
||||
Normalized updates ready to persist.
|
||||
|
||||
Raises:
|
||||
ValueError: If attempting to override invalid/disallowed fields
|
||||
ValueError: If attempting to override invalid/disallowed fields.
|
||||
"""
|
||||
# Normalize keys
|
||||
normalized_updates = normalize_config_dict(updates)
|
||||
@@ -378,7 +403,7 @@ class ConfigResolver:
|
||||
)
|
||||
|
||||
# PERMISSIONS: Check tenant/bank permissions
|
||||
if self.tenant_extension and context:
|
||||
if check_permissions and self.tenant_extension and context:
|
||||
try:
|
||||
allowed_fields = await self.tenant_extension.get_allowed_config_fields(context, bank_id)
|
||||
if allowed_fields is not None: # None means "allow all"
|
||||
@@ -388,7 +413,7 @@ class ConfigResolver:
|
||||
f"Not allowed to modify fields: {sorted(disallowed)}. "
|
||||
f"Your permissions allow: {sorted(list(allowed_fields)[:10])}..."
|
||||
if allowed_fields
|
||||
else "Not allowed to modify fields: {sorted(disallowed)}. "
|
||||
else f"Not allowed to modify fields: {sorted(disallowed)}. "
|
||||
"Your permissions do not allow any config modifications."
|
||||
)
|
||||
except ValueError:
|
||||
@@ -397,6 +422,11 @@ class ConfigResolver:
|
||||
logger.warning(f"Failed to check permissions for bank {bank_id}: {e}")
|
||||
# Continue without permission check (fail open for backward compatibility)
|
||||
|
||||
# Validate every value against its declared field type before the
|
||||
# field-specific checks below, so a wrong-shaped value is reported as such
|
||||
# instead of tripping a structural validator with a confusing message.
|
||||
_validate_config_value_types(normalized_updates)
|
||||
|
||||
# Validate entity_labels structure
|
||||
if "entity_labels" in normalized_updates and normalized_updates["entity_labels"] is not None:
|
||||
from .engine.retain.entity_labels import parse_entity_labels
|
||||
@@ -413,6 +443,16 @@ class ConfigResolver:
|
||||
raise ValueError(
|
||||
"Strategy names must not be empty strings. Remove entries with empty names before saving."
|
||||
)
|
||||
# A strategy's overrides are applied with dataclasses.replace() at retain
|
||||
# time, so a wrong-shaped value there wedges the bank exactly as a
|
||||
# top-level one would. Same contract, same door.
|
||||
for strategy_name, strategy_overrides in normalized_updates["retain_strategies"].items():
|
||||
if not isinstance(strategy_overrides, dict):
|
||||
raise ValueError(f"Invalid retain strategy {strategy_name!r}: must be an object")
|
||||
try:
|
||||
_validate_config_value_types(normalize_config_dict(strategy_overrides))
|
||||
except ValueError as e:
|
||||
raise ValueError(f"Invalid retain strategy {strategy_name!r}: {e}") from e
|
||||
|
||||
# Validate recall budget fields
|
||||
_validate_recall_budget_updates(normalized_updates)
|
||||
@@ -427,7 +467,11 @@ class ConfigResolver:
|
||||
)
|
||||
if chunking_fields_updated:
|
||||
config_dict = await self._resolve_parent_config_dict(bank_id, context)
|
||||
active_bank_overrides = await self._load_bank_config(bank_id)
|
||||
active_bank_overrides = (
|
||||
await self._load_bank_config(bank_id)
|
||||
if projected_bank_overrides is None
|
||||
else dict(projected_bank_overrides)
|
||||
)
|
||||
for key, value in normalized_updates.items():
|
||||
if key not in self._configurable_fields:
|
||||
continue
|
||||
@@ -443,17 +487,26 @@ class ConfigResolver:
|
||||
)
|
||||
_validate_retain_strategy_chunking(base_config, base_config.retain_strategies)
|
||||
|
||||
# Persist the override. Banks are created lazily (on first retain), so a
|
||||
# PATCH that precedes any ingestion would otherwise UPDATE zero rows and
|
||||
# silently no-op while returning 200. Ensure the bank row exists first
|
||||
# (this also creates its per-bank vector indexes), then merge defensively:
|
||||
# COALESCE guards against a NULL config column (NULL || jsonb is NULL),
|
||||
# which would drop the override even when a row is updated.
|
||||
from .engine.retain.fact_storage import ensure_bank_exists
|
||||
return normalized_updates
|
||||
|
||||
async def update_bank_config(
|
||||
self, bank_id: str, updates: dict[str, Any], context: RequestContext | None = None
|
||||
) -> None:
|
||||
"""Validate and persist bank configuration overrides for an existing bank.
|
||||
|
||||
Bank creation belongs to ``MemoryEngine``; this raises ``ValueError`` if
|
||||
the bank does not exist rather than silently discarding the overrides.
|
||||
"""
|
||||
normalized_updates = await self.validate_bank_config_updates(bank_id, updates, context)
|
||||
await self._persist_bank_config(bank_id, normalized_updates)
|
||||
|
||||
async def _persist_bank_config(self, bank_id: str, normalized_updates: dict[str, Any]) -> None:
|
||||
"""Persist already-validated overrides without changing bank lifecycle state."""
|
||||
# Bank lifecycle belongs to MemoryEngine. Callers must create the row
|
||||
# before reaching this persistence step. COALESCE guards against a NULL
|
||||
# config column (NULL || jsonb is NULL), which would drop the override.
|
||||
async with self._backend.acquire() as conn:
|
||||
await ensure_bank_exists(conn, bank_id, ops=self._backend.ops)
|
||||
await conn.execute(
|
||||
result = await conn.execute(
|
||||
f"""
|
||||
UPDATE {fq_table("banks")}
|
||||
SET config = COALESCE(config, '{{}}'::jsonb) || $1::jsonb,
|
||||
@@ -464,6 +517,14 @@ class ConfigResolver:
|
||||
bank_id,
|
||||
)
|
||||
|
||||
# A missing bank row matches zero rows, which would otherwise persist
|
||||
# nothing while reporting success. Fail loudly instead: reaching here
|
||||
# without the row means a caller skipped the engine's provisioning step.
|
||||
# (The Oracle wrapper reshapes rowcount into the same "UPDATE <n>" form.)
|
||||
updated = int(result.split()[-1]) if isinstance(result, str) and result.startswith("UPDATE") else 0
|
||||
if updated == 0:
|
||||
raise BankConfigPersistenceConflictError(bank_id)
|
||||
|
||||
logger.info(f"Updated bank config for {bank_id}: {list(normalized_updates.keys())}")
|
||||
|
||||
async def reset_bank_config(self, bank_id: str) -> None:
|
||||
@@ -487,6 +548,147 @@ class ConfigResolver:
|
||||
logger.info(f"Reset bank config for {bank_id} to defaults")
|
||||
|
||||
|
||||
# Fields whose accepted input shape is deliberately wider than the dataclass
|
||||
# annotation, because a dedicated structural validator normalizes them later.
|
||||
_WIDENED_FIELD_TYPES: dict[str, tuple[type, ...]] = {
|
||||
# parse_entity_labels() accepts both the bare list of label groups and the
|
||||
# {"attributes": [...]} envelope, though the field is annotated `list | None`.
|
||||
"entity_labels": (list, dict),
|
||||
}
|
||||
|
||||
|
||||
def _runtime_types(declared: Any) -> tuple[type, ...]:
|
||||
"""Runtime-checkable base classes for a dataclass field annotation.
|
||||
|
||||
Unwraps unions (``str | None``) and generic aliases (``list[str]`` -> ``list``);
|
||||
``None`` is dropped because callers handle the tombstone separately. Returns an
|
||||
empty tuple for anything not reducible to concrete classes, which the callers
|
||||
read as "no type contract to enforce".
|
||||
"""
|
||||
if declared is type(None):
|
||||
return ()
|
||||
origin = get_origin(declared)
|
||||
if origin in (Union, UnionType):
|
||||
return tuple(t for arg in get_args(declared) for t in _runtime_types(arg))
|
||||
if origin is not None:
|
||||
return (origin,) if isinstance(origin, type) else ()
|
||||
return (declared,) if isinstance(declared, type) else ()
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def _configurable_field_types() -> dict[str, tuple[type, ...]]:
|
||||
"""Map each configurable field to the value types it accepts."""
|
||||
configurable = HindsightConfig.get_configurable_fields()
|
||||
field_types: dict[str, tuple[type, ...]] = {}
|
||||
for field in fields(HindsightConfig):
|
||||
if field.name not in configurable:
|
||||
continue
|
||||
allowed = _WIDENED_FIELD_TYPES.get(field.name) or _runtime_types(field.type)
|
||||
if allowed:
|
||||
field_types[field.name] = allowed
|
||||
return field_types
|
||||
|
||||
|
||||
def _value_matches_type(value: Any, allowed: tuple[type, ...]) -> bool:
|
||||
"""Whether ``value`` satisfies a field's declared type contract."""
|
||||
if isinstance(value, bool):
|
||||
# bool is an int subclass; it must not slip into a numeric field.
|
||||
return bool in allowed
|
||||
if isinstance(value, int) and float in allowed:
|
||||
# JSON draws no int/float distinction: 1 is a valid ratio.
|
||||
return True
|
||||
return isinstance(value, allowed)
|
||||
|
||||
|
||||
# Field types are reported to API clients, so name them the way the JSON payload
|
||||
# reads rather than by their Python class.
|
||||
_TYPE_DESCRIPTIONS: dict[type, str] = {
|
||||
bool: "a boolean",
|
||||
int: "an integer",
|
||||
float: "a number",
|
||||
str: "a string",
|
||||
list: "a list",
|
||||
dict: "an object",
|
||||
}
|
||||
|
||||
|
||||
def _describe_types(allowed: tuple[type, ...]) -> str:
|
||||
return " or ".join(dict.fromkeys(_TYPE_DESCRIPTIONS.get(t, t.__name__) for t in allowed))
|
||||
|
||||
|
||||
def _validate_config_value_types(updates: dict[str, Any]) -> None:
|
||||
"""Reject values whose type contradicts the declared HindsightConfig type.
|
||||
|
||||
Without this, the bank-config API happily stores e.g. a JSON object in
|
||||
``observations_mission``; the write succeeds and the bank then fails every
|
||||
consolidation with ``expected string or bytes-like object, got 'dict'`` from
|
||||
deep inside prompt assembly (issue #3218). Reject at the door instead, naming
|
||||
the field and the expected type.
|
||||
"""
|
||||
field_types = _configurable_field_types()
|
||||
for key, value in updates.items():
|
||||
allowed = field_types.get(key)
|
||||
# None is the "clear this override" tombstone; unknown keys are rejected
|
||||
# elsewhere as non-configurable.
|
||||
if allowed is None or value is None:
|
||||
continue
|
||||
if not _value_matches_type(value, allowed):
|
||||
raise ValueError(f"{key} must be {_describe_types(allowed)}, got {type(value).__name__}")
|
||||
|
||||
|
||||
def _coerce_stored_bank_overrides(bank_id: str, overrides: dict[str, Any], where: str = "") -> dict[str, Any]:
|
||||
"""Make stored bank overrides safe to consume, tolerating pre-validation shapes.
|
||||
|
||||
``_validate_config_value_types`` rejects bad types at write time, but banks
|
||||
configured before that landed can still hold e.g. a JSON object in a
|
||||
string-typed field. Every consumer that treats such a value as text blows up
|
||||
identically on every run (``escape_for_prompt`` -> ``re.sub`` ->
|
||||
"expected string or bytes-like object, got 'dict'"), so the bank's
|
||||
consolidation never recovers on its own (issue #3218).
|
||||
|
||||
String fields are JSON-encoded, which preserves the author's intent — the
|
||||
structure still reaches the prompt, as text. Anything else is dropped so the
|
||||
bank falls back to the tenant/global value rather than wedging.
|
||||
|
||||
``where`` labels the location in warnings; it is set when recursing into a
|
||||
retain strategy, whose overrides reach the same fields via ``apply_strategy``.
|
||||
"""
|
||||
field_types = _configurable_field_types()
|
||||
coerced: dict[str, Any] = {}
|
||||
for key, value in overrides.items():
|
||||
allowed = field_types.get(key)
|
||||
# None passes through: the caller has already dropped top-level tombstones,
|
||||
# and inside a retain strategy a null is a deliberate override to None.
|
||||
if allowed is None or value is None or _value_matches_type(value, allowed):
|
||||
coerced[key] = value
|
||||
continue
|
||||
if str in allowed:
|
||||
coerced[key] = json.dumps(value, ensure_ascii=False)
|
||||
logger.warning(
|
||||
f"Bank {bank_id} config field '{key}'{where} holds a {type(value).__name__} but is a string field; "
|
||||
f"using its JSON encoding. Re-save this field as a string to silence this warning."
|
||||
)
|
||||
else:
|
||||
logger.warning(
|
||||
f"Bank {bank_id} config field '{key}'{where} holds a {type(value).__name__} but must be "
|
||||
f"{_describe_types(allowed)}; ignoring the override and falling back to the server default."
|
||||
)
|
||||
|
||||
# Strategy overrides are spliced onto the resolved config by apply_strategy(),
|
||||
# so a bad value nested there wedges the bank just as a top-level one does.
|
||||
strategies = coerced.get("retain_strategies")
|
||||
if isinstance(strategies, dict):
|
||||
coerced["retain_strategies"] = {
|
||||
name: (
|
||||
_coerce_stored_bank_overrides(bank_id, strategy, where=f" in retain strategy {name!r}")
|
||||
if isinstance(strategy, dict)
|
||||
else strategy
|
||||
)
|
||||
for name, strategy in strategies.items()
|
||||
}
|
||||
return coerced
|
||||
|
||||
|
||||
_RECALL_BUDGET_FIXED_KEYS = (
|
||||
"recall_budget_fixed_low",
|
||||
"recall_budget_fixed_mid",
|
||||
|
||||
@@ -14,10 +14,7 @@ import subprocess
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from typing import IO
|
||||
from typing import IO
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -42,37 +39,28 @@ class IdleTimeoutMiddleware:
|
||||
self.app = app
|
||||
self.idle_timeout = idle_timeout
|
||||
self.last_activity = time.time()
|
||||
self._checker_task = None
|
||||
|
||||
async def __call__(self, scope, receive, send):
|
||||
# Update activity timestamp on each request
|
||||
self.last_activity = time.time()
|
||||
await self.app(scope, receive, send)
|
||||
|
||||
def start_idle_checker(self):
|
||||
"""Start the background task that checks for idle timeout."""
|
||||
self._checker_task = asyncio.create_task(self._check_idle())
|
||||
|
||||
async def _check_idle(self):
|
||||
"""Background task that exits the process after idle timeout."""
|
||||
# If idle_timeout is 0, don't auto-exit
|
||||
"""Exit the daemon after the configured period without requests."""
|
||||
if self.idle_timeout <= 0:
|
||||
return
|
||||
|
||||
while True:
|
||||
await asyncio.sleep(30) # Check every 30 seconds
|
||||
await asyncio.sleep(30)
|
||||
idle_time = time.time() - self.last_activity
|
||||
if idle_time > self.idle_timeout:
|
||||
logger.info(f"Idle timeout reached ({self.idle_timeout}s), shutting down daemon")
|
||||
# Give a moment for any in-flight requests
|
||||
await asyncio.sleep(1)
|
||||
# Send SIGTERM to ourselves to trigger graceful shutdown
|
||||
import signal
|
||||
|
||||
os.kill(os.getpid(), signal.SIGTERM)
|
||||
|
||||
|
||||
def _detach_popen_kwargs(log_handle: "IO[bytes]") -> dict:
|
||||
def _detach_popen_kwargs(log_handle: IO[bytes]) -> dict:
|
||||
"""Cross-platform kwargs to spawn a subprocess detached from the caller.
|
||||
|
||||
On POSIX, ``start_new_session=True`` calls ``setsid(2)`` so the child
|
||||
@@ -169,17 +157,3 @@ def daemonize():
|
||||
subprocess.Popen(cmd, env=env, **_detach_popen_kwargs(log_handle))
|
||||
|
||||
sys.exit(0)
|
||||
|
||||
|
||||
def check_daemon_running(port: int = DEFAULT_DAEMON_PORT) -> bool:
|
||||
"""Check if a daemon is running and responsive on the given port."""
|
||||
import socket
|
||||
|
||||
try:
|
||||
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
||||
sock.settimeout(1)
|
||||
result = sock.connect_ex(("127.0.0.1", port))
|
||||
sock.close()
|
||||
return result == 0
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
@@ -3,7 +3,7 @@ Memory Engine - Core implementation of the memory system.
|
||||
|
||||
This package contains all the implementation details of the memory engine:
|
||||
- MemoryEngine: Main class for memory operations
|
||||
- Utility modules: embedding_utils, link_utils, think_utils, bank_utils
|
||||
- Utility modules: embedding_utils, link_utils, bank_utils
|
||||
- Supporting modules: embeddings, cross_encoder, entity_resolver, etc.
|
||||
"""
|
||||
|
||||
|
||||
@@ -10,7 +10,7 @@ import asyncio
|
||||
import json
|
||||
import logging
|
||||
import uuid
|
||||
from collections.abc import Callable
|
||||
from collections.abc import Awaitable, Callable
|
||||
from contextlib import asynccontextmanager
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime, timezone
|
||||
@@ -19,6 +19,8 @@ from typing import Any
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from ..engine.db_utils import acquire_with_retry
|
||||
from ..models import RequestContext
|
||||
from .schema import fq_table_explicit
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -119,23 +121,60 @@ class AuditLogger:
|
||||
schema_getter: Callable[[], str],
|
||||
enabled: bool,
|
||||
allowed_actions: list[str],
|
||||
bank_enabled_resolver: Callable[[str, RequestContext | None], Awaitable[bool]] | None = None,
|
||||
) -> None:
|
||||
self._pool_getter = pool_getter
|
||||
self._schema_getter = schema_getter
|
||||
self._enabled = enabled
|
||||
self._allowed_actions: frozenset[str] | None = frozenset(allowed_actions) if allowed_actions else None
|
||||
# Resolves the hierarchical ``audit_log_enabled`` for one bank
|
||||
# (env -> tenant -> bank). None means "no per-bank resolution wired",
|
||||
# in which case the global value alone decides.
|
||||
self._bank_enabled_resolver = bank_enabled_resolver
|
||||
|
||||
def is_enabled(self, action: str) -> bool:
|
||||
"""Check if audit logging is enabled for this action."""
|
||||
if not self._enabled:
|
||||
def action_allowed(self, action: str) -> bool:
|
||||
"""Global action-allowlist check. Cheap, synchronous, bank-independent.
|
||||
|
||||
The allowlist is deployment-wide, so this is a valid pre-filter to skip
|
||||
work for actions that can never be audited. It deliberately does NOT
|
||||
consult the enabled flag: that is per-bank overridable, so a bank may
|
||||
turn auditing ON even when the deployment default is off.
|
||||
"""
|
||||
if self._allowed_actions is None:
|
||||
return True
|
||||
return action in self._allowed_actions
|
||||
|
||||
async def should_log(self, action: str, bank_id: str | None, context: RequestContext | None = None) -> bool:
|
||||
"""Full audit decision: action allowlist AND the bank's resolved switch.
|
||||
|
||||
``audit_log_enabled`` is hierarchical (env -> tenant -> bank), so the
|
||||
effective value depends on which bank the action targets. Falls back to
|
||||
the global value when there is no bank in scope or no resolver wired.
|
||||
"""
|
||||
if not self.action_allowed(action):
|
||||
return False
|
||||
if self._allowed_actions is not None:
|
||||
return action in self._allowed_actions
|
||||
return True
|
||||
if bank_id is None or self._bank_enabled_resolver is None:
|
||||
return self._enabled
|
||||
try:
|
||||
return await self._bank_enabled_resolver(bank_id, context)
|
||||
except Exception as e:
|
||||
# Never let a config-resolution failure break the request. Fall back
|
||||
# to the deployment default: a transient DB blip must not silently
|
||||
# create an audit gap for a bank meant to be audited. The tradeoff is
|
||||
# the opt-out direction — a bank that overrode to false under a
|
||||
# default-on deployment will be audited during the outage. We accept
|
||||
# that: a few extra audit rows during a DB blip is the safer failure
|
||||
# than dropping records that compliance may require.
|
||||
logger.warning(f"Audit config resolution failed for bank={bank_id}: {e}; using global default")
|
||||
return self._enabled
|
||||
|
||||
def log_fire_and_forget(self, entry: AuditEntry) -> None:
|
||||
"""Schedule an audit write as a background task."""
|
||||
if not self.is_enabled(entry.action):
|
||||
"""Schedule an audit write as a background task.
|
||||
|
||||
Assumes the caller already made the audit decision via ``should_log``;
|
||||
only the bank-independent allowlist is re-checked here.
|
||||
"""
|
||||
if not self.action_allowed(entry.action):
|
||||
return
|
||||
try:
|
||||
asyncio.create_task(self._safe_log(entry))
|
||||
@@ -150,8 +189,12 @@ class AuditLogger:
|
||||
logger.debug("Audit log skipped: pool not available")
|
||||
return
|
||||
try:
|
||||
schema = self._schema_getter()
|
||||
table = f"{schema}.audit_log"
|
||||
# fq_table_explicit qualifies per dialect: "schema".audit_log on
|
||||
# PostgreSQL, bare audit_log on Oracle (where the schema is set at the
|
||||
# session level). A raw f"{schema}.audit_log" produced public.audit_log
|
||||
# on Oracle, where "public" is a reserved word — every write failed
|
||||
# with ORA-00903 even though the table exists.
|
||||
table = fq_table_explicit("audit_log", self._schema_getter())
|
||||
async with acquire_with_retry(pool, max_retries=1) as conn:
|
||||
await conn.execute(
|
||||
f"""
|
||||
@@ -182,6 +225,7 @@ async def audit_context(
|
||||
bank_id: str | None = None,
|
||||
request: dict[str, Any] | None = None,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
context: RequestContext | None = None,
|
||||
):
|
||||
"""Async context manager that times the operation and writes audit on exit.
|
||||
|
||||
@@ -190,7 +234,7 @@ async def audit_context(
|
||||
result = await do_work()
|
||||
entry.response = result_dict
|
||||
"""
|
||||
if audit_logger is None or not audit_logger.is_enabled(action):
|
||||
if audit_logger is None or not await audit_logger.should_log(action, bank_id, context):
|
||||
entry = AuditEntry(action=action, transport=transport, bank_id=bank_id)
|
||||
yield entry
|
||||
return
|
||||
|
||||
@@ -13,6 +13,8 @@ but operators should opt in with that in mind.
|
||||
|
||||
from typing import Any
|
||||
|
||||
RERANKER_BANK_ID_HEADER = "X-Hindsight-Bank-Id"
|
||||
|
||||
|
||||
def apply_bank_attribution(request: dict[str, Any]) -> None:
|
||||
"""Tag ``request`` with ``user=<bank_id>`` for per-bank cost attribution.
|
||||
@@ -32,3 +34,14 @@ def apply_bank_attribution(request: dict[str, Any]) -> None:
|
||||
bank_id = get_current_bank_id()
|
||||
if bank_id:
|
||||
request["user"] = bank_id
|
||||
|
||||
|
||||
def reranker_bank_attribution_headers() -> dict[str, str]:
|
||||
"""Return the fixed per-bank header for trusted remote reranker endpoints."""
|
||||
from ..config import get_config
|
||||
from .memory_engine import get_current_bank_id
|
||||
|
||||
if not get_config().reranker_send_bank_as_header:
|
||||
return {}
|
||||
bank_id = get_current_bank_id()
|
||||
return {RERANKER_BANK_ID_HEADER: bank_id} if bank_id else {}
|
||||
|
||||
@@ -0,0 +1,70 @@
|
||||
"""Shared causal-link taxonomy.
|
||||
|
||||
Retain writes only the canonical relationship. Transfer import/export also
|
||||
preserves historical relationship types so existing banks keep their graph
|
||||
semantics without allowing new retain output to create those types.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
CANONICAL_CAUSAL_LINK_TYPE = "caused_by"
|
||||
LEGACY_CAUSAL_LINK_TYPE_NAMES = ("causes", "enables", "prevents")
|
||||
|
||||
CANONICAL_CAUSAL_LINK_TYPES = frozenset({CANONICAL_CAUSAL_LINK_TYPE})
|
||||
LEGACY_CAUSAL_LINK_TYPES = frozenset(LEGACY_CAUSAL_LINK_TYPE_NAMES)
|
||||
CAUSAL_LINK_TYPES = (CANONICAL_CAUSAL_LINK_TYPE, *LEGACY_CAUSAL_LINK_TYPE_NAMES)
|
||||
|
||||
DEFAULT_CAUSAL_LINK_WEIGHT = 1.0
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CausalLinkDescriptor:
|
||||
"""One causal edge, parked on the curation archive while an endpoint is invalidated.
|
||||
|
||||
Invalidation moves a fact out of ``memory_units``, so the FK cascade deletes
|
||||
its ``memory_links`` rows — and nothing could recreate a causal edge, which
|
||||
is extraction output rather than derived data. The descriptor is what the
|
||||
archive row stores so revert can rematerialize the edge (#2864).
|
||||
"""
|
||||
|
||||
from_unit_id: str
|
||||
to_unit_id: str
|
||||
link_type: str
|
||||
weight: float = DEFAULT_CAUSAL_LINK_WEIGHT
|
||||
|
||||
def as_json_dict(self) -> dict[str, Any]:
|
||||
"""Serializable form written to ``invalidated_memory_units.causal_links``.
|
||||
|
||||
The key names double as the column list of the ``jsonb_to_recordset``
|
||||
read in ``snapshot_causal_links`` — keep them in sync.
|
||||
"""
|
||||
return {
|
||||
"from_unit_id": self.from_unit_id,
|
||||
"to_unit_id": self.to_unit_id,
|
||||
"link_type": self.link_type,
|
||||
"weight": self.weight,
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_json_dict(cls, raw: Any) -> "CausalLinkDescriptor | None":
|
||||
"""Parse one stored descriptor, or None when it isn't a usable causal edge.
|
||||
|
||||
The archive column is plain JSON with no schema enforcement (a restore
|
||||
from an older backup, or a hand-edited row, can put anything there), and
|
||||
``memory_links`` has a ``link_type`` CHECK constraint — so an unusable
|
||||
entry is skipped rather than allowed to abort the whole revert.
|
||||
"""
|
||||
if not isinstance(raw, dict):
|
||||
return None
|
||||
from_unit_id = raw.get("from_unit_id")
|
||||
to_unit_id = raw.get("to_unit_id")
|
||||
link_type = raw.get("link_type")
|
||||
if not from_unit_id or not to_unit_id or link_type not in CAUSAL_LINK_TYPES:
|
||||
return None
|
||||
return cls(
|
||||
from_unit_id=str(from_unit_id),
|
||||
to_unit_id=str(to_unit_id),
|
||||
link_type=str(link_type),
|
||||
weight=float(raw.get("weight") or DEFAULT_CAUSAL_LINK_WEIGHT),
|
||||
)
|
||||
@@ -109,6 +109,11 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
|
||||
end.replace(hour=23, minute=59, second=59, microsecond=999999),
|
||||
)
|
||||
|
||||
def safe_constraint(start: datetime | None, end: datetime | None) -> DateRange | NoTemporalConstraintSentinel:
|
||||
if start is None or end is None:
|
||||
return NO_TEMPORAL_CONSTRAINT
|
||||
return constraint(start, end)
|
||||
|
||||
def subtract_months(months: int) -> datetime:
|
||||
month_index = reference_date.month - months - 1
|
||||
year = reference_date.year + month_index // 12
|
||||
@@ -126,11 +131,21 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
|
||||
day = min(base_date.day, calendar.monthrange(year, month)[1])
|
||||
return base_date.replace(year=year, month=month, day=day)
|
||||
|
||||
def add_years(base_date: datetime, years: int) -> datetime:
|
||||
def add_years(base_date: datetime, years: int) -> datetime | None:
|
||||
year = base_date.year + years
|
||||
if year < datetime.min.year or year > datetime.max.year:
|
||||
return None
|
||||
day = min(base_date.day, calendar.monthrange(year, base_date.month)[1])
|
||||
return base_date.replace(year=year, day=day)
|
||||
|
||||
def add_days(base_date: datetime | None, days: int) -> datetime | None:
|
||||
if base_date is None:
|
||||
return None
|
||||
try:
|
||||
return base_date + timedelta(days=days)
|
||||
except OverflowError:
|
||||
return None
|
||||
|
||||
def has_chinese_temporal_context(match: re.Match[str]) -> bool:
|
||||
if match.end() >= len(query):
|
||||
return True
|
||||
@@ -438,6 +453,11 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
|
||||
return NO_TEMPORAL_CONSTRAINT
|
||||
return constraint(start, reference_date)
|
||||
|
||||
def safe_since_constraint(start: datetime | None) -> DateRange | NoTemporalConstraintSentinel:
|
||||
if start is None:
|
||||
return NO_TEMPORAL_CONSTRAINT
|
||||
return since_constraint(start)
|
||||
|
||||
def since_from_period(
|
||||
period: DateRange | None,
|
||||
) -> DateRange | NoTemporalConstraintSentinel | None:
|
||||
@@ -450,7 +470,7 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
|
||||
return None
|
||||
return since_constraint(day)
|
||||
|
||||
def relative_offset_datetime(amount: int, unit: str, direction: int) -> datetime:
|
||||
def relative_offset_datetime(amount: int, unit: str, direction: int) -> datetime | None:
|
||||
if unit in ("天", "日"):
|
||||
return reference_date + timedelta(days=direction * amount)
|
||||
if unit in ("周", "星期", "礼拜"):
|
||||
@@ -459,15 +479,15 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
|
||||
return add_months(reference_date, direction * amount)
|
||||
return add_years(reference_date, direction * amount)
|
||||
|
||||
def point_constraint_at_offset(amount: int, unit: str, direction: int) -> DateRange:
|
||||
def point_constraint_at_offset(amount: int, unit: str, direction: int) -> DateRange | NoTemporalConstraintSentinel:
|
||||
d = relative_offset_datetime(amount, unit, direction)
|
||||
return constraint(d, d)
|
||||
return safe_constraint(d, d)
|
||||
|
||||
def window_to_reference(amount: int, unit: str) -> DateRange:
|
||||
return constraint(relative_offset_datetime(amount, unit, -1), reference_date)
|
||||
def window_to_reference(amount: int, unit: str) -> DateRange | NoTemporalConstraintSentinel:
|
||||
return safe_constraint(relative_offset_datetime(amount, unit, -1), reference_date)
|
||||
|
||||
def window_from_reference(amount: int, unit: str) -> DateRange:
|
||||
return constraint(reference_date, relative_offset_datetime(amount, unit, 1))
|
||||
def window_from_reference(amount: int, unit: str) -> DateRange | NoTemporalConstraintSentinel:
|
||||
return safe_constraint(reference_date, relative_offset_datetime(amount, unit, 1))
|
||||
|
||||
# Chinese rule guide
|
||||
#
|
||||
@@ -781,8 +801,8 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
|
||||
if relative_year_fixed_day_since_match:
|
||||
year = relative_year_number(relative_year_fixed_day_since_match.group(1))
|
||||
base = add_years(reference_date, year - reference_date.year)
|
||||
d = base + timedelta(days=fixed_day_offset(relative_year_fixed_day_since_match.group(2)))
|
||||
return since_constraint(d)
|
||||
d = add_days(base, fixed_day_offset(relative_year_fixed_day_since_match.group(2)))
|
||||
return safe_since_constraint(d)
|
||||
|
||||
fixed_day_since_match = chinese_search(
|
||||
rf"(大大后天|大后天|后天|明天|明日|今天|今日|本日|当日|当天|昨天|昨日|大大前天|大前天|前天)"
|
||||
@@ -799,7 +819,7 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
|
||||
amount = parse_chinese_number(exact_relative_since_match.group(1))
|
||||
unit = exact_relative_since_match.group(2)
|
||||
if amount is not None:
|
||||
return since_constraint(relative_offset_datetime(amount, unit, -1))
|
||||
return safe_since_constraint(relative_offset_datetime(amount, unit, -1))
|
||||
|
||||
weekend_since_match = chinese_search(
|
||||
rf"(?<![上下大小每个各隔])"
|
||||
@@ -899,8 +919,8 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
|
||||
if relative_year_daypart_since_match:
|
||||
year = relative_year_number(relative_year_daypart_since_match.group(1))
|
||||
base = add_years(reference_date, year - reference_date.year)
|
||||
d = base + timedelta(days=daypart_day_offset(relative_year_daypart_since_match.group(2)))
|
||||
return since_constraint(d)
|
||||
d = add_days(base, daypart_day_offset(relative_year_daypart_since_match.group(2)))
|
||||
return safe_since_constraint(d)
|
||||
|
||||
daypart_since_match = chinese_search(
|
||||
rf"(昨晚|昨夜|前晚|前夜|今晚|今早|今晨|明早|明晚|明夜){chinese_since_suffix_pattern}"
|
||||
@@ -915,17 +935,17 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
|
||||
if relative_year_daypart_match:
|
||||
year = relative_year_number(relative_year_daypart_match.group(1))
|
||||
base = add_years(reference_date, year - reference_date.year)
|
||||
d = base + timedelta(days=daypart_day_offset(relative_year_daypart_match.group(2)))
|
||||
return constraint(d, d)
|
||||
d = add_days(base, daypart_day_offset(relative_year_daypart_match.group(2)))
|
||||
return safe_constraint(d, d)
|
||||
|
||||
# Day-part abbreviations still resolve only to date granularity.
|
||||
if chinese_search(r"昨晚|昨夜"):
|
||||
d = reference_date + timedelta(days=daypart_day_offset("昨晚"))
|
||||
return constraint(d, d)
|
||||
d = add_days(reference_date, daypart_day_offset("昨晚"))
|
||||
return safe_constraint(d, d)
|
||||
|
||||
if chinese_search(r"前晚|前夜"):
|
||||
d = reference_date + timedelta(days=daypart_day_offset("前晚"))
|
||||
return constraint(d, d)
|
||||
d = add_days(reference_date, daypart_day_offset("前晚"))
|
||||
return safe_constraint(d, d)
|
||||
|
||||
if chinese_search(r"今晚|今早|今晨"):
|
||||
return constraint(reference_date, reference_date)
|
||||
@@ -941,8 +961,8 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
|
||||
if relative_year_fixed_day_match:
|
||||
year = relative_year_number(relative_year_fixed_day_match.group(1))
|
||||
base = add_years(reference_date, year - reference_date.year)
|
||||
d = base + timedelta(days=fixed_day_offset(relative_year_fixed_day_match.group(2)))
|
||||
return constraint(d, d)
|
||||
d = add_days(base, fixed_day_offset(relative_year_fixed_day_match.group(2)))
|
||||
return safe_constraint(d, d)
|
||||
|
||||
if chinese_search(r"昨天|昨日"):
|
||||
d = reference_date - timedelta(days=1)
|
||||
@@ -1085,7 +1105,7 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
|
||||
end_amount = parse_chinese_number(amount_text[-1])
|
||||
unit = adjacent_fuzzy_future_match.group(2)
|
||||
if start_amount is not None and end_amount is not None:
|
||||
return constraint(
|
||||
return safe_constraint(
|
||||
relative_offset_datetime(start_amount, unit, 1),
|
||||
relative_offset_datetime(end_amount, unit, 1),
|
||||
)
|
||||
@@ -1093,7 +1113,7 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
|
||||
few_future_match = chinese_search(rf"[几数]个?(天|日|周|星期|礼拜|月|年){chinese_relative_future_suffix_pattern}")
|
||||
if few_future_match:
|
||||
unit = few_future_match.group(1)
|
||||
return constraint(relative_offset_datetime(2, unit, 1), relative_offset_datetime(5, unit, 1))
|
||||
return safe_constraint(relative_offset_datetime(2, unit, 1), relative_offset_datetime(5, unit, 1))
|
||||
|
||||
exact_future_match = chinese_search(
|
||||
rf"(?<![{_CHINESE_NUMERAL_PREFIX_CHARS}])([0-9]+|[{_CHINESE_NUMERAL_CHARS}]+)个?(天|日|周|星期|礼拜|月|年){chinese_relative_future_suffix_pattern}"
|
||||
@@ -1113,7 +1133,7 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
|
||||
second_amount = parse_chinese_number(adjacent_fuzzy_past_match.group(2))
|
||||
unit = adjacent_fuzzy_past_match.group(3)
|
||||
if first_amount is not None and second_amount is not None and second_amount == first_amount + 1:
|
||||
return constraint(
|
||||
return safe_constraint(
|
||||
relative_offset_datetime(second_amount, unit, -1),
|
||||
relative_offset_datetime(first_amount, unit, -1),
|
||||
)
|
||||
@@ -1144,7 +1164,7 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
|
||||
return constraint(reference_date - timedelta(days=150), reference_date - timedelta(days=60))
|
||||
|
||||
if chinese_search(r"一两年前|[两二]三年前|三两年前"):
|
||||
return constraint(add_years(reference_date, -3), add_years(reference_date, -1))
|
||||
return safe_constraint(add_years(reference_date, -3), add_years(reference_date, -1))
|
||||
|
||||
rolling_this_adjacent_match = chinese_search(
|
||||
r"这(一两|[两二]三|三两|三四|四五|五六|六七|七八|八九|九十)个?(天|日|周|星期|礼拜|月|年)"
|
||||
@@ -1154,7 +1174,7 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
|
||||
end_amount = 3 if amount_text in ("一两", "三两") else parse_chinese_number(amount_text[-1])
|
||||
unit = rolling_this_adjacent_match.group(2)
|
||||
if end_amount is not None:
|
||||
return constraint(relative_offset_datetime(end_amount, unit, -1), reference_date)
|
||||
return safe_constraint(relative_offset_datetime(end_amount, unit, -1), reference_date)
|
||||
|
||||
rolling_this_count_match = chinese_search(rf"这([0-9]+|[{_CHINESE_NUMERAL_CHARS}]+)个?(天|日|周|星期|礼拜|月|年)")
|
||||
if rolling_this_count_match:
|
||||
@@ -1191,7 +1211,7 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
|
||||
end_amount = 3 if amount_text in ("一两", "三两") else parse_chinese_number(amount_text[-1])
|
||||
unit = rolling_past_adjacent_match.group(3)
|
||||
if end_amount is not None:
|
||||
return constraint(relative_offset_datetime(end_amount, unit, -1), reference_date)
|
||||
return safe_constraint(relative_offset_datetime(end_amount, unit, -1), reference_date)
|
||||
|
||||
rolling_past_few_match = chinese_search(r"(过去|近|最近)几个?(天|日|周|星期|礼拜|月|年)")
|
||||
if rolling_past_few_match:
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -17,6 +17,20 @@ _MISSION_PRIORITY_NOTE = (
|
||||
"DECISION GUIDE, or OUTPUT FORMAT below, the MISSION takes priority."
|
||||
)
|
||||
|
||||
# Default language rule — used only when HINDSIGHT_API_LLM_OUTPUT_LANGUAGE is
|
||||
# unset. Without it the whole prompt is English and multilingual models drift:
|
||||
# Chinese source facts intermittently produce English observations. Retain's
|
||||
# fact extraction carries the equivalent rule (see _BASE_FACT_EXTRACTION_PROMPT),
|
||||
# so this makes "preserve the source language" the pipeline-wide default. When an
|
||||
# output language IS configured, this section is omitted and
|
||||
# output_language_directive() takes over — the two must never both be present or
|
||||
# they contradict each other.
|
||||
_DEFAULT_LANGUAGE_RULE = """## LANGUAGE
|
||||
|
||||
Write every observation in the language of its own source facts — never translate them. Per observation, not per batch: when one merges facts of several languages, the majority wins. Proper nouns, identifiers, and units stay verbatim.
|
||||
|
||||
When an existing observation is written in a different language from the new facts updating it, do NOT edit its wording in place — that is what produces an English sentence with a Chinese detail bolted on. Discard the old phrasing and compose the merged observation from scratch in the new facts' language."""
|
||||
|
||||
_PROCESSING_RULES = """## PROCESSING RULES
|
||||
|
||||
1. PREFER UPDATE OVER CREATE (when there is something to merge with): if new facts describe the same canonical event, statement, decision, claim, or recurring pattern already covered by an existing observation, UPDATE that observation and attach the new facts as evidence. Do NOT create a near-duplicate sibling. One canonical observation with many source facts is always better than many siblings with one source fact each. Merge aggressively on: same named event, same diagnostic finding, same architectural decision, same recurring claim. **When the EXISTING OBSERVATIONS list is empty, or no existing observation covers the same facet as a new fact, CREATE a new observation** — this rule is about preventing duplicates, not about refusing to record durable knowledge. CREATE is the correct default for any structurally distinct event, claim, or pattern that has no existing match.
|
||||
@@ -37,19 +51,36 @@ _PROCESSING_RULES = """## PROCESSING RULES
|
||||
|
||||
9. KEEP DISTINCT TOPICS DISTINCT: do not merge observations about different people, entities, or unrelated topics. Merging is for the same canonical fact recurring — not for related-but-distinct claims."""
|
||||
|
||||
# Field-by-field definitions of the input shape used by the cached system
|
||||
# prefix. The call site runs .format(), so these strings must contain no braces.
|
||||
_FACT_FIELDS = """One per line, formatted as `[uuid] fact text (temporal fields)`:
|
||||
- `[uuid]`: the fact's identifier — copy it verbatim into `source_fact_ids`
|
||||
- `occurred_start` / `occurred_end`: when the described event happened. This can be long before the fact was stated — a fact recorded today may describe a 2019 event.
|
||||
- `mentioned_at`: when the source material that states this fact was written. This is the fact's recency: how up to date the statement is, NOT when it was added to memory. A fact taken from an old document keeps its old `mentioned_at` even if it was only just processed."""
|
||||
|
||||
_OBSERVATION_FIELDS = """- `id`: unique identifier — copy this exactly when issuing an UPDATE or DELETE
|
||||
- `text`: the observation content
|
||||
- `proof_count`: how many source facts this observation has already merged
|
||||
- `occurred_start` / `occurred_end`: the span of the events behind the observation — earliest start and latest end across its source facts
|
||||
- `mentioned_at`: the latest of the `mentioned_at` values of its source facts — the most recent point at which this observation was stated
|
||||
- `source_memories`: the supporting facts behind this observation. May be partial or absent for large observations — the count above remains the true total. Each entry carries the same `text` and temporal fields as a new fact, plus:
|
||||
- `context`: optional surrounding context for that fact"""
|
||||
|
||||
# Stable description of the input shape. For the cached split path this lives in
|
||||
# the system prefix (build_consolidation_system_prompt) so it is not re-sent on
|
||||
# every batch; the per-batch user message then carries only the actual data.
|
||||
_INPUT_FORMAT_NOTE = """## INPUT FORMAT
|
||||
_INPUT_FORMAT_NOTE = f"""## INPUT FORMAT
|
||||
|
||||
Each request provides new facts and existing observations:
|
||||
- New facts: one per line, each prefixed with its `[uuid]`, followed by the fact text and optional temporal fields.
|
||||
- Existing observations: a JSON array pooled from recalls across the new facts. Each entry has:
|
||||
- `id`: unique identifier — copy this exactly when issuing an UPDATE or DELETE
|
||||
- `text`: the observation content
|
||||
- `proof_count`: number of supporting memories
|
||||
- `occurred_start` / `occurred_end`: temporal range of source facts
|
||||
- `source_memories`: array of supporting facts with their text and dates"""
|
||||
Each request provides new facts and existing observations. Every temporal field is optional and is omitted when unknown.
|
||||
|
||||
### New facts
|
||||
|
||||
{_FACT_FIELDS}
|
||||
|
||||
### Existing observations
|
||||
|
||||
A JSON array pooled from recalls across the new facts. Each entry has:
|
||||
{_OBSERVATION_FIELDS}"""
|
||||
|
||||
# Per-batch data section for the cached split path — the stable format
|
||||
# explanation above is omitted here (it lives in the cached prefix); only the
|
||||
@@ -64,24 +95,6 @@ _SPLIT_INPUT_SECTION = """## INPUT
|
||||
|
||||
{observations_text}"""
|
||||
|
||||
# Data section — format placeholders {facts_text} and {observations_text} are substituted at call time
|
||||
_INPUT_SECTION = """## INPUT
|
||||
|
||||
### New facts
|
||||
|
||||
{facts_text}
|
||||
|
||||
### Existing observations
|
||||
|
||||
JSON array, pooled from recalls across all new facts above. Each entry has:
|
||||
- `id`: unique identifier — copy this exactly when issuing an UPDATE or DELETE
|
||||
- `text`: the observation content
|
||||
- `proof_count`: number of supporting memories
|
||||
- `occurred_start` / `occurred_end`: temporal range of source facts
|
||||
- `source_memories`: array of supporting facts with their text and dates
|
||||
|
||||
{observations_text}"""
|
||||
|
||||
_DECISION_GUIDE = """## DECISION GUIDE
|
||||
|
||||
- **Same canonical event, decision, claim, or facet as an existing observation → UPDATE** (use `observation_id` + new `source_fact_ids`).
|
||||
@@ -142,39 +155,6 @@ Expected output (UPDATE for the state change; CREATE for the unrelated work-hour
|
||||
- Return `{{"creates": [], "updates": [], "deletes": []}}` if nothing durable is found."""
|
||||
|
||||
|
||||
def build_batch_consolidation_prompt(
|
||||
observations_mission: str | None = None,
|
||||
observation_capacity_note: str | None = None,
|
||||
llm_output_language: str | None = None,
|
||||
) -> str:
|
||||
"""
|
||||
Build the consolidation prompt for batch mode (multiple facts per LLM call).
|
||||
|
||||
The mission defines *what* to track (customisable per bank) and takes
|
||||
priority over the built-in processing rules when the two conflict.
|
||||
Processing rules, decision guide, and output format are always present.
|
||||
When ``llm_output_language`` is set, observations are emitted in that
|
||||
language.
|
||||
"""
|
||||
mission = escape_for_prompt(observations_mission or _DEFAULT_MISSION)
|
||||
|
||||
capacity_section = ""
|
||||
if observation_capacity_note:
|
||||
capacity_section = f"\n\n## CAPACITY CONSTRAINT\n\n{escape_for_prompt(observation_capacity_note)}"
|
||||
|
||||
return (
|
||||
"You are a memory consolidation system. Synthesize new facts into "
|
||||
"observations, merging with existing observations when appropriate.\n\n"
|
||||
f"## MISSION\n\n{mission}\n\n"
|
||||
f"{_MISSION_PRIORITY_NOTE}"
|
||||
f"{capacity_section}\n\n"
|
||||
f"{_PROCESSING_RULES}\n\n"
|
||||
f"{_INPUT_SECTION}\n\n"
|
||||
f"{_DECISION_GUIDE}\n\n"
|
||||
f"{_OUTPUT_SECTION}" + output_language_directive(llm_output_language)
|
||||
)
|
||||
|
||||
|
||||
def build_consolidation_system_prompt(
|
||||
llm_output_language: str | None = None,
|
||||
) -> str:
|
||||
@@ -189,11 +169,17 @@ def build_consolidation_system_prompt(
|
||||
bank and a single CachedContent serves them all. Returns final text
|
||||
(brace-escaped examples already unescaped) for verbatim use as system message
|
||||
and cached prefix.
|
||||
|
||||
``llm_output_language`` picks between two mutually exclusive language rules:
|
||||
unset keeps each observation in the language of its own source facts (the
|
||||
default), set forces every observation into that one configured language.
|
||||
"""
|
||||
language_section = "" if llm_output_language else f"{_DEFAULT_LANGUAGE_RULE}\n\n"
|
||||
template = (
|
||||
"You are a memory consolidation system. Synthesize new facts into "
|
||||
"observations, merging with existing observations when appropriate.\n\n"
|
||||
f"{_MISSION_PRIORITY_NOTE}\n\n"
|
||||
f"{language_section}"
|
||||
f"{_PROCESSING_RULES}\n\n"
|
||||
f"{_INPUT_FORMAT_NOTE}\n\n"
|
||||
f"{_DECISION_GUIDE}\n\n"
|
||||
|
||||
@@ -8,10 +8,10 @@ Configuration via environment variables - see hindsight_api.config for all env v
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
import warnings
|
||||
from abc import ABC, abstractmethod
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
@@ -20,7 +20,6 @@ from ..config import (
|
||||
DEFAULT_RERANKER_ALIBABA_MODEL,
|
||||
DEFAULT_RERANKER_COHERE_MODEL,
|
||||
DEFAULT_RERANKER_FLASHRANK_CACHE_DIR,
|
||||
DEFAULT_RERANKER_FLASHRANK_CPU_MEM_ARENA,
|
||||
DEFAULT_RERANKER_FLASHRANK_MODEL,
|
||||
DEFAULT_RERANKER_GOOGLE_MODEL,
|
||||
DEFAULT_RERANKER_LITELLM_MAX_TOKENS_PER_DOC,
|
||||
@@ -34,58 +33,19 @@ from ..config import (
|
||||
DEFAULT_RERANKER_TEI_MAX_CONCURRENT,
|
||||
DEFAULT_RERANKER_ZEROENTROPY_MODEL,
|
||||
DEFAULT_ZEROENTROPY_BASE_URL,
|
||||
ENV_RERANKER_ALIBABA_API_KEY,
|
||||
ENV_RERANKER_COHERE_API_KEY,
|
||||
ENV_RERANKER_FLASHRANK_CACHE_DIR,
|
||||
ENV_RERANKER_FLASHRANK_CPU_MEM_ARENA,
|
||||
ENV_RERANKER_FLASHRANK_MODEL,
|
||||
ENV_RERANKER_GOOGLE_PROJECT_ID,
|
||||
ENV_RERANKER_PROVIDER,
|
||||
ENV_RERANKER_SILICONFLOW_API_KEY,
|
||||
ENV_RERANKER_TEI_URL,
|
||||
ENV_RERANKER_ZEROENTROPY_API_KEY,
|
||||
RerankerMemberConfig,
|
||||
)
|
||||
from .bank_attribution import reranker_bank_attribution_headers
|
||||
from .local_device import (
|
||||
release_local_inference_memory,
|
||||
resolve_model_device_type,
|
||||
select_local_device,
|
||||
)
|
||||
from .tei_retry import tei_retry_delay
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _resolve_malloc_trim():
|
||||
"""Return a callable that asks glibc to release freed heap pages to the OS.
|
||||
|
||||
Local CPU rerankers (FlashRank/ONNX, SentenceTransformers/torch) allocate
|
||||
large transient numpy/tensor buffers per call. On Linux glibc, those pages
|
||||
are freed at the Python level but kept by the allocator as a high-water
|
||||
mark — RSS grows monotonically across many recalls (see issue #1717).
|
||||
Calling `malloc_trim(0)` after each batch returns those pages to the OS.
|
||||
|
||||
Resolved once at import; returns a no-op on non-glibc platforms (macOS,
|
||||
musl, Windows) where the call is unavailable or unnecessary.
|
||||
"""
|
||||
import sys
|
||||
|
||||
if sys.platform != "linux":
|
||||
return lambda: None
|
||||
|
||||
import ctypes
|
||||
import ctypes.util
|
||||
|
||||
libc_path = ctypes.util.find_library("c")
|
||||
if libc_path is None:
|
||||
return lambda: None
|
||||
try:
|
||||
libc = ctypes.CDLL(libc_path)
|
||||
trim = libc.malloc_trim
|
||||
except (OSError, AttributeError):
|
||||
# Not glibc (musl has no malloc_trim) or libc lookup failed.
|
||||
return lambda: None
|
||||
trim.argtypes = [ctypes.c_size_t]
|
||||
trim.restype = ctypes.c_int
|
||||
return lambda: trim(0)
|
||||
|
||||
|
||||
_malloc_trim = _resolve_malloc_trim()
|
||||
|
||||
|
||||
class CrossEncoderModel(ABC):
|
||||
"""
|
||||
Abstract base class for cross-encoder reranking.
|
||||
@@ -99,6 +59,15 @@ class CrossEncoderModel(ABC):
|
||||
"""Return a human-readable name for this provider (e.g., 'local', 'tei')."""
|
||||
pass
|
||||
|
||||
@property
|
||||
def blocking_init(self) -> bool:
|
||||
"""Whether ``initialize()`` blocks the event loop (loads a model in-process).
|
||||
|
||||
Callers run those in a thread pool. Remote providers leave this False, and
|
||||
so does :class:`MultiCrossEncoder` — it offloads its own members.
|
||||
"""
|
||||
return False
|
||||
|
||||
@abstractmethod
|
||||
async def initialize(self) -> None:
|
||||
"""
|
||||
@@ -150,6 +119,7 @@ class LocalSTCrossEncoder(CrossEncoderModel):
|
||||
fp16: bool = False,
|
||||
bucket_batching: bool = False,
|
||||
batch_size: int = DEFAULT_RERANKER_LOCAL_BATCH_SIZE,
|
||||
allow_mps: bool = False,
|
||||
):
|
||||
"""
|
||||
Initialize local SentenceTransformers cross-encoder.
|
||||
@@ -171,6 +141,9 @@ class LocalSTCrossEncoder(CrossEncoderModel):
|
||||
Default: False (opt-in via env var).
|
||||
batch_size: Batch size for predict() calls. Optimal values vary by
|
||||
hardware and model (MPS: 32, CUDA: 128+). Default: 32.
|
||||
allow_mps: Opt in to the Apple Silicon MPS GPU. Disabled by default
|
||||
because MPS leaks memory under variable-length workloads
|
||||
(see engine/local_device.py). Default: False
|
||||
"""
|
||||
self.model_name = model_name or DEFAULT_RERANKER_LOCAL_MODEL
|
||||
self.force_cpu = force_cpu
|
||||
@@ -178,13 +151,19 @@ class LocalSTCrossEncoder(CrossEncoderModel):
|
||||
self.fp16 = fp16
|
||||
self.bucket_batching = bucket_batching
|
||||
self.batch_size = batch_size
|
||||
self.allow_mps = allow_mps
|
||||
self._model = None
|
||||
self._device_type: str = "cpu"
|
||||
LocalSTCrossEncoder._max_concurrent = max_concurrent
|
||||
|
||||
@property
|
||||
def provider_name(self) -> str:
|
||||
return "local"
|
||||
|
||||
@property
|
||||
def blocking_init(self) -> bool:
|
||||
return True
|
||||
|
||||
async def initialize(self) -> None:
|
||||
"""Load the cross-encoder model and initialize the executor."""
|
||||
if self._model is not None:
|
||||
@@ -200,33 +179,13 @@ class LocalSTCrossEncoder(CrossEncoderModel):
|
||||
|
||||
logger.info(f"Reranker: initializing local provider with model {self.model_name}")
|
||||
|
||||
# Determine device based on hardware availability.
|
||||
# We always set low_cpu_mem_usage=False to prevent lazy loading (meta tensors)
|
||||
# which can cause issues when accelerate is installed but no GPU is available.
|
||||
# Determine device based on hardware availability. We always set
|
||||
# low_cpu_mem_usage=False to prevent lazy loading (meta tensors) which can
|
||||
# cause issues when accelerate is installed but no GPU is available.
|
||||
# Note: We do NOT use device_map because CrossEncoder internally calls .to(device)
|
||||
# after loading, which conflicts with accelerate's device_map handling.
|
||||
import torch
|
||||
|
||||
# Force CPU mode if configured (used in daemon mode to avoid MPS/XPC issues on macOS)
|
||||
if self.force_cpu:
|
||||
device = "cpu"
|
||||
logger.info("Reranker: forcing CPU mode (HINDSIGHT_API_RERANKER_LOCAL_FORCE_CPU=1)")
|
||||
else:
|
||||
# Check for GPU (CUDA), Apple Silicon (MPS), or Intel XPU
|
||||
# Wrap in try-except to gracefully handle any device detection issues
|
||||
# (e.g., in CI environments or when PyTorch is built without GPU support)
|
||||
device = "cpu" # Default to CPU
|
||||
try:
|
||||
has_gpu = torch.cuda.is_available() or (
|
||||
hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
|
||||
)
|
||||
# Intel Arc XPU support — torch.xpu is available when the XPU build is loaded
|
||||
if not has_gpu and hasattr(torch, "xpu"):
|
||||
has_gpu = torch.xpu.is_available()
|
||||
if has_gpu:
|
||||
device = None # Let sentence-transformers auto-detect GPU/MPS/XPU
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to detect GPU/MPS/XPU, falling back to CPU: {e}")
|
||||
# MPS is opt-in (allow_mps) — see engine/local_device.py for why.
|
||||
device = select_local_device(self.force_cpu, self.allow_mps)
|
||||
|
||||
# Patch transformers 5.x compatibility for models using XLM-RoBERTa
|
||||
# (e.g., jina-reranker-v2-base-multilingual). transformers 5.x removed
|
||||
@@ -270,9 +229,11 @@ class LocalSTCrossEncoder(CrossEncoderModel):
|
||||
# Restore original logging level
|
||||
transformers_logger.setLevel(original_level)
|
||||
|
||||
self._device_type = resolve_model_device_type(self._model)
|
||||
|
||||
# FP16 inference: convert model weights to half precision.
|
||||
# Empirically validated: 27-36% faster on MPS, quality-identical (20/20 overlap).
|
||||
if self.fp16 and device != "cpu":
|
||||
if self.fp16 and self._device_type != "cpu":
|
||||
self._model.model.half()
|
||||
logger.info("Reranker: FP16 inference enabled")
|
||||
|
||||
@@ -315,7 +276,7 @@ class LocalSTCrossEncoder(CrossEncoderModel):
|
||||
scores = self._model.predict(pairs, batch_size=self.batch_size, show_progress_bar=False)
|
||||
return scores.tolist() if hasattr(scores, "tolist") else list(scores)
|
||||
finally:
|
||||
_malloc_trim()
|
||||
release_local_inference_memory(self._device_type)
|
||||
|
||||
async def predict(self, pairs: list[tuple[str, str]]) -> list[float]:
|
||||
"""
|
||||
@@ -431,14 +392,20 @@ class RemoteTEICrossEncoder(CrossEncoderModel):
|
||||
await asyncio.sleep(delay)
|
||||
delay *= 2 # Exponential backoff
|
||||
except httpx.HTTPStatusError as e:
|
||||
# Retry on 5xx server errors
|
||||
if e.response.status_code >= 500 and attempt < self.max_retries:
|
||||
# TEI uses 429 as normal overload backpressure. Retry it with
|
||||
# the same bounded budget as transient server errors.
|
||||
if (e.response.status_code == 429 or e.response.status_code >= 500) and attempt < self.max_retries:
|
||||
last_error = e
|
||||
logger.warning(
|
||||
f"TEI server error (attempt {attempt + 1}/{self.max_retries + 1}): {e}. "
|
||||
f"Retrying in {delay}s..."
|
||||
sleep_delay = tei_retry_delay(
|
||||
e.response,
|
||||
delay,
|
||||
request_timeout=self.timeout,
|
||||
)
|
||||
await asyncio.sleep(delay)
|
||||
logger.warning(
|
||||
f"TEI transient error (attempt {attempt + 1}/{self.max_retries + 1}): {e}. "
|
||||
f"Retrying in {sleep_delay:.2f}s..."
|
||||
)
|
||||
await asyncio.sleep(sleep_delay)
|
||||
delay *= 2
|
||||
else:
|
||||
raise
|
||||
@@ -484,6 +451,7 @@ class RemoteTEICrossEncoder(CrossEncoderModel):
|
||||
semaphore,
|
||||
"POST",
|
||||
f"{self.base_url}/rerank",
|
||||
headers=reranker_bank_attribution_headers(),
|
||||
json={
|
||||
"query": query,
|
||||
"texts": texts,
|
||||
@@ -624,7 +592,11 @@ class _CohereCompatibleRerankClient:
|
||||
if self.include_top_n:
|
||||
body["top_n"] = len(texts)
|
||||
|
||||
response = await self._async_client.post(self.rerank_url, json=body)
|
||||
response = await self._async_client.post(
|
||||
self.rerank_url,
|
||||
headers=reranker_bank_attribution_headers(),
|
||||
json=body,
|
||||
)
|
||||
response.raise_for_status()
|
||||
result = response.json()
|
||||
|
||||
@@ -919,6 +891,7 @@ class FlashRankCrossEncoder(CrossEncoderModel):
|
||||
self.max_length = max_length
|
||||
self.cpu_mem_arena = cpu_mem_arena
|
||||
self._ranker = None
|
||||
self._device_type: str = "cpu" # FlashRank runs on CPU via ONNX Runtime
|
||||
FlashRankCrossEncoder._max_concurrent = max_concurrent
|
||||
|
||||
@property
|
||||
@@ -990,11 +963,11 @@ class FlashRankCrossEncoder(CrossEncoderModel):
|
||||
|
||||
def _predict_sync(self, pairs: list[tuple[str, str]]) -> list[float]:
|
||||
"""Synchronous predict - processes each query group."""
|
||||
from flashrank import RerankRequest
|
||||
|
||||
if not pairs:
|
||||
return []
|
||||
|
||||
from flashrank import RerankRequest
|
||||
|
||||
try:
|
||||
# Group pairs by query
|
||||
query_groups: dict[str, list[tuple[int, str]]] = {}
|
||||
@@ -1023,7 +996,7 @@ class FlashRankCrossEncoder(CrossEncoderModel):
|
||||
|
||||
return all_scores
|
||||
finally:
|
||||
_malloc_trim()
|
||||
release_local_inference_memory(self._device_type)
|
||||
|
||||
async def predict(self, pairs: list[tuple[str, str]]) -> list[float]:
|
||||
"""
|
||||
@@ -1151,6 +1124,7 @@ class LiteLLMCrossEncoder(CrossEncoderModel):
|
||||
# LiteLLM /rerank follows Cohere API format
|
||||
response = await self._async_client.post(
|
||||
f"{self.api_base}/rerank",
|
||||
headers=reranker_bank_attribution_headers(),
|
||||
json={
|
||||
"model": self.model,
|
||||
"query": query,
|
||||
@@ -1269,10 +1243,11 @@ class LiteLLMSDKCrossEncoder(CrossEncoderModel):
|
||||
indices = [idx for idx, _ in indexed_texts]
|
||||
|
||||
# Build kwargs for rerank call
|
||||
rerank_kwargs = {
|
||||
rerank_kwargs: dict[str, Any] = {
|
||||
"model": self.model,
|
||||
"query": query,
|
||||
"documents": texts,
|
||||
"headers": reranker_bank_attribution_headers(),
|
||||
}
|
||||
if self.api_key:
|
||||
rerank_kwargs["api_key"] = self.api_key
|
||||
@@ -1281,21 +1256,9 @@ class LiteLLMSDKCrossEncoder(CrossEncoderModel):
|
||||
|
||||
response = await self._litellm.arerank(**rerank_kwargs)
|
||||
|
||||
# Map scores back to original positions
|
||||
# Response format: RerankResponse with results list
|
||||
# Each result is a TypedDict with "index" and "relevance_score"
|
||||
if hasattr(response, "results") and response.results:
|
||||
for result in response.results:
|
||||
# Results are TypedDicts, use dict-style access
|
||||
original_idx = result["index"]
|
||||
score = result.get("relevance_score", result.get("score", 0.0))
|
||||
all_scores[indices[original_idx]] = score
|
||||
elif isinstance(response, list):
|
||||
# Direct list of scores (unlikely but defensive)
|
||||
for i, score in enumerate(response):
|
||||
all_scores[indices[i]] = score
|
||||
else:
|
||||
logger.warning(f"Unexpected response format from LiteLLM rerank: {type(response)}")
|
||||
for result in response.results:
|
||||
original_idx = result["index"]
|
||||
all_scores[indices[original_idx]] = result["relevance_score"]
|
||||
|
||||
return all_scores
|
||||
|
||||
@@ -1614,131 +1577,246 @@ class AlibabaCloudCrossEncoder(CrossEncoderModel):
|
||||
return await self._client.predict(pairs)
|
||||
|
||||
|
||||
def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
"""
|
||||
Create a CrossEncoderModel instance based on configuration.
|
||||
class MultiCrossEncoder(CrossEncoderModel):
|
||||
"""Failover across an ordered chain of cross-encoders.
|
||||
|
||||
Reads configuration via get_config() to ensure consistency across the codebase.
|
||||
Member 0 is the primary (the unindexed ``HINDSIGHT_API_RERANKER_*`` config);
|
||||
members 1..N are the indexed fallbacks. Each ``predict`` tries members in order
|
||||
and returns the first usable set of scores, so an unreachable reranker costs
|
||||
ranking quality (whatever the next member gives) instead of the whole recall.
|
||||
Put ``rrf`` last to degrade to the fusion order rather than failing.
|
||||
|
||||
Each member keeps its own retry budget, so we only advance after a member has
|
||||
exhausted its retries and raised. A member that fails to initialize is not
|
||||
fatal — that is the point of the chain — it is retried lazily on the next
|
||||
request that reaches it.
|
||||
"""
|
||||
|
||||
def __init__(self, members: list[CrossEncoderModel]) -> None:
|
||||
if len(members) < 2:
|
||||
raise ValueError("MultiCrossEncoder requires at least two members")
|
||||
self._members = members
|
||||
self._ready = [False] * len(members)
|
||||
self._locks = [asyncio.Lock() for _ in members]
|
||||
self._active = 0
|
||||
|
||||
@property
|
||||
def provider_name(self) -> str:
|
||||
"""The provider of the member that last served a request (primary before any).
|
||||
|
||||
Callers use this to detect a passthrough reranker, so it has to track the
|
||||
member actually serving rather than name the chain: a chain that has
|
||||
degraded to its ``rrf`` member is passthrough. Concurrent requests share it,
|
||||
so a request that fails over can briefly mislabel a neighbour — this only
|
||||
tunes downstream scoring, never correctness.
|
||||
"""
|
||||
return self._members[self._active].provider_name
|
||||
|
||||
async def _initialize_member(self, index: int) -> None:
|
||||
"""Initialize one member, off the event loop when it loads a model in-process."""
|
||||
member = self._members[index]
|
||||
if member.blocking_init:
|
||||
loop = asyncio.get_running_loop()
|
||||
await loop.run_in_executor(None, lambda: asyncio.run(member.initialize()))
|
||||
else:
|
||||
await member.initialize()
|
||||
self._ready[index] = True
|
||||
|
||||
async def _ensure_member_ready(self, index: int) -> None:
|
||||
async with self._locks[index]:
|
||||
if not self._ready[index]:
|
||||
await self._initialize_member(index)
|
||||
|
||||
async def initialize(self) -> None:
|
||||
"""Initialize every member, tolerating members that are down.
|
||||
|
||||
Members initialize concurrently so one unreachable member cannot eat the
|
||||
startup budget the others need. Failures are logged and retried on use.
|
||||
"""
|
||||
results = await asyncio.gather(
|
||||
*(self._ensure_member_ready(i) for i in range(len(self._members))),
|
||||
return_exceptions=True,
|
||||
)
|
||||
for index, result in enumerate(results):
|
||||
if isinstance(result, BaseException):
|
||||
logger.warning(
|
||||
"Reranker member %d (%s) failed to initialize: %s; it will be retried on use",
|
||||
index,
|
||||
self._members[index].provider_name,
|
||||
result,
|
||||
)
|
||||
if not any(self._ready):
|
||||
logger.error("Reranker: no member of the failover chain initialized; recall will retry them per request")
|
||||
|
||||
async def predict(self, pairs: list[tuple[str, str]]) -> list[float]:
|
||||
"""Score ``pairs`` with the first member that answers usably."""
|
||||
last_exc: BaseException | None = None
|
||||
for index, member in enumerate(self._members):
|
||||
try:
|
||||
if not self._ready[index]:
|
||||
await self._ensure_member_ready(index)
|
||||
scores = await member.predict(pairs)
|
||||
if len(scores) != len(pairs):
|
||||
raise RuntimeError(f"returned {len(scores)} scores for {len(pairs)} pairs")
|
||||
except Exception as e: # noqa: BLE001 - re-raised below if no member answers
|
||||
last_exc = e
|
||||
remaining = len(self._members) - index - 1
|
||||
logger.warning(
|
||||
"Reranker member %d (%s) failed: %s%s",
|
||||
index,
|
||||
member.provider_name,
|
||||
e,
|
||||
f"; trying next member ({remaining} left)" if remaining else "; no members left",
|
||||
)
|
||||
continue
|
||||
if index != self._active:
|
||||
logger.info(
|
||||
"Reranker: now serving from member %d (%s)",
|
||||
index,
|
||||
member.provider_name,
|
||||
)
|
||||
self._active = index
|
||||
return scores
|
||||
# All members failed; surface the last error (loop ran at least once).
|
||||
assert last_exc is not None
|
||||
raise last_exc
|
||||
|
||||
|
||||
def create_cross_encoder(member: RerankerMemberConfig) -> CrossEncoderModel:
|
||||
"""
|
||||
Create a CrossEncoderModel for one member of the reranker chain.
|
||||
|
||||
``member`` is the primary (index 0, the unindexed ``HINDSIGHT_API_RERANKER_*``
|
||||
config) or an indexed fallback. Missing-setting errors name the member's own
|
||||
env var, so a chain misconfiguration points at the exact indexed variable.
|
||||
|
||||
Args:
|
||||
member: Resolved settings for this member
|
||||
|
||||
Returns:
|
||||
Configured CrossEncoderModel instance
|
||||
"""
|
||||
from ..config import get_config
|
||||
|
||||
config = get_config()
|
||||
provider = config.reranker_provider.lower()
|
||||
provider = member.provider.lower()
|
||||
|
||||
if provider == "tei":
|
||||
url = config.reranker_tei_url
|
||||
url = member.tei_url
|
||||
if not url:
|
||||
raise ValueError(f"{ENV_RERANKER_TEI_URL} is required when {ENV_RERANKER_PROVIDER} is 'tei'")
|
||||
raise ValueError(f"{member.env_name('TEI_URL')} is required when {member.env_name('PROVIDER')} is 'tei'")
|
||||
return RemoteTEICrossEncoder(
|
||||
base_url=url,
|
||||
timeout=config.reranker_tei_http_timeout,
|
||||
batch_size=config.reranker_tei_batch_size,
|
||||
max_concurrent=config.reranker_tei_max_concurrent,
|
||||
timeout=member.tei_http_timeout,
|
||||
batch_size=member.tei_batch_size,
|
||||
max_concurrent=member.tei_max_concurrent,
|
||||
)
|
||||
elif provider == "local":
|
||||
return LocalSTCrossEncoder(
|
||||
model_name=config.reranker_local_model,
|
||||
max_concurrent=config.reranker_local_max_concurrent,
|
||||
force_cpu=config.reranker_local_force_cpu,
|
||||
trust_remote_code=config.reranker_local_trust_remote_code,
|
||||
fp16=config.reranker_local_fp16,
|
||||
bucket_batching=config.reranker_local_bucket_batching,
|
||||
batch_size=config.reranker_local_batch_size,
|
||||
model_name=member.local_model,
|
||||
max_concurrent=member.local_max_concurrent,
|
||||
force_cpu=member.local_force_cpu,
|
||||
trust_remote_code=member.local_trust_remote_code,
|
||||
fp16=member.local_fp16,
|
||||
bucket_batching=member.local_bucket_batching,
|
||||
batch_size=member.local_batch_size,
|
||||
allow_mps=member.local_allow_mps,
|
||||
)
|
||||
elif provider == "cohere":
|
||||
api_key = config.reranker_cohere_api_key
|
||||
if not api_key:
|
||||
raise ValueError(f"{ENV_RERANKER_COHERE_API_KEY} is required when {ENV_RERANKER_PROVIDER} is 'cohere'")
|
||||
return CohereCrossEncoder(
|
||||
api_key=api_key,
|
||||
model=config.reranker_cohere_model,
|
||||
base_url=config.reranker_cohere_base_url,
|
||||
timeout=config.reranker_cohere_timeout,
|
||||
)
|
||||
elif provider == "openrouter":
|
||||
api_key = config.reranker_openrouter_api_key
|
||||
api_key = member.cohere_api_key
|
||||
if not api_key:
|
||||
raise ValueError(
|
||||
"HINDSIGHT_API_RERANKER_OPENROUTER_API_KEY, HINDSIGHT_API_OPENROUTER_API_KEY, "
|
||||
f"or HINDSIGHT_API_LLM_API_KEY is required when {ENV_RERANKER_PROVIDER} is 'openrouter'"
|
||||
f"{member.env_name('COHERE_API_KEY')} is required when {member.env_name('PROVIDER')} is 'cohere'"
|
||||
)
|
||||
return CohereCrossEncoder(
|
||||
api_key=api_key,
|
||||
model=config.reranker_openrouter_model,
|
||||
base_url=config.reranker_openrouter_base_url,
|
||||
timeout=config.reranker_openrouter_timeout,
|
||||
model=member.cohere_model,
|
||||
base_url=member.cohere_base_url,
|
||||
timeout=member.cohere_timeout,
|
||||
)
|
||||
elif provider == "openrouter":
|
||||
api_key = member.openrouter_api_key
|
||||
if not api_key:
|
||||
shared = ", HINDSIGHT_API_OPENROUTER_API_KEY, or HINDSIGHT_API_LLM_API_KEY" if member.index == 0 else ""
|
||||
raise ValueError(
|
||||
f"{member.env_name('OPENROUTER_API_KEY')}{shared} is required "
|
||||
f"when {member.env_name('PROVIDER')} is 'openrouter'"
|
||||
)
|
||||
return CohereCrossEncoder(
|
||||
api_key=api_key,
|
||||
model=member.openrouter_model,
|
||||
base_url=member.openrouter_base_url,
|
||||
timeout=member.openrouter_timeout,
|
||||
)
|
||||
elif provider == "flashrank":
|
||||
model = os.environ.get(ENV_RERANKER_FLASHRANK_MODEL, DEFAULT_RERANKER_FLASHRANK_MODEL)
|
||||
cache_dir = os.environ.get(ENV_RERANKER_FLASHRANK_CACHE_DIR, DEFAULT_RERANKER_FLASHRANK_CACHE_DIR)
|
||||
cpu_mem_arena = os.environ.get(
|
||||
ENV_RERANKER_FLASHRANK_CPU_MEM_ARENA, str(DEFAULT_RERANKER_FLASHRANK_CPU_MEM_ARENA)
|
||||
).lower() in ("true", "1", "yes")
|
||||
return FlashRankCrossEncoder(model_name=model, cache_dir=cache_dir, cpu_mem_arena=cpu_mem_arena)
|
||||
return FlashRankCrossEncoder(
|
||||
model_name=member.flashrank_model,
|
||||
cache_dir=member.flashrank_cache_dir,
|
||||
cpu_mem_arena=member.flashrank_cpu_mem_arena,
|
||||
)
|
||||
elif provider == "litellm":
|
||||
return LiteLLMCrossEncoder(
|
||||
api_base=config.reranker_litellm_api_base,
|
||||
api_key=config.reranker_litellm_api_key,
|
||||
model=config.reranker_litellm_model,
|
||||
max_tokens_per_doc=config.reranker_litellm_max_tokens_per_doc,
|
||||
timeout=config.reranker_litellm_timeout,
|
||||
api_base=member.litellm_api_base,
|
||||
api_key=member.litellm_api_key,
|
||||
model=member.litellm_model,
|
||||
max_tokens_per_doc=member.litellm_max_tokens_per_doc,
|
||||
timeout=member.litellm_timeout,
|
||||
)
|
||||
elif provider == "litellm-sdk":
|
||||
return LiteLLMSDKCrossEncoder(
|
||||
api_key=config.reranker_litellm_sdk_api_key or None,
|
||||
model=config.reranker_litellm_sdk_model,
|
||||
api_base=config.reranker_litellm_sdk_api_base,
|
||||
max_tokens_per_doc=config.reranker_litellm_max_tokens_per_doc,
|
||||
timeout=config.reranker_litellm_sdk_timeout,
|
||||
api_key=member.litellm_sdk_api_key or None,
|
||||
model=member.litellm_sdk_model,
|
||||
api_base=member.litellm_sdk_api_base,
|
||||
max_tokens_per_doc=member.litellm_max_tokens_per_doc,
|
||||
timeout=member.litellm_sdk_timeout,
|
||||
)
|
||||
elif provider == "zeroentropy":
|
||||
api_key = config.reranker_zeroentropy_api_key
|
||||
api_key = member.zeroentropy_api_key
|
||||
if not api_key:
|
||||
raise ValueError(
|
||||
f"{ENV_RERANKER_ZEROENTROPY_API_KEY} is required when {ENV_RERANKER_PROVIDER} is 'zeroentropy'"
|
||||
f"{member.env_name('ZEROENTROPY_API_KEY')} is required "
|
||||
f"when {member.env_name('PROVIDER')} is 'zeroentropy'"
|
||||
)
|
||||
return ZeroEntropyCrossEncoder(
|
||||
api_key=api_key,
|
||||
model=config.reranker_zeroentropy_model,
|
||||
base_url=config.reranker_zeroentropy_base_url,
|
||||
timeout=config.reranker_zeroentropy_timeout,
|
||||
model=member.zeroentropy_model,
|
||||
base_url=member.zeroentropy_base_url,
|
||||
timeout=member.zeroentropy_timeout,
|
||||
)
|
||||
elif provider == "siliconflow":
|
||||
api_key = config.reranker_siliconflow_api_key
|
||||
api_key = member.siliconflow_api_key
|
||||
if not api_key:
|
||||
raise ValueError(
|
||||
f"{ENV_RERANKER_SILICONFLOW_API_KEY} is required when {ENV_RERANKER_PROVIDER} is 'siliconflow'"
|
||||
f"{member.env_name('SILICONFLOW_API_KEY')} is required "
|
||||
f"when {member.env_name('PROVIDER')} is 'siliconflow'"
|
||||
)
|
||||
return SiliconFlowCrossEncoder(
|
||||
api_key=api_key,
|
||||
model=config.reranker_siliconflow_model,
|
||||
base_url=config.reranker_siliconflow_base_url,
|
||||
timeout=config.reranker_siliconflow_timeout,
|
||||
model=member.siliconflow_model,
|
||||
base_url=member.siliconflow_base_url,
|
||||
timeout=member.siliconflow_timeout,
|
||||
)
|
||||
elif provider == "google":
|
||||
project_id = config.reranker_google_project_id
|
||||
project_id = member.google_project_id
|
||||
if not project_id:
|
||||
shared = " (or HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID)" if member.index == 0 else ""
|
||||
raise ValueError(
|
||||
f"{ENV_RERANKER_GOOGLE_PROJECT_ID} (or HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID) "
|
||||
f"is required when {ENV_RERANKER_PROVIDER} is 'google'"
|
||||
f"{member.env_name('GOOGLE_PROJECT_ID')}{shared} "
|
||||
f"is required when {member.env_name('PROVIDER')} is 'google'"
|
||||
)
|
||||
return GoogleCrossEncoder(
|
||||
project_id=project_id,
|
||||
model=config.reranker_google_model,
|
||||
service_account_key=config.reranker_google_service_account_key,
|
||||
timeout=config.reranker_google_timeout,
|
||||
model=member.google_model,
|
||||
service_account_key=member.google_service_account_key,
|
||||
timeout=member.google_timeout,
|
||||
)
|
||||
elif provider == "alibaba":
|
||||
api_key = config.reranker_alibaba_api_key
|
||||
api_key = member.alibaba_api_key
|
||||
if not api_key:
|
||||
raise ValueError(f"{ENV_RERANKER_ALIBABA_API_KEY} is required when {ENV_RERANKER_PROVIDER} is 'alibaba'")
|
||||
raise ValueError(
|
||||
f"{member.env_name('ALIBABA_API_KEY')} is required when {member.env_name('PROVIDER')} is 'alibaba'"
|
||||
)
|
||||
return AlibabaCloudCrossEncoder(
|
||||
api_key=api_key,
|
||||
model=config.reranker_alibaba_model,
|
||||
timeout=config.reranker_alibaba_timeout,
|
||||
model=member.alibaba_model,
|
||||
timeout=member.alibaba_timeout,
|
||||
)
|
||||
elif provider == "rrf":
|
||||
return RRFPassthroughCrossEncoder()
|
||||
@@ -1748,3 +1826,23 @@ def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
raise ValueError(
|
||||
f"Unknown reranker provider: {provider}. Supported: 'local', 'tei', 'cohere', 'zeroentropy', 'siliconflow', 'alibaba', 'google', 'flashrank', 'litellm', 'litellm-sdk', 'rrf', 'jina-mlx'"
|
||||
)
|
||||
|
||||
|
||||
def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
"""
|
||||
Create the configured reranker, based on configuration.
|
||||
|
||||
Reads configuration via get_config() to ensure consistency across the codebase.
|
||||
With no ``HINDSIGHT_API_RERANKER_<n>_*`` members configured (the default) this
|
||||
is the single configured reranker; otherwise the chain is wrapped in a
|
||||
:class:`MultiCrossEncoder` that fails over across members in order.
|
||||
|
||||
Returns:
|
||||
Configured CrossEncoderModel instance
|
||||
"""
|
||||
from ..config import get_config
|
||||
|
||||
chain = get_config().reranker_chain()
|
||||
if len(chain) == 1:
|
||||
return create_cross_encoder(chain[0])
|
||||
return MultiCrossEncoder([create_cross_encoder(member) for member in chain])
|
||||
|
||||
@@ -67,16 +67,28 @@ def create_database_backend(backend_type: str) -> DatabaseBackend:
|
||||
return _get_backend_class(backend_type)()
|
||||
|
||||
|
||||
_OPS_CACHE: dict[str, DataAccessOps] = {}
|
||||
|
||||
|
||||
def create_data_access_ops(backend_type: str) -> DataAccessOps:
|
||||
"""Factory: create a DataAccessOps by backend name.
|
||||
"""Factory: the DataAccessOps for a backend name.
|
||||
|
||||
Returns a per-dialect SINGLETON: ``DataAccessOps`` is stateless (it only builds and runs SQL),
|
||||
so one shared instance per dialect is correct — and it means the database backend and the
|
||||
memories store hold the *same* ops object, so a test that patches a method on it (e.g.
|
||||
``enqueue_graph_maintenance``) observes every caller regardless of which layer issued it.
|
||||
|
||||
Args:
|
||||
backend_type: One of "postgresql" or "oracle".
|
||||
|
||||
Returns:
|
||||
A DataAccessOps instance.
|
||||
The shared DataAccessOps instance for that backend.
|
||||
|
||||
Raises:
|
||||
ValueError: If backend_type is not recognized.
|
||||
"""
|
||||
return _get_ops_class(backend_type)()
|
||||
ops = _OPS_CACHE.get(backend_type)
|
||||
if ops is None:
|
||||
ops = _get_ops_class(backend_type)()
|
||||
_OPS_CACHE[backend_type] = ops
|
||||
return ops
|
||||
|
||||
@@ -112,6 +112,23 @@ class DatabaseConnection(ABC):
|
||||
"""
|
||||
...
|
||||
|
||||
async def execute_rows_affected(self, query: str, *args: Any, timeout: float | None = None) -> int:
|
||||
"""Execute a DML statement and return the number of rows it affected.
|
||||
|
||||
Normalizes the dialect-specific execute result into a plain int so callers
|
||||
never hand-parse an ``"UPDATE <n>"`` / ``"DELETE <n>"`` command tag in
|
||||
business logic (mirrors ``parse_json`` above, which normalizes the other
|
||||
dialect-divergent result shape). asyncpg returns the tag directly; the
|
||||
Oracle connection reshapes ``cursor.rowcount`` into the same trailing-count
|
||||
form, so parsing the last token is dialect-safe. Returns 0 when the status
|
||||
has no trailing count (e.g. a non-DML statement).
|
||||
"""
|
||||
status = await self.execute(query, *args, timeout=timeout)
|
||||
if not isinstance(status, str):
|
||||
return 0
|
||||
parts = status.split()
|
||||
return int(parts[-1]) if parts and parts[-1].isdigit() else 0
|
||||
|
||||
@abstractmethod
|
||||
async def executemany(self, query: str, args: list[tuple[Any, ...]], *, timeout: float | None = None) -> None:
|
||||
"""Execute a query for each set of arguments.
|
||||
@@ -307,6 +324,17 @@ class DatabaseBackend(ABC):
|
||||
"""Close the connection pool and release all resources."""
|
||||
...
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def is_ready(self) -> bool:
|
||||
"""Whether the pool exists and can serve connections.
|
||||
|
||||
False before :meth:`initialize` and after :meth:`shutdown`. Best-effort
|
||||
callers (tracing, auditing) check this to skip work during those windows
|
||||
instead of acquiring and interpreting the resulting error.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
@asynccontextmanager
|
||||
async def acquire(self) -> AsyncIterator[DatabaseConnection]:
|
||||
|
||||
@@ -18,12 +18,75 @@ and mirrors Django's ``DatabaseOperations`` architecture.
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from typing import Any
|
||||
|
||||
from .base import DatabaseConnection
|
||||
from .result import ResultRow
|
||||
|
||||
|
||||
def graph_maintenance_bank_serialization_sql(table: str, alias: str) -> str:
|
||||
"""SQL predicate serialising ``graph_maintenance`` claims per bank (#3230).
|
||||
|
||||
Every graph_maintenance run is the same bank-wide sweep — the payload carries
|
||||
only ``bank_id``, and ``run_graph_maintenance_job`` drains the whole queue —
|
||||
so a second concurrent run for one bank adds no work. It is worse than
|
||||
useless: ``claim_graph_maintenance_batch`` locks queue rows ``FOR UPDATE``
|
||||
*without* ``SKIP LOCKED`` (it is written assuming a single runner per bank),
|
||||
so the runs convoy on each other's row locks while each holds a worker slot.
|
||||
|
||||
Same guarantee ``consolidation`` already gets from its ``bank_id != ALL(busy)``
|
||||
exclusion, and the same caveat: a row wedged in 'processing' holds its bank
|
||||
until something releases it (``hindsight-admin recover``, or a restart with a
|
||||
stable ``HINDSIGHT_API_WORKER_ID`` so ``recover_own_tasks`` matches it). That
|
||||
is a general gap in claim recovery, not specific to graph_maintenance.
|
||||
|
||||
Two differences from the consolidation form, both forced by the shape of this
|
||||
problem:
|
||||
|
||||
* It is a **predicate**, not a separate claim phase. Pulling graph_maintenance
|
||||
into its own phase after the generic shared-pool query would drop it below
|
||||
every other operation type: it has no reserved-slot floor
|
||||
(``WORKER_SLOT_TYPE_DEFAULTS`` gives consolidation 2 and graph_maintenance
|
||||
0), and the poller's fairness pass calls ``claim_tasks`` with
|
||||
``shared_limit=1``, so a single pending retain would starve it indefinitely.
|
||||
As a predicate it keeps competing by ``created_at``.
|
||||
* It also suppresses every same-bank row but the oldest **within one batch**.
|
||||
Excluding busy banks alone does not: with several pending rows and nothing
|
||||
yet processing, one batch claims them all — the convoy, unchanged. Several
|
||||
pending rows per bank are reachable through the recovery paths
|
||||
(``_reclaim_own_processing_tasks`` resets *all* of a worker's processing
|
||||
rows in one statement, from ``recover_own_tasks`` at startup and
|
||||
``release_own_tasks`` at shutdown, plus ``_schedule_retry`` /
|
||||
``_defer_operation`` / ``hindsight-admin recover``).
|
||||
|
||||
The candidate row is always 'pending' and the 'pending' branch is
|
||||
strictly-older, so the subquery can never match the candidate itself. The
|
||||
fragment carries no SQL comments on purpose — it is rewritten for Oracle by
|
||||
regex (``db/oracle.py``).
|
||||
|
||||
Args:
|
||||
table: Fully-qualified async_operations table.
|
||||
alias: Alias of the outer candidate row in the calling query.
|
||||
"""
|
||||
return f"""
|
||||
({alias}.operation_type <> 'graph_maintenance' OR NOT EXISTS (
|
||||
SELECT 1 FROM {table} gm_peer
|
||||
WHERE gm_peer.bank_id = {alias}.bank_id
|
||||
AND gm_peer.operation_type = 'graph_maintenance'
|
||||
AND (
|
||||
gm_peer.status = 'processing'
|
||||
OR (gm_peer.status = 'pending'
|
||||
AND gm_peer.task_payload IS NOT NULL
|
||||
AND (gm_peer.next_retry_at IS NULL OR gm_peer.next_retry_at <= NOW())
|
||||
AND (gm_peer.created_at < {alias}.created_at
|
||||
OR (gm_peer.created_at = {alias}.created_at
|
||||
AND gm_peer.operation_id < {alias}.operation_id)))
|
||||
)
|
||||
))
|
||||
"""
|
||||
|
||||
|
||||
@dataclass
|
||||
class TagListingParts:
|
||||
"""Backend-specific SQL fragments for the tag listing query."""
|
||||
@@ -34,6 +97,57 @@ class TagListingParts:
|
||||
bank_prefix: str
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class UpdatedWindow:
|
||||
"""Recall's ``created_after``/``created_before`` bounds, as SQL for graph expansion.
|
||||
|
||||
Recall applies the window to ``updated_at`` — a consolidation touch makes a
|
||||
fact current again — so link expansion has to bound the same column its seed
|
||||
query does. Filtering only the seeds is not enough: a single in-window seed
|
||||
would otherwise drag its whole neighbourhood (shared entities, semantic kNN
|
||||
links, causal links) into the results no matter how old those neighbours are.
|
||||
|
||||
``first_param_index`` is where the bounds land in the owning query's param
|
||||
list, so each call site keeps the placeholder numbering next to the params it
|
||||
binds. Rendering is per-alias because the same window is applied to several
|
||||
correlation names within one query.
|
||||
"""
|
||||
|
||||
after: datetime | None
|
||||
before: datetime | None
|
||||
first_param_index: int
|
||||
|
||||
def clause(self, alias: str) -> str:
|
||||
"""``AND <alias>.updated_at > $n ...`` — empty when the window is unbounded."""
|
||||
parts: list[str] = []
|
||||
index = self.first_param_index
|
||||
if self.after is not None:
|
||||
parts.append(f" AND {alias}.updated_at > ${index}")
|
||||
index += 1
|
||||
if self.before is not None:
|
||||
parts.append(f" AND {alias}.updated_at < ${index}")
|
||||
return "".join(parts)
|
||||
|
||||
@property
|
||||
def params(self) -> list[datetime]:
|
||||
"""The bound values, in placeholder order. Append to the owning param list."""
|
||||
return [bound for bound in (self.after, self.before) if bound is not None]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class LinkExpansionRows:
|
||||
"""The three link-expansion signals, kept apart until they are scored.
|
||||
|
||||
They cannot be concatenated at the SQL layer: each carries a different score
|
||||
scale (shared-entity count, kNN weight, causal weight) and the caller applies
|
||||
a different transformation to each before summing them.
|
||||
"""
|
||||
|
||||
entity: list[ResultRow]
|
||||
semantic: list[ResultRow]
|
||||
causal: list[ResultRow]
|
||||
|
||||
|
||||
class DataAccessOps(ABC):
|
||||
"""Backend-specific multi-statement data access operations.
|
||||
|
||||
@@ -149,9 +263,14 @@ class DataAccessOps(ABC):
|
||||
bank_id: str,
|
||||
entity_names: list[str],
|
||||
entity_dates: list,
|
||||
entity_kinds: list[str],
|
||||
) -> dict[str, str]:
|
||||
"""Bulk insert entities with ON CONFLICT DO NOTHING, returning id-by-lowercase-name.
|
||||
|
||||
``entity_kinds`` ("regular"/"label", parallel to ``entity_names``) is
|
||||
stored on the row so label entities stay out of the partial trigram
|
||||
index (#3208).
|
||||
|
||||
PG uses INSERT ... SELECT FROM unnest() with RETURNING.
|
||||
Non-PG inserts row-by-row then SELECTs.
|
||||
"""
|
||||
@@ -172,6 +291,26 @@ class DataAccessOps(ABC):
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def bulk_reassert_entities(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
bank_id: str,
|
||||
entity_ids: list[str],
|
||||
canonical_names: list[str],
|
||||
entity_kinds: list[str],
|
||||
) -> None:
|
||||
"""Lock resolved parents and re-create any pruned since Phase-1 resolution.
|
||||
|
||||
Closes the retain Phase-1/prune race (#2662): existing rows are locked
|
||||
(PG ``FOR KEY SHARE`` / Oracle ``FOR UPDATE``) so a concurrent
|
||||
``prune_orphan_entities`` blocks until the caller's transaction commits,
|
||||
while rows already deleted are re-inserted idempotently. ``entity_ids``
|
||||
must be sorted by the caller for a stable lock order.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def bulk_insert_unit_entities(
|
||||
self,
|
||||
@@ -230,12 +369,16 @@ class DataAccessOps(ABC):
|
||||
mu_table: str,
|
||||
ue_table: str,
|
||||
per_entity_limit: int,
|
||||
window: UpdatedWindow,
|
||||
) -> str:
|
||||
"""Build entity expansion CTE for link expansion retrieval.
|
||||
|
||||
PG uses DISTINCT ON with CROSS JOIN LATERAL and GROUP BY.
|
||||
Non-PG splits into entity_scores subquery then JOINs for full columns
|
||||
(can't GROUP BY CLOB).
|
||||
|
||||
``window`` narrows candidates *before* the per-entity cap, so out-of-window
|
||||
neighbours don't consume an entity's bounded fan-out.
|
||||
"""
|
||||
...
|
||||
|
||||
@@ -244,6 +387,7 @@ class DataAccessOps(ABC):
|
||||
self,
|
||||
ml_table: str,
|
||||
mu_table: str,
|
||||
window: UpdatedWindow,
|
||||
) -> str:
|
||||
"""Build semantic + causal expansion CTEs.
|
||||
|
||||
@@ -262,7 +406,8 @@ class DataAccessOps(ABC):
|
||||
seed_ids: list,
|
||||
budget: int,
|
||||
per_entity_limit: int,
|
||||
) -> tuple[list[ResultRow], list[ResultRow], list[ResultRow]]:
|
||||
window: UpdatedWindow,
|
||||
) -> LinkExpansionRows:
|
||||
"""Observation-specific graph expansion.
|
||||
|
||||
PG uses native array ops (source_memory_ids column) for performance.
|
||||
@@ -484,6 +629,23 @@ class DataAccessOps(ABC):
|
||||
|
||||
# -- Task claiming operations ------------------------------------------
|
||||
|
||||
@abstractmethod
|
||||
async def prune_terminal_operations(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
cutoff: datetime,
|
||||
*,
|
||||
batch_size: int,
|
||||
) -> int:
|
||||
"""Delete one deterministic batch of terminal operations older than ``cutoff``.
|
||||
|
||||
Implementations must lock candidates without waiting on rows another
|
||||
worker is pruning, never select pending/processing rows, and return the
|
||||
number deleted. The caller provides a transaction around this method.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def claim_tasks(
|
||||
self,
|
||||
@@ -501,6 +663,10 @@ class DataAccessOps(ABC):
|
||||
Oracle implementation uses two-step claims (query busy banks first, then
|
||||
claim excluding them) to avoid ORA-02014.
|
||||
|
||||
Implementations must apply :func:`graph_maintenance_bank_serialization_sql`
|
||||
to every query that can return a ``graph_maintenance`` row, so at most one
|
||||
such row per bank is ever in flight.
|
||||
|
||||
Args:
|
||||
consolidation_bank_priority: Per-bank priority for consolidation scheduling.
|
||||
Maps bank name patterns to integer priorities (higher = claimed first).
|
||||
|
||||
@@ -10,9 +10,17 @@ import uuid as uuid_mod
|
||||
from datetime import UTC, datetime
|
||||
|
||||
from .base import DatabaseConnection
|
||||
from .ops import DataAccessOps, TagListingParts
|
||||
from .ops import (
|
||||
DataAccessOps,
|
||||
LinkExpansionRows,
|
||||
TagListingParts,
|
||||
UpdatedWindow,
|
||||
graph_maintenance_bank_serialization_sql,
|
||||
)
|
||||
from .result import DictResultRow as ResultRow
|
||||
|
||||
ORACLE_IN_LIST_LIMIT = 1000
|
||||
|
||||
|
||||
class OracleOps(DataAccessOps):
|
||||
"""Oracle-specific data access operations."""
|
||||
@@ -172,22 +180,24 @@ class OracleOps(DataAccessOps):
|
||||
bank_id: str,
|
||||
entity_names: list[str],
|
||||
entity_dates: list,
|
||||
entity_kinds: list[str],
|
||||
) -> dict[str, str]:
|
||||
# Row-by-row insert with duplicate suppression.
|
||||
# Can't use RETURNING with ON CONFLICT DO NOTHING reliably,
|
||||
# so INSERT (ignoring dups) then SELECT all IDs at the end.
|
||||
id_by_name: dict[str, str] = {}
|
||||
for name, event_date in zip(entity_names, entity_dates):
|
||||
for name, event_date, kind in zip(entity_names, entity_dates, entity_kinds):
|
||||
ts = event_date if event_date else datetime.now(UTC)
|
||||
await conn.execute(
|
||||
f"""
|
||||
INSERT INTO {table} (bank_id, canonical_name, first_seen, last_seen, mention_count)
|
||||
VALUES ($1, $2, $3, $3, 0)
|
||||
INSERT INTO {table} (bank_id, canonical_name, first_seen, last_seen, mention_count, entity_kind)
|
||||
VALUES ($1, $2, $3, $3, 0, $4)
|
||||
ON CONFLICT (bank_id, LOWER(canonical_name)) DO NOTHING
|
||||
""",
|
||||
bank_id,
|
||||
name,
|
||||
ts,
|
||||
kind,
|
||||
)
|
||||
# Now SELECT all the entities we just inserted (or that already existed)
|
||||
for name in entity_names:
|
||||
@@ -216,7 +226,7 @@ class OracleOps(DataAccessOps):
|
||||
for orig_name in missing_names:
|
||||
row = await conn.fetchrow(
|
||||
f"""
|
||||
SELECT id, LOWER(canonical_name) AS name_lower
|
||||
SELECT id, canonical_name, LOWER(canonical_name) AS name_lower
|
||||
FROM {table}
|
||||
WHERE bank_id = $1 AND LOWER(canonical_name) = LOWER($2)
|
||||
""",
|
||||
@@ -224,10 +234,41 @@ class OracleOps(DataAccessOps):
|
||||
orig_name,
|
||||
)
|
||||
if row:
|
||||
# Wrap in a dict-like to include input_name for downstream compat
|
||||
results.append(row)
|
||||
return results
|
||||
|
||||
async def bulk_reassert_entities(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
bank_id: str,
|
||||
entity_ids: list[str],
|
||||
canonical_names: list[str],
|
||||
entity_kinds: list[str],
|
||||
) -> None:
|
||||
# Oracle has no FOR KEY SHARE; FOR UPDATE is the row-lock equivalent that
|
||||
# blocks a concurrent prune DELETE until this transaction commits. Lock
|
||||
# each surviving parent in the caller's stable id order (pruned ids are
|
||||
# simply absent here), then re-insert any that vanished. The translation
|
||||
# layer rewrites ON CONFLICT DO NOTHING to strip-and-catch ORA-00001, so
|
||||
# a name recreated under a new id is suppressed rather than raising.
|
||||
for entity_id in entity_ids:
|
||||
await conn.fetchrow(
|
||||
f"SELECT id FROM {table} WHERE id = $1 FOR UPDATE",
|
||||
entity_id,
|
||||
)
|
||||
await conn.executemany(
|
||||
f"""
|
||||
INSERT INTO {table} (id, bank_id, canonical_name, entity_kind)
|
||||
VALUES ($1, $2, $3, $4)
|
||||
ON CONFLICT DO NOTHING
|
||||
""",
|
||||
[
|
||||
(entity_id, bank_id, canonical_name, kind)
|
||||
for entity_id, canonical_name, kind in zip(entity_ids, canonical_names, entity_kinds)
|
||||
],
|
||||
)
|
||||
|
||||
async def bulk_insert_unit_entities(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
@@ -253,22 +294,29 @@ class OracleOps(DataAccessOps):
|
||||
) -> None:
|
||||
if not unit_ids:
|
||||
return
|
||||
# Oracle doesn't support ON CONFLICT; rely on the PK and the
|
||||
# IGNORE_ROW_ON_DUPKEY_INDEX hint to skip duplicates server-side.
|
||||
# The hint name must match the PK constraint exactly.
|
||||
# Locking upsert (#3034), the Oracle analogue of the PG
|
||||
# ``ON CONFLICT DO UPDATE``. The old IGNORE_ROW_ON_DUPKEY_INDEX insert
|
||||
# skipped duplicates WITHOUT locking the existing row, so a mutation
|
||||
# re-enqueueing an already-queued unit could not block a worker from
|
||||
# concurrently claiming (deleting) that row and processing the unit's
|
||||
# pre-mutation state — the re-enqueue signal was silently lost. MERGE
|
||||
# WHEN MATCHED takes an exclusive row lock on the existing queue row
|
||||
# (the SET is a deliberate no-op that preserves enqueued_at); WHEN NOT
|
||||
# MATCHED inserts a fresh row. That serialises the mutation against the
|
||||
# worker's claim for the same (bank_id, unit_id).
|
||||
#
|
||||
# Sort to enforce a global lock-acquisition order on the
|
||||
# (bank_id, unit_id) PK. Without this, two concurrent
|
||||
# transactions inserting overlapping unit_id sets in different
|
||||
# orders can deadlock on the unique-check row locks. Sorting
|
||||
# gives every concurrent caller the same lock order, so
|
||||
# conflicting inserts queue cleanly instead of cycling.
|
||||
# Sort to enforce a global (bank_id, unit_id) lock-acquisition order,
|
||||
# matching claim_graph_maintenance_batch's delete order, so overlapping
|
||||
# mutation/worker sets acquire the shared row locks ascending and cannot
|
||||
# cycle.
|
||||
sorted_unit_ids = sorted(unit_ids)
|
||||
await conn.executemany(
|
||||
f"""
|
||||
INSERT /*+ IGNORE_ROW_ON_DUPKEY_INDEX({table}, pk_graph_maintenance_queue) */
|
||||
INTO {table} (bank_id, unit_id)
|
||||
VALUES ($1, $2)
|
||||
MERGE INTO {table} q
|
||||
USING (SELECT $1 AS bank_id, $2 AS unit_id FROM dual) s
|
||||
ON (q.bank_id = s.bank_id AND q.unit_id = s.unit_id)
|
||||
WHEN MATCHED THEN UPDATE SET q.enqueued_at = q.enqueued_at
|
||||
WHEN NOT MATCHED THEN INSERT (bank_id, unit_id) VALUES (s.bank_id, s.unit_id)
|
||||
""",
|
||||
[(bank_id, uid) for uid in sorted_unit_ids],
|
||||
)
|
||||
@@ -293,7 +341,15 @@ class OracleOps(DataAccessOps):
|
||||
bank_id,
|
||||
limit,
|
||||
)
|
||||
claimed = [str(row["unit_id"]) for row in rows]
|
||||
# Ordered locking (#3034): the per-row DELETE takes the queue rows'
|
||||
# exclusive locks in executemany array order. Sort the claimed keys by
|
||||
# unit_id so those locks are acquired in the same (bank_id, unit_id)
|
||||
# order the enqueue MERGE uses — overlapping mutation/worker sets then
|
||||
# lock the shared rows ascending and cannot cycle. (The batch is still
|
||||
# *chosen* oldest-first by enqueued_at above; only the lock/delete order
|
||||
# is normalised.) The Pass 1 retry wrap in run_graph_maintenance_job is
|
||||
# the ORA-00060 backstop for any residual interleaving.
|
||||
claimed = sorted(str(row["unit_id"]) for row in rows)
|
||||
if claimed:
|
||||
await conn.executemany(
|
||||
f"DELETE FROM {table} WHERE bank_id = $1 AND unit_id = $2",
|
||||
@@ -329,6 +385,12 @@ class OracleOps(DataAccessOps):
|
||||
entities_table: str,
|
||||
bank_id: str,
|
||||
) -> int:
|
||||
# NB: the Postgres path additionally selects victims FOR UPDATE in sorted
|
||||
# (entity_id_1, entity_id_2) order to prevent the #2529 deadlock against
|
||||
# retain's sorted cooccurrence upsert. Oracle's DELETE can't carry that
|
||||
# ordered-lock CTE the same way, so here we rely on the Pass 2/3 retry
|
||||
# wrap in run_graph_maintenance_job (retry_with_backoff is ORA-00060
|
||||
# deadlock-aware) to recover instead. Deliberate dialect asymmetry.
|
||||
deleted = await conn.execute(
|
||||
f"""
|
||||
DELETE FROM {ec_table}
|
||||
@@ -430,6 +492,7 @@ class OracleOps(DataAccessOps):
|
||||
mu_table: str,
|
||||
ue_table: str,
|
||||
per_entity_limit: int,
|
||||
window: UpdatedWindow,
|
||||
) -> str:
|
||||
# Oracle: can't GROUP BY CLOB columns (text, context).
|
||||
# Restructure: count entities per unit_id in a subquery, then join to get full columns.
|
||||
@@ -447,6 +510,16 @@ class OracleOps(DataAccessOps):
|
||||
FROM {ue_table} ue_target
|
||||
WHERE ue_target.entity_id = se.entity_id
|
||||
AND ue_target.unit_id != ALL($1::uuid[])
|
||||
-- Filter before applying the cap: candidates from other fact
|
||||
-- types, or outside the recall window, must not consume this
|
||||
-- entity's bounded fan-out.
|
||||
AND EXISTS (
|
||||
SELECT 1
|
||||
FROM {mu_table} mu_target
|
||||
WHERE mu_target.id = ue_target.unit_id
|
||||
AND mu_target.fact_type = $2
|
||||
{window.clause("mu_target")}
|
||||
)
|
||||
ORDER BY ue_target.unit_id DESC
|
||||
FETCH FIRST {per_entity_limit} ROWS ONLY
|
||||
) t
|
||||
@@ -459,7 +532,6 @@ class OracleOps(DataAccessOps):
|
||||
es.score, 'entity' AS source
|
||||
FROM entity_scores es
|
||||
JOIN {mu_table} mu ON mu.id = es.unit_id
|
||||
WHERE mu.fact_type = $2
|
||||
ORDER BY es.score DESC
|
||||
FETCH FIRST $3 ROWS ONLY
|
||||
)"""
|
||||
@@ -468,6 +540,7 @@ class OracleOps(DataAccessOps):
|
||||
self,
|
||||
ml_table: str,
|
||||
mu_table: str,
|
||||
window: UpdatedWindow,
|
||||
) -> str:
|
||||
# Non-PG: can't GROUP BY CLOB columns, no DISTINCT ON.
|
||||
# Restructure semantic: compute max weight per id, then join for full columns.
|
||||
@@ -482,6 +555,7 @@ class OracleOps(DataAccessOps):
|
||||
AND ml.link_type = 'semantic'
|
||||
AND mu.fact_type = $2
|
||||
AND mu.id != ALL($1::uuid[])
|
||||
{window.clause("mu")}
|
||||
UNION ALL
|
||||
SELECT mu.id, ml.weight
|
||||
FROM {ml_table} ml
|
||||
@@ -490,6 +564,7 @@ class OracleOps(DataAccessOps):
|
||||
AND ml.link_type = 'semantic'
|
||||
AND mu.fact_type = $2
|
||||
AND mu.id != ALL($1::uuid[])
|
||||
{window.clause("mu")}
|
||||
) sem_raw
|
||||
GROUP BY id
|
||||
),
|
||||
@@ -516,6 +591,7 @@ class OracleOps(DataAccessOps):
|
||||
WHERE ml.from_unit_id = ANY($1::uuid[])
|
||||
AND ml.link_type IN ('causes', 'caused_by', 'enables', 'prevents')
|
||||
AND mu.fact_type = $2
|
||||
{window.clause("mu")}
|
||||
),
|
||||
causal_expanded AS (
|
||||
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
@@ -534,7 +610,8 @@ class OracleOps(DataAccessOps):
|
||||
seed_ids: list,
|
||||
budget: int,
|
||||
per_entity_limit: int,
|
||||
) -> tuple[list[ResultRow], list[ResultRow], list[ResultRow]]:
|
||||
window: UpdatedWindow,
|
||||
) -> LinkExpansionRows:
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -588,11 +665,13 @@ class OracleOps(DataAccessOps):
|
||||
WHERE os3.observation_id = mu.id
|
||||
AND os3.source_id IN (SELECT source_id FROM connected_sources)
|
||||
)
|
||||
{window.clause("mu")}
|
||||
ORDER BY score DESC
|
||||
FETCH FIRST $2 ROWS ONLY
|
||||
""",
|
||||
seed_ids,
|
||||
budget,
|
||||
*window.params,
|
||||
)
|
||||
logger.debug(f"[LinkExpansion] observation graph (Oracle): found {len(entity_rows)} connected observations")
|
||||
|
||||
@@ -608,12 +687,14 @@ class OracleOps(DataAccessOps):
|
||||
WHERE ml.from_unit_id = ANY($1::uuid[])
|
||||
AND ml.link_type = 'semantic' AND mu.fact_type = 'observation'
|
||||
AND mu.id != ALL($1::uuid[])
|
||||
{window.clause("mu")}
|
||||
UNION ALL
|
||||
SELECT mu.id, ml.weight
|
||||
FROM {ml_table} ml JOIN {mu_table} mu ON mu.id = ml.from_unit_id
|
||||
WHERE ml.to_unit_id = ANY($1::uuid[])
|
||||
AND ml.link_type = 'semantic' AND mu.fact_type = 'observation'
|
||||
AND mu.id != ALL($1::uuid[])
|
||||
{window.clause("mu")}
|
||||
) sem_raw
|
||||
GROUP BY id
|
||||
),
|
||||
@@ -639,6 +720,7 @@ class OracleOps(DataAccessOps):
|
||||
WHERE ml.from_unit_id = ANY($1::uuid[])
|
||||
AND ml.link_type IN ('causes', 'caused_by', 'enables', 'prevents')
|
||||
AND mu.fact_type = 'observation'
|
||||
{window.clause("mu")}
|
||||
),
|
||||
causal_expanded AS (
|
||||
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
@@ -653,11 +735,12 @@ class OracleOps(DataAccessOps):
|
||||
""",
|
||||
seed_ids,
|
||||
budget,
|
||||
*window.params,
|
||||
)
|
||||
|
||||
semantic_rows = [r for r in sem_causal_rows if r["source"] == "semantic"]
|
||||
causal_rows = [r for r in sem_causal_rows if r["source"] == "causal"]
|
||||
return list(entity_rows), semantic_rows, causal_rows
|
||||
return LinkExpansionRows(entity=list(entity_rows), semantic=semantic_rows, causal=causal_rows)
|
||||
|
||||
def build_tag_listing_parts(self, mu_table: str) -> TagListingParts:
|
||||
return TagListingParts(
|
||||
@@ -824,6 +907,157 @@ class OracleOps(DataAccessOps):
|
||||
|
||||
# -- Task claiming operations ------------------------------------------
|
||||
|
||||
async def prune_terminal_operations(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
cutoff: datetime,
|
||||
*,
|
||||
batch_size: int,
|
||||
) -> int:
|
||||
# Oracle rejects a row-limited SELECT ... FOR UPDATE (ORA-02014). Pick
|
||||
# the deterministic bounded IDs first, then lock only that candidate
|
||||
# set and re-check eligibility before deleting in the same transaction.
|
||||
# Clamp to Oracle's 1000-expression IN-list limit because the adapter
|
||||
# expands the candidate UUID list into individual bind variables.
|
||||
# Cancelled children cannot complete parent aggregation, so retain the
|
||||
# parent guard only for completed/failed children. Before removing a
|
||||
# cancelled child, preserve its signal by cancelling a pending parent
|
||||
# in this transaction and refreshing the parent's retention window.
|
||||
# Validate metadata before HEXTORAW: CASE makes malformed UUIDs yield
|
||||
# NULL while keeping the indexed RAW parent.operation_id key unwrapped.
|
||||
effective_batch_size = min(batch_size, ORACLE_IN_LIST_LIMIT)
|
||||
candidates = await conn.fetch(
|
||||
f"""
|
||||
SELECT candidate_operation.operation_id
|
||||
FROM {table} candidate_operation
|
||||
WHERE candidate_operation.status IN ('completed', 'failed', 'cancelled')
|
||||
AND candidate_operation.updated_at < $1
|
||||
AND (
|
||||
candidate_operation.status = 'cancelled'
|
||||
OR NOT EXISTS (
|
||||
SELECT 1
|
||||
FROM {table} parent
|
||||
WHERE parent.operation_id = CASE
|
||||
WHEN REGEXP_LIKE(
|
||||
JSON_VALUE(
|
||||
candidate_operation.result_metadata,
|
||||
'$.parent_operation_id' RETURNING VARCHAR2(36) NULL ON ERROR
|
||||
),
|
||||
'^[0-9A-Fa-f]{{8}}-[0-9A-Fa-f]{{4}}-[0-9A-Fa-f]{{4}}-[0-9A-Fa-f]{{4}}-[0-9A-Fa-f]{{12}}$'
|
||||
)
|
||||
THEN HEXTORAW(REPLACE(
|
||||
JSON_VALUE(
|
||||
candidate_operation.result_metadata,
|
||||
'$.parent_operation_id' RETURNING VARCHAR2(36) NULL ON ERROR
|
||||
),
|
||||
'-',
|
||||
''
|
||||
))
|
||||
ELSE NULL
|
||||
END
|
||||
AND parent.bank_id = candidate_operation.bank_id
|
||||
)
|
||||
)
|
||||
ORDER BY candidate_operation.updated_at, candidate_operation.operation_id
|
||||
LIMIT $2
|
||||
""",
|
||||
cutoff,
|
||||
effective_batch_size,
|
||||
)
|
||||
if not candidates:
|
||||
return 0
|
||||
|
||||
candidate_ids = [row["operation_id"] for row in candidates]
|
||||
locked = await conn.fetch(
|
||||
f"""
|
||||
SELECT candidate_operation.operation_id
|
||||
FROM {table} candidate_operation
|
||||
WHERE candidate_operation.operation_id = ANY($1)
|
||||
AND candidate_operation.status IN ('completed', 'failed', 'cancelled')
|
||||
AND candidate_operation.updated_at < $2
|
||||
AND (
|
||||
candidate_operation.status = 'cancelled'
|
||||
OR NOT EXISTS (
|
||||
SELECT 1
|
||||
FROM {table} parent
|
||||
WHERE parent.operation_id = CASE
|
||||
WHEN REGEXP_LIKE(
|
||||
JSON_VALUE(
|
||||
candidate_operation.result_metadata,
|
||||
'$.parent_operation_id' RETURNING VARCHAR2(36) NULL ON ERROR
|
||||
),
|
||||
'^[0-9A-Fa-f]{{8}}-[0-9A-Fa-f]{{4}}-[0-9A-Fa-f]{{4}}-[0-9A-Fa-f]{{4}}-[0-9A-Fa-f]{{12}}$'
|
||||
)
|
||||
THEN HEXTORAW(REPLACE(
|
||||
JSON_VALUE(
|
||||
candidate_operation.result_metadata,
|
||||
'$.parent_operation_id' RETURNING VARCHAR2(36) NULL ON ERROR
|
||||
),
|
||||
'-',
|
||||
''
|
||||
))
|
||||
ELSE NULL
|
||||
END
|
||||
AND parent.bank_id = candidate_operation.bank_id
|
||||
)
|
||||
)
|
||||
ORDER BY candidate_operation.updated_at, candidate_operation.operation_id
|
||||
FOR UPDATE OF candidate_operation.operation_id SKIP LOCKED
|
||||
""",
|
||||
candidate_ids,
|
||||
cutoff,
|
||||
)
|
||||
if not locked:
|
||||
return 0
|
||||
operation_ids = [row["operation_id"] for row in locked]
|
||||
await conn.execute(
|
||||
f"""
|
||||
UPDATE {table} parent
|
||||
SET status = 'cancelled',
|
||||
updated_at = now(),
|
||||
completed_at = COALESCE(parent.completed_at, now()),
|
||||
error_message = COALESCE(
|
||||
parent.error_message,
|
||||
'Cancelled because a child operation was cancelled'
|
||||
)
|
||||
WHERE parent.status = 'pending'
|
||||
AND EXISTS (
|
||||
SELECT 1
|
||||
FROM {table} candidate_operation
|
||||
WHERE candidate_operation.operation_id = ANY($1)
|
||||
AND candidate_operation.status = 'cancelled'
|
||||
AND candidate_operation.updated_at < $2
|
||||
AND candidate_operation.bank_id = parent.bank_id
|
||||
AND parent.operation_id = CASE
|
||||
WHEN REGEXP_LIKE(
|
||||
JSON_VALUE(
|
||||
candidate_operation.result_metadata,
|
||||
'$.parent_operation_id' RETURNING VARCHAR2(36) NULL ON ERROR
|
||||
),
|
||||
'^[0-9A-Fa-f]{{8}}-[0-9A-Fa-f]{{4}}-[0-9A-Fa-f]{{4}}-[0-9A-Fa-f]{{4}}-[0-9A-Fa-f]{{12}}$'
|
||||
)
|
||||
THEN HEXTORAW(REPLACE(
|
||||
JSON_VALUE(
|
||||
candidate_operation.result_metadata,
|
||||
'$.parent_operation_id' RETURNING VARCHAR2(36) NULL ON ERROR
|
||||
),
|
||||
'-',
|
||||
''
|
||||
))
|
||||
ELSE NULL
|
||||
END
|
||||
)
|
||||
""",
|
||||
operation_ids,
|
||||
cutoff,
|
||||
)
|
||||
await conn.execute(
|
||||
f"DELETE FROM {table} WHERE operation_id = ANY($1)",
|
||||
operation_ids,
|
||||
)
|
||||
return len(operation_ids)
|
||||
|
||||
async def _claim_consolidation_tasks(
|
||||
self,
|
||||
conn,
|
||||
@@ -1104,13 +1338,14 @@ class OracleOps(DataAccessOps):
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = $1
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
ORDER BY created_at
|
||||
SELECT o.operation_id, o.operation_type, o.task_payload, o.retry_count
|
||||
FROM {table} o
|
||||
WHERE o.status = 'pending'
|
||||
AND o.task_payload IS NOT NULL
|
||||
AND o.operation_type = $1
|
||||
AND (o.next_retry_at IS NULL OR o.next_retry_at <= NOW())
|
||||
AND {graph_maintenance_bank_serialization_sql(table, "o")}
|
||||
ORDER BY o.created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
@@ -1125,18 +1360,21 @@ class OracleOps(DataAccessOps):
|
||||
# --- Phase 2: claim from shared pool ---
|
||||
remaining_shared = shared_limit
|
||||
if remaining_shared > 0:
|
||||
# 2a. Non-consolidation tasks
|
||||
# 2a. Non-consolidation tasks. graph_maintenance stays in this
|
||||
# created_at-ordered query — see graph_maintenance_bank_serialization_sql
|
||||
# for why it is a predicate rather than a phase of its own.
|
||||
if claimed_ids:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type != 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND operation_id != ALL($1::uuid[])
|
||||
ORDER BY created_at
|
||||
SELECT o.operation_id, o.operation_type, o.task_payload, o.retry_count
|
||||
FROM {table} o
|
||||
WHERE o.status = 'pending'
|
||||
AND o.task_payload IS NOT NULL
|
||||
AND o.operation_type != 'consolidation'
|
||||
AND (o.next_retry_at IS NULL OR o.next_retry_at <= NOW())
|
||||
AND o.operation_id != ALL($1::uuid[])
|
||||
AND {graph_maintenance_bank_serialization_sql(table, "o")}
|
||||
ORDER BY o.created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
@@ -1146,13 +1384,14 @@ class OracleOps(DataAccessOps):
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type != 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
ORDER BY created_at
|
||||
SELECT o.operation_id, o.operation_type, o.task_payload, o.retry_count
|
||||
FROM {table} o
|
||||
WHERE o.status = 'pending'
|
||||
AND o.task_payload IS NOT NULL
|
||||
AND o.operation_type != 'consolidation'
|
||||
AND (o.next_retry_at IS NULL OR o.next_retry_at <= NOW())
|
||||
AND {graph_maintenance_bank_serialization_sql(table, "o")}
|
||||
ORDER BY o.created_at
|
||||
LIMIT $1
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
|
||||
@@ -4,14 +4,77 @@ Uses unnest(), LATERAL, DISTINCT ON, and native array operations for
|
||||
efficient batch operations.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from datetime import datetime
|
||||
|
||||
from .base import DatabaseConnection
|
||||
from .ops import DataAccessOps, TagListingParts
|
||||
from .ops import (
|
||||
DataAccessOps,
|
||||
LinkExpansionRows,
|
||||
TagListingParts,
|
||||
UpdatedWindow,
|
||||
graph_maintenance_bank_serialization_sql,
|
||||
)
|
||||
from .result import ResultRow
|
||||
|
||||
|
||||
def pg_search_vector_expr(
|
||||
config,
|
||||
*,
|
||||
text_col: str = "text",
|
||||
context_col: str = "context",
|
||||
signals_col: str | None = "text_signals",
|
||||
native_inline: bool = True,
|
||||
) -> str | None:
|
||||
"""SQL expression that builds ``search_vector`` for the configured PG text-search backend.
|
||||
|
||||
Single source of truth shared by ``memory_units`` (the batch insert over the
|
||||
``input_data`` CTE columns and the curation-revert recompute) and
|
||||
``mental_models`` (the knowledge-page writes), so the per-backend tokenization
|
||||
can never drift between the two tables. Returns ``None`` for backends that
|
||||
leave ``search_vector`` unpopulated — pgroonga / pg_textsearch / pg_search
|
||||
index the base text columns directly and keep only a dummy column, so there is
|
||||
nothing to build.
|
||||
|
||||
The ``*_col`` arguments are the SQL for each text source (a column name or a
|
||||
bind placeholder); pass ``signals_col=None`` for a two-column table like
|
||||
``mental_models`` (name + content). Pass ``native_inline=False`` when the
|
||||
table's native ``search_vector`` is a GENERATED column that populates itself
|
||||
(``mental_models``) — writing it inline would fail; only vchord's plain
|
||||
bm25vector column then needs an explicit value.
|
||||
|
||||
``text_search_extension_native_language`` is validated as a PG identifier in
|
||||
``HindsightConfig.validate()``, so embedding it as a SQL literal is safe.
|
||||
"""
|
||||
cols = [text_col, context_col] + ([signals_col] if signals_col is not None else [])
|
||||
combined = " || ' ' || ".join(f"COALESCE({c}, '')" for c in cols)
|
||||
if config.text_search_extension == "vchord":
|
||||
return f"tokenize({combined}, 'llmlingua2')::bm25_catalog.bm25vector"
|
||||
if config.text_search_extension == "native" and native_inline:
|
||||
return f"to_tsvector('{config.text_search_extension_native_language}'::regconfig, {combined})"
|
||||
return None
|
||||
|
||||
|
||||
class PostgreSQLOps(DataAccessOps):
|
||||
"""PostgreSQL-specific data access operations using unnest and LATERAL."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
# Per-table serialization of per-bank vector-index DDL within this
|
||||
# process. Concurrent index DDL on one relation deadlocks by design:
|
||||
# DROP INDEX CONCURRENTLY holds ShareUpdateExclusive while it waits out
|
||||
# every transaction whose snapshot could still see the index — including
|
||||
# other sessions' index DDL queued on that same lock — so many banks
|
||||
# deleted at once form a wait cycle Postgres resolves by killing one.
|
||||
# A session advisory lock would serialize this across processes too, but
|
||||
# advisory locks are banned here (poolers hand sessions around; see the
|
||||
# Database Locking standard). In-process the asyncio lock removes the
|
||||
# cycle outright; across processes the callers' retry-with-backoff
|
||||
# absorbs the (now much rarer) collisions.
|
||||
self._index_ddl_locks: dict[str, asyncio.Lock] = {}
|
||||
|
||||
def _index_ddl_lock(self, table: str) -> asyncio.Lock:
|
||||
return self._index_ddl_locks.setdefault(table, asyncio.Lock())
|
||||
|
||||
@property
|
||||
def uses_observation_sources_table(self) -> bool:
|
||||
return False # PG uses native array ops on source_memory_ids
|
||||
@@ -93,101 +156,39 @@ class PostgreSQLOps(DataAccessOps):
|
||||
config = get_config()
|
||||
table = self._get_mu_table()
|
||||
|
||||
if config.text_search_extension == "vchord":
|
||||
query = f"""
|
||||
WITH input_data AS (
|
||||
SELECT * FROM unnest(
|
||||
$2::text[], $3::vector[], $4::timestamptz[], $5::timestamptz[], $6::timestamptz[], $7::timestamptz[],
|
||||
$8::text[], $9::text[], $10::jsonb[], $11::text[], $12::text[], $13::jsonb[], $14::jsonb[], $15::text[]
|
||||
) AS t(text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, metadata, chunk_id, document_id, tags_json,
|
||||
observation_scopes_json, text_signals)
|
||||
)
|
||||
INSERT INTO {table} (bank_id, text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, metadata, chunk_id, document_id, tags,
|
||||
observation_scopes, text_signals, search_vector)
|
||||
SELECT
|
||||
$1,
|
||||
text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, metadata, chunk_id, document_id,
|
||||
COALESCE(
|
||||
(SELECT array_agg(elem) FROM jsonb_array_elements_text(tags_json) AS elem),
|
||||
'{{}}'::varchar[]
|
||||
),
|
||||
observation_scopes_json,
|
||||
text_signals,
|
||||
tokenize(
|
||||
COALESCE(text, '') || ' ' || COALESCE(context, '') || ' ' || COALESCE(text_signals, ''),
|
||||
'llmlingua2'
|
||||
)::bm25_catalog.bm25vector
|
||||
FROM input_data
|
||||
RETURNING id
|
||||
"""
|
||||
elif config.text_search_extension == "native":
|
||||
# search_vector is a regular tsvector column populated here using the
|
||||
# configured native dictionary. It used to be GENERATED ALWAYS with
|
||||
# a hardcoded 'english', which prevented per-deployment language
|
||||
# configuration. text_search_extension_native_language is validated
|
||||
# in HindsightConfig.validate() as a PG identifier, so embedding it
|
||||
# as a SQL literal is safe.
|
||||
query = f"""
|
||||
WITH input_data AS (
|
||||
SELECT * FROM unnest(
|
||||
$2::text[], $3::vector[], $4::timestamptz[], $5::timestamptz[], $6::timestamptz[], $7::timestamptz[],
|
||||
$8::text[], $9::text[], $10::jsonb[], $11::text[], $12::text[], $13::jsonb[], $14::jsonb[], $15::text[]
|
||||
) AS t(text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, metadata, chunk_id, document_id, tags_json,
|
||||
observation_scopes_json, text_signals)
|
||||
)
|
||||
INSERT INTO {table} (bank_id, text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, metadata, chunk_id, document_id, tags,
|
||||
observation_scopes, text_signals, search_vector)
|
||||
SELECT
|
||||
$1,
|
||||
text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, metadata, chunk_id, document_id,
|
||||
COALESCE(
|
||||
(SELECT array_agg(elem) FROM jsonb_array_elements_text(tags_json) AS elem),
|
||||
'{{}}'::varchar[]
|
||||
),
|
||||
observation_scopes_json,
|
||||
text_signals,
|
||||
to_tsvector(
|
||||
'{config.text_search_extension_native_language}'::regconfig,
|
||||
COALESCE(text, '') || ' ' || COALESCE(context, '') || ' ' || COALESCE(text_signals, '')
|
||||
)
|
||||
FROM input_data
|
||||
RETURNING id
|
||||
"""
|
||||
else:
|
||||
# pg_textsearch, pgroonga, and pg_search: search_vector is a dummy
|
||||
# TEXT column; the actual full-text index operates on the base text
|
||||
# columns directly, so we don't populate search_vector at insert time.
|
||||
query = f"""
|
||||
WITH input_data AS (
|
||||
SELECT * FROM unnest(
|
||||
$2::text[], $3::vector[], $4::timestamptz[], $5::timestamptz[], $6::timestamptz[], $7::timestamptz[],
|
||||
$8::text[], $9::text[], $10::jsonb[], $11::text[], $12::text[], $13::jsonb[], $14::jsonb[], $15::text[]
|
||||
) AS t(text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, metadata, chunk_id, document_id, tags_json,
|
||||
observation_scopes_json, text_signals)
|
||||
)
|
||||
INSERT INTO {table} (bank_id, text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, metadata, chunk_id, document_id, tags,
|
||||
observation_scopes, text_signals)
|
||||
SELECT
|
||||
$1,
|
||||
text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, metadata, chunk_id, document_id,
|
||||
COALESCE(
|
||||
(SELECT array_agg(elem) FROM jsonb_array_elements_text(tags_json) AS elem),
|
||||
'{{}}'::varchar[]
|
||||
),
|
||||
observation_scopes_json,
|
||||
text_signals
|
||||
FROM input_data
|
||||
RETURNING id
|
||||
"""
|
||||
# search_vector is populated inline for backends that store a real vector
|
||||
# (native tsvector, vchord bm25vector). pgroonga / pg_textsearch / pg_search
|
||||
# index the base text columns directly and keep only a dummy column, so the
|
||||
# expression is None and the column is left out of the insert entirely.
|
||||
# Same expression is reused by curation revert (see pg_search_vector_expr).
|
||||
sv_expr = pg_search_vector_expr(config)
|
||||
sv_insert_col = ", search_vector" if sv_expr else ""
|
||||
sv_select_val = f",\n {sv_expr}" if sv_expr else ""
|
||||
query = f"""
|
||||
WITH input_data AS (
|
||||
SELECT * FROM unnest(
|
||||
$2::text[], $3::vector[], $4::timestamptz[], $5::timestamptz[], $6::timestamptz[], $7::timestamptz[],
|
||||
$8::text[], $9::text[], $10::jsonb[], $11::text[], $12::text[], $13::jsonb[], $14::jsonb[], $15::text[]
|
||||
) AS t(text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, metadata, chunk_id, document_id, tags_json,
|
||||
observation_scopes_json, text_signals)
|
||||
)
|
||||
INSERT INTO {table} (bank_id, text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, metadata, chunk_id, document_id, tags,
|
||||
observation_scopes, text_signals{sv_insert_col})
|
||||
SELECT
|
||||
$1,
|
||||
text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, metadata, chunk_id, document_id,
|
||||
COALESCE(
|
||||
(SELECT array_agg(elem) FROM jsonb_array_elements_text(tags_json) AS elem),
|
||||
'{{}}'::varchar[]
|
||||
),
|
||||
observation_scopes_json,
|
||||
text_signals{sv_select_val}
|
||||
FROM input_data
|
||||
RETURNING id
|
||||
"""
|
||||
|
||||
results = await conn.fetch(
|
||||
query,
|
||||
@@ -285,12 +286,22 @@ class PostgreSQLOps(DataAccessOps):
|
||||
bank_id: str,
|
||||
entity_names: list[str],
|
||||
entity_dates: list,
|
||||
entity_kinds: list[str],
|
||||
) -> dict[str, str]:
|
||||
# ORDER BY LOWER(name) so every concurrent batch inserts in the same order
|
||||
# as the conflict target (bank_id, LOWER(canonical_name)). ON CONFLICT DO
|
||||
# NOTHING takes a ShareLock on the inserting transaction of any speculative
|
||||
# row it collides with, so two batches with overlapping names inserting in
|
||||
# different orders deadlock. The caller already sorts by Python's
|
||||
# ``str.lower()``, which agrees with the index for ASCII but not for every
|
||||
# locale (see the Turkish-İ note in entity_resolver) — ordering in SQL makes
|
||||
# the database's own collation the single arbiter for all writers.
|
||||
inserted_rows = await conn.fetch(
|
||||
f"""
|
||||
INSERT INTO {table} (bank_id, canonical_name, first_seen, last_seen, mention_count)
|
||||
SELECT $1, name, COALESCE(event_date, now()), COALESCE(event_date, now()), 0
|
||||
FROM unnest($2::text[], $3::timestamptz[]) AS t(name, event_date)
|
||||
INSERT INTO {table} (bank_id, canonical_name, first_seen, last_seen, mention_count, entity_kind)
|
||||
SELECT $1, name, COALESCE(event_date, now()), COALESCE(event_date, now()), 0, kind
|
||||
FROM unnest($2::text[], $3::timestamptz[], $4::text[]) AS t(name, event_date, kind)
|
||||
ORDER BY LOWER(name)
|
||||
ON CONFLICT (bank_id, LOWER(canonical_name))
|
||||
DO NOTHING
|
||||
RETURNING id, LOWER(canonical_name) AS name_lower
|
||||
@@ -298,6 +309,7 @@ class PostgreSQLOps(DataAccessOps):
|
||||
bank_id,
|
||||
entity_names,
|
||||
entity_dates,
|
||||
entity_kinds,
|
||||
)
|
||||
return {row["name_lower"]: row["id"] for row in inserted_rows}
|
||||
|
||||
@@ -310,7 +322,7 @@ class PostgreSQLOps(DataAccessOps):
|
||||
) -> list[ResultRow]:
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT e.id, LOWER(e.canonical_name) AS name_lower, inputs.input_name
|
||||
SELECT e.id, e.canonical_name, LOWER(e.canonical_name) AS name_lower, inputs.input_name
|
||||
FROM {table} e
|
||||
JOIN (
|
||||
SELECT LOWER(n) AS input_name_lower, n AS input_name
|
||||
@@ -322,6 +334,44 @@ class PostgreSQLOps(DataAccessOps):
|
||||
missing_names,
|
||||
)
|
||||
|
||||
async def bulk_reassert_entities(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
bank_id: str,
|
||||
entity_ids: list[str],
|
||||
canonical_names: list[str],
|
||||
entity_kinds: list[str],
|
||||
) -> None:
|
||||
# One statement, one round-trip (same shape as bulk_insert_links):
|
||||
# * the CTE takes FOR KEY SHARE on every parent that still exists,
|
||||
# held to COMMIT, so a concurrent prune_orphan_entities DELETE blocks
|
||||
# until the caller's unit_entities insert has committed;
|
||||
# * the INSERT re-creates only the parents that were already pruned
|
||||
# (NOT IN locked), carrying the canonical_name resolved in Phase 1.
|
||||
# ON CONFLICT DO NOTHING (no target) keeps the rare case where another
|
||||
# worker recreated the name under a new id from raising — that row stays
|
||||
# absent and its unit link is the sole casualty, never the whole batch.
|
||||
await conn.execute(
|
||||
f"""
|
||||
WITH locked AS (
|
||||
SELECT id FROM {table}
|
||||
WHERE id = ANY($2::uuid[])
|
||||
ORDER BY id
|
||||
FOR KEY SHARE
|
||||
)
|
||||
INSERT INTO {table} (id, bank_id, canonical_name, entity_kind)
|
||||
SELECT t.entity_id, $1, t.canonical_name, t.entity_kind
|
||||
FROM unnest($2::uuid[], $3::text[], $4::text[]) AS t(entity_id, canonical_name, entity_kind)
|
||||
WHERE t.entity_id NOT IN (SELECT id FROM locked)
|
||||
ON CONFLICT DO NOTHING
|
||||
""",
|
||||
bank_id,
|
||||
entity_ids,
|
||||
canonical_names,
|
||||
entity_kinds,
|
||||
)
|
||||
|
||||
async def bulk_insert_unit_entities(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
@@ -357,11 +407,24 @@ class PostgreSQLOps(DataAccessOps):
|
||||
# concurrent caller the same lock order, so conflicting inserts
|
||||
# queue cleanly instead of cycling.
|
||||
sorted_unit_ids = sorted(unit_ids)
|
||||
# DO UPDATE (not DO NOTHING) on a duplicate enqueue — #3034. The SET is a
|
||||
# deliberate no-op that preserves enqueued_at; its only purpose is to take
|
||||
# the existing row's lock. DO NOTHING does NOT lock the conflicting row, so
|
||||
# a mutation that re-enqueues an already-queued unit could not block a
|
||||
# worker from concurrently claiming (deleting) that row and processing the
|
||||
# unit's pre-mutation state; the re-enqueue signal was then silently lost
|
||||
# and the unit's derived links stayed stale with an empty queue. Locking
|
||||
# the row serialises the mutation against the worker's claim for that
|
||||
# (bank_id, unit_id): the worker either waits for the committed post-mutation
|
||||
# state, or (if it claimed first) this INSERT lands a fresh row after the
|
||||
# worker's delete commits. Row locks are acquired in sorted unit_id order,
|
||||
# matching claim_graph_maintenance_batch, so the two never cycle.
|
||||
await conn.execute(
|
||||
f"""
|
||||
INSERT INTO {table} (bank_id, unit_id)
|
||||
SELECT $1, v FROM unnest($2::uuid[]) AS t(v)
|
||||
ON CONFLICT (bank_id, unit_id) DO NOTHING
|
||||
ON CONFLICT (bank_id, unit_id)
|
||||
DO UPDATE SET enqueued_at = {table}.enqueued_at
|
||||
""",
|
||||
bank_id,
|
||||
sorted_unit_ids,
|
||||
@@ -374,16 +437,35 @@ class PostgreSQLOps(DataAccessOps):
|
||||
bank_id: str,
|
||||
limit: int,
|
||||
) -> list[str]:
|
||||
# Ordered locking (#3034). Choose the oldest batch by enqueued_at, but
|
||||
# acquire the row locks in (bank_id, unit_id) order — the same order the
|
||||
# enqueue upsert takes them — so a foreground mutation re-enqueueing an
|
||||
# overlapping unit set can never cycle against a worker draining it. The
|
||||
# `chosen` CTE is MATERIALIZED so the enqueued_at pick is fenced from the
|
||||
# locking clause; `FOR UPDATE OF q ... ORDER BY q.unit_id` then puts
|
||||
# LockRows above the Sort, so locks are taken ascending by unit_id (same
|
||||
# idiom as prune_stale_cooccurrences' #2529 ordered lock). A concurrent
|
||||
# enqueue holding one of these rows blocks this claim until it commits, at
|
||||
# which point the worker deletes and processes the committed state.
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
DELETE FROM {table}
|
||||
WHERE (bank_id, unit_id) IN (
|
||||
WITH chosen AS MATERIALIZED (
|
||||
SELECT bank_id, unit_id FROM {table}
|
||||
WHERE bank_id = $1
|
||||
ORDER BY enqueued_at
|
||||
LIMIT $2
|
||||
),
|
||||
locked AS (
|
||||
SELECT q.bank_id, q.unit_id
|
||||
FROM {table} q
|
||||
JOIN chosen c ON c.bank_id = q.bank_id AND c.unit_id = q.unit_id
|
||||
ORDER BY q.unit_id
|
||||
FOR UPDATE OF q
|
||||
)
|
||||
RETURNING unit_id
|
||||
DELETE FROM {table} q
|
||||
USING locked l
|
||||
WHERE q.bank_id = l.bank_id AND q.unit_id = l.unit_id
|
||||
RETURNING q.unit_id
|
||||
""",
|
||||
bank_id,
|
||||
limit,
|
||||
@@ -425,19 +507,48 @@ class PostgreSQLOps(DataAccessOps):
|
||||
# Scope by joining through entities.bank_id (entity_cooccurrences itself
|
||||
# has no bank_id column — entities don't span banks, so scoping via
|
||||
# entity_id_1 is sufficient).
|
||||
#
|
||||
# Ordered locking (deadlock avoidance, #2529): retain's concurrent
|
||||
# cooccurrence upsert (entity_resolver._flush_pending) locks rows in
|
||||
# sorted (entity_id_1, entity_id_2) order — sorted specifically to give
|
||||
# every writer one consistent lock-acquisition order. A plain
|
||||
# `DELETE ... USING` scans/locks in whatever order the join plan picks,
|
||||
# so it could lock the same rows in the opposite order and cycle. We
|
||||
# instead select the victims in that same sorted order `FOR UPDATE`
|
||||
# first — the locking clause materialises the CTE and places LockRows
|
||||
# above the Sort, so locks are acquired ascending, matching the upsert —
|
||||
# then delete the already-locked rows. Same order on both sides ⇒ no
|
||||
# cycle (the deadlock is prevented, not merely retried). The Pass 2/3
|
||||
# retry wrap in run_graph_maintenance_job stays as a backstop for the
|
||||
# residual paths (FK cascade from prune_orphan_entities, Oracle).
|
||||
#
|
||||
# The staleness predicate is an INTERSECT of the two entities' unit sets
|
||||
# rather than the equivalent `unit_entities u1 JOIN u2 ON u1.unit_id =
|
||||
# u2.unit_id` self-join (#2473): both INTERSECT branches resolve as Index
|
||||
# Only Scans on idx_unit_entities_entity_unit (entity_id, unit_id), so the
|
||||
# per-pair cost is bounded by the two entities' degrees. The self-join let
|
||||
# the planner pick an anti-join that rescanned a high-degree hub entity's
|
||||
# membership set for every pair — 28-30min on a bank with a ~100K-membership
|
||||
# hub, even when zero rows were stale. Don't "simplify" it back.
|
||||
result = await conn.execute(
|
||||
f"""
|
||||
WITH victims AS (
|
||||
SELECT c.entity_id_1, c.entity_id_2
|
||||
FROM {ec_table} c
|
||||
JOIN {entities_table} e ON e.id = c.entity_id_1
|
||||
WHERE e.bank_id = $1
|
||||
AND NOT EXISTS (
|
||||
SELECT unit_id FROM {ue_table} WHERE entity_id = c.entity_id_1
|
||||
INTERSECT
|
||||
SELECT unit_id FROM {ue_table} WHERE entity_id = c.entity_id_2
|
||||
)
|
||||
ORDER BY c.entity_id_1, c.entity_id_2
|
||||
FOR UPDATE OF c
|
||||
)
|
||||
DELETE FROM {ec_table} c
|
||||
USING {entities_table} e
|
||||
WHERE e.id = c.entity_id_1
|
||||
AND e.bank_id = $1
|
||||
AND NOT EXISTS (
|
||||
SELECT 1
|
||||
FROM {ue_table} u1
|
||||
JOIN {ue_table} u2 ON u1.unit_id = u2.unit_id
|
||||
WHERE u1.entity_id = c.entity_id_1
|
||||
AND u2.entity_id = c.entity_id_2
|
||||
)
|
||||
USING victims v
|
||||
WHERE c.entity_id_1 = v.entity_id_1
|
||||
AND c.entity_id_2 = v.entity_id_2
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
@@ -449,11 +560,21 @@ class PostgreSQLOps(DataAccessOps):
|
||||
mu_table: str,
|
||||
unit_ids: list[str],
|
||||
) -> list[ResultRow]:
|
||||
# Cast only canonical UUID text inputs, never the indexed column. The old
|
||||
# ``id::text`` predicate silently ignored malformed, uppercase, braced,
|
||||
# and unhyphenated inputs; filtering before the cast preserves that
|
||||
# behavior while allowing the primary-key index to serve the lookup.
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT id, event_date, fact_type
|
||||
FROM {mu_table}
|
||||
WHERE id::text = ANY($1)
|
||||
WHERE id = ANY(
|
||||
ARRAY(
|
||||
SELECT input.unit_id::uuid
|
||||
FROM unnest($1::text[]) AS input(unit_id)
|
||||
WHERE input.unit_id ~ '^[0-9a-f]{{8}}-[0-9a-f]{{4}}-[0-9a-f]{{4}}-[0-9a-f]{{4}}-[0-9a-f]{{12}}$'
|
||||
)
|
||||
)
|
||||
""",
|
||||
unit_ids,
|
||||
)
|
||||
@@ -522,6 +643,7 @@ class PostgreSQLOps(DataAccessOps):
|
||||
mu_table: str,
|
||||
ue_table: str,
|
||||
per_entity_limit: int,
|
||||
window: UpdatedWindow,
|
||||
) -> str:
|
||||
return f"""
|
||||
seed_entities AS (
|
||||
@@ -541,11 +663,20 @@ class PostgreSQLOps(DataAccessOps):
|
||||
FROM {ue_table} ue_target
|
||||
WHERE ue_target.entity_id = se.entity_id
|
||||
AND ue_target.unit_id != ALL($1::uuid[])
|
||||
-- Filter before applying the cap: candidates from other fact
|
||||
-- types, or outside the recall window, must not consume this
|
||||
-- entity's bounded fan-out.
|
||||
AND EXISTS (
|
||||
SELECT 1
|
||||
FROM {mu_table} mu_target
|
||||
WHERE mu_target.id = ue_target.unit_id
|
||||
AND mu_target.fact_type = $2
|
||||
{window.clause("mu_target")}
|
||||
)
|
||||
ORDER BY ue_target.unit_id DESC
|
||||
LIMIT {per_entity_limit}
|
||||
) t
|
||||
JOIN {mu_table} mu ON mu.id = t.unit_id
|
||||
WHERE mu.fact_type = $2
|
||||
GROUP BY mu.id
|
||||
ORDER BY score DESC
|
||||
LIMIT $3
|
||||
@@ -555,6 +686,7 @@ class PostgreSQLOps(DataAccessOps):
|
||||
self,
|
||||
ml_table: str,
|
||||
mu_table: str,
|
||||
window: UpdatedWindow,
|
||||
) -> str:
|
||||
# Exact v0.5.6 query shape: GROUP BY + MAX(weight) for semantic,
|
||||
# DISTINCT ON for causal.
|
||||
@@ -578,6 +710,7 @@ class PostgreSQLOps(DataAccessOps):
|
||||
AND ml.link_type = 'semantic'
|
||||
AND mu.fact_type = $2
|
||||
AND mu.id != ALL($1::uuid[])
|
||||
{window.clause("mu")}
|
||||
UNION ALL
|
||||
SELECT
|
||||
mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
|
||||
@@ -590,6 +723,7 @@ class PostgreSQLOps(DataAccessOps):
|
||||
AND ml.link_type = 'semantic'
|
||||
AND mu.fact_type = $2
|
||||
AND mu.id != ALL($1::uuid[])
|
||||
{window.clause("mu")}
|
||||
) sem_raw
|
||||
GROUP BY id, text, context, event_date, occurred_start,
|
||||
occurred_end, mentioned_at,
|
||||
@@ -609,6 +743,7 @@ class PostgreSQLOps(DataAccessOps):
|
||||
WHERE ml.from_unit_id = ANY($1::uuid[])
|
||||
AND ml.link_type IN ('causes', 'caused_by', 'enables', 'prevents')
|
||||
AND mu.fact_type = $2
|
||||
{window.clause("mu")}
|
||||
ORDER BY mu.id, ml.weight DESC
|
||||
LIMIT $3
|
||||
)"""
|
||||
@@ -622,8 +757,13 @@ class PostgreSQLOps(DataAccessOps):
|
||||
seed_ids: list,
|
||||
budget: int,
|
||||
per_entity_limit: int,
|
||||
) -> tuple[list[ResultRow], list[ResultRow], list[ResultRow]]:
|
||||
window: UpdatedWindow,
|
||||
) -> LinkExpansionRows:
|
||||
# v0.5.6 array ops: unnest, &&, COUNT(DISTINCT) on source_memory_ids.
|
||||
#
|
||||
# The window bounds the observations that come *back*, not the source facts
|
||||
# traversed to reach them: an observation is in the window when it was itself
|
||||
# written or refreshed there, regardless of how old the facts underneath it are.
|
||||
|
||||
entity_rows = await conn.fetch(
|
||||
f"""
|
||||
@@ -665,11 +805,13 @@ class PostgreSQLOps(DataAccessOps):
|
||||
AND mu.id != ALL($1::uuid[])
|
||||
AND ca.source_ids IS NOT NULL
|
||||
AND mu.source_memory_ids && ca.source_ids
|
||||
{window.clause("mu")}
|
||||
ORDER BY score DESC
|
||||
LIMIT $2
|
||||
""",
|
||||
seed_ids,
|
||||
budget,
|
||||
*window.params,
|
||||
)
|
||||
|
||||
# Exact v0.5.6 query shape: GROUP BY + MAX(weight) for semantic,
|
||||
@@ -691,6 +833,7 @@ class PostgreSQLOps(DataAccessOps):
|
||||
WHERE ml.from_unit_id = ANY($1::uuid[])
|
||||
AND ml.link_type = 'semantic' AND mu.fact_type = 'observation'
|
||||
AND mu.id != ALL($1::uuid[])
|
||||
{window.clause("mu")}
|
||||
UNION ALL
|
||||
SELECT mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
|
||||
mu.occurred_end, mu.mentioned_at, mu.fact_type, mu.document_id,
|
||||
@@ -699,6 +842,7 @@ class PostgreSQLOps(DataAccessOps):
|
||||
WHERE ml.to_unit_id = ANY($1::uuid[])
|
||||
AND ml.link_type = 'semantic' AND mu.fact_type = 'observation'
|
||||
AND mu.id != ALL($1::uuid[])
|
||||
{window.clause("mu")}
|
||||
) sem_raw
|
||||
GROUP BY id, text, context, event_date, occurred_start, occurred_end,
|
||||
mentioned_at, fact_type, document_id, chunk_id, tags, proof_count
|
||||
@@ -713,6 +857,7 @@ class PostgreSQLOps(DataAccessOps):
|
||||
WHERE ml.from_unit_id = ANY($1::uuid[])
|
||||
AND ml.link_type IN ('causes', 'caused_by', 'enables', 'prevents')
|
||||
AND mu.fact_type = 'observation'
|
||||
{window.clause("mu")}
|
||||
ORDER BY mu.id, ml.weight DESC LIMIT $2
|
||||
)
|
||||
SELECT * FROM semantic_expanded
|
||||
@@ -721,11 +866,12 @@ class PostgreSQLOps(DataAccessOps):
|
||||
""",
|
||||
seed_ids,
|
||||
budget,
|
||||
*window.params,
|
||||
)
|
||||
|
||||
semantic_rows = [r for r in sem_causal_rows if r["source"] == "semantic"]
|
||||
causal_rows = [r for r in sem_causal_rows if r["source"] == "causal"]
|
||||
return list(entity_rows), semantic_rows, causal_rows
|
||||
return LinkExpansionRows(entity=list(entity_rows), semantic=semantic_rows, causal=causal_rows)
|
||||
|
||||
def build_tag_listing_parts(self, mu_table: str) -> TagListingParts:
|
||||
return TagListingParts(
|
||||
@@ -745,14 +891,15 @@ class PostgreSQLOps(DataAccessOps):
|
||||
fact_types: dict[str, str],
|
||||
) -> None:
|
||||
escaped = bank_id.replace("'", "''")
|
||||
for ft, suffix in fact_types.items():
|
||||
uid = str(internal_id).replace("-", "")[:16]
|
||||
idx = f"idx_mu_emb_{suffix}_{uid}"
|
||||
await conn.execute(
|
||||
f"CREATE INDEX IF NOT EXISTS {idx} "
|
||||
f"ON {table} {index_clause} "
|
||||
f"WHERE fact_type = '{ft}' AND bank_id = '{escaped}'"
|
||||
)
|
||||
async with self._index_ddl_lock(table):
|
||||
for ft, suffix in fact_types.items():
|
||||
uid = str(internal_id).replace("-", "")[:16]
|
||||
idx = f"idx_mu_emb_{suffix}_{uid}"
|
||||
await conn.execute(
|
||||
f"CREATE INDEX IF NOT EXISTS {idx} "
|
||||
f"ON {table} {index_clause} "
|
||||
f"WHERE fact_type = '{ft}' AND bank_id = '{escaped}'"
|
||||
)
|
||||
|
||||
async def drop_bank_vector_indexes(
|
||||
self,
|
||||
@@ -761,10 +908,19 @@ class PostgreSQLOps(DataAccessOps):
|
||||
internal_id: str,
|
||||
fact_types: dict[str, str],
|
||||
) -> None:
|
||||
for ft, suffix in fact_types.items():
|
||||
uid = str(internal_id).replace("-", "")[:16]
|
||||
idx = f"idx_mu_emb_{suffix}_{uid}"
|
||||
await conn.execute(f"DROP INDEX IF EXISTS {schema}.{idx}")
|
||||
# CONCURRENTLY so the drop takes ShareUpdateExclusive, not ACCESS
|
||||
# EXCLUSIVE, on the shared memory_units table. A plain DROP INDEX blocks
|
||||
# (and deadlocks with) every other bank's concurrent reads/writes on the
|
||||
# table; CONCURRENTLY does not conflict with DML. The caller
|
||||
# (delete_bank) runs this on an autocommit connection after its delete
|
||||
# transaction has committed — CONCURRENTLY cannot run inside a tx.
|
||||
# The lock key must match create_bank_vector_indexes', whose `table`
|
||||
# is the fq name this reconstructs from `schema`.
|
||||
async with self._index_ddl_lock(f"{schema}.memory_units"):
|
||||
for ft, suffix in fact_types.items():
|
||||
uid = str(internal_id).replace("-", "")[:16]
|
||||
idx = f"idx_mu_emb_{suffix}_{uid}"
|
||||
await conn.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}.{idx}")
|
||||
|
||||
def get_entity_resolution_strategy(self) -> str:
|
||||
return "trigram"
|
||||
@@ -894,6 +1050,93 @@ class PostgreSQLOps(DataAccessOps):
|
||||
|
||||
# -- Task claiming operations ------------------------------------------
|
||||
|
||||
async def prune_terminal_operations(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
cutoff: datetime,
|
||||
*,
|
||||
batch_size: int,
|
||||
) -> int:
|
||||
# Lock only the bounded candidate set. SKIP LOCKED lets multiple
|
||||
# workers prune disjoint batches without waiting or double-deleting.
|
||||
# Cancelled children cannot complete parent aggregation, so retain the
|
||||
# parent guard only for completed/failed children. Before removing a
|
||||
# cancelled child, preserve its signal by cancelling a pending parent
|
||||
# in this transaction and refreshing the parent's retention window.
|
||||
candidates = await conn.fetch(
|
||||
f"""
|
||||
SELECT candidate_operation.operation_id
|
||||
FROM {table} candidate_operation
|
||||
WHERE candidate_operation.status IN ('completed', 'failed', 'cancelled')
|
||||
AND candidate_operation.updated_at < $1
|
||||
AND (
|
||||
candidate_operation.status = 'cancelled'
|
||||
OR NOT EXISTS (
|
||||
SELECT 1
|
||||
FROM {table} parent
|
||||
WHERE parent.operation_id = CASE
|
||||
WHEN candidate_operation.result_metadata->>'parent_operation_id'
|
||||
~* '^[0-9a-f]{{8}}-[0-9a-f]{{4}}-[0-9a-f]{{4}}-[0-9a-f]{{4}}-[0-9a-f]{{12}}$'
|
||||
THEN (candidate_operation.result_metadata->>'parent_operation_id')::uuid
|
||||
ELSE NULL
|
||||
END
|
||||
AND parent.bank_id = candidate_operation.bank_id
|
||||
)
|
||||
)
|
||||
ORDER BY candidate_operation.updated_at, candidate_operation.operation_id
|
||||
LIMIT $2
|
||||
FOR UPDATE OF candidate_operation SKIP LOCKED
|
||||
""",
|
||||
cutoff,
|
||||
batch_size,
|
||||
)
|
||||
if not candidates:
|
||||
return 0
|
||||
|
||||
candidate_ids = [row["operation_id"] for row in candidates]
|
||||
await conn.execute(
|
||||
f"""
|
||||
UPDATE {table} parent
|
||||
SET status = 'cancelled',
|
||||
updated_at = now(),
|
||||
completed_at = COALESCE(parent.completed_at, now()),
|
||||
error_message = COALESCE(
|
||||
parent.error_message,
|
||||
'Cancelled because a child operation was cancelled'
|
||||
)
|
||||
WHERE parent.status = 'pending'
|
||||
AND EXISTS (
|
||||
SELECT 1
|
||||
FROM {table} candidate_operation
|
||||
WHERE candidate_operation.operation_id = ANY($1)
|
||||
AND candidate_operation.status = 'cancelled'
|
||||
AND candidate_operation.updated_at < $2
|
||||
AND candidate_operation.bank_id = parent.bank_id
|
||||
AND parent.operation_id = CASE
|
||||
WHEN candidate_operation.result_metadata->>'parent_operation_id'
|
||||
~* '^[0-9a-f]{{8}}-[0-9a-f]{{4}}-[0-9a-f]{{4}}-[0-9a-f]{{4}}-[0-9a-f]{{12}}$'
|
||||
THEN (candidate_operation.result_metadata->>'parent_operation_id')::uuid
|
||||
ELSE NULL
|
||||
END
|
||||
)
|
||||
""",
|
||||
candidate_ids,
|
||||
cutoff,
|
||||
)
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
DELETE FROM {table}
|
||||
WHERE operation_id = ANY($1)
|
||||
AND status IN ('completed', 'failed', 'cancelled')
|
||||
AND updated_at < $2
|
||||
RETURNING operation_id
|
||||
""",
|
||||
candidate_ids,
|
||||
cutoff,
|
||||
)
|
||||
return len(rows)
|
||||
|
||||
async def _claim_consolidation_tasks(
|
||||
self,
|
||||
conn,
|
||||
@@ -1184,13 +1427,14 @@ class PostgreSQLOps(DataAccessOps):
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = $1
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
ORDER BY created_at
|
||||
SELECT o.operation_id, o.operation_type, o.task_payload, o.retry_count
|
||||
FROM {table} o
|
||||
WHERE o.status = 'pending'
|
||||
AND o.task_payload IS NOT NULL
|
||||
AND o.operation_type = $1
|
||||
AND (o.next_retry_at IS NULL OR o.next_retry_at <= NOW())
|
||||
AND {graph_maintenance_bank_serialization_sql(table, "o")}
|
||||
ORDER BY o.created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
@@ -1205,18 +1449,21 @@ class PostgreSQLOps(DataAccessOps):
|
||||
# --- Phase 2: claim from shared pool ---
|
||||
remaining_shared = shared_limit
|
||||
if remaining_shared > 0:
|
||||
# 2a. Non-consolidation tasks
|
||||
# 2a. Non-consolidation tasks. graph_maintenance stays in this
|
||||
# created_at-ordered query — see graph_maintenance_bank_serialization_sql
|
||||
# for why it is a predicate rather than a phase of its own.
|
||||
if claimed_ids:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type != 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND operation_id != ALL($1::uuid[])
|
||||
ORDER BY created_at
|
||||
SELECT o.operation_id, o.operation_type, o.task_payload, o.retry_count
|
||||
FROM {table} o
|
||||
WHERE o.status = 'pending'
|
||||
AND o.task_payload IS NOT NULL
|
||||
AND o.operation_type != 'consolidation'
|
||||
AND (o.next_retry_at IS NULL OR o.next_retry_at <= NOW())
|
||||
AND o.operation_id != ALL($1::uuid[])
|
||||
AND {graph_maintenance_bank_serialization_sql(table, "o")}
|
||||
ORDER BY o.created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
@@ -1226,13 +1473,14 @@ class PostgreSQLOps(DataAccessOps):
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type != 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
ORDER BY created_at
|
||||
SELECT o.operation_id, o.operation_type, o.task_payload, o.retry_count
|
||||
FROM {table} o
|
||||
WHERE o.status = 'pending'
|
||||
AND o.task_payload IS NOT NULL
|
||||
AND o.operation_type != 'consolidation'
|
||||
AND (o.next_retry_at IS NULL OR o.next_retry_at <= NOW())
|
||||
AND {graph_maintenance_bank_serialization_sql(table, "o")}
|
||||
ORDER BY o.created_at
|
||||
LIMIT $1
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
|
||||
@@ -23,6 +23,8 @@ from collections.abc import AsyncIterator
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import Any, NamedTuple
|
||||
|
||||
from .pool_instrumentation import PoolStats, acquire_conn
|
||||
|
||||
|
||||
class _OracleJSONEncoder(json.JSONEncoder):
|
||||
"""JSON encoder that handles datetime and UUID objects."""
|
||||
@@ -78,7 +80,9 @@ _LIKE_ANY_RE = re.compile(r"(\w+)\s+LIKE\s+ANY\s*\(\s*:(\d+)\s*\)", re.IGNORECAS
|
||||
_NOT_LIKE_ALL_RE = re.compile(r"(\w+)\s+NOT\s+LIKE\s+ALL\s*\(\s*:(\d+)\s*\)", re.IGNORECASE)
|
||||
|
||||
_JSON_ARROW_TEXT_RE = re.compile(r'("?\w+"?)\s*->>\s*\'(\w+)\'') # handles both col and "col"
|
||||
_JSON_HAS_KEY_RE = re.compile(r"(\w+)\s*\?\s*'(\w+)'")
|
||||
# Reserved-word columns ("trigger") are already quoted by the time this runs, so the
|
||||
# column group must accept the quoted form too — same shape as the arrow regex above.
|
||||
_JSON_HAS_KEY_RE = re.compile(r"(\"?\w+\"?)\s*\?\s*'(\w+)'")
|
||||
_JSONB_CONTAINS_RE = re.compile(r"(\w+)\s*@>\s*:(\d+)")
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -146,6 +150,7 @@ _JSON_COL_NAMES = {
|
||||
"config",
|
||||
"observation_scopes",
|
||||
"source_memory_ids",
|
||||
"causal_links",
|
||||
"trigger",
|
||||
"http_config",
|
||||
"event_types",
|
||||
@@ -444,9 +449,6 @@ def _rewrite_pg_to_oracle(query: str) -> RewriteResult:
|
||||
if has_for_update:
|
||||
# FOR UPDATE path: use ROWNUM instead of FETCH FIRST.
|
||||
# Extract and remove LIMIT clause, inject ROWNUM into WHERE.
|
||||
def _limit_to_rownum(m):
|
||||
return "" # Remove the LIMIT clause; we'll add ROWNUM below
|
||||
|
||||
limit_val = None
|
||||
limit_match = re.search(r"\bLIMIT\s+(\d+|:\w+)\b", query, re.IGNORECASE)
|
||||
if limit_match:
|
||||
@@ -687,7 +689,6 @@ class OracleConnection(DatabaseConnection):
|
||||
"max_tokens",
|
||||
"priority",
|
||||
"proof_count",
|
||||
"access_count",
|
||||
"importance_score",
|
||||
"decay_factor",
|
||||
"chunk_index",
|
||||
@@ -1242,6 +1243,11 @@ class OracleBackend(DatabaseBackend):
|
||||
def __init__(self) -> None:
|
||||
self._pool: Any = None
|
||||
self._oracledb: Any = None
|
||||
# Oracle pooled sessions retain CURRENT_SCHEMA across checkouts. Cache
|
||||
# SESSION_USER so default-schema acquisitions can explicitly reset a
|
||||
# connection that was previously used for a tenant schema.
|
||||
self._default_schema: str | None = None
|
||||
self._acquire_warn_threshold_s: float = 1.0
|
||||
|
||||
async def initialize(
|
||||
self,
|
||||
@@ -1257,6 +1263,10 @@ class OracleBackend(DatabaseBackend):
|
||||
oracledb = _import_oracledb()
|
||||
self._oracledb = oracledb
|
||||
|
||||
from ...config import get_config
|
||||
|
||||
self._acquire_warn_threshold_s = get_config().db_acquire_warn_threshold_ms / 1000.0
|
||||
|
||||
# Parse URL-format DSN (oracle://user:pass@host:port/service)
|
||||
from urllib.parse import urlparse
|
||||
|
||||
@@ -1277,11 +1287,17 @@ class OracleBackend(DatabaseBackend):
|
||||
logger.info(f"Oracle pool created (min={min_size}, max={max_size})")
|
||||
|
||||
async def shutdown(self) -> None:
|
||||
if self._pool is not None:
|
||||
await self._pool.close(force=True)
|
||||
self._pool = None
|
||||
# Drop the reference before awaiting close() so is_ready flips False for
|
||||
# the whole teardown, not just after it completes (see PostgreSQLBackend).
|
||||
pool, self._pool = self._pool, None
|
||||
if pool is not None:
|
||||
await pool.close(force=True)
|
||||
logger.info("Oracle pool closed")
|
||||
|
||||
@property
|
||||
def is_ready(self) -> bool:
|
||||
return self._pool is not None
|
||||
|
||||
async def _set_session_schema(self, conn: Any) -> None:
|
||||
"""Set the session schema on an Oracle connection.
|
||||
|
||||
@@ -1294,15 +1310,41 @@ class OracleBackend(DatabaseBackend):
|
||||
from ..memory_engine import get_current_schema
|
||||
|
||||
schema = get_current_schema()
|
||||
if schema and schema != "public":
|
||||
cursor = conn.cursor()
|
||||
await cursor.execute(f'ALTER SESSION SET CURRENT_SCHEMA = "{schema}"')
|
||||
await cursor.close()
|
||||
cursor = conn.cursor()
|
||||
try:
|
||||
if self._default_schema is None:
|
||||
await cursor.execute("SELECT SYS_CONTEXT('USERENV', 'SESSION_USER') FROM DUAL")
|
||||
row = await cursor.fetchone()
|
||||
if not row or not row[0]:
|
||||
raise RuntimeError("Oracle did not return SESSION_USER while resetting CURRENT_SCHEMA")
|
||||
self._default_schema = str(row[0])
|
||||
|
||||
target_schema = self._default_schema if not schema or schema == "public" else schema
|
||||
safe_schema = target_schema.replace('"', '""')
|
||||
await cursor.execute(f'ALTER SESSION SET CURRENT_SCHEMA = "{safe_schema}"')
|
||||
finally:
|
||||
# oracledb's AsyncCursor.close() is synchronous (not a coroutine);
|
||||
# awaiting it raises "object NoneType can't be used in 'await'
|
||||
# expression" and aborts every acquire().
|
||||
cursor.close()
|
||||
|
||||
def _pool_stats(self) -> PoolStats | None:
|
||||
"""Snapshot for slow-acquire logs, from oracledb pool attributes."""
|
||||
pool = self._pool
|
||||
if pool is None:
|
||||
return None
|
||||
try:
|
||||
busy = pool.busy
|
||||
return PoolStats(in_use=busy, max=pool.max, idle=pool.opened - busy)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
@asynccontextmanager
|
||||
async def acquire(self) -> AsyncIterator[OracleConnection]:
|
||||
pool = self._ensure_pool()
|
||||
conn = await pool.acquire()
|
||||
conn = await acquire_conn(
|
||||
pool.acquire(), pool_stats=self._pool_stats, warn_threshold_s=self._acquire_warn_threshold_s
|
||||
)
|
||||
try:
|
||||
await self._set_session_schema(conn)
|
||||
yield OracleConnection(conn)
|
||||
@@ -1318,7 +1360,9 @@ class OracleBackend(DatabaseBackend):
|
||||
@asynccontextmanager
|
||||
async def transaction(self) -> AsyncIterator[OracleConnection]:
|
||||
pool = self._ensure_pool()
|
||||
conn = await pool.acquire()
|
||||
conn = await acquire_conn(
|
||||
pool.acquire(), pool_stats=self._pool_stats, warn_threshold_s=self._acquire_warn_threshold_s
|
||||
)
|
||||
try:
|
||||
await self._set_session_schema(conn)
|
||||
yield OracleConnection(conn)
|
||||
|
||||
@@ -0,0 +1,137 @@
|
||||
"""Instrumentation for database connection-pool acquisition.
|
||||
|
||||
asyncpg exposes pool *size* and *idle* counts, but not how many callers are
|
||||
currently **queued waiting** for a connection — and that queue depth is the
|
||||
signal that actually distinguishes a saturated pool from a healthy one. When the
|
||||
pool is exhausted, ``/health`` (which itself acquires a connection to run
|
||||
``SELECT 1``) blocks in ``pool.acquire()`` until a connection frees or the acquire
|
||||
times out, so a liveness probe can fail **with the event loop completely idle**.
|
||||
|
||||
This module tracks the process-wide count of in-flight acquisitions that have not
|
||||
yet obtained a connection, and times each acquire so a slow one logs with full
|
||||
pool stats. It is the DB-side counterpart to ``loop_watchdog`` (which covers loop
|
||||
stalls); together, a stuck ``/health`` can be attributed to either a blocked loop
|
||||
or pool exhaustion from the logs alone.
|
||||
|
||||
The counter is a plain int mutated only from the event-loop thread (asyncpg
|
||||
acquisitions are awaited on the loop), so no lock is needed.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import time
|
||||
from collections.abc import AsyncIterator, Callable
|
||||
from contextlib import asynccontextmanager
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
logger = logging.getLogger("hindsight.db.pool")
|
||||
|
||||
_waiting = 0 # callers currently blocked in pool.acquire(), process-wide
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class PoolStats:
|
||||
"""Point-in-time connection-pool utilization snapshot."""
|
||||
|
||||
in_use: int
|
||||
max: int
|
||||
idle: int
|
||||
|
||||
|
||||
def waiting_count() -> int:
|
||||
"""Number of callers currently blocked waiting to acquire a pooled connection."""
|
||||
return _waiting
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def instrument_acquire(
|
||||
acquire_cm: Any,
|
||||
*,
|
||||
pool_stats: Callable[[], PoolStats | None] | None = None,
|
||||
warn_threshold_s: float,
|
||||
) -> AsyncIterator[Any]:
|
||||
"""Wrap a pool's ``acquire()`` context manager with wait tracking + slow-acquire logging.
|
||||
|
||||
Args:
|
||||
acquire_cm: an async context manager yielding a connection (e.g. the object
|
||||
returned by ``asyncpg.Pool.acquire()``).
|
||||
pool_stats: optional zero-arg callable returning a ``PoolStats`` snapshot for
|
||||
the slow-acquire log line.
|
||||
warn_threshold_s: log a warning when the acquire itself takes at least this long.
|
||||
|
||||
Yields:
|
||||
The acquired connection.
|
||||
"""
|
||||
global _waiting
|
||||
_waiting += 1
|
||||
start = time.monotonic()
|
||||
acquired = False
|
||||
try:
|
||||
async with acquire_cm as conn:
|
||||
acquired = True
|
||||
_waiting -= 1
|
||||
_record_acquire_wait(time.monotonic() - start, pool_stats, warn_threshold_s)
|
||||
yield conn
|
||||
finally:
|
||||
# If __aenter__ raised (acquire timeout / cancellation), we never
|
||||
# decremented above — do it here so the waiter count can't leak.
|
||||
if not acquired:
|
||||
_waiting -= 1
|
||||
|
||||
|
||||
async def acquire_conn(
|
||||
acquire_awaitable: Any,
|
||||
*,
|
||||
pool_stats: Callable[[], PoolStats | None] | None = None,
|
||||
warn_threshold_s: float,
|
||||
) -> Any:
|
||||
"""Await a pool acquire that returns a connection, with wait tracking + slow log.
|
||||
|
||||
For pools whose acquire is ``conn = await pool.acquire()`` (oracledb) rather than
|
||||
an async context manager (asyncpg — use ``instrument_acquire`` for those). The
|
||||
caller is responsible for releasing the returned connection.
|
||||
"""
|
||||
global _waiting
|
||||
_waiting += 1
|
||||
start = time.monotonic()
|
||||
try:
|
||||
conn = await acquire_awaitable
|
||||
finally:
|
||||
_waiting -= 1
|
||||
_record_acquire_wait(time.monotonic() - start, pool_stats, warn_threshold_s)
|
||||
return conn
|
||||
|
||||
|
||||
def _record_acquire_wait(
|
||||
wait_s: float,
|
||||
pool_stats: Callable[[], PoolStats | None] | None,
|
||||
warn_threshold_s: float,
|
||||
) -> None:
|
||||
try:
|
||||
from ...metrics import get_metrics_collector
|
||||
|
||||
get_metrics_collector().record_db_acquire_wait(wait_s)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if wait_s < warn_threshold_s:
|
||||
return
|
||||
|
||||
stats: PoolStats | None = None
|
||||
if pool_stats is not None:
|
||||
try:
|
||||
stats = pool_stats()
|
||||
except Exception:
|
||||
stats = None
|
||||
logger.warning(
|
||||
"slow DB pool acquire: waited %.3fs for a connection "
|
||||
"(in_use=%s max=%s idle=%s waiting=%s). The pool is likely saturated; "
|
||||
"/health can stall on connection acquisition while the event loop is free.",
|
||||
wait_s,
|
||||
stats.in_use if stats else None,
|
||||
stats.max if stats else None,
|
||||
stats.idle if stats else None,
|
||||
_waiting,
|
||||
)
|
||||
@@ -15,6 +15,7 @@ from typing import Any
|
||||
import asyncpg # noqa: F401
|
||||
|
||||
from .base import DatabaseBackend, DatabaseConnection
|
||||
from .pool_instrumentation import PoolStats, instrument_acquire
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -76,6 +77,8 @@ class PostgreSQLBackend(DatabaseBackend):
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._pool: asyncpg.Pool | None = None
|
||||
self._acquire_warn_threshold_s: float = 1.0
|
||||
self._acquire_timeout_s: float | None = None
|
||||
|
||||
async def initialize(
|
||||
self,
|
||||
@@ -88,6 +91,16 @@ class PostgreSQLBackend(DatabaseBackend):
|
||||
statement_cache_size: int = 0,
|
||||
init_callback: Any | None = None,
|
||||
) -> None:
|
||||
from ...config import get_config
|
||||
|
||||
self._acquire_warn_threshold_s = get_config().db_acquire_warn_threshold_ms / 1000.0
|
||||
# Kept for acquire() below: asyncpg's ``timeout`` create_pool kwarg is a
|
||||
# *connect* kwarg (how long establishing a new connection may take), and
|
||||
# ``Pool.acquire()`` defaults to waiting for a free connection forever.
|
||||
# Passing it here alone made HINDSIGHT_API_DB_ACQUIRE_TIMEOUT a no-op for
|
||||
# the wait it names: a pool-exhaustion stall never surfaced as an error,
|
||||
# it just hung (#3002). 0 restores the unbounded behaviour.
|
||||
self._acquire_timeout_s = acquire_timeout if acquire_timeout > 0 else None
|
||||
self._pool = await asyncpg.create_pool(
|
||||
dsn,
|
||||
min_size=min_size,
|
||||
@@ -95,7 +108,12 @@ class PostgreSQLBackend(DatabaseBackend):
|
||||
command_timeout=command_timeout,
|
||||
statement_cache_size=statement_cache_size,
|
||||
timeout=acquire_timeout,
|
||||
# init runs once per new connection; setup runs on every acquire,
|
||||
# after asyncpg's release-time RESET ALL. Passing init_callback as
|
||||
# both keeps the per-connection session GUCs (hnsw.ef_search, etc.)
|
||||
# applied after a connection is reused, not just on first creation.
|
||||
init=init_callback,
|
||||
setup=init_callback,
|
||||
)
|
||||
logger.info(
|
||||
f"PostgreSQL pool created (min={min_size}, max={max_size}, "
|
||||
@@ -103,21 +121,45 @@ class PostgreSQLBackend(DatabaseBackend):
|
||||
)
|
||||
|
||||
async def shutdown(self) -> None:
|
||||
if self._pool is not None:
|
||||
await self._pool.close()
|
||||
self._pool = None
|
||||
# Drop the reference *before* awaiting close(): closing is not
|
||||
# instantaneous, and anything acquiring during that window would
|
||||
# otherwise get an asyncpg "pool is closing" error rather than seeing
|
||||
# is_ready False.
|
||||
pool, self._pool = self._pool, None
|
||||
if pool is not None:
|
||||
await pool.close()
|
||||
logger.info("PostgreSQL pool closed")
|
||||
|
||||
@property
|
||||
def is_ready(self) -> bool:
|
||||
return self._pool is not None
|
||||
|
||||
def _pool_stats(self) -> PoolStats | None:
|
||||
"""Snapshot for slow-acquire logs. in_use = live connections minus idle ones."""
|
||||
pool = self._pool
|
||||
if pool is None:
|
||||
return None
|
||||
idle = pool.get_idle_size()
|
||||
return PoolStats(in_use=pool.get_size() - idle, max=pool.get_max_size(), idle=idle)
|
||||
|
||||
@asynccontextmanager
|
||||
async def acquire(self) -> AsyncIterator[PostgresConnection]:
|
||||
pool = self._ensure_pool()
|
||||
async with pool.acquire() as conn:
|
||||
async with instrument_acquire(
|
||||
pool.acquire(timeout=self._acquire_timeout_s),
|
||||
pool_stats=self._pool_stats,
|
||||
warn_threshold_s=self._acquire_warn_threshold_s,
|
||||
) as conn:
|
||||
yield PostgresConnection(conn)
|
||||
|
||||
@asynccontextmanager
|
||||
async def transaction(self) -> AsyncIterator[PostgresConnection]:
|
||||
pool = self._ensure_pool()
|
||||
async with pool.acquire() as conn:
|
||||
async with instrument_acquire(
|
||||
pool.acquire(timeout=self._acquire_timeout_s),
|
||||
pool_stats=self._pool_stats,
|
||||
warn_threshold_s=self._acquire_warn_threshold_s,
|
||||
) as conn:
|
||||
async with conn.transaction():
|
||||
yield PostgresConnection(conn)
|
||||
|
||||
|
||||
@@ -164,35 +164,6 @@ class BudgetedOperation:
|
||||
"""
|
||||
return BudgetedPool(pool, self)
|
||||
|
||||
async def acquire_many(
|
||||
self,
|
||||
pool: Any,
|
||||
count: int,
|
||||
) -> AsyncIterator[list[Any]]:
|
||||
"""
|
||||
Acquire multiple connections within the budget.
|
||||
|
||||
Note: This acquires connections sequentially to respect the budget.
|
||||
For parallel acquisition, use multiple acquire() calls with asyncio.gather().
|
||||
This method is intended for use with raw asyncpg pools only, not DatabaseBackend.
|
||||
|
||||
Args:
|
||||
pool: asyncpg connection pool (raw pool only)
|
||||
count: Number of connections to acquire
|
||||
|
||||
Yields:
|
||||
List of database connections
|
||||
"""
|
||||
connections = []
|
||||
try:
|
||||
for _ in range(count):
|
||||
conn = await pool.acquire()
|
||||
connections.append(conn)
|
||||
yield connections
|
||||
finally:
|
||||
for conn in connections:
|
||||
await pool.release(conn)
|
||||
|
||||
|
||||
# Global default manager instance
|
||||
_default_manager: ConnectionBudgetManager | None = None
|
||||
|
||||
@@ -4,6 +4,7 @@ Database utility functions for connection management with retry logic.
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import random
|
||||
import time
|
||||
from collections.abc import AsyncIterator
|
||||
from contextlib import AsyncExitStack, asynccontextmanager
|
||||
@@ -16,6 +17,20 @@ DEFAULT_MAX_RETRIES = 3
|
||||
DEFAULT_BASE_DELAY = 0.5 # seconds
|
||||
DEFAULT_MAX_DELAY = 5.0 # seconds
|
||||
|
||||
|
||||
def _backoff_delay(attempt: int, base_delay: float, max_delay: float) -> float:
|
||||
"""Exponential backoff with equal jitter.
|
||||
|
||||
Deterministic backoff makes concurrent retriers wake in lock-step and
|
||||
re-collide on the very same rows, re-triggering the deadlock they just
|
||||
backed off from. "Equal jitter" — half the window fixed, half random —
|
||||
keeps a floor (so we don't hot-spin) while decorrelating the wake-ups, so
|
||||
two contenders that deadlocked together are very unlikely to retry in sync.
|
||||
"""
|
||||
ceil = min(base_delay * (2**attempt), max_delay)
|
||||
return ceil / 2 + random.uniform(0, ceil / 2)
|
||||
|
||||
|
||||
# Retryable exception types (checked by class name to avoid hard imports)
|
||||
_RETRYABLE_EXCEPTION_NAMES = frozenset(
|
||||
{
|
||||
@@ -78,7 +93,7 @@ async def retry_with_backoff(
|
||||
raise
|
||||
last_exception = e
|
||||
if attempt < max_retries:
|
||||
delay = min(base_delay * (2**attempt), max_delay)
|
||||
delay = _backoff_delay(attempt, base_delay, max_delay)
|
||||
if type(e).__name__ == "DeadlockDetectedError" or _is_oracle_deadlock(e):
|
||||
logger.warning(
|
||||
"Deadlock detected during parallel document processing — "
|
||||
@@ -136,7 +151,7 @@ async def acquire_with_retry(backend_or_pool: Any, max_retries: int = DEFAULT_MA
|
||||
if not _is_retryable(e):
|
||||
raise
|
||||
if attempt < max_retries:
|
||||
delay = min(DEFAULT_BASE_DELAY * (2**attempt), DEFAULT_MAX_DELAY)
|
||||
delay = _backoff_delay(attempt, DEFAULT_BASE_DELAY, DEFAULT_MAX_DELAY)
|
||||
logger.warning(
|
||||
f"Database acquire failed (attempt {attempt + 1}/{max_retries + 1}): {e}. "
|
||||
f"Retrying in {delay:.1f}s..."
|
||||
|
||||
@@ -48,6 +48,12 @@ from ..config import (
|
||||
ENV_LLM_API_KEY,
|
||||
)
|
||||
from .bank_attribution import apply_bank_attribution
|
||||
from .local_device import (
|
||||
release_local_inference_memory,
|
||||
resolve_model_device_type,
|
||||
select_local_device,
|
||||
)
|
||||
from .tei_retry import tei_retry_delay
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -136,7 +142,13 @@ class LocalSTEmbeddings(Embeddings):
|
||||
The embedding dimension is auto-detected from the model.
|
||||
"""
|
||||
|
||||
def __init__(self, model_name: str | None = None, force_cpu: bool = False, trust_remote_code: bool = False):
|
||||
def __init__(
|
||||
self,
|
||||
model_name: str | None = None,
|
||||
force_cpu: bool = False,
|
||||
trust_remote_code: bool = False,
|
||||
allow_mps: bool = False,
|
||||
):
|
||||
"""
|
||||
Initialize local SentenceTransformers embeddings.
|
||||
|
||||
@@ -148,12 +160,17 @@ class LocalSTEmbeddings(Embeddings):
|
||||
trust_remote_code: Allow loading models with custom code (security risk).
|
||||
Required for some models with custom architectures.
|
||||
Default: False (disabled for security)
|
||||
allow_mps: Opt in to the Apple Silicon MPS GPU. Disabled by default
|
||||
because MPS leaks memory under variable-length workloads
|
||||
(see engine/local_device.py). Default: False
|
||||
"""
|
||||
self.model_name = model_name or DEFAULT_EMBEDDINGS_LOCAL_MODEL
|
||||
self.force_cpu = force_cpu
|
||||
self.trust_remote_code = trust_remote_code
|
||||
self.allow_mps = allow_mps
|
||||
self._model = None
|
||||
self._dimension: int | None = None
|
||||
self._device_type: str = "cpu"
|
||||
|
||||
@property
|
||||
def provider_name(self) -> str:
|
||||
@@ -180,31 +197,11 @@ class LocalSTEmbeddings(Embeddings):
|
||||
|
||||
logger.info(f"Embeddings: initializing local provider with model {self.model_name}")
|
||||
|
||||
# Determine device based on hardware availability.
|
||||
# We always set low_cpu_mem_usage=False to prevent lazy loading (meta tensors)
|
||||
# which can cause issues when accelerate is installed but no GPU is available.
|
||||
import torch
|
||||
|
||||
# Force CPU mode if configured (used in daemon mode to avoid MPS/XPC issues on macOS)
|
||||
if self.force_cpu:
|
||||
device = "cpu"
|
||||
logger.info("Embeddings: forcing CPU mode")
|
||||
else:
|
||||
# Check for GPU (CUDA), Apple Silicon (MPS), or Intel XPU
|
||||
# Wrap in try-except to gracefully handle any device detection issues
|
||||
# (e.g., in CI environments or when PyTorch is built without GPU support)
|
||||
device = "cpu" # Default to CPU
|
||||
try:
|
||||
has_gpu = torch.cuda.is_available() or (
|
||||
hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
|
||||
)
|
||||
# Intel Arc XPU support — torch.xpu is available when the XPU build is loaded
|
||||
if not has_gpu and hasattr(torch, "xpu"):
|
||||
has_gpu = torch.xpu.is_available()
|
||||
if has_gpu:
|
||||
device = None # Let sentence-transformers auto-detect GPU/MPS/XPU
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to detect GPU/MPS/XPU, falling back to CPU: {e}")
|
||||
# Determine device based on hardware availability. We always set
|
||||
# low_cpu_mem_usage=False to prevent lazy loading (meta tensors) which can
|
||||
# cause issues when accelerate is installed but no GPU is available.
|
||||
# MPS is opt-in (allow_mps) — see engine/local_device.py for why.
|
||||
device = select_local_device(self.force_cpu, self.allow_mps)
|
||||
|
||||
# Suppress verbose transformers warnings during model loading
|
||||
# This suppresses the "UNEXPECTED" warnings from BertModel which are harmless
|
||||
@@ -231,7 +228,8 @@ class LocalSTEmbeddings(Embeddings):
|
||||
transformers_logger.setLevel(original_level)
|
||||
|
||||
self._dimension = self._model.get_sentence_embedding_dimension()
|
||||
logger.info(f"Embeddings: local provider initialized (dim: {self._dimension})")
|
||||
self._device_type = resolve_model_device_type(self._model)
|
||||
logger.info(f"Embeddings: local provider initialized (dim: {self._dimension}, device: {self._device_type})")
|
||||
|
||||
def encode(self, texts: list[str]) -> list[list[float]]:
|
||||
"""
|
||||
@@ -243,11 +241,49 @@ class LocalSTEmbeddings(Embeddings):
|
||||
Returns:
|
||||
List of embedding vectors
|
||||
"""
|
||||
return self._encode_local(texts)
|
||||
|
||||
def encode_query(self, texts: list[str]) -> list[list[float]]:
|
||||
return self._encode_local(texts, input_type="query")
|
||||
|
||||
def encode_documents(self, texts: list[str]) -> list[list[float]]:
|
||||
return self._encode_local(texts, input_type="document")
|
||||
|
||||
def _encode_local(
|
||||
self, texts: list[str], input_type: Literal["query", "document"] | None = None
|
||||
) -> list[list[float]]:
|
||||
if self._model is None:
|
||||
raise RuntimeError("Embeddings not initialized. Call initialize() first.")
|
||||
|
||||
embeddings = self._model.encode(texts, convert_to_numpy=True, show_progress_bar=False)
|
||||
return [emb.tolist() for emb in embeddings]
|
||||
try:
|
||||
# Delegate to SentenceTransformers' own asymmetric entry points rather than
|
||||
# prefixing here: they apply whatever prompts the model ships with (and route
|
||||
# the task for models exposing a Router module), so asymmetric models such as
|
||||
# Qwen3-Embedding get their configured query prompt without Hindsight carrying
|
||||
# per-model prefix config the way the ONNX provider has to. Models that declare
|
||||
# no prompts are unaffected — SentenceTransformers defaults them to empty
|
||||
# strings and skips prompt handling entirely, so this is byte-identical to
|
||||
# encode() for e.g. the default BAAI/bge-small-en-v1.5.
|
||||
# encode_query/encode_document exist only in sentence-transformers >= 5.0,
|
||||
# which is why local-ml pins that floor.
|
||||
if input_type == "query":
|
||||
encode = self._model.encode_query
|
||||
elif input_type == "document":
|
||||
encode = self._model.encode_document
|
||||
else:
|
||||
encode = self._model.encode
|
||||
embeddings = encode(texts, convert_to_numpy=True, show_progress_bar=False)
|
||||
return [emb.tolist() for emb in embeddings]
|
||||
finally:
|
||||
# Only reclaim the GPU allocator pool here, and only when actually on a
|
||||
# GPU (opt-in MPS/CUDA/XPU). encode() runs in tight retain loops, so a
|
||||
# gc.collect()/malloc_trim on every call is too costly on the CPU default
|
||||
# — and unnecessary: refcounting frees the small transient buffers
|
||||
# immediately and the allocator reuses them for the next batch. (The
|
||||
# reranker keeps its per-batch heap trim for the #1717 CPU case; it runs
|
||||
# on the lighter recall path.) See engine/local_device.py.
|
||||
if self._device_type != "cpu":
|
||||
release_local_inference_memory(self._device_type)
|
||||
|
||||
|
||||
class OnnxEmbeddings(Embeddings):
|
||||
@@ -478,7 +514,7 @@ class RemoteTEIEmbeddings(Embeddings):
|
||||
response = self._client.post(url, **kwargs)
|
||||
response.raise_for_status()
|
||||
return response
|
||||
except (httpx.ConnectError, httpx.ReadTimeout, httpx.WriteTimeout) as e:
|
||||
except (httpx.ConnectError, httpx.ConnectTimeout, httpx.ReadTimeout, httpx.WriteTimeout) as e:
|
||||
last_error = e
|
||||
if attempt < self.max_retries:
|
||||
logger.warning(
|
||||
@@ -487,13 +523,20 @@ class RemoteTEIEmbeddings(Embeddings):
|
||||
time.sleep(delay)
|
||||
delay *= 2 # Exponential backoff
|
||||
except httpx.HTTPStatusError as e:
|
||||
# Retry on 5xx server errors
|
||||
if e.response.status_code >= 500 and attempt < self.max_retries:
|
||||
# TEI uses 429 as normal overload backpressure. Retry it with
|
||||
# the same bounded budget as transient server errors.
|
||||
if (e.response.status_code == 429 or e.response.status_code >= 500) and attempt < self.max_retries:
|
||||
last_error = e
|
||||
logger.warning(
|
||||
f"TEI server error (attempt {attempt + 1}/{self.max_retries + 1}): {e}. Retrying in {delay}s..."
|
||||
sleep_delay = tei_retry_delay(
|
||||
e.response,
|
||||
delay,
|
||||
request_timeout=self.timeout,
|
||||
)
|
||||
time.sleep(delay)
|
||||
logger.warning(
|
||||
f"TEI transient error (attempt {attempt + 1}/{self.max_retries + 1}): {e}. "
|
||||
f"Retrying in {sleep_delay:.2f}s..."
|
||||
)
|
||||
time.sleep(sleep_delay)
|
||||
delay *= 2
|
||||
else:
|
||||
raise
|
||||
@@ -1585,6 +1628,7 @@ def create_embeddings_from_env() -> Embeddings:
|
||||
model_name=config.embeddings_local_model,
|
||||
force_cpu=config.embeddings_local_force_cpu,
|
||||
trust_remote_code=config.embeddings_local_trust_remote_code,
|
||||
allow_mps=config.embeddings_local_allow_mps,
|
||||
)
|
||||
elif provider == "onnx":
|
||||
return OnnxEmbeddings(
|
||||
|
||||
@@ -6,13 +6,16 @@ to disambiguate entities across memory units.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import heapq
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
from collections import defaultdict
|
||||
from collections.abc import Iterator
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import UTC, datetime
|
||||
from difflib import SequenceMatcher
|
||||
from typing import Any, Final
|
||||
from typing import Any, Final, cast
|
||||
|
||||
from .db_utils import acquire_with_retry
|
||||
from .memory_engine import fq_table
|
||||
@@ -25,6 +28,7 @@ from .retain.entity_labels import (
|
||||
from .retain.entity_labels import (
|
||||
parse_entity_labels as _parse_entity_labels,
|
||||
)
|
||||
from .retain.types import ResolvedEntity
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -36,6 +40,119 @@ class _EntityToCreate:
|
||||
idx: int
|
||||
name: str
|
||||
event_date: datetime | None
|
||||
# Label entities (from entity_labels config) are never fuzzy-merged in-batch — their
|
||||
# canonical names are user-defined (e.g. "use:use-001") and must stay distinct (GH-1558).
|
||||
# Also stored on the row as entities.entity_kind so label rows stay out of the
|
||||
# partial trigram index (#3208).
|
||||
is_label: bool = False
|
||||
|
||||
|
||||
@dataclass
|
||||
class _SimilarNamePair:
|
||||
"""A pair of in-batch new-entity names judged similar enough to be the same entity."""
|
||||
|
||||
name_a: str
|
||||
name_b: str
|
||||
|
||||
|
||||
# The in-batch dedup pass is O(N^2) over the batch's *new* names. It is sub-millisecond for a
|
||||
# normal retain (a handful of new entities) but scales quadratically — measured on the retain hot
|
||||
# path: ~0.8ms at 100 names, ~5ms at 250, ~22ms at 500, ~81ms at 1000. So skip it past this many
|
||||
# unique new names and log rather than silently degrade; the cap sits well above any realistic
|
||||
# single-retain new-entity count while bounding the tail.
|
||||
_INTRABATCH_MAX_NAMES = 250
|
||||
|
||||
# A pg_trgm "word" is a maximal run of alphanumerics (Unicode letters/digits, underscore excluded);
|
||||
# everything else (space, punctuation, emoji) is a separator. This is why decoration variants like
|
||||
# "Wren <emoji>" collapse to the same trigram set.
|
||||
_TRGM_WORD = re.compile(r"[^\W_]+", re.UNICODE)
|
||||
|
||||
|
||||
def _trigram_set(text: str) -> set[str]:
|
||||
"""Trigrams of ``text`` the way PostgreSQL pg_trgm generates them: lowercase, split into words,
|
||||
pad each word with two leading + one trailing blank, and take every 3-char window."""
|
||||
trigrams: set[str] = set()
|
||||
for word in _TRGM_WORD.findall(text.lower()):
|
||||
padded = f" {word} "
|
||||
for i in range(len(padded) - 2):
|
||||
trigrams.add(padded[i : i + 3])
|
||||
return trigrams
|
||||
|
||||
|
||||
def _trigram_similarity(a: str, b: str) -> float:
|
||||
"""pg_trgm ``similarity(a, b)`` computed in-memory — the Jaccard index of the trigram sets.
|
||||
|
||||
Verified byte-for-byte against Postgres pg_trgm across emoji / accent / CJK / hyphen /
|
||||
apostrophe cases (issue #3107), so the merge cutoff calibrated on pg_trgm transfers exactly.
|
||||
Doing it in Python keeps the in-batch dedup off the retain transaction's DB connection and makes
|
||||
it backend-agnostic (Postgres, Oracle, and the pg_trgm-absent "full" fallback all behave alike).
|
||||
"""
|
||||
ta, tb = _trigram_set(a), _trigram_set(b)
|
||||
intersection = len(ta & tb)
|
||||
union = len(ta) + len(tb) - intersection
|
||||
return intersection / union if union else 0.0
|
||||
|
||||
|
||||
def _find_intrabatch_similar_pairs(names: list[str], threshold: float) -> list[_SimilarNamePair]:
|
||||
"""Every pair of ``names`` whose in-memory trigram similarity meets ``threshold``. O(N^2) over a
|
||||
small, capped set of new names — pure CPU, no DB round-trip."""
|
||||
trigrams = [_trigram_set(n) for n in names]
|
||||
pairs: list[_SimilarNamePair] = []
|
||||
for i in range(len(names)):
|
||||
ti = trigrams[i]
|
||||
for j in range(i + 1, len(names)):
|
||||
tj = trigrams[j]
|
||||
intersection = len(ti & tj)
|
||||
union = len(ti) + len(tj) - intersection
|
||||
if union and intersection / union >= threshold:
|
||||
pairs.append(_SimilarNamePair(name_a=names[i], name_b=names[j]))
|
||||
return pairs
|
||||
|
||||
|
||||
def _cluster_new_entity_names(
|
||||
rep_by_lower: dict[str, str],
|
||||
count_by_lower: dict[str, int],
|
||||
pairs: list[_SimilarNamePair],
|
||||
) -> dict[str, str]:
|
||||
"""Union-find the similar-name pairs into clusters and pick one canonical name each.
|
||||
|
||||
Args:
|
||||
rep_by_lower: lowercase name -> a representative original-case spelling of it.
|
||||
count_by_lower: lowercase name -> how many mentions carry it (for canonical choice).
|
||||
pairs: name pairs judged similar (order/case irrelevant; compared lowercased).
|
||||
|
||||
Returns:
|
||||
lowercase name -> canonical original-case name for its cluster. Singletons map to
|
||||
themselves, so the caller can look up every member uniformly.
|
||||
"""
|
||||
parent: dict[str, str] = {nl: nl for nl in rep_by_lower}
|
||||
|
||||
def find(x: str) -> str:
|
||||
while parent[x] != x:
|
||||
parent[x] = parent[parent[x]] # path halving
|
||||
x = parent[x]
|
||||
return x
|
||||
|
||||
for pair in pairs:
|
||||
a, b = pair.name_a.lower(), pair.name_b.lower()
|
||||
if a in parent and b in parent:
|
||||
ra, rb = find(a), find(b)
|
||||
if ra != rb:
|
||||
parent[ra] = rb
|
||||
|
||||
clusters: dict[str, list[str]] = {}
|
||||
for nl in rep_by_lower:
|
||||
clusters.setdefault(find(nl), []).append(nl)
|
||||
|
||||
canonical_by_member: dict[str, str] = {}
|
||||
for members in clusters.values():
|
||||
# Canonical = most-mentioned, then shortest, then lexicographically smallest — a
|
||||
# deterministic pick that prefers the plainest spelling in the cluster.
|
||||
canonical_lower = min(members, key=lambda nl: (-count_by_lower[nl], len(rep_by_lower[nl]), rep_by_lower[nl]))
|
||||
canonical_name = rep_by_lower[canonical_lower]
|
||||
for nl in members:
|
||||
canonical_by_member[nl] = canonical_name
|
||||
return canonical_by_member
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -75,6 +192,22 @@ def _later_date(a: datetime | None, b: datetime | None) -> datetime | None:
|
||||
return a if a > b else b
|
||||
|
||||
|
||||
def _canonical_cooccurrence_pairs(entity_list: list[str]) -> Iterator[tuple[str, str]]:
|
||||
"""Yield each distinct pair of ``entity_list`` as ``(a, b)`` with ``a < b``.
|
||||
|
||||
Canonical ordering matches the entity_cooccurrences PK and check constraint.
|
||||
The pair is ordered into fresh locals rather than by swapping the loop
|
||||
variables: ``entity_id_1`` is the outer iterate, so swapping it would leak
|
||||
into the remaining inner iterations and build later pairs off the wrong
|
||||
element.
|
||||
"""
|
||||
for i, entity_id_1 in enumerate(entity_list):
|
||||
for entity_id_2 in entity_list[i + 1 :]:
|
||||
if entity_id_1 == entity_id_2:
|
||||
continue
|
||||
yield (entity_id_1, entity_id_2) if entity_id_1 < entity_id_2 else (entity_id_2, entity_id_1)
|
||||
|
||||
|
||||
@dataclass
|
||||
class _CooccurrencePair:
|
||||
"""A (entity_id_1, entity_id_2) pair observed in a retain batch (for post-txn flush)."""
|
||||
@@ -92,6 +225,32 @@ class _CooccurrencePair:
|
||||
_nlp = None
|
||||
|
||||
|
||||
# Candidates scored between cooperative yields to the event loop. Scoring is
|
||||
# synchronous CPU (one SequenceMatcher per candidate, ~50µs), so a batch with a
|
||||
# large candidate set would otherwise hold the loop thread for minutes — health
|
||||
# probes time out and the orchestrator kills the worker mid-op (GH-3211).
|
||||
# 256 candidates ≈ 13ms of work between yields.
|
||||
_SCORING_YIELD_EVERY: Final = 256
|
||||
|
||||
|
||||
def _cheap_rank_key(entity_text_lower: str, candidate: tuple[Any, str, Any, datetime | None, int | None]) -> tuple:
|
||||
"""Ordering key (not a multi-value return) approximating match quality cheaply.
|
||||
|
||||
Used only to truncate oversized candidate sets: the fuzzy strategies already
|
||||
cap and pre-rank in SQL by real similarity, so this is the backstop for sets
|
||||
built without a score (the "full" strategy's substring matching). Ranks an
|
||||
exact match first, then a close name length, then a well-established entity —
|
||||
all O(1) per candidate, unlike the SequenceMatcher pass it protects.
|
||||
"""
|
||||
name_lower = candidate[1].lower()
|
||||
return (
|
||||
0 if name_lower == entity_text_lower else 1,
|
||||
abs(len(name_lower) - len(entity_text_lower)),
|
||||
-(candidate[4] or 0),
|
||||
candidate[1],
|
||||
)
|
||||
|
||||
|
||||
class EntityResolver:
|
||||
"""
|
||||
Resolves entities to canonical IDs with disambiguation.
|
||||
@@ -102,6 +261,8 @@ class EntityResolver:
|
||||
pool: Any,
|
||||
entity_lookup: str = "full",
|
||||
entity_resolution_batch_size: int = 100,
|
||||
intrabatch_merge_similarity: float = 0.5,
|
||||
entity_resolution_max_candidates: int = 200,
|
||||
):
|
||||
"""
|
||||
Initialize entity resolver.
|
||||
@@ -113,12 +274,22 @@ class EntityResolver:
|
||||
similar candidates per entity name (much faster for large banks).
|
||||
entity_resolution_batch_size: Number of unique entity names to include
|
||||
in each pg_trgm candidate lookup query.
|
||||
intrabatch_merge_similarity: pg_trgm similarity at/above which two new
|
||||
names created by the same retain are merged into one entity.
|
||||
entity_resolution_max_candidates: Max candidates scored per entity
|
||||
mention. Scoring is a synchronous SequenceMatcher call per
|
||||
candidate, so an unbounded candidate set turns one resolution
|
||||
batch into minutes of event-loop-blocking CPU (GH-3211).
|
||||
"""
|
||||
self.pool = pool
|
||||
self.entity_lookup = entity_lookup
|
||||
if entity_resolution_batch_size < 1:
|
||||
raise ValueError("entity_resolution_batch_size must be >= 1")
|
||||
self.entity_resolution_batch_size = entity_resolution_batch_size
|
||||
self._intrabatch_merge_similarity = intrabatch_merge_similarity
|
||||
if entity_resolution_max_candidates < 1:
|
||||
raise ValueError("entity_resolution_max_candidates must be >= 1")
|
||||
self.entity_resolution_max_candidates = entity_resolution_max_candidates
|
||||
self._pg_trgm_checked = False
|
||||
# Backend-specific operations — accessed via pool.ops (Django pattern).
|
||||
self._ops = pool.ops if pool is not None else None
|
||||
@@ -222,6 +393,19 @@ class EntityResolver:
|
||||
"""Split values into fixed-size batches."""
|
||||
return [values[i : i + size] for i in range(0, len(values), size)]
|
||||
|
||||
@staticmethod
|
||||
def _label_texts(entity_texts: list[str], taxonomy_lookup: set[str] | None, labels_cfg) -> set[str]:
|
||||
"""Subset of entity_texts that are label entities (resolved by exact match only).
|
||||
|
||||
Only gate on the config, not on the lookup set: text/map groups have no
|
||||
fixed vocabulary, so a config with only those groups builds an EMPTY
|
||||
lookup — its labels are classified by key prefix inside
|
||||
``is_label_entity``, and gating on the set would miss them entirely.
|
||||
"""
|
||||
if not labels_cfg:
|
||||
return set()
|
||||
return {t for t in entity_texts if _is_label_entity(t, labels_cfg, taxonomy_lookup or set())}
|
||||
|
||||
async def resolve_entities_batch(
|
||||
self,
|
||||
bank_id: str,
|
||||
@@ -230,7 +414,7 @@ class EntityResolver:
|
||||
unit_event_date,
|
||||
conn=None,
|
||||
entity_labels: list | None = None,
|
||||
) -> list[str]:
|
||||
) -> list[ResolvedEntity]:
|
||||
"""
|
||||
Resolve multiple entities in batch (MUCH faster than sequential).
|
||||
|
||||
@@ -245,7 +429,8 @@ class EntityResolver:
|
||||
conn: Optional connection to use (if None, acquires from pool)
|
||||
|
||||
Returns:
|
||||
List of entity IDs in same order as input
|
||||
Resolved entity identities (id + stored canonical name) in the same
|
||||
order as input.
|
||||
"""
|
||||
if not entities_data:
|
||||
return []
|
||||
@@ -271,7 +456,7 @@ class EntityResolver:
|
||||
unit_event_date,
|
||||
taxonomy_lookup: set[str] | None = None,
|
||||
labels_cfg=None,
|
||||
) -> list[str]:
|
||||
) -> list[ResolvedEntity]:
|
||||
if self.entity_lookup == "trigram":
|
||||
# Route to backend-specific fuzzy strategy.
|
||||
# Non-PG backends (Oracle) use UTL_MATCH instead of pg_trgm.
|
||||
@@ -311,7 +496,7 @@ class EntityResolver:
|
||||
unit_event_date,
|
||||
taxonomy_lookup: set[str] | None = None,
|
||||
labels_cfg=None,
|
||||
) -> list[str]:
|
||||
) -> list[ResolvedEntity]:
|
||||
"""Original strategy: load all bank entities then match in Python."""
|
||||
# Query ALL candidates for this bank
|
||||
all_entities = await conn.fetch(
|
||||
@@ -395,7 +580,7 @@ class EntityResolver:
|
||||
unit_event_date,
|
||||
taxonomy_lookup: set[str] | None = None,
|
||||
labels_cfg=None,
|
||||
) -> list[str]:
|
||||
) -> list[ResolvedEntity]:
|
||||
"""
|
||||
Trigram strategy: fetch only similar candidates per entity name using pg_trgm.
|
||||
|
||||
@@ -405,40 +590,78 @@ class EntityResolver:
|
||||
"""
|
||||
entity_texts = list(set(e["text"] for e in entities_data))
|
||||
|
||||
# Fetch candidates for unique entity texts in bounded batches.
|
||||
# Label entities resolve by exact match only (their canonical names are
|
||||
# user-defined and must not be fuzzy-merged). Probing them via the trigram
|
||||
# index only returns similar-but-distinct label values that are always
|
||||
# discarded, and that wasted work grows with the number of values a label
|
||||
# accumulates. Resolve label texts with an exact lookup on the unique
|
||||
# (bank_id, LOWER(canonical_name)) index and only fuzzy-match the rest.
|
||||
label_set = self._label_texts(entity_texts, taxonomy_lookup, labels_cfg)
|
||||
label_texts = [t for t in entity_texts if t in label_set]
|
||||
fuzzy_texts = [t for t in entity_texts if t not in label_set]
|
||||
|
||||
rows = []
|
||||
|
||||
# Exact, index-only lookup for label texts.
|
||||
for entity_text_batch in self._chunked(label_texts, self.entity_resolution_batch_size):
|
||||
rows.extend(
|
||||
await conn.fetch(
|
||||
f"""
|
||||
SELECT e.id, e.canonical_name, e.metadata, e.last_seen, e.mention_count,
|
||||
q.query_text
|
||||
FROM unnest($2::text[]) AS q(query_text)
|
||||
JOIN {fq_table("entities")} e ON (
|
||||
e.bank_id = $1
|
||||
AND LOWER(e.canonical_name) = LOWER(q.query_text)
|
||||
)
|
||||
""",
|
||||
bank_id,
|
||||
entity_text_batch,
|
||||
)
|
||||
)
|
||||
|
||||
# Fetch candidates for the remaining texts in bounded batches.
|
||||
# Uses the GIN trigram index on LOWER(canonical_name) for case-insensitive
|
||||
# similarity lookup. Previous version also had LIKE '%...' substring fallbacks,
|
||||
# but those forced full sequential scans of the entities table and caused
|
||||
# TimeoutErrors on banks with 10k+ entities. Lowering the similarity threshold
|
||||
# to 0.15 (from default 0.3) catches most substring relationships while
|
||||
# staying fully index-based.
|
||||
await conn.execute("SET pg_trgm.similarity_threshold = 0.15")
|
||||
try:
|
||||
rows = []
|
||||
for entity_text_batch in self._chunked(entity_texts, self.entity_resolution_batch_size):
|
||||
rows.extend(
|
||||
await conn.fetch(
|
||||
f"""
|
||||
SELECT DISTINCT ON (e.id)
|
||||
e.id, e.canonical_name, e.metadata, e.last_seen, e.mention_count,
|
||||
q.query_text
|
||||
FROM unnest($2::text[]) AS q(query_text)
|
||||
JOIN {fq_table("entities")} e ON (
|
||||
e.bank_id = $1
|
||||
AND LOWER(e.canonical_name) % LOWER(q.query_text)
|
||||
)
|
||||
""",
|
||||
bank_id,
|
||||
entity_text_batch,
|
||||
)
|
||||
# TimeoutErrors on banks with 10k+ entities. The pg_trgm similarity threshold
|
||||
# that governs the `%` operator is applied once at pool-connection setup
|
||||
# (HINDSIGHT_API_ENTITY_TRGM_SIMILARITY_THRESHOLD), so it is not toggled here.
|
||||
# ``entity_kind != 'label'`` matches the predicate of the partial trigram
|
||||
# index (label rows are exact-match-only, so they can never be a
|
||||
# legitimate fuzzy result — without the filter they only inflate the
|
||||
# candidate set and get discarded in the bitmap recheck, #3208). The
|
||||
# clause must textually match the index predicate for the planner to
|
||||
# choose the partial index, so it stays inside the LATERAL's WHERE
|
||||
# alongside the `%` operator rather than moving out to the outer join.
|
||||
#
|
||||
# The LATERAL keeps only the best `max_candidates` per query text: on a bank
|
||||
# with many near-identical names a single probe can otherwise return
|
||||
# thousands of rows, and every one of them costs a SequenceMatcher call in
|
||||
# _resolve_from_candidates (GH-3211). Ranking by pg_trgm similarity — which
|
||||
# the index scan computes anyway — keeps the truncation at the noise end.
|
||||
for entity_text_batch in self._chunked(fuzzy_texts, self.entity_resolution_batch_size):
|
||||
rows.extend(
|
||||
await conn.fetch(
|
||||
f"""
|
||||
SELECT c.id, c.canonical_name, c.metadata, c.last_seen, c.mention_count,
|
||||
q.query_text
|
||||
FROM unnest($2::text[]) AS q(query_text)
|
||||
CROSS JOIN LATERAL (
|
||||
SELECT e.id, e.canonical_name, e.metadata, e.last_seen, e.mention_count
|
||||
FROM {fq_table("entities")} e
|
||||
WHERE e.bank_id = $1
|
||||
AND e.entity_kind != 'label'
|
||||
AND LOWER(e.canonical_name) % LOWER(q.query_text)
|
||||
ORDER BY similarity(LOWER(e.canonical_name), LOWER(q.query_text)) DESC, e.id
|
||||
LIMIT $3
|
||||
) c
|
||||
""",
|
||||
bank_id,
|
||||
entity_text_batch,
|
||||
self.entity_resolution_max_candidates,
|
||||
)
|
||||
finally:
|
||||
# asyncpg returns connections to the pool with session state intact,
|
||||
# so the lowered threshold would leak to future borrowers without RESET.
|
||||
try:
|
||||
await conn.execute("RESET pg_trgm.similarity_threshold")
|
||||
except Exception:
|
||||
logger.warning("Failed to reset pg_trgm similarity threshold after candidate lookup", exc_info=True)
|
||||
)
|
||||
|
||||
# Group candidates by query_text
|
||||
all_candidates: dict[str, list] = {t: [] for t in entity_texts}
|
||||
@@ -499,7 +722,7 @@ class EntityResolver:
|
||||
unit_event_date: datetime | None,
|
||||
taxonomy_lookup: set[str] | None = None,
|
||||
labels_cfg=None,
|
||||
) -> list[str]:
|
||||
) -> list[ResolvedEntity]:
|
||||
"""
|
||||
Oracle strategy: fetch similar candidates using UTL_MATCH.JARO_WINKLER_SIMILARITY.
|
||||
|
||||
@@ -511,6 +734,14 @@ class EntityResolver:
|
||||
entity_texts = list(set(e["text"] for e in entities_data))
|
||||
entities_table = fq_table("entities")
|
||||
|
||||
# Label entities resolve by exact match only, so the fuzzy Jaro-Winkler
|
||||
# join only returns similar-but-distinct label values that are always
|
||||
# discarded. Resolve label texts with an exact lookup on the unique
|
||||
# (bank_id, LOWER(canonical_name)) index and only fuzzy-match the rest.
|
||||
label_set = self._label_texts(entity_texts, taxonomy_lookup, labels_cfg)
|
||||
label_texts = [t for t in entity_texts if t in label_set]
|
||||
fuzzy_texts = [t for t in entity_texts if t not in label_set]
|
||||
|
||||
try:
|
||||
# Batch entity texts into bounded sub-queries using JSON_TABLE to
|
||||
# expand the list into rows. UTL_MATCH.JARO_WINKLER_SIMILARITY
|
||||
@@ -518,7 +749,7 @@ class EntityResolver:
|
||||
# Bounded batches mirror the PG trigram path so very wide retain
|
||||
# batches don't time out a single JOIN on large banks.
|
||||
rows = []
|
||||
for entity_text_batch in self._chunked(entity_texts, self.entity_resolution_batch_size):
|
||||
for entity_text_batch in self._chunked(label_texts, self.entity_resolution_batch_size):
|
||||
rows.extend(
|
||||
await conn.fetch(
|
||||
f"""
|
||||
@@ -527,13 +758,46 @@ class EntityResolver:
|
||||
FROM JSON_TABLE($2, '$[*]' COLUMNS (query_text VARCHAR2(4000) PATH '$')) q
|
||||
JOIN {entities_table} e ON (
|
||||
e.bank_id = $1
|
||||
AND UTL_MATCH.JARO_WINKLER_SIMILARITY(LOWER(e.canonical_name), LOWER(q.query_text)) > 70
|
||||
AND LOWER(e.canonical_name) = LOWER(q.query_text)
|
||||
)
|
||||
""",
|
||||
bank_id,
|
||||
json.dumps(entity_text_batch),
|
||||
)
|
||||
)
|
||||
# Only the best `max_candidates` per query text are returned: each
|
||||
# candidate costs a synchronous SequenceMatcher call downstream, so an
|
||||
# unbounded fuzzy match set blocks the event loop for minutes
|
||||
# (GH-3211). Ranking by the same Jaro-Winkler score the join already
|
||||
# computes keeps the truncation at the noise end.
|
||||
for entity_text_batch in self._chunked(fuzzy_texts, self.entity_resolution_batch_size):
|
||||
rows.extend(
|
||||
await conn.fetch(
|
||||
f"""
|
||||
SELECT id, canonical_name, metadata, last_seen, mention_count, query_text
|
||||
FROM (
|
||||
SELECT e.id, e.canonical_name, e.metadata, e.last_seen, e.mention_count,
|
||||
q.query_text,
|
||||
ROW_NUMBER() OVER (
|
||||
PARTITION BY q.query_text
|
||||
ORDER BY UTL_MATCH.JARO_WINKLER_SIMILARITY(
|
||||
LOWER(e.canonical_name), LOWER(q.query_text)
|
||||
) DESC, e.id
|
||||
) AS rn
|
||||
FROM JSON_TABLE($2, '$[*]' COLUMNS (query_text VARCHAR2(4000) PATH '$')) q
|
||||
JOIN {entities_table} e ON (
|
||||
e.bank_id = $1
|
||||
AND e.entity_kind != 'label'
|
||||
AND UTL_MATCH.JARO_WINKLER_SIMILARITY(LOWER(e.canonical_name), LOWER(q.query_text)) > 70
|
||||
)
|
||||
)
|
||||
WHERE rn <= $3
|
||||
""",
|
||||
bank_id,
|
||||
json.dumps(entity_text_batch),
|
||||
self.entity_resolution_max_candidates,
|
||||
)
|
||||
)
|
||||
except Exception as e:
|
||||
# UTL_MATCH may not be available (ORA-06550, ORA-00904, etc.)
|
||||
# Catch broadly because Oracle error types vary depending on driver.
|
||||
@@ -597,6 +861,38 @@ class EntityResolver:
|
||||
labels_cfg,
|
||||
)
|
||||
|
||||
def _intrabatch_canonical_map(self, entities_to_create: list[_EntityToCreate]) -> dict[str, str]:
|
||||
"""Map each non-label new name (lowercased) to its cluster's canonical spelling.
|
||||
|
||||
Uses in-memory trigram similarity (``_trigram_similarity``, verified equal to Postgres
|
||||
pg_trgm), so it is backend-agnostic — no DB round-trip on the retain hot path, and it runs
|
||||
identically on PostgreSQL, Oracle, and the pg_trgm-absent "full" fallback. Label entities
|
||||
are excluded so distinct label values stay separate (GH-1558).
|
||||
"""
|
||||
rep_by_lower: dict[str, str] = {}
|
||||
count_by_lower: dict[str, int] = {}
|
||||
for e in entities_to_create:
|
||||
if e.is_label:
|
||||
continue
|
||||
name_lower = e.name.lower()
|
||||
rep_by_lower.setdefault(name_lower, e.name)
|
||||
count_by_lower[name_lower] = count_by_lower.get(name_lower, 0) + 1
|
||||
|
||||
if len(rep_by_lower) < 2:
|
||||
return {} # nothing to compare
|
||||
if len(rep_by_lower) > _INTRABATCH_MAX_NAMES:
|
||||
logger.warning(
|
||||
"Skipping in-batch entity dedup: %d unique new names exceeds the %d cap "
|
||||
"(O(N^2) trigram comparison); same-batch surface variants may not be merged.",
|
||||
len(rep_by_lower),
|
||||
_INTRABATCH_MAX_NAMES,
|
||||
)
|
||||
return {}
|
||||
pairs = _find_intrabatch_similar_pairs(list(rep_by_lower.values()), self._intrabatch_merge_similarity)
|
||||
if not pairs:
|
||||
return {}
|
||||
return _cluster_new_entity_names(rep_by_lower, count_by_lower, pairs)
|
||||
|
||||
async def _resolve_from_candidates(
|
||||
self,
|
||||
conn,
|
||||
@@ -607,13 +903,19 @@ class EntityResolver:
|
||||
cooccurrence_map: dict[str, set[str]],
|
||||
taxonomy_lookup: set[str] | None = None,
|
||||
labels_cfg=None,
|
||||
) -> list[str]:
|
||||
) -> list[ResolvedEntity]:
|
||||
"""Shared scoring + upsert logic used by both lookup strategies."""
|
||||
|
||||
# Resolve each entity using pre-fetched candidates
|
||||
entity_ids = [None] * len(entities_data)
|
||||
# Resolve each entity using pre-fetched candidates. A slot stays None
|
||||
# only if find-or-create fails to produce a row for a mention (a DB
|
||||
# inconsistency); it surfaces as a clear error at the reassert boundary
|
||||
# rather than a silent NOT NULL violation deeper in Phase 2.
|
||||
resolved: list[ResolvedEntity | None] = [None] * len(entities_data)
|
||||
entities_to_update: list[_EntityStat] = []
|
||||
entities_to_create: list[_EntityToCreate] = []
|
||||
# Candidates scored since the last yield, counted across mentions so a
|
||||
# batch of many small candidate sets yields as often as one large set.
|
||||
scored_since_yield = 0
|
||||
|
||||
for idx, entity_data in enumerate(entities_data):
|
||||
entity_text = entity_data["text"]
|
||||
@@ -623,41 +925,85 @@ class EntityResolver:
|
||||
|
||||
candidates = all_candidates.get(entity_text, [])
|
||||
|
||||
# Backstop truncation for candidate sets that were not capped at the
|
||||
# source (the "full" strategy matches substrings in Python). The fuzzy
|
||||
# strategies already return at most this many rows per query text, so
|
||||
# this is normally a no-op.
|
||||
if len(candidates) > self.entity_resolution_max_candidates:
|
||||
logger.debug(
|
||||
"Truncating %d candidates to %d for entity text %r",
|
||||
len(candidates),
|
||||
self.entity_resolution_max_candidates,
|
||||
entity_text,
|
||||
)
|
||||
entity_text_lower_for_rank = entity_text.lower()
|
||||
candidates = heapq.nsmallest(
|
||||
self.entity_resolution_max_candidates,
|
||||
candidates,
|
||||
key=lambda c: _cheap_rank_key(entity_text_lower_for_rank, c),
|
||||
)
|
||||
|
||||
# Label entities (from entity_labels config) use exact matching only.
|
||||
# Their canonical names are user-defined (e.g., "use:use-001"),
|
||||
# so fuzzy resolution must NOT merge distinct label values that
|
||||
# happen to be textually similar (GH-1558).
|
||||
is_label = bool(
|
||||
labels_cfg and taxonomy_lookup and _is_label_entity(entity_text, labels_cfg, taxonomy_lookup)
|
||||
)
|
||||
# happen to be textually similar (GH-1558). Don't gate on the
|
||||
# lookup set — it is empty for text/map-only configs, whose labels
|
||||
# classify by key prefix (see _label_texts).
|
||||
is_label = bool(labels_cfg and _is_label_entity(entity_text, labels_cfg, taxonomy_lookup or set()))
|
||||
|
||||
if not candidates:
|
||||
# Will create new entity
|
||||
entities_to_create.append(_EntityToCreate(idx=idx, name=entity_text, event_date=entity_event_date))
|
||||
entities_to_create.append(
|
||||
_EntityToCreate(idx=idx, name=entity_text, event_date=entity_event_date, is_label=is_label)
|
||||
)
|
||||
continue
|
||||
|
||||
if is_label:
|
||||
# Exact case-insensitive match only for label entities
|
||||
exact_match = None
|
||||
exact_match: ResolvedEntity | None = None
|
||||
entity_text_lower = entity_text.lower()
|
||||
for candidate_id, canonical_name, metadata, last_seen, mention_count in candidates:
|
||||
if canonical_name.lower() == entity_text_lower:
|
||||
exact_match = candidate_id
|
||||
exact_match = ResolvedEntity(
|
||||
entity_id=candidate_id, canonical_name=canonical_name, entity_kind="label"
|
||||
)
|
||||
break
|
||||
if exact_match:
|
||||
entity_ids[idx] = exact_match
|
||||
entities_to_update.append(_EntityStat(entity_id=exact_match, event_date=entity_event_date))
|
||||
resolved[idx] = exact_match
|
||||
entities_to_update.append(
|
||||
_EntityStat(entity_id=exact_match.entity_id, event_date=entity_event_date)
|
||||
)
|
||||
else:
|
||||
entities_to_create.append(_EntityToCreate(idx=idx, name=entity_text, event_date=entity_event_date))
|
||||
entities_to_create.append(
|
||||
_EntityToCreate(idx=idx, name=entity_text, event_date=entity_event_date, is_label=True)
|
||||
)
|
||||
continue
|
||||
|
||||
# Score candidates
|
||||
best_candidate = None
|
||||
best_candidate: ResolvedEntity | None = None
|
||||
best_score = 0.0
|
||||
|
||||
nearby_entity_set = {e["text"].lower() for e in nearby_entities if e["text"] != entity_text}
|
||||
|
||||
for candidate_id, canonical_name, metadata, last_seen, mention_count in candidates:
|
||||
# Hand the loop back periodically so /health (and every other task
|
||||
# on this worker) still gets scheduled while a wide batch scores.
|
||||
# Counted before the label skip below, so a candidate list that is
|
||||
# entirely labels still yields — the skip runs _is_label_entity per
|
||||
# row, which is cheap but not free.
|
||||
scored_since_yield += 1
|
||||
if scored_since_yield >= _SCORING_YIELD_EVERY:
|
||||
scored_since_yield = 0
|
||||
await asyncio.sleep(0)
|
||||
|
||||
# A label row can never be a fuzzy-match target (#1558): the
|
||||
# trigram/UTL_MATCH probes exclude them in SQL via entity_kind,
|
||||
# but the "full" fallback strategy loads every bank entity, so
|
||||
# a textually-close label value could still outscore the 0.6
|
||||
# threshold here (e.g. "topic empathy" vs "topic:empathy").
|
||||
if labels_cfg and _is_label_entity(canonical_name, labels_cfg, taxonomy_lookup or set()):
|
||||
continue
|
||||
|
||||
score = 0.0
|
||||
|
||||
# 1. Name similarity (0-0.5)
|
||||
@@ -685,17 +1031,17 @@ class EntityResolver:
|
||||
|
||||
if score > best_score:
|
||||
best_score = score
|
||||
best_candidate = candidate_id
|
||||
best_candidate = ResolvedEntity(entity_id=candidate_id, canonical_name=canonical_name)
|
||||
|
||||
# Apply unified threshold
|
||||
threshold = 0.6
|
||||
|
||||
if best_score > threshold:
|
||||
entity_ids[idx] = best_candidate
|
||||
entities_to_update.append(_EntityStat(entity_id=best_candidate, event_date=entity_event_date))
|
||||
if best_score > threshold and best_candidate is not None:
|
||||
resolved[idx] = best_candidate
|
||||
entities_to_update.append(_EntityStat(entity_id=best_candidate.entity_id, event_date=entity_event_date))
|
||||
else:
|
||||
entities_to_create.append(
|
||||
_EntityToCreate(idx=idx, name=entity_data["text"], event_date=entity_event_date)
|
||||
_EntityToCreate(idx=idx, name=entity_data["text"], event_date=entity_event_date, is_label=is_label)
|
||||
)
|
||||
|
||||
# Existing entities: IDs already known from the candidate SELECT above.
|
||||
@@ -707,24 +1053,46 @@ class EntityResolver:
|
||||
# ON CONFLICT DO NOTHING returns nothing for rows that conflicted; we handle
|
||||
# that rare case with a fallback SELECT.
|
||||
if entities_to_create:
|
||||
# Group by lowercase name — deduplicate within the batch.
|
||||
# Fuzzy-cluster the NON-label names about to be created so same-batch surface
|
||||
# variants (case/emoji/suffix/typo of one name) collapse to a single entity. Without
|
||||
# this, resolution only compares against already-persisted rows, so the first sighting
|
||||
# of each variant in a batch always creates a distinct entity (issue #3107). Labels are
|
||||
# excluded and keep exact grouping.
|
||||
canonical_by_member = self._intrabatch_canonical_map(entities_to_create)
|
||||
|
||||
@dataclass
|
||||
class _NameGroup:
|
||||
name: str
|
||||
event_date: datetime | None
|
||||
is_label: bool
|
||||
indices: list[int] = field(default_factory=list)
|
||||
|
||||
groups: dict[str, _NameGroup] = {}
|
||||
for e in entities_to_create:
|
||||
name_lower = e.name.lower()
|
||||
if name_lower not in groups:
|
||||
groups[name_lower] = _NameGroup(name=e.name, event_date=e.event_date)
|
||||
groups[name_lower].indices.append(e.idx)
|
||||
# Non-label variants fold into their cluster's canonical name; everything else
|
||||
# (labels, singletons) keys on itself, preserving the prior exact-match behavior.
|
||||
canonical = canonical_by_member.get(e.name.lower(), e.name)
|
||||
key = canonical.lower()
|
||||
group = groups.get(key)
|
||||
if group is None:
|
||||
# Labels key on themselves and the dedup pass only clusters
|
||||
# non-label names, so the first member's is_label holds for
|
||||
# every member of the group.
|
||||
group = _NameGroup(name=canonical, event_date=e.event_date, is_label=e.is_label)
|
||||
groups[key] = group
|
||||
elif e.event_date is not None and (group.event_date is None or e.event_date < group.event_date):
|
||||
# Keep the earliest event_date across the cluster ("first seen").
|
||||
group.event_date = e.event_date
|
||||
group.indices.append(e.idx)
|
||||
|
||||
# Sort by lowercase name for deterministic ordering.
|
||||
sorted_groups = sorted(groups.items())
|
||||
entity_names = [g.name for _, g in sorted_groups]
|
||||
entity_dates = [g.event_date for _, g in sorted_groups]
|
||||
entity_kinds = ["label" if g.is_label else "regular" for _, g in sorted_groups]
|
||||
# Stored canonical name per lowercase key, so a resurrected parent
|
||||
# keeps the name it was created/matched with rather than a fallback.
|
||||
canonical_by_name = {name_lower: g.name for name_lower, g in sorted_groups}
|
||||
|
||||
# INSERT ... ON CONFLICT DO NOTHING — no row lock on already-existing entities.
|
||||
# mention_count starts at 0 here; flush_pending_stats() is the sole source of
|
||||
@@ -737,6 +1105,7 @@ class EntityResolver:
|
||||
bank_id,
|
||||
entity_names,
|
||||
entity_dates,
|
||||
entity_kinds,
|
||||
)
|
||||
|
||||
# Fallback SELECT for names that conflicted (another worker won the race).
|
||||
@@ -759,11 +1128,14 @@ class EntityResolver:
|
||||
)
|
||||
for row in existing_rows:
|
||||
id_by_name[row["name_lower"]] = row["id"]
|
||||
canonical_by_name[row["name_lower"]] = row["canonical_name"]
|
||||
# Also index by Python's lower() of the original input name so the
|
||||
# assignment loop (which uses Python-lowercased keys) finds it even
|
||||
# when Python and the database produce different lowercase strings.
|
||||
if "input_name" in row:
|
||||
id_by_name[row["input_name"].lower()] = row["id"]
|
||||
input_name_lower = row["input_name"].lower()
|
||||
id_by_name[input_name_lower] = row["id"]
|
||||
canonical_by_name[input_name_lower] = row["canonical_name"]
|
||||
|
||||
# Assign entity IDs back and queue one stat per original mention so that
|
||||
# flush_pending_stats() increments mention_count by the true mention count,
|
||||
@@ -771,21 +1143,74 @@ class EntityResolver:
|
||||
for name_lower, g in sorted_groups:
|
||||
entity_id = id_by_name.get(name_lower)
|
||||
if entity_id:
|
||||
canonical_name = canonical_by_name.get(name_lower, g.name)
|
||||
kind = "label" if g.is_label else "regular"
|
||||
for original_idx in g.indices:
|
||||
entity_ids[original_idx] = entity_id
|
||||
pending.append(_EntityStat(entity_id=entity_id, event_date=g.event_date))
|
||||
resolved[original_idx] = ResolvedEntity(
|
||||
entity_id=entity_id, canonical_name=canonical_name, entity_kind=kind
|
||||
)
|
||||
pending.append(_EntityStat(entity_id=str(entity_id), event_date=g.event_date))
|
||||
|
||||
# Accumulate into the resolver's pending list; the orchestrator flushes
|
||||
# these with await entity_resolver.flush_pending_stats() after the txn.
|
||||
key = self._task_key()
|
||||
self._pending_stats.setdefault(key, []).extend(pending)
|
||||
|
||||
return entity_ids
|
||||
missing = [i for i, entity in enumerate(resolved) if entity is None]
|
||||
if missing:
|
||||
raise RuntimeError(
|
||||
f"Entity resolution produced no row for {len(missing)} mention(s) "
|
||||
f"(indices {missing[:5]}); refusing to link units to a missing parent."
|
||||
)
|
||||
return cast(list[ResolvedEntity], resolved)
|
||||
|
||||
async def reassert_entities_batch(
|
||||
self,
|
||||
bank_id: str,
|
||||
resolved_entities: list[ResolvedEntity],
|
||||
conn,
|
||||
) -> None:
|
||||
"""Lock (and, if pruned, re-create) resolved parents before linking units.
|
||||
|
||||
Phase-1 resolution and the Phase-2 ``unit_entities`` insert run on
|
||||
different transactions. In the gap, ``prune_orphan_entities`` can delete
|
||||
a just-resolved parent — it legitimately has no ``unit_entities`` row
|
||||
yet — and the Phase-2 FK insert then fails, dropping the whole batch as
|
||||
non-retryable (silent memory loss, #2662).
|
||||
|
||||
Called on the Phase-2 connection immediately before
|
||||
``link_units_to_entities_batch``, this locks the parents that still
|
||||
exist (so the pruner blocks until we commit) and re-inserts any that
|
||||
already vanished, in one round-trip. An entity referenced by a live unit
|
||||
is by definition not an orphan, so resurrecting it is correct.
|
||||
"""
|
||||
# Deduplicate by id and lock in a stable order so concurrent reasserts
|
||||
# acquire row locks consistently (same convention as bulk_insert_links).
|
||||
seen: set[str] = set()
|
||||
unique: list[ResolvedEntity] = []
|
||||
for entity in sorted(resolved_entities, key=lambda e: e.entity_id):
|
||||
if entity.entity_id in seen:
|
||||
continue
|
||||
seen.add(entity.entity_id)
|
||||
unique.append(entity)
|
||||
|
||||
if not unique:
|
||||
return
|
||||
|
||||
await self._ops.bulk_reassert_entities(
|
||||
conn,
|
||||
fq_table("entities"),
|
||||
bank_id,
|
||||
[entity.entity_id for entity in unique],
|
||||
[entity.canonical_name for entity in unique],
|
||||
[entity.entity_kind for entity in unique],
|
||||
)
|
||||
|
||||
async def link_units_to_entities_batch(
|
||||
self,
|
||||
unit_entity_pairs: list[tuple[str, str]] | list[tuple[str, str, datetime | None]],
|
||||
conn=None,
|
||||
bank_id: str | None = None,
|
||||
):
|
||||
"""
|
||||
Link multiple memory units to entities in batch (MUCH faster than sequential).
|
||||
@@ -813,22 +1238,32 @@ class EntityResolver:
|
||||
|
||||
if conn is None:
|
||||
async with acquire_with_retry(self.pool) as conn:
|
||||
return await self._link_units_to_entities_batch_impl(conn, normalized)
|
||||
return await self._link_units_to_entities_batch_impl(conn, normalized, bank_id)
|
||||
else:
|
||||
return await self._link_units_to_entities_batch_impl(conn, normalized)
|
||||
return await self._link_units_to_entities_batch_impl(conn, normalized, bank_id)
|
||||
|
||||
async def _link_units_to_entities_batch_impl(self, conn, unit_entity_pairs: list[tuple[str, str, datetime | None]]):
|
||||
async def _link_units_to_entities_batch_impl(
|
||||
self, conn, unit_entity_pairs: list[tuple[str, str, datetime | None]], bank_id: str | None = None
|
||||
):
|
||||
# Sorted bulk insert to prevent deadlocks from inconsistent lock ordering
|
||||
# across concurrent transactions on the unit_entities unique index.
|
||||
sorted_pairs = sorted(unit_entity_pairs, key=lambda t: (t[0], t[1]))
|
||||
unit_ids = [p[0] for p in sorted_pairs]
|
||||
entity_ids = [p[1] for p in sorted_pairs]
|
||||
|
||||
await self._ops.bulk_insert_unit_entities(
|
||||
conn,
|
||||
fq_table("unit_entities"),
|
||||
unit_ids,
|
||||
entity_ids,
|
||||
# The unit→entity posting belongs to whoever stores the memory, so the
|
||||
# memories store records it. Co-occurrence below is separate and unaffected:
|
||||
# it references only `entities`, which stays in Postgres either way, and is
|
||||
# read by the entity-graph endpoint and by resolution's disambiguation signal.
|
||||
from .memories import get_memories
|
||||
|
||||
await get_memories().record_unit_entities(
|
||||
conn=conn,
|
||||
ops=self._ops,
|
||||
fq_table=fq_table,
|
||||
bank_id=bank_id,
|
||||
unit_ids=unit_ids,
|
||||
entity_ids=entity_ids,
|
||||
)
|
||||
|
||||
# Build maps keyed by unit_id:
|
||||
@@ -853,20 +1288,12 @@ class EntityResolver:
|
||||
for unit_id, entity_ids in unit_to_entities.items():
|
||||
entity_list = list(entity_ids)
|
||||
event_date = unit_event_date.get(unit_id)
|
||||
for i, entity_id_1 in enumerate(entity_list):
|
||||
for entity_id_2 in entity_list[i + 1 :]:
|
||||
if entity_id_1 == entity_id_2:
|
||||
continue
|
||||
# Canonical ordering (entity_id_1 < entity_id_2) matches the
|
||||
# entity_cooccurrences PK and check constraint.
|
||||
if entity_id_1 > entity_id_2:
|
||||
entity_id_1, entity_id_2 = entity_id_2, entity_id_1
|
||||
key = (entity_id_1, entity_id_2)
|
||||
prev = cooccurrence_pairs.get(key, _SENTINEL_MISSING)
|
||||
if prev is _SENTINEL_MISSING:
|
||||
cooccurrence_pairs[key] = event_date
|
||||
else:
|
||||
cooccurrence_pairs[key] = _later_date(prev, event_date)
|
||||
for key in _canonical_cooccurrence_pairs(entity_list):
|
||||
prev = cooccurrence_pairs.get(key, _SENTINEL_MISSING)
|
||||
if prev is _SENTINEL_MISSING:
|
||||
cooccurrence_pairs[key] = event_date
|
||||
else:
|
||||
cooccurrence_pairs[key] = _later_date(prev, event_date)
|
||||
|
||||
# Accumulate co-occurrence pairs for post-transaction flush.
|
||||
# The actual INSERT/UPDATE is deferred to flush_pending_stats() to avoid
|
||||
@@ -878,58 +1305,3 @@ class EntityResolver:
|
||||
_CooccurrencePair(entity_id_1=e1, entity_id_2=e2, event_date=ed)
|
||||
for (e1, e2), ed in cooccurrence_pairs.items()
|
||||
)
|
||||
|
||||
async def get_units_by_entity(self, entity_id: str, limit: int = 100) -> list[str]:
|
||||
"""
|
||||
Get all units that mention an entity.
|
||||
|
||||
Args:
|
||||
entity_id: Entity ID
|
||||
limit: Max results
|
||||
|
||||
Returns:
|
||||
List of unit IDs
|
||||
"""
|
||||
async with acquire_with_retry(self.pool) as conn:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT unit_id
|
||||
FROM {fq_table("unit_entities")}
|
||||
WHERE entity_id = $1
|
||||
ORDER BY unit_id
|
||||
LIMIT $2
|
||||
""",
|
||||
entity_id,
|
||||
limit,
|
||||
)
|
||||
return [row["unit_id"] for row in rows]
|
||||
|
||||
async def get_entity_by_text(
|
||||
self,
|
||||
bank_id: str,
|
||||
entity_text: str,
|
||||
) -> str | None:
|
||||
"""
|
||||
Find an entity by text (for query resolution).
|
||||
|
||||
Args:
|
||||
bank_id: bank ID
|
||||
entity_text: Entity text to search for
|
||||
|
||||
Returns:
|
||||
Entity ID if found, None otherwise
|
||||
"""
|
||||
async with acquire_with_retry(self.pool) as conn:
|
||||
row = await conn.fetchrow(
|
||||
f"""
|
||||
SELECT id FROM {fq_table("entities")}
|
||||
WHERE bank_id = $1
|
||||
AND canonical_name ILIKE $2
|
||||
ORDER BY mention_count DESC
|
||||
LIMIT 1
|
||||
""",
|
||||
bank_id,
|
||||
entity_text,
|
||||
)
|
||||
|
||||
return row["id"] if row else None
|
||||
|
||||
@@ -5,48 +5,54 @@ Three reconciliation passes run together on every worker invocation:
|
||||
1. **Relink top-up.** Drain ``graph_maintenance_queue`` (units whose
|
||||
outgoing temporal/semantic links lost a neighbour to a delete). For
|
||||
each, count current outgoing links per type; if below cap, run the
|
||||
same probes retain uses (:func:`fetch_temporal_neighbors`,
|
||||
:func:`compute_semantic_links_ann`) and insert the missing links.
|
||||
``bulk_insert_links`` has ``ON CONFLICT DO NOTHING`` on the uniqueness
|
||||
key, so we can re-probe freely and the DB de-dupes.
|
||||
same probes retain uses and insert the missing links.
|
||||
|
||||
2. **Orphan entity prune.** Delete ``entities`` rows in the bank that no
|
||||
longer have any ``unit_entities`` references. FK ON DELETE CASCADE on
|
||||
longer have any live memory references. FK ON DELETE CASCADE on
|
||||
``entity_cooccurrences`` then removes any cooccurrence row pointing
|
||||
at the pruned entities.
|
||||
|
||||
3. **Stale cooccurrence prune.** Defensive sweep for cooccurrence rows
|
||||
where both endpoints still exist but no current memory_unit references
|
||||
where both endpoints still exist but no current memory references
|
||||
both of them — the cooccurrence was real at the time it was recorded,
|
||||
but every unit that witnessed it has since been deleted.
|
||||
|
||||
All three passes run on every invocation. The queue is the only source
|
||||
of work for pass 1; passes 2 and 3 are bank-wide sweeps backed by indexes
|
||||
on ``entities(bank_id)`` and ``unit_entities(entity_id)``, so they're
|
||||
cheap when there's nothing to do.
|
||||
Each pass is work the *memories store* owns, because each is a query over
|
||||
`memory_links`, `unit_entities` and `entities` — the slice the store carves
|
||||
out. This module orchestrates them (drain the queue, wrap the sweep in a
|
||||
deadlock-retry) and asks the store to do the part that touches storage. A store
|
||||
whose links travel inside its memories has no `memory_links` to dangle and no
|
||||
join table to sweep, so its relink and cooccurrence passes are no-ops and the
|
||||
job simply prunes the orphan `entities` rows, which stay in Postgres regardless.
|
||||
|
||||
The worker dedupes on bank: a second job for the same bank is dropped
|
||||
while one is pending. Once processing starts, a new job becomes the
|
||||
*next* pending slot — so work enqueued during processing gets picked up
|
||||
by the follow-up run.
|
||||
|
||||
That follow-up run is *deferred*, not parallel: ``claim_tasks`` will not claim a
|
||||
graph_maintenance row for a bank that already has one in flight (#3230). Two
|
||||
concurrent runs would do no extra work anyway — each is this same bank-wide
|
||||
sweep — while convoying on each other's row locks and holding a worker slot
|
||||
each.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import time
|
||||
import uuid as uuid_module
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Any
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from ..models import RequestContext
|
||||
from .db.base import DatabaseConnection
|
||||
from .retain.link_utils import (
|
||||
MAX_TEMPORAL_LINKS_PER_UNIT,
|
||||
_bulk_insert_links,
|
||||
_normalize_datetime,
|
||||
compute_semantic_links_ann,
|
||||
)
|
||||
|
||||
# Re-exported for callers and tests that import the link caps from here; the caps
|
||||
# themselves live with the link builders the relink pass mirrors — the temporal one
|
||||
# with the retain-time builders, the semantic one with the store's relink pass — so
|
||||
# there is a single definition of each and the two cannot drift.
|
||||
from .memories.pg.graph import MAX_SEMANTIC_LINKS_PER_UNIT # noqa: F401
|
||||
from .retain.link_utils import MAX_TEMPORAL_LINKS_PER_UNIT # noqa: F401
|
||||
from .schema import fq_table
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -54,17 +60,19 @@ if TYPE_CHECKING:
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Mirrors the ``top_k`` default in ``compute_semantic_links_ann`` at retain
|
||||
# time. If you change one, change the other — otherwise victims would either
|
||||
# never reach the cap (probe returns less than the cap) or stay perpetually
|
||||
# under it (cap is higher than retain creates).
|
||||
MAX_SEMANTIC_LINKS_PER_UNIT = 50
|
||||
# Retry budget for the idempotent Pass 2/3 entity/cooccurrence sweep. Higher
|
||||
# than db_utils' default (3) because the sweep has no client waiting on it and
|
||||
# is safe to rerun, so we'd rather spend a longer jittered-backoff tail than
|
||||
# drop a maintenance pass and leak stale graph rows (see run_graph_maintenance_job).
|
||||
_SWEEP_MAX_RETRIES = 8
|
||||
|
||||
# Worker fetches this many rows per relink-loop iteration. Bounds
|
||||
# per-iteration probe/insert latency so a 10k-row backlog doesn't hold a
|
||||
# worker slot for minutes. Chosen so the typical iteration runs in well
|
||||
# under 1s.
|
||||
_DRAIN_BATCH_SIZE = 50
|
||||
|
||||
@dataclass
|
||||
class _SweepCounts:
|
||||
"""Prune counts returned by the Pass 2/3 sweep (avoids a bare tuple return)."""
|
||||
|
||||
orphan_entities_pruned: int
|
||||
stale_cooccurrences_pruned: int
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -88,67 +96,52 @@ class JobResult:
|
||||
async def enqueue_relink_victims(
|
||||
conn: DatabaseConnection,
|
||||
bank_id: str,
|
||||
deleted_unit_ids: list[str],
|
||||
ops: Any,
|
||||
affected_unit_ids: list[str],
|
||||
include_affected_units: bool = False,
|
||||
) -> int:
|
||||
"""Enqueue surviving units whose outgoing temporal/semantic links pointed at
|
||||
``deleted_unit_ids`` for later link top-up.
|
||||
``affected_unit_ids`` for later link top-up.
|
||||
|
||||
Must run inside the same transaction that deletes the units, *before* the
|
||||
cascade fires — once the rows are gone, the join that finds the victims
|
||||
returns nothing.
|
||||
Must run inside the same transaction that drops those links, *before* the
|
||||
delete (or cascade) fires — once the rows are gone, the join that finds the
|
||||
victims returns nothing.
|
||||
|
||||
``include_affected_units`` covers the case where the affected units are NOT
|
||||
being removed: an edit deletes every link incident to the edited unit but
|
||||
leaves it live, so the unit needs its own outgoing adjacency rebuilt too.
|
||||
Passing it for a unit that will be gone at commit is harmless but pointless
|
||||
— the drain skips queue rows with no live unit — so callers should only set
|
||||
it when the unit survives the transaction.
|
||||
|
||||
Delegated to the memories store: finding the victims is a `memory_links`
|
||||
query, and a store whose links are inline has none, so it returns 0 and the
|
||||
relink pass has nothing to do. The store resolves the dialect it needs from
|
||||
``conn``.
|
||||
|
||||
Args:
|
||||
conn: Database connection inside the active delete transaction.
|
||||
bank_id: Bank owning the deleted units.
|
||||
deleted_unit_ids: Memory_unit IDs about to be (or being) deleted.
|
||||
ops: ``DataAccessOps`` instance, supplies the dialect-specific
|
||||
bulk-insert path.
|
||||
conn: Database connection inside the active transaction.
|
||||
bank_id: Bank owning the affected units.
|
||||
affected_unit_ids: Memory_unit IDs whose incident temporal/semantic
|
||||
links are about to be (or are being) removed.
|
||||
include_affected_units: Also enqueue ``affected_unit_ids`` themselves,
|
||||
for callers that leave them live.
|
||||
|
||||
Returns:
|
||||
Number of distinct victim units enqueued (after dedup against rows
|
||||
already in the queue).
|
||||
Number of distinct victim units enqueued (0 for a store with no links).
|
||||
"""
|
||||
if not deleted_unit_ids:
|
||||
if not affected_unit_ids:
|
||||
return 0
|
||||
|
||||
deleted_uuids = [uuid_module.UUID(uid) if isinstance(uid, str) else uid for uid in deleted_unit_ids]
|
||||
deleted_str_set = {str(uid) for uid in deleted_uuids}
|
||||
from .memories import get_memories
|
||||
|
||||
# Find units (other than the ones being deleted) that have an outgoing
|
||||
# temporal/semantic link pointing at a doomed unit. Entity links are
|
||||
# intentionally excluded — they're scheduled for removal and would only
|
||||
# add noise to the recompute job.
|
||||
victim_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT DISTINCT from_unit_id
|
||||
FROM {fq_table("memory_links")}
|
||||
WHERE to_unit_id = ANY($1::uuid[])
|
||||
AND bank_id = $2
|
||||
AND link_type IN ('temporal', 'semantic')
|
||||
""",
|
||||
deleted_uuids,
|
||||
bank_id,
|
||||
return await get_memories().enqueue_relink_victims(
|
||||
conn=conn,
|
||||
fq_table=fq_table,
|
||||
bank_id=bank_id,
|
||||
affected_unit_ids=affected_unit_ids,
|
||||
include_affected_units=include_affected_units,
|
||||
)
|
||||
|
||||
victim_ids = [row["from_unit_id"] for row in victim_rows if str(row["from_unit_id"]) not in deleted_str_set]
|
||||
|
||||
if not victim_ids:
|
||||
return 0
|
||||
|
||||
await ops.enqueue_graph_maintenance(
|
||||
conn,
|
||||
fq_table("graph_maintenance_queue"),
|
||||
bank_id,
|
||||
victim_ids,
|
||||
)
|
||||
|
||||
logger.debug(
|
||||
f"[GRAPH_MAINT] Enqueued {len(victim_ids)} relink victims in "
|
||||
f"bank={bank_id} (deleted {len(deleted_unit_ids)} units)"
|
||||
)
|
||||
return len(victim_ids)
|
||||
|
||||
|
||||
async def run_graph_maintenance_job(
|
||||
memory_engine: "MemoryEngine",
|
||||
@@ -163,195 +156,69 @@ async def run_graph_maintenance_job(
|
||||
Per-pass counters from :class:`JobResult`.
|
||||
"""
|
||||
del request_context # accepted for symmetry with other run_*_job helpers
|
||||
from ..config import get_config
|
||||
from .memories import get_memories
|
||||
|
||||
backend = await memory_engine._get_backend()
|
||||
ops = backend.ops
|
||||
store = get_memories()
|
||||
config = get_config()
|
||||
|
||||
result = JobResult()
|
||||
job_start = time.time()
|
||||
|
||||
# --- Pass 1: relink ---
|
||||
# Per-iteration loop: claim → top up → commit. We rely on submit-time
|
||||
# dedup to keep at most one job per bank running, so no need for
|
||||
# SKIP LOCKED.
|
||||
iterations = 0
|
||||
while True:
|
||||
from .memory_engine import acquire_with_retry
|
||||
|
||||
async with acquire_with_retry(backend) as conn:
|
||||
async with conn.transaction():
|
||||
unit_ids = await ops.claim_graph_maintenance_batch(
|
||||
conn,
|
||||
fq_table("graph_maintenance_queue"),
|
||||
bank_id,
|
||||
_DRAIN_BATCH_SIZE,
|
||||
)
|
||||
if not unit_ids:
|
||||
break
|
||||
|
||||
result.relink_links_added += await _relink_batch(conn, bank_id, unit_ids, ops, backend)
|
||||
|
||||
result.relink_units_processed += len(unit_ids)
|
||||
iterations += 1
|
||||
|
||||
if iterations > 10000:
|
||||
# Defensive guard against runaway loops — at 50 units/iter that's
|
||||
# 500k targets, far beyond any realistic single-bank backlog.
|
||||
logger.error(
|
||||
f"[GRAPH_MAINT] bank={bank_id} hit iteration cap ({iterations}); aborting relink ({result.as_dict()})"
|
||||
)
|
||||
break
|
||||
# The store owns the whole drain loop: it is a claim → top-up → commit over
|
||||
# its own link table, so how it batches and re-probes is its business — including
|
||||
# the #3034 serialisation (the claim takes queue rows FOR UPDATE in (bank_id,
|
||||
# unit_id) order against a concurrent re-enqueue), which lives in the store's
|
||||
# claim (`ops.claim_graph_maintenance_batch`). A store with no links returns an
|
||||
# empty dict and this is a no-op.
|
||||
relink = await store.relink_pass(backend=backend, fq_table=fq_table, bank_id=bank_id, config=config)
|
||||
result.relink_units_processed = relink.get("relink_units_processed", 0)
|
||||
result.relink_links_added = relink.get("relink_links_added", 0)
|
||||
|
||||
# --- Pass 2 & 3: entity / cooccurrence sweeps ---
|
||||
# Bank-wide single-statement deletes. Cheap when there's nothing to do.
|
||||
#
|
||||
# Unlike Pass 1's queue claim, these DELETEs aren't protected by any
|
||||
# consistent lock-ordering guarantee: the stale-cooccurrence prune scans
|
||||
# entity_cooccurrences via a join/NOT EXISTS plan, while retain's concurrent
|
||||
# cooccurrence upserts (entity_resolver._flush_pending) lock the same rows in
|
||||
# sorted (entity_id_1, entity_id_2) order. When a sweep and a concurrent
|
||||
# upsert touch overlapping rows in opposite orders, Postgres detects a
|
||||
# genuine circular wait and aborts one side with DeadlockDetectedError. Both
|
||||
# prunes are idempotent bank-wide sweeps — rerunning only deletes what's
|
||||
# still stale — so retrying the whole transaction on deadlock is safe.
|
||||
#
|
||||
# The prunes themselves are the store's: the orphan-`entities` sweep applies
|
||||
# to every store (that registry stays in Postgres), while the cooccurrence
|
||||
# sweep is a no-op for a store that never wrote `unit_entities`.
|
||||
from .db_utils import retry_with_backoff
|
||||
from .memory_engine import acquire_with_retry
|
||||
|
||||
async with acquire_with_retry(backend) as conn:
|
||||
async with conn.transaction():
|
||||
result.orphan_entities_pruned = await ops.prune_orphan_entities(
|
||||
conn,
|
||||
fq_table("entities"),
|
||||
fq_table("unit_entities"),
|
||||
bank_id,
|
||||
)
|
||||
# The orphan prune above cascades cooccurrences via FK. The
|
||||
# explicit cooccurrence pass below catches the *stale-count*
|
||||
# case: both entities still exist but no current unit witnesses
|
||||
# them together.
|
||||
result.stale_cooccurrences_pruned = await ops.prune_stale_cooccurrences(
|
||||
conn,
|
||||
fq_table("entity_cooccurrences"),
|
||||
fq_table("unit_entities"),
|
||||
fq_table("entities"),
|
||||
bank_id,
|
||||
)
|
||||
async def _run_sweep() -> _SweepCounts:
|
||||
async with acquire_with_retry(backend) as conn:
|
||||
async with conn.transaction():
|
||||
orphan_pruned = await store.prune_orphan_entities(conn=conn, fq_table=fq_table, bank_id=bank_id)
|
||||
# The orphan prune above cascades cooccurrences via FK. The
|
||||
# explicit cooccurrence pass below catches the *stale-count*
|
||||
# case: both entities still exist but no current unit witnesses
|
||||
# them together.
|
||||
stale_pruned = await store.prune_stale_cooccurrences(conn=conn, fq_table=fq_table, bank_id=bank_id)
|
||||
return _SweepCounts(orphan_entities_pruned=orphan_pruned, stale_cooccurrences_pruned=stale_pruned)
|
||||
|
||||
# A larger retry budget than the default (3): this is idempotent background
|
||||
# maintenance with no client waiting on it, so a longer retry tail costs
|
||||
# nothing, whereas a dropped sweep silently leaks orphan entities / stale
|
||||
# cooccurrences until the next run. With jittered backoff a single sweep
|
||||
# contending against continuous retain upserts effectively never exhausts
|
||||
# this budget (each retry independently clears with high probability).
|
||||
sweep = await retry_with_backoff(_run_sweep, max_retries=_SWEEP_MAX_RETRIES)
|
||||
result.orphan_entities_pruned = sweep.orphan_entities_pruned
|
||||
result.stale_cooccurrences_pruned = sweep.stale_cooccurrences_pruned
|
||||
|
||||
elapsed = time.time() - job_start
|
||||
logger.info(
|
||||
f"[GRAPH_MAINT] bank={bank_id} done: {result.as_dict()}, elapsed={elapsed:.2f}s, operation_id={operation_id}"
|
||||
)
|
||||
return result.as_dict()
|
||||
|
||||
|
||||
async def _relink_batch(
|
||||
conn: DatabaseConnection,
|
||||
bank_id: str,
|
||||
victim_ids: list[str],
|
||||
ops: Any,
|
||||
backend: Any,
|
||||
) -> int:
|
||||
"""Top up temporal/semantic links for a batch of victim units. Returns rows inserted."""
|
||||
# Load each victim's metadata. Victims whose units were deleted between
|
||||
# enqueue and now silently drop out — exactly the no-op behaviour we want
|
||||
# for stale queue rows.
|
||||
victim_uuids = [uuid_module.UUID(vid) for vid in victim_ids]
|
||||
victim_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT id::text AS id, event_date, fact_type, embedding::text AS embedding
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE id = ANY($1::uuid[])
|
||||
AND bank_id = $2
|
||||
AND fact_type IN ('experience', 'world')
|
||||
""",
|
||||
victim_uuids,
|
||||
bank_id,
|
||||
)
|
||||
|
||||
if not victim_rows:
|
||||
return 0
|
||||
|
||||
alive_uuids = [uuid_module.UUID(row["id"]) for row in victim_rows]
|
||||
|
||||
# Count current outgoing temporal/semantic links per victim so we only
|
||||
# probe for the ones genuinely below cap. Saves the bulk of the work when
|
||||
# most victims still have plenty of links.
|
||||
count_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT from_unit_id, link_type, COUNT(*) AS cnt
|
||||
FROM {fq_table("memory_links")}
|
||||
WHERE from_unit_id = ANY($1::uuid[])
|
||||
AND bank_id = $2
|
||||
AND link_type IN ('temporal', 'semantic')
|
||||
GROUP BY from_unit_id, link_type
|
||||
""",
|
||||
alive_uuids,
|
||||
bank_id,
|
||||
)
|
||||
counts: dict[tuple[str, str], int] = {}
|
||||
for row in count_rows:
|
||||
counts[(str(row["from_unit_id"]), row["link_type"])] = int(row["cnt"])
|
||||
|
||||
# --- Temporal top-up ---
|
||||
temporal_needs = [r for r in victim_rows if counts.get((r["id"], "temporal"), 0) < MAX_TEMPORAL_LINKS_PER_UNIT]
|
||||
new_links: list[tuple] = []
|
||||
|
||||
if temporal_needs:
|
||||
lateral_unit_ids = [uuid_module.UUID(r["id"]) for r in temporal_needs if r["event_date"] is not None]
|
||||
lateral_event_dates = [
|
||||
_normalize_datetime(r["event_date"]) for r in temporal_needs if r["event_date"] is not None
|
||||
]
|
||||
lateral_fact_types = [r["fact_type"] for r in temporal_needs if r["event_date"] is not None]
|
||||
|
||||
if lateral_unit_ids:
|
||||
rows = await ops.fetch_temporal_neighbors(
|
||||
conn,
|
||||
fq_table("memory_units"),
|
||||
bank_id,
|
||||
lateral_unit_ids,
|
||||
lateral_event_dates,
|
||||
lateral_fact_types,
|
||||
MAX_TEMPORAL_LINKS_PER_UNIT,
|
||||
)
|
||||
for row in rows:
|
||||
time_diff_h = float(row["time_diff_hours"])
|
||||
# Mirror the 24h window enforced at retain time. The bidirectional
|
||||
# index scan returns the K closest neighbours regardless of
|
||||
# window, so we filter here.
|
||||
if time_diff_h > 24:
|
||||
continue
|
||||
weight = max(0.3, 1.0 - (time_diff_h / 24))
|
||||
new_links.append((row["from_id"], str(row["id"]), "temporal", weight, None))
|
||||
|
||||
# --- Semantic top-up ---
|
||||
# ANN must run on its own connection: it opens a nested transaction with
|
||||
# SET LOCAL hnsw.ef_search + CREATE TEMP TABLE ON COMMIT DROP, and nesting
|
||||
# that inside our current write transaction would commit our writes early.
|
||||
semantic_needs = [
|
||||
r
|
||||
for r in victim_rows
|
||||
if counts.get((r["id"], "semantic"), 0) < MAX_SEMANTIC_LINKS_PER_UNIT and r["embedding"] is not None
|
||||
]
|
||||
if semantic_needs:
|
||||
from .memory_engine import acquire_with_retry
|
||||
|
||||
seed_ids = [r["id"] for r in semantic_needs]
|
||||
seed_embs = [r["embedding"] for r in semantic_needs]
|
||||
seed_ftypes = [r["fact_type"] for r in semantic_needs]
|
||||
async with acquire_with_retry(backend) as ann_conn:
|
||||
try:
|
||||
ann_links = await compute_semantic_links_ann(
|
||||
ann_conn,
|
||||
bank_id,
|
||||
seed_ids,
|
||||
seed_embs,
|
||||
fact_types=seed_ftypes,
|
||||
)
|
||||
# Strip self-links (rare but possible because the ANN probe
|
||||
# has no exclude list — see the comment in compute_semantic_links_ann).
|
||||
ann_links = [lnk for lnk in ann_links if lnk[0] != lnk[1]]
|
||||
new_links.extend(ann_links)
|
||||
except Exception as e:
|
||||
# ANN uses PG-specific HNSW syntax; on dialects/configs where
|
||||
# it isn't available we still want the temporal top-up to land.
|
||||
logger.warning(f"[GRAPH_MAINT] Semantic top-up failed for bank={bank_id}: {type(e).__name__}: {e}")
|
||||
|
||||
if not new_links:
|
||||
return 0
|
||||
|
||||
await _bulk_insert_links(
|
||||
conn,
|
||||
new_links,
|
||||
bank_id=bank_id,
|
||||
skip_exists_check=False,
|
||||
ops=ops,
|
||||
)
|
||||
return len(new_links)
|
||||
|
||||
@@ -6,6 +6,7 @@ authentication when a TenantExtension is configured.
|
||||
"""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
@@ -13,9 +14,26 @@ if TYPE_CHECKING:
|
||||
from hindsight_api.engine.memory_engine import BankLlmHealthInfo, Budget
|
||||
from hindsight_api.engine.response_models import RecallResult, ReflectResult
|
||||
from hindsight_api.engine.search.tags import TagsMatch
|
||||
from hindsight_api.extensions import BankWriteOperation
|
||||
from hindsight_api.models import RequestContext
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class BankConfigState:
|
||||
"""Resolved bank configuration and its bank-level overrides."""
|
||||
|
||||
config: dict[str, Any]
|
||||
overrides: dict[str, Any]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class BankTemplateImportWrite:
|
||||
"""One bank-write decision reserved for a specific imported resource."""
|
||||
|
||||
operation: "BankWriteOperation"
|
||||
target: str | None = None
|
||||
|
||||
|
||||
class MemoryEngineInterface(ABC):
|
||||
"""
|
||||
Abstract interface for the Memory Engine.
|
||||
@@ -180,6 +198,37 @@ class MemoryEngineInterface(ABC):
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def get_bank_config(
|
||||
self,
|
||||
bank_id: str,
|
||||
*,
|
||||
request_context: "RequestContext",
|
||||
) -> BankConfigState:
|
||||
"""Return resolved configuration after authenticating and authorizing the read."""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def update_bank_config(
|
||||
self,
|
||||
bank_id: str,
|
||||
updates: dict[str, Any],
|
||||
*,
|
||||
request_context: "RequestContext",
|
||||
) -> BankConfigState:
|
||||
"""Create a bank if needed and persist validated configuration overrides."""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def reset_bank_config(
|
||||
self,
|
||||
bank_id: str,
|
||||
*,
|
||||
request_context: "RequestContext",
|
||||
) -> BankConfigState:
|
||||
"""Remove all bank configuration overrides after authorization."""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def update_bank_disposition(
|
||||
self,
|
||||
@@ -275,6 +324,8 @@ class MemoryEngineInterface(ABC):
|
||||
*,
|
||||
fact_type: str | None = None,
|
||||
search_query: str | None = None,
|
||||
entity_id: str | None = None,
|
||||
created_before: datetime | None = None,
|
||||
limit: int = 100,
|
||||
offset: int = 0,
|
||||
request_context: "RequestContext",
|
||||
@@ -286,6 +337,8 @@ class MemoryEngineInterface(ABC):
|
||||
bank_id: The memory bank ID.
|
||||
fact_type: Filter by fact type.
|
||||
search_query: Full-text search query.
|
||||
entity_id: Filter to memory units linked to this entity ID.
|
||||
created_before: Keep units with ``created_at`` before this instant.
|
||||
limit: Maximum results.
|
||||
offset: Pagination offset.
|
||||
request_context: Request context for authentication.
|
||||
@@ -478,11 +531,15 @@ class MemoryEngineInterface(ABC):
|
||||
Get consolidation freshness for a bank.
|
||||
|
||||
Cheap alternative to get_bank_stats when callers only need
|
||||
last_consolidated_at / pending_consolidation / failed_consolidation.
|
||||
last_consolidated_at / last_memory_write_at / pending_consolidation /
|
||||
failed_consolidation.
|
||||
|
||||
Returns:
|
||||
Dict with last_consolidated_at (ISO-8601 string or None),
|
||||
pending_consolidation (int), and failed_consolidation (int).
|
||||
Dict with last_consolidated_at and last_memory_write_at (ISO-8601
|
||||
strings or None), pending_consolidation (int), and
|
||||
failed_consolidation (int). last_memory_write_at is the newest write
|
||||
across the bank's memories — a mental model refreshed at or after it
|
||||
cannot be stale, whatever its scope.
|
||||
"""
|
||||
...
|
||||
|
||||
@@ -565,6 +622,30 @@ class MemoryEngineInterface(ABC):
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def delete_operation(
|
||||
self,
|
||||
bank_id: str,
|
||||
operation_id: str,
|
||||
*,
|
||||
request_context: "RequestContext",
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Delete a terminal async operation record.
|
||||
|
||||
Args:
|
||||
bank_id: The memory bank ID.
|
||||
operation_id: The operation ID to delete.
|
||||
request_context: Request context for authentication.
|
||||
|
||||
Returns:
|
||||
Dict with success status and message.
|
||||
|
||||
Raises:
|
||||
ValueError: If operation not found.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def update_bank(
|
||||
self,
|
||||
@@ -572,6 +653,8 @@ class MemoryEngineInterface(ABC):
|
||||
*,
|
||||
name: str | None = None,
|
||||
mission: str | None = None,
|
||||
config_updates: dict[str, Any] | None = None,
|
||||
create_if_missing: bool = True,
|
||||
request_context: "RequestContext",
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
@@ -581,6 +664,9 @@ class MemoryEngineInterface(ABC):
|
||||
bank_id: The memory bank ID.
|
||||
name: New bank name (optional).
|
||||
mission: New mission text (optional, replaces existing).
|
||||
config_updates: Bank configuration overrides to apply with the profile update.
|
||||
create_if_missing: Create a missing bank when True; otherwise raise
|
||||
a 404 operation error.
|
||||
request_context: Request context for authentication.
|
||||
|
||||
Returns:
|
||||
|
||||
@@ -6,12 +6,54 @@ enabling support for multiple LLM backends (OpenAI, Anthropic, Gemini, Codex, et
|
||||
"""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from contextlib import AbstractAsyncContextManager
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from typing import Any
|
||||
from enum import StrEnum
|
||||
from typing import Any, Callable, Self
|
||||
|
||||
from .response_models import LLMToolCallResult
|
||||
|
||||
|
||||
class LLMToolChoiceMode(StrEnum):
|
||||
"""Canonical tool-selection modes shared by every LLM provider."""
|
||||
|
||||
AUTO = "auto"
|
||||
NONE = "none"
|
||||
REQUIRED = "required"
|
||||
NAMED = "named"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class LLMToolChoice:
|
||||
"""Typed internal tool selection serialized only at provider boundaries."""
|
||||
|
||||
mode: LLMToolChoiceMode
|
||||
function_name: str | None = None
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if self.mode is LLMToolChoiceMode.NAMED:
|
||||
if self.function_name is None or not self.function_name or self.function_name != self.function_name.strip():
|
||||
raise ValueError("Named tool choice requires a non-empty canonical function name")
|
||||
elif self.function_name is not None:
|
||||
raise ValueError(f"Tool choice mode {self.mode.value!r} cannot include a function name")
|
||||
|
||||
@classmethod
|
||||
def named(cls, function_name: str) -> Self:
|
||||
return cls(mode=LLMToolChoiceMode.NAMED, function_name=function_name)
|
||||
|
||||
@property
|
||||
def selected_function_name(self) -> str:
|
||||
if self.function_name is None:
|
||||
raise ValueError("Tool choice does not select a named function")
|
||||
return self.function_name
|
||||
|
||||
|
||||
LLM_TOOL_CHOICE_AUTO = LLMToolChoice(mode=LLMToolChoiceMode.AUTO)
|
||||
LLM_TOOL_CHOICE_NONE = LLMToolChoice(mode=LLMToolChoiceMode.NONE)
|
||||
LLM_TOOL_CHOICE_REQUIRED = LLMToolChoice(mode=LLMToolChoiceMode.REQUIRED)
|
||||
|
||||
|
||||
class LLMInterface(ABC):
|
||||
"""
|
||||
Abstract interface for LLM providers.
|
||||
@@ -71,6 +113,7 @@ class LLMInterface(ABC):
|
||||
strict_schema: bool = False,
|
||||
return_usage: bool = False,
|
||||
cached_prefix: str | None = None,
|
||||
attempt_context: Callable[[], AbstractAsyncContextManager[None]] | None = None,
|
||||
) -> Any:
|
||||
"""
|
||||
Make an LLM API call with retry logic.
|
||||
@@ -92,6 +135,11 @@ class LLMInterface(ABC):
|
||||
cached_prefix: Opaque handle from ``get_or_create_cached_prefix`` for the
|
||||
cacheable system prefix, or None. Providers without explicit prompt
|
||||
caching ignore it (and the wrapper only forwards it when set).
|
||||
attempt_context: Factory for an async context manager holding the shared
|
||||
concurrency permits. Passed only when the provider declares
|
||||
``supports_attempt_scoped_concurrency()``; the provider must enter it
|
||||
around each individual upstream request so retry backoff never
|
||||
occupies a permit.
|
||||
|
||||
Returns:
|
||||
If return_usage=False: Parsed response if response_format is provided, otherwise text content.
|
||||
@@ -114,8 +162,10 @@ class LLMInterface(ABC):
|
||||
max_retries: int = 5,
|
||||
initial_backoff: float = 1.0,
|
||||
max_backoff: float = 30.0,
|
||||
tool_choice: str | dict[str, Any] = "auto",
|
||||
tool_choice: LLMToolChoice = LLM_TOOL_CHOICE_AUTO,
|
||||
cached_prefix: str | None = None,
|
||||
cached_prefix_message_count: int = 0,
|
||||
attempt_context: Callable[[], AbstractAsyncContextManager[None]] | None = None,
|
||||
) -> LLMToolCallResult:
|
||||
"""
|
||||
Make an LLM API call with tool/function calling support.
|
||||
@@ -129,7 +179,9 @@ class LLMInterface(ABC):
|
||||
max_retries: Maximum retry attempts.
|
||||
initial_backoff: Initial backoff time in seconds.
|
||||
max_backoff: Maximum backoff time in seconds.
|
||||
tool_choice: How to choose tools - "auto", "none", "required", or specific function.
|
||||
tool_choice: Canonical tool-selection policy.
|
||||
attempt_context: Factory for an async context manager holding the shared
|
||||
concurrency permits — see ``call``.
|
||||
|
||||
Returns:
|
||||
LLMToolCallResult with content and/or tool_calls.
|
||||
@@ -145,6 +197,10 @@ class LLMInterface(ABC):
|
||||
"""
|
||||
return False
|
||||
|
||||
def supports_attempt_scoped_concurrency(self) -> bool:
|
||||
"""Whether retries can acquire concurrency permits per upstream attempt."""
|
||||
return False
|
||||
|
||||
# ── Prompt prefix caching (optional, per-provider) ─────────────────────────
|
||||
|
||||
def supports_prompt_caching(self) -> bool:
|
||||
@@ -185,6 +241,45 @@ class LLMInterface(ABC):
|
||||
"""
|
||||
return None
|
||||
|
||||
# ── Step-by-step incremental prompt caching (optional) ─────────────────────
|
||||
#
|
||||
# For agentic loops (reflect) the dominant cost is the conversation prefix
|
||||
# re-sent every turn, not the static system prefix. Providers that can cache
|
||||
# a *growing* prefix implement these: the caller rolls one cache per step
|
||||
# (each covering the previous step's full input), passes its handle plus the
|
||||
# message count it covers to ``call_with_tools`` so only the new turns are
|
||||
# sent fresh, and tears the caches down when the loop ends. Default no-ops so
|
||||
# non-supporting providers transparently run uncached.
|
||||
|
||||
def supports_incremental_prompt_cache(self) -> bool:
|
||||
"""Whether this provider can cache a growing multi-turn conversation prefix."""
|
||||
return False
|
||||
|
||||
async def create_incremental_cache(
|
||||
self,
|
||||
*,
|
||||
session_id: str,
|
||||
messages: list[dict[str, Any]],
|
||||
tools: list[dict[str, Any]] | None = None,
|
||||
) -> str | None:
|
||||
"""Cache ``system + tools + messages`` and return an opaque handle, or None.
|
||||
|
||||
The handle is passed back to ``call_with_tools(cached_prefix=...,
|
||||
cached_prefix_message_count=len(messages))``. Caches are grouped under
|
||||
``session_id`` for teardown via ``delete_cache_session``. Returns None
|
||||
when caching is unavailable or the prefix is too small — caller falls
|
||||
back to an uncached call.
|
||||
"""
|
||||
return None
|
||||
|
||||
async def delete_cached_prefix(self, name: str) -> None:
|
||||
"""Best-effort delete of a single cache handle (a superseded step)."""
|
||||
return None
|
||||
|
||||
async def delete_cache_session(self, session_id: str) -> None:
|
||||
"""Best-effort teardown of every cache created under ``session_id``."""
|
||||
return None
|
||||
|
||||
async def submit_batch(
|
||||
self,
|
||||
requests: list[dict[str, Any]],
|
||||
|
||||
@@ -36,6 +36,21 @@ from .db_utils import acquire_with_retry
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _llm_requests_persistable() -> bool:
|
||||
"""Whether the ``llm_requests`` table exists on the active backend.
|
||||
|
||||
``llm_requests`` is PostgreSQL-only: its migration is ``run_for_dialect(pg=...)``
|
||||
with the Oracle slot intentionally absent, and MaintenanceLoop skips its
|
||||
retention sweep on Oracle for the same reason. On Oracle the table does not
|
||||
exist, so best-effort trace writes must be skipped rather than attempted —
|
||||
otherwise every LLM call fires an INSERT that fails with ORA-00903 and spams
|
||||
the error log. Mirrors the ``_is_oracle()`` gate in MaintenanceLoop.start.
|
||||
"""
|
||||
from .schema import _is_oracle
|
||||
|
||||
return not _is_oracle()
|
||||
|
||||
|
||||
# ── bank/operation attribution (carried across the async call chain) ──────────
|
||||
|
||||
|
||||
@@ -376,10 +391,32 @@ class LLMTraceRecorder:
|
||||
# INSERTs it patches — but it must not block on unrelated operations).
|
||||
self._pending: dict[str | None, set[asyncio.Task]] = {}
|
||||
|
||||
def _writable(self) -> Any | None:
|
||||
"""Return the pool to write through, or None if writing isn't possible.
|
||||
|
||||
Covers the two lifecycle windows in which best-effort trace writes must
|
||||
be skipped rather than attempted: before the backend pool is created
|
||||
(``initialize()`` verifies the LLM before the DB is up) and during/after
|
||||
shutdown. Writes already in flight need no handling — the pools close
|
||||
gracefully, waiting for their connections to be released.
|
||||
"""
|
||||
pool = self._pool_getter()
|
||||
if pool is None:
|
||||
return None
|
||||
# Backends declare readiness explicitly; a raw pool (some callers pass
|
||||
# one directly) has no lifecycle flag and is assumed usable.
|
||||
from .db.base import DatabaseBackend
|
||||
|
||||
if isinstance(pool, DatabaseBackend) and not pool.is_ready:
|
||||
return None
|
||||
return pool
|
||||
|
||||
def is_enabled(self, scope: str) -> bool:
|
||||
"""Whether tracing is active for the given call scope."""
|
||||
if not self._enabled:
|
||||
return False
|
||||
if not _llm_requests_persistable():
|
||||
return False
|
||||
if self._allowed_scopes is not None:
|
||||
return scope in self._allowed_scopes
|
||||
return True
|
||||
@@ -473,7 +510,7 @@ class LLMTraceRecorder:
|
||||
|
||||
async def _safe_write(self, record: LLMRequestRecord) -> None:
|
||||
"""Write a trace row. Errors are logged, never raised."""
|
||||
pool = self._pool_getter()
|
||||
pool = self._writable()
|
||||
if pool is None:
|
||||
logger.debug("LLM trace skipped: pool not available")
|
||||
return
|
||||
@@ -546,7 +583,7 @@ class LLMTraceRecorder:
|
||||
ids are snapshotted synchronously here because the caller may reset the
|
||||
context immediately after.
|
||||
"""
|
||||
if not self._enabled or trace_ctx is None or not trace_ctx.trace_id:
|
||||
if not self._enabled or not _llm_requests_persistable() or trace_ctx is None or not trace_ctx.trace_id:
|
||||
return
|
||||
created_ids = list(dict.fromkeys([*(created or []), *trace_ctx.created_memory_ids]))
|
||||
source_ids = list(dict.fromkeys([*(source or []), *trace_ctx.source_memory_ids]))
|
||||
@@ -568,8 +605,9 @@ class LLMTraceRecorder:
|
||||
# so the UPDATE patches rows that already exist rather than racing ahead
|
||||
# of them (without blocking on unrelated operations' pending writes).
|
||||
await self._flush_pending(trace_id)
|
||||
pool = self._pool_getter()
|
||||
pool = self._writable()
|
||||
if pool is None:
|
||||
logger.debug("LLM trace memory_id attach skipped: pool not available")
|
||||
return
|
||||
try:
|
||||
schema = self._schema_getter()
|
||||
|
||||
@@ -9,9 +9,11 @@ import os
|
||||
import re
|
||||
import time
|
||||
import uuid
|
||||
from contextlib import AsyncExitStack
|
||||
from contextlib import AsyncExitStack, asynccontextmanager
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from json_repair import repair_json
|
||||
|
||||
# Vertex AI imports (conditional - for LLMProvider to pass credentials to GeminiLLM)
|
||||
try:
|
||||
from google.oauth2 import service_account
|
||||
@@ -27,13 +29,18 @@ from ..config import (
|
||||
ENV_REFLECT_LLM_MAX_CONCURRENT,
|
||||
ENV_RETAIN_LLM_MAX_CONCURRENT,
|
||||
)
|
||||
from .llm_interface import (
|
||||
LLM_TOOL_CHOICE_AUTO,
|
||||
LLMToolChoice,
|
||||
LLMToolChoiceMode,
|
||||
)
|
||||
from .llm_interface import (
|
||||
OutputTooLongError as OutputTooLongError,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .response_models import LLMToolCallResult
|
||||
|
||||
# Seed applied to every Groq request for deterministic behavior.
|
||||
DEFAULT_LLM_SEED = 4242
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Disable httpx logging
|
||||
@@ -107,13 +114,34 @@ def _semaphores_for_scope(scope: str) -> list[asyncio.Semaphore]:
|
||||
return [per_op, _global_llm_semaphore]
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def _attempt_permits(scope: str):
|
||||
"""Hold configured LLM concurrency permits for one upstream attempt."""
|
||||
from ..worker.stage import get_stage, set_stage
|
||||
|
||||
async with AsyncExitStack() as stack:
|
||||
for sem in _semaphores_for_scope(scope):
|
||||
await stack.enter_async_context(sem)
|
||||
try:
|
||||
yield
|
||||
except BaseException:
|
||||
# A failed attempt exits here with its permits released while the
|
||||
# provider classifies the error and sleeps out its backoff. Suffix
|
||||
# the stage so `attempt=N` always means "permits held, request in
|
||||
# flight" (#3002); the next attempt re-stamps after re-acquiring.
|
||||
stage = get_stage()
|
||||
if stage is not None and not stage.endswith(".backoff"):
|
||||
set_stage(f"{stage}.backoff")
|
||||
raise
|
||||
|
||||
|
||||
def _request_params(
|
||||
*,
|
||||
max_completion_tokens: int | None = None,
|
||||
temperature: float | None = None,
|
||||
scope: str | None = None,
|
||||
response_format: Any | None = None,
|
||||
tool_choice: str | dict[str, Any] | None = None,
|
||||
tool_choice: LLMToolChoice | None = None,
|
||||
) -> dict[str, Any] | None:
|
||||
"""Build the requested-params bag for tracing — only values the caller set.
|
||||
|
||||
@@ -128,8 +156,8 @@ def _request_params(
|
||||
params["temperature"] = temperature
|
||||
if response_format is not None:
|
||||
params["response_schema"] = getattr(response_format, "__name__", None) or "structured"
|
||||
if tool_choice is not None and tool_choice != "auto":
|
||||
params["tool_choice"] = tool_choice if isinstance(tool_choice, str) else "named"
|
||||
if tool_choice is not None and tool_choice.mode is not LLMToolChoiceMode.AUTO:
|
||||
params["tool_choice"] = tool_choice.function_name or tool_choice.mode.value
|
||||
return params or None
|
||||
|
||||
|
||||
@@ -164,16 +192,11 @@ def sanitize_text(text: str | None) -> str | None:
|
||||
sanitize_llm_output = sanitize_text
|
||||
|
||||
|
||||
class OutputTooLongError(Exception):
|
||||
"""
|
||||
Bridge exception raised when LLM output exceeds token limits.
|
||||
|
||||
This wraps provider-specific errors (e.g., OpenAI's LengthFinishReasonError)
|
||||
to allow callers to handle output length issues without depending on
|
||||
provider-specific implementations.
|
||||
"""
|
||||
|
||||
pass
|
||||
# ``OutputTooLongError`` is re-exported from ``llm_interface`` (the canonical
|
||||
# definition the providers raise) so that ``fact_extraction`` and ``multi_llm``,
|
||||
# which import it from here, catch/inspect the very same class. Do NOT redefine
|
||||
# it locally: a shadow class silently breaks ``except OutputTooLongError`` on the
|
||||
# real provider path (see issue #3172).
|
||||
|
||||
|
||||
def parse_llm_json(raw: str) -> Any:
|
||||
@@ -184,6 +207,14 @@ def parse_llm_json(raw: str) -> Any:
|
||||
1. Markdown code fences (```json ... ```) — strip them before parsing.
|
||||
2. Embedded control characters (\\x00-\\x1f, \\x7f) — replace with space
|
||||
and retry if the initial parse fails.
|
||||
3. Structural malformation (trailing commas, unterminated strings, single
|
||||
quotes, invalid ``\\escape`` sequences) — repaired as a last resort via
|
||||
``json_repair`` (#2547/#2544).
|
||||
|
||||
The repair pass is purely *structural*: it fixes JSON that ``json.loads``
|
||||
cannot parse at all. It deliberately does NOT touch content semantics —
|
||||
degenerate-but-valid JSON (repetition loops or leaked scaffolding inside
|
||||
string values) parses fine here and is out of scope for this helper.
|
||||
|
||||
Args:
|
||||
raw: Raw text returned by the LLM.
|
||||
@@ -192,7 +223,8 @@ def parse_llm_json(raw: str) -> Any:
|
||||
Parsed Python object (dict, list, etc.).
|
||||
|
||||
Raises:
|
||||
json.JSONDecodeError: If the text cannot be parsed even after cleanup.
|
||||
json.JSONDecodeError: If the text cannot be parsed even after cleanup
|
||||
and structural repair (e.g. repair yields an empty result).
|
||||
"""
|
||||
text = raw.strip()
|
||||
|
||||
@@ -209,7 +241,19 @@ def parse_llm_json(raw: str) -> Any:
|
||||
# Some models (e.g. Gemini) embed raw control characters inside JSON
|
||||
# string values. Replacing them with a space usually produces valid JSON.
|
||||
cleaned = re.sub(r"[\x00-\x1f\x7f]", " ", text)
|
||||
|
||||
try:
|
||||
return json.loads(cleaned)
|
||||
except json.JSONDecodeError:
|
||||
# Last resort: structural repair of malformed JSON. ``repair_json`` never
|
||||
# raises — unrecoverable input yields an empty result ("" / {} / []). Keep
|
||||
# failing loudly in that case rather than let an empty object masquerade
|
||||
# as a successful parse: callers (retry ladders, the #1833 fail-loud path)
|
||||
# rely on JSONDecodeError to retry or surface the failure.
|
||||
repaired = repair_json(cleaned, return_objects=True)
|
||||
if not repaired:
|
||||
raise
|
||||
return repaired
|
||||
|
||||
|
||||
_PROVIDERS_WITHOUT_API_KEY = frozenset(
|
||||
@@ -235,6 +279,17 @@ def requires_api_key(provider: str) -> bool:
|
||||
return provider.lower() not in _PROVIDERS_WITHOUT_API_KEY
|
||||
|
||||
|
||||
def _validate_ollama_num_ctx(value: Any) -> int | None:
|
||||
"""Validate a native Ollama context-window override."""
|
||||
if value is None:
|
||||
return None
|
||||
if isinstance(value, bool) or not isinstance(value, int):
|
||||
raise ValueError(f"ollama_num_ctx must be a positive integer, got {value!r}")
|
||||
if value < 1:
|
||||
raise ValueError(f"ollama_num_ctx must be >= 1, got {value}")
|
||||
return value
|
||||
|
||||
|
||||
def create_llm_provider(
|
||||
provider: str,
|
||||
api_key: str,
|
||||
@@ -254,6 +309,7 @@ def create_llm_provider(
|
||||
litellmrouter_config: dict[str, Any] | None = None,
|
||||
gemini_service_tier: str | None = None,
|
||||
timeout: float | None = None,
|
||||
ollama_num_ctx: int | None = None,
|
||||
) -> Any: # Returns LLMInterface
|
||||
"""
|
||||
Factory function to create the appropriate LLM provider implementation.
|
||||
@@ -268,6 +324,8 @@ def create_llm_provider(
|
||||
openai_service_tier: OpenAI service tier (for OpenAI provider) - None (default) or "flex" (50% cheaper).
|
||||
bedrock_service_tier: Bedrock service tier (for Bedrock provider) - None (default), "flex", "priority", or "reserved".
|
||||
gemini_service_tier: Gemini service tier (for Gemini provider) - None (default) or "flex" (50% cheaper).
|
||||
ollama_num_ctx: Native Ollama context window override. None lets Ollama use the
|
||||
model/server default.
|
||||
extra_body: Extra request-body params merged into the provider's native
|
||||
call. Threaded into OpenAI-compatible, Fireworks, Anthropic, Gemini/
|
||||
VertexAI and LiteLLM providers (each merges them in its own parameter
|
||||
@@ -291,6 +349,8 @@ def create_llm_provider(
|
||||
Returns:
|
||||
LLMInterface implementation for the specified provider.
|
||||
"""
|
||||
ollama_num_ctx = _validate_ollama_num_ctx(ollama_num_ctx)
|
||||
|
||||
from .providers import (
|
||||
AnthropicLLM,
|
||||
ClaudeCodeLLM,
|
||||
@@ -303,6 +363,7 @@ def create_llm_provider(
|
||||
MockLLM,
|
||||
NoneLLM,
|
||||
OpenAICompatibleLLM,
|
||||
OpenAIResponsesLLM,
|
||||
)
|
||||
|
||||
provider_lower = provider.lower()
|
||||
@@ -320,6 +381,7 @@ def create_llm_provider(
|
||||
base_url=base_url,
|
||||
model=model,
|
||||
reasoning_effort=reasoning_effort,
|
||||
extra_body=extra_body,
|
||||
)
|
||||
|
||||
elif provider_lower == "claude-code":
|
||||
@@ -470,6 +532,22 @@ def create_llm_provider(
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
elif provider_lower == "openai-responses":
|
||||
# OpenAI Responses API (/v1/responses). Unlike chat/completions, it
|
||||
# supports reasoning + function tools together, so reflect's tool loop
|
||||
# can run with a real reasoning_effort. See OpenAIResponsesLLM.
|
||||
return OpenAIResponsesLLM(
|
||||
provider=provider,
|
||||
api_key=api_key,
|
||||
base_url=base_url,
|
||||
model=model,
|
||||
reasoning_effort=reasoning_effort,
|
||||
openai_service_tier=openai_service_tier,
|
||||
extra_body=extra_body,
|
||||
default_headers=default_headers,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
elif provider_lower in (
|
||||
"openai",
|
||||
"groq",
|
||||
@@ -494,6 +572,7 @@ def create_llm_provider(
|
||||
groq_service_tier=groq_service_tier,
|
||||
openai_service_tier=openai_service_tier,
|
||||
extra_body=extra_body,
|
||||
ollama_num_ctx=ollama_num_ctx,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
@@ -531,6 +610,7 @@ class LLMProvider:
|
||||
max_retries: int | None = None,
|
||||
initial_backoff: float | None = None,
|
||||
max_backoff: float | None = None,
|
||||
ollama_num_ctx: int | None = None,
|
||||
):
|
||||
"""
|
||||
Initialize LLM provider.
|
||||
@@ -545,6 +625,8 @@ class LLMProvider:
|
||||
openai_service_tier: OpenAI service tier (None or "flex") - from config.
|
||||
bedrock_service_tier: Bedrock service tier (None, "flex", "priority", "reserved") - from config.
|
||||
gemini_service_tier: Gemini service tier (None or "flex") - from config.
|
||||
ollama_num_ctx: Native Ollama context window override. ``None`` lets Ollama
|
||||
use the model/server default.
|
||||
gemini_safety_settings: Safety settings for Gemini/VertexAI providers.
|
||||
extra_body: Extra request-body params merged into the provider's native call
|
||||
(OpenAI-compatible, Fireworks, Anthropic, Gemini/VertexAI, LiteLLM).
|
||||
@@ -598,6 +680,7 @@ class LLMProvider:
|
||||
self.openai_service_tier = openai_service_tier
|
||||
self.bedrock_service_tier = bedrock_service_tier
|
||||
self.gemini_service_tier = gemini_service_tier
|
||||
self.ollama_num_ctx = _validate_ollama_num_ctx(ollama_num_ctx)
|
||||
# Gemini safety settings (instance default; can be overridden per-request via context var)
|
||||
self.gemini_safety_settings = gemini_safety_settings
|
||||
# Gemini prompt caching: when True, retain extraction (and any future
|
||||
@@ -615,6 +698,7 @@ class LLMProvider:
|
||||
# Validate provider
|
||||
valid_providers = [
|
||||
"openai",
|
||||
"openai-responses",
|
||||
"groq",
|
||||
"ollama",
|
||||
"ollama-cloud",
|
||||
@@ -742,6 +826,7 @@ class LLMProvider:
|
||||
gemini_safety_settings=self.gemini_safety_settings,
|
||||
prompt_cache_enabled=self.prompt_cache_enabled,
|
||||
litellmrouter_config=router_config,
|
||||
ollama_num_ctx=self.ollama_num_ctx,
|
||||
timeout=self.timeout,
|
||||
)
|
||||
|
||||
@@ -803,7 +888,7 @@ class LLMProvider:
|
||||
initial_backoff: float | None = None,
|
||||
max_backoff: float | None = None,
|
||||
skip_validation: bool = False,
|
||||
strict_schema: bool = False,
|
||||
strict_schema: bool | None = None,
|
||||
return_usage: bool = False,
|
||||
cached_prefix: str | None = None,
|
||||
) -> Any:
|
||||
@@ -824,9 +909,10 @@ class LLMProvider:
|
||||
configured default (``llm_max_backoff``), else 60.0.
|
||||
skip_validation: Return raw JSON without Pydantic validation.
|
||||
strict_schema: Per-call override requesting grammar-enforced (json_schema strict)
|
||||
structured output instead of the soft json_object path. The server-level
|
||||
HINDSIGHT_API_LLM_STRICT_SCHEMA flag is OR-ed in here so it applies to every call;
|
||||
providers without a strict mode ignore it.
|
||||
structured output instead of the soft json_object path. None (the default)
|
||||
inherits the server-level HINDSIGHT_API_LLM_STRICT_SCHEMA flag; an explicit
|
||||
True or False wins over it, so a caller can force strict output on -- or off --
|
||||
for its own scope. Providers without a strict mode ignore it.
|
||||
return_usage: If True, return tuple (result, TokenUsage) instead of just result.
|
||||
|
||||
Returns:
|
||||
@@ -844,7 +930,13 @@ class LLMProvider:
|
||||
from ..worker.stage import set_stage
|
||||
|
||||
structured = "+structured" if response_format is not None else ""
|
||||
set_stage(f"llm.{self.provider}.{scope}{structured}")
|
||||
# `.queued` until the concurrency permits are in hand — see the acquire
|
||||
# below. Without it, a call waiting on a saturated semaphore is
|
||||
# indistinguishable from one the provider is actively running, and the
|
||||
# label points at the provider (#3002: an operator lost an hour to
|
||||
# "llm.bedrock.*" for tasks that had never reached Bedrock).
|
||||
base_stage = f"llm.{self.provider}.{scope}{structured}"
|
||||
set_stage(f"{base_stage}.queued")
|
||||
|
||||
# Resolve the retry policy: explicit per-call arg wins, else the provider's
|
||||
# configured per-operation/global default, else this method's own fallback.
|
||||
@@ -861,14 +953,18 @@ class LLMProvider:
|
||||
)
|
||||
|
||||
# Resolve strict-schema once, here, rather than in each provider: the
|
||||
# per-call argument OR the server-level HINDSIGHT_API_LLM_STRICT_SCHEMA
|
||||
# flag. Providers with a json_schema response_format (OpenAI-compatible,
|
||||
# per-call argument, falling back to the server-level
|
||||
# HINDSIGHT_API_LLM_STRICT_SCHEMA flag when the caller expressed no
|
||||
# preference. Providers with a json_schema response_format (OpenAI-compatible,
|
||||
# LiteLLM) then grammar-enforce structured output instead of the fragile
|
||||
# soft json_object path; Gemini already enforces its native response_schema,
|
||||
# and providers without a strict mode simply ignore the flag.
|
||||
from ..config import get_config
|
||||
|
||||
strict_schema = strict_schema or get_config().llm_strict_schema
|
||||
# An explicit per-call value wins in BOTH directions -- `or` would have made a
|
||||
# per-call False indistinguishable from "unset", silently ignoring any caller
|
||||
# that opts out while the global flag is on.
|
||||
strict_schema = strict_schema if strict_schema is not None else get_config().llm_strict_schema
|
||||
|
||||
# LLM call observability flows through the OTel GenAI recorder
|
||||
# (tracing.get_span_recorder().record_llm_call). Provider implementations
|
||||
@@ -897,9 +993,18 @@ class LLMProvider:
|
||||
# hand so the error path below can attach it if parsing/validation fails.
|
||||
usage_token = set_response_usage(None)
|
||||
try:
|
||||
# Providers that own retry loops acquire the shared permits for each
|
||||
# upstream attempt so backoff never occupies request capacity.
|
||||
attempt_gated = self._provider_impl.supports_attempt_scoped_concurrency()
|
||||
async with AsyncExitStack() as stack:
|
||||
for sem in _semaphores_for_scope(scope):
|
||||
await stack.enter_async_context(sem)
|
||||
if not attempt_gated:
|
||||
for sem in _semaphores_for_scope(scope):
|
||||
await stack.enter_async_context(sem)
|
||||
# Permits in hand — only now leave `.queued`. Attempt-gated
|
||||
# providers acquire permits per attempt instead, so they keep
|
||||
# `.queued` until their first `attempt=N` stamp lands after
|
||||
# the permit acquire inside attempt_context (#3002).
|
||||
set_stage(base_stage)
|
||||
|
||||
# cached_prefix is only set for providers that returned a handle
|
||||
# from get_or_create_cached_prefix() (e.g. Gemini); it's None for
|
||||
@@ -908,6 +1013,7 @@ class LLMProvider:
|
||||
cache_kwarg = {"cached_prefix": cached_prefix} if cached_prefix is not None else {}
|
||||
try:
|
||||
# Delegate to provider implementation
|
||||
attempt_kwarg = {"attempt_context": lambda: _attempt_permits(scope)} if attempt_gated else {}
|
||||
result = await self._provider_impl.call(
|
||||
messages=messages,
|
||||
response_format=response_format,
|
||||
@@ -921,6 +1027,7 @@ class LLMProvider:
|
||||
strict_schema=strict_schema,
|
||||
return_usage=return_usage,
|
||||
**cache_kwarg,
|
||||
**attempt_kwarg,
|
||||
)
|
||||
except Exception as e:
|
||||
# The provider call may have succeeded (and incurred token
|
||||
@@ -965,8 +1072,9 @@ class LLMProvider:
|
||||
max_retries: int | None = None,
|
||||
initial_backoff: float | None = None,
|
||||
max_backoff: float | None = None,
|
||||
tool_choice: str | dict[str, Any] = "auto",
|
||||
tool_choice: LLMToolChoice = LLM_TOOL_CHOICE_AUTO,
|
||||
cached_prefix: str | None = None,
|
||||
cached_prefix_message_count: int = 0,
|
||||
) -> "LLMToolCallResult":
|
||||
"""
|
||||
Make an LLM API call with tool/function calling support.
|
||||
@@ -983,14 +1091,16 @@ class LLMProvider:
|
||||
configured default (``llm_initial_backoff``), else 1.0.
|
||||
max_backoff: Maximum backoff time in seconds. ``None`` uses the provider's
|
||||
configured default (``llm_max_backoff``), else 30.0.
|
||||
tool_choice: How to choose tools - "auto", "none", "required", or {"type": "function", "function": {"name": "..."}}
|
||||
tool_choice: Canonical tool-selection policy.
|
||||
|
||||
Returns:
|
||||
LLMToolCallResult with content and/or tool_calls.
|
||||
"""
|
||||
from ..worker.stage import set_stage
|
||||
|
||||
set_stage(f"llm.{self.provider}.{scope}+tools")
|
||||
# `.queued` until the permits are held — see the structured path above.
|
||||
base_stage = f"llm.{self.provider}.{scope}+tools"
|
||||
set_stage(f"{base_stage}.queued")
|
||||
|
||||
# Resolve the retry policy: explicit per-call arg wins, else the provider's
|
||||
# configured per-operation/global default, else this method's own fallback.
|
||||
@@ -1029,16 +1139,28 @@ class LLMProvider:
|
||||
# hand so the error path below can attach it if parsing/validation fails.
|
||||
usage_token = set_response_usage(None)
|
||||
try:
|
||||
attempt_gated = self._provider_impl.supports_attempt_scoped_concurrency()
|
||||
async with AsyncExitStack() as stack:
|
||||
for sem in _semaphores_for_scope(scope):
|
||||
await stack.enter_async_context(sem)
|
||||
if not attempt_gated:
|
||||
for sem in _semaphores_for_scope(scope):
|
||||
await stack.enter_async_context(sem)
|
||||
# Permits in hand — only now leave `.queued`; attempt-gated
|
||||
# providers stay `.queued` until their first post-acquire
|
||||
# `attempt=N` stamp (see call() above, #3002).
|
||||
set_stage(base_stage)
|
||||
|
||||
# cached_prefix is only set for providers that returned a handle
|
||||
# from get_or_create_cached_prefix(); forward it only when present
|
||||
# so non-caching providers keep their signature (same as call()).
|
||||
cache_kwarg = {"cached_prefix": cached_prefix} if cached_prefix is not None else {}
|
||||
# from get_or_create_cached_prefix() / create_incremental_cache();
|
||||
# forward it (plus how many leading messages it covers) only when
|
||||
# present so non-caching providers keep their signature.
|
||||
cache_kwarg = (
|
||||
{"cached_prefix": cached_prefix, "cached_prefix_message_count": cached_prefix_message_count}
|
||||
if cached_prefix is not None
|
||||
else {}
|
||||
)
|
||||
try:
|
||||
# Delegate to provider implementation
|
||||
attempt_kwarg = {"attempt_context": lambda: _attempt_permits(scope)} if attempt_gated else {}
|
||||
result = await self._provider_impl.call_with_tools(
|
||||
messages=messages,
|
||||
tools=tools,
|
||||
@@ -1050,6 +1172,7 @@ class LLMProvider:
|
||||
max_backoff=max_backoff,
|
||||
tool_choice=tool_choice,
|
||||
**cache_kwarg,
|
||||
**attempt_kwarg,
|
||||
)
|
||||
except Exception as e:
|
||||
# The provider call may have succeeded (and incurred token
|
||||
@@ -1260,6 +1383,7 @@ class LLMProvider:
|
||||
ENV_LLM_GROQ_SERVICE_TIER,
|
||||
ENV_LLM_LITELLMROUTER_CONFIG,
|
||||
ENV_LLM_MODEL,
|
||||
ENV_LLM_OLLAMA_NUM_CTX,
|
||||
ENV_LLM_OPENAI_SERVICE_TIER,
|
||||
ENV_LLM_PROMPT_CACHE_ENABLED,
|
||||
ENV_LLM_PROVIDER,
|
||||
@@ -1270,6 +1394,7 @@ class LLMProvider:
|
||||
ENV_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY,
|
||||
_get_default_model_for_provider,
|
||||
_parse_llm_router_config,
|
||||
_parse_optional_positive_int,
|
||||
parse_gemini_service_tier,
|
||||
)
|
||||
|
||||
@@ -1314,6 +1439,7 @@ class LLMProvider:
|
||||
),
|
||||
gemini_safety_settings=json.loads(os.getenv(ENV_LLM_GEMINI_SAFETY_SETTINGS, "null")),
|
||||
prompt_cache_enabled=prompt_cache_enabled,
|
||||
ollama_num_ctx=_parse_optional_positive_int(ENV_LLM_OLLAMA_NUM_CTX, os.getenv(ENV_LLM_OLLAMA_NUM_CTX)),
|
||||
litellmrouter_config=_parse_llm_router_config(ENV_LLM_LITELLMROUTER_CONFIG),
|
||||
vertexai_project_id=os.getenv(ENV_LLM_VERTEXAI_PROJECT_ID) or None,
|
||||
vertexai_region=os.getenv(ENV_LLM_VERTEXAI_REGION) or None,
|
||||
|
||||
@@ -0,0 +1,172 @@
|
||||
"""Device selection and post-inference memory release for local (in-process)
|
||||
SentenceTransformer / CrossEncoder models.
|
||||
|
||||
Two concerns live here, both about keeping a local API instance's memory flat:
|
||||
|
||||
**1. Device selection — MPS is opt-in.**
|
||||
On Apple Silicon the PyTorch **MPS** (Metal) backend caches a distinct compiled
|
||||
kernel graph *and* allocator pool per unique input tensor shape, and never
|
||||
releases them. Under the variable-length, high-volume recall/rerank/embed traffic
|
||||
the engine generates (documents and candidate sets of every size), that per-shape
|
||||
cache grows without bound. A local instance was observed idling at ~20 GB — ~9.4 GB
|
||||
of Metal graphics memory plus ~8 GB of native heap, essentially all of it stale
|
||||
per-shape MPS cache. CPU inference has no per-shape cache: the same workload holds
|
||||
flat at a few hundred MB, with negligible latency cost for the small default
|
||||
models (and MPS actually *slows down* over time as it recompiles graphs for new
|
||||
shapes). So MPS is excluded from auto-detection and must be opted into explicitly;
|
||||
CUDA and Intel XPU still auto-select.
|
||||
|
||||
This is a confirmed, still-open PyTorch bug in the MPSGraph compilation cache
|
||||
(keyed on tensor shape, no eviction path). We are tracking it upstream:
|
||||
- https://github.com/pytorch/pytorch/issues/181213
|
||||
([MPS] unbounded RSS growth with varying-shape inference — our exact case)
|
||||
- https://github.com/pytorch/pytorch/issues/164299 (graphCache identified as
|
||||
the primary leak culprit)
|
||||
- https://github.com/pytorch/pytorch/issues/182815 (proposes, but has not yet
|
||||
shipped, a torch.mps.invalidate_graph_cache() API / PYTORCH_MPS_DISABLE_GRAPH_CACHE
|
||||
env var that would let us keep MPS)
|
||||
No released mitigation exists today: empty_cache(), synchronize(),
|
||||
PYTORCH_MPS_HIGH_WATERMARK_RATIO, and autorelease pools were all confirmed
|
||||
ineffective upstream. Revisit MPS-as-default once one of those knobs lands.
|
||||
|
||||
**2. Memory release after each batch.**
|
||||
Local CPU inference allocates large transient numpy/tensor buffers per call. The
|
||||
allocator keeps those freed pages as a high-water mark, so RSS grows monotonically
|
||||
across many calls (issue #1717). We return them to the OS after each batch —
|
||||
``malloc_trim`` on glibc/Linux, ``malloc_zone_pressure_relief`` on macOS (the
|
||||
original #1717 fix covered only Linux). When the model ran on a GPU we also empty
|
||||
that backend's allocator pool via ``torch.<backend>.empty_cache()``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import ctypes
|
||||
import ctypes.util
|
||||
import gc
|
||||
import logging
|
||||
import sys
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def select_local_device(force_cpu: bool, allow_mps: bool) -> str | None:
|
||||
"""Choose the device for a local SentenceTransformer / CrossEncoder.
|
||||
|
||||
Returns a value suitable to pass as the model's ``device`` argument:
|
||||
|
||||
- ``"cpu"`` — forced CPU, or the only accelerator is MPS and it is not allowed.
|
||||
- ``None`` — let sentence-transformers auto-detect (picks CUDA / XPU,
|
||||
handling multi-GPU correctly).
|
||||
- ``"mps"`` — Apple Silicon GPU, only when ``allow_mps`` is set.
|
||||
|
||||
MPS is never auto-selected because its per-shape cache leaks unbounded memory
|
||||
under the engine's variable-length workload (see the module docstring). Set the
|
||||
matching ``*_ALLOW_MPS`` config flag to opt back in.
|
||||
"""
|
||||
if force_cpu:
|
||||
return "cpu"
|
||||
|
||||
try:
|
||||
import torch
|
||||
|
||||
if torch.cuda.is_available():
|
||||
return None # auto-detect CUDA
|
||||
if hasattr(torch, "xpu") and torch.xpu.is_available():
|
||||
return None # auto-detect Intel XPU
|
||||
|
||||
mps_available = hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
|
||||
if mps_available:
|
||||
if allow_mps:
|
||||
return "mps"
|
||||
logger.info(
|
||||
"Local model: MPS (Apple Silicon GPU) is available but disabled by "
|
||||
"default because its per-shape cache leaks memory under variable-length "
|
||||
"workloads; running on CPU. Set the *_ALLOW_MPS flag to opt in."
|
||||
)
|
||||
return "cpu"
|
||||
return "cpu"
|
||||
except Exception as e: # pragma: no cover - defensive
|
||||
logger.warning("Local device detection failed, falling back to CPU: %s", e)
|
||||
return "cpu"
|
||||
|
||||
|
||||
def resolve_model_device_type(model: object) -> str:
|
||||
"""Best-effort device *type* ("cpu" / "cuda" / "mps" / "xpu") of a loaded model.
|
||||
|
||||
Used to decide which GPU allocator pool to empty after inference. Falls back to
|
||||
``"cpu"`` (the safe no-op choice for release) if the device can't be read.
|
||||
"""
|
||||
device = getattr(model, "device", None)
|
||||
if device is None:
|
||||
inner = getattr(model, "model", None) # CrossEncoder wraps the HF model
|
||||
device = getattr(inner, "device", None)
|
||||
try:
|
||||
return device.type if device is not None else "cpu"
|
||||
except Exception: # pragma: no cover - defensive
|
||||
return "cpu"
|
||||
|
||||
|
||||
def _resolve_heap_trim():
|
||||
"""Return a callable that asks the C allocator to release freed pages to the OS.
|
||||
|
||||
glibc (Linux) exposes ``malloc_trim``; macOS exposes
|
||||
``malloc_zone_pressure_relief``. Resolved once at import; returns a no-op on
|
||||
platforms where neither is available (musl, Windows).
|
||||
"""
|
||||
if sys.platform == "linux":
|
||||
libc_path = ctypes.util.find_library("c")
|
||||
if libc_path is None:
|
||||
return lambda: None
|
||||
try:
|
||||
libc = ctypes.CDLL(libc_path)
|
||||
trim = libc.malloc_trim
|
||||
except (OSError, AttributeError):
|
||||
# Not glibc (musl has no malloc_trim) or libc lookup failed.
|
||||
return lambda: None
|
||||
trim.argtypes = [ctypes.c_size_t]
|
||||
trim.restype = ctypes.c_int
|
||||
return lambda: trim(0)
|
||||
|
||||
if sys.platform == "darwin":
|
||||
try:
|
||||
libc = ctypes.CDLL("/usr/lib/libSystem.dylib")
|
||||
default_zone = libc.malloc_default_zone
|
||||
default_zone.restype = ctypes.c_void_p
|
||||
relief = libc.malloc_zone_pressure_relief
|
||||
relief.argtypes = [ctypes.c_void_p, ctypes.c_size_t]
|
||||
relief.restype = ctypes.c_size_t
|
||||
except (OSError, AttributeError):
|
||||
return lambda: None
|
||||
# pressure_relief(zone, goal=0) reclaims as much as possible.
|
||||
return lambda: relief(default_zone(), 0)
|
||||
|
||||
return lambda: None
|
||||
|
||||
|
||||
_heap_trim = _resolve_heap_trim()
|
||||
|
||||
|
||||
def _empty_gpu_cache(device_type: str | None) -> None:
|
||||
"""Empty the allocator pool of the GPU backend the model ran on, if any."""
|
||||
if not device_type or device_type == "cpu":
|
||||
return
|
||||
try:
|
||||
import torch
|
||||
|
||||
backend = getattr(torch, device_type, None) # torch.cuda / torch.mps / torch.xpu
|
||||
if backend is not None and hasattr(backend, "empty_cache"):
|
||||
backend.empty_cache()
|
||||
except Exception: # pragma: no cover - defensive
|
||||
pass
|
||||
|
||||
|
||||
def release_local_inference_memory(device_type: str | None = None) -> None:
|
||||
"""Release transient heap (and GPU allocator) memory after a local inference batch.
|
||||
|
||||
Frees Python objects, returns freed native pages to the OS, and empties the GPU
|
||||
allocator pool when the model ran on a GPU. Safe to call on every platform and
|
||||
device; the pieces that don't apply are cheap no-ops.
|
||||
"""
|
||||
gc.collect()
|
||||
_heap_trim()
|
||||
_empty_gpu_cache(device_type)
|
||||
@@ -21,9 +21,16 @@ from one place, so we don't spawn a separate ``asyncio`` task per concern:
|
||||
The loop wakes on a short fixed tick and runs each job when its own
|
||||
``last_run + interval`` is due (run-at-start, then on interval), so adding jobs
|
||||
with different cadences doesn't burst CPU. Cross-tenant discovery goes through
|
||||
server-side PL/pgSQL routines (``public.schemas_with_expired_rows`` and
|
||||
``public.banks_needing_consolidation``) — one round-trip each — instead of a
|
||||
per-schema query storm, which matters at thousands of tenants.
|
||||
server-side PL/pgSQL routines (``schemas_with_expired_rows`` and
|
||||
``banks_needing_consolidation``, in the configured schema — see ``fq_routine``) —
|
||||
one round-trip each — instead of a per-schema query storm, which matters at
|
||||
thousands of tenants.
|
||||
|
||||
The loop runs in *every* API/worker process with no leader election, so a job that
|
||||
enqueues work must make that enqueue idempotent or the fleet queues one wave per
|
||||
process. Retention and operation cleanup are deletes; the consolidation reconcile
|
||||
and the scheduled mental model refresh both dedupe against in-flight operations
|
||||
inside the inserting transaction (see ``_submit_async_operation``).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -32,13 +39,13 @@ import asyncio
|
||||
import logging
|
||||
import time
|
||||
from collections.abc import Coroutine
|
||||
from datetime import datetime, timezone
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from ..config import HindsightConfig, get_config
|
||||
from ..models import RequestContext
|
||||
from .db_utils import acquire_with_retry
|
||||
from .schema import _is_oracle, fq_table
|
||||
from .schema import _is_oracle, fq_routine, fq_table, fq_table_explicit
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .memory_engine import MemoryEngine
|
||||
@@ -49,6 +56,18 @@ logger = logging.getLogger(__name__)
|
||||
_TICK_SECONDS = 60
|
||||
# Retention sweeps are not time-sensitive; hourly matches the previous per-sweep cadence.
|
||||
_RETENTION_INTERVAL_SECONDS = 3600
|
||||
# Operation cleanup deletes one bounded batch per schema per run, so its cadence
|
||||
# sets the drain rate for a backlog. Kept at one-per-tick (the value it used while
|
||||
# it rode the worker's poll loop) so throughput is unchanged by the move.
|
||||
_OPERATION_CLEANUP_INTERVAL_SECONDS = 60
|
||||
# Cross-store txn recovery (only when the memories store keeps its rows outside SQL): a backstop
|
||||
# for a writer that crashed between its external writes and the decide. The happy path decides
|
||||
# inline after commit, so this rarely finds work; five minutes bounds how long a crashed txn stalls
|
||||
# its namespace's fold.
|
||||
_TXN_RECOVERY_INTERVAL_SECONDS = 300
|
||||
# A pending txn is left alone for this long from first sighting before the sweep aborts an
|
||||
# unwitnessed one — the writer may still be mid-flight (PendingTxn carries no timestamp).
|
||||
_TXN_RECOVERY_GRACE_SECONDS = 300
|
||||
|
||||
|
||||
class MaintenanceLoop:
|
||||
@@ -60,6 +79,9 @@ class MaintenanceLoop:
|
||||
self._stop = asyncio.Event()
|
||||
# Monotonic timestamps of the last run per job, keyed by job name.
|
||||
self._last_run: dict[str, float] = {}
|
||||
# Cross-store txn recovery: first-sighting time per pending txn_id, so an unwitnessed
|
||||
# txn gets a grace period before the sweep aborts it. Persists across ticks.
|
||||
self._txn_first_seen: dict[str, float] = {}
|
||||
|
||||
# ── lifecycle ──────────────────────────────────────────────────────────
|
||||
|
||||
@@ -97,10 +119,32 @@ class MaintenanceLoop:
|
||||
def _any_job_enabled() -> bool:
|
||||
cfg = get_config()
|
||||
reconcile_on = cfg.consolidation_reconcile_interval_seconds > 0
|
||||
audit_on = cfg.audit_log_enabled and cfg.audit_log_retention_days > 0
|
||||
# Not gated on audit_log_enabled: that is per-bank overridable, so rows
|
||||
# can exist even when the deployment default is off. Retention is driven
|
||||
# purely by the (server-level) window.
|
||||
audit_on = cfg.audit_log_retention_days > 0
|
||||
llm_on = cfg.llm_trace_enabled and cfg.llm_trace_retention_days > 0
|
||||
mm_refresh_on = cfg.mental_model_refresh_tick_seconds > 0
|
||||
return reconcile_on or audit_on or llm_on or mm_refresh_on
|
||||
op_cleanup_on = cfg.operation_retention_days > 0
|
||||
return (
|
||||
reconcile_on
|
||||
or audit_on
|
||||
or llm_on
|
||||
or mm_refresh_on
|
||||
or op_cleanup_on
|
||||
or MaintenanceLoop._cross_store_recovery_enabled()
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _cross_store_recovery_enabled() -> bool:
|
||||
"""True when the memories store keeps memories outside SQL and therefore has
|
||||
cross-store write-group txns a crashed writer could leave undecided."""
|
||||
try:
|
||||
from .memories import get_memories
|
||||
|
||||
return not get_memories().writes_memory_rows_in_sql
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
# ── loop ───────────────────────────────────────────────────────────────
|
||||
|
||||
@@ -134,6 +178,10 @@ class MaintenanceLoop:
|
||||
mm_interval = cfg.mental_model_refresh_tick_seconds
|
||||
if mm_interval > 0 and self._is_due("mm_refresh", mm_interval):
|
||||
await self._run_timed("scheduled mental model refresh", self._run_scheduled_mm_refresh())
|
||||
if cfg.operation_retention_days > 0 and self._is_due("operation_cleanup", _OPERATION_CLEANUP_INTERVAL_SECONDS):
|
||||
await self._run_timed("operation cleanup", self._run_operation_cleanup(cfg))
|
||||
if self._cross_store_recovery_enabled() and self._is_due("txn_recovery", _TXN_RECOVERY_INTERVAL_SECONDS):
|
||||
await self._run_timed("cross-store txn recovery", self._run_txn_recovery())
|
||||
|
||||
async def _run_timed(self, name: str, coro: Coroutine[Any, Any, None]) -> None:
|
||||
"""Run a maintenance job and emit one timing line for it.
|
||||
@@ -152,7 +200,10 @@ class MaintenanceLoop:
|
||||
async def _run_retention(self, cfg: HindsightConfig) -> None:
|
||||
# Retention days are static server-level config, so one global cutoff
|
||||
# applies to every tenant schema (the routine sweeps them all).
|
||||
if cfg.audit_log_enabled and cfg.audit_log_retention_days > 0:
|
||||
# Not gated on audit_log_enabled: it is per-bank overridable, so a bank
|
||||
# may be writing audit rows while the deployment default is off. Gating
|
||||
# the purge on the global flag would let those rows accumulate forever.
|
||||
if cfg.audit_log_retention_days > 0:
|
||||
await self._purge_expired("audit_log", "started_at", cfg.audit_log_retention_days)
|
||||
if cfg.llm_trace_enabled and cfg.llm_trace_retention_days > 0:
|
||||
await self._purge_expired("llm_requests", "started_at", cfg.llm_trace_retention_days)
|
||||
@@ -163,7 +214,7 @@ class MaintenanceLoop:
|
||||
try:
|
||||
async with acquire_with_retry(backend, max_retries=1) as conn:
|
||||
rows = await conn.fetch(
|
||||
"SELECT * FROM public.schemas_with_expired_rows($1, $2, $3)", table, ts_col, days
|
||||
f"SELECT * FROM {fq_routine('schemas_with_expired_rows')}($1, $2, $3)", table, ts_col, days
|
||||
)
|
||||
for row in rows:
|
||||
schema = row[0]
|
||||
@@ -178,6 +229,109 @@ class MaintenanceLoop:
|
||||
except Exception as e:
|
||||
logger.warning(f"Retention sweep failed for {table}: {e}")
|
||||
|
||||
# ── terminal operation cleanup ─────────────────────────────────────────
|
||||
|
||||
async def _run_operation_cleanup(self, cfg: HindsightConfig) -> None:
|
||||
"""Prune one bounded batch of expired terminal operations per tenant schema.
|
||||
|
||||
Previously this rode the worker's task-claiming loop, so it only fired
|
||||
when that loop happened to iterate and was interleaved with claiming. It
|
||||
is a periodic housekeeping sweep like the retention jobs above, so it
|
||||
belongs on the same schedule.
|
||||
|
||||
Discovery is one cross-tenant round-trip (``schemas_with_expired_operations``)
|
||||
rather than a connection + prune transaction per tenant; pending and
|
||||
processing rows are never prunable, so a schema holding only in-flight
|
||||
work is correctly reported as having nothing to do.
|
||||
"""
|
||||
engine = self._engine
|
||||
backend = engine._backend
|
||||
try:
|
||||
async with acquire_with_retry(backend, max_retries=1) as conn:
|
||||
rows = await conn.fetch(
|
||||
f"SELECT * FROM {fq_routine('schemas_with_expired_operations')}($1)",
|
||||
cfg.operation_retention_days,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"Operation cleanup discovery failed: {e}")
|
||||
return
|
||||
if not rows:
|
||||
return
|
||||
|
||||
# Prune only schemas the deployment actually serves. The routine reports
|
||||
# every schema owning an async_operations table, including ones tenant
|
||||
# discovery doesn't claim.
|
||||
try:
|
||||
tenants = await engine._tenant_extension.list_tenants()
|
||||
except Exception as e:
|
||||
logger.warning(f"Operation cleanup tenant discovery failed: {e}")
|
||||
return
|
||||
known = {t.schema for t in tenants} | {get_config().database_schema}
|
||||
|
||||
from .memory_engine import _current_schema
|
||||
|
||||
cutoff = datetime.now(timezone.utc) - timedelta(days=cfg.operation_retention_days)
|
||||
pruned = 0
|
||||
for row in rows:
|
||||
schema = row[0]
|
||||
if schema not in known:
|
||||
continue
|
||||
# Oracle resolves unqualified names from a context-bound session
|
||||
# schema; on PostgreSQL this is harmless and fq_table stays explicit.
|
||||
token = _current_schema.set(schema)
|
||||
try:
|
||||
table = fq_table_explicit("async_operations", schema)
|
||||
async with acquire_with_retry(backend, max_retries=1) as conn:
|
||||
async with conn.transaction():
|
||||
deleted = await backend.ops.prune_terminal_operations(
|
||||
conn, table, cutoff, batch_size=cfg.operation_cleanup_batch_size
|
||||
)
|
||||
if deleted:
|
||||
pruned += deleted
|
||||
logger.info(f"Operation cleanup pruned {deleted} expired terminal operations from {schema}")
|
||||
except Exception as e:
|
||||
logger.warning(f"Operation cleanup failed for schema {schema}: {e}")
|
||||
finally:
|
||||
_current_schema.reset(token)
|
||||
if pruned:
|
||||
logger.info(f"Operation cleanup: pruned {pruned} operation(s) total")
|
||||
|
||||
# ── cross-store txn recovery ─────────────────────────────────────────────
|
||||
|
||||
async def _run_txn_recovery(self) -> None:
|
||||
"""Resolve write-group txns a crashed writer left undecided, for a store that keeps its
|
||||
rows outside SQL.
|
||||
|
||||
For each bank, the store lists its namespace's pending txns and decides each against the
|
||||
Postgres witness table (present ⇒ commit, absent past the grace ⇒ abort — never on
|
||||
assumption), then reaps expired witness rows. A no-op for the SQL stores. Best-effort: a
|
||||
failure here only delays a stalled fold until the next tick.
|
||||
"""
|
||||
from .memories import get_memories
|
||||
|
||||
store = get_memories()
|
||||
if store.writes_memory_rows_in_sql:
|
||||
return
|
||||
backend = self._engine._backend
|
||||
try:
|
||||
async with acquire_with_retry(backend, max_retries=1) as conn:
|
||||
bank_ids = [r[0] for r in await conn.fetch(f"SELECT bank_id FROM {fq_table('banks')}")]
|
||||
if not bank_ids:
|
||||
return
|
||||
decided = await store.recover_pending_txns(
|
||||
conn=conn,
|
||||
fq_table=fq_table,
|
||||
bank_ids=bank_ids,
|
||||
first_seen=self._txn_first_seen,
|
||||
now=time.monotonic(),
|
||||
grace_seconds=_TXN_RECOVERY_GRACE_SECONDS,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"Cross-store txn recovery failed: {e}")
|
||||
return
|
||||
if decided:
|
||||
logger.info(f"Cross-store txn recovery: decided {decided} undecided txn(s)")
|
||||
|
||||
# ── consolidation reconcile ──────────────────────────────────────────────
|
||||
|
||||
async def _run_reconcile(self) -> None:
|
||||
@@ -185,7 +339,9 @@ class MaintenanceLoop:
|
||||
engine = self._engine
|
||||
try:
|
||||
async with acquire_with_retry(engine._backend, max_retries=1) as conn:
|
||||
rows = await conn.fetch("SELECT schema_name, bank_id FROM public.banks_needing_consolidation()")
|
||||
rows = await conn.fetch(
|
||||
f"SELECT schema_name, bank_id FROM {fq_routine('banks_needing_consolidation')}()"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"Consolidation reconcile discovery failed: {e}")
|
||||
return
|
||||
@@ -244,7 +400,7 @@ class MaintenanceLoop:
|
||||
|
||||
Discovery (the set of cron-scheduled models, minus any with an in-flight
|
||||
refresh) is one cross-tenant round-trip via
|
||||
``public.mental_models_with_cron()``. Cron *due-ness* is evaluated here in
|
||||
``mental_models_with_cron()``. Cron *due-ness* is evaluated here in
|
||||
Python — a scheduled fire has elapsed when the most recent cron boundary at
|
||||
or before now is later than ``last_refreshed_at`` — because cron arithmetic
|
||||
isn't expressible in plain SQL. Each due model is refreshed only when it is
|
||||
@@ -256,7 +412,7 @@ class MaintenanceLoop:
|
||||
async with acquire_with_retry(engine._backend, max_retries=1) as conn:
|
||||
rows = await conn.fetch(
|
||||
"SELECT schema_name, bank_id, mental_model_id, refresh_cron, last_refreshed_at "
|
||||
"FROM public.mental_models_with_cron()"
|
||||
f"FROM {fq_routine('mental_models_with_cron')}()"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"Scheduled mental model refresh discovery failed: {e}")
|
||||
@@ -300,6 +456,7 @@ class MaintenanceLoop:
|
||||
submitted = 0
|
||||
skipped_unknown = 0
|
||||
skipped_fresh = 0
|
||||
skipped_in_flight = 0
|
||||
for row in due:
|
||||
schema = row["schema_name"]
|
||||
bank_id = row["bank_id"]
|
||||
@@ -330,18 +487,28 @@ class MaintenanceLoop:
|
||||
if not is_stale:
|
||||
skipped_fresh += 1
|
||||
continue
|
||||
await engine.submit_async_refresh_mental_model(
|
||||
bank_id=bank_id, mental_model_id=mm_id, request_context=context
|
||||
# skip_if_in_flight makes the enqueue itself idempotent. The discovery
|
||||
# routine already excludes models with a pending/processing refresh,
|
||||
# but that exclusion is a *read*: this loop runs in every process, so
|
||||
# every process saw the same "nothing in flight" snapshot and inserted
|
||||
# its own operation — one queued wave per process (#3210). The insert
|
||||
# now carries the check, so a second one is never created.
|
||||
result = await engine.submit_async_refresh_mental_model(
|
||||
bank_id=bank_id, mental_model_id=mm_id, request_context=context, skip_if_in_flight=True
|
||||
)
|
||||
submitted += 1
|
||||
if result.get("deduplicated"):
|
||||
skipped_in_flight += 1
|
||||
else:
|
||||
submitted += 1
|
||||
except Exception as e:
|
||||
logger.warning(f"Scheduled mental model refresh failed for {mm_id} in {schema}: {e}")
|
||||
finally:
|
||||
_current_schema.reset(token)
|
||||
|
||||
if submitted or skipped_unknown or skipped_fresh:
|
||||
if submitted or skipped_unknown or skipped_fresh or skipped_in_flight:
|
||||
logger.info(
|
||||
f"Scheduled mental model refresh: scheduled {submitted} model(s)"
|
||||
+ (f", {skipped_fresh} up-to-date" if skipped_fresh else "")
|
||||
+ (f", {skipped_in_flight} already in flight" if skipped_in_flight else "")
|
||||
+ (f", skipped {skipped_unknown} in unrecognized schema(s)" if skipped_unknown else "")
|
||||
)
|
||||
|
||||
@@ -0,0 +1,86 @@
|
||||
"""The memories store: which one is installed, and how the engine reaches it.
|
||||
|
||||
Resolved through the ordinary extension loader — ``HINDSIGHT_API_MEMORIES_EXTENSION``
|
||||
names a ``module:Class``, and ``HINDSIGHT_API_MEMORIES_*`` becomes its config — so
|
||||
this behaves like every other extension point. Unset (the normal case) means
|
||||
:class:`~hindsight_api.engine.memories.postgres.PostgresMemories`: rows in
|
||||
`memory_units`, links in `memory_links` / `unit_entities`, retrieval as SQL.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
|
||||
from .base import (
|
||||
META_CHUNK_ID,
|
||||
CausalEdgeRecord,
|
||||
DeletePredicate,
|
||||
FactRecord,
|
||||
MemoriesExtension,
|
||||
MemoryPatch,
|
||||
ScanPage,
|
||||
StoredMemory,
|
||||
build_fact_records,
|
||||
build_text_signals,
|
||||
source_key,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_memories: MemoriesExtension | None = None
|
||||
|
||||
|
||||
def create_memories(context=None) -> MemoriesExtension:
|
||||
"""Build the configured memories store, or the Postgres default."""
|
||||
from ...extensions.loader import load_extension
|
||||
|
||||
loaded = load_extension("MEMORIES", MemoriesExtension, context=context)
|
||||
if loaded is not None:
|
||||
logger.info("[memories] store=%s (memory rows do not go to postgres)", loaded.name)
|
||||
return loaded
|
||||
|
||||
from .postgres import PostgresMemories
|
||||
|
||||
return PostgresMemories({})
|
||||
|
||||
|
||||
def get_memories() -> MemoriesExtension:
|
||||
"""The process-wide memories store, built on first use.
|
||||
|
||||
Retrieval and the retain pipeline reach it through call chains that do not
|
||||
carry the engine, so it is resolved here rather than threaded through every
|
||||
signature.
|
||||
"""
|
||||
global _memories
|
||||
if _memories is None:
|
||||
_memories = create_memories()
|
||||
return _memories
|
||||
|
||||
|
||||
def set_memories(memories: MemoriesExtension | None) -> None:
|
||||
"""Override the store (tests, and engine startup after initialize())."""
|
||||
global _memories
|
||||
_memories = memories
|
||||
# The graph arm's retriever is chosen from the store and then cached, so it
|
||||
# has to be re-resolved whenever the store changes.
|
||||
from ..search.retrieval import set_default_graph_retriever
|
||||
|
||||
set_default_graph_retriever(None)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"META_CHUNK_ID",
|
||||
"CausalEdgeRecord",
|
||||
"DeletePredicate",
|
||||
"FactRecord",
|
||||
"MemoriesExtension",
|
||||
"MemoryPatch",
|
||||
"ScanPage",
|
||||
"StoredMemory",
|
||||
"build_fact_records",
|
||||
"build_text_signals",
|
||||
"create_memories",
|
||||
"get_memories",
|
||||
"set_memories",
|
||||
"source_key",
|
||||
]
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,20 @@
|
||||
"""The Postgres memories implementation, split by what calls it.
|
||||
|
||||
:class:`~hindsight_api.engine.memories.postgres.PostgresMemories` is a thin class
|
||||
over these modules; the queries live here, grouped by concern rather than piled
|
||||
behind one object:
|
||||
|
||||
* :mod:`counts` — the stats/admin aggregates (freshness, per-doc, timeseries, scopes)
|
||||
* :mod:`curation` — the memory/entity list and detail views
|
||||
* :mod:`graph` — the graph view, entity postings, and the maintenance passes
|
||||
* :mod:`reads` — addressed reads: get, scan, count, tags, consolidation state
|
||||
* :mod:`writes` — inserts, deletes, and observation invalidation
|
||||
|
||||
Every function here takes the live connection and Hindsight's ``fq_table``
|
||||
resolver rather than reaching for globals, so each is callable from a
|
||||
transaction the caller already owns.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
__all__ = ["counts", "curation", "graph", "reads", "writes"]
|
||||
@@ -0,0 +1,157 @@
|
||||
"""The count/aggregate surfaces: consolidation freshness, per-document counts,
|
||||
ingestion over time, observation scopes.
|
||||
|
||||
Each is one ``GROUP BY`` (or filtered ``COUNT``) over `memory_units`. They back
|
||||
the stats and admin views, not retrieval, so they are grouped here away from the
|
||||
addressed reads. The SQL is lifted verbatim from the engine methods that used to
|
||||
carry it; only the connection and ``fq_table`` resolver are now parameters.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable
|
||||
from datetime import datetime
|
||||
from typing import Any
|
||||
|
||||
|
||||
async def consolidation_freshness(*, conn, fq_table: Callable[[str], str], bank_id: str) -> dict[str, Any]:
|
||||
"""Last consolidation time, the pending / failed fact counts, and the write watermark, in one scan.
|
||||
|
||||
All four come from a single pass so keeping ``failed`` — part of the
|
||||
published contract — costs nothing over reflect()'s ``pending`` read, and
|
||||
``last_memory_write_at`` (the newest ``updated_at`` anywhere in the bank)
|
||||
rides along for free. That watermark is what lets a caller decide a mental
|
||||
model is up to date without running its own scoped scan: nothing in the bank
|
||||
changed since the refresh, so nothing in the model's scope did either.
|
||||
"""
|
||||
row = await conn.fetchrow(
|
||||
f"""
|
||||
SELECT
|
||||
MAX(consolidated_at) AS last_consolidated_at,
|
||||
MAX(updated_at) AS last_memory_write_at,
|
||||
COUNT(*) FILTER (WHERE consolidated_at IS NULL AND fact_type IN ('experience', 'world')) AS pending,
|
||||
COUNT(*) FILTER (WHERE consolidation_failed_at IS NOT NULL AND fact_type IN ('experience', 'world')) AS failed
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $1
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
if row is None:
|
||||
return {"last_consolidated_at": None, "last_memory_write_at": None, "pending": 0, "failed": 0}
|
||||
return {
|
||||
"last_consolidated_at": row["last_consolidated_at"],
|
||||
"last_memory_write_at": row["last_memory_write_at"],
|
||||
"pending": row["pending"] or 0,
|
||||
"failed": row["failed"] or 0,
|
||||
}
|
||||
|
||||
|
||||
async def link_counts(*, conn, fq_table: Callable[[str], str], bank_id: str) -> dict[str, int]:
|
||||
"""``{link_type: count}`` of live links in a bank.
|
||||
|
||||
Non-entity links (temporal / semantic / caused_by) are a single ``GROUP BY`` over
|
||||
``memory_links``. Entity links are no longer stored there — they are derived on demand
|
||||
from ``unit_entities``, replicating the historical writer cap of ``MAX_LINKS_PER_ENTITY``
|
||||
bidirectional edges per shared entity — so they are aggregated to one ``entity`` scalar.
|
||||
"""
|
||||
max_links_per_entity = 10
|
||||
non_entity_link_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT link_type, COUNT(*) as count
|
||||
FROM {fq_table("memory_links")}
|
||||
WHERE bank_id = $1
|
||||
GROUP BY link_type
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
entity_total_row = await conn.fetchrow(
|
||||
f"""
|
||||
WITH per_entity AS (
|
||||
SELECT ue.entity_id, COUNT(*) AS n
|
||||
FROM {fq_table("unit_entities")} ue
|
||||
JOIN {fq_table("memory_units")} mu ON mu.id = ue.unit_id
|
||||
WHERE mu.bank_id = $1
|
||||
GROUP BY ue.entity_id
|
||||
)
|
||||
SELECT COALESCE(SUM(LEAST(n - 1, $2)), 0)::bigint AS count
|
||||
FROM per_entity
|
||||
""",
|
||||
bank_id,
|
||||
max_links_per_entity,
|
||||
)
|
||||
entity_link_total = int(entity_total_row["count"] or 0) if entity_total_row else 0
|
||||
counts: dict[str, int] = {row["link_type"]: row["count"] for row in non_entity_link_rows}
|
||||
if entity_link_total > 0:
|
||||
counts["entity"] = entity_link_total
|
||||
return counts
|
||||
|
||||
|
||||
async def document_memory_counts(
|
||||
*, conn, fq_table: Callable[[str], str], bank_id: str, document_ids: list[str]
|
||||
) -> dict[str, int]:
|
||||
"""Live memory count per document id, for the ids given."""
|
||||
if not document_ids:
|
||||
return {}
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT document_id, COUNT(*) AS unit_count
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $1 AND document_id = ANY($2::text[])
|
||||
GROUP BY document_id
|
||||
""",
|
||||
bank_id,
|
||||
list(document_ids),
|
||||
)
|
||||
return {row["document_id"]: row["unit_count"] for row in rows}
|
||||
|
||||
|
||||
async def memories_timeseries(
|
||||
*, conn, fq_table: Callable[[str], str], bank_id: str, time_field: str, trunc: str, since: datetime
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Memories bucketed by ``time_field`` (truncated to ``trunc``) and fact_type.
|
||||
|
||||
``time_field`` is whitelisted by the caller before it reaches here — it is
|
||||
interpolated into SQL. Event-time fields fall back to ``created_at`` per row so
|
||||
rows without an event timestamp still appear.
|
||||
"""
|
||||
bucket_expr = time_field if time_field == "created_at" else f"COALESCE({time_field}, created_at)"
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT date_trunc('{trunc}', {bucket_expr} AT TIME ZONE 'UTC') AS bucket,
|
||||
fact_type, COUNT(*) AS count
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $1 AND {bucket_expr} >= $2
|
||||
GROUP BY bucket, fact_type
|
||||
ORDER BY bucket
|
||||
""",
|
||||
bank_id,
|
||||
since,
|
||||
)
|
||||
return [{"bucket": r["bucket"], "fact_type": r["fact_type"], "count": r["count"]} for r in rows]
|
||||
|
||||
|
||||
async def observation_scope_counts(*, conn, fq_table: Callable[[str], str], bank_id: str) -> list[dict[str, Any]]:
|
||||
"""Observations grouped by scope (their sorted tag set), most-populous first."""
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT scope, COUNT(*) AS count
|
||||
FROM (
|
||||
SELECT COALESCE(ARRAY(SELECT unnest(tags) ORDER BY 1), '{{}}'::text[]) AS scope
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $1 AND fact_type = 'observation'
|
||||
) s
|
||||
GROUP BY scope
|
||||
ORDER BY count DESC, scope
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
return [{"tags": list(r["scope"]), "count": r["count"]} for r in rows]
|
||||
|
||||
|
||||
__all__ = [
|
||||
"consolidation_freshness",
|
||||
"document_memory_counts",
|
||||
"link_counts",
|
||||
"memories_timeseries",
|
||||
"observation_scope_counts",
|
||||
]
|
||||
@@ -0,0 +1,507 @@
|
||||
"""Curation reads: the memory list, the memory detail view, and the entity list.
|
||||
|
||||
These back the curation UI — the table of memories a bank holds, the detail panel
|
||||
for one of them, and the entity roster beside it. They are paged and filtered
|
||||
rather than ranked: nothing here scores anything, and nothing walks the corpus.
|
||||
|
||||
Two things separate them from the addressed reads in :mod:`reads`. They render
|
||||
*view* dicts (ISO strings, joined entity names, a ``state`` discriminator) rather
|
||||
than :class:`~hindsight_api.engine.memories.base.StoredMemory`, because the HTTP
|
||||
layer serialises what comes back verbatim. And they read the archive as well as
|
||||
the live table: curation moves an invalidated fact to `invalidated_memory_units`,
|
||||
so "show me the invalidated ones" is a different table, not a different predicate.
|
||||
|
||||
Authentication, operation validation and audit stay with the engine methods that
|
||||
call these — only the queries and their row rendering live here.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from datetime import datetime
|
||||
from typing import Any
|
||||
|
||||
from ...search.tags import build_tags_where_clause
|
||||
|
||||
|
||||
def _entity_rows_for_units_sql(*, ops, fq_table, unit_ids_placeholder: int) -> str:
|
||||
"""SQL SELECT producing ``(unit_id, entity_id, canonical_name)`` rows for
|
||||
the given unit IDs.
|
||||
|
||||
Direct rows come from ``unit_entities``. Observations rarely carry
|
||||
direct rows there; their entity association lives transitively through
|
||||
their source memories (``source_memory_ids`` on PG, the
|
||||
``observation_sources`` junction on Oracle). When an observation has
|
||||
no direct entity rows the SELECT inherits its source memories'
|
||||
entities, so the result is the same set callers would get from
|
||||
``get_memory_unit``.
|
||||
|
||||
``unit_ids_placeholder`` is the 1-based parameter index that holds the
|
||||
``uuid[]`` of unit IDs. The placeholder is referenced twice — both
|
||||
sides of the UNION need it — so callers should not reuse the slot.
|
||||
"""
|
||||
ue = fq_table("unit_entities")
|
||||
ents = fq_table("entities")
|
||||
mu = fq_table("memory_units")
|
||||
p = unit_ids_placeholder
|
||||
|
||||
direct = (
|
||||
f"SELECT ue.unit_id, e.id AS entity_id, e.canonical_name "
|
||||
f"FROM {ue} ue "
|
||||
f"JOIN {ents} e ON e.id = ue.entity_id "
|
||||
f"WHERE ue.unit_id = ANY(${p}::uuid[])"
|
||||
)
|
||||
|
||||
if ops.uses_observation_sources_table:
|
||||
os_t = fq_table("observation_sources")
|
||||
inherited = (
|
||||
f"SELECT os.observation_id AS unit_id, e.id AS entity_id, e.canonical_name "
|
||||
f"FROM {os_t} os "
|
||||
f"JOIN {ue} src_ue ON src_ue.unit_id = os.source_id "
|
||||
f"JOIN {ents} e ON e.id = src_ue.entity_id "
|
||||
f"WHERE os.observation_id = ANY(${p}::uuid[]) "
|
||||
f"AND NOT EXISTS (SELECT 1 FROM {ue} d WHERE d.unit_id = os.observation_id)"
|
||||
)
|
||||
else:
|
||||
inherited = (
|
||||
f"SELECT obs.id AS unit_id, e.id AS entity_id, e.canonical_name "
|
||||
f"FROM {mu} obs "
|
||||
f"CROSS JOIN LATERAL unnest(obs.source_memory_ids) AS src_id "
|
||||
f"JOIN {ue} src_ue ON src_ue.unit_id = src_id "
|
||||
f"JOIN {ents} e ON e.id = src_ue.entity_id "
|
||||
f"WHERE obs.id = ANY(${p}::uuid[]) "
|
||||
f"AND obs.fact_type = 'observation' "
|
||||
f"AND obs.source_memory_ids IS NOT NULL "
|
||||
f"AND NOT EXISTS (SELECT 1 FROM {ue} d WHERE d.unit_id = obs.id)"
|
||||
)
|
||||
|
||||
return f"({direct}) UNION ({inherited})"
|
||||
|
||||
|
||||
async def list_memory_units(
|
||||
*,
|
||||
conn,
|
||||
ops,
|
||||
fq_table,
|
||||
bank_id: str,
|
||||
fact_type: str | None = None,
|
||||
search_query: str | None = None,
|
||||
consolidation_state: str | None = None,
|
||||
state: str | None = None,
|
||||
document_id: str | None = None,
|
||||
entity_id: str | None = None,
|
||||
tags: list[str] | None = None,
|
||||
tags_match: str = "any",
|
||||
created_before: datetime | None = None,
|
||||
limit: int = 100,
|
||||
offset: int = 0,
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
List memory units for table view with optional full-text search.
|
||||
|
||||
Args:
|
||||
conn: Open database connection (the caller owns the transaction).
|
||||
ops: Dialect ops. Unused by this query; part of the interface signature.
|
||||
fq_table: Table-name resolver.
|
||||
bank_id: Filter by bank ID
|
||||
fact_type: Filter by fact type (world, experience)
|
||||
search_query: Full-text search query (searches text and context fields)
|
||||
document_id: Optional filter to a single source document.
|
||||
tags: Optional list of tag names to filter by. When omitted, no tag
|
||||
filtering is applied (except tags_match='exact', which then selects
|
||||
the untagged/global scope).
|
||||
tags_match: How to combine tags (same modes as recall): 'any' (OR,
|
||||
default) or 'all' (AND) both also include untagged units;
|
||||
'any_strict'/'all_strict' exclude untagged units; 'exact' matches
|
||||
units whose tag set equals the given tags exactly.
|
||||
state: Optional curation-state filter ('valid' or 'invalidated').
|
||||
Invalidated facts live in a separate archive table; 'invalidated'
|
||||
reads that archive. Omitted/('valid') lists live facts.
|
||||
consolidation_state: Optional filter on consolidation state. One of
|
||||
'failed' (consolidation permanently failed and awaiting recovery),
|
||||
'pending' (not yet consolidated, no failure), or
|
||||
'done' (successfully consolidated). Only applies to source memory
|
||||
types (world/experience).
|
||||
limit: Maximum number of results to return
|
||||
offset: Offset for pagination
|
||||
|
||||
Returns:
|
||||
Dict with items (list of memory units) and total count
|
||||
"""
|
||||
if state is not None and state not in ("valid", "invalidated"):
|
||||
raise ValueError(f"Invalid state '{state}': expected 'valid' or 'invalidated'.")
|
||||
if entity_id is not None:
|
||||
import uuid as _uuid
|
||||
|
||||
try:
|
||||
_uuid.UUID(entity_id)
|
||||
except ValueError:
|
||||
raise ValueError(f"Invalid entity_id: '{entity_id}' is not a valid UUID") from None
|
||||
# Invalidated facts live in a separate archive table; pick the source
|
||||
# accordingly. Default (state is None) lists live facts.
|
||||
is_archived = state == "invalidated"
|
||||
source_table = fq_table("invalidated_memory_units") if is_archived else fq_table("memory_units")
|
||||
|
||||
# Build query conditions
|
||||
query_conditions = []
|
||||
query_params = []
|
||||
param_count = 0
|
||||
|
||||
if bank_id:
|
||||
param_count += 1
|
||||
query_conditions.append(f"bank_id = ${param_count}")
|
||||
query_params.append(bank_id)
|
||||
|
||||
if fact_type:
|
||||
param_count += 1
|
||||
query_conditions.append(f"fact_type = ${param_count}")
|
||||
query_params.append(fact_type)
|
||||
|
||||
if document_id:
|
||||
param_count += 1
|
||||
query_conditions.append(f"document_id = ${param_count}")
|
||||
query_params.append(document_id)
|
||||
|
||||
if entity_id:
|
||||
# Reverse lookup via the stored entity links. Entity links reference live memory units, so
|
||||
# this yields nothing against the invalidated archive (documented on the method).
|
||||
param_count += 1
|
||||
query_conditions.append(
|
||||
f"id IN (SELECT unit_id FROM {fq_table('unit_entities')} WHERE entity_id = ${param_count}::uuid)"
|
||||
)
|
||||
query_params.append(entity_id)
|
||||
|
||||
if search_query:
|
||||
# Full-text search on text and context fields using ILIKE
|
||||
param_count += 1
|
||||
query_conditions.append(f"(text ILIKE ${param_count} OR context ILIKE ${param_count})")
|
||||
query_params.append(f"%{search_query}%")
|
||||
|
||||
if consolidation_state:
|
||||
# Named apart from `state`, which the engine method used to shadow here;
|
||||
# `is_archived` was already resolved above, so behaviour is unchanged.
|
||||
wanted = consolidation_state.lower()
|
||||
if wanted == "failed":
|
||||
query_conditions.append("consolidation_failed_at IS NOT NULL AND fact_type IN ('experience', 'world')")
|
||||
elif wanted == "pending":
|
||||
query_conditions.append(
|
||||
"consolidated_at IS NULL AND consolidation_failed_at IS NULL AND fact_type IN ('experience', 'world')"
|
||||
)
|
||||
elif wanted == "done":
|
||||
query_conditions.append("consolidated_at IS NOT NULL AND fact_type IN ('experience', 'world')")
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Invalid consolidation_state '{consolidation_state}': expected 'failed', 'pending', or 'done'."
|
||||
)
|
||||
|
||||
if tags:
|
||||
tags_clause, tags_params, next_param = build_tags_where_clause(tags, param_count + 1, "", tags_match)
|
||||
if tags_clause:
|
||||
query_conditions.append(tags_clause.removeprefix("AND "))
|
||||
query_params.extend(tags_params)
|
||||
param_count = next_param - 1
|
||||
elif tags_match == "exact":
|
||||
# Exact match with no tags is the "global" scope: rows that carry no
|
||||
# tags at all. (Other match modes treat empty tags as "no filter".)
|
||||
query_conditions.append("(tags IS NULL OR tags = '{}')")
|
||||
|
||||
if created_before is not None:
|
||||
param_count += 1
|
||||
query_conditions.append(f"created_at < ${param_count}")
|
||||
query_params.append(created_before)
|
||||
|
||||
where_clause = "WHERE " + " AND ".join(query_conditions) if query_conditions else ""
|
||||
|
||||
# Get total count
|
||||
count_query = f"""
|
||||
SELECT COUNT(*) as total
|
||||
FROM {source_table}
|
||||
{where_clause}
|
||||
"""
|
||||
count_result = await conn.fetchrow(count_query, *query_params)
|
||||
total = count_result["total"]
|
||||
|
||||
# Get units with limit and offset
|
||||
param_count += 1
|
||||
limit_param = f"${param_count}"
|
||||
query_params.append(limit)
|
||||
|
||||
param_count += 1
|
||||
offset_param = f"${param_count}"
|
||||
query_params.append(offset)
|
||||
|
||||
# The archive carries invalidation bookkeeping; the live table doesn't.
|
||||
curation_cols = (
|
||||
"invalidation_reason, invalidated_at"
|
||||
if is_archived
|
||||
else "NULL::text AS invalidation_reason, NULL::timestamptz AS invalidated_at"
|
||||
)
|
||||
units = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, text, event_date, context, fact_type, document_id,
|
||||
mentioned_at, occurred_start, occurred_end, chunk_id, proof_count,
|
||||
tags, metadata, consolidated_at, consolidation_failed_at, edited_at, {curation_cols}
|
||||
FROM {source_table}
|
||||
{where_clause}
|
||||
ORDER BY mentioned_at DESC NULLS LAST, created_at DESC
|
||||
LIMIT {limit_param} OFFSET {offset_param}
|
||||
""",
|
||||
*query_params,
|
||||
)
|
||||
|
||||
# Get entity information for these units
|
||||
if units:
|
||||
unit_ids = [row["id"] for row in units]
|
||||
unit_entities = await conn.fetch(
|
||||
f"""
|
||||
SELECT ue.unit_id, e.canonical_name
|
||||
FROM {fq_table("unit_entities")} ue
|
||||
JOIN {fq_table("entities")} e ON ue.entity_id = e.id
|
||||
WHERE ue.unit_id = ANY($1::uuid[])
|
||||
ORDER BY ue.unit_id
|
||||
""",
|
||||
unit_ids,
|
||||
)
|
||||
else:
|
||||
unit_entities = []
|
||||
|
||||
# Build entity mapping
|
||||
entity_map: dict[Any, list[str]] = {}
|
||||
for row in unit_entities:
|
||||
unit_id = row["unit_id"]
|
||||
entity_name = row["canonical_name"]
|
||||
if unit_id not in entity_map:
|
||||
entity_map[unit_id] = []
|
||||
entity_map[unit_id].append(entity_name)
|
||||
|
||||
# Build result items
|
||||
items = []
|
||||
for row in units:
|
||||
unit_id = row["id"]
|
||||
entities = entity_map.get(unit_id, [])
|
||||
|
||||
items.append(
|
||||
{
|
||||
"id": str(unit_id),
|
||||
"text": row["text"],
|
||||
"context": row["context"] if row["context"] else "",
|
||||
"date": row["event_date"].isoformat() if row["event_date"] else "",
|
||||
"fact_type": row["fact_type"],
|
||||
"document_id": row["document_id"],
|
||||
"mentioned_at": row["mentioned_at"].isoformat() if row["mentioned_at"] else None,
|
||||
"occurred_start": row["occurred_start"].isoformat() if row["occurred_start"] else None,
|
||||
"occurred_end": row["occurred_end"].isoformat() if row["occurred_end"] else None,
|
||||
"entities": ", ".join(entities) if entities else "",
|
||||
"chunk_id": row["chunk_id"] if row["chunk_id"] else None,
|
||||
"proof_count": row["proof_count"] if row["proof_count"] is not None else 1,
|
||||
"tags": list(row["tags"]) if row["tags"] else [],
|
||||
"metadata": conn.parse_json(row["metadata"]) if row["metadata"] is not None else {},
|
||||
"consolidated_at": row["consolidated_at"].isoformat() if row["consolidated_at"] else None,
|
||||
"consolidation_failed_at": (
|
||||
row["consolidation_failed_at"].isoformat() if row["consolidation_failed_at"] else None
|
||||
),
|
||||
"state": "invalidated" if is_archived else "valid",
|
||||
"invalidation_reason": row["invalidation_reason"],
|
||||
"invalidated_at": row["invalidated_at"].isoformat() if row["invalidated_at"] else None,
|
||||
"edited_at": row["edited_at"].isoformat() if row["edited_at"] else None,
|
||||
}
|
||||
)
|
||||
|
||||
return {"items": items, "total": total, "limit": limit, "offset": offset}
|
||||
|
||||
|
||||
async def get_memory_unit(*, conn, ops, fq_table, bank_id: str, unit_id: str) -> dict[str, Any] | None:
|
||||
"""
|
||||
Get a single memory unit by ID.
|
||||
|
||||
Args:
|
||||
conn: Open database connection (the caller owns the transaction).
|
||||
ops: Dialect ops, for the observation→source entity inheritance shape.
|
||||
fq_table: Table-name resolver.
|
||||
bank_id: Bank ID
|
||||
unit_id: Memory unit ID (the caller validates it is a UUID)
|
||||
|
||||
Returns:
|
||||
Dict with memory unit data or None if not found
|
||||
"""
|
||||
# Get the memory unit (include source_memory_ids for mental models).
|
||||
# Curation moves invalidated facts to invalidated_memory_units, so fall
|
||||
# back to the archive (with its invalidation bookkeeping) on a miss.
|
||||
select_cols = (
|
||||
"id, text, context, event_date, occurred_start, occurred_end, "
|
||||
"mentioned_at, fact_type, document_id, chunk_id, tags, metadata, source_memory_ids, "
|
||||
"observation_scopes, edited_at"
|
||||
)
|
||||
row = await conn.fetchrow(
|
||||
f"SELECT {select_cols}, NULL::text AS invalidation_reason, NULL::timestamptz AS invalidated_at "
|
||||
f"FROM {fq_table('memory_units')} WHERE id = $1 AND bank_id = $2",
|
||||
unit_id,
|
||||
bank_id,
|
||||
)
|
||||
unit_state = "valid"
|
||||
if not row:
|
||||
row = await conn.fetchrow(
|
||||
f"SELECT {select_cols}, invalidation_reason, invalidated_at "
|
||||
f"FROM {fq_table('invalidated_memory_units')} WHERE id = $1 AND bank_id = $2",
|
||||
unit_id,
|
||||
bank_id,
|
||||
)
|
||||
unit_state = "invalidated"
|
||||
|
||||
if not row:
|
||||
return None
|
||||
|
||||
# Get entity information. _entity_rows_for_units_sql handles the
|
||||
# observation→source_memory_ids inheritance fallback in SQL, so a
|
||||
# single query covers direct rows and inherited ones.
|
||||
entities_rows = await conn.fetch(
|
||||
_entity_rows_for_units_sql(ops=ops, fq_table=fq_table, unit_ids_placeholder=1),
|
||||
[row["id"]],
|
||||
)
|
||||
entities = [r["canonical_name"] for r in entities_rows]
|
||||
|
||||
result: dict[str, Any] = {
|
||||
"id": str(row["id"]),
|
||||
"text": row["text"],
|
||||
"context": row["context"] if row["context"] else "",
|
||||
"date": row["event_date"].isoformat() if row["event_date"] else "",
|
||||
"type": row["fact_type"],
|
||||
"mentioned_at": row["mentioned_at"].isoformat() if row["mentioned_at"] else None,
|
||||
"occurred_start": row["occurred_start"].isoformat() if row["occurred_start"] else None,
|
||||
"occurred_end": row["occurred_end"].isoformat() if row["occurred_end"] else None,
|
||||
"entities": entities,
|
||||
"document_id": row["document_id"] if row["document_id"] else None,
|
||||
"chunk_id": str(row["chunk_id"]) if row["chunk_id"] else None,
|
||||
"tags": row["tags"] if row["tags"] else [],
|
||||
"metadata": conn.parse_json(row["metadata"]) if row["metadata"] is not None else {},
|
||||
"observation_scopes": (
|
||||
conn.parse_json(row["observation_scopes"]) if row["observation_scopes"] is not None else None
|
||||
),
|
||||
"state": unit_state,
|
||||
"invalidation_reason": row["invalidation_reason"],
|
||||
"invalidated_at": row["invalidated_at"].isoformat() if row["invalidated_at"] else None,
|
||||
"edited_at": row["edited_at"].isoformat() if row["edited_at"] else None,
|
||||
}
|
||||
|
||||
# For observations, include source_memory_ids
|
||||
# history is deprecated here - use GET /memories/{id}/history instead
|
||||
if row["fact_type"] == "observation":
|
||||
result["history"] = []
|
||||
|
||||
if row["fact_type"] == "observation" and row["source_memory_ids"]:
|
||||
source_ids = row["source_memory_ids"]
|
||||
result["source_memory_ids"] = [str(sid) for sid in source_ids]
|
||||
|
||||
# Fetch source memories
|
||||
source_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, text, fact_type, context, occurred_start, mentioned_at
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE id = ANY($1::uuid[])
|
||||
ORDER BY mentioned_at DESC NULLS LAST
|
||||
""",
|
||||
source_ids,
|
||||
)
|
||||
result["source_memories"] = [
|
||||
{
|
||||
"id": str(r["id"]),
|
||||
"text": r["text"],
|
||||
"type": r["fact_type"],
|
||||
"context": r["context"],
|
||||
"occurred_start": r["occurred_start"].isoformat() if r["occurred_start"] else None,
|
||||
"mentioned_at": r["mentioned_at"].isoformat() if r["mentioned_at"] else None,
|
||||
}
|
||||
for r in source_rows
|
||||
]
|
||||
|
||||
return result
|
||||
|
||||
|
||||
async def list_entities(
|
||||
*,
|
||||
conn,
|
||||
fq_table,
|
||||
bank_id: str,
|
||||
search: str | None = None,
|
||||
limit: int = 100,
|
||||
offset: int = 0,
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
List all entities for a bank with pagination.
|
||||
|
||||
Args:
|
||||
conn: Open database connection (the caller owns the transaction).
|
||||
fq_table: Table-name resolver.
|
||||
bank_id: bank IDentifier
|
||||
search: Optional case-insensitive substring match on canonical_name.
|
||||
limit: Maximum number of entities to return
|
||||
offset: Offset for pagination
|
||||
|
||||
Returns:
|
||||
Dict with items, total, limit, offset
|
||||
"""
|
||||
conditions = ["bank_id = $1"]
|
||||
params: list[Any] = [bank_id]
|
||||
if search:
|
||||
# Substring match, same ILIKE shape entity lookup uses elsewhere. Applied
|
||||
# to the count too, so the UI pages over the filtered set.
|
||||
params.append(f"%{search}%")
|
||||
conditions.append(f"canonical_name ILIKE ${len(params)}")
|
||||
where_clause = " AND ".join(conditions)
|
||||
|
||||
# Get total count
|
||||
total_row = await conn.fetchrow(
|
||||
f"""
|
||||
SELECT COUNT(*) as total
|
||||
FROM {fq_table("entities")}
|
||||
WHERE {where_clause}
|
||||
""",
|
||||
*params,
|
||||
)
|
||||
total = total_row["total"] if total_row else 0
|
||||
|
||||
# Get paginated entities
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, canonical_name, mention_count, first_seen, last_seen, metadata
|
||||
FROM {fq_table("entities")}
|
||||
WHERE {where_clause}
|
||||
ORDER BY mention_count DESC, last_seen DESC, id ASC
|
||||
LIMIT ${len(params) + 1} OFFSET ${len(params) + 2}
|
||||
""",
|
||||
*params,
|
||||
limit,
|
||||
offset,
|
||||
)
|
||||
|
||||
entities = []
|
||||
for row in rows:
|
||||
# Handle metadata - may be dict, JSON string, or None
|
||||
metadata = row["metadata"]
|
||||
if metadata is None:
|
||||
metadata = {}
|
||||
elif isinstance(metadata, str):
|
||||
try:
|
||||
metadata = json.loads(metadata)
|
||||
except json.JSONDecodeError:
|
||||
metadata = {}
|
||||
|
||||
entities.append(
|
||||
{
|
||||
"id": str(row["id"]),
|
||||
"canonical_name": row["canonical_name"],
|
||||
"mention_count": row["mention_count"],
|
||||
"first_seen": row["first_seen"].isoformat() if row["first_seen"] else None,
|
||||
"last_seen": row["last_seen"].isoformat() if row["last_seen"] else None,
|
||||
"metadata": metadata,
|
||||
}
|
||||
)
|
||||
return {
|
||||
"items": entities,
|
||||
"total": total,
|
||||
"limit": limit,
|
||||
"offset": offset,
|
||||
}
|
||||
|
||||
|
||||
__all__ = ["get_memory_unit", "list_entities", "list_memory_units"]
|
||||
@@ -0,0 +1,784 @@
|
||||
"""Graph-shaped reads and the link-maintenance passes, in SQL.
|
||||
|
||||
Everything here is a query over the *joins* around `memory_units` rather than
|
||||
over the memories themselves: `unit_entities` (which entities a memory mentions)
|
||||
and `memory_links` (memory-to-memory temporal/semantic/causal edges).
|
||||
|
||||
Two groups of callers:
|
||||
|
||||
* **The graph view.** :func:`graph_units`, :func:`graph_entity_rows` and
|
||||
:func:`graph_direct_links` return raw rows; the engine still owns the
|
||||
filtering, the observation inheritance, the derived entity edges, the
|
||||
colouring and the response assembly. These functions answer only "which
|
||||
memories", "which entity postings" and "which stored edges".
|
||||
* **The graph-maintenance job.** :func:`enqueue_relink_victims` runs inside the
|
||||
delete transaction; :func:`relink_pass`, :func:`prune_orphan_entities` and
|
||||
:func:`prune_stale_cooccurrences` are the three reconciliation passes the job
|
||||
drives. The job keeps the orchestration (pass ordering, the deadlock retry
|
||||
around the sweeps, the timing log); each function here does the pass's work.
|
||||
|
||||
:func:`entity_memory_counts` and :func:`entities_for_units` are the two entity
|
||||
postings reads that are not part of the graph view but read the same join table.
|
||||
|
||||
A store whose links travel inside the memory has nothing to relink and no join
|
||||
table to sweep, which is why these are methods on the interface at all: it
|
||||
answers them with zeroes rather than with SQL.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import uuid as uuid_module
|
||||
from collections.abc import Callable
|
||||
from typing import Any
|
||||
|
||||
from ....config import get_config
|
||||
from ...db.base import DatabaseConnection
|
||||
from ...retain.link_utils import (
|
||||
MAX_TEMPORAL_LINKS_PER_UNIT,
|
||||
_bulk_insert_links,
|
||||
_normalize_datetime,
|
||||
compute_semantic_links_ann,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Mirrors the ``top_k`` default in ``compute_semantic_links_ann`` at retain
|
||||
# time. If you change one, change the other — otherwise victims would either
|
||||
# never reach the cap (probe returns less than the cap) or stay perpetually
|
||||
# under it (cap is higher than retain creates).
|
||||
MAX_SEMANTIC_LINKS_PER_UNIT = 50
|
||||
|
||||
# Worker fetches this many rows per relink-loop iteration. Bounds
|
||||
# per-iteration probe/insert latency so a 10k-row backlog doesn't hold a
|
||||
# worker slot for minutes. Chosen so the typical iteration runs in well
|
||||
# under 1s.
|
||||
_DRAIN_BATCH_SIZE = 50
|
||||
|
||||
# Defensive guard against runaway relink loops — at _DRAIN_BATCH_SIZE units per
|
||||
# iteration that's 500k targets, far beyond any realistic single-bank backlog.
|
||||
_RELINK_ITERATION_CAP = 10000
|
||||
|
||||
# Cap at 10k edges — the UI can't usefully render more, and uncapped queries
|
||||
# on highly-connected graphs (e.g. 1000 nodes with 500k+ edges) are too slow.
|
||||
_GRAPH_MAX_EDGES = 10000
|
||||
|
||||
# Columns the graph view renders: nodes take id/text/date/context/entities,
|
||||
# the table rows take the rest, and `source_memory_ids` is what lets the caller
|
||||
# inherit an observation's links and entities from the facts behind it.
|
||||
_GRAPH_UNIT_COLUMNS = (
|
||||
"id, text, event_date, context, occurred_start, occurred_end, mentioned_at, "
|
||||
"document_id, chunk_id, fact_type, tags, created_at, proof_count, source_memory_ids"
|
||||
)
|
||||
|
||||
|
||||
def _ops_for(conn: DatabaseConnection) -> Any:
|
||||
"""The ``DataAccessOps`` matching the connection's SQL dialect.
|
||||
|
||||
This is the SQL memories store, and SQL means Postgres *or* Oracle — the two
|
||||
speak different dialects (Oracle inherits entity links through the
|
||||
``observation_sources`` junction, Postgres through ``source_memory_ids``
|
||||
arrays), so the ops must follow the connection rather than assume Postgres.
|
||||
The ops go by ``conn.backend_type`` — the connection objects carry the dialect
|
||||
but not the backend's ``ops`` handle, so resolve through the per-dialect cache
|
||||
of ``create_data_access_ops`` (a dict lookup after the first call, and the same
|
||||
instance the backend holds). The default covers callers that hand in a bare
|
||||
asyncpg connection with no dialect to report.
|
||||
"""
|
||||
from ...db import create_data_access_ops
|
||||
|
||||
return create_data_access_ops(getattr(conn, "backend_type", "postgresql"))
|
||||
|
||||
|
||||
def _as_uuids(unit_ids: list) -> list:
|
||||
"""Coerce a mixed list of uuid strings / UUIDs to UUIDs for a ``uuid[]`` bind."""
|
||||
return [uuid_module.UUID(uid) if isinstance(uid, str) else uid for uid in unit_ids]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------- graph view
|
||||
|
||||
|
||||
def _observations_via_source_match(
|
||||
fq_table: Callable[[str], str],
|
||||
ops: Any,
|
||||
source_column: str,
|
||||
source_placeholder: int,
|
||||
bank_placeholder: int | None,
|
||||
) -> str:
|
||||
"""A predicate matching observations whose *sources* satisfy ``<col> = $n``.
|
||||
|
||||
Observations carry no `document_id` / `chunk_id` of their own; the link to a
|
||||
source row lives in `source_memory_ids` (native array) or the
|
||||
`observation_sources` junction, depending on the dialect.
|
||||
"""
|
||||
if ops.uses_observation_sources_table:
|
||||
bank_clause = f" AND src.bank_id = ${bank_placeholder}" if bank_placeholder else ""
|
||||
return (
|
||||
f"id IN (SELECT os.observation_id "
|
||||
f"FROM {fq_table('observation_sources')} os "
|
||||
f"JOIN {fq_table('memory_units')} src ON src.id = os.source_id "
|
||||
f"WHERE src.{source_column} = ${source_placeholder}{bank_clause})"
|
||||
)
|
||||
bank_clause = f" AND bank_id = ${bank_placeholder}" if bank_placeholder else ""
|
||||
return (
|
||||
f"source_memory_ids && (SELECT array_agg(id) "
|
||||
f"FROM {fq_table('memory_units')} "
|
||||
f"WHERE {source_column} = ${source_placeholder}{bank_clause})"
|
||||
)
|
||||
|
||||
|
||||
async def graph_units(
|
||||
*,
|
||||
conn: DatabaseConnection,
|
||||
fq_table: Callable[[str], str],
|
||||
bank_id: str | None = None,
|
||||
fact_type: str | None = None,
|
||||
search_query: str | None = None,
|
||||
document_id: str | None = None,
|
||||
chunk_id: str | None = None,
|
||||
tags: list[str] | None = None,
|
||||
tags_match: str = "all_strict",
|
||||
limit: int = 1000,
|
||||
) -> dict[str, Any]:
|
||||
"""Memory nodes for the graph view, plus the total matching count.
|
||||
|
||||
Returns ``{"units": [...], "total": int}``: ``units`` is the page (newest
|
||||
first, capped at ``limit``); ``total`` is how many match the filters, which
|
||||
the UI shows alongside the page. ``document_id`` / ``chunk_id`` also match an
|
||||
observation whose *sources* carry them, since observations have neither of
|
||||
their own.
|
||||
"""
|
||||
from ...search.tags import build_tags_where_clause_simple
|
||||
|
||||
ops = _ops_for(conn)
|
||||
conditions: list[str] = []
|
||||
params: list[Any] = []
|
||||
|
||||
bank_placeholder: int | None = None
|
||||
if bank_id:
|
||||
params.append(bank_id)
|
||||
bank_placeholder = len(params)
|
||||
conditions.append(f"bank_id = ${bank_placeholder}")
|
||||
|
||||
if fact_type:
|
||||
params.append(fact_type)
|
||||
conditions.append(f"fact_type = ${len(params)}")
|
||||
|
||||
if document_id:
|
||||
params.append(document_id)
|
||||
obs = _observations_via_source_match(fq_table, ops, "document_id", len(params), bank_placeholder)
|
||||
conditions.append(f"(document_id = ${len(params)} OR (fact_type = 'observation' AND {obs}))")
|
||||
|
||||
if chunk_id:
|
||||
params.append(chunk_id)
|
||||
obs = _observations_via_source_match(fq_table, ops, "chunk_id", len(params), bank_placeholder)
|
||||
conditions.append(f"(chunk_id = ${len(params)} OR (fact_type = 'observation' AND {obs}))")
|
||||
|
||||
if search_query:
|
||||
params.append(f"%{search_query}%")
|
||||
conditions.append(f"(text ILIKE ${len(params)} OR context ILIKE ${len(params)})")
|
||||
|
||||
if tags:
|
||||
tag_clause = build_tags_where_clause_simple(tags, len(params) + 1, match=tags_match)
|
||||
if tag_clause:
|
||||
conditions.append(tag_clause.removeprefix("AND "))
|
||||
params.append(tags)
|
||||
elif tags_match == "exact":
|
||||
# Exact match with no tags is the "global" scope: rows carrying no tags at
|
||||
# all. (Other modes treat empty tags as "no filter".)
|
||||
conditions.append("(tags IS NULL OR tags = '{}')")
|
||||
|
||||
where_clause = "WHERE " + " AND ".join(conditions) if conditions else ""
|
||||
|
||||
total_row = await conn.fetchrow(
|
||||
f"SELECT COUNT(*) AS total FROM {fq_table('memory_units')} {where_clause}",
|
||||
*params,
|
||||
)
|
||||
total = total_row["total"] if total_row else 0
|
||||
|
||||
params.append(limit)
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT {_GRAPH_UNIT_COLUMNS}
|
||||
FROM {fq_table("memory_units")}
|
||||
{where_clause}
|
||||
ORDER BY mentioned_at DESC NULLS LAST, event_date DESC
|
||||
LIMIT ${len(params)}
|
||||
""",
|
||||
*params,
|
||||
)
|
||||
return {"units": [dict(row) for row in rows], "total": total}
|
||||
|
||||
|
||||
async def graph_entity_rows(
|
||||
*,
|
||||
conn: DatabaseConnection,
|
||||
fq_table: Callable[[str], str],
|
||||
bank_id: str,
|
||||
unit_ids: list[str],
|
||||
) -> list[dict[str, Any]]:
|
||||
"""``(unit_id, entity_id, canonical_name)`` rows for the graph view's entity edges.
|
||||
|
||||
Direct `unit_entities` postings only. An observation's entities are inherited
|
||||
from its source memories by the caller, which is why the ids it passes here
|
||||
are the visible units *plus* their source memories.
|
||||
|
||||
Scoped by unit id rather than by bank: the ids already came from a
|
||||
bank-scoped :func:`graph_units`, and `unit_entities` carries no bank column.
|
||||
"""
|
||||
if not unit_ids:
|
||||
return []
|
||||
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT ue.unit_id, e.id AS entity_id, e.canonical_name
|
||||
FROM {fq_table("unit_entities")} ue
|
||||
JOIN {fq_table("entities")} e ON ue.entity_id = e.id
|
||||
WHERE ue.unit_id = ANY($1::uuid[])
|
||||
ORDER BY ue.unit_id
|
||||
""",
|
||||
_as_uuids(unit_ids),
|
||||
)
|
||||
return [dict(row) for row in rows]
|
||||
|
||||
|
||||
async def graph_direct_links(
|
||||
*,
|
||||
conn: DatabaseConnection,
|
||||
fq_table: Callable[[str], str],
|
||||
bank_id: str,
|
||||
unit_ids: list[str],
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Memory-to-memory edges with *both* endpoints in ``unit_ids``.
|
||||
|
||||
Entity edges are derived by the caller from `unit_entities` so we don't
|
||||
materialize them in `memory_links` anymore (dropped in migration
|
||||
e9b2c7d1f3a4) — no link_type filter is needed. ``entity_name`` is selected as
|
||||
NULL so the row shape matches the derived edges the caller mixes these with.
|
||||
|
||||
Pass the visible units *and* the source memories they inherit from: the
|
||||
caller copies a source memory's links onto the observations built on it.
|
||||
"""
|
||||
if not unit_ids:
|
||||
return []
|
||||
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT ml.from_unit_id,
|
||||
ml.to_unit_id,
|
||||
ml.link_type,
|
||||
ml.weight,
|
||||
NULL::text AS entity_name
|
||||
FROM {fq_table("memory_links")} ml
|
||||
WHERE ml.from_unit_id = ANY($1::uuid[])
|
||||
AND ml.to_unit_id = ANY($1::uuid[])
|
||||
ORDER BY ml.weight DESC NULLS LAST
|
||||
LIMIT $2
|
||||
""",
|
||||
_as_uuids(unit_ids),
|
||||
_GRAPH_MAX_EDGES,
|
||||
)
|
||||
return [dict(row) for row in rows]
|
||||
|
||||
|
||||
# ------------------------------------------------------------ entity postings
|
||||
|
||||
|
||||
async def entity_memory_counts(
|
||||
*,
|
||||
conn: DatabaseConnection,
|
||||
fq_table: Callable[[str], str],
|
||||
bank_id: str,
|
||||
entity_ids: list[str] | None = None,
|
||||
) -> dict[str, int]:
|
||||
"""Live memory count per entity id, for the entities in ``bank_id``.
|
||||
|
||||
The GROUP BY is what makes this an orphan test: an entity with no surviving
|
||||
`unit_entities` row produces no group, so it is simply absent from the
|
||||
result rather than present with a zero.
|
||||
|
||||
Scoped through ``memory_units.bank_id`` — `unit_entities` has no bank column,
|
||||
and joining is what keeps the count to *live* memories (deleted units take
|
||||
their postings with them via ON DELETE CASCADE).
|
||||
"""
|
||||
params: list[Any] = [bank_id]
|
||||
entity_filter = ""
|
||||
if entity_ids is not None:
|
||||
if not entity_ids:
|
||||
return {}
|
||||
params.append(_as_uuids(entity_ids))
|
||||
entity_filter = f"AND ue.entity_id = ANY(${len(params)}::uuid[])"
|
||||
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT ue.entity_id, COUNT(*) AS n
|
||||
FROM {fq_table("unit_entities")} ue
|
||||
JOIN {fq_table("memory_units")} mu ON mu.id = ue.unit_id
|
||||
WHERE mu.bank_id = $1
|
||||
{entity_filter}
|
||||
GROUP BY ue.entity_id
|
||||
""",
|
||||
*params,
|
||||
)
|
||||
return {str(row["entity_id"]): int(row["n"]) for row in rows}
|
||||
|
||||
|
||||
def _entity_rows_for_units_sql(
|
||||
fq_table: Callable[[str], str],
|
||||
ops: Any,
|
||||
unit_ids_placeholder: int,
|
||||
) -> str:
|
||||
"""SQL SELECT producing ``(unit_id, entity_id, canonical_name)`` rows for
|
||||
the given unit IDs.
|
||||
|
||||
Direct rows come from ``unit_entities``. Observations rarely carry
|
||||
direct rows there; their entity association lives transitively through
|
||||
their source memories (``source_memory_ids`` on PG, the
|
||||
``observation_sources`` junction on Oracle). When an observation has
|
||||
no direct entity rows the SELECT inherits its source memories'
|
||||
entities, so the result is the same set callers would get from
|
||||
``get_memory_unit``.
|
||||
|
||||
``unit_ids_placeholder`` is the 1-based parameter index that holds the
|
||||
``uuid[]`` of unit IDs. The placeholder is referenced twice — both
|
||||
sides of the UNION need it — so callers should not reuse the slot.
|
||||
"""
|
||||
ue = fq_table("unit_entities")
|
||||
ents = fq_table("entities")
|
||||
mu = fq_table("memory_units")
|
||||
p = unit_ids_placeholder
|
||||
|
||||
direct = (
|
||||
f"SELECT ue.unit_id, e.id AS entity_id, e.canonical_name "
|
||||
f"FROM {ue} ue "
|
||||
f"JOIN {ents} e ON e.id = ue.entity_id "
|
||||
f"WHERE ue.unit_id = ANY(${p}::uuid[])"
|
||||
)
|
||||
|
||||
if ops.uses_observation_sources_table:
|
||||
os_t = fq_table("observation_sources")
|
||||
inherited = (
|
||||
f"SELECT os.observation_id AS unit_id, e.id AS entity_id, e.canonical_name "
|
||||
f"FROM {os_t} os "
|
||||
f"JOIN {ue} src_ue ON src_ue.unit_id = os.source_id "
|
||||
f"JOIN {ents} e ON e.id = src_ue.entity_id "
|
||||
f"WHERE os.observation_id = ANY(${p}::uuid[]) "
|
||||
f"AND NOT EXISTS (SELECT 1 FROM {ue} d WHERE d.unit_id = os.observation_id)"
|
||||
)
|
||||
else:
|
||||
inherited = (
|
||||
f"SELECT obs.id AS unit_id, e.id AS entity_id, e.canonical_name "
|
||||
f"FROM {mu} obs "
|
||||
f"CROSS JOIN LATERAL unnest(obs.source_memory_ids) AS src_id "
|
||||
f"JOIN {ue} src_ue ON src_ue.unit_id = src_id "
|
||||
f"JOIN {ents} e ON e.id = src_ue.entity_id "
|
||||
f"WHERE obs.id = ANY(${p}::uuid[]) "
|
||||
f"AND obs.fact_type = 'observation' "
|
||||
f"AND obs.source_memory_ids IS NOT NULL "
|
||||
f"AND NOT EXISTS (SELECT 1 FROM {ue} d WHERE d.unit_id = obs.id)"
|
||||
)
|
||||
|
||||
return f"({direct}) UNION ({inherited})"
|
||||
|
||||
|
||||
async def entities_for_units(
|
||||
*,
|
||||
conn: DatabaseConnection,
|
||||
fq_table: Callable[[str], str],
|
||||
bank_id: str,
|
||||
unit_ids: list[str],
|
||||
) -> dict[str, list[str]]:
|
||||
"""The entity ids each unit carries, keyed by unit id.
|
||||
|
||||
Observations inherit their source memories' entities when they carry no
|
||||
direct postings of their own — see :func:`_entity_rows_for_units_sql`. Units
|
||||
with no entities are absent rather than mapped to an empty list.
|
||||
"""
|
||||
if not unit_ids:
|
||||
return {}
|
||||
|
||||
rows = await conn.fetch(
|
||||
_entity_rows_for_units_sql(fq_table, _ops_for(conn), unit_ids_placeholder=1),
|
||||
_as_uuids(unit_ids),
|
||||
)
|
||||
|
||||
# UNION already de-duplicates whole rows, but a unit can reach the same
|
||||
# entity through more than one source memory, so dedupe per unit while
|
||||
# preserving the order the rows arrived in.
|
||||
by_unit: dict[str, list[str]] = {}
|
||||
for row in rows:
|
||||
unit_key = str(row["unit_id"])
|
||||
entity_id = str(row["entity_id"])
|
||||
bucket = by_unit.setdefault(unit_key, [])
|
||||
if entity_id not in bucket:
|
||||
bucket.append(entity_id)
|
||||
return by_unit
|
||||
|
||||
|
||||
async def entity_map_for_units(
|
||||
*,
|
||||
conn: DatabaseConnection,
|
||||
fq_table: Callable[[str], str],
|
||||
bank_id: str,
|
||||
unit_ids: list[str],
|
||||
) -> dict[str, list[dict[str, str]]]:
|
||||
"""``{unit_id: [{entity_id, canonical_name}]}`` — the recall/curation shape.
|
||||
|
||||
The named twin of :func:`entities_for_units`: recall renders the entity name
|
||||
on each fact, so it needs the label, not just the id. Observation-via-source
|
||||
inheritance and the per-unit dedupe are identical.
|
||||
"""
|
||||
if not unit_ids:
|
||||
return {}
|
||||
|
||||
rows = await conn.fetch(
|
||||
_entity_rows_for_units_sql(fq_table, _ops_for(conn), unit_ids_placeholder=1),
|
||||
_as_uuids(unit_ids),
|
||||
)
|
||||
by_unit: dict[str, list[dict[str, str]]] = {}
|
||||
for row in rows:
|
||||
unit_key = str(row["unit_id"])
|
||||
entity_id = str(row["entity_id"])
|
||||
bucket = by_unit.setdefault(unit_key, [])
|
||||
if not any(existing["entity_id"] == entity_id for existing in bucket):
|
||||
bucket.append({"entity_id": entity_id, "canonical_name": row["canonical_name"]})
|
||||
return by_unit
|
||||
|
||||
|
||||
# --------------------------------------------------------------- maintenance
|
||||
|
||||
|
||||
async def enqueue_relink_victims(
|
||||
*,
|
||||
conn: DatabaseConnection,
|
||||
fq_table: Callable[[str], str],
|
||||
bank_id: str,
|
||||
affected_unit_ids: list,
|
||||
include_affected_units: bool = False,
|
||||
) -> int:
|
||||
"""Enqueue surviving units whose outgoing temporal/semantic links pointed at
|
||||
``affected_unit_ids`` for later link top-up.
|
||||
|
||||
Must run inside the same transaction that drops those links, *before* the
|
||||
delete (or cascade) fires — once the rows are gone, the join that finds the
|
||||
victims returns nothing.
|
||||
|
||||
Args:
|
||||
conn: Database connection inside the active transaction.
|
||||
fq_table: Schema-qualifying table-name resolver.
|
||||
bank_id: Bank owning the affected units.
|
||||
affected_unit_ids: Memory_unit IDs whose incident temporal/semantic links
|
||||
are about to be (or are being) removed.
|
||||
include_affected_units: Also enqueue ``affected_unit_ids`` themselves — for
|
||||
an edit that deletes a unit's links but leaves the unit live, so its own
|
||||
outgoing adjacency is rebuilt too. One combined insert keeps the queue's
|
||||
sorted lock ordering intact.
|
||||
|
||||
Returns:
|
||||
Number of distinct victim units enqueued (after dedup against rows
|
||||
already in the queue).
|
||||
"""
|
||||
if not affected_unit_ids:
|
||||
return 0
|
||||
|
||||
ops = _ops_for(conn)
|
||||
affected_uuids = _as_uuids(affected_unit_ids)
|
||||
affected_str_set = {str(uid) for uid in affected_uuids}
|
||||
|
||||
# Find units (other than the affected ones) that have an outgoing
|
||||
# temporal/semantic link pointing at an affected unit. Entity links are
|
||||
# intentionally excluded — they're scheduled for removal and would only
|
||||
# add noise to the recompute job.
|
||||
victim_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT DISTINCT from_unit_id
|
||||
FROM {fq_table("memory_links")}
|
||||
WHERE to_unit_id = ANY($1::uuid[])
|
||||
AND bank_id = $2
|
||||
AND link_type IN ('temporal', 'semantic')
|
||||
""",
|
||||
affected_uuids,
|
||||
bank_id,
|
||||
)
|
||||
|
||||
victim_ids = {row["from_unit_id"] for row in victim_rows if str(row["from_unit_id"]) not in affected_str_set}
|
||||
if include_affected_units:
|
||||
victim_ids.update(affected_uuids)
|
||||
|
||||
if not victim_ids:
|
||||
return 0
|
||||
|
||||
await ops.enqueue_graph_maintenance(
|
||||
conn,
|
||||
fq_table("graph_maintenance_queue"),
|
||||
bank_id,
|
||||
list(victim_ids),
|
||||
)
|
||||
|
||||
logger.debug(
|
||||
f"[GRAPH_MAINT] Enqueued {len(victim_ids)} relink victims in "
|
||||
f"bank={bank_id} ({len(affected_unit_ids)} units affected)"
|
||||
)
|
||||
return len(victim_ids)
|
||||
|
||||
|
||||
async def relink_pass(
|
||||
*,
|
||||
backend: Any,
|
||||
fq_table: Callable[[str], str],
|
||||
bank_id: str,
|
||||
config: Any,
|
||||
) -> dict:
|
||||
"""Drain ``graph_maintenance_queue`` for ``bank_id``, topping up lost links.
|
||||
|
||||
Per-iteration loop: claim → top up → commit. We rely on at most one job per
|
||||
bank running, so no need for SKIP LOCKED. Submit-time dedup alone does NOT
|
||||
give that — it only inspects 'pending' rows — so the guarantee comes from
|
||||
``claim_tasks``, which refuses to claim a graph_maintenance row for a bank
|
||||
that already has one in flight (``graph_maintenance_bank_serialization_sql``,
|
||||
#3230). Without it these claims convoy: they lock queue rows ``FOR UPDATE``
|
||||
with no ``SKIP LOCKED``, so a second run blocks on the first while holding a
|
||||
worker slot.
|
||||
|
||||
Takes ``backend`` rather than a connection because the loop spans several
|
||||
transactions — one per claimed batch, plus a separate connection for the ANN
|
||||
probe — so it has to acquire its own.
|
||||
|
||||
``config`` is the caller's resolved configuration. The Postgres pass takes
|
||||
its caps from retain's link_utils (so relink and retain agree on what "full"
|
||||
means) and never reads it; it is accepted so a store that *does* tune its
|
||||
relinking gets it.
|
||||
|
||||
Returns:
|
||||
``{"relink_units_processed": int, "relink_links_added": int}``.
|
||||
"""
|
||||
del config # accepted for symmetry with stores that tune their own relinking
|
||||
ops = backend.ops
|
||||
|
||||
units_processed = 0
|
||||
links_added = 0
|
||||
iterations = 0
|
||||
while True:
|
||||
from ...memory_engine import acquire_with_retry
|
||||
|
||||
async with acquire_with_retry(backend) as conn:
|
||||
async with conn.transaction():
|
||||
unit_ids = await ops.claim_graph_maintenance_batch(
|
||||
conn,
|
||||
fq_table("graph_maintenance_queue"),
|
||||
bank_id,
|
||||
_DRAIN_BATCH_SIZE,
|
||||
)
|
||||
if not unit_ids:
|
||||
break
|
||||
|
||||
links_added += await _relink_batch(conn, fq_table, bank_id, unit_ids, ops, backend)
|
||||
|
||||
units_processed += len(unit_ids)
|
||||
iterations += 1
|
||||
|
||||
if iterations > _RELINK_ITERATION_CAP:
|
||||
# Defensive guard against runaway loops — at 50 units/iter that's
|
||||
# 500k targets, far beyond any realistic single-bank backlog.
|
||||
logger.error(
|
||||
f"[GRAPH_MAINT] bank={bank_id} hit iteration cap ({iterations}); aborting relink "
|
||||
f"(units_processed={units_processed}, links_added={links_added})"
|
||||
)
|
||||
break
|
||||
|
||||
return {"relink_units_processed": units_processed, "relink_links_added": links_added}
|
||||
|
||||
|
||||
async def _relink_batch(
|
||||
conn: DatabaseConnection,
|
||||
fq_table: Callable[[str], str],
|
||||
bank_id: str,
|
||||
victim_ids: list[str],
|
||||
ops: Any,
|
||||
backend: Any,
|
||||
) -> int:
|
||||
"""Top up temporal/semantic links for a batch of victim units. Returns rows inserted."""
|
||||
# Load each victim's metadata. Victims whose units were deleted between
|
||||
# enqueue and now silently drop out — exactly the no-op behaviour we want
|
||||
# for stale queue rows.
|
||||
victim_uuids = [uuid_module.UUID(vid) for vid in victim_ids]
|
||||
victim_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT id::text AS id, event_date, fact_type, embedding::text AS embedding
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE id = ANY($1::uuid[])
|
||||
AND bank_id = $2
|
||||
AND fact_type IN ('experience', 'world')
|
||||
""",
|
||||
victim_uuids,
|
||||
bank_id,
|
||||
)
|
||||
|
||||
if not victim_rows:
|
||||
return 0
|
||||
|
||||
alive_uuids = [uuid_module.UUID(row["id"]) for row in victim_rows]
|
||||
|
||||
# Count current outgoing temporal/semantic links per victim so we only
|
||||
# probe for the ones genuinely below cap. Saves the bulk of the work when
|
||||
# most victims still have plenty of links.
|
||||
count_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT from_unit_id, link_type, COUNT(*) AS cnt
|
||||
FROM {fq_table("memory_links")}
|
||||
WHERE from_unit_id = ANY($1::uuid[])
|
||||
AND bank_id = $2
|
||||
AND link_type IN ('temporal', 'semantic')
|
||||
GROUP BY from_unit_id, link_type
|
||||
""",
|
||||
alive_uuids,
|
||||
bank_id,
|
||||
)
|
||||
counts: dict[tuple[str, str], int] = {}
|
||||
for row in count_rows:
|
||||
counts[(str(row["from_unit_id"]), row["link_type"])] = int(row["cnt"])
|
||||
|
||||
# --- Temporal top-up ---
|
||||
temporal_needs = [r for r in victim_rows if counts.get((r["id"], "temporal"), 0) < MAX_TEMPORAL_LINKS_PER_UNIT]
|
||||
new_links: list[tuple] = []
|
||||
|
||||
if temporal_needs:
|
||||
lateral_unit_ids = [uuid_module.UUID(r["id"]) for r in temporal_needs if r["event_date"] is not None]
|
||||
lateral_event_dates = [
|
||||
_normalize_datetime(r["event_date"]) for r in temporal_needs if r["event_date"] is not None
|
||||
]
|
||||
lateral_fact_types = [r["fact_type"] for r in temporal_needs if r["event_date"] is not None]
|
||||
|
||||
if lateral_unit_ids:
|
||||
rows = await ops.fetch_temporal_neighbors(
|
||||
conn,
|
||||
fq_table("memory_units"),
|
||||
bank_id,
|
||||
lateral_unit_ids,
|
||||
lateral_event_dates,
|
||||
lateral_fact_types,
|
||||
MAX_TEMPORAL_LINKS_PER_UNIT,
|
||||
)
|
||||
for row in rows:
|
||||
time_diff_h = float(row["time_diff_hours"])
|
||||
# Mirror the 24h window enforced at retain time. The bidirectional
|
||||
# index scan returns the K closest neighbours regardless of
|
||||
# window, so we filter here.
|
||||
if time_diff_h > 24:
|
||||
continue
|
||||
weight = max(0.3, 1.0 - (time_diff_h / 24))
|
||||
new_links.append((row["from_id"], str(row["id"]), "temporal", weight, None))
|
||||
|
||||
# --- Semantic top-up ---
|
||||
# ANN must run on its own connection: it opens a nested transaction with
|
||||
# SET LOCAL hnsw.ef_search + CREATE TEMP TABLE ON COMMIT DROP, and nesting
|
||||
# that inside our current write transaction would commit our writes early.
|
||||
semantic_needs = [
|
||||
r
|
||||
for r in victim_rows
|
||||
if counts.get((r["id"], "semantic"), 0) < MAX_SEMANTIC_LINKS_PER_UNIT and r["embedding"] is not None
|
||||
]
|
||||
if semantic_needs:
|
||||
from ...memory_engine import acquire_with_retry
|
||||
|
||||
seed_ids = [r["id"] for r in semantic_needs]
|
||||
seed_embs = [r["embedding"] for r in semantic_needs]
|
||||
seed_ftypes = [r["fact_type"] for r in semantic_needs]
|
||||
async with acquire_with_retry(backend) as ann_conn:
|
||||
try:
|
||||
ann_links = await compute_semantic_links_ann(
|
||||
ann_conn,
|
||||
bank_id,
|
||||
seed_ids,
|
||||
seed_embs,
|
||||
fact_types=seed_ftypes,
|
||||
threshold=get_config().semantic_link_min_similarity,
|
||||
)
|
||||
# Strip self-links (rare but possible because the ANN probe
|
||||
# has no exclude list — see the comment in compute_semantic_links_ann).
|
||||
ann_links = [lnk for lnk in ann_links if lnk[0] != lnk[1]]
|
||||
new_links.extend(ann_links)
|
||||
except Exception as e:
|
||||
# ANN uses PG-specific HNSW syntax; on dialects/configs where
|
||||
# it isn't available we still want the temporal top-up to land.
|
||||
logger.warning(f"[GRAPH_MAINT] Semantic top-up failed for bank={bank_id}: {type(e).__name__}: {e}")
|
||||
|
||||
if not new_links:
|
||||
return 0
|
||||
|
||||
await _bulk_insert_links(
|
||||
conn,
|
||||
new_links,
|
||||
bank_id=bank_id,
|
||||
skip_exists_check=False,
|
||||
ops=ops,
|
||||
)
|
||||
return len(new_links)
|
||||
|
||||
|
||||
async def prune_orphan_entities(
|
||||
*,
|
||||
conn: DatabaseConnection,
|
||||
fq_table: Callable[[str], str],
|
||||
bank_id: str,
|
||||
) -> int:
|
||||
"""Delete ``entities`` rows in the bank with no remaining ``unit_entities``
|
||||
references. Returns the number pruned.
|
||||
|
||||
FK ON DELETE CASCADE on ``entity_cooccurrences`` then removes any
|
||||
cooccurrence row pointing at the pruned entities — which is why this runs
|
||||
before :func:`prune_stale_cooccurrences` rather than after.
|
||||
|
||||
A bank-wide single-statement delete, cheap when there's nothing to do. It is
|
||||
idempotent (rerunning only deletes what is still orphaned), so the caller is
|
||||
free to retry the whole transaction on deadlock.
|
||||
"""
|
||||
ops = _ops_for(conn)
|
||||
return await ops.prune_orphan_entities(
|
||||
conn,
|
||||
fq_table("entities"),
|
||||
fq_table("unit_entities"),
|
||||
bank_id,
|
||||
)
|
||||
|
||||
|
||||
async def prune_stale_cooccurrences(
|
||||
*,
|
||||
conn: DatabaseConnection,
|
||||
fq_table: Callable[[str], str],
|
||||
bank_id: str,
|
||||
) -> int:
|
||||
"""Delete cooccurrence rows no current memory witnesses. Returns the count.
|
||||
|
||||
Defensive sweep for rows where both endpoints still exist but no current
|
||||
memory_unit references both of them — the cooccurrence was real at the time
|
||||
it was recorded, but every unit that witnessed it has since been deleted.
|
||||
:func:`prune_orphan_entities` cascades the *missing-entity* case via FK; this
|
||||
pass catches the *stale-count* case it cannot see.
|
||||
|
||||
Like the orphan prune, a bank-wide idempotent sweep backed by indexes, so
|
||||
it's cheap when there's nothing to do and safe for the caller to retry.
|
||||
"""
|
||||
ops = _ops_for(conn)
|
||||
return await ops.prune_stale_cooccurrences(
|
||||
conn,
|
||||
fq_table("entity_cooccurrences"),
|
||||
fq_table("unit_entities"),
|
||||
fq_table("entities"),
|
||||
bank_id,
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"MAX_SEMANTIC_LINKS_PER_UNIT",
|
||||
"enqueue_relink_victims",
|
||||
"entities_for_units",
|
||||
"entity_map_for_units",
|
||||
"entity_memory_counts",
|
||||
"graph_direct_links",
|
||||
"graph_entity_rows",
|
||||
"graph_units",
|
||||
"prune_orphan_entities",
|
||||
"prune_stale_cooccurrences",
|
||||
"relink_pass",
|
||||
]
|
||||
@@ -0,0 +1,539 @@
|
||||
"""Addressed reads over `memory_units`: get, scan, count, tags, consolidation state.
|
||||
|
||||
Not retrieval — nothing here ranks. These are the queries behind the curation
|
||||
detail view, export, the bank-stats panel and the consolidation queue, lifted out
|
||||
of the call sites that used to issue them inline (``memory_engine``,
|
||||
``transfer/export``, ``consolidation/consolidator``) so
|
||||
:class:`~hindsight_api.engine.memories.postgres.PostgresMemories` can delegate
|
||||
rather than embed SQL.
|
||||
|
||||
Every function takes the live connection and Hindsight's ``fq_table`` resolver, so
|
||||
each one runs inside whatever transaction the caller already holds; none of them
|
||||
acquires a connection of its own.
|
||||
|
||||
**Cursor semantics.** ``scan_memories``'s ``page_token`` is opaque to callers, and
|
||||
for Postgres it is simply a *numeric offset rendered as a decimal string* against
|
||||
the scan's fixed ``ORDER BY created_at, id``. An empty token means "start at the
|
||||
beginning", and an empty token comes back once the walk is exhausted (i.e. the
|
||||
final short page). An offset cursor is a position rather than a snapshot — exactly
|
||||
the guarantee :class:`~hindsight_api.engine.memories.base.ScanPage` documents:
|
||||
rows written or deleted mid-walk can shift later pages, so a scan is
|
||||
eventually-complete browsing rather than a consistent iterator. ``skip`` is applied
|
||||
*on top of* the decoded cursor, so a caller that pages with both should pass
|
||||
``skip`` only on the first call — the returned token already accounts for it.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import uuid
|
||||
from collections.abc import Callable
|
||||
from datetime import datetime
|
||||
from typing import Any
|
||||
|
||||
from ...search.tags import (
|
||||
build_tag_groups_where_clause,
|
||||
build_tags_where_clause,
|
||||
build_tags_where_clause_simple,
|
||||
)
|
||||
from ..base import ScanPage, StoredMemory
|
||||
|
||||
# The `memory_units` projection every read here shares. Superset of the by-id
|
||||
# SELECT the recall source-facts path used (text/fact_type/context/timestamps/
|
||||
# document_id/chunk_id/tags/metadata), plus the observation bookkeeping columns
|
||||
# `StoredMemory` carries: source_memory_ids and consolidated_at.
|
||||
_MEMORY_COLUMNS = """
|
||||
id, text, fact_type, context, document_id, chunk_id, tags, metadata,
|
||||
proof_count, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
created_at, source_memory_ids, consolidated_at, observation_scopes
|
||||
"""
|
||||
|
||||
# The scan's order. Fixed (created_at, id) like the export loader's, because an
|
||||
# offset cursor is only meaningful against a total order.
|
||||
_SCAN_ORDER = "ORDER BY created_at, id"
|
||||
|
||||
|
||||
def _as_json(value: Any) -> Any:
|
||||
"""Coerce an asyncpg JSONB column (str or already-decoded) to a Python object.
|
||||
|
||||
Connections differ in whether a JSONB codec is registered, so the column
|
||||
arrives either as text or as the decoded object.
|
||||
"""
|
||||
if value is None:
|
||||
return None
|
||||
if isinstance(value, str):
|
||||
try:
|
||||
return json.loads(value)
|
||||
except json.JSONDecodeError:
|
||||
# A valid scalar such as `"combined"` arrives already decoded on
|
||||
# connections that do register a decoder.
|
||||
return value
|
||||
return value
|
||||
|
||||
|
||||
def _as_uuids(unit_ids: list[Any]) -> list[uuid.UUID]:
|
||||
"""Unit ids as UUIDs, dropping anything unparseable.
|
||||
|
||||
A malformed id is treated the same way a deleted one is — simply absent from
|
||||
the result — rather than failing the whole read.
|
||||
"""
|
||||
out: list[uuid.UUID] = []
|
||||
for unit_id in unit_ids or []:
|
||||
if isinstance(unit_id, uuid.UUID):
|
||||
out.append(unit_id)
|
||||
continue
|
||||
try:
|
||||
out.append(uuid.UUID(str(unit_id)))
|
||||
except (ValueError, AttributeError, TypeError):
|
||||
continue
|
||||
return out
|
||||
|
||||
|
||||
def _column(row: Any, name: str, default: Any = None) -> Any:
|
||||
"""One column of an asyncpg Record, tolerating a narrower projection."""
|
||||
try:
|
||||
return row[name]
|
||||
except (KeyError, IndexError):
|
||||
return default
|
||||
|
||||
|
||||
def _stored_from_row(row: Any) -> StoredMemory:
|
||||
"""Map a `memory_units` row onto :class:`StoredMemory`.
|
||||
|
||||
Shared by every read in this module so the row → dataclass mapping exists
|
||||
once. ``entity_ids`` stays empty: the unit→entity posting lives in
|
||||
`unit_entities` and is served by ``entities_for_units``, not by a join here.
|
||||
"""
|
||||
source_ids = _column(row, "source_memory_ids") or []
|
||||
return StoredMemory(
|
||||
unit_id=str(row["id"]),
|
||||
text=row["text"],
|
||||
fact_type=row["fact_type"],
|
||||
context=_column(row, "context"),
|
||||
document_id=_column(row, "document_id"),
|
||||
chunk_id=str(_column(row, "chunk_id")) if _column(row, "chunk_id") else None,
|
||||
tags=list(_column(row, "tags") or []),
|
||||
metadata=_as_json(_column(row, "metadata")),
|
||||
proof_count=_column(row, "proof_count") or 1,
|
||||
event_date=_column(row, "event_date"),
|
||||
occurred_start=_column(row, "occurred_start"),
|
||||
occurred_end=_column(row, "occurred_end"),
|
||||
mentioned_at=_column(row, "mentioned_at"),
|
||||
created_at=_column(row, "created_at"),
|
||||
source_memory_ids=[str(sid) for sid in source_ids],
|
||||
consolidated_at=_column(row, "consolidated_at"),
|
||||
# Consolidation routes a candidate by its scopes, so this has to survive
|
||||
# the trip through the store rather than being re-queried per memory.
|
||||
observation_scopes=_as_json(_column(row, "observation_scopes")),
|
||||
)
|
||||
|
||||
|
||||
def _decode_page_token(page_token: str) -> int:
|
||||
"""Decode the offset cursor. Empty, malformed or negative all mean "start"."""
|
||||
if not page_token:
|
||||
return 0
|
||||
try:
|
||||
offset = int(page_token)
|
||||
except (TypeError, ValueError):
|
||||
return 0
|
||||
return offset if offset > 0 else 0
|
||||
|
||||
|
||||
async def get_memories(
|
||||
*, conn, fq_table: Callable[[str], str], bank_id: str, unit_ids: list[str]
|
||||
) -> list[StoredMemory]:
|
||||
"""Fetch memories by id. Missing or deleted ids are simply absent."""
|
||||
ids = _as_uuids(unit_ids)
|
||||
if not ids:
|
||||
return []
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT {_MEMORY_COLUMNS}
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $1 AND id = ANY($2::uuid[])
|
||||
""",
|
||||
bank_id,
|
||||
ids,
|
||||
)
|
||||
return [_stored_from_row(row) for row in rows]
|
||||
|
||||
|
||||
async def _semantic_edges(
|
||||
*, conn, fq_table: Callable[[str], str], bank_id: str, unit_ids: list[uuid.UUID]
|
||||
) -> dict[str, list[tuple[str, float]]]:
|
||||
"""Derived kNN edges for ``unit_ids``, keyed by unit id.
|
||||
|
||||
Walked in both directions, like the graph arm's semantic expansion: a
|
||||
`memory_links` row is written once, so a unit's neighbourhood is the union of
|
||||
the edges leaving it and those arriving at it.
|
||||
"""
|
||||
if not unit_ids:
|
||||
return {}
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT from_unit_id AS unit_id, to_unit_id AS target_id, weight
|
||||
FROM {fq_table("memory_links")}
|
||||
WHERE bank_id = $1 AND link_type = 'semantic' AND from_unit_id = ANY($2::uuid[])
|
||||
UNION ALL
|
||||
SELECT to_unit_id AS unit_id, from_unit_id AS target_id, weight
|
||||
FROM {fq_table("memory_links")}
|
||||
WHERE bank_id = $1 AND link_type = 'semantic' AND to_unit_id = ANY($2::uuid[])
|
||||
""",
|
||||
bank_id,
|
||||
unit_ids,
|
||||
)
|
||||
edges: dict[str, list[tuple[str, float]]] = {}
|
||||
for row in rows:
|
||||
edges.setdefault(str(row["unit_id"]), []).append((str(row["target_id"]), float(row["weight"] or 0.0)))
|
||||
return edges
|
||||
|
||||
|
||||
async def scan_memories(
|
||||
*,
|
||||
conn,
|
||||
fq_table: Callable[[str], str],
|
||||
bank_id: str,
|
||||
fact_types: list[str] | None = None,
|
||||
limit: int = 100,
|
||||
page_token: str = "",
|
||||
tags: list[str] | None = None,
|
||||
tags_match: str = "any",
|
||||
tag_groups: list | None = None,
|
||||
document_id: str | None = None,
|
||||
metadata_equals: dict[str, str] | None = None,
|
||||
skip: int = 0,
|
||||
include_edges: bool = False,
|
||||
) -> ScanPage:
|
||||
"""Page through stored memories. A full walk — for browsing and export only.
|
||||
|
||||
See the module docstring for the ``page_token`` (offset) cursor semantics.
|
||||
"""
|
||||
if limit is None or limit <= 0:
|
||||
return ScanPage()
|
||||
|
||||
where: list[str] = ["bank_id = $1"]
|
||||
params: list[Any] = [bank_id]
|
||||
|
||||
if fact_types:
|
||||
params.append(list(fact_types))
|
||||
where.append(f"fact_type = ANY(${len(params)})")
|
||||
|
||||
if document_id is not None:
|
||||
# A real column here, which is why it is not folded into
|
||||
# `metadata_equals`: only a store without the column keeps it in the bag.
|
||||
params.append(document_id)
|
||||
where.append(f"document_id = ${len(params)}")
|
||||
|
||||
if metadata_equals:
|
||||
# str→str equality across every key, which is exactly JSONB containment.
|
||||
params.append(json.dumps(metadata_equals))
|
||||
where.append(f"metadata @> ${len(params)}::jsonb")
|
||||
|
||||
# The tags clause owns its own `AND` prefix and, per the helper's contract,
|
||||
# only consumes a bind param when `tags` is non-empty (match="exact" with no
|
||||
# tags is the untagged/global scope and needs none).
|
||||
tags_clause = build_tags_where_clause_simple(tags, len(params) + 1, match=tags_match)
|
||||
if tags:
|
||||
params.append(list(tags))
|
||||
|
||||
# Compound tag groups (AND/OR/NOT trees), AND-ed on. Also owns its `AND` prefix and appends
|
||||
# one bind param per leaf; empty/absent groups yield no clause and no params.
|
||||
groups_clause, group_params, _ = build_tag_groups_where_clause(tag_groups, param_offset=len(params) + 1)
|
||||
params.extend(group_params)
|
||||
|
||||
offset = _decode_page_token(page_token) + max(int(skip or 0), 0)
|
||||
params.append(limit)
|
||||
limit_idx = len(params)
|
||||
params.append(offset)
|
||||
offset_idx = len(params)
|
||||
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT {_MEMORY_COLUMNS}
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE {" AND ".join(where)} {tags_clause} {groups_clause}
|
||||
{_SCAN_ORDER}
|
||||
LIMIT ${limit_idx} OFFSET ${offset_idx}
|
||||
""",
|
||||
*params,
|
||||
)
|
||||
|
||||
memories = [_stored_from_row(row) for row in rows]
|
||||
if include_edges and memories:
|
||||
edges = await _semantic_edges(
|
||||
conn=conn, fq_table=fq_table, bank_id=bank_id, unit_ids=_as_uuids([m.unit_id for m in memories])
|
||||
)
|
||||
for memory in memories:
|
||||
memory.semantic_edges = edges.get(memory.unit_id, [])
|
||||
|
||||
# A short page means the walk is exhausted, so the cursor goes empty.
|
||||
next_token = str(offset + len(rows)) if len(rows) == limit else ""
|
||||
return ScanPage(memories=memories, next_page_token=next_token)
|
||||
|
||||
|
||||
async def count_memories(*, conn, fq_table: Callable[[str], str], bank_id: str) -> dict[str, int]:
|
||||
"""Live memory count per fact_type. The bank-stats node counts."""
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT fact_type, COUNT(*) as count
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $1
|
||||
GROUP BY fact_type
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
return {row["fact_type"]: int(row["count"]) for row in rows}
|
||||
|
||||
|
||||
async def list_tags(
|
||||
*,
|
||||
conn,
|
||||
fq_table: Callable[[str], str],
|
||||
bank_id: str,
|
||||
pattern: str | None = None,
|
||||
limit: int = 100,
|
||||
offset: int = 0,
|
||||
) -> dict[str, Any]:
|
||||
"""One page of a bank's tag histogram: ``{"items": [{tag, count}], "total", "limit", "offset"}``.
|
||||
|
||||
``memory_units`` lives in SQL for this store, so the wildcard filter, the
|
||||
``count DESC, tag ASC`` ordering and the paging all run in SQL — the whole
|
||||
histogram never crosses the wire. The dialect fragments come from
|
||||
``build_tag_listing_parts`` (``unnest`` on Postgres, ``JSON_TABLE`` on Oracle):
|
||||
this module backs both dialects, so it must not inline either one's SQL.
|
||||
"""
|
||||
from ...db import create_data_access_ops
|
||||
|
||||
ops = create_data_access_ops(getattr(conn, "backend_type", "postgresql"))
|
||||
tag_parts = ops.build_tag_listing_parts(fq_table("memory_units"))
|
||||
tag_source = tag_parts.tag_source
|
||||
non_empty_check = tag_parts.non_empty_check
|
||||
tag_col = tag_parts.tag_col
|
||||
bank_prefix = tag_parts.bank_prefix
|
||||
|
||||
params: list[Any] = [bank_id]
|
||||
pattern_clause = ""
|
||||
if pattern:
|
||||
# '*' is the wildcard, matched case-insensitively — same anchored ILIKE semantics as before.
|
||||
params.append(pattern.replace("*", "%"))
|
||||
pattern_clause = f"AND {tag_col} ILIKE $2"
|
||||
|
||||
total_row = await conn.fetchrow(
|
||||
f"""
|
||||
SELECT COUNT(DISTINCT {tag_col}) as total
|
||||
FROM {tag_source}
|
||||
WHERE {bank_prefix}bank_id = $1 {non_empty_check}
|
||||
{pattern_clause}
|
||||
""",
|
||||
*params,
|
||||
)
|
||||
total = int(total_row["total"]) if total_row else 0
|
||||
|
||||
limit_param = len(params) + 1
|
||||
offset_param = len(params) + 2
|
||||
params.extend([limit, offset])
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT {tag_col} as tag, COUNT(*) as count
|
||||
FROM {tag_source}
|
||||
WHERE {bank_prefix}bank_id = $1 {non_empty_check}
|
||||
{pattern_clause}
|
||||
GROUP BY {tag_col}
|
||||
ORDER BY count DESC, {tag_col} ASC
|
||||
LIMIT ${limit_param} OFFSET ${offset_param}
|
||||
""",
|
||||
*params,
|
||||
)
|
||||
return {
|
||||
"items": [{"tag": row["tag"], "count": int(row["count"])} for row in rows],
|
||||
"total": total,
|
||||
"limit": limit,
|
||||
"offset": offset,
|
||||
}
|
||||
|
||||
|
||||
async def find_unconsolidated(
|
||||
*,
|
||||
conn,
|
||||
fq_table: Callable[[str], str],
|
||||
bank_id: str,
|
||||
fact_types: list[str],
|
||||
limit: int,
|
||||
scope_tags: list[str] | None = None,
|
||||
) -> list[StoredMemory]:
|
||||
"""Memories not yet folded into an observation, oldest first.
|
||||
|
||||
The consolidator's candidate query: never consolidated, never *failed* to
|
||||
consolidate (a memory the LLM could not handle must not be retried forever),
|
||||
ordered by ``created_at`` so the queue drains in arrival order. ``scope_tags``
|
||||
is the same ``tags @> scope`` containment the job's scope filter uses — the
|
||||
job ORs several scopes together; one scope is passed here.
|
||||
"""
|
||||
where = [
|
||||
"bank_id = $1",
|
||||
"consolidated_at IS NULL",
|
||||
"consolidation_failed_at IS NULL",
|
||||
]
|
||||
params: list[Any] = [bank_id]
|
||||
if fact_types:
|
||||
params.append(list(fact_types))
|
||||
where.append(f"fact_type = ANY(${len(params)})")
|
||||
if scope_tags:
|
||||
params.append(list(scope_tags))
|
||||
where.append(f"tags @> ${len(params)}::varchar[]")
|
||||
params.append(limit)
|
||||
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT {_MEMORY_COLUMNS}
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE {" AND ".join(where)}
|
||||
ORDER BY created_at ASC
|
||||
LIMIT ${len(params)}
|
||||
""",
|
||||
*params,
|
||||
)
|
||||
return [_stored_from_row(row) for row in rows]
|
||||
|
||||
|
||||
async def count_unconsolidated(
|
||||
*,
|
||||
conn,
|
||||
fq_table: Callable[[str], str],
|
||||
bank_id: str,
|
||||
fact_types: list[str],
|
||||
scopes: list[list[str] | None],
|
||||
limit: int,
|
||||
) -> int:
|
||||
"""Bounded ``COUNT(*)`` of unconsolidated candidates matching any scope — the cheap counterpart
|
||||
to :func:`find_unconsolidated` that never ships a row.
|
||||
|
||||
Same predicates as ``find_unconsolidated`` (never consolidated, never failed, matching
|
||||
fact_type), with the scopes OR'd as ``tags @> scope`` containment. ``id`` is the PK so each row
|
||||
counts once; the inner ``LIMIT`` floors the count at ``limit`` exactly as walking that many rows
|
||||
would, so a huge backlog stays a single index count instead of a 17-column fetch.
|
||||
"""
|
||||
where = ["bank_id = $1", "consolidated_at IS NULL", "consolidation_failed_at IS NULL"]
|
||||
params: list[Any] = [bank_id]
|
||||
if fact_types:
|
||||
params.append(list(fact_types))
|
||||
where.append(f"fact_type = ANY(${len(params)})")
|
||||
# An unscoped entry (None) matches every row, collapsing the OR to no tag filter at all.
|
||||
scope_clauses: list[str] = []
|
||||
unscoped = any(scope is None for scope in scopes)
|
||||
if not unscoped:
|
||||
for scope in scopes:
|
||||
params.append(list(scope or []))
|
||||
scope_clauses.append(f"tags @> ${len(params)}::varchar[]")
|
||||
if scope_clauses:
|
||||
where.append("(" + " OR ".join(scope_clauses) + ")")
|
||||
params.append(limit)
|
||||
row = await conn.fetchrow(
|
||||
f"""
|
||||
SELECT COUNT(*) AS c FROM (
|
||||
SELECT 1 FROM {fq_table("memory_units")}
|
||||
WHERE {" AND ".join(where)}
|
||||
LIMIT ${len(params)}
|
||||
) sub
|
||||
""",
|
||||
*params,
|
||||
)
|
||||
return int(row["c"]) if row else 0
|
||||
|
||||
|
||||
async def mark_consolidated(
|
||||
*,
|
||||
conn,
|
||||
fq_table: Callable[[str], str],
|
||||
bank_id: str,
|
||||
unit_ids: list[str],
|
||||
when: datetime | None,
|
||||
failed: bool = False,
|
||||
) -> None:
|
||||
"""Stamp (or clear, with ``when=None``) the consolidated marker on sources.
|
||||
|
||||
``failed`` writes ``consolidation_failed_at`` instead of ``consolidated_at``,
|
||||
which is what keeps a memory the LLM could not consolidate out of the queue.
|
||||
|
||||
``when=None`` clears the column rather than stamping it — that is how a source
|
||||
is requeued once the observation built on it is deleted. The clear keeps the
|
||||
``fact_type IN ('experience', 'world')`` guard the requeue sites carry:
|
||||
observations are never themselves consolidated, so nothing about them should
|
||||
be reset by a requeue.
|
||||
|
||||
``updated_at`` is deliberately left alone, matching the consolidator's own
|
||||
statements: consolidation bookkeeping is not an edit to the memory, and
|
||||
bumping it would make every consolidation pass look like a write to the
|
||||
staleness check below.
|
||||
"""
|
||||
ids = _as_uuids(unit_ids)
|
||||
if not ids:
|
||||
return
|
||||
column = "consolidation_failed_at" if failed else "consolidated_at"
|
||||
guard = "" if when is not None else " AND fact_type IN ('experience', 'world')"
|
||||
await conn.execute(
|
||||
f"""
|
||||
UPDATE {fq_table("memory_units")}
|
||||
SET {column} = $1
|
||||
WHERE bank_id = $2 AND id = ANY($3::uuid[]){guard}
|
||||
""",
|
||||
when,
|
||||
bank_id,
|
||||
ids,
|
||||
)
|
||||
|
||||
|
||||
async def any_memory_updated_since(
|
||||
*,
|
||||
conn,
|
||||
fq_table: Callable[[str], str],
|
||||
bank_id: str,
|
||||
since: datetime,
|
||||
fact_types: list[str] | None = None,
|
||||
tags: list[str] | None = None,
|
||||
tags_match: str = "any",
|
||||
tag_groups: list | None = None,
|
||||
) -> bool:
|
||||
"""Whether any memory in ``bank_id``'s scope was written after ``since``.
|
||||
|
||||
Backs the mental-model staleness check, so it is a bounded existence test —
|
||||
``LIMIT 1``, never a COUNT: the answer is "is there one", and the planner can
|
||||
stop at the first hit. The scope is the mental model's: its flat tags (or the
|
||||
compound ``tag_groups``) plus an optional ``fact_types`` restriction. This is
|
||||
where the staleness query's WHERE lives, so the same scope that gates a
|
||||
refresh decides whether one is due.
|
||||
"""
|
||||
params: list[Any] = [bank_id, since]
|
||||
where = ["bank_id = $1", "updated_at > $2"]
|
||||
|
||||
tag_clause, tag_params, next_param = build_tags_where_clause(tags, param_offset=len(params) + 1, match=tags_match)
|
||||
if tag_clause:
|
||||
where.append(tag_clause.removeprefix("AND "))
|
||||
params.extend(tag_params)
|
||||
|
||||
group_clause, group_params, _ = build_tag_groups_where_clause(tag_groups, param_offset=next_param)
|
||||
if group_clause:
|
||||
where.append(group_clause.removeprefix("AND "))
|
||||
params.extend(group_params)
|
||||
# Untagged, no tag_groups → no tag constraint, matching any memory in the bank.
|
||||
|
||||
if fact_types:
|
||||
params.append(list(fact_types))
|
||||
where.append(f"fact_type = ANY(${len(params)}::text[])")
|
||||
|
||||
row = await conn.fetchval(
|
||||
f"SELECT 1 FROM {fq_table('memory_units')} WHERE {' AND '.join(where)} LIMIT 1",
|
||||
*params,
|
||||
)
|
||||
return row is not None
|
||||
|
||||
|
||||
__all__ = [
|
||||
"any_memory_updated_since",
|
||||
"count_memories",
|
||||
"find_unconsolidated",
|
||||
"get_memories",
|
||||
"list_tags",
|
||||
"mark_consolidated",
|
||||
"scan_memories",
|
||||
]
|
||||
@@ -0,0 +1,584 @@
|
||||
"""Writes against `memory_units`: the fact insert, the deletes, and observation invalidation.
|
||||
|
||||
Everything here mutates the memories slice and nothing else. The document row,
|
||||
the chunks, the entity registry and the link tables stay with their own callers —
|
||||
what lands in this module is only the statements that touch `memory_units` (and,
|
||||
on backends that keep one, the `observation_sources` junction that hangs off it).
|
||||
|
||||
Each function takes the live connection and Hindsight's ``fq_table`` resolver, so
|
||||
it runs inside whatever transaction the caller already holds; ``ops`` is the
|
||||
dialect ops object, which is what lets the same code serve the PG (native array)
|
||||
and Oracle (junction table) shapes of the observation→source relation.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import uuid
|
||||
from collections.abc import Callable
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from ....config import get_config
|
||||
from ..base import StoredMemory
|
||||
|
||||
if TYPE_CHECKING: # pragma: no cover - typing only
|
||||
from ...retain.types import ProcessedFact
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
async def insert_facts(
|
||||
*,
|
||||
conn,
|
||||
ops,
|
||||
bank_id: str,
|
||||
facts: list[ProcessedFact],
|
||||
document_id: str | None = None,
|
||||
) -> list[str]:
|
||||
"""Insert facts into the database in batch.
|
||||
|
||||
Args:
|
||||
conn: Database connection
|
||||
bank_id: Bank identifier
|
||||
facts: List of ProcessedFact objects to insert
|
||||
document_id: Optional document ID to associate with facts
|
||||
|
||||
Returns:
|
||||
List of unit IDs (UUIDs as strings) for the inserted facts, in the same
|
||||
order as ``facts``.
|
||||
"""
|
||||
if not facts:
|
||||
return []
|
||||
|
||||
# Imported here: `retain` reaches back into the engine for `fq_table`, so a
|
||||
# module-level import would close the cycle once the engine imports this store.
|
||||
from ...retain.fact_extraction import _sanitize_text
|
||||
|
||||
# Prepare data for batch insert
|
||||
fact_texts = []
|
||||
embeddings = []
|
||||
event_dates = []
|
||||
occurred_starts = []
|
||||
occurred_ends = []
|
||||
mentioned_ats = []
|
||||
contexts = []
|
||||
fact_types = []
|
||||
metadata_jsons = []
|
||||
chunk_ids = []
|
||||
document_ids = []
|
||||
tags_list = []
|
||||
observation_scopes_list = []
|
||||
text_signals_list = []
|
||||
|
||||
for fact in facts:
|
||||
fact_texts.append(_sanitize_text(fact.fact_text))
|
||||
# Convert embedding to string for asyncpg vector type
|
||||
embeddings.append(str(fact.embedding))
|
||||
# event_date: Use occurred_start if available, otherwise use mentioned_at
|
||||
# This maintains backward compatibility while handling None occurred_start
|
||||
event_dates.append(fact.occurred_start if fact.occurred_start is not None else fact.mentioned_at)
|
||||
occurred_starts.append(fact.occurred_start)
|
||||
occurred_ends.append(fact.occurred_end)
|
||||
mentioned_ats.append(fact.mentioned_at)
|
||||
contexts.append(_sanitize_text(fact.context))
|
||||
fact_types.append(fact.fact_type)
|
||||
metadata_jsons.append(json.dumps(fact.metadata))
|
||||
chunk_ids.append(fact.chunk_id)
|
||||
# Use per-fact document_id if available, otherwise fallback to batch-level document_id
|
||||
document_ids.append(fact.document_id if fact.document_id else document_id)
|
||||
# Convert tags to JSON string for proper batch insertion (PostgreSQL unnest doesn't handle 2D arrays well)
|
||||
tags_list.append(json.dumps(fact.tags if fact.tags else []))
|
||||
# observation_scopes: stored as JSONB (string or 2D array), None if not provided
|
||||
observation_scopes_list.append(
|
||||
json.dumps(fact.observation_scopes) if fact.observation_scopes is not None else None
|
||||
)
|
||||
# Build text_signals: entity names + date tokens for enriched BM25 indexing
|
||||
signal_parts = []
|
||||
if fact.entities:
|
||||
signal_parts.extend(e.name for e in fact.entities)
|
||||
if fact.occurred_start:
|
||||
try:
|
||||
signal_parts.append(fact.occurred_start.strftime("%B %d %Y").lstrip("0").replace(" 0", " "))
|
||||
except (ValueError, AttributeError):
|
||||
pass
|
||||
if fact.occurred_end and fact.occurred_end != fact.occurred_start:
|
||||
try:
|
||||
signal_parts.append(fact.occurred_end.strftime("%B %d %Y").lstrip("0").replace(" 0", " "))
|
||||
except (ValueError, AttributeError):
|
||||
pass
|
||||
text_signals_list.append(" ".join(signal_parts) if signal_parts else None)
|
||||
|
||||
# Batch insert all facts — delegates to DataAccessOps which handles
|
||||
# unnest (PG) vs row-by-row (Oracle) transparently.
|
||||
config = get_config()
|
||||
|
||||
return await ops.insert_facts_batch(
|
||||
conn,
|
||||
bank_id,
|
||||
fact_texts,
|
||||
embeddings,
|
||||
event_dates,
|
||||
occurred_starts,
|
||||
occurred_ends,
|
||||
mentioned_ats,
|
||||
contexts,
|
||||
fact_types,
|
||||
metadata_jsons,
|
||||
chunk_ids,
|
||||
document_ids,
|
||||
tags_list,
|
||||
observation_scopes_list,
|
||||
text_signals_list,
|
||||
text_search_extension=config.text_search_extension,
|
||||
)
|
||||
|
||||
|
||||
async def delete_document(*, conn, fq_table: Callable[[str], str], bank_id: str, document_id: str) -> None:
|
||||
"""Delete every memory unit belonging to ``document_id``.
|
||||
|
||||
Explicitly delete memory_units by document_id BEFORE deleting the
|
||||
document row. The CASCADE from documents→chunks→memory_units only
|
||||
catches units that have a non-NULL chunk_id FK. Units with chunk_id=NULL
|
||||
(e.g. from partial writes or edge cases) would survive the cascade.
|
||||
This explicit delete ensures complete cleanup.
|
||||
|
||||
Called when a document is replaced, so it races the replacement's writes: it
|
||||
must remove only what was written *before* this call, never the facts
|
||||
arriving moments later — which the ``document_id``/``bank_id`` predicate
|
||||
gives for free inside the caller's transaction.
|
||||
"""
|
||||
await conn.execute(
|
||||
f"DELETE FROM {fq_table('memory_units')} WHERE document_id = $1 AND bank_id = $2",
|
||||
document_id,
|
||||
bank_id,
|
||||
)
|
||||
|
||||
|
||||
async def delete_observations(*, conn, fq_table: Callable[[str], str], bank_id: str) -> None:
|
||||
"""Delete all observations in a bank, leaving the facts behind them.
|
||||
|
||||
Only the observation rows: requeuing the surviving sources (clearing
|
||||
``consolidated_at``) and resetting the bank's consolidation timestamp belong
|
||||
to the caller, which owns the bank row.
|
||||
"""
|
||||
await conn.execute(
|
||||
f"DELETE FROM {fq_table('memory_units')} WHERE bank_id = $1 AND fact_type = 'observation'",
|
||||
bank_id,
|
||||
)
|
||||
|
||||
|
||||
async def observations_for_sources(
|
||||
*,
|
||||
conn,
|
||||
ops,
|
||||
fq_table: Callable[[str], str],
|
||||
bank_id: str,
|
||||
unit_ids: list[str | uuid.UUID],
|
||||
) -> list[StoredMemory]:
|
||||
"""Observations consolidated from any of ``unit_ids``.
|
||||
|
||||
Only ``unit_id`` and ``source_memory_ids`` are populated — the caller uses
|
||||
them to delete the observations and to work out which sources survive, and
|
||||
the rest of the row is about to be deleted anyway.
|
||||
"""
|
||||
if not unit_ids:
|
||||
return []
|
||||
|
||||
fact_uuids = [uuid.UUID(str(fid)) if not isinstance(fid, uuid.UUID) else fid for fid in unit_ids]
|
||||
|
||||
if ops is not None and not ops.uses_observation_sources_table:
|
||||
# PG: use native array overlap operator
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, source_memory_ids
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $1
|
||||
AND fact_type = 'observation'
|
||||
AND source_memory_ids && $2::uuid[]
|
||||
""",
|
||||
bank_id,
|
||||
fact_uuids,
|
||||
)
|
||||
else:
|
||||
# Oracle / default: use observation_sources junction table
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT mu.id, mu.source_memory_ids
|
||||
FROM {fq_table("memory_units")} mu
|
||||
WHERE mu.bank_id = $1
|
||||
AND mu.fact_type = 'observation'
|
||||
AND EXISTS (
|
||||
SELECT 1 FROM {fq_table("observation_sources")} os
|
||||
WHERE os.observation_id = mu.id
|
||||
AND os.source_id = ANY($2::uuid[])
|
||||
)
|
||||
""",
|
||||
bank_id,
|
||||
fact_uuids,
|
||||
)
|
||||
|
||||
return [
|
||||
StoredMemory(
|
||||
unit_id=str(row["id"]),
|
||||
text="",
|
||||
fact_type="observation",
|
||||
source_memory_ids=[str(src_id) for src_id in (row["source_memory_ids"] or [])],
|
||||
)
|
||||
for row in rows
|
||||
]
|
||||
|
||||
|
||||
async def delete_stale_observations(
|
||||
*,
|
||||
conn,
|
||||
ops,
|
||||
fq_table: Callable[[str], str],
|
||||
bank_id: str,
|
||||
fact_ids: list[str | uuid.UUID],
|
||||
) -> int:
|
||||
"""Delete observations whose source memories are about to be removed.
|
||||
|
||||
Mirrors the cleanup performed by ``MemoryEngine.delete_document`` so that
|
||||
every code path that removes ``memory_units`` also removes the
|
||||
observations derived from them. Without this, ingesting a fresh version
|
||||
of a document via the retain pipeline (which does a full-replace
|
||||
``DELETE FROM documents`` cascade) used to leave orphan observations
|
||||
pointing at memory IDs that no longer existed.
|
||||
|
||||
For each observation referencing any of ``fact_ids``:
|
||||
1. Delete the observation row (its text is stale once even one source
|
||||
memory disappears).
|
||||
2. Reset ``consolidated_at = NULL`` on the surviving source memories so
|
||||
they get re-consolidated under fresh observations on the next run.
|
||||
|
||||
Must be called within an active transaction, before the source memories
|
||||
are deleted.
|
||||
|
||||
Returns the number of observations deleted.
|
||||
"""
|
||||
if not fact_ids:
|
||||
return 0
|
||||
|
||||
fact_uuids = [uuid.UUID(str(fid)) if not isinstance(fid, uuid.UUID) else fid for fid in fact_ids]
|
||||
|
||||
affected_obs = await observations_for_sources(
|
||||
conn=conn, ops=ops, fq_table=fq_table, bank_id=bank_id, unit_ids=fact_uuids
|
||||
)
|
||||
if not affected_obs:
|
||||
return 0
|
||||
|
||||
deleted_set = {str(uid) for uid in fact_uuids}
|
||||
obs_ids = [uuid.UUID(obs.unit_id) for obs in affected_obs]
|
||||
seen_remaining: set[str] = set()
|
||||
remaining_source_ids: list[uuid.UUID] = []
|
||||
for obs in affected_obs:
|
||||
for src_str in obs.source_memory_ids:
|
||||
if src_str not in deleted_set and src_str not in seen_remaining:
|
||||
remaining_source_ids.append(uuid.UUID(src_str))
|
||||
seen_remaining.add(src_str)
|
||||
|
||||
await conn.execute(
|
||||
f"DELETE FROM {fq_table('memory_units')} WHERE id = ANY($1::uuid[])",
|
||||
obs_ids,
|
||||
)
|
||||
|
||||
# Their history is keyed by observation_id and no longer cascades from memory_units (that FK
|
||||
# was dropped so history can be recorded for observations kept outside SQL), so drop the
|
||||
# deleted observations' snapshots explicitly rather than leaving them to accumulate.
|
||||
await conn.execute(
|
||||
f"DELETE FROM {fq_table('observation_history')} WHERE bank_id = $1 AND observation_id = ANY($2::uuid[])",
|
||||
bank_id,
|
||||
obs_ids,
|
||||
)
|
||||
|
||||
if remaining_source_ids:
|
||||
await conn.execute(
|
||||
f"""
|
||||
UPDATE {fq_table("memory_units")}
|
||||
SET consolidated_at = NULL
|
||||
WHERE id = ANY($1::uuid[])
|
||||
AND fact_type IN ('experience', 'world')
|
||||
""",
|
||||
remaining_source_ids,
|
||||
)
|
||||
|
||||
logger.info(
|
||||
f"[OBSERVATIONS] Deleted {len(obs_ids)} observations, reset {len(remaining_source_ids)} "
|
||||
f"source memories for re-consolidation in bank {bank_id}"
|
||||
)
|
||||
return len(obs_ids)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------- curation archive
|
||||
#
|
||||
# Invalidation moves a rejected memory between two tables rather than flagging it,
|
||||
# so recall / consolidation / graph never carry a "valid?" predicate: live facts
|
||||
# live in `memory_units`, invalidated ones in `invalidated_memory_units`. The
|
||||
# archive is cold storage — no index, so it drops the `embedding` and
|
||||
# `search_vector` columns, which are recomputed on the way back.
|
||||
|
||||
# The two recall-surface columns the archive omits. Both follow server config
|
||||
# (embedding dimension, search backend), so keeping them out of the INSERT…SELECT
|
||||
# round-trip makes a model or text-backend switch structurally unable to trip a
|
||||
# type/dimension mismatch (#2209, #2503); each is recomputed on revert.
|
||||
_ARCHIVE_OMITTED = ('"embedding"', '"search_vector"')
|
||||
|
||||
|
||||
async def _memory_unit_columns(conn, fq_table: Callable[[str], str]) -> str:
|
||||
"""The quoted, ordinal column list of `memory_units`.
|
||||
|
||||
Read from the catalog rather than hardcoded so a schema migration cannot make
|
||||
the archive round-trip drift from the live table (the archive is created via
|
||||
``LIKE memory_units``, so the lists line up).
|
||||
"""
|
||||
rows = await conn.fetch(
|
||||
f"SELECT a.attname FROM pg_attribute a "
|
||||
f"WHERE a.attrelid = '{fq_table('memory_units')}'::regclass "
|
||||
f"AND a.attnum > 0 AND NOT a.attisdropped ORDER BY a.attnum"
|
||||
)
|
||||
return ", ".join(f'"{r["attname"]}"' for r in rows)
|
||||
|
||||
|
||||
async def _archive_columns(conn, fq_table: Callable[[str], str]) -> str:
|
||||
"""`_memory_unit_columns` minus the two the archive does not carry."""
|
||||
collist = await _memory_unit_columns(conn, fq_table)
|
||||
return ", ".join(c for c in (s.strip() for s in collist.split(",")) if c not in _ARCHIVE_OMITTED)
|
||||
|
||||
|
||||
_ARCHIVE_SELECT = (
|
||||
"id, text, fact_type, context, occurred_start, occurred_end, mentioned_at, "
|
||||
"document_id, chunk_id, tags, metadata, proof_count, event_date, created_at, "
|
||||
"consolidated_at, entity_ids"
|
||||
)
|
||||
|
||||
|
||||
def _archived_stored(row: Any) -> StoredMemory:
|
||||
"""Map an `invalidated_memory_units` row onto :class:`StoredMemory`."""
|
||||
return StoredMemory(
|
||||
unit_id=str(row["id"]),
|
||||
text=row["text"],
|
||||
fact_type=row["fact_type"],
|
||||
context=row["context"],
|
||||
document_id=row["document_id"],
|
||||
chunk_id=str(row["chunk_id"]) if row["chunk_id"] else None,
|
||||
tags=list(row["tags"] or []),
|
||||
metadata=row["metadata"] if isinstance(row["metadata"], dict) else None,
|
||||
proof_count=row["proof_count"] or 1,
|
||||
event_date=row["event_date"],
|
||||
occurred_start=row["occurred_start"],
|
||||
occurred_end=row["occurred_end"],
|
||||
mentioned_at=row["mentioned_at"],
|
||||
created_at=row["created_at"],
|
||||
consolidated_at=row["consolidated_at"],
|
||||
entity_ids=[str(e) for e in (row["entity_ids"] or [])],
|
||||
)
|
||||
|
||||
|
||||
async def get_archived_memory(*, conn, fq_table, bank_id: str, unit_id: str) -> StoredMemory | None:
|
||||
row = await conn.fetchrow(
|
||||
f"SELECT {_ARCHIVE_SELECT} FROM {fq_table('invalidated_memory_units')} WHERE id = $1 AND bank_id = $2",
|
||||
str(unit_id),
|
||||
bank_id,
|
||||
)
|
||||
return _archived_stored(row) if row else None
|
||||
|
||||
|
||||
async def invalidate_memory(*, conn, fq_table, bank_id: str, unit_id: str, reason: str | None) -> bool:
|
||||
mu = fq_table("memory_units")
|
||||
arch = fq_table("invalidated_memory_units")
|
||||
ue = fq_table("unit_entities")
|
||||
arch_cols = await _archive_columns(conn, fq_table)
|
||||
|
||||
# Snapshot the entity ids before the delete cascade takes `unit_entities`, so
|
||||
# revert can restore the postings the move is about to drop.
|
||||
entity_ids = [
|
||||
r["entity_id"] for r in await conn.fetch(f"SELECT entity_id FROM {ue} WHERE unit_id = $1", str(unit_id))
|
||||
]
|
||||
# Causal edges are retain-time extraction output the FK cascade would destroy for good —
|
||||
# unlike temporal/semantic links they can't be recomputed, so snapshot their descriptors onto
|
||||
# the archive row and revert rematerializes them (#2864).
|
||||
from ...retain.link_utils import snapshot_causal_links
|
||||
|
||||
causal_links = await snapshot_causal_links(conn, bank_id, str(unit_id))
|
||||
inserted = await conn.fetchval(
|
||||
f"INSERT INTO {arch} ({arch_cols}, invalidation_reason, invalidated_at, entity_ids, causal_links) "
|
||||
f"SELECT {arch_cols}, $2, now(), $3::uuid[], $5::jsonb FROM {mu} WHERE id = $1 AND bank_id = $4 "
|
||||
f"RETURNING id",
|
||||
str(unit_id),
|
||||
reason,
|
||||
entity_ids,
|
||||
bank_id,
|
||||
json.dumps([descriptor.as_json_dict() for descriptor in causal_links]),
|
||||
)
|
||||
if inserted is None:
|
||||
return False
|
||||
# The cascade prunes `unit_entities` and `memory_links` with the row.
|
||||
await conn.execute(f"DELETE FROM {mu} WHERE id = $1 AND bank_id = $2", str(unit_id), bank_id)
|
||||
return True
|
||||
|
||||
|
||||
async def set_invalidation_reason(*, conn, fq_table, bank_id: str, unit_id: str, reason: str | None) -> None:
|
||||
await conn.execute(
|
||||
f"UPDATE {fq_table('invalidated_memory_units')} SET invalidation_reason = $3 WHERE id = $1 AND bank_id = $2",
|
||||
str(unit_id),
|
||||
bank_id,
|
||||
reason,
|
||||
)
|
||||
|
||||
|
||||
async def restore_memory(*, conn, fq_table, bank_id: str, unit_id: str) -> StoredMemory | None:
|
||||
mu = fq_table("memory_units")
|
||||
arch = fq_table("invalidated_memory_units")
|
||||
ue = fq_table("unit_entities")
|
||||
ent = fq_table("entities")
|
||||
arch_cols = await _archive_columns(conn, fq_table)
|
||||
|
||||
arch_row = await conn.fetchrow(
|
||||
f"SELECT {_ARCHIVE_SELECT} FROM {arch} WHERE id = $1 AND bank_id = $2", str(unit_id), bank_id
|
||||
)
|
||||
if arch_row is None:
|
||||
return None
|
||||
|
||||
# Move the row back. The archive omits embedding/search_vector, so both default
|
||||
# to NULL here; search_vector is rebuilt now, the embedding by the caller.
|
||||
await conn.execute(
|
||||
f"INSERT INTO {mu} ({arch_cols}) SELECT {arch_cols} FROM {arch} WHERE id = $1 AND bank_id = $2",
|
||||
str(unit_id),
|
||||
bank_id,
|
||||
)
|
||||
# Rebuild search_vector with the *current* backend, so a backend change while
|
||||
# the fact sat archived cannot leave a stale/wrong-type vector (#2503). None
|
||||
# means the backend indexes base columns directly and leaves it empty.
|
||||
from ...db.ops_postgresql import pg_search_vector_expr
|
||||
|
||||
sv_expr = pg_search_vector_expr(get_config())
|
||||
if sv_expr is not None:
|
||||
await conn.execute(
|
||||
f"UPDATE {mu} SET search_vector = {sv_expr} WHERE id = $1 AND bank_id = $2", str(unit_id), bank_id
|
||||
)
|
||||
# Re-consolidate from scratch; links are rebuilt by graph maintenance.
|
||||
await conn.execute(
|
||||
f"UPDATE {mu} SET consolidated_at = NULL, consolidation_failed_at = NULL, updated_at = now() "
|
||||
f"WHERE id = $1 AND bank_id = $2",
|
||||
str(unit_id),
|
||||
bank_id,
|
||||
)
|
||||
# Restore the entity postings for entities that still exist — some may have
|
||||
# been swept as orphans while the memory was archived.
|
||||
if arch_row["entity_ids"]:
|
||||
await conn.execute(
|
||||
f"INSERT INTO {ue} (unit_id, entity_id) "
|
||||
f"SELECT $1, eid FROM unnest($2::uuid[]) AS eid "
|
||||
f"WHERE EXISTS (SELECT 1 FROM {ent} e WHERE e.id = eid AND e.bank_id = $3) "
|
||||
f"ON CONFLICT DO NOTHING",
|
||||
str(unit_id),
|
||||
arch_row["entity_ids"],
|
||||
bank_id,
|
||||
)
|
||||
# Rematerialize the causal edges parked at invalidation (#2864). Edges whose peer is still
|
||||
# archived or permanently deleted are skipped — the peer keeps its own copy and recreates the
|
||||
# edge when it reverts, so the restore is order-independent and idempotent.
|
||||
from ...retain.link_utils import rematerialize_causal_links
|
||||
from .graph import _ops_for
|
||||
|
||||
causal_json = await conn.fetchval(
|
||||
f"SELECT causal_links FROM {arch} WHERE id = $1 AND bank_id = $2", str(unit_id), bank_id
|
||||
)
|
||||
if causal_json:
|
||||
await rematerialize_causal_links(conn, bank_id, conn.parse_json(causal_json) or [], ops=_ops_for(conn))
|
||||
# Invalidation cascaded away this unit's derived outgoing links; queue it so graph maintenance
|
||||
# rebuilds them (the drain only touches queued units — it never scans for missing adjacency).
|
||||
await _ops_for(conn).enqueue_graph_maintenance(
|
||||
conn, fq_table("graph_maintenance_queue"), bank_id, [uuid.UUID(str(unit_id))]
|
||||
)
|
||||
await conn.execute(f"DELETE FROM {arch} WHERE id = $1 AND bank_id = $2", str(unit_id), bank_id)
|
||||
return _archived_stored(arch_row)
|
||||
|
||||
|
||||
async def set_memory_embedding(*, conn, fq_table, bank_id: str, unit_id: str, embedding) -> None:
|
||||
await conn.execute(
|
||||
f"UPDATE {fq_table('memory_units')} SET embedding = $3::vector WHERE id = $1 AND bank_id = $2",
|
||||
str(unit_id),
|
||||
bank_id,
|
||||
embedding,
|
||||
)
|
||||
|
||||
|
||||
async def clear_unit_entities(*, conn, fq_table, bank_id: str, unit_id: str) -> None:
|
||||
await conn.execute(f"DELETE FROM {fq_table('unit_entities')} WHERE unit_id = $1", str(unit_id))
|
||||
|
||||
|
||||
async def apply_edit(
|
||||
*,
|
||||
conn,
|
||||
fq_table,
|
||||
bank_id: str,
|
||||
unit_id: str,
|
||||
text: str,
|
||||
context: str | None,
|
||||
fact_type: str,
|
||||
occurred_start,
|
||||
occurred_end,
|
||||
event_date,
|
||||
mentioned_at,
|
||||
entity_ids: list[str] | None,
|
||||
) -> None:
|
||||
# `entity_ids` and `mentioned_at` are unused here: the entity postings are
|
||||
# re-linked into `unit_entities` by the caller, and an edit does not move the
|
||||
# mention time. Both are on the signature for a store that carries entities on
|
||||
# the memory and rebuilds it wholesale.
|
||||
from ...causal_links import CAUSAL_LINK_TYPES
|
||||
from ...db.ops_postgresql import pg_search_vector_expr
|
||||
|
||||
mu = fq_table("memory_units")
|
||||
ml = fq_table("memory_links")
|
||||
# The caller enqueues the relink victims (and the edited unit itself, via
|
||||
# ``include_affected_units``) before invoking this — one combined queue insert keeps the
|
||||
# graph-maintenance queue's lock ordering intact.
|
||||
# Keep the stored text-search vector in sync with the edited text/context.
|
||||
# Reference the bind parameters, not the columns: PostgreSQL evaluates the
|
||||
# UPDATE's RHS before the sibling SET assignments land, so a column reference
|
||||
# would see the pre-edit values.
|
||||
sv_expr = pg_search_vector_expr(get_config(), text_col="$3", context_col="$4")
|
||||
sv_clause = f", search_vector = {sv_expr}" if sv_expr else ""
|
||||
await conn.execute(
|
||||
f"""
|
||||
UPDATE {mu}
|
||||
SET text = $3, context = $4, fact_type = $5, occurred_start = $6, occurred_end = $7,
|
||||
event_date = $8, consolidated_at = NULL, consolidation_failed_at = NULL,
|
||||
edited_at = now(), updated_at = now(){sv_clause}
|
||||
WHERE id = $1 AND bank_id = $2
|
||||
""",
|
||||
str(unit_id),
|
||||
bank_id,
|
||||
text,
|
||||
context,
|
||||
fact_type,
|
||||
occurred_start,
|
||||
occurred_end,
|
||||
event_date,
|
||||
)
|
||||
# Drop only the DERIVED links — graph maintenance recomputes temporal/semantic. Causal edges
|
||||
# are retain-time extraction output that nothing recreates, so an edit preserves them (#2864).
|
||||
await conn.execute(
|
||||
f"DELETE FROM {ml} WHERE (from_unit_id = $1 OR to_unit_id = $1) AND NOT (link_type = ANY($2::text[]))",
|
||||
str(unit_id),
|
||||
list(CAUSAL_LINK_TYPES),
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"apply_edit",
|
||||
"clear_unit_entities",
|
||||
"delete_document",
|
||||
"delete_observations",
|
||||
"delete_stale_observations",
|
||||
"get_archived_memory",
|
||||
"insert_facts",
|
||||
"invalidate_memory",
|
||||
"observations_for_sources",
|
||||
"restore_memory",
|
||||
"set_invalidation_reason",
|
||||
"set_memory_embedding",
|
||||
]
|
||||
@@ -0,0 +1,509 @@
|
||||
"""The default memories store: Postgres holds the memories and the links.
|
||||
|
||||
This is the behaviour Hindsight has always had, stated as an implementation of
|
||||
:class:`~hindsight_api.engine.memories.base.MemoriesExtension` rather than as the
|
||||
absence of one. Rows go in `memory_units`, the joins around it are `memory_links`
|
||||
and `unit_entities`, and every read is SQL — writing a row *is* indexing it, so
|
||||
:meth:`index_facts` has nothing left to do.
|
||||
|
||||
The class is deliberately thin. Each method delegates to a plain function in
|
||||
:mod:`hindsight_api.engine.memories.pg`, split by what calls it — curation,
|
||||
graph, reads, writes — so a change to one area is a change to one file, and the
|
||||
SQL is grouped by concern rather than piled behind a class. The two retrieval
|
||||
arms delegate further out still, to the query functions that already own them in
|
||||
:mod:`hindsight_api.engine.search.retrieval`.
|
||||
|
||||
Keeping this as an explicit store (rather than an ``if store is None`` branch at
|
||||
each call site) means the default path is the one the whole test suite exercises,
|
||||
and a second implementation cannot change it by accident.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime
|
||||
from typing import Any
|
||||
|
||||
from .base import DeletePredicate, MemoriesExtension, MemoryPatch, ScanPage, StoredMemory
|
||||
from .pg import counts, curation, graph, reads, writes
|
||||
|
||||
|
||||
class PostgresMemories(MemoriesExtension):
|
||||
"""Memories in `memory_units`, links in `memory_links` / `unit_entities`."""
|
||||
|
||||
name = "postgres"
|
||||
|
||||
# ------------------------------------------------------------------ writes
|
||||
|
||||
async def insert_facts(
|
||||
self,
|
||||
*,
|
||||
conn,
|
||||
ops,
|
||||
bank_id: str,
|
||||
facts: list,
|
||||
document_id: str | None = None,
|
||||
defer_index: bool = False,
|
||||
txn=None,
|
||||
) -> list[str]:
|
||||
# `txn` is ignored: Postgres memories live in the caller's own transaction, so the
|
||||
# write is already atomic with it — there is no separate store to hold invisible.
|
||||
# `defer_index` is meaningless here: the INSERT that returns the ids is
|
||||
# also what indexes the facts, so there is nothing to defer.
|
||||
return await writes.insert_facts(conn=conn, ops=ops, bank_id=bank_id, facts=facts, document_id=document_id)
|
||||
|
||||
async def delete_facts(self, bank_id: str, unit_ids: list[str], *, txn=None) -> None:
|
||||
"""No-op: the caller's `memory_units` DELETE (or its FK cascade) removed them."""
|
||||
|
||||
async def delete_where(self, bank_id: str, predicate: DeletePredicate, txn=None) -> int:
|
||||
"""No-op: predicate deletes are issued as SQL by the caller that owns the transaction."""
|
||||
return 0
|
||||
|
||||
async def delete_document(self, *, conn, fq_table, bank_id: str, document_id: str, txn=None) -> None:
|
||||
# `txn` ignored: Postgres memories are covered by the caller's own transaction.
|
||||
await writes.delete_document(conn=conn, fq_table=fq_table, bank_id=bank_id, document_id=document_id)
|
||||
|
||||
async def drop_bank_storage(self, bank_id: str) -> None:
|
||||
"""No-op: deleting the bank cascades to its memories."""
|
||||
|
||||
async def delete_observations(self, *, conn, fq_table, bank_id: str, txn=None) -> None:
|
||||
await writes.delete_observations(conn=conn, fq_table=fq_table, bank_id=bank_id)
|
||||
|
||||
async def update_memories(self, bank_id: str, patches: list[MemoryPatch], txn=None) -> None:
|
||||
"""No-op: the caller's UPDATE already wrote the row it holds open."""
|
||||
|
||||
# ------------------------------------------------------------------ recall arms
|
||||
|
||||
async def search(
|
||||
self,
|
||||
*,
|
||||
conn,
|
||||
bank_id: str,
|
||||
fact_types: list[str],
|
||||
query_embedding: str,
|
||||
query_text: str,
|
||||
limit: int,
|
||||
tags: list[str] | None = None,
|
||||
tags_match: str = "any",
|
||||
tag_groups: list | None = None,
|
||||
created_after: datetime | None = None,
|
||||
created_before: datetime | None = None,
|
||||
min_semantic: float | None = None,
|
||||
min_keyword: float | None = None,
|
||||
graph_seed_min_similarity: float | None = None,
|
||||
) -> "dict[str, SemanticBm25Result]":
|
||||
# Imported here: retrieval imports this package, so a module-level import
|
||||
# would close the cycle.
|
||||
from ..search.retrieval import retrieve_semantic_bm25_combined_sql
|
||||
|
||||
return await retrieve_semantic_bm25_combined_sql(
|
||||
conn,
|
||||
query_embedding,
|
||||
query_text,
|
||||
bank_id,
|
||||
fact_types,
|
||||
limit,
|
||||
tags=tags,
|
||||
tags_match=tags_match,
|
||||
tag_groups=tag_groups,
|
||||
created_after=created_after,
|
||||
created_before=created_before,
|
||||
min_semantic=min_semantic,
|
||||
min_keyword=min_keyword,
|
||||
graph_seed_min_similarity=graph_seed_min_similarity,
|
||||
)
|
||||
|
||||
async def temporal_search(
|
||||
self,
|
||||
*,
|
||||
conn,
|
||||
bank_id: str,
|
||||
fact_types: list[str],
|
||||
query_embedding: str,
|
||||
start_date: datetime,
|
||||
end_date: datetime,
|
||||
limit: int,
|
||||
semantic_threshold: float = 0.1,
|
||||
tags: list[str] | None = None,
|
||||
tags_match: str = "any",
|
||||
tag_groups: list | None = None,
|
||||
created_after: datetime | None = None,
|
||||
created_before: datetime | None = None,
|
||||
) -> dict[str, list]:
|
||||
from ..search.retrieval import retrieve_temporal_combined_sql
|
||||
|
||||
return await retrieve_temporal_combined_sql(
|
||||
conn,
|
||||
query_embedding,
|
||||
bank_id,
|
||||
fact_types,
|
||||
start_date,
|
||||
end_date,
|
||||
limit,
|
||||
semantic_threshold=semantic_threshold,
|
||||
tags=tags,
|
||||
tags_match=tags_match,
|
||||
tag_groups=tag_groups,
|
||||
created_after=created_after,
|
||||
created_before=created_before,
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------ addressed reads
|
||||
|
||||
async def get_memories(self, *, conn, fq_table, bank_id: str, unit_ids: list[str]) -> list[StoredMemory]:
|
||||
return await reads.get_memories(conn=conn, fq_table=fq_table, bank_id=bank_id, unit_ids=unit_ids)
|
||||
|
||||
async def scan_memories(
|
||||
self,
|
||||
*,
|
||||
conn,
|
||||
fq_table,
|
||||
bank_id: str,
|
||||
fact_types: list[str] | None = None,
|
||||
limit: int = 100,
|
||||
page_token: str = "",
|
||||
tags: list[str] | None = None,
|
||||
tags_match: str = "any",
|
||||
tag_groups: list | None = None,
|
||||
document_id: str | None = None,
|
||||
metadata_equals: dict[str, str] | None = None,
|
||||
skip: int = 0,
|
||||
include_edges: bool = False,
|
||||
) -> ScanPage:
|
||||
return await reads.scan_memories(
|
||||
conn=conn,
|
||||
fq_table=fq_table,
|
||||
bank_id=bank_id,
|
||||
fact_types=fact_types,
|
||||
limit=limit,
|
||||
page_token=page_token,
|
||||
tags=tags,
|
||||
tags_match=tags_match,
|
||||
tag_groups=tag_groups,
|
||||
document_id=document_id,
|
||||
metadata_equals=metadata_equals,
|
||||
skip=skip,
|
||||
include_edges=include_edges,
|
||||
)
|
||||
|
||||
async def count_memories(self, *, conn, fq_table, bank_id: str) -> dict[str, int]:
|
||||
return await reads.count_memories(conn=conn, fq_table=fq_table, bank_id=bank_id)
|
||||
|
||||
async def list_tags(
|
||||
self,
|
||||
*,
|
||||
conn,
|
||||
fq_table,
|
||||
bank_id: str,
|
||||
pattern: str | None = None,
|
||||
limit: int = 100,
|
||||
offset: int = 0,
|
||||
) -> dict[str, Any]:
|
||||
return await reads.list_tags(
|
||||
conn=conn, fq_table=fq_table, bank_id=bank_id, pattern=pattern, limit=limit, offset=offset
|
||||
)
|
||||
|
||||
async def find_unconsolidated(
|
||||
self,
|
||||
*,
|
||||
conn,
|
||||
fq_table,
|
||||
bank_id: str,
|
||||
fact_types: list[str],
|
||||
limit: int,
|
||||
scope_tags: list[str] | None = None,
|
||||
) -> list[StoredMemory]:
|
||||
return await reads.find_unconsolidated(
|
||||
conn=conn,
|
||||
fq_table=fq_table,
|
||||
bank_id=bank_id,
|
||||
fact_types=fact_types,
|
||||
limit=limit,
|
||||
scope_tags=scope_tags,
|
||||
)
|
||||
|
||||
async def count_unconsolidated(
|
||||
self,
|
||||
*,
|
||||
conn,
|
||||
fq_table,
|
||||
bank_id: str,
|
||||
fact_types: list[str],
|
||||
scopes: list[list[str] | None],
|
||||
limit: int,
|
||||
) -> int:
|
||||
return await reads.count_unconsolidated(
|
||||
conn=conn, fq_table=fq_table, bank_id=bank_id, fact_types=fact_types, scopes=scopes, limit=limit
|
||||
)
|
||||
|
||||
async def mark_consolidated(
|
||||
self,
|
||||
*,
|
||||
conn,
|
||||
fq_table,
|
||||
bank_id: str,
|
||||
unit_ids: list[str],
|
||||
when: datetime | None,
|
||||
failed: bool = False,
|
||||
txn=None,
|
||||
) -> None:
|
||||
await reads.mark_consolidated(
|
||||
conn=conn, fq_table=fq_table, bank_id=bank_id, unit_ids=unit_ids, when=when, failed=failed
|
||||
)
|
||||
|
||||
async def any_memory_updated_since(
|
||||
self,
|
||||
*,
|
||||
conn,
|
||||
fq_table,
|
||||
bank_id: str,
|
||||
since: datetime,
|
||||
fact_types: list[str] | None = None,
|
||||
tags: list[str] | None = None,
|
||||
tags_match: str = "any",
|
||||
tag_groups: list | None = None,
|
||||
) -> bool:
|
||||
return await reads.any_memory_updated_since(
|
||||
conn=conn,
|
||||
fq_table=fq_table,
|
||||
bank_id=bank_id,
|
||||
since=since,
|
||||
fact_types=fact_types,
|
||||
tags=tags,
|
||||
tags_match=tags_match,
|
||||
tag_groups=tag_groups,
|
||||
)
|
||||
|
||||
# -- count surfaces --
|
||||
|
||||
async def consolidation_freshness(self, *, conn, fq_table, bank_id: str) -> dict[str, Any]:
|
||||
return await counts.consolidation_freshness(conn=conn, fq_table=fq_table, bank_id=bank_id)
|
||||
|
||||
async def document_memory_counts(self, *, conn, fq_table, bank_id: str, document_ids: list[str]) -> dict[str, int]:
|
||||
return await counts.document_memory_counts(
|
||||
conn=conn, fq_table=fq_table, bank_id=bank_id, document_ids=document_ids
|
||||
)
|
||||
|
||||
async def link_counts(self, *, conn, fq_table, bank_id: str) -> dict[str, int]:
|
||||
return await counts.link_counts(conn=conn, fq_table=fq_table, bank_id=bank_id)
|
||||
|
||||
async def memories_timeseries(
|
||||
self, *, conn, fq_table, bank_id: str, time_field: str, trunc: str, since: datetime
|
||||
) -> list[dict[str, Any]]:
|
||||
return await counts.memories_timeseries(
|
||||
conn=conn, fq_table=fq_table, bank_id=bank_id, time_field=time_field, trunc=trunc, since=since
|
||||
)
|
||||
|
||||
async def observation_scope_counts(self, *, conn, fq_table, bank_id: str) -> list[dict[str, Any]]:
|
||||
return await counts.observation_scope_counts(conn=conn, fq_table=fq_table, bank_id=bank_id)
|
||||
|
||||
# ------------------------------------------------------------------ observations
|
||||
|
||||
async def upsert_observation(self, *, conn, bank_id: str, record, txn=None) -> None:
|
||||
"""No-op: the observation was written as a `memory_units` row by the caller."""
|
||||
|
||||
async def observations_for_sources(
|
||||
self, *, conn, ops, fq_table, bank_id: str, unit_ids: list[str]
|
||||
) -> list[StoredMemory]:
|
||||
return await writes.observations_for_sources(
|
||||
conn=conn, ops=ops, fq_table=fq_table, bank_id=bank_id, unit_ids=unit_ids
|
||||
)
|
||||
|
||||
async def delete_stale_observations(self, *, conn, ops, fq_table, bank_id: str, fact_ids: list) -> int:
|
||||
return await writes.delete_stale_observations(
|
||||
conn=conn, ops=ops, fq_table=fq_table, bank_id=bank_id, fact_ids=fact_ids
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------ curation reads
|
||||
|
||||
async def list_memory_units(
|
||||
self,
|
||||
*,
|
||||
conn,
|
||||
ops,
|
||||
fq_table,
|
||||
bank_id: str,
|
||||
fact_type: str | None = None,
|
||||
search_query: str | None = None,
|
||||
consolidation_state: str | None = None,
|
||||
state: str | None = None,
|
||||
document_id: str | None = None,
|
||||
entity_id: str | None = None,
|
||||
tags: list[str] | None = None,
|
||||
tags_match: str = "any",
|
||||
created_before: datetime | None = None,
|
||||
limit: int = 100,
|
||||
offset: int = 0,
|
||||
) -> dict[str, Any]:
|
||||
return await curation.list_memory_units(
|
||||
conn=conn,
|
||||
ops=ops,
|
||||
fq_table=fq_table,
|
||||
bank_id=bank_id,
|
||||
fact_type=fact_type,
|
||||
search_query=search_query,
|
||||
consolidation_state=consolidation_state,
|
||||
state=state,
|
||||
document_id=document_id,
|
||||
entity_id=entity_id,
|
||||
tags=tags,
|
||||
tags_match=tags_match,
|
||||
created_before=created_before,
|
||||
limit=limit,
|
||||
offset=offset,
|
||||
)
|
||||
|
||||
async def get_memory_unit(self, *, conn, ops, fq_table, bank_id: str, unit_id: str) -> dict[str, Any] | None:
|
||||
return await curation.get_memory_unit(conn=conn, ops=ops, fq_table=fq_table, bank_id=bank_id, unit_id=unit_id)
|
||||
|
||||
# -- curation archive --
|
||||
|
||||
async def get_archived_memory(self, *, conn, fq_table, bank_id: str, unit_id: str) -> StoredMemory | None:
|
||||
return await writes.get_archived_memory(conn=conn, fq_table=fq_table, bank_id=bank_id, unit_id=unit_id)
|
||||
|
||||
async def invalidate_memory(
|
||||
self, *, conn, fq_table, bank_id: str, unit_id: str, reason: str | None, txn=None
|
||||
) -> bool:
|
||||
return await writes.invalidate_memory(
|
||||
conn=conn, fq_table=fq_table, bank_id=bank_id, unit_id=unit_id, reason=reason
|
||||
)
|
||||
|
||||
async def set_invalidation_reason(self, *, conn, fq_table, bank_id: str, unit_id: str, reason: str | None) -> None:
|
||||
await writes.set_invalidation_reason(
|
||||
conn=conn, fq_table=fq_table, bank_id=bank_id, unit_id=unit_id, reason=reason
|
||||
)
|
||||
|
||||
async def restore_memory(self, *, conn, fq_table, bank_id: str, unit_id: str, txn=None) -> StoredMemory | None:
|
||||
return await writes.restore_memory(conn=conn, fq_table=fq_table, bank_id=bank_id, unit_id=unit_id)
|
||||
|
||||
async def set_memory_embedding(self, *, conn, fq_table, bank_id: str, unit_id: str, embedding, txn=None) -> None:
|
||||
await writes.set_memory_embedding(
|
||||
conn=conn, fq_table=fq_table, bank_id=bank_id, unit_id=unit_id, embedding=embedding
|
||||
)
|
||||
|
||||
async def clear_unit_entities(self, *, conn, fq_table, bank_id: str, unit_id: str) -> None:
|
||||
await writes.clear_unit_entities(conn=conn, fq_table=fq_table, bank_id=bank_id, unit_id=unit_id)
|
||||
|
||||
async def apply_edit(
|
||||
self,
|
||||
*,
|
||||
conn,
|
||||
fq_table,
|
||||
bank_id: str,
|
||||
unit_id: str,
|
||||
text: str,
|
||||
context: str | None,
|
||||
fact_type: str,
|
||||
occurred_start,
|
||||
occurred_end,
|
||||
event_date,
|
||||
mentioned_at,
|
||||
entity_ids: list[str] | None,
|
||||
txn=None,
|
||||
) -> None:
|
||||
await writes.apply_edit(
|
||||
conn=conn,
|
||||
fq_table=fq_table,
|
||||
bank_id=bank_id,
|
||||
unit_id=unit_id,
|
||||
text=text,
|
||||
context=context,
|
||||
fact_type=fact_type,
|
||||
occurred_start=occurred_start,
|
||||
occurred_end=occurred_end,
|
||||
event_date=event_date,
|
||||
mentioned_at=mentioned_at,
|
||||
entity_ids=entity_ids,
|
||||
)
|
||||
|
||||
async def list_entities(
|
||||
self,
|
||||
*,
|
||||
conn,
|
||||
fq_table,
|
||||
bank_id: str,
|
||||
search: str | None = None,
|
||||
limit: int = 100,
|
||||
offset: int = 0,
|
||||
) -> dict[str, Any]:
|
||||
return await curation.list_entities(
|
||||
conn=conn, fq_table=fq_table, bank_id=bank_id, search=search, limit=limit, offset=offset
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------ graph
|
||||
|
||||
async def graph_units(
|
||||
self,
|
||||
*,
|
||||
conn,
|
||||
fq_table,
|
||||
bank_id: str,
|
||||
fact_type: str | None = None,
|
||||
search_query: str | None = None,
|
||||
document_id: str | None = None,
|
||||
chunk_id: str | None = None,
|
||||
tags: list[str] | None = None,
|
||||
tags_match: str = "all_strict",
|
||||
limit: int = 1000,
|
||||
) -> dict[str, Any]:
|
||||
return await graph.graph_units(
|
||||
conn=conn,
|
||||
fq_table=fq_table,
|
||||
bank_id=bank_id,
|
||||
fact_type=fact_type,
|
||||
search_query=search_query,
|
||||
document_id=document_id,
|
||||
chunk_id=chunk_id,
|
||||
tags=tags,
|
||||
tags_match=tags_match,
|
||||
limit=limit,
|
||||
)
|
||||
|
||||
async def graph_entity_rows(self, *, conn, fq_table, bank_id: str, unit_ids: list[str]) -> list[dict[str, Any]]:
|
||||
return await graph.graph_entity_rows(conn=conn, fq_table=fq_table, bank_id=bank_id, unit_ids=unit_ids)
|
||||
|
||||
async def graph_direct_links(self, *, conn, fq_table, bank_id: str, unit_ids: list[str]) -> list[dict[str, Any]]:
|
||||
return await graph.graph_direct_links(conn=conn, fq_table=fq_table, bank_id=bank_id, unit_ids=unit_ids)
|
||||
|
||||
async def entity_memory_counts(
|
||||
self, *, conn, fq_table, bank_id: str, entity_ids: list[str] | None = None
|
||||
) -> dict[str, int]:
|
||||
return await graph.entity_memory_counts(conn=conn, fq_table=fq_table, bank_id=bank_id, entity_ids=entity_ids)
|
||||
|
||||
async def entities_for_units(self, *, conn, fq_table, bank_id: str, unit_ids: list[str]) -> dict[str, list[str]]:
|
||||
return await graph.entities_for_units(conn=conn, fq_table=fq_table, bank_id=bank_id, unit_ids=unit_ids)
|
||||
|
||||
async def entity_map_for_units(
|
||||
self, *, conn, fq_table, bank_id: str, unit_ids: list[str]
|
||||
) -> dict[str, list[dict[str, str]]]:
|
||||
return await graph.entity_map_for_units(conn=conn, fq_table=fq_table, bank_id=bank_id, unit_ids=unit_ids)
|
||||
|
||||
# ------------------------------------------------------------------ maintenance
|
||||
|
||||
async def record_unit_entities(
|
||||
self, *, conn, ops, fq_table, bank_id: str | None = None, unit_ids: list[Any], entity_ids: list[Any]
|
||||
) -> None:
|
||||
# The join is keyed by global unit id, so bank_id is not needed here.
|
||||
await ops.bulk_insert_unit_entities(conn, fq_table("unit_entities"), unit_ids, entity_ids)
|
||||
|
||||
async def enqueue_relink_victims(
|
||||
self, *, conn, fq_table, bank_id: str, affected_unit_ids: list, include_affected_units: bool = False
|
||||
) -> int:
|
||||
return await graph.enqueue_relink_victims(
|
||||
conn=conn,
|
||||
fq_table=fq_table,
|
||||
bank_id=bank_id,
|
||||
affected_unit_ids=affected_unit_ids,
|
||||
include_affected_units=include_affected_units,
|
||||
)
|
||||
|
||||
async def relink_pass(self, *, backend, fq_table, bank_id: str, config) -> dict:
|
||||
return await graph.relink_pass(backend=backend, fq_table=fq_table, bank_id=bank_id, config=config)
|
||||
|
||||
async def prune_orphan_entities(self, *, conn, fq_table, bank_id: str) -> int:
|
||||
return await graph.prune_orphan_entities(conn=conn, fq_table=fq_table, bank_id=bank_id)
|
||||
|
||||
async def prune_stale_cooccurrences(self, *, conn, fq_table, bank_id: str) -> int:
|
||||
return await graph.prune_stale_cooccurrences(conn=conn, fq_table=fq_table, bank_id=bank_id)
|
||||
|
||||
|
||||
__all__ = ["PostgresMemories"]
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,242 @@
|
||||
"""Models describing what a mental model refresh did.
|
||||
|
||||
A refresh resolves a scope, picks full-vs-delta, runs reflect over a bounded
|
||||
snapshot, and (in delta mode) applies structured operations to the existing
|
||||
document. Every one of those steps can quietly produce a document that isn't
|
||||
what the user expected, and until now the reasoning behind each only ever
|
||||
reached a log line.
|
||||
|
||||
These models carry that reasoning out to callers, so both the dry run (preview,
|
||||
nothing persisted) and ``trigger.keep_trace`` (recorded on every real refresh,
|
||||
including the cron- and consolidation-driven ones no human is watching) can
|
||||
report it.
|
||||
|
||||
Kept out of ``response_models`` on purpose: these reference the tag-group types
|
||||
from ``search.tags``, and ``response_models`` is imported early enough in the
|
||||
engine's import graph that pulling the search package in from there is a cycle.
|
||||
"""
|
||||
|
||||
from datetime import datetime
|
||||
from typing import Any, Literal
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
|
||||
from .response_models import LLMCallTrace, TokenUsage
|
||||
from .search.tags import TagGroup, TagsMatch
|
||||
|
||||
RefreshMode = Literal["full", "delta"]
|
||||
|
||||
ModeFallbackReason = Literal[
|
||||
"no_baseline_content",
|
||||
"source_query_changed",
|
||||
"structured_doc_unreadable",
|
||||
"delta_ops_failed",
|
||||
"delta_ops_all_skipped",
|
||||
]
|
||||
|
||||
RefreshOutcome = Literal[
|
||||
"content_written",
|
||||
"content_preserved_no_new_facts",
|
||||
"refresh_failed_empty_candidate",
|
||||
"refresh_failed_delta_not_applied",
|
||||
]
|
||||
|
||||
|
||||
class MentalModelRefreshScope(BaseModel):
|
||||
"""The memory scope a refresh actually resolved to.
|
||||
|
||||
A model's stored ``tags`` are not what filters memories — ``tags_match``
|
||||
defaults to ``all_strict`` when tags are present, and ``tag_groups``
|
||||
override flat tags entirely. This reports the resolved result.
|
||||
"""
|
||||
|
||||
tags: list[str] | None = Field(default=None, description="Flat tags used to filter memories (null when unused).")
|
||||
tags_match: TagsMatch = Field(description="Resolved tag match mode.")
|
||||
tag_groups: list[TagGroup] | None = Field(
|
||||
default=None, description="Compound tag expressions used instead of flat tags, when set."
|
||||
)
|
||||
fact_types: list[str] | None = Field(default=None, description="Fact types retrieved (null means all).")
|
||||
exclude_mental_models: bool = Field(description="Whether other mental models were excluded from the reflect loop.")
|
||||
exclude_mental_model_ids: list[str] = Field(
|
||||
default_factory=list, description="Mental models excluded by ID (always includes the model being refreshed)."
|
||||
)
|
||||
|
||||
|
||||
class MentalModelRefreshWindow(BaseModel):
|
||||
"""The time window a refresh read memories from."""
|
||||
|
||||
created_after: datetime | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Lower bound on memory creation time. Set only in delta mode, where it is the model's "
|
||||
"last_refreshed_at — so a delta refresh only sees memories newer than the last one."
|
||||
),
|
||||
)
|
||||
created_before: datetime = Field(
|
||||
description=(
|
||||
"Database-time snapshot bounding the refresh. Memories committed after this are not read, "
|
||||
"so they stay newer than the persisted watermark and are caught by the next refresh."
|
||||
)
|
||||
)
|
||||
watermark: datetime | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"The last_refreshed_at a real refresh would persist: the newest in-scope memory visible at "
|
||||
"the snapshot, not now(). Null means no in-scope memory was visible."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class MentalModelFactCounts(BaseModel):
|
||||
"""Facts the refresh saw, keyed by fact type.
|
||||
|
||||
``retrieved`` and ``used`` diverging is the single most common cause of a
|
||||
disappointing refresh: recall found plenty, but the reflect agent declared
|
||||
none of it relevant to the topic, so none of it reached the document.
|
||||
"""
|
||||
|
||||
retrieved: dict[str, int] = Field(
|
||||
default_factory=dict, description="Facts the reflect agent's tool calls returned, by fact type."
|
||||
)
|
||||
used: dict[str, int] = Field(
|
||||
default_factory=dict, description="Facts the agent declared it actually based the answer on, by fact type."
|
||||
)
|
||||
|
||||
|
||||
class MentalModelDeltaOperations(BaseModel):
|
||||
"""Structured operations a delta refresh emitted against the existing document."""
|
||||
|
||||
applied: list[dict[str, Any]] = Field(
|
||||
default_factory=list, description="Operations applied to the document, in order."
|
||||
)
|
||||
skipped: list[dict[str, Any]] = Field(
|
||||
default_factory=list, description="Operations dropped as invalid, each with a reason."
|
||||
)
|
||||
|
||||
|
||||
class MentalModelTraceToolCall(BaseModel):
|
||||
"""One reflect tool call made during a refresh.
|
||||
|
||||
``output`` is carried only by the dry run, which persists nothing. The trace
|
||||
stored on the model row keeps ``result_count`` instead: it is re-read on every
|
||||
fetch, so embedding full recall payloads there would bloat the row without
|
||||
bound. Raw prompts and responses are available separately via LLM request
|
||||
tracing.
|
||||
"""
|
||||
|
||||
tool: str = Field(description="Tool name: recall, search_observations, get_mental_model, expand, …")
|
||||
reason: str | None = Field(default=None, description="The agent's stated reason for the call.")
|
||||
input: dict[str, Any] = Field(default_factory=dict, description="Tool input parameters.")
|
||||
output: dict[str, Any] | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"What the tool returned. Present on a dry run, which stores nothing; omitted from the "
|
||||
"trace persisted by a real refresh to keep that row bounded."
|
||||
),
|
||||
)
|
||||
updated_at: datetime | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"The refresh window's lower bound as given to this call — the delta watermark. Named "
|
||||
"for what it actually filters: the predicate is on the memory's updated_at, so a "
|
||||
"memory merely touched since the last refresh qualifies. Null means the tool applies "
|
||||
"no time bound at all, so its results are not limited to the window (mental-model "
|
||||
"lookup and chunk expansion behave this way)."
|
||||
),
|
||||
)
|
||||
result_count: int | None = Field(default=None, description="Number of items the tool returned, when countable.")
|
||||
duration_ms: int = Field(description="Execution time in milliseconds.")
|
||||
iteration: int = Field(default=0, description="Agent loop iteration (1-based) this call belongs to.")
|
||||
|
||||
|
||||
class MentalModelRefreshTrace(BaseModel):
|
||||
"""Execution trace of a mental model refresh, recorded when trigger.keep_trace is on.
|
||||
|
||||
Deliberately shaped like reflect's trace — the calls the agent made, plus the
|
||||
refresh-specific decision — and nothing more. This is persisted on the mental
|
||||
model row and re-read on every fetch, so anything derivable from elsewhere is
|
||||
left out: the evidence lives in ``reflect_response.based_on``, and the
|
||||
resolved scope and snapshot window are reported by the dry run.
|
||||
"""
|
||||
|
||||
recorded_at: datetime | None = Field(default=None, description="When this trace was recorded.")
|
||||
effective_mode: RefreshMode = Field(description="Whether the refresh ran as full or delta.")
|
||||
mode_fallback_reason: ModeFallbackReason | None = Field(
|
||||
default=None, description="Why delta was requested but not applied, if that happened."
|
||||
)
|
||||
outcome: RefreshOutcome = Field(description="What the refresh did with the document.")
|
||||
tool_calls: list[MentalModelTraceToolCall] = Field(
|
||||
default_factory=list, description="Reflect tool calls made during the refresh."
|
||||
)
|
||||
llm_calls: list[LLMCallTrace] = Field(default_factory=list, description="LLM calls made during the refresh.")
|
||||
delta_operations: MentalModelDeltaOperations | None = Field(
|
||||
default=None, description="Structured operations emitted, in delta mode."
|
||||
)
|
||||
usage: TokenUsage | None = Field(default=None, description="Token usage across the refresh's LLM calls.")
|
||||
duration_ms: int = Field(default=0, description="Wall-clock duration of the refresh.")
|
||||
warnings: list[str] = Field(
|
||||
default_factory=list, description="Conditions worth a human's attention, in plain language."
|
||||
)
|
||||
|
||||
|
||||
class MentalModelDryRunRefreshResult(BaseModel):
|
||||
"""Preview of what a mental model refresh would do, having changed nothing.
|
||||
|
||||
Runs the real pipeline — same scope resolution, same reflect call, same
|
||||
delta operations — then reports the result instead of persisting it. The
|
||||
model's content, structured content, watermark, and last_refreshed_at are
|
||||
all left untouched, so a delta dry run is repeatable: it reads the same
|
||||
window the next real refresh would.
|
||||
"""
|
||||
|
||||
model_config = ConfigDict(
|
||||
json_schema_extra={
|
||||
"example": {
|
||||
"mental_model_id": "coding-style",
|
||||
"name": "Coding Style",
|
||||
"requested_mode": "delta",
|
||||
"effective_mode": "full",
|
||||
"mode_fallback_reason": "source_query_changed",
|
||||
"outcome": "content_written",
|
||||
"would_persist": True,
|
||||
"facts": {"retrieved": {"observation": 12}, "used": {"observation": 4}},
|
||||
"warnings": [],
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
mental_model_id: str = Field(description="The mental model previewed.")
|
||||
name: str = Field(description="Display name of the mental model.")
|
||||
requested_mode: RefreshMode = Field(description="The mode asked for (from the model's trigger, or overridden).")
|
||||
effective_mode: RefreshMode = Field(description="The mode the refresh actually ran in.")
|
||||
mode_fallback_reason: ModeFallbackReason | None = Field(
|
||||
default=None, description="Why delta was requested but not applied, if that happened."
|
||||
)
|
||||
outcome: RefreshOutcome = Field(description="What a real refresh would do with the document.")
|
||||
would_persist: bool = Field(description="Whether a real refresh would write new content.")
|
||||
scope: MentalModelRefreshScope = Field(description="The resolved memory scope.")
|
||||
window: MentalModelRefreshWindow = Field(description="The snapshot window read from.")
|
||||
facts: MentalModelFactCounts = Field(description="Facts retrieved versus actually used.")
|
||||
based_on: dict[str, list[dict[str, Any]]] = Field(
|
||||
default_factory=dict,
|
||||
description=(
|
||||
"The evidence this run would ground the document on, keyed by fact type — the same "
|
||||
"shape a refresh persists under reflect_response.based_on. Returned so a preview can "
|
||||
"show its sources without having to write them anywhere."
|
||||
),
|
||||
)
|
||||
current_content: str = Field(description="The model's content as it stands now.")
|
||||
candidate_content: str = Field(description="Raw reflect synthesis, before any delta operations.")
|
||||
preview_content: str = Field(
|
||||
description="The content a real refresh would store: the delta-edited document, or the candidate in full mode."
|
||||
)
|
||||
diff: str = Field(description="Unified diff from current_content to preview_content. Empty when identical.")
|
||||
delta_operations: MentalModelDeltaOperations | None = Field(
|
||||
default=None, description="Structured operations emitted, in delta mode."
|
||||
)
|
||||
trace: MentalModelRefreshTrace = Field(description="Execution trace of the run, always included for a dry run.")
|
||||
usage: TokenUsage = Field(default_factory=TokenUsage, description="Token usage across the run's LLM calls.")
|
||||
duration_ms: int = Field(default=0, description="Wall-clock duration of the run.")
|
||||
warnings: list[str] = Field(
|
||||
default_factory=list, description="Conditions worth a human's attention, in plain language."
|
||||
)
|
||||
@@ -19,10 +19,18 @@ class BatchRetainParentMetadata:
|
||||
total_tokens: int
|
||||
num_sub_batches: int
|
||||
is_parent: bool = True
|
||||
# Set only when the whole batch targets a single document, so the operations
|
||||
# list surfaces which document an in-flight retain is (re)writing. The
|
||||
# documents UI cross-checks this to badge rows as "updating". Multi-document
|
||||
# batches leave it None and are matched per single-document child instead.
|
||||
document_id: str | None = None
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
"""Convert to dict for JSON serialization."""
|
||||
return asdict(self)
|
||||
"""Convert to dict for JSON serialization, omitting document_id when unset."""
|
||||
data = asdict(self)
|
||||
if data.get("document_id") is None:
|
||||
data.pop("document_id", None)
|
||||
return data
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -33,10 +41,15 @@ class BatchRetainChildMetadata:
|
||||
parent_operation_id: str
|
||||
sub_batch_index: int
|
||||
total_sub_batches: int
|
||||
# Set only when this child processes a single document (see the parent's note).
|
||||
document_id: str | None = None
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
"""Convert to dict for JSON serialization."""
|
||||
return asdict(self)
|
||||
"""Convert to dict for JSON serialization, omitting document_id when unset."""
|
||||
data = asdict(self)
|
||||
if data.get("document_id") is None:
|
||||
data.pop("document_id", None)
|
||||
return data
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -142,3 +155,27 @@ class RefreshMentalModelMetadata:
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
"""Convert to dict for JSON serialization."""
|
||||
return asdict(self)
|
||||
|
||||
|
||||
@dataclass
|
||||
class RefreshMentalModelOutcomeMetadata:
|
||||
"""Machine-readable outcome metadata for a completed refresh_mental_model operation.
|
||||
|
||||
Refresh parity with RetainOutcomeMetadata (#2605): lets a monitoring layer
|
||||
distinguish "refreshed with real content" from "refreshed empty" by reading
|
||||
result_metadata alone, without a follow-up content fetch.
|
||||
"""
|
||||
|
||||
content_len: int
|
||||
populated_content: bool
|
||||
based_on_counts: dict[str, int] = field(default_factory=dict)
|
||||
# Delta operations the model emitted, as applied vs rejected. A refresh whose
|
||||
# ops are routinely rejected still completes successfully with a plausible
|
||||
# document, so the count is the only signal that some of this run's new facts
|
||||
# never reached it. Both are 0 for a full-mode refresh, which emits no ops.
|
||||
delta_ops_applied: int = 0
|
||||
delta_ops_skipped: int = 0
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
"""Convert to dict for JSON serialization."""
|
||||
return asdict(self)
|
||||
|
||||
@@ -72,6 +72,7 @@ class MarkitdownParser(FileParser):
|
||||
ocr_base_url: str | None = None,
|
||||
ocr_model: str | None = None,
|
||||
ocr_prompt: str | None = None,
|
||||
ocr_default_headers: dict[str, str] | None = None,
|
||||
):
|
||||
"""Initialize markitdown parser."""
|
||||
# Lazy import to avoid requiring markitdown for all users
|
||||
@@ -89,6 +90,7 @@ class MarkitdownParser(FileParser):
|
||||
base_url=ocr_base_url,
|
||||
model=ocr_model,
|
||||
prompt=ocr_prompt,
|
||||
default_headers=ocr_default_headers,
|
||||
)
|
||||
self._markitdown = MarkItDown(
|
||||
llm_client=ocr_options.llm_client,
|
||||
@@ -105,6 +107,7 @@ class MarkitdownParser(FileParser):
|
||||
base_url: str | None,
|
||||
model: str | None,
|
||||
prompt: str | None,
|
||||
default_headers: dict[str, str] | None,
|
||||
) -> MarkitdownOcrOptions:
|
||||
"""Build MarkItDown options for OpenAI-compatible image OCR."""
|
||||
if not model or not model.strip():
|
||||
@@ -129,8 +132,15 @@ class MarkitdownParser(FileParser):
|
||||
except ImportError as e:
|
||||
raise RuntimeError("openai package is required when Markitdown OCR is enabled.") from e
|
||||
|
||||
client_kwargs: dict[str, object] = {
|
||||
"api_key": api_key,
|
||||
"base_url": base_url.strip(),
|
||||
}
|
||||
if default_headers:
|
||||
client_kwargs["default_headers"] = default_headers
|
||||
|
||||
return MarkitdownOcrOptions(
|
||||
llm_client=OpenAI(api_key=api_key, base_url=base_url.strip()),
|
||||
llm_client=OpenAI(**client_kwargs),
|
||||
llm_model=model.strip(),
|
||||
llm_prompt=prompt or DEFAULT_FILE_PARSER_MARKITDOWN_OCR_PROMPT,
|
||||
)
|
||||
@@ -186,7 +196,11 @@ class MarkitdownParser(FileParser):
|
||||
if Path(filename).suffix.lower() not in _TEXT_EXTENSIONS:
|
||||
return None
|
||||
try:
|
||||
file_data.decode("utf-8")
|
||||
# file_data may arrive as a non-``bytes`` buffer (e.g. a memoryview or
|
||||
# a native/Rust-backed buffer object) that has no ``.decode``; coerce
|
||||
# through the buffer protocol before the UTF-8 probe. The ``tmp.write``
|
||||
# in the caller already relies only on the same buffer protocol.
|
||||
bytes(file_data).decode("utf-8")
|
||||
except UnicodeDecodeError:
|
||||
return None
|
||||
from markitdown import StreamInfo
|
||||
|
||||
@@ -15,6 +15,7 @@ from .llamacpp_llm import LlamaCppLLM
|
||||
from .mock_llm import MockLLM
|
||||
from .none_llm import NoneLLM
|
||||
from .openai_compatible_llm import OpenAICompatibleLLM
|
||||
from .openai_responses_llm import OpenAIResponsesLLM
|
||||
|
||||
__all__ = [
|
||||
"AnthropicLLM",
|
||||
@@ -28,4 +29,5 @@ __all__ = [
|
||||
"MockLLM",
|
||||
"NoneLLM",
|
||||
"OpenAICompatibleLLM",
|
||||
"OpenAIResponsesLLM",
|
||||
]
|
||||
|
||||
@@ -12,12 +12,15 @@ import asyncio
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
from typing import Any
|
||||
from contextlib import AbstractAsyncContextManager, nullcontext
|
||||
from typing import Any, Callable
|
||||
|
||||
from hindsight_api.engine.llm_interface import LLMInterface
|
||||
from hindsight_api.engine.llm_interface import LLM_TOOL_CHOICE_AUTO, LLMInterface, LLMToolChoice
|
||||
from hindsight_api.engine.llm_trace import LLMResponseUsage, stash_response_usage
|
||||
from hindsight_api.engine.providers.llm_debug import dump_request_on_4xx
|
||||
from hindsight_api.engine.response_models import LLMToolCall, LLMToolCallResult, TokenUsage
|
||||
from hindsight_api.metrics import get_metrics_collector
|
||||
from hindsight_api.worker.stage import set_stage
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -34,6 +37,43 @@ def _usage_from_anthropic_response(response: Any) -> LLMResponseUsage:
|
||||
)
|
||||
|
||||
|
||||
_EPHEMERAL_CACHE = {"type": "ephemeral"}
|
||||
|
||||
|
||||
def _cached_system_blocks(system_prompt: str) -> list[dict[str, Any]]:
|
||||
"""Render the system prompt as a block list with a cache_control marker.
|
||||
|
||||
Anthropic prompt caching is a prefix match: marking the (single) system
|
||||
block caches tools + system together. The system prompt is stable per
|
||||
scope — fact extraction reuses it across every chunk, reflect and
|
||||
consolidation keep their stable instructions there — so repeat calls read
|
||||
it at ~10% of the base input price. Markers below the model's minimum
|
||||
cacheable prefix are silently ignored (no write premium), so marking is
|
||||
safe unconditionally. This is the "inline-marker provider" strategy that
|
||||
``LLMInterface.get_or_create_cached_prefix`` documents for Anthropic.
|
||||
"""
|
||||
return [{"type": "text", "text": system_prompt, "cache_control": _EPHEMERAL_CACHE}]
|
||||
|
||||
|
||||
def _mark_last_message_for_caching(messages: list[dict[str, Any]]) -> None:
|
||||
"""Add a cache_control marker to the final content block, in place.
|
||||
|
||||
Used on the multi-turn (tool-calling) path: the reflect agent loop resends
|
||||
the entire growing conversation each iteration, so this request's
|
||||
end-marker becomes the next iteration's cache read point. Together with
|
||||
the system marker this uses 2 of the 4 allowed breakpoints.
|
||||
"""
|
||||
if not messages:
|
||||
return
|
||||
last = messages[-1]
|
||||
content = last.get("content")
|
||||
if isinstance(content, str):
|
||||
if content.strip(): # the API rejects empty text blocks
|
||||
last["content"] = [{"type": "text", "text": content, "cache_control": _EPHEMERAL_CACHE}]
|
||||
elif isinstance(content, list) and content and isinstance(content[-1], dict):
|
||||
content[-1]["cache_control"] = _EPHEMERAL_CACHE
|
||||
|
||||
|
||||
class AnthropicLLM(LLMInterface):
|
||||
"""
|
||||
LLM provider using Anthropic's Claude models.
|
||||
@@ -135,6 +175,7 @@ class AnthropicLLM(LLMInterface):
|
||||
skip_validation: bool = False,
|
||||
strict_schema: bool = False,
|
||||
return_usage: bool = False,
|
||||
attempt_context: Callable[[], AbstractAsyncContextManager[None]] | None = None,
|
||||
) -> Any:
|
||||
"""
|
||||
Make an LLM API call with retry logic.
|
||||
@@ -206,7 +247,9 @@ class AnthropicLLM(LLMInterface):
|
||||
}
|
||||
|
||||
if system_prompt:
|
||||
call_params["system"] = system_prompt
|
||||
# One-shot calls share only the system prompt with each other, so
|
||||
# that is the sole cache breakpoint on this path.
|
||||
call_params["system"] = _cached_system_blocks(system_prompt)
|
||||
|
||||
if use_forced_tool:
|
||||
# Single tool whose input_schema IS the response schema; force the model to
|
||||
@@ -223,7 +266,9 @@ class AnthropicLLM(LLMInterface):
|
||||
|
||||
for attempt in range(max_retries + 1):
|
||||
try:
|
||||
response = await self._client.messages.create(**call_params)
|
||||
async with attempt_context() if attempt_context is not None else nullcontext():
|
||||
set_stage(f"llm.{self.provider}.{scope}.attempt={attempt + 1}/{max_retries + 1}")
|
||||
response = await self._client.messages.create(**call_params)
|
||||
# Stash usage before parse/validate, which may raise locally
|
||||
# even though the provider charged for these tokens (#2387).
|
||||
stash_response_usage(_usage_from_anthropic_response(response))
|
||||
@@ -346,6 +391,9 @@ class AnthropicLLM(LLMInterface):
|
||||
logger.error(f"Anthropic auth error (HTTP {e.status_code}), not retrying: {str(e)}")
|
||||
raise
|
||||
|
||||
# Diagnostic dump (opt-in) of the exact request behind any 4xx.
|
||||
dump_request_on_4xx(scope=scope, provider=self.provider, model=self.model, err=e, request=call_params)
|
||||
|
||||
last_exception = e
|
||||
if attempt < max_retries:
|
||||
# Check if it's a rate limit or server error
|
||||
@@ -380,7 +428,8 @@ class AnthropicLLM(LLMInterface):
|
||||
max_retries: int = 5,
|
||||
initial_backoff: float = 1.0,
|
||||
max_backoff: float = 30.0,
|
||||
tool_choice: str | dict[str, Any] = "auto",
|
||||
tool_choice: LLMToolChoice = LLM_TOOL_CHOICE_AUTO,
|
||||
attempt_context: Callable[[], AbstractAsyncContextManager[None]] | None = None,
|
||||
) -> LLMToolCallResult:
|
||||
"""
|
||||
Make an LLM API call with tool/function calling support.
|
||||
@@ -450,6 +499,11 @@ class AnthropicLLM(LLMInterface):
|
||||
else:
|
||||
anthropic_messages.append({"role": role, "content": content})
|
||||
|
||||
# Multi-turn tool loop: cache the stable prefix (tools + system) via
|
||||
# the system marker, and the growing conversation via an end-marker
|
||||
# that the next iteration reads back.
|
||||
_mark_last_message_for_caching(anthropic_messages)
|
||||
|
||||
call_params: dict[str, Any] = {
|
||||
"model": self.model,
|
||||
"messages": anthropic_messages,
|
||||
@@ -457,7 +511,7 @@ class AnthropicLLM(LLMInterface):
|
||||
"max_tokens": max_completion_tokens or 4096,
|
||||
}
|
||||
if system_prompt:
|
||||
call_params["system"] = system_prompt
|
||||
call_params["system"] = _cached_system_blocks(system_prompt)
|
||||
|
||||
if self._extra_body:
|
||||
call_params["extra_body"] = self._extra_body
|
||||
@@ -465,7 +519,9 @@ class AnthropicLLM(LLMInterface):
|
||||
last_exception = None
|
||||
for attempt in range(max_retries + 1):
|
||||
try:
|
||||
response = await self._client.messages.create(**call_params)
|
||||
async with attempt_context() if attempt_context is not None else nullcontext():
|
||||
set_stage(f"llm.{self.provider}.{scope}.attempt={attempt + 1}/{max_retries + 1}")
|
||||
response = await self._client.messages.create(**call_params)
|
||||
stash_response_usage(_usage_from_anthropic_response(response))
|
||||
|
||||
# Extract content and tool calls
|
||||
@@ -533,6 +589,8 @@ class AnthropicLLM(LLMInterface):
|
||||
except (APIConnectionError, APIStatusError) as e:
|
||||
if isinstance(e, APIStatusError) and e.status_code in (401, 403):
|
||||
raise
|
||||
# Diagnostic dump (opt-in) of the exact request behind any 4xx.
|
||||
dump_request_on_4xx(scope=scope, provider=self.provider, model=self.model, err=e, request=call_params)
|
||||
last_exception = e
|
||||
if attempt < max_retries:
|
||||
await asyncio.sleep(min(initial_backoff * (2**attempt), max_backoff))
|
||||
@@ -543,7 +601,221 @@ class AnthropicLLM(LLMInterface):
|
||||
raise last_exception
|
||||
raise RuntimeError("Anthropic tool call failed")
|
||||
|
||||
# ── Message Batches API (50% token discount) ─────────────────────────────
|
||||
|
||||
_BATCH_TOOL_NAME = "structured_response"
|
||||
|
||||
async def supports_batch_api(self) -> bool:
|
||||
"""Anthropic supports batch operations via the Message Batches API."""
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def _map_batch_status(processing_status: str) -> str:
|
||||
"""Map Anthropic ``processing_status`` onto the OpenAI vocabulary.
|
||||
|
||||
The engine's poll loop breaks on "completed" and hard-fails on
|
||||
"failed"/"expired"/"cancelled"; anything else keeps polling. Anthropic
|
||||
batches only end as "ended" (per-request failures surface in the
|
||||
results, mirroring OpenAI's "completed"-with-errors semantics), so
|
||||
"ended" maps to "completed" and the non-terminal states pass through.
|
||||
"""
|
||||
return "completed" if processing_status == "ended" else processing_status
|
||||
|
||||
def _translate_batch_body(self, body: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Translate one OpenAI-shaped request body into Messages API params.
|
||||
|
||||
Mirrors the conversion rules of ``call()``: system messages fold into
|
||||
the ``system`` param; ``max_completion_tokens`` becomes ``max_tokens``
|
||||
(default 4096); ``temperature`` is dropped (the sync path never sends
|
||||
it either — current Claude models reject non-default sampling params);
|
||||
an OpenAI ``response_format`` json_schema becomes a single forced
|
||||
tool_use tool when strict (native constrained decoding, issue #1002),
|
||||
else the schema is injected into the system prompt.
|
||||
|
||||
The system prompt carries the same cache_control marker as the sync
|
||||
one-shot path (its sole breakpoint): every request in a retain batch
|
||||
shares the fact-extraction system prompt, so the first item's cache
|
||||
write serves the remaining items as best-effort reads — and the
|
||||
cache-read discount stacks with the 50% batch discount.
|
||||
"""
|
||||
system_prompt: str | None = None
|
||||
messages: list[dict[str, Any]] = []
|
||||
for msg in body.get("messages", []):
|
||||
role = msg.get("role", "user")
|
||||
content = msg.get("content", "")
|
||||
if role == "system":
|
||||
system_prompt = (system_prompt + "\n\n" + content) if system_prompt else content
|
||||
else:
|
||||
messages.append({"role": role, "content": content})
|
||||
|
||||
params: dict[str, Any] = {
|
||||
"model": body.get("model") or self.model,
|
||||
"messages": messages,
|
||||
"max_tokens": body.get("max_completion_tokens") or 4096,
|
||||
}
|
||||
|
||||
json_schema = (body.get("response_format") or {}).get("json_schema") or {}
|
||||
schema = json_schema.get("schema")
|
||||
if schema is not None:
|
||||
if json_schema.get("strict"):
|
||||
params["tools"] = [
|
||||
{
|
||||
"name": self._BATCH_TOOL_NAME,
|
||||
"description": "Return the structured response.",
|
||||
"input_schema": schema,
|
||||
}
|
||||
]
|
||||
params["tool_choice"] = {"type": "tool", "name": self._BATCH_TOOL_NAME}
|
||||
else:
|
||||
schema_msg = "\n\nYou must respond with valid JSON matching this schema:\n" + json.dumps(
|
||||
schema, indent=2, ensure_ascii=False
|
||||
)
|
||||
system_prompt = (system_prompt + schema_msg) if system_prompt else schema_msg
|
||||
|
||||
if system_prompt:
|
||||
params["system"] = _cached_system_blocks(system_prompt)
|
||||
|
||||
# Batch params ARE the raw Messages body, so operator-configured extra
|
||||
# body params merge directly (the sync path routes them through the
|
||||
# SDK's extra_body, which does the same merge server-side).
|
||||
if self._extra_body:
|
||||
params.update(self._extra_body)
|
||||
|
||||
return params
|
||||
|
||||
def _translate_batch_message(self, message: Any) -> dict[str, Any]:
|
||||
"""Render an Anthropic Message as the OpenAI response body the engine parses.
|
||||
|
||||
The engine reads ``choices[0].message.content`` (json.loads'ing it when
|
||||
a schema was requested) and sums ``usage`` under the OpenAI key names.
|
||||
Forced-tool responses carry their JSON in the tool_use block's input,
|
||||
so that is re-serialized as the content string.
|
||||
"""
|
||||
content = ""
|
||||
tool_input = None
|
||||
for block in message.content:
|
||||
if block.type == "tool_use" and block.name == self._BATCH_TOOL_NAME:
|
||||
tool_input = block.input or {}
|
||||
elif block.type == "text":
|
||||
content += block.text
|
||||
if tool_input is not None:
|
||||
content = json.dumps(tool_input, ensure_ascii=False)
|
||||
|
||||
usage = getattr(message, "usage", None)
|
||||
input_tokens = (usage.input_tokens or 0) if usage else 0
|
||||
output_tokens = (usage.output_tokens or 0) if usage else 0
|
||||
|
||||
return {
|
||||
"choices": [
|
||||
{
|
||||
"message": {"role": "assistant", "content": content},
|
||||
"finish_reason": getattr(message, "stop_reason", None),
|
||||
}
|
||||
],
|
||||
"usage": {
|
||||
"prompt_tokens": input_tokens,
|
||||
"completion_tokens": output_tokens,
|
||||
"total_tokens": input_tokens + output_tokens,
|
||||
},
|
||||
}
|
||||
|
||||
async def submit_batch(
|
||||
self,
|
||||
requests: list[dict[str, Any]],
|
||||
endpoint: str = "/v1/chat/completions",
|
||||
completion_window: str = "24h",
|
||||
) -> dict[str, Any]:
|
||||
"""Submit a batch of requests to the Message Batches API.
|
||||
|
||||
Accepts the engine's OpenAI-JSONL-shaped entries. ``endpoint`` and
|
||||
``completion_window`` belong to that shared shape and have no Anthropic
|
||||
equivalent (batches always resolve within 24 hours); both are ignored.
|
||||
"""
|
||||
batch_requests = [
|
||||
{
|
||||
"custom_id": req["custom_id"],
|
||||
"params": self._translate_batch_body(req.get("body") or {}),
|
||||
}
|
||||
for req in requests
|
||||
]
|
||||
|
||||
logger.info(f"Submitting Anthropic message batch with {len(batch_requests)} requests")
|
||||
batch = await self._client.messages.batches.create(requests=batch_requests)
|
||||
logger.info(f"Anthropic batch submitted: {batch.id}, status={batch.processing_status}")
|
||||
|
||||
return {
|
||||
"batch_id": batch.id,
|
||||
"status": self._map_batch_status(batch.processing_status),
|
||||
"created_at": batch.created_at,
|
||||
"request_count": len(batch_requests),
|
||||
}
|
||||
|
||||
async def get_batch_status(self, batch_id: str) -> dict[str, Any]:
|
||||
"""Get batch status in the shape the engine's poll loop expects."""
|
||||
batch = await self._client.messages.batches.retrieve(batch_id)
|
||||
|
||||
counts = batch.request_counts
|
||||
processing = getattr(counts, "processing", 0) or 0
|
||||
succeeded = getattr(counts, "succeeded", 0) or 0
|
||||
errored = getattr(counts, "errored", 0) or 0
|
||||
canceled = getattr(counts, "canceled", 0) or 0
|
||||
expired = getattr(counts, "expired", 0) or 0
|
||||
resolved = succeeded + errored + canceled + expired
|
||||
|
||||
result: dict[str, Any] = {
|
||||
"batch_id": batch.id,
|
||||
"status": self._map_batch_status(batch.processing_status),
|
||||
"created_at": batch.created_at,
|
||||
"request_counts": {
|
||||
"total": processing + resolved,
|
||||
"completed": resolved,
|
||||
"failed": errored,
|
||||
},
|
||||
}
|
||||
|
||||
ended_at = getattr(batch, "ended_at", None)
|
||||
if ended_at:
|
||||
result["completed_at"] = ended_at
|
||||
|
||||
return result
|
||||
|
||||
async def retrieve_batch_results(self, batch_id: str) -> list[dict[str, Any]]:
|
||||
"""Retrieve completed batch results, translated to the OpenAI shape.
|
||||
|
||||
Succeeded entries become ``{"custom_id", "response": {"body": ...}}``;
|
||||
errored/canceled/expired entries become ``{"custom_id", "error": ...}``
|
||||
so the engine's per-result error handling applies unchanged.
|
||||
"""
|
||||
batch = await self._client.messages.batches.retrieve(batch_id)
|
||||
if batch.processing_status != "ended":
|
||||
raise ValueError(f"Batch {batch_id} is not completed yet (status: {batch.processing_status})")
|
||||
|
||||
decoder = await self._client.messages.batches.results(batch_id)
|
||||
results: list[dict[str, Any]] = []
|
||||
async for entry in decoder:
|
||||
outcome = entry.result
|
||||
if outcome.type == "succeeded":
|
||||
results.append(
|
||||
{
|
||||
"custom_id": entry.custom_id,
|
||||
"response": {"body": self._translate_batch_message(outcome.message)},
|
||||
}
|
||||
)
|
||||
else:
|
||||
error = getattr(outcome, "error", None)
|
||||
if error is not None:
|
||||
detail = f"{getattr(error, 'type', 'error')}: {getattr(error, 'message', error)}"
|
||||
else:
|
||||
detail = f"batch request {outcome.type}"
|
||||
results.append({"custom_id": entry.custom_id, "error": detail})
|
||||
|
||||
logger.info(f"Retrieved {len(results)} results for Anthropic batch {batch_id}")
|
||||
return results
|
||||
|
||||
async def cleanup(self) -> None:
|
||||
"""Clean up resources (close Anthropic client connections)."""
|
||||
if hasattr(self, "_client") and self._client:
|
||||
await self._client.close()
|
||||
|
||||
def supports_attempt_scoped_concurrency(self) -> bool:
|
||||
return True
|
||||
|
||||
@@ -11,14 +11,16 @@ import json
|
||||
import logging
|
||||
import tempfile
|
||||
import time
|
||||
from typing import Any
|
||||
from contextlib import AbstractAsyncContextManager, nullcontext
|
||||
from typing import Any, Callable
|
||||
|
||||
from pydantic import ValidationError
|
||||
|
||||
from hindsight_api.engine.llm_interface import LLMInterface
|
||||
from hindsight_api.engine.llm_interface import LLM_TOOL_CHOICE_AUTO, LLMInterface, LLMToolChoice, LLMToolChoiceMode
|
||||
from hindsight_api.engine.llm_trace import LLMResponseUsage, stash_response_usage
|
||||
from hindsight_api.engine.response_models import LLMToolCall, LLMToolCallResult, TokenUsage
|
||||
from hindsight_api.metrics import get_metrics_collector
|
||||
from hindsight_api.worker.stage import set_stage
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -49,6 +51,20 @@ def _get_isolated_claude_env() -> dict[str, str]:
|
||||
return _isolated_claude_env
|
||||
|
||||
|
||||
def _result_error_detail(message: Any) -> str:
|
||||
"""Build an actionable error string from an ``is_error`` ResultMessage.
|
||||
|
||||
The CLI can report a failure with ``is_error=True`` while ``subtype``
|
||||
still reads ``"success"``, putting the real detail in ``result`` (e.g.
|
||||
quota exhaustion: ``You've hit your weekly limit · resets ...`` with
|
||||
``api_error_status: 429``). The SDK's own fallback exception surfaces
|
||||
only the subtype, producing the misleading "Claude Code returned an
|
||||
error result: success" (issue #2702) — so prefer ``result``.
|
||||
"""
|
||||
detail = (message.result or "").strip() or message.subtype or "unknown error"
|
||||
return f"Claude Code reported an error: {detail}"
|
||||
|
||||
|
||||
class ClaudeCodeLLM(LLMInterface):
|
||||
"""
|
||||
LLM provider using Claude Code authentication.
|
||||
@@ -148,6 +164,7 @@ class ClaudeCodeLLM(LLMInterface):
|
||||
skip_validation: bool = False,
|
||||
strict_schema: bool = False,
|
||||
return_usage: bool = False,
|
||||
attempt_context: Callable[[], AbstractAsyncContextManager[None]] | None = None,
|
||||
) -> Any:
|
||||
"""
|
||||
Make an LLM API call with retry logic.
|
||||
@@ -176,6 +193,7 @@ class ClaudeCodeLLM(LLMInterface):
|
||||
from claude_agent_sdk import ( # type: ignore[unresolved-import]
|
||||
AssistantMessage,
|
||||
ClaudeAgentOptions,
|
||||
ResultMessage,
|
||||
TextBlock,
|
||||
query,
|
||||
)
|
||||
@@ -209,10 +227,27 @@ class ClaudeCodeLLM(LLMInterface):
|
||||
user_content += schema_instruction
|
||||
|
||||
# Configure SDK options
|
||||
#
|
||||
# tools=[] is required here for the same reason call_with_tools() below
|
||||
# already sets it: with `tools` left at its default (None -> full
|
||||
# "claude_code" built-in preset), allowed_tools=[] alone does not stop
|
||||
# the CLI from loading the full built-in toolset and deferring into
|
||||
# ToolSearch before answering, which burns the single max_turns=1
|
||||
# budget on a tool-deferral step instead of a text response. Without
|
||||
# this, single-turn calls intermittently fail with "Reached maximum
|
||||
# number of turns (1)" even though the prompt itself needs no tools.
|
||||
options = ClaudeAgentOptions(
|
||||
system_prompt=system_prompt if system_prompt else None,
|
||||
max_turns=1, # Single-turn for API-style interactions
|
||||
tools=[], # Disable built-in tools so nothing forces a ToolSearch deferral
|
||||
allowed_tools=[], # Disable tools for standard LLM calls
|
||||
# Pin the configured model (issue #2881). Without this the spawned CLI
|
||||
# runs its own default model — an Opus-class model on Pro/Max OAuth —
|
||||
# regardless of HINDSIGHT_API_*_LLM_MODEL, while metrics/logs still print
|
||||
# self.model, so the mismatch is invisible. The isolated CLAUDE_CONFIG_DIR
|
||||
# (fresh temp dir) means a host settings.json can't reach the CLI either,
|
||||
# so passing it through here is the only channel.
|
||||
model=self.model or None,
|
||||
env=_get_isolated_claude_env(),
|
||||
)
|
||||
|
||||
@@ -223,11 +258,18 @@ class ClaudeCodeLLM(LLMInterface):
|
||||
# Collect streaming response
|
||||
full_text = ""
|
||||
|
||||
async for message in query(prompt=user_content, options=options):
|
||||
if isinstance(message, AssistantMessage):
|
||||
for block in message.content:
|
||||
if isinstance(block, TextBlock):
|
||||
full_text += block.text
|
||||
async with attempt_context() if attempt_context is not None else nullcontext():
|
||||
set_stage(f"llm.claude_code.{scope}.attempt={attempt + 1}/{max_retries + 1}")
|
||||
async for message in query(prompt=user_content, options=options):
|
||||
if isinstance(message, AssistantMessage):
|
||||
for block in message.content:
|
||||
if isinstance(block, TextBlock):
|
||||
full_text += block.text
|
||||
elif isinstance(message, ResultMessage) and message.is_error:
|
||||
# Surface the CLI's actual error text (e.g. quota
|
||||
# exhaustion) instead of the SDK's subtype-based
|
||||
# fallback exception (issue #2702).
|
||||
raise RuntimeError(_result_error_detail(message))
|
||||
|
||||
# The Claude Agent SDK doesn't report exact counts; stash the same
|
||||
# char/4 estimate the success path traces so a later parse/validate
|
||||
@@ -287,7 +329,7 @@ class ClaudeCodeLLM(LLMInterface):
|
||||
|
||||
# Record trace span
|
||||
try:
|
||||
from hindsight_api.tracing import get_span_recorder
|
||||
from hindsight_api.tracing import _serialize_for_span, get_span_recorder
|
||||
|
||||
span_recorder = get_span_recorder()
|
||||
span_recorder.record_llm_call(
|
||||
@@ -295,15 +337,17 @@ class ClaudeCodeLLM(LLMInterface):
|
||||
model=self.model,
|
||||
scope=scope,
|
||||
messages=messages,
|
||||
response_content=result if isinstance(result, str) else result.model_dump_json(),
|
||||
response_content=_serialize_for_span(result),
|
||||
input_tokens=estimated_input,
|
||||
output_tokens=estimated_output,
|
||||
duration=duration,
|
||||
finish_reason=None,
|
||||
error=None,
|
||||
)
|
||||
except Exception:
|
||||
pass # logging failure must never affect the operation
|
||||
except Exception as span_error:
|
||||
# Tracing must remain best-effort, but expose instrumentation
|
||||
# bugs that would otherwise silently erase spans (#3025).
|
||||
logger.debug("Claude Code span recording failed: %s", span_error, exc_info=True)
|
||||
|
||||
# Log slow calls
|
||||
if duration > 10.0:
|
||||
@@ -362,7 +406,8 @@ class ClaudeCodeLLM(LLMInterface):
|
||||
max_retries: int = 5,
|
||||
initial_backoff: float = 1.0,
|
||||
max_backoff: float = 30.0,
|
||||
tool_choice: str | dict[str, Any] = "auto",
|
||||
tool_choice: LLMToolChoice = LLM_TOOL_CHOICE_AUTO,
|
||||
attempt_context: Callable[[], AbstractAsyncContextManager[None]] | None = None,
|
||||
) -> LLMToolCallResult:
|
||||
"""
|
||||
Make an LLM API call with tool/function calling support using Claude Agent SDK.
|
||||
@@ -380,7 +425,7 @@ class ClaudeCodeLLM(LLMInterface):
|
||||
max_retries: Maximum retry attempts.
|
||||
initial_backoff: Initial backoff time in seconds.
|
||||
max_backoff: Maximum backoff time in seconds.
|
||||
tool_choice: How to choose tools - "auto", "none", "required", or specific function dict.
|
||||
tool_choice: Canonical tool-selection policy.
|
||||
- "auto": Model decides whether to call tools (default)
|
||||
- "required": Model must call at least one tool
|
||||
- "none": Model must not call any tools
|
||||
@@ -393,6 +438,7 @@ class ClaudeCodeLLM(LLMInterface):
|
||||
AssistantMessage,
|
||||
ClaudeAgentOptions,
|
||||
ClaudeSDKClient,
|
||||
ResultMessage,
|
||||
SdkMcpTool,
|
||||
TextBlock,
|
||||
ToolUseBlock,
|
||||
@@ -473,30 +519,27 @@ class ClaudeCodeLLM(LLMInterface):
|
||||
mcp_servers_config = {"hindsight_tools": mcp_server} if sdk_tools else {}
|
||||
|
||||
# Process tool_choice
|
||||
if isinstance(tool_choice, dict) and tool_choice.get("type") == "function":
|
||||
if tool_choice.mode is LLMToolChoiceMode.NAMED:
|
||||
# Force a specific tool: filter allowed_tools to only that tool and add instruction
|
||||
forced_name = tool_choice.get("function", {}).get("name")
|
||||
if forced_name:
|
||||
# Filter to only the forced tool (with MCP prefix)
|
||||
forced_tool_mcp_name = f"mcp__hindsight_tools__{forced_name}"
|
||||
if forced_tool_mcp_name in allowed_tool_names:
|
||||
allowed_tool_names = [forced_tool_mcp_name]
|
||||
# Add strong instruction to system prompt
|
||||
force_instruction = (
|
||||
f"\n\nIMPORTANT: You MUST call the '{forced_name}' tool. Do not respond with text only."
|
||||
)
|
||||
system_prompt += force_instruction
|
||||
logger.debug(f"Claude Code: Forcing tool call to '{forced_name}'")
|
||||
else:
|
||||
logger.warning(f"Claude Code: Forced tool '{forced_name}' not found in available tools")
|
||||
elif tool_choice == "required":
|
||||
forced_name = tool_choice.selected_function_name
|
||||
forced_tool_mcp_name = f"mcp__hindsight_tools__{forced_name}"
|
||||
if forced_tool_mcp_name in allowed_tool_names:
|
||||
allowed_tool_names = [forced_tool_mcp_name]
|
||||
force_instruction = (
|
||||
f"\n\nIMPORTANT: You MUST call the '{forced_name}' tool. Do not respond with text only."
|
||||
)
|
||||
system_prompt += force_instruction
|
||||
logger.debug(f"Claude Code: Forcing tool call to '{forced_name}'")
|
||||
else:
|
||||
logger.warning(f"Claude Code: Forced tool '{forced_name}' not found in available tools")
|
||||
elif tool_choice.mode is LLMToolChoiceMode.REQUIRED:
|
||||
# Must call at least one tool
|
||||
tool_instruction = (
|
||||
"\n\nIMPORTANT: You MUST call at least one of the available tools. Do not respond with text only."
|
||||
)
|
||||
system_prompt += tool_instruction
|
||||
logger.debug("Claude Code: Tool call required")
|
||||
elif tool_choice == "none":
|
||||
elif tool_choice.mode is LLMToolChoiceMode.NONE:
|
||||
# No tools should be called - disable all tools
|
||||
allowed_tool_names = []
|
||||
mcp_servers_config = {}
|
||||
@@ -504,16 +547,33 @@ class ClaudeCodeLLM(LLMInterface):
|
||||
# else: tool_choice == "auto" or unspecified - use default behavior (no changes needed)
|
||||
|
||||
# Configure SDK options with MCP server
|
||||
#
|
||||
# tools=[] disables built-in CLI tools (Read, Write, Bash, ToolSearch, etc.)
|
||||
# Without this, Claude Code CLI defers MCP tools when too many built-in tools
|
||||
# are loaded, forcing Claude to use ToolSearch first — which wastes the max_turns
|
||||
# are loaded, forcing Claude to use ToolSearch first — which wastes the turn
|
||||
# budget and prevents direct MCP tool calls.
|
||||
#
|
||||
# max_turns=1 is critical (issue #2966). call_with_tools() is one *round* of
|
||||
# an agentic loop the caller drives: the model proposes tool calls, we return
|
||||
# them, and the orchestrator (reflect/agent.py) executes the REAL tools and
|
||||
# feeds the results back on the next call. The SDK, however, runs its own
|
||||
# in-process loop: it invokes our SDK MCP handlers — which are deliberate
|
||||
# placeholders returning "[Tool <name> called successfully]" (no real data) —
|
||||
# and lets the model react. With max_turns >= 2 the model calls recall, sees
|
||||
# the empty placeholder, re-queries with reworded searches, exhausts the turn
|
||||
# budget, and the run ends in error_max_turns with its tool calls discarded —
|
||||
# exactly the "0 tool calls / no information" failure in #2966. Capping at a
|
||||
# single turn stops the SDK from acting on the placeholder results: the model
|
||||
# emits its first tool call (or a text answer) and we return that to the caller
|
||||
# unchanged, matching how every other provider's call_with_tools() behaves.
|
||||
options = ClaudeAgentOptions(
|
||||
system_prompt=system_prompt if system_prompt else None,
|
||||
tools=[], # Disable built-in tools so MCP tools load eagerly
|
||||
max_turns=2, # Allow tool call + tool result round-trip
|
||||
max_turns=1, # One round: propose tool calls (or answer); caller drives the loop
|
||||
mcp_servers=mcp_servers_config,
|
||||
allowed_tools=allowed_tool_names,
|
||||
# Pin the configured model (issue #2881) — see the call() options block.
|
||||
model=self.model or None,
|
||||
env=_get_isolated_claude_env(),
|
||||
)
|
||||
|
||||
@@ -524,32 +584,49 @@ class ClaudeCodeLLM(LLMInterface):
|
||||
full_text = ""
|
||||
tool_calls: list[LLMToolCall] = []
|
||||
|
||||
# Use ClaudeSDKClient for tool calling support
|
||||
# Note: query() does NOT support custom tools, only ClaudeSDKClient does
|
||||
async with ClaudeSDKClient(options=options) as client:
|
||||
# Send the query
|
||||
await client.query(user_content)
|
||||
async with attempt_context() if attempt_context is not None else nullcontext():
|
||||
set_stage(f"llm.claude_code.tools.attempt={attempt + 1}/{max_retries + 1}")
|
||||
# Use ClaudeSDKClient for tool calling support
|
||||
# Note: query() does NOT support custom tools, only ClaudeSDKClient does
|
||||
async with ClaudeSDKClient(options=options) as client:
|
||||
# Send the query
|
||||
await client.query(user_content)
|
||||
|
||||
# Receive response
|
||||
async for message in client.receive_response():
|
||||
if isinstance(message, AssistantMessage):
|
||||
for block in message.content:
|
||||
if isinstance(block, TextBlock):
|
||||
full_text += block.text
|
||||
elif isinstance(block, ToolUseBlock):
|
||||
# SDK returns tool names with MCP prefix (mcp__hindsight_tools__{name})
|
||||
# Strip the prefix to return original tool name expected by caller
|
||||
tool_name = block.name
|
||||
if tool_name.startswith("mcp__hindsight_tools__"):
|
||||
tool_name = tool_name.replace("mcp__hindsight_tools__", "", 1)
|
||||
# Receive response
|
||||
async for message in client.receive_response():
|
||||
if isinstance(message, AssistantMessage):
|
||||
for block in message.content:
|
||||
if isinstance(block, TextBlock):
|
||||
full_text += block.text
|
||||
elif isinstance(block, ToolUseBlock):
|
||||
# SDK returns tool names with MCP prefix (mcp__hindsight_tools__{name})
|
||||
# Strip the prefix to return original tool name expected by caller
|
||||
tool_name = block.name
|
||||
if tool_name.startswith("mcp__hindsight_tools__"):
|
||||
tool_name = tool_name.replace("mcp__hindsight_tools__", "", 1)
|
||||
|
||||
tool_calls.append(
|
||||
LLMToolCall(
|
||||
id=block.id,
|
||||
name=tool_name,
|
||||
arguments=block.input,
|
||||
tool_calls.append(
|
||||
LLMToolCall(
|
||||
id=block.id,
|
||||
name=tool_name,
|
||||
arguments=block.input,
|
||||
)
|
||||
)
|
||||
)
|
||||
if tool_calls:
|
||||
# This round proposed tool call(s). Stop consuming the
|
||||
# stream so the SDK does not run another turn against our
|
||||
# placeholder handlers (issue #2966) — the caller executes
|
||||
# the real tools and calls us again with the results.
|
||||
break
|
||||
elif isinstance(message, ResultMessage) and message.is_error:
|
||||
# With max_turns=1 the CLI reports error_max_turns whenever
|
||||
# the model spent its single turn issuing a tool call (there
|
||||
# was no follow-up turn to emit final text). That is expected
|
||||
# here and not a failure: we already captured the tool call
|
||||
# above and break before reaching this branch. Only a genuine
|
||||
# error with nothing to return should surface (issue #2702).
|
||||
if not tool_calls:
|
||||
raise RuntimeError(_result_error_detail(message))
|
||||
|
||||
# Record metrics
|
||||
duration = time.time() - start_time
|
||||
@@ -611,3 +688,6 @@ class ClaudeCodeLLM(LLMInterface):
|
||||
async def cleanup(self) -> None:
|
||||
"""Clean up resources (no HTTP client to close for Claude Agent SDK)."""
|
||||
pass
|
||||
|
||||
def supports_attempt_scoped_concurrency(self) -> bool:
|
||||
return True
|
||||
|
||||
@@ -19,6 +19,7 @@ from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import binascii
|
||||
import contextlib
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
@@ -31,6 +32,11 @@ from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
try:
|
||||
import fcntl
|
||||
except ImportError: # pragma: no cover - Windows
|
||||
fcntl = None # type: ignore[assignment]
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@@ -58,6 +64,9 @@ _CODEX_TOKEN_REFRESH_SKEW_SECONDS = 60
|
||||
_CODEX_TERMINAL_REFRESH_ERROR_CODES = frozenset(
|
||||
{"refresh_token_expired", "refresh_token_reused", "refresh_token_invalidated"}
|
||||
)
|
||||
_CODEX_AUTH_LOCK_TIMEOUT_SECONDS = 20.0
|
||||
_CODEX_AUTH_LOCKS_GUARD = threading.Lock()
|
||||
_CODEX_AUTH_LOCKS: dict[Path, threading.Lock] = {}
|
||||
|
||||
|
||||
def default_codex_auth_file() -> Path:
|
||||
@@ -76,6 +85,44 @@ def default_codex_auth_file() -> Path:
|
||||
return Path.home() / ".codex" / "auth.json"
|
||||
|
||||
|
||||
def _path_scoped_lock(auth_file: Path) -> threading.Lock:
|
||||
key = auth_file.expanduser().resolve(strict=False)
|
||||
with _CODEX_AUTH_LOCKS_GUARD:
|
||||
lock = _CODEX_AUTH_LOCKS.get(key)
|
||||
if lock is None:
|
||||
lock = threading.Lock()
|
||||
_CODEX_AUTH_LOCKS[key] = lock
|
||||
return lock
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _codex_auth_lock(auth_file: Path, timeout_seconds: float = _CODEX_AUTH_LOCK_TIMEOUT_SECONDS):
|
||||
"""Cross-process advisory lock for one Codex auth store."""
|
||||
with _path_scoped_lock(auth_file):
|
||||
if fcntl is None: # pragma: no cover - Windows
|
||||
logger.debug("fcntl unavailable; Codex refresh proceeds without a cross-process lock.")
|
||||
yield
|
||||
return
|
||||
|
||||
lock_path = auth_file.with_suffix(".lock")
|
||||
lock_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with open(lock_path, "a+") as lock_file:
|
||||
deadline = time.monotonic() + max(1.0, timeout_seconds)
|
||||
while True:
|
||||
try:
|
||||
fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
|
||||
break
|
||||
except (BlockingIOError, OSError):
|
||||
if time.monotonic() >= deadline:
|
||||
raise TimeoutError("Timed out waiting for the Codex auth store lock") from None
|
||||
time.sleep(0.05)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
with contextlib.suppress(OSError):
|
||||
fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN)
|
||||
|
||||
|
||||
class CodexRefreshExpiredError(RuntimeError):
|
||||
"""Raised when the Codex refresh_token itself is no longer valid.
|
||||
|
||||
@@ -192,6 +239,34 @@ class CodexAuthManager:
|
||||
return None
|
||||
return data.get("tokens", {}).get("refresh_token")
|
||||
|
||||
@staticmethod
|
||||
def _load_tokens_from_file(auth_file: Path) -> dict[str, Any] | None:
|
||||
try:
|
||||
with open(auth_file) as f:
|
||||
data = json.load(f)
|
||||
except (OSError, json.JSONDecodeError):
|
||||
return None
|
||||
tokens = data.get("tokens")
|
||||
return tokens if isinstance(tokens, dict) else None
|
||||
|
||||
def _adopt_tokens(self, tokens: dict[str, Any]) -> bool:
|
||||
"""Adopt a newer on-disk Codex token set if present."""
|
||||
access_token = tokens.get("access_token")
|
||||
refresh_token = tokens.get("refresh_token")
|
||||
account_id = tokens.get("account_id")
|
||||
|
||||
changed = False
|
||||
if isinstance(access_token, str) and access_token and access_token != self.access_token:
|
||||
self.access_token = access_token
|
||||
changed = True
|
||||
if isinstance(refresh_token, str) and refresh_token and refresh_token != self.refresh_token:
|
||||
self.refresh_token = refresh_token
|
||||
changed = True
|
||||
if isinstance(account_id, str) and account_id and account_id != self.account_id:
|
||||
self.account_id = account_id
|
||||
changed = True
|
||||
return changed
|
||||
|
||||
@staticmethod
|
||||
def _decode_jwt_exp_unixtime(token: str) -> int | None:
|
||||
"""Return the JWT ``exp`` claim as a unix timestamp, or None on parse failure.
|
||||
@@ -225,6 +300,11 @@ class CodexAuthManager:
|
||||
return False
|
||||
return exp <= int(time.time()) + skew_seconds
|
||||
|
||||
def _token_is_fresh_with_known_expiry(self, skew_seconds: int = _CODEX_TOKEN_REFRESH_SKEW_SECONDS) -> bool:
|
||||
"""True only when the cached token has a known expiry outside the skew window."""
|
||||
exp = self._decode_jwt_exp_unixtime(self.access_token)
|
||||
return exp is not None and exp > int(time.time()) + skew_seconds
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Persistence
|
||||
# ------------------------------------------------------------------
|
||||
@@ -339,78 +419,93 @@ class CodexAuthManager:
|
||||
if not self._token_is_stale():
|
||||
return
|
||||
|
||||
if not self.refresh_token:
|
||||
raise RuntimeError(
|
||||
"Codex access_token is expired but no refresh_token is available. "
|
||||
"Run 'codex auth login' to re-authenticate."
|
||||
)
|
||||
with _codex_auth_lock(self._auth_file):
|
||||
disk_tokens = self._load_tokens_from_file(self._auth_file)
|
||||
if disk_tokens and self._adopt_tokens(disk_tokens):
|
||||
if force or self._token_is_fresh_with_known_expiry():
|
||||
return
|
||||
|
||||
log_reason = f" ({reason})" if reason else ""
|
||||
logger.info(f"Refreshing Codex OAuth access_token{log_reason}")
|
||||
|
||||
request_body = {
|
||||
"client_id": _CODEX_CLIENT_ID,
|
||||
"grant_type": "refresh_token",
|
||||
"refresh_token": self.refresh_token,
|
||||
}
|
||||
try:
|
||||
response = self._http_client.post(
|
||||
_CODEX_REFRESH_TOKEN_URL,
|
||||
json=request_body,
|
||||
headers={"Content-Type": "application/json"},
|
||||
timeout=30.0,
|
||||
)
|
||||
except httpx.RequestError as e:
|
||||
raise RuntimeError(f"Codex OAuth refresh network error: {type(e).__name__}") from e
|
||||
|
||||
if response.status_code == 401:
|
||||
error_code = self._extract_oauth_error_code(response)
|
||||
if error_code in _CODEX_TERMINAL_REFRESH_ERROR_CODES:
|
||||
raise CodexRefreshExpiredError(
|
||||
f"Codex refresh_token is permanently invalid (error.code={error_code}). "
|
||||
if not self.refresh_token:
|
||||
raise RuntimeError(
|
||||
"Codex access_token is expired but no refresh_token is available. "
|
||||
"Run 'codex auth login' to re-authenticate."
|
||||
)
|
||||
raise CodexRefreshExpiredError(
|
||||
f"Codex OAuth refresh returned 401 with unrecognized error code "
|
||||
f"({error_code or 'none'}). Run 'codex auth login' to re-authenticate."
|
||||
)
|
||||
|
||||
if response.status_code >= 400:
|
||||
raise RuntimeError(f"Codex OAuth refresh failed with HTTP {response.status_code}")
|
||||
log_reason = f" ({reason})" if reason else ""
|
||||
logger.info(f"Refreshing Codex OAuth access_token{log_reason}")
|
||||
|
||||
try:
|
||||
body = response.json()
|
||||
except json.JSONDecodeError as e:
|
||||
raise RuntimeError(f"Codex OAuth refresh returned non-JSON body: {e}") from e
|
||||
request_access_token = self.access_token
|
||||
request_refresh_token = self.refresh_token
|
||||
request_body = {
|
||||
"client_id": _CODEX_CLIENT_ID,
|
||||
"grant_type": "refresh_token",
|
||||
"refresh_token": request_refresh_token,
|
||||
}
|
||||
try:
|
||||
response = self._http_client.post(
|
||||
_CODEX_REFRESH_TOKEN_URL,
|
||||
json=request_body,
|
||||
headers={"Content-Type": "application/json"},
|
||||
timeout=30.0,
|
||||
)
|
||||
except httpx.RequestError as e:
|
||||
raise RuntimeError(f"Codex OAuth refresh network error: {type(e).__name__}") from e
|
||||
|
||||
new_access = body.get("access_token")
|
||||
if not new_access:
|
||||
raise RuntimeError("Codex OAuth refresh returned no access_token")
|
||||
if response.status_code == 401:
|
||||
error_code = self._extract_oauth_error_code(response)
|
||||
disk_tokens = self._load_tokens_from_file(self._auth_file)
|
||||
if disk_tokens and (
|
||||
disk_tokens.get("access_token") != request_access_token
|
||||
or disk_tokens.get("refresh_token") != request_refresh_token
|
||||
):
|
||||
self._adopt_tokens(disk_tokens)
|
||||
return
|
||||
if error_code in _CODEX_TERMINAL_REFRESH_ERROR_CODES:
|
||||
raise CodexRefreshExpiredError(
|
||||
f"Codex refresh_token is permanently invalid (error.code={error_code}). "
|
||||
"Run 'codex auth login' to re-authenticate."
|
||||
)
|
||||
raise CodexRefreshExpiredError(
|
||||
f"Codex OAuth refresh returned 401 with unrecognized error code "
|
||||
f"({error_code or 'none'}). Run 'codex auth login' to re-authenticate."
|
||||
)
|
||||
|
||||
new_refresh = body.get("refresh_token") or self.refresh_token
|
||||
new_id_token = body.get("id_token")
|
||||
if response.status_code >= 400:
|
||||
raise RuntimeError(f"Codex OAuth refresh failed with HTTP {response.status_code}")
|
||||
|
||||
# Update in-memory state first so waiters see fresh credentials
|
||||
# immediately, even if disk write fails.
|
||||
self.access_token = new_access
|
||||
self.refresh_token = new_refresh
|
||||
try:
|
||||
body = response.json()
|
||||
except json.JSONDecodeError as e:
|
||||
raise RuntimeError(f"Codex OAuth refresh returned non-JSON body: {e}") from e
|
||||
|
||||
persisted: dict[str, Any] = {
|
||||
"access_token": new_access,
|
||||
"refresh_token": new_refresh,
|
||||
}
|
||||
if new_id_token:
|
||||
persisted["id_token"] = new_id_token
|
||||
new_access = body.get("access_token")
|
||||
if not new_access:
|
||||
raise RuntimeError("Codex OAuth refresh returned no access_token")
|
||||
|
||||
try:
|
||||
self._persist_auth_atomic(persisted)
|
||||
except OSError as e:
|
||||
logger.warning(
|
||||
f"Codex OAuth refresh succeeded but persisting auth.json failed: {type(e).__name__}. "
|
||||
"In-memory credentials are up to date; on-disk file is stale."
|
||||
)
|
||||
new_refresh = body.get("refresh_token") or self.refresh_token
|
||||
new_id_token = body.get("id_token")
|
||||
|
||||
logger.info("Codex OAuth access_token refreshed successfully")
|
||||
# Update in-memory state first so waiters see fresh credentials
|
||||
# immediately, even if disk write fails.
|
||||
self.access_token = new_access
|
||||
self.refresh_token = new_refresh
|
||||
|
||||
persisted: dict[str, Any] = {
|
||||
"access_token": new_access,
|
||||
"refresh_token": new_refresh,
|
||||
}
|
||||
if new_id_token:
|
||||
persisted["id_token"] = new_id_token
|
||||
|
||||
try:
|
||||
self._persist_auth_atomic(persisted)
|
||||
except OSError as e:
|
||||
logger.warning(
|
||||
f"Codex OAuth refresh succeeded but persisting auth.json failed: {type(e).__name__}. "
|
||||
"In-memory credentials are up to date; on-disk file is stale."
|
||||
)
|
||||
|
||||
logger.info("Codex OAuth access_token refreshed successfully")
|
||||
|
||||
def ensure_fresh_token(self) -> None:
|
||||
"""Proactively refresh the access_token if it is near or past expiry.
|
||||
|
||||
@@ -20,15 +20,19 @@ import json
|
||||
import logging
|
||||
import time
|
||||
import uuid
|
||||
from contextlib import AbstractAsyncContextManager, nullcontext
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from typing import Any, Callable
|
||||
|
||||
import httpx
|
||||
|
||||
from hindsight_api.engine.llm_interface import LLMInterface
|
||||
from hindsight_api.engine.llm_interface import LLM_TOOL_CHOICE_AUTO, LLMInterface, LLMToolChoice, LLMToolChoiceMode
|
||||
from hindsight_api.engine.llm_trace import LLMResponseUsage, stash_response_usage
|
||||
from hindsight_api.engine.providers.llm_debug import dump_request_on_4xx
|
||||
from hindsight_api.engine.response_models import LLMToolCall, LLMToolCallResult, TokenUsage
|
||||
from hindsight_api.engine.structured_output import strict_json_schema
|
||||
from hindsight_api.metrics import get_metrics_collector
|
||||
from hindsight_api.worker.stage import set_stage
|
||||
|
||||
from .codex_auth import (
|
||||
_CODEX_CLIENT_ID,
|
||||
@@ -53,6 +57,59 @@ __all__ = [
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Newer Codex models are gated on the first-party client identity; the previous
|
||||
# browser-shaped User-Agent returned "Model not found" for Luna (#2643).
|
||||
# Use a neutral version because Hindsight must not claim a specific Codex release.
|
||||
_CODEX_ORIGINATOR = "codex_cli_rs"
|
||||
_CODEX_USER_AGENT = "codex_cli_rs/0.0.0 (Hindsight)"
|
||||
|
||||
# Name of the single forced function tool used to carry structured output when
|
||||
# strict_schema is on. The Codex backend speaks the OpenAI Responses API, so a
|
||||
# forced function call gives us constrained decoding straight into the response
|
||||
# schema — no prompt-injected schema, no raw json.loads on free-form model text,
|
||||
# no invalid-\escape retry storm (issue #2504, same class as #1002 / #2339).
|
||||
_STRUCTURED_TOOL_NAME = "structured_response"
|
||||
|
||||
# Valid JSON string escape characters (the char that may follow a backslash).
|
||||
_VALID_JSON_ESCAPE_CHARS = set('"\\/bfnrtu')
|
||||
|
||||
|
||||
def _repair_invalid_json_escapes(text: str) -> str:
|
||||
"""Best-effort repair of invalid ``\\escape`` sequences in a JSON string.
|
||||
|
||||
Escape-heavy content (code, serial/CLI commands, Windows paths, regexes)
|
||||
makes weaker models emit backslashes that aren't valid JSON escapes (e.g.
|
||||
``\\d``, ``\\s``, ``C:\\Users``), so ``json.loads`` fails deterministically
|
||||
and every retry re-fails the same way (issue #2504). This doubles any
|
||||
backslash that isn't part of a valid escape so the payload parses. It is a
|
||||
lenient fallback only — the strict_schema forced-tool path is the real fix.
|
||||
"""
|
||||
result: list[str] = []
|
||||
i = 0
|
||||
n = len(text)
|
||||
while i < n:
|
||||
ch = text[i]
|
||||
if ch == "\\" and i + 1 < n:
|
||||
nxt = text[i + 1]
|
||||
if nxt in _VALID_JSON_ESCAPE_CHARS:
|
||||
# Preserve the valid escape (both chars) verbatim.
|
||||
result.append(ch)
|
||||
result.append(nxt)
|
||||
i += 2
|
||||
continue
|
||||
# Invalid escape: escape the lone backslash so JSON parses.
|
||||
result.append("\\\\")
|
||||
i += 1
|
||||
continue
|
||||
if ch == "\\" and i + 1 == n:
|
||||
# Trailing lone backslash — escape it.
|
||||
result.append("\\\\")
|
||||
i += 1
|
||||
continue
|
||||
result.append(ch)
|
||||
i += 1
|
||||
return "".join(result)
|
||||
|
||||
|
||||
class CodexLLM(LLMInterface):
|
||||
"""
|
||||
@@ -70,6 +127,7 @@ class CodexLLM(LLMInterface):
|
||||
base_url: str,
|
||||
model: str,
|
||||
reasoning_effort: str = "low",
|
||||
extra_body: dict[str, Any] | None = None,
|
||||
**kwargs: Any,
|
||||
):
|
||||
"""Initialize Codex LLM provider."""
|
||||
@@ -117,9 +175,10 @@ class CodexLLM(LLMInterface):
|
||||
if self.model.startswith("openai/"):
|
||||
self.model = self.model[len("openai/") :]
|
||||
|
||||
# Map reasoning effort to Codex reasoning summary format
|
||||
# Codex supports: "auto", "concise", "detailed"
|
||||
# Reasoning summary controls presentation separately from the backend's
|
||||
# reasoning effort, which is sent unchanged in each request payload.
|
||||
self.reasoning_summary = self._map_reasoning_effort(reasoning_effort)
|
||||
self._extra_body = dict(extra_body or {})
|
||||
|
||||
# HTTP client for SSE streaming
|
||||
self._client = httpx.AsyncClient(timeout=120.0)
|
||||
@@ -140,6 +199,18 @@ class CodexLLM(LLMInterface):
|
||||
def account_id(self) -> str:
|
||||
return self._auth_manager.account_id
|
||||
|
||||
def _build_request_headers(self) -> httpx.Headers:
|
||||
return httpx.Headers(
|
||||
{
|
||||
"Authorization": f"Bearer {self.access_token}",
|
||||
"Content-Type": "application/json",
|
||||
"OpenAI-Account-ID": self.account_id,
|
||||
"User-Agent": _CODEX_USER_AGENT,
|
||||
"Origin": "https://chatgpt.com",
|
||||
"originator": _CODEX_ORIGINATOR,
|
||||
}
|
||||
)
|
||||
|
||||
@property
|
||||
def refresh_token(self) -> str | None:
|
||||
return self._auth_manager.refresh_token
|
||||
@@ -276,32 +347,6 @@ class CodexLLM(LLMInterface):
|
||||
}
|
||||
return mapping.get(effort.lower(), "auto")
|
||||
|
||||
def _normalize_tool_choice(self, tool_choice: str | dict[str, Any]) -> str | dict[str, Any]:
|
||||
"""Normalize forced function tool choice for the Codex Responses API.
|
||||
|
||||
Older agent paths may still pass OpenAI chat-completions style named
|
||||
tool choice payloads such as:
|
||||
|
||||
{"type": "function", "function": {"name": "recall"}}
|
||||
|
||||
Codex Responses expects the named function at the top level instead:
|
||||
|
||||
{"type": "function", "name": "recall"}
|
||||
"""
|
||||
if not isinstance(tool_choice, dict):
|
||||
return tool_choice
|
||||
if str(tool_choice.get("type") or "").strip() != "function":
|
||||
return tool_choice
|
||||
function_payload = tool_choice.get("function")
|
||||
if isinstance(function_payload, dict):
|
||||
function_name = str(function_payload.get("name") or "").strip()
|
||||
if function_name:
|
||||
return {"type": "function", "name": function_name}
|
||||
function_name = str(tool_choice.get("name") or "").strip()
|
||||
if function_name:
|
||||
return {"type": "function", "name": function_name}
|
||||
return tool_choice
|
||||
|
||||
async def verify_connection(self) -> None:
|
||||
"""Verify Codex connection by making a simple test call."""
|
||||
try:
|
||||
@@ -335,8 +380,20 @@ class CodexLLM(LLMInterface):
|
||||
skip_validation: bool = False,
|
||||
strict_schema: bool = False,
|
||||
return_usage: bool = False,
|
||||
attempt_context: Callable[[], AbstractAsyncContextManager[None]] | None = None,
|
||||
) -> Any:
|
||||
"""Make API call to Codex backend with SSE streaming."""
|
||||
"""Make API call to Codex backend with SSE streaming.
|
||||
|
||||
Args:
|
||||
strict_schema: Route structured output through a single forced
|
||||
function tool (constrained decoding) instead of prompt-injecting
|
||||
the schema and parsing free-form text. The Codex backend speaks
|
||||
the OpenAI Responses API, so the forced function call emits the
|
||||
response schema directly as tool arguments — eliminating the
|
||||
invalid-``\\escape`` retry storm (issue #2504). When False, falls
|
||||
back to schema-in-prompt + JSON parse, now hardened with a lenient
|
||||
invalid-escape repair before giving up.
|
||||
"""
|
||||
start_time = time.time()
|
||||
|
||||
# Proactively refresh the OAuth access_token if it's near expiry.
|
||||
@@ -361,11 +418,22 @@ class CodexLLM(LLMInterface):
|
||||
else:
|
||||
user_messages.append(msg)
|
||||
|
||||
# Add JSON schema instruction if response_format is provided
|
||||
# Structured output: prefer a single forced function tool (constrained
|
||||
# decoding) over text-injecting the schema and parsing the reply. The
|
||||
# forced tool guarantees schema-shaped JSON in the tool arguments,
|
||||
# eliminating the invalid-\escape retry storm (issue #2504). When
|
||||
# strict_schema is off we keep the schema-in-prompt + json.loads
|
||||
# fallback (now hardened with a lenient escape repair) for callers that
|
||||
# can't force tools.
|
||||
schema = None
|
||||
use_forced_tool = False
|
||||
if response_format is not None and hasattr(response_format, "model_json_schema"):
|
||||
schema = response_format.model_json_schema()
|
||||
schema_msg = f"\n\nYou must respond with valid JSON matching this schema:\n{json.dumps(schema, indent=2, ensure_ascii=False)}"
|
||||
system_instruction += schema_msg
|
||||
schema = strict_json_schema(response_format) if strict_schema else response_format.model_json_schema()
|
||||
if strict_schema:
|
||||
use_forced_tool = True
|
||||
else:
|
||||
schema_msg = f"\n\nYou must respond with valid JSON matching this schema:\n{json.dumps(schema, indent=2, ensure_ascii=False)}"
|
||||
system_instruction += schema_msg
|
||||
|
||||
# gpt-5.2-codex only supports "detailed" reasoning summary
|
||||
reasoning_summary = "detailed" if "5.2" in self.model else self.reasoning_summary
|
||||
@@ -385,20 +453,29 @@ class CodexLLM(LLMInterface):
|
||||
"tools": [],
|
||||
"tool_choice": "auto",
|
||||
"parallel_tool_calls": True,
|
||||
"reasoning": {"summary": reasoning_summary},
|
||||
"reasoning": {"effort": self.reasoning_effort, "summary": reasoning_summary},
|
||||
"store": False, # Codex uses stateless mode
|
||||
"stream": True, # SSE streaming
|
||||
"include": ["reasoning.encrypted_content"],
|
||||
"prompt_cache_key": str(uuid.uuid4()),
|
||||
}
|
||||
payload.update(self._extra_body)
|
||||
|
||||
headers = {
|
||||
"Authorization": f"Bearer {self.access_token}",
|
||||
"Content-Type": "application/json",
|
||||
"OpenAI-Account-ID": self.account_id,
|
||||
"User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7)",
|
||||
"Origin": "https://chatgpt.com",
|
||||
}
|
||||
if use_forced_tool and schema is not None:
|
||||
# Single function tool whose parameters ARE the response schema;
|
||||
# force it via tool_choice so the backend does constrained decoding.
|
||||
payload["tools"] = [
|
||||
{
|
||||
"type": "function",
|
||||
"name": _STRUCTURED_TOOL_NAME,
|
||||
"description": "Return the structured response.",
|
||||
"parameters": schema,
|
||||
}
|
||||
]
|
||||
payload["tool_choice"] = {"type": "function", "name": _STRUCTURED_TOOL_NAME}
|
||||
payload["parallel_tool_calls"] = False
|
||||
|
||||
headers = self._build_request_headers()
|
||||
|
||||
url = f"{self.base_url}/codex/responses"
|
||||
|
||||
@@ -409,11 +486,20 @@ class CodexLLM(LLMInterface):
|
||||
attempt = 0
|
||||
while True:
|
||||
try:
|
||||
response = await self._client.post(url, json=payload, headers=headers, timeout=120.0)
|
||||
response.raise_for_status()
|
||||
async with attempt_context() if attempt_context is not None else nullcontext():
|
||||
set_stage(f"llm.codex.{scope}.attempt={attempt + 1}/{max_retries + 1}")
|
||||
response = await self._client.post(url, json=payload, headers=headers, timeout=120.0)
|
||||
response.raise_for_status()
|
||||
|
||||
# Parse SSE stream
|
||||
content = await self._parse_sse_stream(response)
|
||||
# Forced-tool path: read structured output from the function-call
|
||||
# arguments (already a JSON string in a dedicated channel) rather
|
||||
# than from free-form assistant text.
|
||||
if use_forced_tool:
|
||||
text_content, tool_calls = await self._parse_sse_tool_stream(response)
|
||||
content = text_content or ""
|
||||
else:
|
||||
tool_calls = []
|
||||
content = await self._parse_sse_stream(response)
|
||||
|
||||
# Codex SSE carries no usage block; stash the same char/4 estimate
|
||||
# the success path traces so a later parse/validate failure records
|
||||
@@ -426,7 +512,28 @@ class CodexLLM(LLMInterface):
|
||||
)
|
||||
|
||||
# Handle structured output
|
||||
if response_format is not None:
|
||||
if use_forced_tool:
|
||||
tool_input = None
|
||||
for tc in tool_calls:
|
||||
if tc.name == _STRUCTURED_TOOL_NAME:
|
||||
tool_input = tc.arguments if isinstance(tc.arguments, dict) else None
|
||||
break
|
||||
if tool_input is None:
|
||||
# Model ignored the forced tool (rare — e.g. a gateway that
|
||||
# drops tool_choice). Retry so we don't hard-fail.
|
||||
logger.warning(
|
||||
f"Codex forced structured tool missing from response "
|
||||
f"(attempt {attempt + 1}/{max_retries + 1})"
|
||||
)
|
||||
if attempt < max_retries:
|
||||
backoff = min(initial_backoff * (2**attempt), max_backoff)
|
||||
await asyncio.sleep(backoff)
|
||||
attempt += 1
|
||||
continue
|
||||
raise RuntimeError("Codex did not return the forced structured_response tool call")
|
||||
content = json.dumps(tool_input)
|
||||
result = tool_input if skip_validation else response_format.model_validate(tool_input)
|
||||
elif response_format is not None:
|
||||
# Models may wrap JSON in markdown
|
||||
clean_content = content
|
||||
if "```json" in content:
|
||||
@@ -437,13 +544,20 @@ class CodexLLM(LLMInterface):
|
||||
try:
|
||||
json_data = json.loads(clean_content)
|
||||
except json.JSONDecodeError as e:
|
||||
logger.warning(f"Codex JSON parse error (attempt {attempt + 1}/{max_retries + 1}): {e}")
|
||||
if attempt < max_retries:
|
||||
backoff = min(initial_backoff * (2**attempt), max_backoff)
|
||||
await asyncio.sleep(backoff)
|
||||
attempt += 1
|
||||
continue
|
||||
raise
|
||||
# Escape-heavy content deterministically re-fails every
|
||||
# retry (issue #2504). Try a lenient invalid-escape repair
|
||||
# before burning a retry / re-raising.
|
||||
try:
|
||||
json_data = json.loads(_repair_invalid_json_escapes(clean_content))
|
||||
logger.info("Codex JSON parsed after repairing invalid escape sequences")
|
||||
except json.JSONDecodeError:
|
||||
logger.warning(f"Codex JSON parse error (attempt {attempt + 1}/{max_retries + 1}): {e}")
|
||||
if attempt < max_retries:
|
||||
backoff = min(initial_backoff * (2**attempt), max_backoff)
|
||||
await asyncio.sleep(backoff)
|
||||
attempt += 1
|
||||
continue
|
||||
raise
|
||||
|
||||
if skip_validation:
|
||||
result = json_data
|
||||
@@ -467,7 +581,7 @@ class CodexLLM(LLMInterface):
|
||||
|
||||
# Record trace span
|
||||
try:
|
||||
from hindsight_api.tracing import get_span_recorder
|
||||
from hindsight_api.tracing import _serialize_for_span, get_span_recorder
|
||||
|
||||
# Estimate tokens for tracing
|
||||
estimated_input = sum(len(m.get("content", "")) for m in messages) // 4
|
||||
@@ -478,15 +592,17 @@ class CodexLLM(LLMInterface):
|
||||
model=self.model,
|
||||
scope=scope,
|
||||
messages=messages,
|
||||
response_content=result if isinstance(result, str) else result.model_dump_json(),
|
||||
response_content=_serialize_for_span(result),
|
||||
input_tokens=estimated_input,
|
||||
output_tokens=estimated_output,
|
||||
duration=duration,
|
||||
finish_reason=None,
|
||||
error=None,
|
||||
)
|
||||
except Exception:
|
||||
pass # logging failure must never affect the operation
|
||||
except Exception as span_error:
|
||||
# Tracing must remain best-effort, but expose instrumentation
|
||||
# bugs that would otherwise silently erase spans (#3025).
|
||||
logger.debug("Codex span recording failed: %s", span_error, exc_info=True)
|
||||
|
||||
if return_usage:
|
||||
# Codex doesn't provide token counts, estimate based on content
|
||||
@@ -542,6 +658,9 @@ class CodexLLM(LLMInterface):
|
||||
"Run 'codex auth login' to re-authenticate."
|
||||
) from e
|
||||
|
||||
# Diagnostic dump (opt-in) of the exact request behind any 4xx.
|
||||
dump_request_on_4xx(scope=scope, provider=self.provider, model=self.model, err=e, request=payload)
|
||||
|
||||
# Log the actual error message from the API
|
||||
error_detail = e.response.text[:500] if hasattr(e.response, "text") else str(e)
|
||||
|
||||
@@ -637,7 +756,8 @@ class CodexLLM(LLMInterface):
|
||||
max_retries: int = 5,
|
||||
initial_backoff: float = 1.0,
|
||||
max_backoff: float = 30.0,
|
||||
tool_choice: str | dict[str, Any] = "auto",
|
||||
tool_choice: LLMToolChoice = LLM_TOOL_CHOICE_AUTO,
|
||||
attempt_context: Callable[[], AbstractAsyncContextManager[None]] | None = None,
|
||||
) -> LLMToolCallResult:
|
||||
"""
|
||||
Make API call with tool calling support.
|
||||
@@ -654,7 +774,7 @@ class CodexLLM(LLMInterface):
|
||||
max_retries: Maximum retry attempts.
|
||||
initial_backoff: Initial backoff time in seconds.
|
||||
max_backoff: Maximum backoff time in seconds.
|
||||
tool_choice: How to choose tools - "auto", "none", "required", or a specific function.
|
||||
tool_choice: Canonical tool-selection policy.
|
||||
|
||||
Returns:
|
||||
LLMToolCallResult with content and/or tool_calls.
|
||||
@@ -716,22 +836,21 @@ class CodexLLM(LLMInterface):
|
||||
"instructions": system_instruction,
|
||||
"input": user_messages,
|
||||
"tools": codex_tools,
|
||||
"tool_choice": self._normalize_tool_choice(tool_choice),
|
||||
"tool_choice": (
|
||||
{"type": "function", "name": tool_choice.selected_function_name}
|
||||
if tool_choice.mode is LLMToolChoiceMode.NAMED
|
||||
else tool_choice.mode.value
|
||||
),
|
||||
"parallel_tool_calls": True,
|
||||
"reasoning": {"summary": reasoning_summary},
|
||||
"reasoning": {"effort": self.reasoning_effort, "summary": reasoning_summary},
|
||||
"store": False,
|
||||
"stream": True,
|
||||
"include": ["reasoning.encrypted_content"],
|
||||
"prompt_cache_key": str(uuid.uuid4()),
|
||||
}
|
||||
payload.update(self._extra_body)
|
||||
|
||||
headers = {
|
||||
"Authorization": f"Bearer {self.access_token}",
|
||||
"Content-Type": "application/json",
|
||||
"OpenAI-Account-ID": self.account_id,
|
||||
"User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7)",
|
||||
"Origin": "https://chatgpt.com",
|
||||
}
|
||||
headers = self._build_request_headers()
|
||||
|
||||
url = f"{self.base_url}/codex/responses"
|
||||
|
||||
@@ -744,10 +863,28 @@ class CodexLLM(LLMInterface):
|
||||
# surfaces immediately to keep behavior identical for callers.
|
||||
attempted_refresh_after_auth_error = False
|
||||
|
||||
try:
|
||||
response = await self._client.post(url, json=payload, headers=headers, timeout=120.0)
|
||||
async def _request_attempt(attempt: int) -> tuple[str | None, list[LLMToolCall]]:
|
||||
async with attempt_context() if attempt_context is not None else nullcontext():
|
||||
set_stage(f"llm.codex.tools.attempt={attempt}/2")
|
||||
response = await self._client.post(url, json=payload, headers=headers, timeout=120.0)
|
||||
if response.status_code != 200:
|
||||
# 401/403 on the first attempt may still be recovered by the
|
||||
# reactive token refresh below — don't log those as errors yet.
|
||||
detail = f"Codex API error {response.status_code}: {response.text[:500]}"
|
||||
if response.status_code in (401, 403) and not attempted_refresh_after_auth_error:
|
||||
logger.warning(f"{detail} (will attempt token refresh)")
|
||||
else:
|
||||
logger.error(detail)
|
||||
response.raise_for_status()
|
||||
return await self._parse_sse_tool_stream(response)
|
||||
|
||||
if response.status_code in (401, 403) and not attempted_refresh_after_auth_error:
|
||||
try:
|
||||
try:
|
||||
content, tool_calls = await _request_attempt(1)
|
||||
except httpx.HTTPStatusError as auth_error:
|
||||
response = auth_error.response
|
||||
if response.status_code not in (401, 403) or attempted_refresh_after_auth_error:
|
||||
raise
|
||||
attempted_refresh_after_auth_error = True
|
||||
try:
|
||||
await self._refresh_oauth_tokens(
|
||||
@@ -756,7 +893,7 @@ class CodexLLM(LLMInterface):
|
||||
)
|
||||
headers["Authorization"] = f"Bearer {self.access_token}"
|
||||
logger.info("Codex auth refreshed after auth error; retrying tool-call request once")
|
||||
response = await self._client.post(url, json=payload, headers=headers, timeout=120.0)
|
||||
content, tool_calls = await _request_attempt(2)
|
||||
except CodexRefreshExpiredError as refresh_err:
|
||||
logger.error(
|
||||
"Codex refresh_token is permanently invalid; cannot recover from auth error in tool-call path"
|
||||
@@ -769,16 +906,7 @@ class CodexLLM(LLMInterface):
|
||||
logger.error(
|
||||
f"Codex token refresh attempt failed in tool-call path: {type(refresh_err).__name__}: {refresh_err}"
|
||||
)
|
||||
# Fall through to the normal error path below.
|
||||
|
||||
# Log response details on error
|
||||
if response.status_code != 200:
|
||||
logger.error(f"Codex API error {response.status_code}: {response.text[:500]}")
|
||||
|
||||
response.raise_for_status()
|
||||
|
||||
# Parse SSE for tool calls and content
|
||||
content, tool_calls = await self._parse_sse_tool_stream(response)
|
||||
raise auth_error
|
||||
|
||||
duration = time.time() - start_time
|
||||
metrics = get_metrics_collector()
|
||||
@@ -828,6 +956,8 @@ class CodexLLM(LLMInterface):
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
# Diagnostic dump (opt-in) of the exact request behind any 4xx.
|
||||
dump_request_on_4xx(scope=scope, provider=self.provider, model=self.model, err=e, request=payload)
|
||||
logger.error(f"Codex tool call error: {e}")
|
||||
raise
|
||||
|
||||
@@ -872,8 +1002,13 @@ class CodexLLM(LLMInterface):
|
||||
try:
|
||||
arguments = json.loads(arguments_str)
|
||||
except json.JSONDecodeError:
|
||||
logger.warning(f"Failed to parse tool arguments: {arguments_str}")
|
||||
arguments = {}
|
||||
# Escape-heavy content can emit invalid \escape
|
||||
# sequences (issue #2504); repair before giving up.
|
||||
try:
|
||||
arguments = json.loads(_repair_invalid_json_escapes(arguments_str))
|
||||
except json.JSONDecodeError:
|
||||
logger.warning(f"Failed to parse tool arguments: {arguments_str}")
|
||||
arguments = {}
|
||||
|
||||
tool_calls.append(
|
||||
LLMToolCall(
|
||||
@@ -892,3 +1027,6 @@ class CodexLLM(LLMInterface):
|
||||
"""Clean up HTTP clients."""
|
||||
await self._client.aclose()
|
||||
self._auth_manager.close()
|
||||
|
||||
def supports_attempt_scoped_concurrency(self) -> bool:
|
||||
return True
|
||||
|
||||
@@ -56,6 +56,14 @@ _DEFAULT_REFRESH_MARGIN_SECONDS = 5 * 60
|
||||
# to None and callers proceed uncached, rather than stalling the whole batch.
|
||||
_DEFAULT_CREATE_TIMEOUT_SECONDS = 30.0
|
||||
|
||||
# TTL for the per-step reflect caches created by ``create_incremental``. These
|
||||
# live only for the duration of one reflect (seconds), so the TTL is just a
|
||||
# storage backstop in case the explicit ``delete_session`` at reflect end is
|
||||
# missed (crash / event-loop teardown). Short so orphaned caches age out fast —
|
||||
# storage is billed per token-hour, so a 5-minute cap keeps the cost of a leaked
|
||||
# cache negligible.
|
||||
_DEFAULT_INCREMENTAL_TTL_SECONDS = 5 * 60
|
||||
|
||||
|
||||
@dataclass
|
||||
class _CacheEntry:
|
||||
@@ -92,6 +100,10 @@ class GeminiCacheManager:
|
||||
self._create_timeout_seconds = create_timeout_seconds
|
||||
self._entries: dict[str, _CacheEntry] = {}
|
||||
self._lock = asyncio.Lock()
|
||||
# session_id -> CachedContent names created via ``create_incremental``.
|
||||
# A reflect creates a fresh rolling cache per step under one session id;
|
||||
# ``delete_session`` tears them all down when the reflect finishes.
|
||||
self._sessions: dict[str, list[str]] = {}
|
||||
|
||||
@staticmethod
|
||||
def fingerprint(
|
||||
@@ -229,18 +241,99 @@ class GeminiCacheManager:
|
||||
if entry.name == name:
|
||||
self._entries.pop(key, None)
|
||||
|
||||
async def create_incremental(
|
||||
self,
|
||||
*,
|
||||
session_id: str,
|
||||
model: str,
|
||||
system_instruction: str,
|
||||
contents: list[Any],
|
||||
tools: list[dict[str, Any]] | None = None,
|
||||
) -> str | None:
|
||||
"""Create a fresh CachedContent holding ``system + tools + contents`` and
|
||||
track it under ``session_id`` for later teardown.
|
||||
|
||||
Unlike ``get_or_create``, this does NOT deduplicate by fingerprint: each
|
||||
step of a reflect grows the conversation prefix, so every call is a
|
||||
distinct, single-use cache. The reflect loop creates one per step (each
|
||||
covering the previous step's full input) and reuses it for exactly the
|
||||
next model turn, then supersedes it. All caches for the session are
|
||||
deleted by ``delete_session`` when the reflect ends; the short TTL is
|
||||
only a backstop.
|
||||
|
||||
Returns the cache resource name, or ``None`` when caching is disabled,
|
||||
the prefix is below the model minimum, or the create otherwise fails —
|
||||
callers MUST fall back to an uncached call in that case.
|
||||
"""
|
||||
try:
|
||||
name = await self._create_cache(
|
||||
model=model,
|
||||
system_instruction=system_instruction,
|
||||
tools=tools,
|
||||
contents=contents,
|
||||
ttl_seconds=_DEFAULT_INCREMENTAL_TTL_SECONDS,
|
||||
)
|
||||
except _CacheNotEligible as e:
|
||||
logger.debug(
|
||||
"GeminiCacheManager: incremental prefix not eligible (model=%s, reason=%s) — caller falls back",
|
||||
model,
|
||||
e,
|
||||
)
|
||||
return None
|
||||
except Exception:
|
||||
logger.exception(
|
||||
"GeminiCacheManager: failed to create incremental cache (model=%s); caller falls back",
|
||||
model,
|
||||
)
|
||||
return None
|
||||
if name is not None:
|
||||
self._sessions.setdefault(session_id, []).append(name)
|
||||
return name
|
||||
|
||||
async def delete(self, name: str) -> None:
|
||||
"""Best-effort server-side delete of a single CachedContent.
|
||||
|
||||
Swallows all errors: a failed delete just means the cache ages out on
|
||||
its TTL. Also drops any matching in-process entry.
|
||||
"""
|
||||
self.invalidate(name)
|
||||
try:
|
||||
await self._client.aio.caches.delete(name=name)
|
||||
except Exception:
|
||||
logger.debug("GeminiCacheManager: delete of cache %s failed (will age out on TTL)", name, exc_info=True)
|
||||
|
||||
async def delete_session(self, session_id: str) -> None:
|
||||
"""Delete every CachedContent created for ``session_id`` (reflect teardown).
|
||||
|
||||
Deletes concurrently and best-effort — a reflect must never fail because
|
||||
a cache couldn't be torn down; the short TTL is the backstop.
|
||||
"""
|
||||
names = self._sessions.pop(session_id, [])
|
||||
if not names:
|
||||
return
|
||||
await asyncio.gather(*(self.delete(n) for n in names), return_exceptions=True)
|
||||
|
||||
async def _create_cache(
|
||||
self,
|
||||
*,
|
||||
model: str,
|
||||
system_instruction: str,
|
||||
tools: list[dict[str, Any]] | None = None,
|
||||
contents: list[Any] | None = None,
|
||||
ttl_seconds: int | None = None,
|
||||
) -> str | None:
|
||||
"""Wrap ``client.aio.caches.create`` with the config we want.
|
||||
|
||||
The SDK surface differs slightly across google-genai versions;
|
||||
this implementation targets the >=1.0.0 line where caches live
|
||||
under ``client.aio.caches``.
|
||||
|
||||
``contents`` (already-converted ``genai_types.Content`` turns) is
|
||||
appended after the system_instruction/tools so the cache can hold a
|
||||
growing multi-turn conversation prefix, not just the static prefix —
|
||||
this is what the step-by-step reflect cache relies on. ``ttl_seconds``
|
||||
overrides the manager default (used to give per-step reflect caches a
|
||||
short backstop TTL).
|
||||
"""
|
||||
# Lazy import so this module doesn't require the SDK at import time.
|
||||
from google.genai import types as genai_types
|
||||
@@ -254,8 +347,10 @@ class GeminiCacheManager:
|
||||
# still part of the fingerprint so a schema change keys a fresh cache.
|
||||
config_kwargs: dict[str, Any] = {
|
||||
"system_instruction": system_instruction,
|
||||
"ttl": f"{self._ttl_seconds}s",
|
||||
"ttl": f"{ttl_seconds if ttl_seconds is not None else self._ttl_seconds}s",
|
||||
}
|
||||
if contents:
|
||||
config_kwargs["contents"] = contents
|
||||
if tools:
|
||||
# OpenAI-style {"function": {...}} entries must be converted to
|
||||
# Gemini's Tool/FunctionDeclaration shape before caching.
|
||||
|
||||
@@ -12,16 +12,19 @@ import io
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
from contextlib import AbstractAsyncContextManager, nullcontext
|
||||
from contextvars import ContextVar
|
||||
from typing import Any
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Callable
|
||||
|
||||
from google import genai
|
||||
from google.genai import errors as genai_errors
|
||||
from google.genai import types as genai_types
|
||||
|
||||
from hindsight_api.engine.llm_interface import LLMInterface
|
||||
from hindsight_api.engine.llm_interface import LLM_TOOL_CHOICE_AUTO, LLMInterface, LLMToolChoice, LLMToolChoiceMode
|
||||
from hindsight_api.engine.llm_trace import LLMResponseUsage, stash_response_usage
|
||||
from hindsight_api.engine.llm_wrapper import parse_llm_json
|
||||
from hindsight_api.engine.providers.llm_debug import dump_request_on_4xx
|
||||
from hindsight_api.engine.response_models import LLMToolCall, LLMToolCallResult, TokenUsage
|
||||
from hindsight_api.metrics import get_metrics_collector
|
||||
from hindsight_api.worker.stage import set_stage
|
||||
@@ -63,6 +66,99 @@ def _usage_from_gemini_response(response: Any) -> LLMResponseUsage:
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _GeminiConversation:
|
||||
"""A message list converted to Gemini's request shape."""
|
||||
|
||||
system_instruction: str | None
|
||||
contents: list["genai_types.Content"]
|
||||
|
||||
|
||||
def _convert_messages_to_gemini(msg_list: list[dict[str, Any]]) -> _GeminiConversation:
|
||||
"""Convert OpenAI-style messages to a Gemini (system_instruction, contents) pair.
|
||||
|
||||
Shared by ``call_with_tools`` (request body) and the incremental cache
|
||||
builder so a cached prefix and the live request serialise turns identically —
|
||||
any drift would fingerprint differently and defeat the cache. Consecutive
|
||||
``role="tool"`` messages are grouped into a single ``user`` Content with
|
||||
multiple FunctionResponse parts, matching Gemini's multi-turn requirement.
|
||||
"""
|
||||
system_instruction: str | None = None
|
||||
gemini_contents: list[genai_types.Content] = []
|
||||
pending_tool_names_by_call_id: dict[str, str] = {}
|
||||
i = 0
|
||||
while i < len(msg_list):
|
||||
msg = msg_list[i]
|
||||
role = msg.get("role", "user")
|
||||
content = msg.get("content", "")
|
||||
|
||||
if role != "tool" and pending_tool_names_by_call_id:
|
||||
missing_ids = ", ".join(sorted(pending_tool_names_by_call_id))
|
||||
raise ValueError(f"Gemini assistant tool calls require results before the next message: {missing_ids}")
|
||||
|
||||
if role == "system":
|
||||
system_instruction = (system_instruction + "\n\n" + content) if system_instruction else content
|
||||
i += 1
|
||||
elif role == "tool":
|
||||
parts = []
|
||||
while i < len(msg_list) and msg_list[i].get("role") == "tool":
|
||||
tool_msg = msg_list[i]
|
||||
tool_content = tool_msg.get("content", "")
|
||||
tool_call_id = tool_msg["tool_call_id"]
|
||||
tool_name = pending_tool_names_by_call_id.pop(tool_call_id, None)
|
||||
if tool_name is None:
|
||||
raise ValueError(f"Gemini tool result references unknown tool_call_id {tool_call_id!r}")
|
||||
parts.append(
|
||||
genai_types.Part(
|
||||
function_response=genai_types.FunctionResponse(
|
||||
name=tool_name,
|
||||
response={"result": tool_content},
|
||||
)
|
||||
)
|
||||
)
|
||||
i += 1
|
||||
if pending_tool_names_by_call_id:
|
||||
missing_ids = ", ".join(sorted(pending_tool_names_by_call_id))
|
||||
raise ValueError(f"Gemini assistant tool calls are missing results: {missing_ids}")
|
||||
gemini_contents.append(genai_types.Content(role="user", parts=parts))
|
||||
elif role == "assistant":
|
||||
tool_calls_in_msg = msg.get("tool_calls", [])
|
||||
if tool_calls_in_msg:
|
||||
parts = []
|
||||
if content:
|
||||
parts.append(genai_types.Part(text=content))
|
||||
for tc in tool_calls_in_msg:
|
||||
tool_call_id = tc["id"]
|
||||
fn = tc["function"]
|
||||
fn_name = fn["name"]
|
||||
if tool_call_id in pending_tool_names_by_call_id:
|
||||
raise ValueError(
|
||||
f"Gemini assistant tool call id {tool_call_id!r} must be unique within its turn"
|
||||
)
|
||||
pending_tool_names_by_call_id[tool_call_id] = fn_name
|
||||
fn_args_str = fn.get("arguments", "{}")
|
||||
fn_args = parse_llm_json(fn_args_str)
|
||||
thought_signature = tc.get("thought_signature")
|
||||
fc_kwargs: dict[str, Any] = {"name": fn_name, "args": fn_args}
|
||||
part_kwargs: dict[str, Any] = {"function_call": genai_types.FunctionCall(**fc_kwargs)}
|
||||
if thought_signature:
|
||||
part_kwargs["thought_signature"] = base64.b64decode(thought_signature)
|
||||
parts.append(genai_types.Part(**part_kwargs))
|
||||
gemini_contents.append(genai_types.Content(role="model", parts=parts))
|
||||
else:
|
||||
gemini_contents.append(genai_types.Content(role="model", parts=[genai_types.Part(text=content)]))
|
||||
i += 1
|
||||
else:
|
||||
gemini_contents.append(genai_types.Content(role="user", parts=[genai_types.Part(text=content)]))
|
||||
i += 1
|
||||
|
||||
if pending_tool_names_by_call_id:
|
||||
missing_ids = ", ".join(sorted(pending_tool_names_by_call_id))
|
||||
raise ValueError(f"Gemini assistant tool calls are missing results: {missing_ids}")
|
||||
|
||||
return _GeminiConversation(system_instruction=system_instruction, contents=gemini_contents)
|
||||
|
||||
|
||||
class GeminiLLM(LLMInterface):
|
||||
"""
|
||||
LLM provider for Google Gemini and Vertex AI.
|
||||
@@ -216,6 +312,7 @@ class GeminiLLM(LLMInterface):
|
||||
strict_schema: bool = False,
|
||||
return_usage: bool = False,
|
||||
cached_prefix: str | None = None,
|
||||
attempt_context: Callable[[], AbstractAsyncContextManager[None]] | None = None,
|
||||
) -> Any:
|
||||
"""
|
||||
Make a Gemini/VertexAI API call with retry logic.
|
||||
@@ -329,17 +426,17 @@ class GeminiLLM(LLMInterface):
|
||||
last_exception = None
|
||||
|
||||
for attempt in range(max_retries + 1):
|
||||
if attempt > 0:
|
||||
set_stage(f"llm.gemini.{scope}.attempt={attempt + 1}/{max_retries + 1}")
|
||||
try:
|
||||
response = await asyncio.wait_for(
|
||||
self._client.aio.models.generate_content(
|
||||
model=self.model,
|
||||
contents=gemini_contents,
|
||||
config=generation_config,
|
||||
),
|
||||
timeout=90.0, # Safety net for network hangs; valid slow responses are <90s
|
||||
)
|
||||
async with attempt_context() if attempt_context is not None else nullcontext():
|
||||
set_stage(f"llm.gemini.{scope}.attempt={attempt + 1}/{max_retries + 1}")
|
||||
response = await asyncio.wait_for(
|
||||
self._client.aio.models.generate_content(
|
||||
model=self.model,
|
||||
contents=gemini_contents,
|
||||
config=generation_config,
|
||||
),
|
||||
timeout=90.0, # Safety net for network hangs; valid slow responses are <90s
|
||||
)
|
||||
# Stash usage before parse/validate, which may raise locally
|
||||
# even though the provider charged for these tokens (#2387).
|
||||
stash_response_usage(_usage_from_gemini_response(response))
|
||||
@@ -479,6 +576,17 @@ class GeminiLLM(LLMInterface):
|
||||
logger.error(f"Gemini auth error (HTTP {e.code}), not retrying: {str(e)}")
|
||||
raise
|
||||
|
||||
# Diagnostic dump (opt-in) of the exact request behind any 4xx, captured
|
||||
# before the cache-drop retry below rebuilds the config so we see what failed.
|
||||
dump_request_on_4xx(
|
||||
scope=scope,
|
||||
provider=self.provider,
|
||||
model=self.model,
|
||||
err=e,
|
||||
request=generation_config,
|
||||
messages=gemini_contents,
|
||||
)
|
||||
|
||||
# Cached-request safety net: a stale/invalid/expired CachedContent
|
||||
# (or an incompatibility like cache + tool_config) surfaces as a 400.
|
||||
# Retrying the same cached request can't recover, so on the first
|
||||
@@ -525,8 +633,10 @@ class GeminiLLM(LLMInterface):
|
||||
max_retries: int = 5,
|
||||
initial_backoff: float = 1.0,
|
||||
max_backoff: float = 30.0,
|
||||
tool_choice: str | dict[str, Any] = "auto",
|
||||
tool_choice: LLMToolChoice = LLM_TOOL_CHOICE_AUTO,
|
||||
cached_prefix: str | None = None,
|
||||
cached_prefix_message_count: int = 0,
|
||||
attempt_context: Callable[[], AbstractAsyncContextManager[None]] | None = None,
|
||||
) -> LLMToolCallResult:
|
||||
"""
|
||||
Make a Gemini/VertexAI API call with tool/function calling support.
|
||||
@@ -540,15 +650,22 @@ class GeminiLLM(LLMInterface):
|
||||
max_retries: Maximum retry attempts.
|
||||
initial_backoff: Initial backoff time in seconds.
|
||||
max_backoff: Maximum backoff time in seconds.
|
||||
tool_choice: How to choose tools (Gemini uses "auto" only).
|
||||
tool_choice: Canonical tool-selection policy.
|
||||
cached_prefix: Optional CachedContent resource name (from
|
||||
``GeminiCacheManager.get_or_create`` with ``tools=...``). When
|
||||
set, the system_instruction and tool definitions are assumed
|
||||
``GeminiCacheManager.get_or_create`` or ``create_incremental``).
|
||||
When set, the system_instruction and tool definitions are assumed
|
||||
to live in the cache; this call will skip resending them and
|
||||
the cached prefix is billed at the cached-input rate. The
|
||||
``tools`` argument is still required (the caller may pass
|
||||
an empty list when the cache holds them) so existing call
|
||||
sites don't break.
|
||||
cached_prefix_message_count: Number of leading ``messages`` already
|
||||
baked into ``cached_prefix`` (the step-by-step reflect cache holds
|
||||
a growing conversation prefix, not just system+tools). Only the
|
||||
messages AFTER this index are sent as request contents — the rest
|
||||
come from the cache and bill at the cached rate. 0 means the cache
|
||||
holds only the static prefix (system+tools), so the full
|
||||
conversation is still sent (legacy behaviour).
|
||||
|
||||
Returns:
|
||||
LLMToolCallResult with content and/or tool_calls.
|
||||
@@ -556,86 +673,45 @@ class GeminiLLM(LLMInterface):
|
||||
start_time = time.time()
|
||||
using_cache = cached_prefix is not None
|
||||
|
||||
# Convert tools to Gemini format. When the cache is in use, the
|
||||
# tool definitions are baked into the CachedContent at create time
|
||||
# and the SDK rejects re-sending them alongside ``cached_content``.
|
||||
# Convert tools to Gemini format. While the cache is in use the tool
|
||||
# definitions live in the CachedContent and the SDK rejects re-sending
|
||||
# them alongside ``cached_content`` (see ``_build_tools_config``), but we
|
||||
# still build them unconditionally so the cached-call-failed fallback —
|
||||
# which drops the cache and re-sends prefix + tools inline — has real
|
||||
# tools to send rather than an empty list.
|
||||
gemini_tools = []
|
||||
if not using_cache:
|
||||
for tool in tools:
|
||||
func = tool.get("function", {})
|
||||
gemini_tools.append(
|
||||
genai_types.Tool(
|
||||
function_declarations=[
|
||||
genai_types.FunctionDeclaration(
|
||||
name=func.get("name", ""),
|
||||
description=func.get("description", ""),
|
||||
parameters=func.get("parameters"),
|
||||
)
|
||||
]
|
||||
)
|
||||
)
|
||||
|
||||
# Convert messages
|
||||
system_instruction = None
|
||||
gemini_contents = []
|
||||
msg_list = list(messages)
|
||||
i = 0
|
||||
while i < len(msg_list):
|
||||
msg = msg_list[i]
|
||||
role = msg.get("role", "user")
|
||||
content = msg.get("content", "")
|
||||
|
||||
if role == "system":
|
||||
# Always capture system_instruction. _build_tools_config omits it
|
||||
# (and tools) from the request while the cache carries the prefix,
|
||||
# but it must be available so the cached-call-failed safety net can
|
||||
# re-send the prefix + tools inline.
|
||||
system_instruction = (system_instruction + "\n\n" + content) if system_instruction else content
|
||||
i += 1
|
||||
elif role == "tool":
|
||||
# Gemini requires ALL tool responses for a given model turn to be grouped
|
||||
# into a single Content with multiple FunctionResponse parts.
|
||||
# Consecutive role="tool" messages correspond to one model turn's tool calls.
|
||||
parts = []
|
||||
while i < len(msg_list) and msg_list[i].get("role") == "tool":
|
||||
tool_msg = msg_list[i]
|
||||
tool_content = tool_msg.get("content", "")
|
||||
parts.append(
|
||||
genai_types.Part(
|
||||
function_response=genai_types.FunctionResponse(
|
||||
name=tool_msg.get("name", ""),
|
||||
response={"result": tool_content},
|
||||
)
|
||||
for tool in tools:
|
||||
func = tool.get("function", {})
|
||||
gemini_tools.append(
|
||||
genai_types.Tool(
|
||||
function_declarations=[
|
||||
genai_types.FunctionDeclaration(
|
||||
name=func.get("name", ""),
|
||||
description=func.get("description", ""),
|
||||
parameters=func.get("parameters"),
|
||||
)
|
||||
)
|
||||
i += 1
|
||||
gemini_contents.append(genai_types.Content(role="user", parts=parts))
|
||||
elif role == "assistant":
|
||||
tool_calls_in_msg = msg.get("tool_calls", [])
|
||||
if tool_calls_in_msg:
|
||||
# Convert OpenAI-style tool_calls to Gemini function_call parts
|
||||
# This is required for proper multi-turn conversation history
|
||||
parts = []
|
||||
if content:
|
||||
parts.append(genai_types.Part(text=content))
|
||||
for tc in tool_calls_in_msg:
|
||||
fn = tc.get("function", {})
|
||||
fn_name = fn.get("name", "")
|
||||
fn_args_str = fn.get("arguments", "{}")
|
||||
fn_args = parse_llm_json(fn_args_str)
|
||||
thought_signature = tc.get("thought_signature")
|
||||
fc_kwargs: dict[str, Any] = {"name": fn_name, "args": fn_args}
|
||||
part_kwargs: dict[str, Any] = {"function_call": genai_types.FunctionCall(**fc_kwargs)}
|
||||
if thought_signature:
|
||||
part_kwargs["thought_signature"] = base64.b64decode(thought_signature)
|
||||
parts.append(genai_types.Part(**part_kwargs))
|
||||
gemini_contents.append(genai_types.Content(role="model", parts=parts))
|
||||
else:
|
||||
gemini_contents.append(genai_types.Content(role="model", parts=[genai_types.Part(text=content)]))
|
||||
i += 1
|
||||
else:
|
||||
gemini_contents.append(genai_types.Content(role="user", parts=[genai_types.Part(text=content)]))
|
||||
i += 1
|
||||
]
|
||||
)
|
||||
)
|
||||
|
||||
# Convert messages. ``system_instruction`` and the FULL contents are always
|
||||
# computed: _build_tools_config omits system/tools from the request while
|
||||
# the cache carries the prefix, but the cached-call-failed safety net must
|
||||
# be able to re-send the whole prefix + tools inline.
|
||||
converted = _convert_messages_to_gemini(list(messages))
|
||||
system_instruction = converted.system_instruction
|
||||
full_contents = converted.contents
|
||||
|
||||
# Step-by-step reflect cache: when the cache already holds the first
|
||||
# ``cached_prefix_message_count`` messages, send ONLY the newer turns as
|
||||
# request contents — the cached prefix supplies the rest at the cached
|
||||
# rate. The split is always at a whole-turn boundary (the reflect loop
|
||||
# advances the cache one completed turn at a time), so slicing the raw
|
||||
# messages before conversion never splits a grouped tool turn.
|
||||
if using_cache and cached_prefix_message_count > 0:
|
||||
delta_contents = _convert_messages_to_gemini(list(messages)[cached_prefix_message_count:]).contents
|
||||
else:
|
||||
delta_contents = full_contents
|
||||
|
||||
# Apply safety settings: context var (per-request bank override) takes precedence over instance default
|
||||
effective_safety_settings = _safety_settings_ctx.get()
|
||||
@@ -665,22 +741,20 @@ class GeminiLLM(LLMInterface):
|
||||
config_kwargs["max_output_tokens"] = max_completion_tokens
|
||||
|
||||
# Map OpenAI-style tool_choice to Gemini FunctionCallingConfig
|
||||
if tool_choice == "required":
|
||||
if tool_choice.mode is LLMToolChoiceMode.REQUIRED:
|
||||
config_kwargs["tool_config"] = genai_types.ToolConfig(
|
||||
function_calling_config=genai_types.FunctionCallingConfig(
|
||||
mode="ANY",
|
||||
)
|
||||
)
|
||||
elif isinstance(tool_choice, dict) and tool_choice.get("type") == "function":
|
||||
fn_name = tool_choice.get("function", {}).get("name")
|
||||
if fn_name:
|
||||
config_kwargs["tool_config"] = genai_types.ToolConfig(
|
||||
function_calling_config=genai_types.FunctionCallingConfig(
|
||||
mode="ANY",
|
||||
allowed_function_names=[fn_name],
|
||||
)
|
||||
elif tool_choice.mode is LLMToolChoiceMode.NAMED:
|
||||
config_kwargs["tool_config"] = genai_types.ToolConfig(
|
||||
function_calling_config=genai_types.FunctionCallingConfig(
|
||||
mode="ANY",
|
||||
allowed_function_names=[tool_choice.selected_function_name],
|
||||
)
|
||||
elif tool_choice == "none":
|
||||
)
|
||||
elif tool_choice.mode is LLMToolChoiceMode.NONE:
|
||||
config_kwargs["tool_config"] = genai_types.ToolConfig(
|
||||
function_calling_config=genai_types.FunctionCallingConfig(mode="NONE")
|
||||
)
|
||||
@@ -698,17 +772,21 @@ class GeminiLLM(LLMInterface):
|
||||
|
||||
last_exception = None
|
||||
for attempt in range(max_retries + 1):
|
||||
if attempt > 0:
|
||||
set_stage(f"llm.gemini.tools.attempt={attempt + 1}/{max_retries + 1}")
|
||||
try:
|
||||
response = await asyncio.wait_for(
|
||||
self._client.aio.models.generate_content(
|
||||
model=self.model,
|
||||
contents=gemini_contents,
|
||||
config=config,
|
||||
),
|
||||
timeout=90.0, # Safety net for network hangs; valid slow responses are <90s
|
||||
)
|
||||
# With the cache active, send only the un-cached tail (delta);
|
||||
# on the uncached fallback path send the full conversation so the
|
||||
# re-inlined system+tools prefix has its whole context.
|
||||
active_contents = delta_contents if cache_active else full_contents
|
||||
async with attempt_context() if attempt_context is not None else nullcontext():
|
||||
set_stage(f"llm.gemini.tools.attempt={attempt + 1}/{max_retries + 1}")
|
||||
response = await asyncio.wait_for(
|
||||
self._client.aio.models.generate_content(
|
||||
model=self.model,
|
||||
contents=active_contents,
|
||||
config=config,
|
||||
),
|
||||
timeout=90.0, # Safety net for network hangs; valid slow responses are <90s
|
||||
)
|
||||
stash_response_usage(_usage_from_gemini_response(response))
|
||||
|
||||
# Extract content and tool calls
|
||||
@@ -807,6 +885,17 @@ class GeminiLLM(LLMInterface):
|
||||
logger.error(f"Gemini auth error (HTTP {e.code}), not retrying: {str(e)}")
|
||||
raise
|
||||
|
||||
# Diagnostic dump (opt-in) of the exact request behind any 4xx, captured
|
||||
# before the cache-drop retry below rebuilds the config so we see what failed.
|
||||
dump_request_on_4xx(
|
||||
scope=scope,
|
||||
provider=self.provider,
|
||||
model=self.model,
|
||||
err=e,
|
||||
request=config,
|
||||
messages=active_contents,
|
||||
)
|
||||
|
||||
# Cached-request safety net (see ``call``): a stale/invalid cache or
|
||||
# a cache+tool_config conflict surfaces as a 400. Drop the cache,
|
||||
# invalidate it for later operations, and retry THIS call inline
|
||||
@@ -883,6 +972,56 @@ class GeminiLLM(LLMInterface):
|
||||
tools=tools,
|
||||
)
|
||||
|
||||
# ── Step-by-step incremental prompt caching (reflect tool loop) ──────────
|
||||
|
||||
def supports_incremental_prompt_cache(self) -> bool:
|
||||
"""True when explicit caching is on — the reflect loop can then roll a
|
||||
per-step CachedContent that grows with the conversation."""
|
||||
return self._prompt_cache_enabled
|
||||
|
||||
def _ensure_cache_manager(self) -> Any:
|
||||
if self._cache_manager is None:
|
||||
from hindsight_api.engine.providers.gemini_cache import GeminiCacheManager
|
||||
|
||||
self._cache_manager = GeminiCacheManager(self._client)
|
||||
return self._cache_manager
|
||||
|
||||
async def create_incremental_cache(
|
||||
self,
|
||||
*,
|
||||
session_id: str,
|
||||
messages: list[dict[str, Any]],
|
||||
tools: list[dict[str, Any]] | None = None,
|
||||
) -> str | None:
|
||||
"""Cache ``system + tools + messages`` as a conversation prefix and return
|
||||
its resource name (or ``None`` — caller falls back to an uncached call).
|
||||
|
||||
The reflect loop calls this once per step with the growing message list so
|
||||
each step's cache entirely contains the previous step's input; the next
|
||||
model turn then references it and re-sends only its own delta. Caches are
|
||||
tracked under ``session_id`` and torn down by ``delete_cache_session``.
|
||||
"""
|
||||
if not self._prompt_cache_enabled or self._client is None:
|
||||
return None
|
||||
converted = _convert_messages_to_gemini(list(messages))
|
||||
return await self._ensure_cache_manager().create_incremental(
|
||||
session_id=session_id,
|
||||
model=self.model,
|
||||
system_instruction=converted.system_instruction or "",
|
||||
contents=converted.contents,
|
||||
tools=tools,
|
||||
)
|
||||
|
||||
async def delete_cached_prefix(self, name: str) -> None:
|
||||
"""Best-effort delete of a single CachedContent (superseded reflect step)."""
|
||||
if self._cache_manager is not None:
|
||||
await self._cache_manager.delete(name)
|
||||
|
||||
async def delete_cache_session(self, session_id: str) -> None:
|
||||
"""Tear down every CachedContent created for a reflect session."""
|
||||
if self._cache_manager is not None:
|
||||
await self._cache_manager.delete_session(session_id)
|
||||
|
||||
# ── Batch API (Gemini API only — not Vertex AI) ─────────────────────────
|
||||
#
|
||||
# Google's Gemini Batch API gives a flat 50% discount on input + output
|
||||
@@ -1030,7 +1169,7 @@ class GeminiLLM(LLMInterface):
|
||||
Mirrors the synchronous ``call`` path: system messages become
|
||||
``systemInstruction``; a ``response_format`` json_schema forces JSON
|
||||
output (``responseMimeType``), appends the schema as a textual hint, and
|
||||
grammar-enforces via ``responseJsonSchema`` when ``strict`` is set.
|
||||
grammar-enforces via ``responseJsonSchema`` whenever a schema is present.
|
||||
"""
|
||||
system_texts: list[str] = []
|
||||
contents: list[dict[str, Any]] = []
|
||||
@@ -1059,8 +1198,13 @@ class GeminiLLM(LLMInterface):
|
||||
system_texts.append(
|
||||
"You must respond with valid JSON matching this schema:\n" + json.dumps(schema, ensure_ascii=False)
|
||||
)
|
||||
if json_schema.get("strict"):
|
||||
generation_config["responseJsonSchema"] = schema
|
||||
# #2699: Gemini always grammar-enforces structured output via its native
|
||||
# response_schema (``strict`` is an OpenAI concept, meaningless here). Set
|
||||
# the native schema whenever one is present so the batch path mirrors the
|
||||
# interactive path; otherwise batch requests at default config
|
||||
# (HINDSIGHT_API_LLM_STRICT_SCHEMA=False) get only a textual hint and
|
||||
# intermittently emit malformed JSON, losing every fact in the chunk.
|
||||
generation_config["responseJsonSchema"] = schema
|
||||
|
||||
request: dict[str, Any] = {"contents": contents}
|
||||
if system_texts:
|
||||
@@ -1163,3 +1307,6 @@ class GeminiLLM(LLMInterface):
|
||||
"""Clean up resources (close connections, etc.)."""
|
||||
# Gemini client doesn't require explicit cleanup
|
||||
pass
|
||||
|
||||
def supports_attempt_scoped_concurrency(self) -> bool:
|
||||
return True
|
||||
|
||||
@@ -17,14 +17,24 @@ import json
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from typing import Any
|
||||
from contextlib import AbstractAsyncContextManager, nullcontext
|
||||
from typing import Any, Callable
|
||||
|
||||
from litellm.exceptions import Timeout as LiteLLMTimeout
|
||||
|
||||
from hindsight_api.config import DEFAULT_LLM_TIMEOUT, ENV_LLM_TIMEOUT
|
||||
from hindsight_api.engine.llm_interface import LLMInterface, OutputTooLongError
|
||||
from hindsight_api.engine.llm_interface import (
|
||||
LLM_TOOL_CHOICE_AUTO,
|
||||
LLMInterface,
|
||||
LLMToolChoice,
|
||||
LLMToolChoiceMode,
|
||||
OutputTooLongError,
|
||||
)
|
||||
from hindsight_api.engine.llm_trace import LLMResponseUsage, stash_response_usage
|
||||
from hindsight_api.engine.llm_wrapper import parse_llm_json
|
||||
from hindsight_api.engine.providers.llm_debug import dump_request_on_4xx
|
||||
from hindsight_api.engine.response_models import LLMToolCall, LLMToolCallResult, TokenUsage
|
||||
from hindsight_api.engine.structured_output import strict_json_schema
|
||||
from hindsight_api.metrics import get_metrics_collector
|
||||
from hindsight_api.worker.stage import set_stage
|
||||
|
||||
@@ -226,6 +236,7 @@ class LiteLLMLLM(LLMInterface):
|
||||
skip_validation: bool = False,
|
||||
strict_schema: bool = False,
|
||||
return_usage: bool = False,
|
||||
attempt_context: Callable[[], AbstractAsyncContextManager[None]] | None = None,
|
||||
) -> Any:
|
||||
start_time = time.time()
|
||||
|
||||
@@ -233,7 +244,7 @@ class LiteLLMLLM(LLMInterface):
|
||||
|
||||
# Add JSON schema response format if provided
|
||||
if response_format is not None and hasattr(response_format, "model_json_schema"):
|
||||
schema = response_format.model_json_schema()
|
||||
schema = strict_json_schema(response_format) if strict_schema else response_format.model_json_schema()
|
||||
call_kwargs["response_format"] = {
|
||||
"type": "json_schema",
|
||||
"json_schema": {
|
||||
@@ -246,13 +257,13 @@ class LiteLLMLLM(LLMInterface):
|
||||
last_exception = None
|
||||
|
||||
for attempt in range(max_retries + 1):
|
||||
if attempt > 0:
|
||||
set_stage(f"llm.{self._stage_label}.{scope}.attempt={attempt + 1}/{max_retries + 1}")
|
||||
try:
|
||||
response = await asyncio.wait_for(
|
||||
self._acompletion(**call_kwargs),
|
||||
timeout=self.timeout,
|
||||
)
|
||||
async with attempt_context() if attempt_context is not None else nullcontext():
|
||||
set_stage(f"llm.{self._stage_label}.{scope}.attempt={attempt + 1}/{max_retries + 1}")
|
||||
response = await asyncio.wait_for(
|
||||
self._acompletion(**call_kwargs),
|
||||
timeout=self.timeout,
|
||||
)
|
||||
# Stash usage before the length check and parse/validate below,
|
||||
# which may raise locally even though the provider charged for
|
||||
# these tokens (#2387).
|
||||
@@ -277,7 +288,17 @@ class LiteLLMLLM(LLMInterface):
|
||||
try:
|
||||
json_data = json.loads(clean_content)
|
||||
except json.JSONDecodeError:
|
||||
json_data = json.loads(content)
|
||||
try:
|
||||
json_data = json.loads(content)
|
||||
except json.JSONDecodeError:
|
||||
if attempt < max_retries:
|
||||
# Prefer a clean re-roll first — a fresh generation
|
||||
# usually beats repairing a malformed one.
|
||||
raise
|
||||
# Retry budget spent: structural repair as a last
|
||||
# resort (#2547/#2544). Raises again if unrecoverable,
|
||||
# which the outer handler surfaces loudly.
|
||||
json_data = parse_llm_json(content)
|
||||
|
||||
if skip_validation:
|
||||
result = json_data
|
||||
@@ -378,6 +399,9 @@ class LiteLLMLLM(LLMInterface):
|
||||
logger.error(f"LiteLLM auth error, not retrying: {e}")
|
||||
raise
|
||||
|
||||
# Diagnostic dump (opt-in) of the exact request behind any 4xx.
|
||||
dump_request_on_4xx(scope=scope, provider=self.provider, model=self.model, err=e, request=call_kwargs)
|
||||
|
||||
last_exception = e
|
||||
if attempt < max_retries:
|
||||
# Retry on rate limits, connection errors, server errors
|
||||
@@ -408,23 +432,31 @@ class LiteLLMLLM(LLMInterface):
|
||||
max_retries: int = 5,
|
||||
initial_backoff: float = 1.0,
|
||||
max_backoff: float = 30.0,
|
||||
tool_choice: str | dict[str, Any] = "auto",
|
||||
tool_choice: LLMToolChoice = LLM_TOOL_CHOICE_AUTO,
|
||||
attempt_context: Callable[[], AbstractAsyncContextManager[None]] | None = None,
|
||||
) -> LLMToolCallResult:
|
||||
start_time = time.time()
|
||||
|
||||
call_kwargs = self._build_common_kwargs(messages, max_completion_tokens, temperature)
|
||||
call_kwargs["tools"] = tools
|
||||
call_kwargs["tool_choice"] = tool_choice
|
||||
call_kwargs["tool_choice"] = (
|
||||
{
|
||||
"type": "function",
|
||||
"function": {"name": tool_choice.selected_function_name},
|
||||
}
|
||||
if tool_choice.mode is LLMToolChoiceMode.NAMED
|
||||
else tool_choice.mode.value
|
||||
)
|
||||
|
||||
last_exception = None
|
||||
for attempt in range(max_retries + 1):
|
||||
if attempt > 0:
|
||||
set_stage(f"llm.{self._stage_label}.tools.attempt={attempt + 1}/{max_retries + 1}")
|
||||
try:
|
||||
response = await asyncio.wait_for(
|
||||
self._acompletion(**call_kwargs),
|
||||
timeout=self.timeout,
|
||||
)
|
||||
async with attempt_context() if attempt_context is not None else nullcontext():
|
||||
set_stage(f"llm.{self._stage_label}.tools.attempt={attempt + 1}/{max_retries + 1}")
|
||||
response = await asyncio.wait_for(
|
||||
self._acompletion(**call_kwargs),
|
||||
timeout=self.timeout,
|
||||
)
|
||||
# Stash usage before the tool-call argument parse below, which
|
||||
# can raise json.JSONDecodeError locally even though the provider
|
||||
# already billed for these tokens; without this the error trace
|
||||
@@ -524,6 +556,9 @@ class LiteLLMLLM(LLMInterface):
|
||||
if "401" in error_str or "403" in error_str or "unauthorized" in error_str:
|
||||
raise
|
||||
|
||||
# Diagnostic dump (opt-in) of the exact request behind any 4xx.
|
||||
dump_request_on_4xx(scope=scope, provider=self.provider, model=self.model, err=e, request=call_kwargs)
|
||||
|
||||
last_exception = e
|
||||
if attempt < max_retries:
|
||||
is_retryable = any(
|
||||
@@ -544,3 +579,6 @@ class LiteLLMLLM(LLMInterface):
|
||||
async def cleanup(self) -> None:
|
||||
"""Clean up resources."""
|
||||
pass
|
||||
|
||||
def supports_attempt_scoped_concurrency(self) -> bool:
|
||||
return True
|
||||
|
||||
@@ -19,10 +19,11 @@ import socket
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
from contextlib import AbstractAsyncContextManager
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from typing import Any, Callable
|
||||
|
||||
from hindsight_api.engine.llm_interface import LLMInterface
|
||||
from hindsight_api.engine.llm_interface import LLM_TOOL_CHOICE_AUTO, LLMInterface, LLMToolChoice
|
||||
from hindsight_api.engine.response_models import LLMToolCallResult
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -367,6 +368,7 @@ class LlamaCppLLM(LLMInterface):
|
||||
skip_validation: bool = False,
|
||||
strict_schema: bool = False,
|
||||
return_usage: bool = False,
|
||||
attempt_context: Callable[[], AbstractAsyncContextManager[None]] | None = None,
|
||||
) -> Any:
|
||||
"""Delegate call to the OpenAI-compatible API."""
|
||||
await self._ensure_initialized()
|
||||
@@ -382,6 +384,7 @@ class LlamaCppLLM(LLMInterface):
|
||||
skip_validation=skip_validation,
|
||||
strict_schema=strict_schema,
|
||||
return_usage=return_usage,
|
||||
attempt_context=attempt_context,
|
||||
)
|
||||
|
||||
async def call_with_tools(
|
||||
@@ -394,7 +397,8 @@ class LlamaCppLLM(LLMInterface):
|
||||
max_retries: int = 5,
|
||||
initial_backoff: float = 1.0,
|
||||
max_backoff: float = 30.0,
|
||||
tool_choice: str | dict[str, Any] = "auto",
|
||||
tool_choice: LLMToolChoice = LLM_TOOL_CHOICE_AUTO,
|
||||
attempt_context: Callable[[], AbstractAsyncContextManager[None]] | None = None,
|
||||
) -> LLMToolCallResult:
|
||||
"""Delegate tool calls to the OpenAI-compatible API."""
|
||||
await self._ensure_initialized()
|
||||
@@ -408,8 +412,12 @@ class LlamaCppLLM(LLMInterface):
|
||||
initial_backoff=initial_backoff,
|
||||
max_backoff=max_backoff,
|
||||
tool_choice=tool_choice,
|
||||
attempt_context=attempt_context,
|
||||
)
|
||||
|
||||
def supports_attempt_scoped_concurrency(self) -> bool:
|
||||
return True
|
||||
|
||||
async def cleanup(self) -> None:
|
||||
"""Stop the shared llama.cpp server."""
|
||||
global _shared_server
|
||||
|
||||
@@ -0,0 +1,168 @@
|
||||
"""Opt-in diagnostic: dump the exact request behind an LLM 4xx rejection.
|
||||
|
||||
Some ``400 INVALID_ARGUMENT`` / ``400 Bad Request`` rejections of structured-output
|
||||
calls are not reproducible by reconstructing the request after the fact — the failing
|
||||
factor lives in the request as it was actually assembled at runtime. Reconstructed
|
||||
replays of the same inputs return ``200``, so the only reliable way to see what the
|
||||
model rejected is to capture the real request at the moment it fails.
|
||||
|
||||
This helper is provider-agnostic. Every provider's error handler calls
|
||||
``dump_request_on_4xx`` with whatever it assembled — a Pydantic config
|
||||
(google-genai ``GenerateContentConfig``), a kwargs dict (OpenAI / Anthropic /
|
||||
LiteLLM ``**call_params``), etc. — plus the raised error. The helper self-gates:
|
||||
it is a no-op unless the ``llm_debug_dump_4xx`` config flag
|
||||
(``HINDSIGHT_API_LLM_DEBUG_DUMP_4XX``) is enabled AND the error carries a 4xx
|
||||
status, so callers can drop one unconditional call into each ``except`` block.
|
||||
|
||||
Safety / scope:
|
||||
- Off by default — the config flag is unset in normal operation.
|
||||
- The serialized config omits message bodies (the ``messages``/``contents``/``input``
|
||||
keys are stripped); message previews are length-capped, so an enabled dump can't
|
||||
flood logs or spill large bodies.
|
||||
- Never raises — diagnostics must not break the request path (falls back to ``repr``).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Top-level request keys whose values are message bodies. Stripped from the config
|
||||
# view so the dump never spills large user content — previews are logged separately.
|
||||
_CONTENT_KEYS = ("messages", "contents", "input")
|
||||
|
||||
_PREVIEW_CHARS = 1500
|
||||
_CONFIG_REPR_CAP = 8000
|
||||
_ERR_CAP = 200
|
||||
|
||||
|
||||
def _enabled() -> bool:
|
||||
from hindsight_api.config import get_config
|
||||
|
||||
return bool(get_config().llm_debug_dump_4xx)
|
||||
|
||||
|
||||
def status_code_of(err: Any) -> int | None:
|
||||
"""Best-effort HTTP status of a provider error, across SDK error shapes.
|
||||
|
||||
OpenAI/Anthropic expose ``status_code``; google-genai uses ``code``; some wrap the
|
||||
status on a ``response``. Returns None when no integer status is discoverable.
|
||||
"""
|
||||
for attr in ("status_code", "code", "http_status"):
|
||||
value = getattr(err, attr, None)
|
||||
if isinstance(value, int):
|
||||
return value
|
||||
response = getattr(err, "response", None)
|
||||
if response is not None:
|
||||
value = getattr(response, "status_code", None)
|
||||
if isinstance(value, int):
|
||||
return value
|
||||
return None
|
||||
|
||||
|
||||
def _serialize_config(request: Any) -> str:
|
||||
"""Render the request config to a string without message bodies, never raising."""
|
||||
try:
|
||||
if request is None:
|
||||
return "null"
|
||||
# Pydantic models (google-genai GenerateContentConfig, SDK params objects).
|
||||
dump = getattr(request, "model_dump_json", None)
|
||||
if callable(dump):
|
||||
return dump(exclude_none=True)
|
||||
if isinstance(request, dict):
|
||||
view = {k: v for k, v in request.items() if k not in _CONTENT_KEYS}
|
||||
return json.dumps(view, ensure_ascii=False, default=str)
|
||||
return repr(request)[:_CONFIG_REPR_CAP]
|
||||
except Exception:
|
||||
return repr(request)[:_CONFIG_REPR_CAP]
|
||||
|
||||
|
||||
@dataclass
|
||||
class _MessagePreview:
|
||||
"""A message rendered for the dump: role + extracted text (not yet length-capped)."""
|
||||
|
||||
role: str
|
||||
text: str
|
||||
|
||||
|
||||
def _message_preview(msg: Any) -> _MessagePreview:
|
||||
"""Extract role + text from a message across dict and provider-object shapes."""
|
||||
# OpenAI / Anthropic dict: {"role": ..., "content": str | list[block]}
|
||||
if isinstance(msg, dict):
|
||||
role = str(msg.get("role", "?"))
|
||||
content = msg.get("content")
|
||||
if isinstance(content, str):
|
||||
return _MessagePreview(role, content)
|
||||
if isinstance(content, list):
|
||||
text = ""
|
||||
for block in content:
|
||||
if isinstance(block, dict):
|
||||
text += block.get("text") or ""
|
||||
else:
|
||||
text += getattr(block, "text", "") or ""
|
||||
return _MessagePreview(role, text)
|
||||
return _MessagePreview(role, "" if content is None else str(content))
|
||||
# google-genai Content: role + parts[].text
|
||||
role = str(getattr(msg, "role", "?"))
|
||||
text = ""
|
||||
for part in getattr(msg, "parts", None) or []:
|
||||
text += getattr(part, "text", None) or ""
|
||||
if not text:
|
||||
text = getattr(msg, "content", "") or ""
|
||||
return _MessagePreview(role, text)
|
||||
|
||||
|
||||
def _resolve_messages(request: Any, messages: Any) -> Any:
|
||||
"""Where per-message previews come from: explicit ``messages``, else inside ``request``."""
|
||||
if messages is not None:
|
||||
return messages
|
||||
if isinstance(request, dict):
|
||||
for key in _CONTENT_KEYS:
|
||||
if key in request:
|
||||
return request[key]
|
||||
return []
|
||||
|
||||
|
||||
def dump_request_on_4xx(
|
||||
*,
|
||||
scope: str,
|
||||
provider: str,
|
||||
model: str,
|
||||
err: Any,
|
||||
request: Any = None,
|
||||
messages: Any = None,
|
||||
) -> None:
|
||||
"""Log the exact request behind an LLM 4xx when the diagnostic is enabled.
|
||||
|
||||
No-op unless ``HINDSIGHT_API_LLM_DEBUG_DUMP_4XX`` is truthy and ``err`` carries a
|
||||
4xx status. ``request`` is whatever the provider assembled (a Pydantic config, a
|
||||
kwargs dict, ...); ``messages`` overrides where the per-message previews come from
|
||||
(defaults to the message list found inside ``request``).
|
||||
"""
|
||||
if not _enabled():
|
||||
return
|
||||
code = status_code_of(err)
|
||||
if code is None or not (400 <= code < 500):
|
||||
return
|
||||
try:
|
||||
cfg_repr = _serialize_config(request)
|
||||
summary = []
|
||||
for msg in _resolve_messages(request, messages) or []:
|
||||
m = _message_preview(msg)
|
||||
summary.append({"role": m.role, "chars": len(m.text), "preview": m.text[:_PREVIEW_CHARS]})
|
||||
logger.error(
|
||||
"[LLM_4XX_DUMP] provider=%s model=%s scope=%s code=%s err=%s config=%s contents=%s",
|
||||
provider,
|
||||
model,
|
||||
scope,
|
||||
code,
|
||||
str(err)[:_ERR_CAP],
|
||||
cfg_repr,
|
||||
json.dumps(summary, ensure_ascii=False),
|
||||
)
|
||||
except Exception as dump_exc: # never let diagnostics break the request path
|
||||
logger.warning("[LLM_4XX_DUMP] failed to serialize rejected request: %s", dump_exc)
|
||||
@@ -7,9 +7,10 @@ without making actual API calls to external LLM services.
|
||||
|
||||
import logging
|
||||
from collections.abc import Callable
|
||||
from contextlib import AbstractAsyncContextManager
|
||||
from typing import Any
|
||||
|
||||
from ..llm_interface import LLMInterface
|
||||
from ..llm_interface import LLM_TOOL_CHOICE_AUTO, LLMInterface, LLMToolChoice, LLMToolChoiceMode
|
||||
from ..response_models import LLMToolCall, LLMToolCallResult, TokenUsage
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -91,6 +92,7 @@ class MockLLM(LLMInterface):
|
||||
skip_validation: bool = False,
|
||||
strict_schema: bool = False,
|
||||
return_usage: bool = False,
|
||||
attempt_context: Callable[[], AbstractAsyncContextManager[None]] | None = None,
|
||||
) -> Any:
|
||||
"""
|
||||
Make a mock LLM API call.
|
||||
@@ -200,7 +202,8 @@ class MockLLM(LLMInterface):
|
||||
max_retries: int = 5,
|
||||
initial_backoff: float = 1.0,
|
||||
max_backoff: float = 30.0,
|
||||
tool_choice: str | dict[str, Any] = "auto",
|
||||
tool_choice: LLMToolChoice = LLM_TOOL_CHOICE_AUTO,
|
||||
attempt_context: Callable[[], AbstractAsyncContextManager[None]] | None = None,
|
||||
) -> LLMToolCallResult:
|
||||
"""
|
||||
Make a mock LLM API call with tool/function calling support.
|
||||
@@ -266,7 +269,7 @@ class MockLLM(LLMInterface):
|
||||
else:
|
||||
result = LLMToolCallResult(content="mock response", finish_reason="stop")
|
||||
else:
|
||||
result = LLMToolCallResult(content="mock response", finish_reason="stop")
|
||||
result = self._compliant_tool_call(tools, tool_choice, messages)
|
||||
|
||||
# Set mock token usage on result if not already set
|
||||
if result.input_tokens == 0:
|
||||
@@ -297,6 +300,61 @@ class MockLLM(LLMInterface):
|
||||
|
||||
return result
|
||||
|
||||
@staticmethod
|
||||
def _compliant_tool_call(
|
||||
tools: list[dict[str, Any]],
|
||||
tool_choice: LLMToolChoice,
|
||||
messages: list[dict[str, Any]],
|
||||
) -> LLMToolCallResult:
|
||||
"""Default tool response: simulate a compliant tool-calling model.
|
||||
|
||||
Real providers drive the reflect loop entirely through tool calls -- they
|
||||
honor a forced tool choice, then finish via ``done`` -- and the reflect
|
||||
agent now rejects a turn that yields no tool call at all (a transport that
|
||||
can't tool-call raises ReflectToolCallError). So the mock must behave like a
|
||||
working provider here rather than returning bare "mock response" prose,
|
||||
which used to be salvaged as the answer. Only this default path is affected;
|
||||
tests that script turns via ``_response_callback`` / ``_mock_response`` are not.
|
||||
"""
|
||||
tool_names = {t.get("function", {}).get("name") for t in tools}
|
||||
|
||||
def _mock_query() -> str:
|
||||
for message in reversed(messages):
|
||||
content = message.get("content")
|
||||
if message.get("role") == "user" and isinstance(content, str) and content.strip():
|
||||
return content[:200]
|
||||
return "mock query"
|
||||
|
||||
# Honor a forced retrieval tool so the loop actually runs recall/search and
|
||||
# gathers evidence (populates based_on for tests that assert on it).
|
||||
if tool_choice.mode is LLMToolChoiceMode.NAMED and tool_choice.function_name in {
|
||||
"search_mental_models",
|
||||
"search_observations",
|
||||
"recall",
|
||||
}:
|
||||
return LLMToolCallResult(
|
||||
tool_calls=[
|
||||
LLMToolCall(
|
||||
id="mock_forced",
|
||||
name=tool_choice.function_name,
|
||||
arguments={"reason": "mock", "query": _mock_query()},
|
||||
)
|
||||
],
|
||||
finish_reason="tool_calls",
|
||||
)
|
||||
|
||||
# Auto turn: finish via the done tool, mirroring a model that has gathered
|
||||
# enough. The reflect evidence guardrail handles the empty-bank case (no
|
||||
# evidence -> forced text synthesis on the final iteration).
|
||||
if "done" in tool_names:
|
||||
return LLMToolCallResult(
|
||||
tool_calls=[LLMToolCall(id="mock_done", name="done", arguments={"answer": "mock response"})],
|
||||
finish_reason="tool_calls",
|
||||
)
|
||||
|
||||
# No done tool offered (non-reflect tool call): fall back to plain text.
|
||||
return LLMToolCallResult(content="mock response", finish_reason="stop")
|
||||
|
||||
@staticmethod
|
||||
def _build_mock_facts(messages: list[dict]) -> dict:
|
||||
"""Build a canned fact extraction response from the user message text.
|
||||
|
||||
@@ -8,9 +8,10 @@ it raises a clear error instead of a confusing connection failure.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import Any
|
||||
from contextlib import AbstractAsyncContextManager
|
||||
from typing import Any, Callable
|
||||
|
||||
from ..llm_interface import LLMInterface
|
||||
from ..llm_interface import LLM_TOOL_CHOICE_AUTO, LLMInterface, LLMToolChoice
|
||||
from ..response_models import LLMToolCallResult
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -48,6 +49,7 @@ class NoneLLM(LLMInterface):
|
||||
skip_validation: bool = False,
|
||||
strict_schema: bool = False,
|
||||
return_usage: bool = False,
|
||||
attempt_context: Callable[[], AbstractAsyncContextManager[None]] | None = None,
|
||||
) -> Any:
|
||||
"""Raise LLMNotAvailableError — no LLM is configured."""
|
||||
raise LLMNotAvailableError(
|
||||
@@ -65,7 +67,8 @@ class NoneLLM(LLMInterface):
|
||||
max_retries: int = 5,
|
||||
initial_backoff: float = 1.0,
|
||||
max_backoff: float = 30.0,
|
||||
tool_choice: str | dict[str, Any] = "auto",
|
||||
tool_choice: LLMToolChoice = LLM_TOOL_CHOICE_AUTO,
|
||||
attempt_context: Callable[[], AbstractAsyncContextManager[None]] | None = None,
|
||||
) -> LLMToolCallResult:
|
||||
"""Raise LLMNotAvailableError — no LLM is configured."""
|
||||
raise LLMNotAvailableError(
|
||||
|
||||
@@ -26,9 +26,10 @@ import logging
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
from contextlib import AbstractAsyncContextManager, nullcontext
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from email.utils import parsedate_to_datetime
|
||||
from typing import Any
|
||||
from typing import Any, Callable
|
||||
from urllib.parse import parse_qs, urlparse, urlunparse
|
||||
|
||||
import httpx
|
||||
@@ -36,9 +37,18 @@ from openai import APIConnectionError, APIStatusError, AsyncOpenAI, LengthFinish
|
||||
|
||||
from hindsight_api.config import DEFAULT_LLM_TIMEOUT, ENV_LLM_TIMEOUT
|
||||
from hindsight_api.engine.bank_attribution import apply_bank_attribution
|
||||
from hindsight_api.engine.llm_interface import LLMInterface, OutputTooLongError, ProviderRateLimitResetError
|
||||
from hindsight_api.engine.llm_interface import (
|
||||
LLM_TOOL_CHOICE_AUTO,
|
||||
LLMInterface,
|
||||
LLMToolChoice,
|
||||
LLMToolChoiceMode,
|
||||
OutputTooLongError,
|
||||
ProviderRateLimitResetError,
|
||||
)
|
||||
from hindsight_api.engine.llm_trace import LLMResponseUsage, stash_response_usage
|
||||
from hindsight_api.engine.providers.llm_debug import dump_request_on_4xx
|
||||
from hindsight_api.engine.response_models import LLMToolCall, LLMToolCallResult, TokenUsage
|
||||
from hindsight_api.engine.structured_output import strict_json_schema
|
||||
from hindsight_api.metrics import get_metrics_collector
|
||||
from hindsight_api.worker.stage import set_stage
|
||||
|
||||
@@ -47,17 +57,57 @@ logger = logging.getLogger(__name__)
|
||||
# Seed applied to every Groq request for deterministic behavior
|
||||
DEFAULT_LLM_SEED = 4242
|
||||
JSON_MODE_USER_HINT = "Return valid json only."
|
||||
DEFAULT_VERIFICATION_MAX_COMPLETION_TOKENS = 512
|
||||
|
||||
# Self-hosted OpenAI-compatible servers that advertise tool_choice="required"
|
||||
|
||||
def _validate_ollama_num_ctx(value: Any) -> int | None:
|
||||
"""Validate a native Ollama context-window override."""
|
||||
if value is None:
|
||||
return None
|
||||
if isinstance(value, bool) or not isinstance(value, int):
|
||||
raise ValueError(f"ollama_num_ctx must be a positive integer, got {value!r}")
|
||||
if value < 1:
|
||||
raise ValueError(f"ollama_num_ctx must be >= 1, got {value}")
|
||||
return value
|
||||
|
||||
|
||||
# Provider implementations that advertise tool_choice="required"
|
||||
# but silently ignore it: instead of forcing a tool call they return
|
||||
# finish_reason "stop"/"tool_calls" with an EMPTY tool_calls array and no error.
|
||||
# Reflect's agent loop then sees no tool call, runs synthesis with no retrieval,
|
||||
# and answers "I don't have information" even when the bank holds the answer.
|
||||
# See issues #1563 (LM Studio), #1179 (LM Studio + Qwen), #1877 (vLLM with
|
||||
# --enable-auto-tool-choice). llama-server (the "llamacpp" provider) honors
|
||||
# "required" correctly and is intentionally excluded (#1179).
|
||||
# --enable-auto-tool-choice). The generic OpenAI provider is intentionally not
|
||||
# inferred from its URL: custom OpenAI-compatible endpoints can implement the
|
||||
# required-tool contract, and silently downgrading them changes request semantics.
|
||||
# llama-server (the "llamacpp" provider) honors "required" correctly and is
|
||||
# intentionally excluded (#1179).
|
||||
_TOOL_CHOICE_REQUIRED_UNSUPPORTED_PROVIDERS = frozenset({"lmstudio", "ollama"})
|
||||
|
||||
# Local providers whose OpenAI-compatible surface always lives under a `/v1`
|
||||
# path (LM Studio: http://localhost:1234/v1, Ollama: http://localhost:11434/v1).
|
||||
# For these we know the exact endpoint shape, so a bare host base URL can be
|
||||
# normalized safely. Cloud/proxy endpoints are left untouched — their path is
|
||||
# provider-specific and must be supplied verbatim.
|
||||
_V1_PATH_LOCAL_PROVIDERS = frozenset({"lmstudio", "ollama"})
|
||||
|
||||
|
||||
def _ensure_v1_base_url(base_url: str) -> str:
|
||||
"""Append the OpenAI-compatible ``/v1`` prefix to a bare local base URL.
|
||||
|
||||
LM Studio's server UI advertises its address as ``http://localhost:1234``,
|
||||
so users commonly set ``HINDSIGHT_API_LLM_BASE_URL`` to that bare host. The
|
||||
OpenAI SDK then POSTs to ``<host>/chat/completions`` and LM Studio rejects it
|
||||
with ``Unexpected endpoint or method`` — its OpenAI-compatible routes live
|
||||
under ``/v1``. Only a base URL with no meaningful path (bare host or a lone
|
||||
trailing slash) is rewritten; anything with an explicit path (e.g. a reverse
|
||||
proxy mount or an already-correct ``/v1``) is returned unchanged. See #2922.
|
||||
"""
|
||||
parsed = urlparse(base_url)
|
||||
if parsed.path.strip("/"):
|
||||
return base_url
|
||||
return urlunparse(parsed._replace(path="/v1"))
|
||||
|
||||
|
||||
class ProviderResponseError(RuntimeError):
|
||||
"""Raised when a provider returns a success response without usable content."""
|
||||
@@ -67,23 +117,68 @@ class ProviderResponseError(RuntimeError):
|
||||
self.retryable = retryable
|
||||
|
||||
|
||||
def _is_json(text: str) -> bool:
|
||||
"""True if ``text`` parses as a JSON value."""
|
||||
try:
|
||||
json.loads(text)
|
||||
except (json.JSONDecodeError, ValueError):
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def _outer_json_span(content: str) -> str | None:
|
||||
"""Return the outermost ``{...}`` / ``[...]`` span if it parses as JSON, else None.
|
||||
|
||||
Fallback for responses where fences are partial/absent or the model wrapped
|
||||
the JSON in surrounding prose. Only returned when it is valid JSON so callers
|
||||
never receive a worse candidate than the raw content.
|
||||
"""
|
||||
starts = [i for i in (content.find("{"), content.find("[")) if i >= 0]
|
||||
ends = [i for i in (content.rfind("}"), content.rfind("]")) if i >= 0]
|
||||
if not starts or not ends:
|
||||
return None
|
||||
start, end = min(starts), max(ends)
|
||||
if end <= start:
|
||||
return None
|
||||
candidate = content[start : end + 1].strip()
|
||||
return candidate if _is_json(candidate) else None
|
||||
|
||||
|
||||
def _strip_code_fences(content: str) -> str:
|
||||
"""Strip markdown code fences from LLM response if present.
|
||||
|
||||
Many LLM providers (MiniMax, some Ollama models, Claude via proxies)
|
||||
wrap JSON responses in ```json ... ``` fences even when json_object
|
||||
response format is requested. This strips the fences while preserving
|
||||
the JSON content inside. Returns the original content unchanged if
|
||||
no fences are detected.
|
||||
response format is requested. Fences are detected by line (a closing
|
||||
``` must sit alone on its line) so triple-backticks *inside* JSON string
|
||||
values do not truncate the payload. When the stripped candidate is not
|
||||
valid JSON (partial fence, prose-wrapped output, truncated response), fall
|
||||
back to the outermost parseable JSON span. Returns the original content
|
||||
unchanged if no better candidate is found.
|
||||
"""
|
||||
if "```" not in content:
|
||||
return content
|
||||
try:
|
||||
if "```json" in content:
|
||||
return content.split("```json")[1].split("```")[0].strip()
|
||||
return content.split("```")[1].split("```")[0].strip()
|
||||
except (IndexError, ValueError):
|
||||
return content
|
||||
candidate = content
|
||||
if "```" in content:
|
||||
lines = content.split("\n")
|
||||
# Find first line that starts a code fence (``` optionally followed by language)
|
||||
fence_start = next((i for i, line in enumerate(lines) if line.startswith("```")), None)
|
||||
if fence_start is not None:
|
||||
# Find matching closing fence (``` alone or with trailing whitespace)
|
||||
fence_end = next(
|
||||
(j for j in range(fence_start + 1, len(lines)) if lines[j].strip() == "```"),
|
||||
None,
|
||||
)
|
||||
if fence_end is not None:
|
||||
candidate = "\n".join(lines[fence_start + 1 : fence_end]).strip()
|
||||
|
||||
if _is_json(candidate):
|
||||
return candidate
|
||||
|
||||
# Fence stripping did not yield valid JSON — try to recover the outer JSON span.
|
||||
span = _outer_json_span(content)
|
||||
if span is not None:
|
||||
return span
|
||||
|
||||
return candidate
|
||||
|
||||
|
||||
# Reasoning/thinking tags emitted by extended-thinking models. Some providers
|
||||
@@ -435,6 +530,8 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
timeout: float | None = None,
|
||||
groq_service_tier: str | None = None,
|
||||
extra_body: dict[str, Any] | None = None,
|
||||
*,
|
||||
ollama_num_ctx: int | None = None,
|
||||
**kwargs: Any,
|
||||
):
|
||||
"""
|
||||
@@ -445,10 +542,15 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
api_key: API key (optional for ollama/lmstudio).
|
||||
base_url: Base URL for the API (uses defaults for groq/ollama/lmstudio if empty).
|
||||
model: Model name.
|
||||
reasoning_effort: Reasoning effort level for supported models ("low", "medium", "high").
|
||||
reasoning_effort: Reasoning effort level for supported models
|
||||
("none", "low", "medium", "high"). "none" is required when calling
|
||||
function tools on some reasoning models, which reject every other
|
||||
value — including omitting the parameter entirely.
|
||||
timeout: Request timeout in seconds (uses env var or 120s default).
|
||||
groq_service_tier: Groq service tier ("on_demand", "flex", "auto").
|
||||
extra_body: Extra body params merged into every API call.
|
||||
ollama_num_ctx: Native Ollama context window override. None lets Ollama use
|
||||
the model/server default.
|
||||
**kwargs: Additional provider-specific parameters.
|
||||
"""
|
||||
super().__init__(provider, api_key, base_url, model, reasoning_effort, **kwargs)
|
||||
@@ -503,6 +605,11 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
# lives on a separate control-plane host — see FireworksLLM.
|
||||
self.base_url = "https://api.fireworks.ai/inference/v1"
|
||||
|
||||
# Normalize bare local base URLs (e.g. a user pasting the address shown
|
||||
# in the LM Studio UI) so the OpenAI SDK targets the `/v1` routes. See #2922.
|
||||
if self.provider in _V1_PATH_LOCAL_PROVIDERS and self.base_url:
|
||||
self.base_url = _ensure_v1_base_url(self.base_url)
|
||||
|
||||
# For ollama/lmstudio, use dummy key if not provided
|
||||
if self.provider in ("ollama", "lmstudio") and not self.api_key:
|
||||
self.api_key = "local"
|
||||
@@ -529,6 +636,7 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
# Service tier configuration (from config, not env vars)
|
||||
self.groq_service_tier = groq_service_tier
|
||||
self.openai_service_tier = kwargs.get("openai_service_tier")
|
||||
self.ollama_num_ctx = _validate_ollama_num_ctx(ollama_num_ctx)
|
||||
# User-configured extra body params (merged into every API call)
|
||||
self._config_extra_body = extra_body or {}
|
||||
|
||||
@@ -559,17 +667,17 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
def _drops_tool_choice_required(self) -> bool:
|
||||
"""Whether this endpoint silently ignores ``tool_choice="required"``.
|
||||
|
||||
True for self-hosted OpenAI-compatible servers known to return an empty
|
||||
tool_calls array for "required" instead of forcing a call (#1563/#1179/
|
||||
#1877). Covers LM Studio / Ollama directly, plus any server reached via
|
||||
the generic "openai" provider with a custom ``base_url`` (e.g. a local
|
||||
vLLM endpoint). The real OpenAI API (no base_url override) honors
|
||||
"required", and cloud providers keep their own default base_urls, so both
|
||||
are left untouched.
|
||||
Only explicitly identified provider implementations are classified as
|
||||
unsupported. A custom base URL does not identify endpoint capabilities:
|
||||
an OpenAI-compatible endpoint may correctly enforce required tool calls,
|
||||
and replacing ``required`` with ``auto`` would violate the caller's named
|
||||
tool choice after the tools list has been narrowed.
|
||||
"""
|
||||
if self.provider in _TOOL_CHOICE_REQUIRED_UNSUPPORTED_PROVIDERS:
|
||||
return True
|
||||
return self.provider == "openai" and bool(self.base_url)
|
||||
return self.provider in _TOOL_CHOICE_REQUIRED_UNSUPPORTED_PROVIDERS
|
||||
|
||||
def _verification_max_completion_tokens(self) -> int:
|
||||
"""Return the startup verification budget for OpenAI-compatible gateways."""
|
||||
return DEFAULT_VERIFICATION_MAX_COMPLETION_TOKENS
|
||||
|
||||
async def verify_connection(self) -> None:
|
||||
"""
|
||||
@@ -582,7 +690,7 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
logger.info(f"Verifying connection: {self.provider}/{self.model}")
|
||||
await self.call(
|
||||
messages=[{"role": "user", "content": "Say 'ok'"}],
|
||||
max_completion_tokens=100,
|
||||
max_completion_tokens=self._verification_max_completion_tokens(),
|
||||
max_retries=2,
|
||||
initial_backoff=0.5,
|
||||
max_backoff=2.0,
|
||||
@@ -644,6 +752,11 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
# use the widely-supported max_tokens
|
||||
return "max_tokens"
|
||||
|
||||
def _apply_provider_extra_body_defaults(self, extra_body: dict[str, Any]) -> None:
|
||||
"""Apply provider-specific extra_body defaults while preserving user overrides."""
|
||||
if self.provider == "minimax":
|
||||
extra_body.setdefault("thinking", {"type": "disabled"})
|
||||
|
||||
async def call(
|
||||
self,
|
||||
messages: list[dict[str, str]],
|
||||
@@ -657,6 +770,7 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
skip_validation: bool = False,
|
||||
strict_schema: bool = False,
|
||||
return_usage: bool = False,
|
||||
attempt_context: Callable[[], AbstractAsyncContextManager[None]] | None = None,
|
||||
) -> Any:
|
||||
"""
|
||||
Make an LLM API call with retry logic.
|
||||
@@ -697,6 +811,7 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
skip_validation=skip_validation,
|
||||
scope=scope,
|
||||
return_usage=return_usage,
|
||||
attempt_context=attempt_context,
|
||||
)
|
||||
|
||||
start_time = time.time()
|
||||
@@ -731,6 +846,7 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
|
||||
# Provider-specific parameters
|
||||
extra_body: dict[str, Any] = {**self._config_extra_body}
|
||||
self._apply_provider_extra_body_defaults(extra_body)
|
||||
if self.provider == "groq":
|
||||
call_params["seed"] = DEFAULT_LLM_SEED
|
||||
# Add service_tier if configured
|
||||
@@ -746,7 +862,7 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
if response_format is not None:
|
||||
schema = None
|
||||
if hasattr(response_format, "model_json_schema"):
|
||||
schema = response_format.model_json_schema()
|
||||
schema = strict_json_schema(response_format) if strict_schema else response_format.model_json_schema()
|
||||
|
||||
if strict_schema and schema is not None:
|
||||
# Use OpenAI's strict JSON schema enforcement
|
||||
@@ -789,11 +905,11 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
# Surface attempt count in worker stage so JSON-schema retry loops
|
||||
# are visible from logs (small models on strict structured output
|
||||
# often loop here). Cheap no-op outside worker context.
|
||||
if attempt > 0:
|
||||
set_stage(f"llm.{self.provider}.{scope}.attempt={attempt + 1}/{max_retries + 1}")
|
||||
try:
|
||||
if response_format is not None:
|
||||
response = await self._client.chat.completions.create(**call_params)
|
||||
async with attempt_context() if attempt_context is not None else nullcontext():
|
||||
set_stage(f"llm.{self.provider}.{scope}.attempt={attempt + 1}/{max_retries + 1}")
|
||||
response = await self._client.chat.completions.create(**call_params)
|
||||
# Stash usage before parse/validate, which may raise locally
|
||||
# even though the provider charged for these tokens (#2387).
|
||||
stash_response_usage(_usage_from_openai_response(response))
|
||||
@@ -849,7 +965,9 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
else:
|
||||
result = response_format.model_validate(json_data)
|
||||
else:
|
||||
response = await self._client.chat.completions.create(**call_params)
|
||||
async with attempt_context() if attempt_context is not None else nullcontext():
|
||||
set_stage(f"llm.{self.provider}.{scope}.attempt={attempt + 1}/{max_retries + 1}")
|
||||
response = await self._client.chat.completions.create(**call_params)
|
||||
stash_response_usage(_usage_from_openai_response(response))
|
||||
result, first_choice = _content_or_error(
|
||||
response,
|
||||
@@ -886,7 +1004,9 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
output_tokens = max(0, output_tokens - thoughts_tokens)
|
||||
total_tokens = max(0, total_tokens - thoughts_tokens)
|
||||
|
||||
# Record LLM metrics
|
||||
# Record LLM metrics. ``output_tokens`` is visible-only by now, so
|
||||
# ``thoughts_tokens`` has to be recorded alongside it or the reasoning
|
||||
# half of the billed output reaches no counter at all.
|
||||
metrics = get_metrics_collector()
|
||||
metrics.record_llm_call(
|
||||
provider=self.provider,
|
||||
@@ -896,6 +1016,8 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
success=True,
|
||||
cached_input_tokens=cached_tokens,
|
||||
thoughts_tokens=thoughts_tokens,
|
||||
)
|
||||
|
||||
# Record trace span
|
||||
@@ -964,6 +1086,9 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
logger.error(f"Auth error (HTTP {e.status_code}), not retrying: {str(e)}")
|
||||
raise
|
||||
|
||||
# Diagnostic dump (opt-in) of the exact request behind any 4xx.
|
||||
dump_request_on_4xx(scope=scope, provider=self.provider, model=self.model, err=e, request=call_params)
|
||||
|
||||
_raise_provider_quota_defer(
|
||||
e, provider=self.provider, model=self.model, scope=scope, max_backoff=max_backoff
|
||||
)
|
||||
@@ -1054,7 +1179,8 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
max_retries: int = 5,
|
||||
initial_backoff: float = 1.0,
|
||||
max_backoff: float = 30.0,
|
||||
tool_choice: str | dict[str, Any] = "auto",
|
||||
tool_choice: LLMToolChoice = LLM_TOOL_CHOICE_AUTO,
|
||||
attempt_context: Callable[[], AbstractAsyncContextManager[None]] | None = None,
|
||||
) -> LLMToolCallResult:
|
||||
"""
|
||||
Make an LLM API call with tool/function calling support.
|
||||
@@ -1068,51 +1194,43 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
max_retries: Maximum retry attempts.
|
||||
initial_backoff: Initial backoff time in seconds.
|
||||
max_backoff: Maximum backoff time in seconds.
|
||||
tool_choice: How to choose tools - "auto", "none", "required", or specific function.
|
||||
tool_choice: Canonical tool-selection policy.
|
||||
|
||||
Returns:
|
||||
LLMToolCallResult with content and/or tool_calls.
|
||||
"""
|
||||
start_time = time.time()
|
||||
|
||||
request_tool_choice: str | dict[str, Any] | None = tool_choice
|
||||
|
||||
# Normalize named tool_choice dicts to "required" + filter tools.
|
||||
# Some providers (e.g. LM Studio, Ollama) reject the OpenAI named format
|
||||
# {"type": "function", "function": {"name": "..."}}. The semantics are
|
||||
# identical to tool_choice="required" with the tools list restricted to
|
||||
# just the requested tool, so we apply that transformation where supported.
|
||||
if isinstance(request_tool_choice, dict) and request_tool_choice.get("type") == "function":
|
||||
forced_name = request_tool_choice.get("function", {}).get("name")
|
||||
if forced_name:
|
||||
filtered = [t for t in tools if t.get("function", {}).get("name") == forced_name]
|
||||
if filtered:
|
||||
tools = filtered
|
||||
request_tool_choice = "required"
|
||||
request_tool_choice: str | None
|
||||
if tool_choice.mode is LLMToolChoiceMode.NAMED:
|
||||
forced_name = tool_choice.selected_function_name
|
||||
filtered = [tool for tool in tools if tool.get("function", {}).get("name") == forced_name]
|
||||
if len(filtered) != 1:
|
||||
raise ValueError(
|
||||
f"Named tool_choice must reference exactly one declared tool; "
|
||||
f"found {len(filtered)} definitions for {forced_name!r}"
|
||||
)
|
||||
tools = filtered
|
||||
request_tool_choice = LLMToolChoiceMode.REQUIRED.value
|
||||
elif tool_choice.mode is LLMToolChoiceMode.AUTO:
|
||||
request_tool_choice = None
|
||||
else:
|
||||
request_tool_choice = tool_choice.mode.value
|
||||
|
||||
# DeepSeek accepts tool calls but rejects explicit required/named
|
||||
# tool_choice values. The tools list has already been narrowed for
|
||||
# forced calls, so omitting tool_choice preserves the practical behavior.
|
||||
if "deepseek" in self.model.lower() and request_tool_choice != "auto":
|
||||
if "deepseek" in self.model.lower() and tool_choice.mode is not LLMToolChoiceMode.AUTO:
|
||||
request_tool_choice = None
|
||||
|
||||
# "auto" is the OpenAI API default — omitting tool_choice is semantically
|
||||
# identical. Some providers (e.g. DeepSeek's reasoner pathway, which
|
||||
# deepseek-v4-flash falls into when thinking mode is enabled) reject the
|
||||
# parameter outright, returning HTTP 400 even for value "auto". Sending it
|
||||
# only when the caller asks for a non-default behaviour avoids those 400s
|
||||
# without changing semantics for compliant providers.
|
||||
if request_tool_choice == "auto":
|
||||
request_tool_choice = None
|
||||
|
||||
# vLLM (--enable-auto-tool-choice), LM Studio, Ollama and similar
|
||||
# self-hosted servers silently drop tool_choice="required", returning an
|
||||
# empty tool_calls array instead of forcing a call (#1563/#1179/#1877).
|
||||
# LM Studio and Ollama silently drop tool_choice="required", returning an
|
||||
# empty tool_calls array instead of forcing a call (#1563/#1179).
|
||||
# Downgrade to auto (None) so the model still gets to call a tool. Named
|
||||
# tool_choice dicts were already normalized to "required" + a single
|
||||
# filtered tool above, so the call stays practically forced even under
|
||||
# auto. The real OpenAI API honors "required" and is left untouched.
|
||||
if request_tool_choice == "required" and self._drops_tool_choice_required():
|
||||
# auto. Generic OpenAI-compatible endpoints retain the canonical
|
||||
# ``required`` contract regardless of whether they use a custom base URL.
|
||||
if request_tool_choice == LLMToolChoiceMode.REQUIRED.value and self._drops_tool_choice_required():
|
||||
request_tool_choice = None
|
||||
|
||||
# DeepSeek tool-call replies can carry provider-specific reasoning_content.
|
||||
@@ -1147,8 +1265,16 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
temperature = max(0.01, min(temperature, 1.0))
|
||||
call_params["temperature"] = temperature
|
||||
|
||||
# Set reasoning_effort for reasoning models, matching call(). Omitting it
|
||||
# here is not a neutral default: OpenAI rejects function tools on a
|
||||
# reasoning model unless reasoning_effort is present and set to "none",
|
||||
# so leaving it out fails exactly like sending an unsupported value.
|
||||
if self._supports_reasoning_model():
|
||||
call_params["reasoning_effort"] = self.reasoning_effort
|
||||
|
||||
# Provider-specific parameters
|
||||
extra_body: dict[str, Any] = {**self._config_extra_body}
|
||||
self._apply_provider_extra_body_defaults(extra_body)
|
||||
if self.provider == "groq":
|
||||
call_params["seed"] = DEFAULT_LLM_SEED
|
||||
if extra_body:
|
||||
@@ -1159,10 +1285,10 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
last_exception = None
|
||||
|
||||
for attempt in range(max_retries + 1):
|
||||
if attempt > 0:
|
||||
set_stage(f"llm.{self.provider}.tools.attempt={attempt + 1}/{max_retries + 1}")
|
||||
try:
|
||||
response = await self._client.chat.completions.create(**call_params)
|
||||
async with attempt_context() if attempt_context is not None else nullcontext():
|
||||
set_stage(f"llm.{self.provider}.tools.attempt={attempt + 1}/{max_retries + 1}")
|
||||
response = await self._client.chat.completions.create(**call_params)
|
||||
|
||||
message = response.choices[0].message
|
||||
finish_reason = response.choices[0].finish_reason
|
||||
@@ -1196,6 +1322,8 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
if thoughts_tokens:
|
||||
output_tokens = max(0, output_tokens - thoughts_tokens)
|
||||
|
||||
# See ``call()``: record the reasoning and cached counts too, so no
|
||||
# billed token is dropped from the metrics counters.
|
||||
metrics = get_metrics_collector()
|
||||
metrics.record_llm_call(
|
||||
provider=self.provider,
|
||||
@@ -1205,6 +1333,8 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
success=True,
|
||||
cached_input_tokens=cached_tokens,
|
||||
thoughts_tokens=thoughts_tokens,
|
||||
)
|
||||
|
||||
# Record OpenTelemetry span
|
||||
@@ -1266,6 +1396,10 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
f"not retrying: {_summarize_status_error(e)}"
|
||||
)
|
||||
raise
|
||||
|
||||
# Diagnostic dump (opt-in) of the exact request behind any 4xx.
|
||||
dump_request_on_4xx(scope=scope, provider=self.provider, model=self.model, err=e, request=call_params)
|
||||
|
||||
_raise_provider_quota_defer(
|
||||
e, provider=self.provider, model=self.model, scope=scope, max_backoff=max_backoff
|
||||
)
|
||||
@@ -1302,6 +1436,7 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
skip_validation: bool,
|
||||
scope: str = "memory",
|
||||
return_usage: bool = False,
|
||||
attempt_context: Callable[[], AbstractAsyncContextManager[None]] | None = None,
|
||||
) -> Any:
|
||||
"""
|
||||
Call Ollama using native API with JSON schema enforcement.
|
||||
@@ -1337,9 +1472,10 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
|
||||
# Add optional parameters with optimized defaults for Ollama
|
||||
options: dict[str, Any] = {
|
||||
"num_ctx": 16384, # 16k context window for larger prompts
|
||||
"num_batch": 512, # Optimal batch size for prompt processing
|
||||
}
|
||||
if self.ollama_num_ctx is not None:
|
||||
options["num_ctx"] = self.ollama_num_ctx
|
||||
if max_completion_tokens:
|
||||
options["num_predict"] = max_completion_tokens
|
||||
if temperature is not None:
|
||||
@@ -1355,10 +1491,10 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
|
||||
async with httpx.AsyncClient(timeout=300.0) as client:
|
||||
for attempt in range(max_retries + 1):
|
||||
if attempt > 0:
|
||||
set_stage(f"llm.ollama_native.{scope}.attempt={attempt + 1}/{max_retries + 1}")
|
||||
try:
|
||||
response = await client.post(native_url, json=payload, headers=headers)
|
||||
async with attempt_context() if attempt_context is not None else nullcontext():
|
||||
set_stage(f"llm.ollama_native.{scope}.attempt={attempt + 1}/{max_retries + 1}")
|
||||
response = await client.post(native_url, json=payload, headers=headers)
|
||||
response.raise_for_status()
|
||||
|
||||
result = response.json()
|
||||
@@ -1595,3 +1731,6 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
"""Clean up resources (close OpenAI client connections)."""
|
||||
if hasattr(self, "_client") and self._client:
|
||||
await self._client.close()
|
||||
|
||||
def supports_attempt_scoped_concurrency(self) -> bool:
|
||||
return True
|
||||
|
||||
@@ -0,0 +1,639 @@
|
||||
"""OpenAI Responses API provider (``client.responses.create`` → ``/v1/responses``).
|
||||
|
||||
This provider **always** uses the Responses API — never chat/completions. The
|
||||
chat/completions path rejects ``reasoning_effort`` combined with function tools on
|
||||
some reasoning models — gpt-5.6-terra returns HTTP 400 unless ``reasoning_effort``
|
||||
is exactly ``"none"`` (see #2983). Reflect is a tool-calling search loop, so that
|
||||
constraint forces the whole reflect operation, including the final synthesis, to
|
||||
run with reasoning disabled.
|
||||
|
||||
The Responses API models the chain-of-thought as a first-class reasoning item, so
|
||||
reasoning and function tools coexist. This provider issues requests against
|
||||
``responses.create``, translating the engine's chat-shaped inputs:
|
||||
|
||||
- chat messages → Responses ``input`` items (assistant ``tool_calls`` →
|
||||
``function_call`` items, ``role="tool"`` results → ``function_call_output``),
|
||||
- nested ``{"type":"function","function":{...}}`` tools → flattened
|
||||
``{"type":"function","name":...,"parameters":...}``,
|
||||
- the flat ``reasoning_effort`` scalar → a ``reasoning={"effort": ...}`` object,
|
||||
- ``response_format`` → ``text={"format": {...}}``,
|
||||
|
||||
and reads ``response.output_text`` plus ``function_call`` items out of
|
||||
``response.output``.
|
||||
|
||||
It is OpenAI-only (unlike the multi-vendor ``OpenAICompatibleLLM``, it carries no
|
||||
groq/ollama/deepseek special-casing) and stands alone on ``LLMInterface`` so the
|
||||
"always Responses, never completions" guarantee is structural.
|
||||
|
||||
The conversation is replayed statelessly on every turn (``store=False``, no
|
||||
``previous_response_id``); the model re-derives reasoning each turn rather than
|
||||
resuming a server-side chain. Server-side reasoning reuse across turns is a
|
||||
possible future optimization.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from contextlib import AbstractAsyncContextManager, nullcontext
|
||||
from typing import Any, Callable
|
||||
from urllib.parse import parse_qs, urlparse, urlunparse
|
||||
|
||||
from openai import APIConnectionError, APIStatusError, AsyncOpenAI
|
||||
|
||||
from hindsight_api.config import DEFAULT_LLM_TIMEOUT, ENV_LLM_TIMEOUT
|
||||
from hindsight_api.engine.bank_attribution import apply_bank_attribution
|
||||
from hindsight_api.engine.llm_interface import (
|
||||
LLM_TOOL_CHOICE_AUTO,
|
||||
LLMInterface,
|
||||
LLMToolChoice,
|
||||
LLMToolChoiceMode,
|
||||
OutputTooLongError,
|
||||
)
|
||||
from hindsight_api.engine.llm_trace import LLMResponseUsage, stash_response_usage
|
||||
from hindsight_api.engine.providers.llm_debug import dump_request_on_4xx
|
||||
|
||||
# Provider-agnostic pure helpers (text cleanup, quota-defer parsing, json-mode
|
||||
# hint). These are module-level utilities, not chat/completions behavior.
|
||||
from hindsight_api.engine.providers.openai_compatible_llm import (
|
||||
_ensure_json_word_in_user_message,
|
||||
_raise_provider_quota_defer,
|
||||
_strip_code_fences,
|
||||
_strip_reasoning_tags,
|
||||
)
|
||||
from hindsight_api.engine.response_models import LLMToolCall, LLMToolCallResult, TokenUsage
|
||||
from hindsight_api.engine.structured_output import strict_json_schema
|
||||
from hindsight_api.metrics import get_metrics_collector
|
||||
from hindsight_api.worker.stage import set_stage
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _tools_to_responses(tools: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
||||
"""Convert chat/completions tool schemas to the flattened Responses shape.
|
||||
|
||||
Chat nests the function under a ``function`` key; Responses lifts ``name`` /
|
||||
``description`` / ``parameters`` to the top level. Anything already flattened
|
||||
(or a built-in tool) is passed through unchanged.
|
||||
"""
|
||||
converted: list[dict[str, Any]] = []
|
||||
for tool in tools:
|
||||
function = tool.get("function") if isinstance(tool, dict) else None
|
||||
if tool.get("type") == "function" and isinstance(function, dict):
|
||||
item: dict[str, Any] = {
|
||||
"type": "function",
|
||||
"name": function.get("name"),
|
||||
"parameters": function.get("parameters", {}),
|
||||
}
|
||||
if function.get("description") is not None:
|
||||
item["description"] = function["description"]
|
||||
if "strict" in function:
|
||||
item["strict"] = function["strict"]
|
||||
converted.append(item)
|
||||
else:
|
||||
converted.append(tool)
|
||||
return converted
|
||||
|
||||
|
||||
def _messages_to_responses_input(messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
||||
"""Convert chat/completions messages to Responses ``input`` items.
|
||||
|
||||
- ``role="tool"`` result → ``function_call_output`` keyed by ``call_id``,
|
||||
- assistant ``tool_calls`` → ``function_call`` items (plus a message item when
|
||||
the assistant turn also carried text),
|
||||
- every other message → a role/content message item.
|
||||
"""
|
||||
items: list[dict[str, Any]] = []
|
||||
for message in messages:
|
||||
role = message.get("role")
|
||||
|
||||
if role == "tool":
|
||||
content = message.get("content")
|
||||
items.append(
|
||||
{
|
||||
"type": "function_call_output",
|
||||
"call_id": message.get("tool_call_id"),
|
||||
"output": content if isinstance(content, str) else json.dumps(content, ensure_ascii=False),
|
||||
}
|
||||
)
|
||||
continue
|
||||
|
||||
if role == "assistant" and message.get("tool_calls"):
|
||||
text = message.get("content")
|
||||
if isinstance(text, str) and text:
|
||||
items.append({"role": "assistant", "content": text})
|
||||
for tool_call in message["tool_calls"]:
|
||||
function = tool_call.get("function", {})
|
||||
arguments = function.get("arguments")
|
||||
items.append(
|
||||
{
|
||||
"type": "function_call",
|
||||
"call_id": tool_call.get("id"),
|
||||
"name": function.get("name"),
|
||||
"arguments": arguments if isinstance(arguments, str) else json.dumps(arguments or {}),
|
||||
}
|
||||
)
|
||||
continue
|
||||
|
||||
content = message.get("content")
|
||||
items.append({"role": role, "content": content if content is not None else ""})
|
||||
return items
|
||||
|
||||
|
||||
def _inject_schema_into_input(input_items: list[dict[str, Any]], schema: dict[str, Any]) -> list[dict[str, Any]]:
|
||||
"""Soft structured-output path: append the JSON schema to the first message.
|
||||
|
||||
A leading ``system`` message gets the schema appended, otherwise it is
|
||||
prepended to the first message. Only string-content message items are touched.
|
||||
"""
|
||||
schema_msg = (
|
||||
f"\n\nYou must respond with valid JSON matching this schema:\n"
|
||||
f"{json.dumps(schema, indent=2, ensure_ascii=False)}"
|
||||
)
|
||||
items = [dict(item) for item in input_items]
|
||||
for item in items:
|
||||
if not isinstance(item.get("content"), str):
|
||||
continue
|
||||
if item.get("role") == "system":
|
||||
item["content"] += schema_msg
|
||||
else:
|
||||
item["content"] = schema_msg + "\n\n" + item["content"]
|
||||
return items
|
||||
return items
|
||||
|
||||
|
||||
class OpenAIResponsesLLM(LLMInterface):
|
||||
"""OpenAI provider that talks exclusively to the Responses API.
|
||||
|
||||
OpenAI-only and standalone (does not inherit the multi-vendor
|
||||
chat/completions provider), so it never routes through ``chat.completions``.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
provider: str,
|
||||
api_key: str,
|
||||
base_url: str,
|
||||
model: str,
|
||||
reasoning_effort: str = "low",
|
||||
timeout: float | None = None,
|
||||
extra_body: dict[str, Any] | None = None,
|
||||
default_headers: dict[str, str] | None = None,
|
||||
**kwargs: Any,
|
||||
):
|
||||
"""Initialize the Responses provider.
|
||||
|
||||
Args:
|
||||
provider: Provider name (``"openai-responses"``).
|
||||
api_key: OpenAI API key (required).
|
||||
base_url: Optional custom endpoint; empty uses the OpenAI default.
|
||||
model: Model name (reasoning models — gpt-5.x, o-series — benefit most).
|
||||
reasoning_effort: Effort forwarded as ``reasoning={"effort": ...}`` for
|
||||
reasoning models. Unlike chat/completions, this is sent alongside
|
||||
function tools.
|
||||
timeout: Request timeout in seconds (env var or 120s default).
|
||||
extra_body: Extra body params merged into every request.
|
||||
default_headers: Custom headers passed to the OpenAI SDK client (for
|
||||
operators routing through proxies / request-tracing middleware).
|
||||
**kwargs: Additional provider-specific parameters (e.g. ``openai_service_tier``).
|
||||
"""
|
||||
super().__init__(provider, api_key, base_url, model, reasoning_effort, **kwargs)
|
||||
|
||||
if not self.api_key:
|
||||
raise ValueError(f"API key is required for {self.provider}")
|
||||
|
||||
# OpenAI service tier ("flex" = 50% cheaper / variable latency, "priority",
|
||||
# "default"). Applied as the native Responses ``service_tier`` param.
|
||||
self.openai_service_tier = kwargs.get("openai_service_tier")
|
||||
self._config_extra_body = extra_body or {}
|
||||
self.default_headers = default_headers
|
||||
self.timeout = timeout or float(os.getenv(ENV_LLM_TIMEOUT, str(DEFAULT_LLM_TIMEOUT)))
|
||||
|
||||
# Manual retries (max_retries=0). Extract query params from base_url so an
|
||||
# Azure-style ``?api-version=`` is forwarded as a default query param.
|
||||
client_kwargs: dict[str, Any] = {"api_key": self.api_key, "max_retries": 0}
|
||||
if self.default_headers:
|
||||
client_kwargs["default_headers"] = self.default_headers
|
||||
if self.base_url:
|
||||
parsed = urlparse(self.base_url)
|
||||
if parsed.query:
|
||||
clean_url = urlunparse(parsed._replace(query=""))
|
||||
client_kwargs["base_url"] = clean_url
|
||||
client_kwargs["default_query"] = {k: v[0] for k, v in parse_qs(parsed.query).items()}
|
||||
self.base_url = clean_url
|
||||
else:
|
||||
client_kwargs["base_url"] = self.base_url
|
||||
if self.timeout:
|
||||
client_kwargs["timeout"] = self.timeout
|
||||
|
||||
self._client = AsyncOpenAI(**client_kwargs)
|
||||
logger.info(
|
||||
f"OpenAI Responses client initialized: provider={self.provider}, model={self.model}, "
|
||||
f"base_url={self.base_url or 'default'}"
|
||||
)
|
||||
|
||||
def _supports_reasoning_model(self) -> bool:
|
||||
"""Whether the model is an OpenAI reasoning model (gpt-5.x, o1, o3)."""
|
||||
model_lower = self.model.lower()
|
||||
return any(x in model_lower for x in ["gpt-5", "o1", "o3"])
|
||||
|
||||
async def verify_connection(self) -> None:
|
||||
"""Verify configuration with a minimal Responses call."""
|
||||
try:
|
||||
logger.info(f"Verifying connection: {self.provider}/{self.model}")
|
||||
await self.call(
|
||||
messages=[{"role": "user", "content": "Say 'ok'"}],
|
||||
max_completion_tokens=512,
|
||||
max_retries=2,
|
||||
initial_backoff=0.5,
|
||||
max_backoff=2.0,
|
||||
scope="verification",
|
||||
)
|
||||
logger.info(f"Connection verified: {self.provider}/{self.model}")
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Connection verification failed for {self.provider}/{self.model}: {e}") from e
|
||||
|
||||
async def cleanup(self) -> None:
|
||||
"""Close the OpenAI client connections."""
|
||||
if hasattr(self, "_client") and self._client:
|
||||
await self._client.close()
|
||||
|
||||
@staticmethod
|
||||
def _extract_usage(response: Any) -> TokenUsage:
|
||||
"""Read Responses token usage, splitting reasoning out of visible output.
|
||||
|
||||
The provider folds reasoning tokens into ``output_tokens``; the
|
||||
``TokenUsage`` contract treats output as visible-only and surfaces
|
||||
reasoning separately, so subtract to avoid double-counting.
|
||||
"""
|
||||
usage = getattr(response, "usage", None)
|
||||
input_tokens = int(getattr(usage, "input_tokens", 0) or 0)
|
||||
raw_output = int(getattr(usage, "output_tokens", 0) or 0)
|
||||
cached = int(getattr(getattr(usage, "input_tokens_details", None), "cached_tokens", 0) or 0)
|
||||
reasoning = int(getattr(getattr(usage, "output_tokens_details", None), "reasoning_tokens", 0) or 0)
|
||||
visible_output = max(0, raw_output - reasoning)
|
||||
return TokenUsage(
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=visible_output,
|
||||
total_tokens=input_tokens + visible_output,
|
||||
cached_tokens=cached,
|
||||
thoughts_tokens=reasoning,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _raise_if_truncated(response: Any) -> None:
|
||||
"""Map a ``max_output_tokens`` truncation to ``OutputTooLongError``."""
|
||||
if getattr(response, "status", None) != "incomplete":
|
||||
return
|
||||
details = getattr(response, "incomplete_details", None)
|
||||
if getattr(details, "reason", None) == "max_output_tokens":
|
||||
raise OutputTooLongError(
|
||||
"LLM output exceeded token limits. Input may need to be split into smaller chunks."
|
||||
)
|
||||
|
||||
def _record_success(
|
||||
self,
|
||||
*,
|
||||
scope: str,
|
||||
usage: TokenUsage,
|
||||
duration: float,
|
||||
span_messages: list[dict[str, Any]],
|
||||
response_content: Any,
|
||||
finish_reason: str,
|
||||
tool_calls_dict: list[dict[str, Any]] | None = None,
|
||||
) -> None:
|
||||
"""Record metrics + a trace span for a successful call (shared by both paths)."""
|
||||
get_metrics_collector().record_llm_call(
|
||||
provider=self.provider,
|
||||
model=self.model,
|
||||
scope=scope,
|
||||
duration=duration,
|
||||
input_tokens=usage.input_tokens,
|
||||
output_tokens=usage.output_tokens,
|
||||
success=True,
|
||||
cached_input_tokens=usage.cached_tokens,
|
||||
thoughts_tokens=usage.thoughts_tokens,
|
||||
)
|
||||
|
||||
from hindsight_api.tracing import _serialize_for_span, get_span_recorder
|
||||
|
||||
get_span_recorder().record_llm_call(
|
||||
provider=self.provider,
|
||||
model=self.model,
|
||||
scope=scope,
|
||||
messages=span_messages,
|
||||
response_content=_serialize_for_span(response_content),
|
||||
input_tokens=usage.input_tokens,
|
||||
output_tokens=usage.output_tokens,
|
||||
duration=duration,
|
||||
finish_reason=finish_reason,
|
||||
error=None,
|
||||
cached_tokens=usage.cached_tokens,
|
||||
tool_calls=tool_calls_dict,
|
||||
)
|
||||
|
||||
async def _run_with_retries(
|
||||
self,
|
||||
params: dict[str, Any],
|
||||
*,
|
||||
scope: str,
|
||||
parse: Any,
|
||||
max_retries: int,
|
||||
initial_backoff: float,
|
||||
max_backoff: float,
|
||||
attempt_context: Callable[[], AbstractAsyncContextManager[None]] | None = None,
|
||||
) -> Any:
|
||||
"""Call ``responses.create`` with retries and hand the response to ``parse``.
|
||||
|
||||
``parse(response)`` extracts the result; raising ``json.JSONDecodeError``
|
||||
signals a bad structured reply and triggers a retry. Connection and
|
||||
non-auth 4xx/5xx errors back off and retry; 401/403 fail fast. Provider
|
||||
usage is stashed before ``parse`` runs so an error trace can still attach
|
||||
real token counts (see #2387).
|
||||
"""
|
||||
last_exception: Exception | None = None
|
||||
for attempt in range(max_retries + 1):
|
||||
try:
|
||||
async with attempt_context() if attempt_context is not None else nullcontext():
|
||||
set_stage(f"llm.{self.provider}.{scope}.attempt={attempt + 1}/{max_retries + 1}")
|
||||
response = await self._client.responses.create(**params)
|
||||
usage = self._extract_usage(response)
|
||||
stash_response_usage(
|
||||
LLMResponseUsage(
|
||||
input_tokens=usage.input_tokens,
|
||||
output_tokens=usage.output_tokens,
|
||||
cached_tokens=usage.cached_tokens,
|
||||
)
|
||||
)
|
||||
return parse(response)
|
||||
|
||||
except json.JSONDecodeError as e:
|
||||
last_exception = e
|
||||
logger.warning(
|
||||
f"JSON parse error from Responses reply "
|
||||
f"({self.provider}/{self.model}, scope={scope}, attempt {attempt + 1}/{max_retries + 1}): "
|
||||
f"{str(e)[:200]}"
|
||||
)
|
||||
if attempt < max_retries:
|
||||
await asyncio.sleep(min(initial_backoff * (2**attempt), max_backoff))
|
||||
continue
|
||||
raise
|
||||
|
||||
except APIConnectionError as e:
|
||||
last_exception = e
|
||||
status_code = getattr(e, "status_code", None) or getattr(
|
||||
getattr(e, "response", None), "status_code", None
|
||||
)
|
||||
logger.warning(
|
||||
f"APIConnectionError ({self.provider}/{self.model}, scope={scope}, HTTP {status_code}, "
|
||||
f"attempt {attempt + 1}/{max_retries + 1}): {str(e)[:200]}"
|
||||
)
|
||||
if attempt < max_retries:
|
||||
await asyncio.sleep(min(initial_backoff * (2**attempt), max_backoff))
|
||||
continue
|
||||
raise
|
||||
|
||||
except APIStatusError as e:
|
||||
if e.status_code in (401, 403):
|
||||
logger.error(f"Auth error (HTTP {e.status_code}, {self.provider}/{self.model}), not retrying")
|
||||
raise
|
||||
|
||||
dump_request_on_4xx(scope=scope, provider=self.provider, model=self.model, err=e, request=params)
|
||||
_raise_provider_quota_defer(
|
||||
e, provider=self.provider, model=self.model, scope=scope, max_backoff=max_backoff
|
||||
)
|
||||
|
||||
last_exception = e
|
||||
if attempt < max_retries:
|
||||
logger.warning(
|
||||
f"APIStatusError ({self.provider}/{self.model}, scope={scope}, "
|
||||
f"attempt {attempt + 1}/{max_retries + 1}): HTTP {e.status_code}"
|
||||
)
|
||||
backoff = min(initial_backoff * (2**attempt), max_backoff)
|
||||
jitter = backoff * 0.2 * (2 * (time.time() % 1) - 1)
|
||||
await asyncio.sleep(backoff + jitter)
|
||||
continue
|
||||
logger.error(
|
||||
f"API error after {max_retries + 1} attempts "
|
||||
f"({self.provider}/{self.model}, scope={scope}): HTTP {e.status_code}"
|
||||
)
|
||||
raise
|
||||
|
||||
if last_exception:
|
||||
raise last_exception
|
||||
raise RuntimeError("LLM call failed after all retries with no exception captured")
|
||||
|
||||
async def call(
|
||||
self,
|
||||
messages: list[dict[str, str]],
|
||||
response_format: Any | None = None,
|
||||
max_completion_tokens: int | None = None,
|
||||
temperature: float | None = None,
|
||||
scope: str = "memory",
|
||||
max_retries: int = 10,
|
||||
initial_backoff: float = 1.0,
|
||||
max_backoff: float = 60.0,
|
||||
skip_validation: bool = False,
|
||||
strict_schema: bool = False,
|
||||
return_usage: bool = False,
|
||||
cached_prefix: str | None = None,
|
||||
attempt_context: Callable[[], AbstractAsyncContextManager[None]] | None = None,
|
||||
) -> Any:
|
||||
"""Make a Responses API call with retry logic (see ``LLMInterface.call``)."""
|
||||
start_time = time.time()
|
||||
is_reasoning_model = self._supports_reasoning_model()
|
||||
|
||||
input_items = _messages_to_responses_input(messages)
|
||||
params: dict[str, Any] = {"model": self.model, "input": input_items, "store": False}
|
||||
|
||||
if max_completion_tokens is not None:
|
||||
# Reasoning models spend part of the budget on hidden reasoning; keep
|
||||
# headroom so the visible answer isn't truncated.
|
||||
if is_reasoning_model and max_completion_tokens < 16000:
|
||||
max_completion_tokens = 16000
|
||||
params["max_output_tokens"] = max_completion_tokens
|
||||
if temperature is not None and not is_reasoning_model:
|
||||
params["temperature"] = temperature
|
||||
if is_reasoning_model:
|
||||
params["reasoning"] = {"effort": self.reasoning_effort}
|
||||
if self.openai_service_tier:
|
||||
params["service_tier"] = self.openai_service_tier
|
||||
if self._config_extra_body:
|
||||
params["extra_body"] = {**self._config_extra_body}
|
||||
|
||||
# Structured output: strict grammar-enforced schema, or the soft
|
||||
# schema-in-prompt + json_object fallback.
|
||||
if response_format is not None and hasattr(response_format, "model_json_schema"):
|
||||
if strict_schema:
|
||||
params["text"] = {
|
||||
"format": {
|
||||
"type": "json_schema",
|
||||
"name": "response",
|
||||
"strict": True,
|
||||
"schema": strict_json_schema(response_format),
|
||||
}
|
||||
}
|
||||
else:
|
||||
schema = response_format.model_json_schema()
|
||||
params["input"] = _ensure_json_word_in_user_message(_inject_schema_into_input(input_items, schema))
|
||||
params["text"] = {"format": {"type": "json_object"}}
|
||||
|
||||
apply_bank_attribution(params)
|
||||
|
||||
def parse(response: Any) -> Any:
|
||||
self._raise_if_truncated(response)
|
||||
usage = self._extract_usage(response)
|
||||
duration = time.time() - start_time
|
||||
content = _strip_reasoning_tags(response.output_text or "")
|
||||
|
||||
if response_format is not None:
|
||||
json_data = json.loads(_strip_code_fences(content))
|
||||
result: Any = json_data if skip_validation else response_format.model_validate(json_data)
|
||||
else:
|
||||
result = content
|
||||
|
||||
self._record_success(
|
||||
scope=scope,
|
||||
usage=usage,
|
||||
duration=duration,
|
||||
span_messages=params["input"],
|
||||
response_content=result,
|
||||
finish_reason=getattr(response, "status", "completed"),
|
||||
)
|
||||
if duration > 10.0:
|
||||
logger.info(
|
||||
f"slow llm call: scope={scope}, model={self.provider}/{self.model}, "
|
||||
f"input_tokens={usage.input_tokens}, output_tokens={usage.output_tokens}, time={duration:.3f}s"
|
||||
)
|
||||
return (result, usage) if return_usage else result
|
||||
|
||||
return await self._run_with_retries(
|
||||
params,
|
||||
scope=scope,
|
||||
parse=parse,
|
||||
max_retries=max_retries,
|
||||
initial_backoff=initial_backoff,
|
||||
max_backoff=max_backoff,
|
||||
attempt_context=attempt_context,
|
||||
)
|
||||
|
||||
async def call_with_tools(
|
||||
self,
|
||||
messages: list[dict[str, Any]],
|
||||
tools: list[dict[str, Any]],
|
||||
max_completion_tokens: int | None = None,
|
||||
temperature: float | None = None,
|
||||
scope: str = "tools",
|
||||
max_retries: int = 5,
|
||||
initial_backoff: float = 1.0,
|
||||
max_backoff: float = 30.0,
|
||||
tool_choice: LLMToolChoice = LLM_TOOL_CHOICE_AUTO,
|
||||
cached_prefix: str | None = None,
|
||||
cached_prefix_message_count: int = 0,
|
||||
attempt_context: Callable[[], AbstractAsyncContextManager[None]] | None = None,
|
||||
) -> LLMToolCallResult:
|
||||
"""Make a Responses API call with tools (see ``LLMInterface.call_with_tools``).
|
||||
|
||||
Unlike chat/completions, ``reasoning`` and ``tools`` are sent together —
|
||||
the reason this provider exists.
|
||||
"""
|
||||
start_time = time.time()
|
||||
is_reasoning_model = self._supports_reasoning_model()
|
||||
|
||||
responses_tools = _tools_to_responses(tools)
|
||||
|
||||
request_tool_choice: str | dict[str, Any] | None
|
||||
if tool_choice.mode is LLMToolChoiceMode.NAMED:
|
||||
forced_name = tool_choice.selected_function_name
|
||||
filtered = [tool for tool in responses_tools if tool.get("name") == forced_name]
|
||||
if len(filtered) != 1:
|
||||
raise ValueError(
|
||||
f"Named tool_choice must reference exactly one declared tool; "
|
||||
f"found {len(filtered)} definitions for {forced_name!r}"
|
||||
)
|
||||
responses_tools = filtered
|
||||
request_tool_choice = {"type": "function", "name": forced_name}
|
||||
elif tool_choice.mode is LLMToolChoiceMode.AUTO:
|
||||
request_tool_choice = None
|
||||
else:
|
||||
request_tool_choice = tool_choice.mode.value
|
||||
|
||||
params: dict[str, Any] = {
|
||||
"model": self.model,
|
||||
"input": _messages_to_responses_input(messages),
|
||||
"tools": responses_tools,
|
||||
"store": False,
|
||||
}
|
||||
if request_tool_choice is not None:
|
||||
params["tool_choice"] = request_tool_choice
|
||||
if max_completion_tokens is not None:
|
||||
params["max_output_tokens"] = max_completion_tokens
|
||||
if temperature is not None and not is_reasoning_model:
|
||||
params["temperature"] = temperature
|
||||
if is_reasoning_model:
|
||||
params["reasoning"] = {"effort": self.reasoning_effort}
|
||||
if self.openai_service_tier:
|
||||
params["service_tier"] = self.openai_service_tier
|
||||
if self._config_extra_body:
|
||||
params["extra_body"] = {**self._config_extra_body}
|
||||
|
||||
apply_bank_attribution(params)
|
||||
|
||||
def parse(response: Any) -> LLMToolCallResult:
|
||||
self._raise_if_truncated(response)
|
||||
usage = self._extract_usage(response)
|
||||
duration = time.time() - start_time
|
||||
|
||||
tool_calls: list[LLMToolCall] = []
|
||||
for item in getattr(response, "output", None) or []:
|
||||
if getattr(item, "type", None) != "function_call":
|
||||
continue
|
||||
raw_args = getattr(item, "arguments", "") or ""
|
||||
try:
|
||||
arguments = json.loads(raw_args) if raw_args else {}
|
||||
except json.JSONDecodeError:
|
||||
arguments = {"_raw": raw_args}
|
||||
tool_calls.append(
|
||||
LLMToolCall(id=getattr(item, "call_id", None), name=getattr(item, "name", ""), arguments=arguments)
|
||||
)
|
||||
|
||||
content = response.output_text or None
|
||||
finish_reason = "tool_calls" if tool_calls else getattr(response, "status", "completed")
|
||||
|
||||
self._record_success(
|
||||
scope=scope,
|
||||
usage=usage,
|
||||
duration=duration,
|
||||
span_messages=params["input"],
|
||||
response_content=content,
|
||||
finish_reason=finish_reason,
|
||||
tool_calls_dict=(
|
||||
[{"id": tc.id, "name": tc.name, "arguments": tc.arguments} for tc in tool_calls]
|
||||
if tool_calls
|
||||
else None
|
||||
),
|
||||
)
|
||||
|
||||
return LLMToolCallResult(
|
||||
content=content,
|
||||
tool_calls=tool_calls,
|
||||
finish_reason=finish_reason,
|
||||
input_tokens=usage.input_tokens,
|
||||
output_tokens=usage.output_tokens,
|
||||
cached_tokens=usage.cached_tokens,
|
||||
thoughts_tokens=usage.thoughts_tokens,
|
||||
)
|
||||
|
||||
return await self._run_with_retries(
|
||||
params,
|
||||
scope=scope,
|
||||
parse=parse,
|
||||
max_retries=max_retries,
|
||||
initial_backoff=initial_backoff,
|
||||
max_backoff=max_backoff,
|
||||
attempt_context=attempt_context,
|
||||
)
|
||||
|
||||
def supports_attempt_scoped_concurrency(self) -> bool:
|
||||
return True
|
||||
@@ -6,6 +6,7 @@ structured information like temporal constraints.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import re
|
||||
from abc import ABC, abstractmethod
|
||||
from datetime import datetime, timedelta
|
||||
|
||||
@@ -19,6 +20,103 @@ from hindsight_api.engine.temporal_periods import (
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# dateparser.search_dates over-matches: short common words that happen to be
|
||||
# weekday/month abbreviations in *some* language ("we"/"me"/"did" -> a weekday,
|
||||
# "do" -> Sunday) come back as bogus dates. When such a false positive appears
|
||||
# *before* the real date in the query, taking the first match (or a hard-coded
|
||||
# blacklist of such words) silently produces a wrong temporal window — worse
|
||||
# than none, because the constraint is non-null so nothing downstream can tell
|
||||
# extraction failed. See issue #2768.
|
||||
#
|
||||
# Instead of blacklisting words one at a time (a moving target — every short
|
||||
# word dateparser resolves is a new instance of the same bug), we score each
|
||||
# match by the date signal it actually carries and keep only matches with a
|
||||
# real signal, preferring the strongest. A bare weekday abbreviation carries no
|
||||
# day/month/year and scores zero, so it is rejected regardless of language or
|
||||
# dateparser version.
|
||||
_TOKEN_RE = re.compile(r"[a-z0-9]+")
|
||||
_MONTH_WORDS = {
|
||||
"january",
|
||||
"february",
|
||||
"march",
|
||||
"april",
|
||||
"may",
|
||||
"june",
|
||||
"july",
|
||||
"august",
|
||||
"september",
|
||||
"october",
|
||||
"november",
|
||||
"december",
|
||||
}
|
||||
_RELATIVE_WORDS = {"today", "yesterday", "tomorrow", "tonight", "now"}
|
||||
_WEEKDAY_WORDS = {
|
||||
"monday",
|
||||
"tuesday",
|
||||
"wednesday",
|
||||
"thursday",
|
||||
"friday",
|
||||
"saturday",
|
||||
"sunday",
|
||||
}
|
||||
_PERIOD_WORDS = {
|
||||
"last",
|
||||
"next",
|
||||
"this",
|
||||
"past",
|
||||
"coming",
|
||||
"ago",
|
||||
"week",
|
||||
"weeks",
|
||||
"month",
|
||||
"months",
|
||||
"year",
|
||||
"years",
|
||||
"day",
|
||||
"days",
|
||||
"hour",
|
||||
"hours",
|
||||
"minute",
|
||||
"minutes",
|
||||
"quarter",
|
||||
"decade",
|
||||
"century",
|
||||
"weekend",
|
||||
"morning",
|
||||
"afternoon",
|
||||
"evening",
|
||||
"night",
|
||||
"noon",
|
||||
"midnight",
|
||||
}
|
||||
|
||||
|
||||
def _date_match_score(text: str) -> int:
|
||||
"""Score how strong a temporal signal a matched span carries.
|
||||
|
||||
A score of 0 means the span is a bare token with no explicit date content
|
||||
(the false-positive class from issue #2768) and should be rejected. Higher
|
||||
scores mean a stronger, less ambiguous date reference. A digit is the
|
||||
strongest signal (day/year/ISO date); an explicit English month/relative
|
||||
word next; weekday names and period words weakest but still explicit.
|
||||
"""
|
||||
tokens = _TOKEN_RE.findall(text.lower())
|
||||
if not tokens:
|
||||
return 0
|
||||
score = 0
|
||||
if any(any(ch.isdigit() for ch in tok) for tok in tokens):
|
||||
score += 100
|
||||
token_set = set(tokens)
|
||||
if token_set & _MONTH_WORDS:
|
||||
score += 50
|
||||
if token_set & _RELATIVE_WORDS:
|
||||
score += 50
|
||||
if token_set & _WEEKDAY_WORDS:
|
||||
score += 30
|
||||
if token_set & _PERIOD_WORDS:
|
||||
score += 20
|
||||
return score
|
||||
|
||||
|
||||
class TemporalConstraint(BaseModel):
|
||||
"""
|
||||
@@ -92,9 +190,25 @@ class DateparserQueryAnalyzer(QueryAnalyzer):
|
||||
- No model loading required (lazy import on first use)
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize dateparser query analyzer."""
|
||||
def __init__(self, languages: list[str] | None = None):
|
||||
"""Initialize dateparser query analyzer.
|
||||
|
||||
Args:
|
||||
languages: Restrict dateparser's language detection to these codes
|
||||
(e.g. ["en"]). None (default) keeps full auto-detection across
|
||||
all 200+ locales — unchanged behavior.
|
||||
"""
|
||||
self._search_dates = None
|
||||
self._languages = languages
|
||||
|
||||
def _search_kwargs(self) -> dict:
|
||||
"""Extra kwargs for search_dates, shared by load() and analyze().
|
||||
|
||||
Both call sites must use the same locale set: warming up under
|
||||
auto-detection while running restricted (or vice versa) leaves part of
|
||||
the lazy-load cost on the first real query.
|
||||
"""
|
||||
return {} if self._languages is None else {"languages": self._languages}
|
||||
|
||||
def load(self) -> None:
|
||||
"""Load dateparser and warm up internal data structures.
|
||||
@@ -107,7 +221,7 @@ class DateparserQueryAnalyzer(QueryAnalyzer):
|
||||
|
||||
self._search_dates = search_dates
|
||||
# Warm up: fire a dummy call to trigger lazy-loaded internal tables.
|
||||
self._search_dates("today")
|
||||
self._search_dates("today", **self._search_kwargs())
|
||||
|
||||
def analyze(self, query: str, reference_date: datetime | None = None) -> QueryAnalysis:
|
||||
"""
|
||||
@@ -152,7 +266,7 @@ class DateparserQueryAnalyzer(QueryAnalyzer):
|
||||
# treat any failure as "no temporal constraint found" so the caller
|
||||
# can fall back to non-temporal retrieval.
|
||||
try:
|
||||
results = self._search_dates(query, settings=settings)
|
||||
results = self._search_dates(query, settings=settings, **self._search_kwargs())
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
"dateparser raised %s on query (treating as no temporal constraint): %s",
|
||||
@@ -164,20 +278,23 @@ class DateparserQueryAnalyzer(QueryAnalyzer):
|
||||
if not results:
|
||||
return QueryAnalysis(temporal_constraint=None)
|
||||
|
||||
# Filter out false positives (common words parsed as dates)
|
||||
false_positives = {"do", "may", "march", "will", "can", "sat", "sun", "mon", "tue", "wed", "thu", "fri"}
|
||||
valid_results = [
|
||||
(text, date)
|
||||
# Score each match by the date signal it carries and keep only those
|
||||
# with a real signal, rejecting bare weekday/month-abbreviation false
|
||||
# positives ("we"/"me"/"did"). Prefer the strongest match, breaking ties
|
||||
# by longest span, so an explicit date ("in May", "2026-06-10") always
|
||||
# beats an earlier weak word regardless of position. See issue #2768.
|
||||
scored_results = [
|
||||
(_date_match_score(text), len(text), date)
|
||||
for text, date in results
|
||||
if (text.lower() not in false_positives or len(text) > 3)
|
||||
and not is_embedded_cjk_dateparser_match(query, text)
|
||||
if not is_embedded_cjk_dateparser_match(query, text)
|
||||
]
|
||||
scored_results = [entry for entry in scored_results if entry[0] > 0]
|
||||
|
||||
if not valid_results:
|
||||
if not scored_results:
|
||||
return QueryAnalysis(temporal_constraint=None)
|
||||
|
||||
# Use the first valid date found
|
||||
_, parsed_date = valid_results[0]
|
||||
# Highest signal score wins; ties broken by the longest matched span.
|
||||
_, _, parsed_date = max(scored_results, key=lambda entry: (entry[0], entry[1]))
|
||||
|
||||
# Create constraint for single day
|
||||
start_date = parsed_date.replace(hour=0, minute=0, second=0, microsecond=0)
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user