Compare commits
125
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
aa4d7028ac | ||
|
|
acd565c506 | ||
|
|
1c237c3cc8 | ||
|
|
15e153f340 | ||
|
|
0f281e1d1b | ||
|
|
66ec76c722 | ||
|
|
c917dd5108 | ||
|
|
aae106df1d | ||
|
|
3d0546672c | ||
|
|
891465edaa | ||
|
|
90aef22052 | ||
|
|
edd0d0c5bb | ||
|
|
5f137bf391 | ||
|
|
6692e38c80 | ||
|
|
df8ac42b52 | ||
|
|
d98990b46e | ||
|
|
2e2dfe1309 | ||
|
|
19318ac088 | ||
|
|
e256a7409b | ||
|
|
0dd32a3d6d | ||
|
|
27e4b188d7 | ||
|
|
a04f4eed3a | ||
|
|
e181fb75c3 | ||
|
|
3e6a812f71 | ||
|
|
6ea34d830f | ||
|
|
f8b3988cf8 | ||
|
|
d69c53739c | ||
|
|
997f27f1bf | ||
|
|
b64943d195 | ||
|
|
b46f9694f6 | ||
|
|
803a45171d | ||
|
|
3770d62e34 | ||
|
|
3650a6e173 | ||
|
|
7850efcef6 | ||
|
|
b919b97df3 | ||
|
|
58d97444b3 | ||
|
|
f970174213 | ||
|
|
b8947a24e9 | ||
|
|
435f1640bb | ||
|
|
9da3dc9eb8 | ||
|
|
8fbdc6bf79 | ||
|
|
0480f171e4 | ||
|
|
ec9cc702ec | ||
|
|
205e47b4e9 | ||
|
|
396f63aafc | ||
|
|
a6c2d90eec | ||
|
|
2e8c221c54 | ||
|
|
28760f62d4 | ||
|
|
32b90cc982 | ||
|
|
a37257ede5 | ||
|
|
f660a4ad85 | ||
|
|
6956dc0661 | ||
|
|
217fa475e4 | ||
|
|
e32e6f2468 | ||
|
|
41b292e574 | ||
|
|
3c24979a40 | ||
|
|
449c9c9eed | ||
|
|
e5b49eb672 | ||
|
|
6c58578139 | ||
|
|
df8244873a | ||
|
|
f28a8458ed | ||
|
|
7edfd6f1d7 | ||
|
|
a763109d88 | ||
|
|
77bdda7042 | ||
|
|
39ee97a61c | ||
|
|
9f2899a936 | ||
|
|
f36a462d1d | ||
|
|
00acde0159 | ||
|
|
5d6da3cf7e | ||
|
|
27d4386deb | ||
|
|
6c8f4be6e7 | ||
|
|
c43a3ebfab | ||
|
|
4afe43e63d | ||
|
|
bf943e4ab7 | ||
|
|
e8c42f2a74 | ||
|
|
32c5d65773 | ||
|
|
e6fb5d4799 | ||
|
|
383b1d2017 | ||
|
|
257df73d5c | ||
|
|
650da8c1eb | ||
|
|
4e5b4ceecc | ||
|
|
0c2e384287 | ||
|
|
3e1b4611cf | ||
|
|
5e846b4917 | ||
|
|
bad06de5ad | ||
|
|
6c65c374cc | ||
|
|
5cef095960 | ||
|
|
45713adcc6 | ||
|
|
06b208fd99 | ||
|
|
ec4b68ebf1 | ||
|
|
a4c1593340 | ||
|
|
229eefb0c9 | ||
|
|
da8077adeb | ||
|
|
681a79e6b4 | ||
|
|
500ab70f13 | ||
|
|
e5025aaf5b | ||
|
|
77bbb93940 | ||
|
|
64ede36172 | ||
|
|
8b33bfdd28 | ||
|
|
098362d450 | ||
|
|
c094ac27d2 | ||
|
|
8f16473a6a | ||
|
|
8391296af5 | ||
|
|
9032ed9c95 | ||
|
|
dca0025511 | ||
|
|
4fa25f47d2 | ||
|
|
423e25db44 | ||
|
|
4e4b87b445 | ||
|
|
aaa58f2472 | ||
|
|
82859af02b | ||
|
|
e8b817fc6e | ||
|
|
83a080f90a | ||
|
|
d3f97da1e4 | ||
|
|
ba365c723b | ||
|
|
29cd6c6ab0 | ||
|
|
490cc52432 | ||
|
|
07aec3b95d | ||
|
|
1171ca276a | ||
|
|
96bd69c7bd | ||
|
|
5781d28d8f | ||
|
|
dbe0ffb989 | ||
|
|
11a4eda803 | ||
|
|
a133cc1495 | ||
|
|
25c7cd0449 | ||
|
|
1d0dcbce25 |
@@ -77,6 +77,18 @@ results = await asyncio.gather(*tasks, return_exceptions=True)
|
||||
- **No direct database access in `api/http.py`** (or any API router). HTTP handlers must not build SQL, call `acquire_with_retry` / `conn.fetch` / `conn.fetchrow` / `conn.execute`, or reference `fq_table(...)`. All persistence and queries live in `MemoryEngine` (the engine layer). A handler parses/validates the request, calls an engine method, shapes the HTTP response, and maps domain results to status codes (e.g. a `None` return → 404).
|
||||
- **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).
|
||||
- **Every list endpoint paginates, following the existing ones.** A `GET` that returns a collection whose size grows with the data (banks, documents, memories, entities, operations, webhook deliveries, audit logs, …) must take `limit`/`offset` and bound its result — an unbounded list is an unbounded payload plus unbounded per-row work (per-item counts, config resolution, embedding hydration). Copy the shape `list_documents` uses, don't invent a new one: `limit: int = Query(default=100, ge=0)` and `offset: int = Query(default=0, ge=0)` on the handler, matching keyword args on the engine method, and a response carrying the page **plus `total`, `limit`, `offset`** so a client knows when to stop. Add a `q` search param when the collection is something a user picks from in a UI — client-side filtering only ever sees the loaded page. Bounded-by-construction endpoints are the exception, not the rule: a tree/export that is whole-structure by design, or a table capped at write time (e.g. `observation_history` / `mental_model_history`, trimmed to `*_max_entries` on insert). If it isn't bounded, paginate it.
|
||||
|
||||
### Bank/Tenant Isolation in Queries
|
||||
- **Bank isolation is a hard security invariant: no query may read, count, update, or delete another bank's rows.** Tenant isolation is enforced at the schema level (the resolved `search_path` / `fq_table(...)` qualifier, gated by `_authenticate_tenant`); bank isolation is enforced *within* a schema by a `bank_id` predicate on every statement that touches a multi-bank table.
|
||||
- **Every SQL statement against a multi-bank table must be constrained by `bank_id`** — directly in the `WHERE`, or transitively (see below). Multi-bank tables carry a `bank_id` column: `memory_units`, `documents`, `entities`, `entity_links`, `mental_models`, `knowledge_pages`, `memory_links`, `observation_history`, and similar.
|
||||
- **The trap: filtering by a caller-supplied, non-globally-unique key without `bank_id`.** Keys like `document_id` and `mental_models.id` are unique only *per bank* (their PK is composite, e.g. `(id, bank_id)`), so the *same* id legally exists in every bank. A statement like `UPDATE memory_units SET tags = $1 WHERE document_id = $2` — no `bank_id` — silently reads/writes **every** bank's rows that share the id. This is the exact defect from #3429/#3430. Adding `AND bank_id = $n` fixes it.
|
||||
- **Three ways a statement is legitimately scoped** (accept these; flag anything that fits none):
|
||||
1. **Explicit** `WHERE ... AND bank_id = $n`.
|
||||
2. **Globally-unique single-column PK.** Filtering by a global uuid PK (`memory_units.id`, `entities.id`, `knowledge_pages.id`) or a bank-encoded key (`chunks.chunk_id` is `{bank_id}_{document_id}_{idx}`) cannot collide across banks. Contrast the *composite*-PK ids (`documents.id`/`document_id`, `mental_models.id`) — those are dangerous and MUST carry `bank_id`.
|
||||
3. **Transitive.** Junction tables without a `bank_id` column (`unit_entities`, `entity_cooccurrences`, `observation_sources`) are safe only when reached through globally-unique unit/entity ids that were themselves selected from a bank-scoped query in the same call, and edges are intra-bank by construction. If the id set could contain another bank's ids, it is not scoped.
|
||||
- **Watch two smells:** (a) a caller-supplied id used in the `WHERE` with no adjacent `bank_id`, while a *neighbouring* statement in the same method does carry `bank_id` (asymmetry is the tell); (b) a `bank_id` predicate applied only under `if bank_id:` with a `bank_id: str | None = None` default — latent even if all current callers pass one.
|
||||
- **Cross-bank by design must rewrite `bank_id` to the destination.** The transfer/import path is the only one that legitimately crosses banks; verify every write pins the *destination* `bank_id` and never inherits a source row's `bank_id`.
|
||||
|
||||
### 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.
|
||||
@@ -152,6 +164,28 @@ Flag any new logic that lacks test coverage.
|
||||
|
||||
See CLAUDE.md → Key Conventions → Testing for the full pattern.
|
||||
|
||||
### 6a. Check tests assert memory state via the engine API, not raw SQL
|
||||
|
||||
Tests must verify what a retain / recall / consolidation produced by calling the public
|
||||
`MemoryEngine` read API — `list_memory_units` (units and their `metadata` / `tags`; counts via
|
||||
`total`; `document_id` / `fact_type` / `entity_id` filters), `list_entities` (canonical names,
|
||||
mention counts), `get_graph_data` (nodes/edges), `get_bank_stats`, `recall_async` — **not** by
|
||||
reaching into the memory tables (`memory_units`, `memory_links`, `unit_entities`) with raw SQL via
|
||||
`pool.acquire()` / `conn.fetch*`. Asserting on those tables couples the test to a storage-layer
|
||||
detail and checks a proxy instead of the observable property (see **General Principles** → tests
|
||||
assert the property, and the handler rule in **7b**).
|
||||
|
||||
**Flag as should fix** any added or changed test whose assertion runs a `SELECT` / `COUNT` against
|
||||
`memory_units` / `memory_links` / `unit_entities` where an engine read method returns the same
|
||||
fact. Prime tell: `async with pool.acquire() as conn:` followed by `SELECT ... FROM memory_units`
|
||||
inside a test body; a `fetchval("SELECT count(*) FROM memory_units ...")` that `list_memory_units`
|
||||
`["total"]` would return; a `canonical_name` query that `list_entities` covers.
|
||||
|
||||
Direct SQL on those tables is legitimate **only** when it forces or inspects internal state the
|
||||
public API cannot express — e.g. an `UPDATE documents SET updated_at` that forges a race, or a
|
||||
raw `memory_links` row-count that the deduped `get_graph_data` edge list cannot reproduce. Those
|
||||
must carry a comment saying why the direct access is necessary; flag any that do not.
|
||||
|
||||
### 7. Check API consistency
|
||||
|
||||
If any files in `hindsight-api-slim/hindsight_api/api/` were changed:
|
||||
@@ -177,6 +211,26 @@ For each changed handler in `hindsight-api-slim/hindsight_api/api/` (e.g. `http.
|
||||
- **Flag any direct DB access in the handler** — `acquire_with_retry`, `conn.fetch` / `fetchrow` / `execute`, raw SQL strings, or `fq_table(...)`. These are a **must fix**: the query must be moved into a `MemoryEngine` method that returns a typed model, and the handler must call that method.
|
||||
- **Verify authentication is enforced in the engine** — the handler must delegate to an engine method that authenticates via `request_context` (`_authenticate_tenant`, typically through `get_bank_profile`). A handler that reads/writes tenant-scoped data without an engine method enforcing auth is a **must fix** (tenant data could leak across schemas).
|
||||
|
||||
### 7c. Check bank/tenant query scoping
|
||||
|
||||
For **every SQL statement added or changed** in the diff (grep the diff for `conn.fetch`, `conn.fetchrow`, `conn.fetchval`, `conn.execute`, `executemany`, and any raw `SELECT`/`INSERT`/`UPDATE`/`DELETE` f-strings, including multi-line ones), verify it cannot touch another bank's rows — see **Bank/Tenant Isolation in Queries** above.
|
||||
|
||||
For each statement against a multi-bank table (`memory_units`, `documents`, `entities`, `entity_links`, `mental_models`, `knowledge_pages`, `memory_links`, `observation_history`, …), confirm it is scoped by one of the three legitimate mechanisms:
|
||||
1. explicit `AND bank_id = $n`;
|
||||
2. a globally-unique single-column PK (`*.id` uuid, or the bank-encoded `chunks.chunk_id`) — **not** a composite-PK id like `documents.id`/`document_id` or `mental_models.id`;
|
||||
3. transitively, through a globally-unique id set that was itself selected from a bank-scoped query in the same call.
|
||||
|
||||
**Flag as a must fix** any statement filtering a multi-bank table by a caller-supplied, non-globally-unique key (`document_id`, `mental_models.id`, an entity name, …) with **no** `bank_id` predicate — construct the concrete two-bank scenario (two banks share the id; the statement reads/counts/updates/deletes the wrong bank's rows or over-reports) to confirm it's real before flagging. Prime tells: a `bank_id`-carrying sibling statement right next to a `bank_id`-less one; a `WHERE bank_id` guarded by `if bank_id:` with a `None` default; an import/transfer write that inherits a source `bank_id` instead of pinning the destination.
|
||||
|
||||
### 7d. Check list endpoints paginate
|
||||
|
||||
For every added or changed `GET` handler that returns a collection, confirm it takes `limit`/`offset` and returns `total` — see **API Layer & Data Access** above for the exact shape. Then check the fix is real end to end, since a param that nothing enforces is worse than none:
|
||||
|
||||
- **The bound reaches the work, not just the response.** Verify the page size actually limits the expensive part — the SQL `LIMIT`/`OFFSET`, or (when paging must happen after an in-process filter, as in `list_banks` where the `filter_bank_list` extension hook can drop any bank) an explicit slice with the per-item work — live store counts, `get_bank_configs`, re-embedding — done for the page only. Paging in SQL *before* a filter that can drop rows is a **must fix**: it hands back short or empty pages and a `total` that counts rows the caller can't see.
|
||||
- **Every in-repo consumer pages.** A new default `limit` silently truncates callers that used to get everything: the control plane (`src/lib/api.ts` + the `src/app/api/` proxy route + any context/selector that holds the full list), the CLI (`hindsight-cli/src/api.rs`), MCP tools, and the Zapier dynamic dropdowns. Each must either page through to completion or expose paging in its UI — flag any consumer left on a single default-sized page.
|
||||
- **Search moves server-side with it.** A picker that filtered client-side over the full list now only filters the loaded page. If the endpoint gained `q`, the UI must send it (and disable its local filtering, e.g. cmdk's `shouldFilter={false}`); if it didn't, say why the collection is small enough not to need it.
|
||||
- **Tests that look up their own row must not depend on landing on page 1** — they should search or pass an explicit `limit`, not rely on default ordering.
|
||||
|
||||
### 8. Check code comments
|
||||
|
||||
For each non-trivial change:
|
||||
@@ -193,6 +247,48 @@ If any files in `hindsight-integrations/` were added or changed, verify:
|
||||
- **Docs gallery + sidebar entry** — the integration must have an entry in `hindsight-docs/src/data/integrations.json`. This file is the **single source of truth** that drives both the integrations gallery and the docs sidebar (the sidebar category is injected from it at render time across all docs versions). The entry needs an internal `/sdks/integrations/<slug>` `link` and a matching page at `hindsight-docs/docs-integrations/<slug>.md(x)`. The `hindsight-docs/scripts/check-integrations.mjs` build step enforces both directions — forward: every internal JSON entry has a doc page; reverse: every released tag (`integrations/<name>/vX.Y.Z`) appears in the JSON (private infra like `cloudflare-oauth-proxy` is in the script's `EXCLUDED` set). Flag any integration that is released (or being released) but missing from `integrations.json`, and any JSON entry without a doc page. Do **not** hand-edit `versioned_sidebars/*.json` to add integration links — they are positional placeholders filled from the JSON.
|
||||
- **Code standards** — the integration code must follow all Python style rules (type hints, no raw dicts, no tuple returns, etc.).
|
||||
|
||||
### 9a. Check parity across sibling implementations
|
||||
|
||||
Whenever the same capability is implemented once per *variant* — one per harness, per language, per
|
||||
dialect, per provider — the new or changed variant is where a capability silently goes missing. The
|
||||
defect never looks like a bug in the diff: the code that's wrong is the code that **isn't there**,
|
||||
and every existing test still passes because the sibling that forgot is by definition the one nobody
|
||||
wrote a test for. That is how dsh shipped in daemon mode without ever starting a daemon (#3524):
|
||||
`ensureDaemon` sat in the hook-only wrappers, so all five persistent-plugin harnesses lacked it.
|
||||
|
||||
Known sibling families in this repo (this is not the whole list — the rule is about the *shape*):
|
||||
|
||||
| Family | Where |
|
||||
|---|---|
|
||||
| Coding-agent harnesses | `hindsight-integrations/coding-agents/src/` (hook harnesses vs. persistent-plugin harnesses: dsh, opencode, Kilo, Cline, Prime Agent) |
|
||||
| Wrapper SDK clients | TypeScript + Python wrappers — see step 7a |
|
||||
| Alembic migrations | `_pg_upgrade` / `_oracle_upgrade` in every migration |
|
||||
| Dataplane ↔ control plane | `api/http.py` params vs `hindsight-control-plane/src/app/api/**` proxy routes + `lib/api.ts` |
|
||||
| LLM providers | per-provider branches in `engine/llm_wrapper.py` |
|
||||
|
||||
**Procedure — do this by hand; no linter catches it.** When the diff adds a new sibling, or changes
|
||||
one sibling of a family:
|
||||
|
||||
1. **Enumerate the family.** List every existing sibling (`ls` the directory, grep the registry).
|
||||
2. **Diff the capability list, not the code.** For each capability the *other* siblings have —
|
||||
lifecycle hooks called, setup/teardown performed, config flags honoured, opt-outs respected,
|
||||
registry/installer/docs entries — confirm the changed sibling has it, or that its absence is
|
||||
deliberate and commented. Grep is the tool: `grep -rn ensureDaemon src` proves who calls it.
|
||||
3. **Prefer hoisting over copying.** If the capability now exists in N places, the fix is usually to
|
||||
move it into the one path every sibling already shares (e.g. `RuntimeCore`, `buildHookOutput`),
|
||||
not to paste an Nth copy that the N+1th sibling will forget again.
|
||||
4. **Demand a structural guard, not just a unit test.** A test for the sibling that forgot doesn't
|
||||
exist by construction, so ask for a test that asserts *over the whole family*: enumerate the
|
||||
siblings from the filesystem/registry and assert each satisfies the contract, with an explicit,
|
||||
commented exemption list. Precedents: `registry covers every installable harness`
|
||||
(`harness/registry.test.ts`), `every harness entrypoint reaches a daemon` (`core/daemon.test.ts`),
|
||||
`test_backup_tables_covers_entire_schema`, `test_migration_shape.py`.
|
||||
|
||||
Flag a capability present in every sibling but one as a **must fix** — state which siblings have it,
|
||||
which doesn't, and what the user-visible symptom is (for #3524: every `hindsight_*` tool call fails
|
||||
with ECONNREFUSED and nothing ever starts the daemon). A new sibling family member landing with no
|
||||
family-wide guard test is a **should fix**.
|
||||
|
||||
### 10. Check MCP tool registration completeness
|
||||
|
||||
If any new MCP tools were added or existing tools renamed in `hindsight-api-slim/hindsight_api/mcp_tools.py`:
|
||||
@@ -254,9 +350,12 @@ Present a clear summary organized by severity:
|
||||
- Missing tests for new endpoints
|
||||
- Direct DB access (raw SQL / `acquire_with_retry` / `fq_table`) in an `api/` handler instead of a `MemoryEngine` method
|
||||
- Tenant-scoped data accessed without authentication enforced in the engine (`_authenticate_tenant` / `get_bank_profile`)
|
||||
- A SQL statement against a multi-bank table filtered by a caller-supplied, non-globally-unique key without a `bank_id` predicate (cross-bank read/write leak — see step 7c)
|
||||
- New integration missing tests, CI job, or release-integration.sh entry
|
||||
- Released/added integration missing from `hindsight-docs/src/data/integrations.json`, or a JSON entry with no `docs-integrations/<slug>` page (fails the docs build via `check-integrations.mjs`)
|
||||
- New PostgreSQL table missing from `BACKUP_TABLES` in `admin/cli.py` (silent data loss on restore)
|
||||
- A capability every sibling implementation has except the one in the diff (see step 9a) — a
|
||||
harness, dialect, provider or language variant that skips a lifecycle step the others perform
|
||||
|
||||
**Should fix** — issues that hurt code quality:
|
||||
- Dead code / unused imports missed by linter
|
||||
|
||||
+58
-6
@@ -7,7 +7,10 @@ HINDSIGHT_API_LLM_PROVIDER=openai
|
||||
HINDSIGHT_API_LLM_API_KEY=your-api-key-here
|
||||
HINDSIGHT_API_LLM_MODEL=gpt-4o-mini
|
||||
HINDSIGHT_API_LLM_BASE_URL=https://api.openai.com/v1
|
||||
# Reasoning effort for providers/models that support it. Examples: low, medium, high, xhigh.
|
||||
# Reasoning effort for providers/models that support it. Examples: none, low, medium, high, xhigh.
|
||||
# Set it and the value is sent as given, whatever the model is called — use `none` to stop a
|
||||
# self-hosted reasoning model (vLLM, Ollama, llama.cpp, TGI) emitting thinking blocks. Unset,
|
||||
# no reasoning parameter is sent at all and each model runs at its own default effort.
|
||||
# HINDSIGHT_API_LLM_REASONING_EFFORT=low
|
||||
|
||||
# Sampling temperature for internal LLM calls. Set a number in [0.0, 2.0], or `none`
|
||||
@@ -54,6 +57,13 @@ HINDSIGHT_API_LLM_BASE_URL=https://api.openai.com/v1
|
||||
# the same model in us-east-1 accepts response_format and needs nothing here.
|
||||
# HINDSIGHT_API_LLM_STRUCTURED_OUTPUT_FORCED_TOOL=false
|
||||
|
||||
# Transport-level output cap for reflect's final synthesis call. Unset = uncapped:
|
||||
# the model runs to a natural stop and the reflect/mental-model max_tokens governs
|
||||
# visible length via a prompt directive + a post-hoc rewrite (not by truncating the
|
||||
# provider call, which on thinking models is eaten by reasoning tokens and cuts pages
|
||||
# off mid-word). Set an integer only to enforce a hard cost ceiling on the call.
|
||||
# HINDSIGHT_API_REFLECT_MAX_COMPLETION_TOKENS=16000
|
||||
|
||||
# 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.
|
||||
@@ -159,6 +169,7 @@ HINDSIGHT_API_LOG_LEVEL=info
|
||||
# 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_DB_SESSION_SETUP_ON_ACQUIRE=true # Re-apply the per-connection session settings (statement_timeout, planner parallelism, trigram threshold, vector-search tuning, and the vchord search path) every time a connection is taken from the pool, not just when it is opened. Releasing a connection resets it to the server defaults, so turn this off only when the same settings are pinned on the role/database (ALTER ROLE ... SET) — then it is a pure round trip per acquire, worth reclaiming behind a transaction-mode pooler. On the vchord text-search backend the search path is in that set and losing it fails recall outright, so pin it too. application_name is always re-applied regardless.
|
||||
# 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.
|
||||
@@ -166,12 +177,23 @@ HINDSIGHT_API_LOG_LEVEL=info
|
||||
# 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.
|
||||
|
||||
# Background maintenance cadences (Optional)
|
||||
# Each sweep begins with one cross-tenant discovery call that probes every schema holding the relevant
|
||||
# table, in every API/worker process — so its cost scales with tenant count while the work it finds does
|
||||
# not. On deployments with thousands of tenants these intervals are the knob to raise.
|
||||
# HINDSIGHT_API_RETENTION_SWEEP_INTERVAL_SECONDS=3600 # How often expired audit_log / llm_requests rows are deleted across all tenant schemas. Retention is counted in days, so this only sets how promptly they disappear; 0 disables the sweeps.
|
||||
# HINDSIGHT_API_OPERATION_CLEANUP_INTERVAL_SECONDS=900 # How often expired terminal operation rows are pruned; with the batch size above this sets the drain rate for a backlog. 0 disables the job.
|
||||
# HINDSIGHT_API_MAINTENANCE_START_JITTER_SECONDS=60 # Upper bound on a random delay before a process runs its FIRST maintenance tick. Every job is due on that tick, so without an offset a fleet started together runs every sweep in every process at once. 0 disables the jitter.
|
||||
|
||||
# Vector Extension (Optional - uses pgvector by default)
|
||||
# Options: "pgvector" (default), "vchord", "pgvectorscale" (DiskANN)
|
||||
# HINDSIGHT_API_VECTOR_EXTENSION=pgvector
|
||||
# For Azure PostgreSQL with DiskANN:
|
||||
# HINDSIGHT_API_VECTOR_EXTENSION=pgvectorscale # Auto-detects pg_diskann on Azure
|
||||
|
||||
# Per-bank vector indexes (pgvector / pgvectorscale / vchord only; ScaNN and Oracle use one global index)
|
||||
# HINDSIGHT_API_VECTOR_INDEX_MIN_ROWS=0 # Memories a bank needs in one fact type before that fact type gets its own vector index. 0 (default) = no minimum, every bank holding memories is indexed. Set ~10000 on deployments with thousands of banks: every index lives on the shared memory_units table and is planned against by every OTHER bank's queries, so unconditional per-bank indexes put a ceiling on bank count. Smaller banks then use exact search, which is faster AND exact.
|
||||
|
||||
# Text Search Extension (Optional - uses native PostgreSQL full-text search by default)
|
||||
# Backend options: "native" (default), "vchord", "pg_textsearch", "pgroonga", "pg_search"
|
||||
# HINDSIGHT_API_TEXT_SEARCH_EXTENSION=native
|
||||
@@ -184,11 +206,18 @@ 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
|
||||
# Cap on the number of terms in the native PostgreSQL BM25 tsquery. Long queries
|
||||
# OR-join every normalized token, and native ranking (no IDF, re-ranks every
|
||||
# match) can then scan a large fraction of the bank and time out. Over the cap,
|
||||
# the most selective terms are kept — lowest tenant-wide document frequency, read
|
||||
# for free from pg_stats (no reindex). 0 restores the uncapped behavior; the cap
|
||||
# bounds only the native backend (other BM25 backends get the raw query).
|
||||
# HINDSIGHT_API_BM25_MAX_QUERY_TERMS=16
|
||||
# When the cap above trims a query, keep the most selective terms (lowest
|
||||
# document frequency, from pg_stats) instead of the first N. true is strictly
|
||||
# better for recall at no extra cost when stats exist; set false to opt out of
|
||||
# the catalog read and cap by position. Ignored when the cap is 0.
|
||||
# HINDSIGHT_API_BM25_SELECTIVE_TERMS=true
|
||||
|
||||
# File Parser (Optional - uses markitdown by default)
|
||||
# HINDSIGHT_API_FILE_PARSER=markitdown
|
||||
@@ -234,6 +263,12 @@ HINDSIGHT_API_LOG_LEVEL=info
|
||||
# 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
|
||||
# Asymmetric models (E5, google/embeddinggemma-300m, ...) expect a different instruction
|
||||
# in front of a search than in front of stored text. Providers that only accept plain text
|
||||
# (tei, openai-compatible, litellm) need it applied client-side; local/zeroentropy handle it
|
||||
# themselves and ignore these. Unset = text sent as-is.
|
||||
# HINDSIGHT_API_EMBEDDINGS_QUERY_PREFIX="task: search result | query: "
|
||||
# HINDSIGHT_API_EMBEDDINGS_PASSAGE_PREFIX="title: none | text: "
|
||||
# For TEI provider:
|
||||
# HINDSIGHT_API_EMBEDDINGS_TEI_URL=http://localhost:8080
|
||||
# For OpenAI-compatible embeddings:
|
||||
@@ -290,6 +325,10 @@ HINDSIGHT_API_LOG_LEVEL=info
|
||||
# HINDSIGHT_API_RERANKER_LOCAL_ALLOW_MPS=false
|
||||
# For TEI provider:
|
||||
# HINDSIGHT_API_RERANKER_TEI_URL=http://localhost:8081
|
||||
# For flashrank provider: passages scored per ONNX forward pass. Each pass
|
||||
# allocates attention tensors sized batch * heads * seq^2, so raising this
|
||||
# raises peak memory quadratically in passage length:
|
||||
# HINDSIGHT_API_RERANKER_FLASHRANK_BATCH_SIZE=32
|
||||
# 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
|
||||
@@ -337,6 +376,19 @@ HINDSIGHT_API_LOG_LEVEL=info
|
||||
# HINDSIGHT_API_LOOP_WATCHDOG_POLL_INTERVAL_MS=250
|
||||
# HINDSIGHT_API_DB_ACQUIRE_WARN_THRESHOLD_MS=1000
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Extensions (Optional)
|
||||
# -----------------------------------------------------------------------------
|
||||
|
||||
# Request headers copied into RequestContext.extra_headers so a custom
|
||||
# TenantExtension / OperationValidatorExtension can read them. Comma-separated,
|
||||
# matched case-insensitively. Unset by default: extensions see only the
|
||||
# Authorization header. Use this when the bearer token identifies a proxy rather
|
||||
# than the caller, and per-caller identity arrives in a separate header. A listed
|
||||
# header that arrives more than once is dropped, so only list headers the proxy
|
||||
# in front of Hindsight sets itself (stripping any client-supplied copy).
|
||||
# HINDSIGHT_API_EXTENSION_PASSTHROUGH_HEADERS=x-user-assertion
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Webhooks (Optional)
|
||||
# -----------------------------------------------------------------------------
|
||||
|
||||
File diff suppressed because one or more lines are too long
|
Before Width: | Height: | Size: 40 KiB After Width: | Height: | Size: 41 KiB |
@@ -1016,6 +1016,34 @@
|
||||
{
|
||||
"date": "2026-08-11",
|
||||
"stars": 19506
|
||||
},
|
||||
{
|
||||
"date": "2026-08-12",
|
||||
"stars": 19698
|
||||
},
|
||||
{
|
||||
"date": "2026-08-13",
|
||||
"stars": 19841
|
||||
},
|
||||
{
|
||||
"date": "2026-08-14",
|
||||
"stars": 19920
|
||||
},
|
||||
{
|
||||
"date": "2026-08-15",
|
||||
"stars": 19959
|
||||
},
|
||||
{
|
||||
"date": "2026-08-16",
|
||||
"stars": 20011
|
||||
},
|
||||
{
|
||||
"date": "2026-08-17",
|
||||
"stars": 20055
|
||||
},
|
||||
{
|
||||
"date": "2026-08-18",
|
||||
"stars": 20114
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
@@ -327,17 +327,50 @@ jobs:
|
||||
working-directory: hindsight-cli
|
||||
run: cargo build --release --target ${{ matrix.target }}
|
||||
|
||||
- name: Install cargo-about
|
||||
if: matrix.asset_name == 'hindsight-linux-amd64'
|
||||
uses: taiki-e/install-action@v2
|
||||
with:
|
||||
tool: [email protected]
|
||||
|
||||
- name: Verify cargo-about
|
||||
if: matrix.asset_name == 'hindsight-linux-amd64'
|
||||
run: cargo about --version
|
||||
|
||||
# The build above only fetches crates for this target; cargo-about resolves
|
||||
# the graph for every target platform, so fetch for all of them (that is what
|
||||
# `cargo fetch` without --target does) before the --offline generate.
|
||||
- name: Fetch crate sources for the license scan
|
||||
if: matrix.asset_name == 'hindsight-linux-amd64'
|
||||
working-directory: hindsight-cli
|
||||
run: cargo fetch
|
||||
|
||||
- name: Generate license manifest
|
||||
if: matrix.asset_name == 'hindsight-linux-amd64'
|
||||
working-directory: hindsight-cli
|
||||
run: mkdir -p ../artifacts && cargo about generate --offline --manifest-path Cargo.toml --config about.toml about.hbs --output-file ../artifacts/THIRD_PARTY_LICENSES.txt
|
||||
|
||||
- name: Verify license files
|
||||
if: matrix.asset_name == 'hindsight-linux-amd64'
|
||||
run: |
|
||||
test -s LICENSE
|
||||
test -s artifacts/THIRD_PARTY_LICENSES.txt
|
||||
grep -Fq "THIRD-PARTY SOFTWARE LICENSES" artifacts/THIRD_PARTY_LICENSES.txt
|
||||
|
||||
- name: Prepare artifact
|
||||
run: |
|
||||
mkdir -p artifacts
|
||||
cp hindsight-cli/target/${{ matrix.target }}/release/${{ matrix.artifact_name }} artifacts/${{ matrix.asset_name }}
|
||||
if [ "${{ matrix.asset_name }}" = "hindsight-linux-amd64" ]; then
|
||||
cp LICENSE artifacts/LICENSE
|
||||
fi
|
||||
chmod +x artifacts/${{ matrix.asset_name }}
|
||||
|
||||
- name: Upload artifacts
|
||||
uses: actions/upload-artifact@v7
|
||||
with:
|
||||
name: rust-cli-${{ matrix.asset_name }}
|
||||
path: artifacts/${{ matrix.asset_name }}
|
||||
path: artifacts/*
|
||||
retention-days: 1
|
||||
|
||||
release-docker-images:
|
||||
@@ -601,6 +634,11 @@ jobs:
|
||||
cp artifacts/rust-cli-linux-arm64/hindsight-linux-arm64 release-assets/ || true
|
||||
cp artifacts/rust-cli-darwin-amd64/hindsight-darwin-amd64 release-assets/ || true
|
||||
cp artifacts/rust-cli-darwin-arm64/hindsight-darwin-arm64 release-assets/ || true
|
||||
# Rust CLI license files (shared by all four platform binaries)
|
||||
cp artifacts/rust-cli-linux/LICENSE release-assets/
|
||||
cp artifacts/rust-cli-linux/THIRD_PARTY_LICENSES.txt release-assets/
|
||||
test -s release-assets/LICENSE
|
||||
test -s release-assets/THIRD_PARTY_LICENSES.txt
|
||||
# Helm chart
|
||||
cp artifacts/helm-chart/*.tgz release-assets/ || true
|
||||
ls -la release-assets/
|
||||
|
||||
@@ -32,6 +32,7 @@ jobs:
|
||||
integration-tests: ${{ steps.filter.outputs.integration-tests }}
|
||||
integrations-openclaw: ${{ steps.filter.outputs.integrations-openclaw }}
|
||||
integrations-ai-sdk: ${{ steps.filter.outputs.integrations-ai-sdk }}
|
||||
integrations-eliza: ${{ steps.filter.outputs.integrations-eliza }}
|
||||
integrations-agent-framework: ${{ steps.filter.outputs.integrations-agent-framework }}
|
||||
integrations-composio: ${{ steps.filter.outputs.integrations-composio }}
|
||||
integrations-chat: ${{ steps.filter.outputs.integrations-chat }}
|
||||
@@ -43,6 +44,7 @@ jobs:
|
||||
integrations-continue: ${{ steps.filter.outputs.integrations-continue }}
|
||||
integrations-cursor-cli: ${{ steps.filter.outputs.integrations-cursor-cli }}
|
||||
integrations-zcode: ${{ steps.filter.outputs.integrations-zcode }}
|
||||
integrations-agent-plugin: ${{ steps.filter.outputs.integrations-agent-plugin }}
|
||||
integrations-copilot-cli: ${{ steps.filter.outputs.integrations-copilot-cli }}
|
||||
integrations-crewai: ${{ steps.filter.outputs.integrations-crewai }}
|
||||
integrations-litellm: ${{ steps.filter.outputs.integrations-litellm }}
|
||||
@@ -142,6 +144,8 @@ jobs:
|
||||
- 'hindsight-integrations/openclaw/**'
|
||||
integrations-ai-sdk:
|
||||
- 'hindsight-integrations/ai-sdk/**'
|
||||
integrations-eliza:
|
||||
- 'hindsight-integrations/eliza/**'
|
||||
integrations-agent-framework:
|
||||
- 'hindsight-integrations/agent-framework/**'
|
||||
integrations-composio:
|
||||
@@ -194,6 +198,8 @@ jobs:
|
||||
- 'hindsight-integrations/zed/**'
|
||||
integrations-zcode:
|
||||
- 'hindsight-integrations/zcode/**'
|
||||
integrations-agent-plugin:
|
||||
- 'hindsight-integrations/agent-plugin/**'
|
||||
integrations-n8n:
|
||||
- 'hindsight-integrations/n8n/**'
|
||||
integrations-zapier:
|
||||
@@ -796,6 +802,29 @@ jobs:
|
||||
working-directory: ./hindsight-integrations/zcode
|
||||
run: uv run pytest tests -v
|
||||
|
||||
test-agent-plugin-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
(github.event_name == 'workflow_dispatch' ||
|
||||
needs.detect-changes.outputs.integrations-agent-plugin == '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 || '' }}
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v6
|
||||
with:
|
||||
python-version-file: ".python-version"
|
||||
|
||||
- name: Validate Agent Plugin manifests
|
||||
working-directory: ./hindsight-integrations/agent-plugin
|
||||
run: python3 validate.py
|
||||
|
||||
build-ai-sdk-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
@@ -859,6 +888,37 @@ jobs:
|
||||
working-directory: ./hindsight-integrations/ai-sdk
|
||||
run: npm run test:deno
|
||||
|
||||
build-eliza-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
(github.event_name == 'workflow_dispatch' ||
|
||||
needs.detect-changes.outputs.integrations-eliza == '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: Set up Node.js
|
||||
uses: actions/setup-node@v6
|
||||
with:
|
||||
node-version: '22'
|
||||
|
||||
- name: Install dependencies
|
||||
working-directory: ./hindsight-integrations/eliza
|
||||
run: npm ci
|
||||
|
||||
- name: Run tests
|
||||
working-directory: ./hindsight-integrations/eliza
|
||||
run: npm test
|
||||
|
||||
- name: Build
|
||||
working-directory: ./hindsight-integrations/eliza
|
||||
run: npm run build
|
||||
|
||||
test-opencode-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
@@ -1395,6 +1455,30 @@ jobs:
|
||||
hindsight-cli/target
|
||||
key: ${{ runner.os }}-cargo-${{ hashFiles('**/Cargo.lock') }}
|
||||
|
||||
- name: Install cargo-about
|
||||
uses: taiki-e/install-action@v2
|
||||
with:
|
||||
tool: [email protected]
|
||||
|
||||
- name: Verify cargo-about
|
||||
run: cargo about --version
|
||||
|
||||
# cargo-about resolves the dependency graph for every target platform, so it
|
||||
# needs crates this host never builds (e.g. the Android-only
|
||||
# android_system_properties). Populate the registry first: `cargo fetch`
|
||||
# without --target downloads for all targets, and Cargo.lock is not checked
|
||||
# in, so nothing is cached from a previous step.
|
||||
- name: Fetch crate sources for the license scan
|
||||
working-directory: hindsight-cli
|
||||
run: cargo fetch
|
||||
|
||||
- name: Generate and verify license manifest
|
||||
working-directory: hindsight-cli
|
||||
run: |
|
||||
cargo about generate --offline --manifest-path Cargo.toml --config about.toml about.hbs --output-file /tmp/THIRD_PARTY_LICENSES.txt
|
||||
test -s /tmp/THIRD_PARTY_LICENSES.txt
|
||||
grep -Fq "THIRD-PARTY SOFTWARE LICENSES" /tmp/THIRD_PARTY_LICENSES.txt
|
||||
|
||||
- name: Run unit tests
|
||||
working-directory: hindsight-cli
|
||||
run: cargo test
|
||||
@@ -4560,6 +4644,8 @@ jobs:
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 45
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
@@ -4584,6 +4670,8 @@ jobs:
|
||||
${{ runner.os }}-huggingface-
|
||||
|
||||
- name: Run Hermes compatibility test
|
||||
env:
|
||||
GITHUB_TOKEN: ${{ github.token }}
|
||||
run: ./scripts/test-hermes-compat.sh
|
||||
|
||||
- name: Collect embedded daemon logs on failure
|
||||
@@ -5217,6 +5305,7 @@ jobs:
|
||||
- test-codex-integration
|
||||
- test-cursor-cli-integration
|
||||
- test-zcode-integration
|
||||
- test-agent-plugin-integration
|
||||
- build-ai-sdk-integration
|
||||
- test-ai-sdk-integration-deno
|
||||
- test-opencode-integration
|
||||
|
||||
@@ -20,6 +20,9 @@ wheels/
|
||||
|
||||
# Node
|
||||
node_modules/
|
||||
# Without this, the pattern above matches directories only — a node_modules SYMLINK (what you get
|
||||
# pointing a scratch worktree at an installed one) is a file, slips past it, and can be committed.
|
||||
node_modules
|
||||
|
||||
# Environment variables and local config
|
||||
.env
|
||||
|
||||
@@ -65,7 +65,7 @@ services:
|
||||
environment:
|
||||
# LLM Configuration
|
||||
HINDSIGHT_API_LLM_PROVIDER: ${HINDSIGHT_API_LLM_PROVIDER:-openai}
|
||||
HINDSIGHT_API_LLM_API_KEY: ${OPENAI_API_KEY:-your-api-key}
|
||||
HINDSIGHT_API_LLM_API_KEY: ${HINDSIGHT_API_LLM_API_KEY:-}
|
||||
|
||||
# Database Configuration
|
||||
HINDSIGHT_API_DATABASE_URL: postgresql://${HINDSIGHT_DB_USER:-hindsight_user}:${HINDSIGHT_DB_PASSWORD:-hindsight_password}@db:5432/${HINDSIGHT_DB_NAME:-hindsight_db}
|
||||
|
||||
@@ -27,7 +27,7 @@ needed in the image.
|
||||
## Quick start
|
||||
|
||||
```bash
|
||||
export OPENAI_API_KEY=sk-xxx
|
||||
export HINDSIGHT_API_LLM_API_KEY=sk-xxx
|
||||
|
||||
docker compose -f docker/docker-compose/custom-models/docker-compose.yaml up --build
|
||||
```
|
||||
|
||||
@@ -3,11 +3,12 @@ name: hindsight-custom-models
|
||||
# in at build time, so pod startup does not depend on HuggingFace at runtime.
|
||||
#
|
||||
# Quick start:
|
||||
# export OPENAI_API_KEY=sk-xxx
|
||||
# export HINDSIGHT_API_LLM_API_KEY=sk-xxx
|
||||
# docker compose -f docker/docker-compose/custom-models/docker-compose.yaml up --build
|
||||
#
|
||||
# Required environment variables:
|
||||
# - OPENAI_API_KEY (or configure another LLM provider via HINDSIGHT_API_LLM_*)
|
||||
# - HINDSIGHT_API_LLM_API_KEY (pair it with HINDSIGHT_API_LLM_PROVIDER to use
|
||||
# a provider other than the default openai)
|
||||
|
||||
services:
|
||||
hindsight:
|
||||
@@ -25,7 +26,7 @@ services:
|
||||
- "9999:9999"
|
||||
environment:
|
||||
HINDSIGHT_API_LLM_PROVIDER: ${HINDSIGHT_API_LLM_PROVIDER:-openai}
|
||||
HINDSIGHT_API_LLM_API_KEY: ${OPENAI_API_KEY:-your-api-key}
|
||||
HINDSIGHT_API_LLM_API_KEY: ${HINDSIGHT_API_LLM_API_KEY:-}
|
||||
|
||||
# Point Hindsight at the models baked into the image above.
|
||||
HINDSIGHT_API_EMBEDDINGS_PROVIDER: local
|
||||
|
||||
@@ -39,7 +39,7 @@ services:
|
||||
- "8888:8888"
|
||||
- "9999:9999"
|
||||
environment:
|
||||
- HINDSIGHT_API_LLM_API_KEY=${OPENAI_API_KEY?Please set the OPENAI_API_KEY env variable}
|
||||
- HINDSIGHT_API_LLM_API_KEY=${HINDSIGHT_API_LLM_API_KEY:?Please set the HINDSIGHT_API_LLM_API_KEY env variable}
|
||||
- HINDSIGHT_API_DATABASE_URL=postgresql://${HINDSIGHT_DB_USER:-hindsight_user}:${HINDSIGHT_DB_PASSWORD:?Please set the HINDSIGHT_DB_PASSWORD env variable}@db:5432/${HINDSIGHT_DB_NAME:-hindsight_db}
|
||||
depends_on:
|
||||
- db
|
||||
|
||||
@@ -72,7 +72,7 @@ services:
|
||||
environment:
|
||||
# LLM Configuration
|
||||
HINDSIGHT_API_LLM_PROVIDER: ${HINDSIGHT_API_LLM_PROVIDER:-openai}
|
||||
HINDSIGHT_API_LLM_API_KEY: ${OPENAI_API_KEY:-your-api-key}
|
||||
HINDSIGHT_API_LLM_API_KEY: ${HINDSIGHT_API_LLM_API_KEY:-}
|
||||
|
||||
# Database Configuration
|
||||
HINDSIGHT_API_DATABASE_URL: postgresql://${HINDSIGHT_DB_USER:-hindsight_user}:${HINDSIGHT_DB_PASSWORD:-hindsight_password}@db:5432/${HINDSIGHT_DB_NAME:-hindsight_db}
|
||||
|
||||
@@ -68,7 +68,7 @@ services:
|
||||
environment:
|
||||
# LLM Configuration
|
||||
HINDSIGHT_API_LLM_PROVIDER: ${HINDSIGHT_API_LLM_PROVIDER:-openai}
|
||||
HINDSIGHT_API_LLM_API_KEY: ${OPENAI_API_KEY:-your-api-key}
|
||||
HINDSIGHT_API_LLM_API_KEY: ${HINDSIGHT_API_LLM_API_KEY:-}
|
||||
|
||||
# Database Configuration
|
||||
HINDSIGHT_API_DATABASE_URL: postgresql://${HINDSIGHT_DB_USER:-hindsight_user}:${HINDSIGHT_DB_PASSWORD:-hindsight_password}@db:5432/${HINDSIGHT_DB_NAME:-hindsight_db}
|
||||
|
||||
@@ -68,7 +68,7 @@ services:
|
||||
environment:
|
||||
# LLM Configuration
|
||||
HINDSIGHT_API_LLM_PROVIDER: ${HINDSIGHT_API_LLM_PROVIDER:-openai}
|
||||
HINDSIGHT_API_LLM_API_KEY: ${OPENAI_API_KEY:-your-api-key}
|
||||
HINDSIGHT_API_LLM_API_KEY: ${HINDSIGHT_API_LLM_API_KEY:-}
|
||||
|
||||
# Database Configuration
|
||||
HINDSIGHT_API_DATABASE_URL: postgresql://${HINDSIGHT_DB_USER:-hindsight_user}:${HINDSIGHT_DB_PASSWORD:-hindsight_password}@db:5432/${HINDSIGHT_DB_NAME:-hindsight_db}
|
||||
|
||||
@@ -59,7 +59,7 @@ services:
|
||||
- "8888:8888"
|
||||
- "9999:9999"
|
||||
environment:
|
||||
- HINDSIGHT_API_LLM_API_KEY=${OPENAI_API_KEY?Please set the OPENAI_API_KEY env variable}
|
||||
- HINDSIGHT_API_LLM_API_KEY=${HINDSIGHT_API_LLM_API_KEY:?Please set the HINDSIGHT_API_LLM_API_KEY env variable}
|
||||
- HINDSIGHT_API_DATABASE_URL=postgresql://${HINDSIGHT_DB_USER:-hindsight_user}:${HINDSIGHT_DB_PASSWORD:?Please set the HINDSIGHT_DB_PASSWORD env variable}@db:5432/${HINDSIGHT_DB_NAME:-hindsight_db}
|
||||
# S3 file storage configuration (SeaweedFS)
|
||||
- HINDSIGHT_API_FILE_STORAGE_TYPE=s3
|
||||
|
||||
@@ -0,0 +1,107 @@
|
||||
# Hindsight with TEI embeddings + reranker
|
||||
|
||||
Example Docker Compose setup that serves **embeddings and reranking from two
|
||||
[HuggingFace Text Embeddings Inference (TEI)](https://github.com/huggingface/text-embeddings-inference)
|
||||
sidecars** instead of the in-process local models.
|
||||
|
||||
Because embeddings and reranking run outside the API, Hindsight itself needs
|
||||
no baked-in models, so this uses the **slim** image
|
||||
(`ghcr.io/vectorize-io/hindsight:latest-slim`). Only the LLM — used for
|
||||
retain/recall/reflect — still needs a provider and API key.
|
||||
|
||||
## When to use this
|
||||
|
||||
- You want embeddings/reranking on a dedicated, independently scalable
|
||||
inference server (e.g. a GPU node) rather than in the API process.
|
||||
- You run the **slim** image and pull embeddings/reranking from an external
|
||||
service.
|
||||
- You want a self-hosted, offline-capable alternative to a cloud embeddings
|
||||
provider (OpenAI, Cohere, ...).
|
||||
|
||||
If you just want local models in-process, use the default full image — no
|
||||
sidecars required.
|
||||
|
||||
## What it runs
|
||||
|
||||
| Service | Image | Model |
|
||||
| --------------- | ------------------------------------------------------ | ---------------------------------------- |
|
||||
| `tei-embedding` | `ghcr.io/huggingface/text-embeddings-inference:cpu-1.8.3` | `BAAI/bge-small-en-v1.5` (384-dim) |
|
||||
| `tei-reranker` | `ghcr.io/huggingface/text-embeddings-inference:cpu-1.8.3` | `BAAI/bge-reranker-base` |
|
||||
| `hindsight` | `ghcr.io/vectorize-io/hindsight:latest-slim` | — (slim; talks to the sidecars) |
|
||||
|
||||
This is a prod-like configuration: the embedding model is Hindsight's default
|
||||
(`bge-small-en-v1.5`), the reranker is the `bge-reranker-base` cross-encoder
|
||||
commonly paired with it on dedicated inference servers, and both services carry
|
||||
throughput flags (`--max-concurrent-requests`, `--max-batch-tokens`,
|
||||
`--max-client-batch-size`) tuned for sustained multi-client load instead of
|
||||
TEI's bare defaults. The API points at the sidecars with:
|
||||
|
||||
```
|
||||
HINDSIGHT_API_EMBEDDINGS_PROVIDER=tei
|
||||
HINDSIGHT_API_EMBEDDINGS_TEI_URL=http://tei-embedding:80
|
||||
HINDSIGHT_API_RERANKER_PROVIDER=tei
|
||||
HINDSIGHT_API_RERANKER_TEI_URL=http://tei-reranker:80
|
||||
```
|
||||
|
||||
## Quick start
|
||||
|
||||
```bash
|
||||
export HINDSIGHT_API_LLM_API_KEY=sk-xxx
|
||||
|
||||
docker compose -f docker/docker-compose/tei/docker-compose.yaml up
|
||||
```
|
||||
|
||||
- API: http://localhost:8888
|
||||
- Control Plane: http://localhost:9999
|
||||
- TEI embedding server: http://localhost:8080 (exposed for debugging)
|
||||
- TEI reranker server: http://localhost:8081 (exposed for debugging)
|
||||
|
||||
`hindsight` waits (via `depends_on: service_healthy`) until both TEI servers
|
||||
report healthy, so the first boot pauses while each model downloads into its
|
||||
`tei_*_cache` volume. Subsequent boots reuse the cached models.
|
||||
|
||||
To use an LLM provider other than the default `openai`:
|
||||
|
||||
```bash
|
||||
export HINDSIGHT_API_LLM_PROVIDER=gemini
|
||||
export HINDSIGHT_API_LLM_API_KEY=...
|
||||
docker compose -f docker/docker-compose/tei/docker-compose.yaml up
|
||||
```
|
||||
|
||||
## Using your own models
|
||||
|
||||
Change the `--model-id` in each service's `command` to any TEI-supported
|
||||
model. The embedding dimension is **auto-detected from the server** and the
|
||||
pgvector schema is adjusted to match on first boot — no dimension env var to
|
||||
set. (If you switch the embedding model after data already exists, start from
|
||||
a fresh `pg_data` volume, since the stored vectors were built for the old
|
||||
dimension.)
|
||||
|
||||
## Verifying the servers
|
||||
|
||||
```bash
|
||||
# Health
|
||||
curl 127.0.0.1:8080/health && curl 127.0.0.1:8081/health
|
||||
|
||||
# Embedding (returns a 384-length vector for the default model)
|
||||
curl 127.0.0.1:8080/embed -H 'content-type: application/json' \
|
||||
-d '{"inputs":"hello world"}'
|
||||
|
||||
# Rerank
|
||||
curl 127.0.0.1:8081/rerank -H 'content-type: application/json' \
|
||||
-d '{"query":"what is the capital of France?","texts":["Paris is the capital of France.","Bananas are yellow."]}'
|
||||
```
|
||||
|
||||
## Apple Silicon / arm64
|
||||
|
||||
The `cpu-1.8.3` TEI images are published for `linux/amd64` only. On an
|
||||
Apple Silicon Mac, run under emulation:
|
||||
|
||||
```bash
|
||||
export DOCKER_DEFAULT_PLATFORM=linux/amd64
|
||||
docker compose -f docker/docker-compose/tei/docker-compose.yaml up
|
||||
```
|
||||
|
||||
Emulated startup is slow (model load takes a few minutes). For production,
|
||||
run on `amd64` hosts — or a GPU node with the CUDA-tagged TEI image and a GPU
|
||||
reservation.
|
||||
@@ -0,0 +1,113 @@
|
||||
name: hindsight-tei
|
||||
# Example: run Hindsight with embeddings and reranking served by two
|
||||
# HuggingFace Text Embeddings Inference (TEI) sidecars instead of the
|
||||
# in-process local models.
|
||||
#
|
||||
# Because embeddings and reranking are external, Hindsight itself needs no
|
||||
# baked-in models — this uses the **slim** image
|
||||
# (`ghcr.io/vectorize-io/hindsight:latest-slim`). Only the LLM (used for
|
||||
# retain/recall/reflect) still needs a provider + API key.
|
||||
#
|
||||
# The two TEI services here run a prod-like configuration:
|
||||
# `BAAI/bge-small-en-v1.5` embeddings (384-dim, Hindsight's default) and the
|
||||
# `BAAI/bge-reranker-base` cross-encoder, with the batching/concurrency flags
|
||||
# tuned for sustained multi-client load rather than TEI's bare defaults. Swap
|
||||
# the `--model-id` args to serve any TEI-supported model — the embedding
|
||||
# dimension is auto-detected from the server, and the pgvector schema is
|
||||
# adjusted to match on first boot.
|
||||
#
|
||||
# Quick start:
|
||||
# export HINDSIGHT_API_LLM_API_KEY=sk-xxx
|
||||
# docker compose -f docker/docker-compose/tei/docker-compose.yaml up
|
||||
#
|
||||
# First boot downloads the two models into the `tei_*_cache` volumes;
|
||||
# subsequent boots reuse them.
|
||||
|
||||
services:
|
||||
tei-embedding:
|
||||
image: ghcr.io/huggingface/text-embeddings-inference:cpu-1.8.3
|
||||
container_name: hindsight-tei-embedding
|
||||
# Prod-like tuning: high request concurrency with bounded batch sizes.
|
||||
command:
|
||||
[
|
||||
"--model-id", "BAAI/bge-small-en-v1.5",
|
||||
"--max-concurrent-requests", "512",
|
||||
"--max-batch-tokens", "16384",
|
||||
"--max-client-batch-size", "32",
|
||||
"--auto-truncate",
|
||||
]
|
||||
environment:
|
||||
# TEI listens on port 80 inside the container by default.
|
||||
PORT: "80"
|
||||
ports:
|
||||
# Exposed on the host so you can curl the server directly, e.g.
|
||||
# curl 127.0.0.1:8080/embed -H 'content-type: application/json' \
|
||||
# -d '{"inputs":"hello world"}'
|
||||
- "8080:80"
|
||||
volumes:
|
||||
- tei_embedding_cache:/data
|
||||
healthcheck:
|
||||
# The TEI image ships curl; hit its /health endpoint so Hindsight only
|
||||
# starts once the model is loaded and serving.
|
||||
test: ["CMD", "curl", "-fsS", "http://localhost:80/health"]
|
||||
interval: 10s
|
||||
timeout: 5s
|
||||
retries: 60
|
||||
start_period: 30s
|
||||
|
||||
tei-reranker:
|
||||
image: ghcr.io/huggingface/text-embeddings-inference:cpu-1.8.3
|
||||
container_name: hindsight-tei-reranker
|
||||
# Prod-like tuning: reranking batches are larger than embedding batches
|
||||
# (rerank inputs are query+document pairs scored in bulk during recall).
|
||||
command:
|
||||
[
|
||||
"--model-id", "BAAI/bge-reranker-base",
|
||||
"--max-concurrent-requests", "512",
|
||||
"--max-batch-tokens", "32768",
|
||||
"--max-client-batch-size", "128",
|
||||
"--auto-truncate",
|
||||
]
|
||||
environment:
|
||||
PORT: "80"
|
||||
ports:
|
||||
- "8081:80"
|
||||
volumes:
|
||||
- tei_reranker_cache:/data
|
||||
healthcheck:
|
||||
test: ["CMD", "curl", "-fsS", "http://localhost:80/health"]
|
||||
interval: 10s
|
||||
timeout: 5s
|
||||
retries: 60
|
||||
start_period: 30s
|
||||
|
||||
hindsight:
|
||||
image: ghcr.io/vectorize-io/hindsight:${HINDSIGHT_VERSION:-latest-slim}
|
||||
container_name: hindsight-tei
|
||||
depends_on:
|
||||
tei-embedding:
|
||||
condition: service_healthy
|
||||
tei-reranker:
|
||||
condition: service_healthy
|
||||
ports:
|
||||
- "8888:8888"
|
||||
- "9999:9999"
|
||||
environment:
|
||||
# LLM still runs through a provider — bring your own key. Pair with
|
||||
# HINDSIGHT_API_LLM_PROVIDER to use a provider other than openai.
|
||||
HINDSIGHT_API_LLM_PROVIDER: ${HINDSIGHT_API_LLM_PROVIDER:-openai}
|
||||
HINDSIGHT_API_LLM_API_KEY: ${HINDSIGHT_API_LLM_API_KEY:-}
|
||||
|
||||
# Embeddings + reranking served by the TEI sidecars above. Use the
|
||||
# in-cluster service DNS names, not localhost.
|
||||
HINDSIGHT_API_EMBEDDINGS_PROVIDER: tei
|
||||
HINDSIGHT_API_EMBEDDINGS_TEI_URL: http://tei-embedding:80
|
||||
HINDSIGHT_API_RERANKER_PROVIDER: tei
|
||||
HINDSIGHT_API_RERANKER_TEI_URL: http://tei-reranker:80
|
||||
volumes:
|
||||
- pg_data:/home/hindsight/.pg0
|
||||
|
||||
volumes:
|
||||
pg_data:
|
||||
tei_embedding_cache:
|
||||
tei_reranker_cache:
|
||||
@@ -8,17 +8,12 @@ HINDSIGHT_VERSION=latest
|
||||
|
||||
# LLM Configuration
|
||||
HINDSIGHT_API_LLM_PROVIDER=openai
|
||||
OPENAI_API_KEY=your-openai-api-key-here
|
||||
HINDSIGHT_API_LLM_API_KEY=your-openai-api-key-here
|
||||
|
||||
# Alternative LLM providers (uncomment and configure as needed):
|
||||
# Alternative LLM providers (uncomment and set the key above accordingly):
|
||||
# HINDSIGHT_API_LLM_PROVIDER=anthropic
|
||||
# ANTHROPIC_API_KEY=your-anthropic-api-key
|
||||
|
||||
# HINDSIGHT_API_LLM_PROVIDER=gemini
|
||||
# GEMINI_API_KEY=your-gemini-api-key
|
||||
|
||||
# HINDSIGHT_API_LLM_PROVIDER=groq
|
||||
# GROQ_API_KEY=your-groq-api-key
|
||||
|
||||
# Vector and Text Search (already configured in docker-compose.yaml)
|
||||
# HINDSIGHT_API_VECTOR_EXTENSION=pgvectorscale
|
||||
|
||||
@@ -9,14 +9,14 @@ Both extensions are from [Timescale](https://github.com/timescale) and provide p
|
||||
## Prerequisites
|
||||
|
||||
- Docker and Docker Compose installed
|
||||
- OpenAI API key (or another LLM provider)
|
||||
- An OpenAI API key (or a key for another LLM provider)
|
||||
|
||||
## Quick Start
|
||||
|
||||
```bash
|
||||
# Set environment variables
|
||||
export HINDSIGHT_DB_PASSWORD="your-secure-password"
|
||||
export OPENAI_API_KEY="your-openai-api-key"
|
||||
export HINDSIGHT_API_LLM_API_KEY="your-openai-api-key"
|
||||
|
||||
# Build and start
|
||||
docker compose -f docker/docker-compose/timescale/docker-compose.yaml up -d --build
|
||||
@@ -50,7 +50,7 @@ docker compose -f docker/docker-compose/timescale/docker-compose.yaml down -v
|
||||
| `HINDSIGHT_DB_USER` | PostgreSQL username | `hindsight_user` |
|
||||
| `HINDSIGHT_DB_NAME` | Database name | `hindsight_db` |
|
||||
| `HINDSIGHT_VERSION` | Hindsight Docker image version | `latest` |
|
||||
| `OPENAI_API_KEY` | OpenAI API key | (required) |
|
||||
| `HINDSIGHT_API_LLM_API_KEY` | API key for the LLM provider | (required) |
|
||||
| `HINDSIGHT_API_LLM_PROVIDER` | LLM provider | `openai` |
|
||||
|
||||
### Why Timescale Extensions?
|
||||
|
||||
@@ -8,7 +8,8 @@ name: hindsight
|
||||
#
|
||||
# Required environment variables:
|
||||
# - HINDSIGHT_DB_PASSWORD: Password for the PostgreSQL user
|
||||
# - OPENAI_API_KEY (or configure another LLM provider)
|
||||
# - HINDSIGHT_API_LLM_API_KEY (pair it with HINDSIGHT_API_LLM_PROVIDER to use
|
||||
# a provider other than the default openai)
|
||||
#
|
||||
# Optional environment variables with defaults:
|
||||
# - HINDSIGHT_VERSION: Hindsight application version (default: latest)
|
||||
@@ -80,7 +81,7 @@ services:
|
||||
environment:
|
||||
# LLM Configuration
|
||||
HINDSIGHT_API_LLM_PROVIDER: ${HINDSIGHT_API_LLM_PROVIDER:-openai}
|
||||
HINDSIGHT_API_LLM_API_KEY: ${OPENAI_API_KEY:-your-api-key}
|
||||
HINDSIGHT_API_LLM_API_KEY: ${HINDSIGHT_API_LLM_API_KEY:-}
|
||||
|
||||
# Database Configuration
|
||||
HINDSIGHT_API_DATABASE_URL: postgresql://${HINDSIGHT_DB_USER:-hindsight_user}:${HINDSIGHT_DB_PASSWORD:-hindsight_password}@db:5432/${HINDSIGHT_DB_NAME:-hindsight_db}
|
||||
|
||||
@@ -70,7 +70,7 @@ services:
|
||||
# LLM Configuration (uses OpenAI for testing vchord)
|
||||
# LLM configuration
|
||||
HINDSIGHT_API_LLM_PROVIDER: ${HINDSIGHT_API_LLM_PROVIDER:-openai}
|
||||
HINDSIGHT_API_LLM_API_KEY: ${OPENAI_API_KEY:-your-api-key}
|
||||
HINDSIGHT_API_LLM_API_KEY: ${HINDSIGHT_API_LLM_API_KEY:-}
|
||||
|
||||
# Database Configuration
|
||||
HINDSIGHT_API_DATABASE_URL: postgresql://${HINDSIGHT_DB_USER:-hindsight_user}:${HINDSIGHT_DB_PASSWORD:-hindsight_password}@db:5432/${HINDSIGHT_DB_NAME:-hindsight_db}
|
||||
|
||||
@@ -1,341 +0,0 @@
|
||||
# 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.
|
||||
@@ -1,135 +0,0 @@
|
||||
# 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.
|
||||
@@ -1,139 +0,0 @@
|
||||
# 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.9.0
|
||||
appVersion: "0.9.0"
|
||||
version: 0.9.1
|
||||
appVersion: "0.9.1"
|
||||
keywords:
|
||||
- ai
|
||||
- memory
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@vectorize-io/hindsight-all",
|
||||
"version": "0.9.0",
|
||||
"version": "0.9.1",
|
||||
"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",
|
||||
|
||||
@@ -4,15 +4,15 @@ build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "hindsight-all-slim"
|
||||
version = "0.9.0"
|
||||
version = "0.9.1"
|
||||
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.9.0",
|
||||
"hindsight-api-slim==0.9.1",
|
||||
"hindsight-client>=0.0.7",
|
||||
"hindsight-embed==0.9.0",
|
||||
"hindsight-embed==0.9.1",
|
||||
]
|
||||
|
||||
[tool.uv.sources]
|
||||
|
||||
@@ -4,15 +4,15 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "hindsight-all"
|
||||
version = "0.9.0"
|
||||
version = "0.9.1"
|
||||
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.9.0",
|
||||
"hindsight-api-slim[all]==0.9.1",
|
||||
"hindsight-client>=0.0.7",
|
||||
"hindsight-embed==0.9.0",
|
||||
"hindsight-embed==0.9.1",
|
||||
]
|
||||
|
||||
[tool.uv.sources]
|
||||
@@ -22,7 +22,7 @@ hindsight-embed = { workspace = true }
|
||||
|
||||
[project.optional-dependencies]
|
||||
local-llm = [
|
||||
"hindsight-api-slim[local-llm]==0.9.0",
|
||||
"hindsight-api-slim[local-llm]==0.9.1",
|
||||
]
|
||||
test = [
|
||||
"pytest>=7.0.0",
|
||||
|
||||
@@ -53,4 +53,4 @@ __all__ = [
|
||||
"RemoteTEICrossEncoder",
|
||||
"LLMConfig",
|
||||
]
|
||||
__version__ = "0.9.0"
|
||||
__version__ = "0.9.1"
|
||||
|
||||
@@ -175,6 +175,83 @@ def uses_per_bank_vector_indexes(ext: str) -> bool:
|
||||
return _normalize_resolved(ext) != "scann"
|
||||
|
||||
|
||||
def per_bank_index_min_rows() -> int:
|
||||
"""Rows a (bank, fact_type) needs before it earns its own partial vector index.
|
||||
|
||||
Distinct from :func:`minimum_rows_for_index`, which is ScaNN's *build*
|
||||
requirement for its single global index (AlloyDB cannot construct one below
|
||||
a floor). This is a cost policy for the per-bank backends: the indexes sit on
|
||||
the shared ``memory_units`` table, so each one is enumerated and locked at
|
||||
plan time by queries belonging to every *other* bank, and opened by every DML
|
||||
statement against the table. A small bank's index cannot repay that — the
|
||||
``(bank_id, fact_type)`` B-tree plus a top-N sort answers the same query
|
||||
exactly and faster. See issue #3485.
|
||||
|
||||
Read from config rather than passed in because the write path's pre-check,
|
||||
the maintenance operation and the admin command must all apply the same
|
||||
number; a threshold that differed between the one deciding to queue work and
|
||||
the one deciding what to do would either oscillate or never converge.
|
||||
"""
|
||||
from .config import get_config
|
||||
|
||||
return get_config().vector_index_min_rows
|
||||
|
||||
|
||||
def per_bank_index_drop_rows() -> int:
|
||||
"""Row count below which an existing per-bank vector index is dropped.
|
||||
|
||||
Strictly below :func:`per_bank_index_min_rows` so the build and drop
|
||||
decisions cannot both be true at one row count. Without the gap, a bank
|
||||
hovering at the threshold — consolidation prunes a few facts, retain adds
|
||||
them back — would rebuild and drop the same ANN index on alternating sweeps.
|
||||
"""
|
||||
from .config import VECTOR_INDEX_DROP_RATIO
|
||||
|
||||
return int(per_bank_index_min_rows() * VECTOR_INDEX_DROP_RATIO)
|
||||
|
||||
|
||||
def should_keep_per_bank_index(row_count: int) -> bool:
|
||||
"""Whether an *existing* index on a partition of ``row_count`` rows is kept.
|
||||
|
||||
The counterpart to :func:`qualifies_for_per_bank_index`, and deliberately a
|
||||
separate, lower bound: keeping starts below building, so a partition
|
||||
hovering at the threshold does not rebuild and drop the same ANN index on
|
||||
alternating writes.
|
||||
|
||||
The ``row_count > 0`` term is not redundant with the ratio. At the default
|
||||
threshold of 0 the drop floor is also 0, so a bare ``row_count >= floor``
|
||||
keeps an index over an *emptied* partition forever — every bank ever written
|
||||
to and then cleared would hold three indexes over nothing, which is the
|
||||
accumulation the threshold exists to prevent. An emptied partition loses its
|
||||
index at every threshold.
|
||||
"""
|
||||
return row_count > 0 and row_count >= per_bank_index_drop_rows()
|
||||
|
||||
|
||||
def qualifies_for_per_bank_index(row_count: int) -> bool:
|
||||
"""Whether a (bank, fact_type) holding ``row_count`` rows should have an index.
|
||||
|
||||
At the default threshold of 0 this is true for every partition that holds
|
||||
any rows at all, which is the behaviour before the threshold existed.
|
||||
|
||||
An empty partition is excluded explicitly rather than by arithmetic: at a
|
||||
threshold of 0, ``row_count >= minimum`` alone is true for zero rows, so
|
||||
every bank in the deployment would be entitled to three indexes over nothing
|
||||
the moment it was created — the exact index explosion the threshold exists
|
||||
to prevent, reintroduced by its own default.
|
||||
|
||||
Only the build side: an existing index is kept until the count falls under
|
||||
:func:`per_bank_index_drop_rows`, so callers reconciling live state must
|
||||
consult both bounds rather than treating this as the full policy.
|
||||
|
||||
Takes no extension: the backend question is settled before any reconcile
|
||||
runs (``uses_per_bank_vector_indexes`` gates the maintenance operation and
|
||||
``_vector_index_clause`` gates the admin command), so re-asking it here
|
||||
would be a second, weaker copy of a decision already made.
|
||||
"""
|
||||
return row_count > 0 and row_count >= per_bank_index_min_rows()
|
||||
|
||||
|
||||
def bootstrap_extension(conn: Connection, ext: str) -> None:
|
||||
"""Install the configured vector extension and any prerequisites if possible."""
|
||||
normalized = validate_extension(ext)
|
||||
|
||||
@@ -23,7 +23,12 @@ 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 ..engine.vector_index_health import (
|
||||
BankIndexResult,
|
||||
drop_orphaned_bank_indexes,
|
||||
list_bank_ids,
|
||||
reconcile_bank_vector_indexes,
|
||||
)
|
||||
from ..extensions import TenantExtension, load_extension
|
||||
from ..pg0 import parse_pg0_url, resolve_database_url
|
||||
|
||||
@@ -68,6 +73,7 @@ BACKUP_TABLES = [
|
||||
"audit_log",
|
||||
"llm_requests",
|
||||
"graph_maintenance_queue",
|
||||
"entity_maintenance_queue",
|
||||
]
|
||||
|
||||
MANIFEST_VERSION = "2"
|
||||
@@ -617,11 +623,14 @@ async def _run_repair_bank(
|
||||
schema: str | None,
|
||||
bank_id: str | None,
|
||||
dry_run: bool,
|
||||
) -> list[SchemaVectorIndexResult]:
|
||||
) -> list[BankIndexResult]:
|
||||
"""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.
|
||||
cannot run inside a transaction block.
|
||||
|
||||
Deliberately unbudgeted, unlike the background operation: this is an operator
|
||||
asking for convergence now, across as many banks as they named.
|
||||
"""
|
||||
schemas = [schema] if schema else await _resolve_schemas(base_schema)
|
||||
index_clause = _vector_index_clause()
|
||||
@@ -630,13 +639,38 @@ async def _run_repair_bank(
|
||||
assert index_clause is not None
|
||||
|
||||
conn = await _admin_connect(db_url)
|
||||
results: list[BankIndexResult] = []
|
||||
try:
|
||||
results = await repair_vector_indexes(conn, schemas, index_clause, dry_run=dry_run, bank_id=bank_id)
|
||||
for result in results:
|
||||
for target_schema in schemas:
|
||||
try:
|
||||
bank_ids = [bank_id] if bank_id else await list_bank_ids(conn, target_schema)
|
||||
except Exception as exc: # noqa: BLE001 — one bad schema must not abort the sweep
|
||||
typer.echo(f" schema '{target_schema}': skipped ({exc})", err=True)
|
||||
continue
|
||||
schema_results = [
|
||||
await reconcile_bank_vector_indexes(conn, target_schema, bid, index_clause, dry_run=dry_run)
|
||||
for bid in bank_ids
|
||||
]
|
||||
results.extend(schema_results)
|
||||
# Only in --all mode: an index whose bank row is gone is unreachable
|
||||
# from every bank-scoped path, so this is the one place that can
|
||||
# collect it. Normally finds nothing — delete_bank drops a bank's
|
||||
# indexes while it still knows their names — but a deployment that
|
||||
# hit the #3485 wall could not run delete_bank at all.
|
||||
orphans = [] if bank_id else await drop_orphaned_bank_indexes(conn, target_schema, dry_run=dry_run)
|
||||
if orphans:
|
||||
typer.echo(
|
||||
f" schema '{target_schema}': {len(orphans)} orphaned index(es) "
|
||||
f"{'to drop (dry-run)' if dry_run else 'dropped'} (no matching bank)"
|
||||
)
|
||||
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"
|
||||
f" schema '{target_schema}': {len(bank_ids)} bank(s) scanned, "
|
||||
f"{sum(r.already_present for r in schema_results)} present, "
|
||||
f"{sum(r.created for r in schema_results)} created, "
|
||||
f"{sum(r.dropped for r in schema_results)} dropped, "
|
||||
f"{sum(r.skipped for r in schema_results)} to-create (dry-run), "
|
||||
f"{sum(r.would_drop for r in schema_results)} to-drop (dry-run), "
|
||||
f"{sum(r.failed for r in schema_results)} failed"
|
||||
)
|
||||
return results
|
||||
finally:
|
||||
@@ -668,17 +702,23 @@ def repair_bank(
|
||||
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.
|
||||
"""Reconcile per-(bank, fact_type) vector index coverage against the size threshold.
|
||||
|
||||
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.
|
||||
A (bank, fact_type) earns a partial vector index once it holds
|
||||
HINDSIGHT_API_VECTOR_INDEX_MIN_ROWS rows; below that the planner answers the
|
||||
same query exactly, and faster, from the (bank_id, fact_type) B-tree plus a
|
||||
top-N sort. This command builds what qualifies and drops what no longer does
|
||||
— including indexes orphaned by a deleted bank — detecting invalid coverage
|
||||
too (an INVALID leftover, or an index whose access method drifted after a
|
||||
backend switch, counts as missing). All DDL is CONCURRENTLY, so it never
|
||||
blocks the live fleet.
|
||||
|
||||
Writes keep this converged on their own — every insert that could move a bank
|
||||
across the threshold queues a vector_index_maintenance operation. Reach for
|
||||
the command when you want convergence without waiting for a write: after a
|
||||
restore or upgrade, after a backend switch, or to shed indexes in bulk on a
|
||||
deployment recovering from lock-table exhaustion (#3485). Idempotent and safe
|
||||
to re-run.
|
||||
"""
|
||||
if bool(bank_id) == all_banks:
|
||||
typer.echo("Error: pass exactly one of --bank <id> or --all.", err=True)
|
||||
@@ -712,15 +752,18 @@ def repair_bank(
|
||||
)
|
||||
)
|
||||
|
||||
total_banks = sum(r.banks_scanned for r in results)
|
||||
total_banks = len(results)
|
||||
total_present = sum(r.already_present for r in results)
|
||||
total_created = sum(r.created for r in results)
|
||||
total_dropped = sum(r.dropped for r in results)
|
||||
total_skipped = sum(r.skipped for r in results)
|
||||
total_would_drop = sum(r.would_drop 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"
|
||||
f"{total_present} already present, {total_created} created, {total_dropped} dropped, "
|
||||
f"{total_skipped} to-create (dry-run), {total_would_drop} to-drop (dry-run), "
|
||||
f"{total_failed} failed"
|
||||
)
|
||||
if total_failed:
|
||||
failed_names = [name for r in results for name in r.failed_indexes]
|
||||
|
||||
+118
@@ -0,0 +1,118 @@
|
||||
"""Add entity_maintenance_queue table (+ seed it with every existing entity)
|
||||
|
||||
Queue of entities whose unit references may have gone away — the input to the
|
||||
graph_maintenance job's orphan-entity and stale-cooccurrence prunes.
|
||||
|
||||
Those two prunes used to be bank-wide single statements re-evaluated on every
|
||||
run: the orphan prune probed once per entity in the bank, and the cooccurrence
|
||||
prune evaluated an INTERSECT per cooccurrence row in the bank, whether or not
|
||||
anything had changed. Their cost tracked the size of the bank rather than the
|
||||
size of the delete, so past a few million rows they blew asyncpg's command
|
||||
timeout on every run and the job could never complete (#3222).
|
||||
|
||||
With a queue the prunes only examine entities a delete actually touched, the
|
||||
same way ``graph_maintenance_queue`` already scopes the relink pass.
|
||||
|
||||
Deliberately NOT seeded with the existing entities. Backfilling them would
|
||||
reclaim whatever a bank accumulated while its sweep was failing, but it writes
|
||||
one row per entity inside a migration that runs at API startup, and then charges
|
||||
a prune check for every one of them — a slow upgrade plus a large self-inflicted
|
||||
backlog, to collect rows that cost a bank nothing. The queue starts empty and
|
||||
fills from real deletes; historical strays stay until something touches them.
|
||||
|
||||
Revision ID: c4f7a91b2d38
|
||||
Revises: d9c1a7b4e2f6
|
||||
Create Date: 2026-08-11
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "c4f7a91b2d38"
|
||||
down_revision: str | Sequence[str] | None = "d9c1a7b4e2f6"
|
||||
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()
|
||||
# Composite PK gives ON CONFLICT DO NOTHING dedup when the same entity is
|
||||
# enqueued from overlapping deletes. No FK to entities: the prune's whole
|
||||
# job is to delete the entity, and a cascade would race it away mid-drain.
|
||||
# A queue row naming an entity that no longer exists is a no-op.
|
||||
op.execute(
|
||||
f"""
|
||||
CREATE TABLE IF NOT EXISTS {schema}entity_maintenance_queue (
|
||||
bank_id TEXT NOT NULL,
|
||||
entity_id UUID NOT NULL,
|
||||
enqueued_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
||||
PRIMARY KEY (bank_id, entity_id)
|
||||
)
|
||||
"""
|
||||
)
|
||||
op.execute(
|
||||
f"""
|
||||
CREATE INDEX IF NOT EXISTS idx_entity_maintenance_queue_bank_enqueued
|
||||
ON {schema}entity_maintenance_queue (bank_id, enqueued_at)
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema = _pg_schema_prefix()
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_entity_maintenance_queue_bank_enqueued")
|
||||
op.execute(f"DROP TABLE IF EXISTS {schema}entity_maintenance_queue")
|
||||
|
||||
|
||||
def _oracle_execute_ignoring_955(sql: str) -> None:
|
||||
"""Run a CREATE statement and swallow ORA-00955 (object already exists).
|
||||
|
||||
Mirrors the helper in the graph_maintenance_queue migration so reruns stay
|
||||
safe on a database where the table was created by an earlier partial run.
|
||||
"""
|
||||
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.strip()})
|
||||
|
||||
|
||||
def _oracle_upgrade() -> None:
|
||||
_oracle_execute_ignoring_955(
|
||||
"""
|
||||
CREATE TABLE entity_maintenance_queue (
|
||||
bank_id VARCHAR2(256) NOT NULL,
|
||||
entity_id RAW(16) NOT NULL,
|
||||
enqueued_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
|
||||
CONSTRAINT pk_entity_maintenance_queue PRIMARY KEY (bank_id, entity_id)
|
||||
)
|
||||
"""
|
||||
)
|
||||
_oracle_execute_ignoring_955(
|
||||
"CREATE INDEX idx_entity_maintenance_queue_bank_enqueued ON entity_maintenance_queue (bank_id, enqueued_at)"
|
||||
)
|
||||
|
||||
|
||||
def _oracle_downgrade() -> None:
|
||||
op.execute("DROP INDEX idx_entity_maintenance_queue_bank_enqueued")
|
||||
op.execute("DROP TABLE entity_maintenance_queue")
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade, oracle=_oracle_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade, oracle=_oracle_downgrade)
|
||||
+300
@@ -0,0 +1,300 @@
|
||||
"""Make the cross-schema maintenance routines skip a schema under concurrent DDL.
|
||||
|
||||
``banks_needing_consolidation()``, ``mental_models_with_cron()``,
|
||||
``schemas_with_expired_rows(...)`` and ``schemas_with_expired_operations(...)``
|
||||
snapshot the schemas owning a target table from ``pg_class`` and then query each
|
||||
schema in turn, inside one transaction. Every such query takes AccessShareLock on
|
||||
two or three relations, and those locks are held until the caller commits.
|
||||
|
||||
``c7e9f1a3b5d2`` already handles the schema *vanishing* mid-scan. The same race
|
||||
has a second outcome: the schema is not gone, it is being rewritten, and its DDL
|
||||
holds — or is queued for — AccessExclusiveLock. A queued AccessExclusiveLock
|
||||
blocks later AccessShareLock requests, so::
|
||||
|
||||
routine holds AccessShare(memory_units) -> wants AccessShare(banks)
|
||||
dropper queued AccessExclusive(banks) -> wants AccessExclusive(memory_units)
|
||||
|
||||
is a cycle, and PostgreSQL breaks it by killing one side. When it picks the
|
||||
routine the whole scan aborts, so one tenant being dropped takes out an entire
|
||||
maintenance pass. Observed as a recurring ``DeadlockDetectedError`` in the test
|
||||
suite, where xdist workers create and drop schemas continuously while
|
||||
``test_maintenance_routines`` calls the routines against the same database; in
|
||||
production the background maintenance loop races tenant deletion and migration
|
||||
the same way.
|
||||
|
||||
Fix the routine's side of the cycle: give each per-schema query a short
|
||||
``lock_timeout`` so it abandons the wait long before the deadlock detector runs,
|
||||
and skip that schema. A schema mid-DDL has nothing useful to report anyway, and
|
||||
the maintenance loop runs on a ticker, so it is picked up on the next pass. Locks
|
||||
already held from earlier schemas stay until the caller commits — that is fine,
|
||||
the point is only that this routine stops *waiting* on the other party.
|
||||
|
||||
``lock_timeout`` is set via ``set_config(..., is_local => true)`` rather than
|
||||
``SET LOCAL``: PL/pgSQL rejects the ``SET`` command inside a non-volatile
|
||||
function, and these are all ``STABLE``. The previous value is restored before
|
||||
returning so the caller's transaction is left as it was found. Only conflicting
|
||||
DDL can trigger it — AccessShareLock does not conflict with ordinary DML — so
|
||||
this never fires on a merely busy table.
|
||||
|
||||
Downgrade is a no-op: the bodies here are the ones from ``b6d2f8a4c1e7`` /
|
||||
``d7b2f8a1c934`` plus strictly-additive resilience, with identical signatures and
|
||||
results, so leaving them in place is harmless. Downgrading past those migrations
|
||||
restores or drops them as they define.
|
||||
|
||||
Revision ID: c8b4e2a71f95
|
||||
Revises: e7c3a91f4b62
|
||||
Create Date: 2026-08-17
|
||||
"""
|
||||
|
||||
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 = "c8b4e2a71f95"
|
||||
down_revision: str | Sequence[str] | None = "e7c3a91f4b62"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
# Short enough to abandon the wait before PostgreSQL's deadlock detector runs
|
||||
# (deadlock_timeout defaults to 1s), long enough to ride out a brief DDL
|
||||
# statement rather than skipping a healthy schema.
|
||||
_LOCK_TIMEOUT = "250ms"
|
||||
|
||||
# Both outcomes of the same race, kept as separate arms so each reason is legible
|
||||
# at the point it is handled.
|
||||
_SKIP_ARMS = """
|
||||
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;
|
||||
-- Schema is mid-DDL and holds (or has queued) an
|
||||
-- AccessExclusiveLock. Skip it rather than wait: waiting is
|
||||
-- what closes the deadlock cycle. deadlock_detected is
|
||||
-- belt-and-braces for a cycle formed before lock_timeout.
|
||||
WHEN lock_not_available OR deadlock_detected THEN
|
||||
CONTINUE;
|
||||
"""
|
||||
|
||||
|
||||
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 (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 _pg_upgrade() -> None:
|
||||
# Tenant schemas carry no copy of these routines; only the configured
|
||||
# schema's copy is ever called. Non-install runs have nothing to replace —
|
||||
# and unlike b6d2f8a4c1e7 there are no stray per-tenant copies to clean up,
|
||||
# that migration already did it.
|
||||
if not _is_install_run():
|
||||
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;
|
||||
prev_lock_timeout text;
|
||||
BEGIN
|
||||
prev_lock_timeout := current_setting('lock_timeout');
|
||||
PERFORM set_config('lock_timeout', '{_LOCK_TIMEOUT}', true);
|
||||
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);
|
||||
{_SKIP_ARMS} END;
|
||||
END LOOP;
|
||||
PERFORM set_config('lock_timeout', prev_lock_timeout, true);
|
||||
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;
|
||||
prev_lock_timeout text;
|
||||
BEGIN
|
||||
IF p_days IS NULL OR p_days <= 0 THEN
|
||||
RETURN;
|
||||
END IF;
|
||||
prev_lock_timeout := current_setting('lock_timeout');
|
||||
PERFORM set_config('lock_timeout', '{_LOCK_TIMEOUT}', true);
|
||||
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;
|
||||
{_SKIP_ARMS} END;
|
||||
IF has_expired THEN
|
||||
RETURN NEXT sch;
|
||||
END IF;
|
||||
END LOOP;
|
||||
PERFORM set_config('lock_timeout', prev_lock_timeout, true);
|
||||
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;
|
||||
prev_lock_timeout text;
|
||||
BEGIN
|
||||
prev_lock_timeout := current_setting('lock_timeout');
|
||||
PERFORM set_config('lock_timeout', '{_LOCK_TIMEOUT}', true);
|
||||
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);
|
||||
{_SKIP_ARMS} END;
|
||||
END LOOP;
|
||||
PERFORM set_config('lock_timeout', prev_lock_timeout, true);
|
||||
END;
|
||||
$fn$;
|
||||
"""
|
||||
)
|
||||
|
||||
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;
|
||||
prev_lock_timeout text;
|
||||
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;
|
||||
prev_lock_timeout := current_setting('lock_timeout');
|
||||
PERFORM set_config('lock_timeout', '{_LOCK_TIMEOUT}', true);
|
||||
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;
|
||||
{_SKIP_ARMS} END;
|
||||
IF has_expired THEN
|
||||
RETURN NEXT sch;
|
||||
END IF;
|
||||
END LOOP;
|
||||
PERFORM set_config('lock_timeout', prev_lock_timeout, true);
|
||||
END;
|
||||
$fn$;
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
# No-op by design — see the module docstring. The bodies installed here are
|
||||
# the previous ones plus a skip arm; dropping the routines would strand the
|
||||
# migrations that claim to own them, and re-installing the old bodies would
|
||||
# duplicate their definitions here.
|
||||
return
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
# Oracle slot intentionally absent: these routines are PostgreSQL-only, and
|
||||
# the Oracle worker keeps its per-schema sweep.
|
||||
run_for_dialect(pg=_pg_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade)
|
||||
+95
@@ -0,0 +1,95 @@
|
||||
"""Add async_operations.serialization_key for per-document retain serialization.
|
||||
|
||||
``update_mode="append"`` is a read-modify-write over the whole document: the
|
||||
retain reads ``documents.original_text``, concatenates the new content onto it,
|
||||
and reprocesses the result. Two appends to one document whose read→write
|
||||
windows overlap therefore lose an update — the loser's turn is content nobody
|
||||
else has.
|
||||
|
||||
The orchestrator now detects that at write time and fails the loser instead of
|
||||
committing over it, but detection alone turns lost data into wasted extraction.
|
||||
This column lets the worker's claim query keep a document to one in-flight
|
||||
retain at a time, so the conflict is avoided rather than paid for: a second
|
||||
retain for the same document simply is not claimed until the first finishes,
|
||||
and the waiting operation holds no worker slot while it waits.
|
||||
|
||||
It carries the single document an operation targets (NULL when it targets none
|
||||
or several), so the claim predicate can compare it without digging into
|
||||
``task_payload`` — a shape both dialects index cheaply and which the Oracle
|
||||
rewrite of the claim SQL can handle.
|
||||
|
||||
The partial index covers only live rows: claims never look at terminal
|
||||
operations, and retain queues are dominated by completed history.
|
||||
|
||||
Revision ID: d9c1a7b4e2f6
|
||||
Revises: b3e8d1c6f4a9
|
||||
Create Date: 2026-08-11
|
||||
"""
|
||||
|
||||
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 = "d9c1a7b4e2f6"
|
||||
down_revision: str | Sequence[str] | None = "b3e8d1c6f4a9"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
_INDEX = "idx_async_operations_serialization_key"
|
||||
|
||||
|
||||
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 = context.config.get_main_option("target_schema")
|
||||
op.add_column(
|
||||
"async_operations",
|
||||
sa.Column("serialization_key", sa.Text(), nullable=True),
|
||||
schema=schema or None,
|
||||
)
|
||||
prefix = _pg_schema_prefix()
|
||||
op.execute(
|
||||
f"CREATE INDEX IF NOT EXISTS {_INDEX} ON {prefix}async_operations "
|
||||
f"(bank_id, serialization_key) "
|
||||
f"WHERE serialization_key IS NOT NULL AND status IN ('pending', 'processing')"
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
prefix = _pg_schema_prefix()
|
||||
op.execute(f"DROP INDEX IF EXISTS {prefix}{_INDEX}")
|
||||
op.drop_column("async_operations", "serialization_key", schema=schema or None)
|
||||
|
||||
|
||||
def _oracle_upgrade() -> None:
|
||||
op.add_column("async_operations", sa.Column("serialization_key", sa.String(4000), nullable=True))
|
||||
# Oracle has no partial indexes. A function-based index on the same
|
||||
# predicate gets the equivalent selectivity: terminal rows collapse to NULL
|
||||
# and Oracle does not store all-NULL entries, so the index only holds the
|
||||
# live rows the claim query looks at.
|
||||
op.get_bind().exec_driver_sql(
|
||||
f"CREATE INDEX {_INDEX} ON async_operations ("
|
||||
f" CASE WHEN status IN ('pending', 'processing') THEN bank_id END,"
|
||||
f" CASE WHEN status IN ('pending', 'processing') THEN serialization_key END)"
|
||||
)
|
||||
|
||||
|
||||
def _oracle_downgrade() -> None:
|
||||
op.get_bind().exec_driver_sql(f"DROP INDEX {_INDEX}")
|
||||
op.drop_column("async_operations", "serialization_key")
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade, oracle=_oracle_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade, oracle=_oracle_downgrade)
|
||||
+99
@@ -0,0 +1,99 @@
|
||||
"""Add last_memory_seen_at to mental_models, splitting it from last_refreshed_at.
|
||||
|
||||
``last_refreshed_at`` carried two meanings at once: the wall-clock time of the
|
||||
last refresh, and the source-data watermark (the newest in-scope memory the
|
||||
refresh saw) that staleness keys off. A refresh persisted the watermark into it,
|
||||
and the watermark is clamped so it never regresses — so on a model whose scope
|
||||
gained no new memories the refresh wrote back the value already there. The
|
||||
document was rewritten, the timestamp never moved, and a client asking
|
||||
"have I already refreshed this?" refreshed it again on every tick.
|
||||
|
||||
``last_memory_seen_at`` takes over the watermark meaning; ``last_refreshed_at``
|
||||
goes back to being what its name says. The new column is backfilled from
|
||||
``last_refreshed_at`` — which today holds the watermark — so staleness decides
|
||||
exactly as it did before the migration and no bank mass-refreshes on deploy.
|
||||
Nullable, so consumers COALESCE back to ``last_refreshed_at`` for any row a
|
||||
refresh has not stamped yet.
|
||||
|
||||
Revision ID: e7c3a91f4b62
|
||||
Revises: c4f7a91b2d38
|
||||
Create Date: 2026-08-17
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "e7c3a91f4b62"
|
||||
down_revision: str | Sequence[str] | None = "c4f7a91b2d38"
|
||||
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}mental_models
|
||||
ADD COLUMN IF NOT EXISTS last_memory_seen_at TIMESTAMP WITH TIME ZONE
|
||||
"""
|
||||
)
|
||||
# last_refreshed_at currently holds the watermark, so copying it carries each
|
||||
# model's staleness decision across the cutover unchanged. Only stamp rows
|
||||
# still NULL, so re-running the migration is a no-op rather than a rollback of
|
||||
# watermarks that refreshes have since advanced.
|
||||
op.execute(
|
||||
f"""
|
||||
UPDATE {schema}mental_models
|
||||
SET last_memory_seen_at = last_refreshed_at
|
||||
WHERE last_memory_seen_at IS NULL
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema = _pg_schema_prefix()
|
||||
op.execute(f"ALTER TABLE {schema}mental_models DROP COLUMN IF EXISTS last_memory_seen_at")
|
||||
|
||||
|
||||
def _oracle_upgrade() -> None:
|
||||
# Oracle has no ADD COLUMN IF NOT EXISTS; guard on the data dictionary so a
|
||||
# re-run doesn't fail with ORA-01430 (column already exists).
|
||||
op.get_bind().exec_driver_sql(
|
||||
"""
|
||||
DECLARE
|
||||
n NUMBER;
|
||||
BEGIN
|
||||
SELECT COUNT(*) INTO n FROM user_tab_columns
|
||||
WHERE table_name = 'MENTAL_MODELS'
|
||||
AND column_name = 'LAST_MEMORY_SEEN_AT';
|
||||
IF n = 0 THEN
|
||||
EXECUTE IMMEDIATE
|
||||
'ALTER TABLE mental_models ADD (last_memory_seen_at TIMESTAMP WITH TIME ZONE)';
|
||||
END IF;
|
||||
END;
|
||||
"""
|
||||
)
|
||||
op.get_bind().exec_driver_sql(
|
||||
"UPDATE mental_models SET last_memory_seen_at = last_refreshed_at WHERE last_memory_seen_at IS NULL"
|
||||
)
|
||||
|
||||
|
||||
def _oracle_downgrade() -> None:
|
||||
op.get_bind().exec_driver_sql("ALTER TABLE mental_models DROP COLUMN last_memory_seen_at")
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade, oracle=_oracle_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade, oracle=_oracle_downgrade)
|
||||
+72
@@ -0,0 +1,72 @@
|
||||
"""Add a partial index for the cron-scheduled mental model discovery sweep.
|
||||
|
||||
``mental_models_with_cron()`` (``f4d1c2b3a5e6``, currently installed by
|
||||
``c8b4e2a71f95``) is a cross-tenant discovery routine: it loops over every schema
|
||||
holding a ``mental_models`` table and, for each, selects the models carrying a
|
||||
non-empty ``trigger->>'refresh_cron'``. No index covers that predicate, so each
|
||||
per-schema probe is a **sequential scan** of that tenant's ``mental_models``
|
||||
table — paid on every maintenance tick, in every API/worker process, whether or
|
||||
not the tenant has a single cron-scheduled model.
|
||||
|
||||
Cron-scheduled models are rare by construction (the trigger defaults to
|
||||
``{"refresh_after_consolidation": false}``), so at thousands of tenants the sweep
|
||||
spends essentially all of its time proving that tenants have nothing to do. A
|
||||
partial index whose predicate matches the routine's WHERE clause exactly turns a
|
||||
tenant with no cron-scheduled models into an empty index scan.
|
||||
|
||||
``bank_id`` is the indexed column so the routine's projection stays on the
|
||||
leading column of the index; the predicate is what does the work here.
|
||||
|
||||
PostgreSQL only: the maintenance loop and its discovery routines are PG-only
|
||||
(the Oracle slot is intentionally absent, mirroring ``f4d1c2b3a5e6``), so an
|
||||
Oracle deployment never runs the scan this index exists to avoid.
|
||||
|
||||
Revision ID: f2a7c9d4b168
|
||||
Revises: c8b4e2a71f95
|
||||
Create Date: 2026-08-17
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "f2a7c9d4b168"
|
||||
down_revision: str | Sequence[str] | None = "c8b4e2a71f95"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
_INDEX = "idx_mental_models_cron"
|
||||
|
||||
|
||||
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()
|
||||
# Plain (non-CONCURRENT) build: mental_models holds one row per mental model,
|
||||
# so this is a sub-second SHARE lock even on large installations — unlike the
|
||||
# async_operations indexes in a8c1e4f7b0d3, which needed CONCURRENTLY.
|
||||
# The predicate is character-for-character the routine's WHERE clause, which
|
||||
# is what lets the planner match the partial index.
|
||||
op.execute(
|
||||
f"CREATE INDEX IF NOT EXISTS {_INDEX} ON {schema}mental_models (bank_id) "
|
||||
"WHERE COALESCE(\"trigger\"->>'refresh_cron', '') <> ''"
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema = _pg_schema_prefix()
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}{_INDEX}")
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade)
|
||||
@@ -21,6 +21,7 @@ from fastapi.responses import JSONResponse
|
||||
|
||||
from hindsight_api.api import page_markdown
|
||||
from hindsight_api.api.disconnect import ClientDisconnectCancellationMiddleware, get_scope_cancellation_token
|
||||
from hindsight_api.api.passthrough_headers import collect_passthrough_headers
|
||||
from hindsight_api.cancellation import OperationCancelledError
|
||||
from hindsight_api.engine.audit import (
|
||||
AuditEntry,
|
||||
@@ -598,6 +599,14 @@ class RecallResponse(BaseModel):
|
||||
source_facts: dict[str, RecallResult] | None = Field(
|
||||
default=None, description="Source facts for observation-type results, keyed by fact ID"
|
||||
)
|
||||
source_facts_truncated: bool | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Whether the source_facts map was cut short by the token budget. When true, some IDs in "
|
||||
"results[].source_fact_ids have no entry in source_facts — the budget ran out, the "
|
||||
"references are not dangling. Only set when source facts were requested."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class EntityInput(BaseModel):
|
||||
@@ -645,6 +654,17 @@ class MemoryItem(BaseModel):
|
||||
default=None,
|
||||
description="Optional entities to combine with auto-extracted entities.",
|
||||
)
|
||||
resolve_entities: bool = Field(
|
||||
default=True,
|
||||
description="Whether the names in 'entities' are resolved against the entities already in "
|
||||
"the bank. True (default) matches each name to a similar existing entity when it scores "
|
||||
"above the match threshold, so a name close to one already in the bank may resolve to that "
|
||||
"one instead of the one you wrote. False takes your names literally — an existing entity is "
|
||||
"reused only on a case-insensitive name match, any other name creates a new entity, and "
|
||||
"your names are never merged with each other. This applies only to the entities you supply "
|
||||
"here; auto-extracted entities are always resolved, since they are the extractor's guess at "
|
||||
"a name rather than yours. Ignored when 'entities' is omitted.",
|
||||
)
|
||||
tags: list[str] | None = Field(
|
||||
default=None,
|
||||
description="Optional tags for visibility scoping. Memories with tags can be filtered during recall.",
|
||||
@@ -967,17 +987,21 @@ class ReflectRequest(BaseModel):
|
||||
)
|
||||
tags: list[str] | None = Field(
|
||||
default=None,
|
||||
description="Filter memories by tags during reflection. If not specified, all memories are considered.",
|
||||
description="Scope raw facts, observations, mental models, and tagged directives during reflection. "
|
||||
"With no tags, memory retrieval is unfiltered while only untagged/global directives are loaded. "
|
||||
"Use tags=[] with tags_match='exact' to select the untagged/global scope.",
|
||||
)
|
||||
tags_match: TagsMatch = Field(
|
||||
default="any",
|
||||
description="How to match tags: 'any' (OR, includes untagged), 'all' (AND, includes untagged), "
|
||||
"'any_strict' (OR, excludes untagged), 'all_strict' (AND, excludes untagged).",
|
||||
"'any_strict' (OR, excludes untagged), 'all_strict' (AND, excludes untagged), or "
|
||||
"'exact' (set equality). Untagged directives remain global in every mode.",
|
||||
)
|
||||
tag_groups: list[TagGroup] | None = Field(
|
||||
default=None,
|
||||
description="Compound tag filter using boolean groups. Groups in the list are AND-ed. "
|
||||
"Each group is a leaf {tags, match} or compound {and: [...]}, {or: [...]}, {not: ...}.",
|
||||
"Each group is a leaf {tags, match} or compound {and: [...]}, {or: [...]}, {not: ...}. "
|
||||
"Mutually exclusive with tags.",
|
||||
)
|
||||
apply_all_directives: bool = Field(
|
||||
default=False,
|
||||
@@ -1269,7 +1293,7 @@ class BankListItem(BaseModel):
|
||||
|
||||
|
||||
class BankListResponse(BaseModel):
|
||||
"""Response model for listing all banks."""
|
||||
"""Response model for listing banks, one page at a time."""
|
||||
|
||||
model_config = ConfigDict(
|
||||
json_schema_extra={
|
||||
@@ -1286,12 +1310,18 @@ class BankListResponse(BaseModel):
|
||||
"last_document_at": "2024-01-16T14:20:00Z",
|
||||
"last_write_at": "2024-01-17T09:05:00Z",
|
||||
}
|
||||
]
|
||||
],
|
||||
"total": 50,
|
||||
"limit": 100,
|
||||
"offset": 0,
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
banks: list[BankListItem]
|
||||
total: int = Field(description="Total number of banks visible to the caller, ignoring `limit`/`offset`.")
|
||||
limit: int
|
||||
offset: int
|
||||
|
||||
|
||||
class CreateBankRequest(BaseModel):
|
||||
@@ -1438,8 +1468,7 @@ class BankConfigUpdate(BaseModel):
|
||||
json_schema_extra={
|
||||
"example": {
|
||||
"updates": {
|
||||
"llm_model": "claude-sonnet-4-5",
|
||||
"retain_extraction_mode": "verbose",
|
||||
"retain_extraction_mode": "custom",
|
||||
"retain_custom_instructions": "Extract technical details carefully",
|
||||
}
|
||||
}
|
||||
@@ -1447,8 +1476,8 @@ class BankConfigUpdate(BaseModel):
|
||||
)
|
||||
|
||||
updates: dict[str, Any] = Field(
|
||||
description="Configuration overrides. Keys can be in Python field format (llm_provider) "
|
||||
"or environment variable format (HINDSIGHT_API_LLM_PROVIDER). "
|
||||
description="Configuration overrides. Keys can be in Python field format (retain_extraction_mode) "
|
||||
"or environment variable format (HINDSIGHT_API_RETAIN_EXTRACTION_MODE). "
|
||||
"Only hierarchical fields can be overridden per-bank."
|
||||
)
|
||||
|
||||
@@ -1461,12 +1490,10 @@ class BankConfigResponse(BaseModel):
|
||||
"example": {
|
||||
"bank_id": "my-bank",
|
||||
"config": {
|
||||
"llm_provider": "openai",
|
||||
"llm_model": "gpt-4",
|
||||
"retain_extraction_mode": "verbose",
|
||||
"retain_chunk_size": 3000,
|
||||
},
|
||||
"overrides": {
|
||||
"llm_model": "gpt-4",
|
||||
"retain_extraction_mode": "verbose",
|
||||
},
|
||||
}
|
||||
@@ -1783,8 +1810,19 @@ class UpdateMemoryRequest(BaseModel):
|
||||
)
|
||||
entities: list[str] | None = Field(
|
||||
default=None,
|
||||
description="Replace the fact's entities. Names are resolved/find-or-created "
|
||||
"the same way retain does; '[]' detaches all entities. Omit to leave unchanged.",
|
||||
description="Replace the fact's entities. How each name is matched to an entity is "
|
||||
"governed by 'resolve_entities'. '[]' detaches all entities. Omit to leave unchanged.",
|
||||
)
|
||||
resolve_entities: bool = Field(
|
||||
default=True,
|
||||
description="Whether the names in 'entities' are resolved against the entities already in "
|
||||
"the bank. True (default) is what retain does: a similar existing entity is reused when it "
|
||||
"scores above the match threshold, so a name close to one already in the bank may resolve "
|
||||
"to that one instead of the one you wrote. False takes the names literally — an existing "
|
||||
"entity is reused only on a case-insensitive name match, any other name creates a new "
|
||||
"entity, and names in the same request are never merged with each other. Use False for "
|
||||
"hand-authored corrections, where the name you sent is the answer rather than a guess. "
|
||||
"Ignored when 'entities' is omitted.",
|
||||
)
|
||||
state: str | None = Field(
|
||||
default=None,
|
||||
@@ -1980,12 +2018,19 @@ class BankStatsResponse(BaseModel):
|
||||
default=None,
|
||||
description=(
|
||||
"When a memory was last written in this bank — stored, edited, or consolidated (ISO format). "
|
||||
"Null if the bank has no memories. A mental model whose `last_refreshed_at` is at or after this "
|
||||
"is up to date whatever its tags; an older one may need a refresh, which only the single "
|
||||
"Null if the bank has no memories. A mental model whose `last_memory_seen_at` is at or after "
|
||||
"this is up to date whatever its tags; an older one may need a refresh, which only the single "
|
||||
"mental-model read can confirm."
|
||||
),
|
||||
)
|
||||
pending_consolidation: int = Field(default=0, description="Number of memories not yet processed into observations")
|
||||
pending_consolidation: int = Field(
|
||||
default=0,
|
||||
description=(
|
||||
"Number of source memories (world/experience) still queued for consolidation into "
|
||||
"observations. Excludes memories whose consolidation permanently failed — those are "
|
||||
"counted only in failed_consolidation — so this drains to 0 when the consolidator catches up."
|
||||
),
|
||||
)
|
||||
failed_consolidation: int = Field(
|
||||
default=0,
|
||||
description="Number of source memories (world/experience) whose consolidation permanently failed and can be retried via the consolidation recovery endpoint.",
|
||||
@@ -2093,6 +2138,9 @@ class DirectiveListResponse(BaseModel):
|
||||
"""Response model for listing directives."""
|
||||
|
||||
items: list[DirectiveResponse]
|
||||
total: int = Field(description="Total number of directives matching the filter (not just this page)")
|
||||
limit: int = Field(description="Page size that was applied")
|
||||
offset: int = Field(description="Offset that was applied")
|
||||
|
||||
|
||||
class CreateDirectiveRequest(BaseModel):
|
||||
@@ -2102,7 +2150,10 @@ class CreateDirectiveRequest(BaseModel):
|
||||
content: str = Field(description="The directive text to inject into prompts")
|
||||
priority: int = Field(default=0, description="Higher priority directives are injected first")
|
||||
is_active: bool = Field(default=True, description="Whether this directive is active")
|
||||
tags: list[str] = FieldWithDefault(list, description="Tags for filtering")
|
||||
tags: list[str] = FieldWithDefault(
|
||||
list,
|
||||
description="Directive execution scope. Empty means global; non-empty requires a matching reflect scope.",
|
||||
)
|
||||
|
||||
|
||||
class UpdateDirectiveRequest(BaseModel):
|
||||
@@ -2273,7 +2324,26 @@ class MentalModelResponse(BaseModel):
|
||||
tags: list[str] = FieldWithDefault(list)
|
||||
max_tokens: int | None = Field(default=None)
|
||||
trigger: MentalModelTrigger | None = Field(default=None)
|
||||
last_refreshed_at: str | None = None
|
||||
last_refreshed_at: str | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"When a refresh last finished for this model — wall-clock, in ISO format. Advances on "
|
||||
"every refresh that completes, including one that found nothing new and preserved the "
|
||||
"content, and on a direct edit of `content`. A refresh that failed leaves it alone. "
|
||||
"This is the field to answer 'have I already refreshed this?'; it says nothing about "
|
||||
"whether the model is behind the data, which is `last_memory_seen_at` / `is_stale`."
|
||||
),
|
||||
)
|
||||
last_memory_seen_at: str | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"How far through the bank's memories this model is written — the newest in-scope memory "
|
||||
"the last refresh saw, in ISO format. Stands still when nothing in the model's scope has "
|
||||
"been written, however often it is refreshed. Compare against `last_memory_write_at` "
|
||||
"from GET /stats to flag a whole list cheaply: at or after it means up to date, older "
|
||||
"means it may need a refresh. Null for a model no refresh has stamped yet."
|
||||
),
|
||||
)
|
||||
created_at: str | None = None
|
||||
reflect_response: dict | None = Field(
|
||||
default=None,
|
||||
@@ -2283,9 +2353,9 @@ class MentalModelResponse(BaseModel):
|
||||
default=None,
|
||||
description=(
|
||||
"True when memories matching this mental model's tag/fact_type scope have been written "
|
||||
"since last_refreshed_at. Exact, and costly to compute, so it is populated only by the "
|
||||
"since last_memory_seen_at. Exact, and costly to compute, so it is populated only by the "
|
||||
"single mental-model read at detail=full — never when listing. For a whole list, compare "
|
||||
"each `last_refreshed_at` against the bank's `last_memory_write_at` from GET /stats: "
|
||||
"each `last_memory_seen_at` against the bank's `last_memory_write_at` from GET /stats: "
|
||||
"at or after it means up to date, older means it may need a refresh."
|
||||
),
|
||||
)
|
||||
@@ -2295,6 +2365,9 @@ class MentalModelListResponse(BaseModel):
|
||||
"""Response model for listing mental models."""
|
||||
|
||||
items: list[MentalModelResponse]
|
||||
total: int = Field(description="Total number of mental models matching the filter (not just this page)")
|
||||
limit: int = Field(description="Page size that was applied")
|
||||
offset: int = Field(description="Offset that was applied")
|
||||
|
||||
|
||||
# =========================================================================
|
||||
@@ -2326,6 +2399,15 @@ class KnowledgeNode(BaseModel):
|
||||
"was written, but possibly outside the page's tags. Read the page's mental model for the exact answer. "
|
||||
"Shares the bank-stats freshness, so it can lag a just-written memory by up to a minute.",
|
||||
)
|
||||
trigger: MentalModelTrigger | None = Field(
|
||||
default=None,
|
||||
description="Pages only: the page's refresh settings — when it rebuilds itself "
|
||||
"(`refresh_after_consolidation` or `refresh_cron`), in which mode, and over which facts. "
|
||||
"This is the EFFECTIVE policy: a setting the page never stored is reported at its default, "
|
||||
"so compare the fields you care about rather than the whole object against a patch you "
|
||||
"sent. Absent on folders, which have no backing mental model, and on a page with no "
|
||||
"trigger stored.",
|
||||
)
|
||||
children: list["KnowledgeNode"] = FieldWithDefault(list)
|
||||
|
||||
|
||||
@@ -2364,6 +2446,16 @@ class UpdateNodeRequest(BaseModel):
|
||||
source_query: str | None = None
|
||||
tags: list[str] | None = None
|
||||
max_tokens: int | None = None
|
||||
trigger: MentalModelTrigger | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Refresh settings to change. Applied as a patch: only the fields present in this "
|
||||
"object are updated, and the rest keep the page's current values — so moving a page "
|
||||
"onto a schedule does not reset how it refreshes. Setting refresh_cron clears "
|
||||
"refresh_after_consolidation and vice versa, since a page refreshes on one or the "
|
||||
"other, never both."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class CreateKnowledgePageResponse(BaseModel):
|
||||
@@ -2432,6 +2524,7 @@ def _knowledge_node_model(node: dict[str, Any]) -> KnowledgeNode:
|
||||
tags=list(node.get("tags") or []) if is_page else [],
|
||||
timestamp=(node.get("last_refreshed_at") if is_page else node.get("updated_at")),
|
||||
is_stale=node.get("is_stale") if is_page else None,
|
||||
trigger=node.get("trigger") if is_page else None,
|
||||
)
|
||||
|
||||
|
||||
@@ -2901,17 +2994,21 @@ async def apply_bank_template_manifest(
|
||||
projected_mental_model_ids = {item.id for item in default_mental_models} & imported_mental_model_ids
|
||||
projected_directive_names = {item.name for item in default_directives} & imported_directive_names
|
||||
|
||||
# limit=None throughout the import path: a create/update decision per imported
|
||||
# resource is only correct against the bank's *whole* set. Under the default
|
||||
# page size a bank with more than 100 models would look like it lacked the
|
||||
# ones past the first page, and the import would create duplicates.
|
||||
existing_by_id: dict[str, dict[str, Any]] = {}
|
||||
if bank_exists and manifest.mental_models:
|
||||
existing = await memory.list_mental_models(bank_id=bank_id, request_context=request_context)
|
||||
existing_by_id = {m["id"]: m for m in existing}
|
||||
existing = await memory.list_mental_models(bank_id=bank_id, limit=None, request_context=request_context)
|
||||
existing_by_id = {m["id"]: m for m in existing.items}
|
||||
|
||||
existing_by_name: dict[str, dict[str, Any]] = {}
|
||||
if bank_exists and manifest.directives:
|
||||
existing_directives = await memory.list_directives(
|
||||
bank_id=bank_id, active_only=False, request_context=request_context
|
||||
bank_id=bank_id, active_only=False, limit=None, request_context=request_context
|
||||
)
|
||||
existing_by_name = {d["name"]: d for d in existing_directives}
|
||||
existing_by_name = {d["name"]: d for d in existing_directives.items}
|
||||
|
||||
bank_writes: list[BankTemplateImportWrite] = []
|
||||
if config_updates:
|
||||
@@ -2955,9 +3052,10 @@ async def apply_bank_template_manifest(
|
||||
if projected_mental_model_ids:
|
||||
provisioned = await memory.list_mental_models(
|
||||
bank_id=bank_id,
|
||||
limit=None,
|
||||
request_context=request_context,
|
||||
)
|
||||
provisioned_by_id = {item["id"]: item for item in provisioned}
|
||||
provisioned_by_id = {item["id"]: item for item in provisioned.items}
|
||||
existing_by_id.update(
|
||||
{
|
||||
item_id: provisioned_by_id[item_id]
|
||||
@@ -2969,9 +3067,10 @@ async def apply_bank_template_manifest(
|
||||
provisioned = await memory.list_directives(
|
||||
bank_id=bank_id,
|
||||
active_only=False,
|
||||
limit=None,
|
||||
request_context=request_context,
|
||||
)
|
||||
provisioned_by_name = {item["name"]: item for item in provisioned}
|
||||
provisioned_by_name = {item["name"]: item for item in provisioned.items}
|
||||
existing_by_name.update(
|
||||
{name: provisioned_by_name[name] for name in projected_directive_names & provisioned_by_name.keys()}
|
||||
)
|
||||
@@ -2999,17 +3098,18 @@ async def apply_default_bank_template_resources(
|
||||
"""Apply only the resources from a server-owned default template."""
|
||||
existing_by_id: dict[str, dict[str, Any]] = {}
|
||||
if manifest.mental_models:
|
||||
existing = await memory.list_mental_models(bank_id=bank_id, request_context=request_context)
|
||||
existing_by_id = {model["id"]: model for model in existing}
|
||||
existing = await memory.list_mental_models(bank_id=bank_id, limit=None, request_context=request_context)
|
||||
existing_by_id = {model["id"]: model for model in existing.items}
|
||||
|
||||
existing_by_name: dict[str, dict[str, Any]] = {}
|
||||
if manifest.directives:
|
||||
existing_directives = await memory.list_directives(
|
||||
bank_id=bank_id,
|
||||
active_only=False,
|
||||
limit=None,
|
||||
request_context=request_context,
|
||||
)
|
||||
existing_by_name = {directive["name"]: directive for directive in existing_directives}
|
||||
existing_by_name = {directive["name"]: directive for directive in existing_directives.items}
|
||||
|
||||
await _apply_bank_template_resources(
|
||||
memory,
|
||||
@@ -3167,6 +3267,15 @@ class OperationResponse(BaseModel):
|
||||
default=None,
|
||||
description="Original filename for file-conversion operations (file_convert_retain); null for other task types.",
|
||||
)
|
||||
mental_model_id: str | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Mental model this operation acted on (refresh_mental_model); null for other task types. "
|
||||
"Without it the list cannot say which model an operation refreshed — `document_id` is null "
|
||||
"for these, and the list carries no result_metadata. The single-operation read exposes the "
|
||||
"same value under `result_metadata`."
|
||||
),
|
||||
)
|
||||
created_at: str
|
||||
updated_at: str | None = Field(
|
||||
default=None,
|
||||
@@ -3228,6 +3337,7 @@ class OperationsListResponse(BaseModel):
|
||||
{
|
||||
"id": "550e8400-e29b-41d4-a716-446655440000",
|
||||
"task_type": "retain",
|
||||
"items_count": 5,
|
||||
"created_at": "2024-01-15T10:30:00Z",
|
||||
"status": "pending",
|
||||
"error_message": None,
|
||||
@@ -3316,7 +3426,7 @@ class OperationStatusResponse(BaseModel):
|
||||
"example": {
|
||||
"operation_id": "550e8400-e29b-41d4-a716-446655440000",
|
||||
"status": "completed",
|
||||
"operation_type": "refresh_mental_models",
|
||||
"operation_type": "refresh_mental_model",
|
||||
"created_at": "2024-01-15T10:30:00Z",
|
||||
"updated_at": "2024-01-15T10:31:30Z",
|
||||
"completed_at": "2024-01-15T10:31:30Z",
|
||||
@@ -3402,15 +3512,19 @@ class VersionResponse(BaseModel):
|
||||
model_config = ConfigDict(
|
||||
json_schema_extra={
|
||||
"example": {
|
||||
"api_version": "0.4.0",
|
||||
"api_version": "0.9.0",
|
||||
"features": {
|
||||
"observations": False,
|
||||
"mcp": True,
|
||||
"worker": True,
|
||||
"bank_config_api": False,
|
||||
"bank_llm_health": True,
|
||||
"file_upload_api": True,
|
||||
"document_export_api": True,
|
||||
"document_import_api": True,
|
||||
"audit_log": False,
|
||||
"llm_trace": False,
|
||||
"store_document_text": True,
|
||||
},
|
||||
}
|
||||
}
|
||||
@@ -3954,15 +4068,20 @@ def _register_routes(app: FastAPI):
|
||||
# Create audit decorator bound to this app's audit logger
|
||||
audited = _make_audited_http(lambda: getattr(app.state, "audit_logger", None))
|
||||
|
||||
def get_request_context(authorization: str | None = Header(default=None)) -> RequestContext:
|
||||
def get_request_context(request: Request, authorization: str | None = Header(default=None)) -> RequestContext:
|
||||
"""
|
||||
Extract request context from Authorization header.
|
||||
Extract request context from the Authorization header.
|
||||
|
||||
Supports:
|
||||
- Bearer token: "Bearer <api_key>"
|
||||
- Direct API key: "<api_key>"
|
||||
|
||||
Returns RequestContext with extracted API key (may be None if no auth header).
|
||||
|
||||
Any header named in HINDSIGHT_API_EXTENSION_PASSTHROUGH_HEADERS is also
|
||||
copied into ``extra_headers`` for extensions to read. That allowlist is
|
||||
empty by default, so no other header reaches extension code unless an
|
||||
operator opts in.
|
||||
"""
|
||||
api_key = None
|
||||
if authorization:
|
||||
@@ -3970,7 +4089,8 @@ def _register_routes(app: FastAPI):
|
||||
api_key = authorization[7:].strip()
|
||||
else:
|
||||
api_key = authorization.strip()
|
||||
return RequestContext(api_key=api_key)
|
||||
extra_headers = collect_passthrough_headers(request.headers.raw, get_config().extension_passthrough_headers)
|
||||
return RequestContext(api_key=api_key, extra_headers=extra_headers)
|
||||
|
||||
def precheck_for(operation: PrecheckOperation):
|
||||
"""
|
||||
@@ -4042,6 +4162,19 @@ def _register_routes(app: FastAPI):
|
||||
content={"detail": str(exc)},
|
||||
)
|
||||
|
||||
# A bank briefly closed to writes — a store migrating it between backends holds it for a few
|
||||
# seconds. 503 + Retry-After rather than a 500: nothing is broken, and the difference decides
|
||||
# whether a client retries or reports a failure to the user.
|
||||
from ..engine.memories.base import StoreWriteUnavailable
|
||||
|
||||
@app.exception_handler(StoreWriteUnavailable)
|
||||
async def store_write_unavailable_handler(request, exc: StoreWriteUnavailable):
|
||||
return JSONResponse(
|
||||
status_code=503,
|
||||
content={"detail": str(exc)},
|
||||
headers={"Retry-After": str(getattr(exc, "retry_after", 30))},
|
||||
)
|
||||
|
||||
async def _readiness_response() -> JSONResponse:
|
||||
"""Shared body of /health and /health/ready: 200 if healthy, 503 if not."""
|
||||
health = await app.state.memory.health_check()
|
||||
@@ -4419,6 +4552,7 @@ def _register_routes(app: FastAPI):
|
||||
occurred_end=occurred_end,
|
||||
new_fact_type=request.fact_type,
|
||||
entities=request.entities,
|
||||
resolve_entities=request.resolve_entities,
|
||||
state=request.state,
|
||||
reason=request.reason,
|
||||
request_context=request_context,
|
||||
@@ -4637,6 +4771,7 @@ def _register_routes(app: FastAPI):
|
||||
entities=entities_response,
|
||||
chunks=chunks_response,
|
||||
source_facts=source_facts_response,
|
||||
source_facts_truncated=core_result.source_facts_truncated,
|
||||
)
|
||||
|
||||
handler_duration = time.time() - handler_start
|
||||
@@ -4831,16 +4966,26 @@ def _register_routes(app: FastAPI):
|
||||
@app.get(
|
||||
"/v1/default/banks",
|
||||
response_model=BankListResponse,
|
||||
summary="List all memory banks",
|
||||
description="Get a list of all agents with their profiles",
|
||||
summary="List memory banks",
|
||||
description=(
|
||||
"List banks with their profiles and summary stats, most recently written first "
|
||||
"(`last_write_at` descending), with pagination and optional search."
|
||||
),
|
||||
operation_id="list_banks",
|
||||
tags=["Banks"],
|
||||
)
|
||||
async def api_list_banks(request_context: RequestContext = Depends(get_request_context)):
|
||||
"""Get list of all banks with their profiles."""
|
||||
async def api_list_banks(
|
||||
q: str | None = Query(None, description="Case-insensitive substring filter on bank ID or name (e.g. 'alice')"),
|
||||
limit: int = Query(default=100, ge=0, description="Maximum number of banks to return"),
|
||||
offset: int = Query(default=0, ge=0, description="Offset for pagination"),
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
):
|
||||
"""Get one page of banks with their profiles."""
|
||||
try:
|
||||
banks = await app.state.memory.list_banks(request_context=request_context)
|
||||
return BankListResponse(banks=banks)
|
||||
data = await app.state.memory.list_banks(
|
||||
search_query=q, limit=limit, offset=offset, request_context=request_context
|
||||
)
|
||||
return BankListResponse(**data)
|
||||
except (AuthenticationError, HTTPException):
|
||||
raise
|
||||
except Exception as e:
|
||||
@@ -5144,7 +5289,7 @@ def _register_routes(app: FastAPI):
|
||||
):
|
||||
"""List mental models for a bank."""
|
||||
try:
|
||||
mental_models = await app.state.memory.list_mental_models(
|
||||
page = await app.state.memory.list_mental_models(
|
||||
bank_id=bank_id,
|
||||
tags=tags_filter,
|
||||
tags_match=tags_match,
|
||||
@@ -5153,7 +5298,12 @@ def _register_routes(app: FastAPI):
|
||||
offset=offset,
|
||||
request_context=request_context,
|
||||
)
|
||||
return MentalModelListResponse(items=[MentalModelResponse(**m) for m in mental_models])
|
||||
return MentalModelListResponse(
|
||||
items=[MentalModelResponse(**m) for m in page.items],
|
||||
total=page.total,
|
||||
limit=limit,
|
||||
offset=offset,
|
||||
)
|
||||
except OperationValidationError as e:
|
||||
raise HTTPException(status_code=e.status_code, detail=e.reason)
|
||||
except (AuthenticationError, HTTPException):
|
||||
@@ -5607,7 +5757,11 @@ def _register_routes(app: FastAPI):
|
||||
parent_id=body.parent_id,
|
||||
tags=body.tags if body.tags else None,
|
||||
max_tokens=body.max_tokens,
|
||||
trigger=body.trigger.model_dump() if body.trigger else None,
|
||||
# Only what the client actually set: the engine merges these over the
|
||||
# knowledge-page defaults, and a full dump would drown them in this model's
|
||||
# own field defaults (mode="full", exclude_mental_models=False) — which is
|
||||
# how every page created with a trigger lost its delta refresh (#3506).
|
||||
trigger=body.trigger.model_dump(exclude_unset=True) if body.trigger else None,
|
||||
request_context=request_context,
|
||||
)
|
||||
if node is None:
|
||||
@@ -5755,8 +5909,10 @@ def _register_routes(app: FastAPI):
|
||||
response_model=KnowledgeNode,
|
||||
summary="Rename/move a knowledge-base node or update a page's options",
|
||||
description="Rename a node (set `name`), move it under another folder (set `parent_id`, null "
|
||||
"for the root), and/or update a page's options (`source_query`, `tags`, `max_tokens`). "
|
||||
"Changing `source_query` schedules an async refresh so the page rebuilds against the new question.",
|
||||
"for the root), and/or update a page's options (`source_query`, `tags`, `max_tokens`, `trigger`). "
|
||||
"Changing `source_query` schedules an async refresh so the page rebuilds against the new question. "
|
||||
"`trigger` is applied as a patch: the fields you send are updated and the rest keep the page's "
|
||||
"current values.",
|
||||
operation_id="update_knowledge_node",
|
||||
tags=["Knowledge Base"],
|
||||
)
|
||||
@@ -5784,7 +5940,7 @@ def _register_routes(app: FastAPI):
|
||||
)
|
||||
# Page options live on the backing mental model; each applies only when
|
||||
# present in the body (so tags=[] clears, distinct from "not provided").
|
||||
page_fields = {"source_query", "tags", "max_tokens"} & body.model_fields_set
|
||||
page_fields = {"source_query", "tags", "max_tokens", "trigger"} & body.model_fields_set
|
||||
if page_fields:
|
||||
did_change = True
|
||||
updated = await app.state.memory.update_knowledge_page(
|
||||
@@ -5793,6 +5949,10 @@ def _register_routes(app: FastAPI):
|
||||
source_query=body.source_query if "source_query" in page_fields else None,
|
||||
tags=body.tags if "tags" in page_fields else None,
|
||||
max_tokens=body.max_tokens if "max_tokens" in page_fields else None,
|
||||
# Only the trigger fields the client stated: the engine patches them over
|
||||
# the page's current trigger, and a full dump would carry this model's own
|
||||
# defaults (mode="full", exclude_mental_models=False) into every update.
|
||||
trigger=(body.trigger.model_dump(exclude_unset=True) if body.trigger else None),
|
||||
request_context=request_context,
|
||||
)
|
||||
# A new source query means the content is stale — rebuild it.
|
||||
@@ -5804,7 +5964,8 @@ def _register_routes(app: FastAPI):
|
||||
)
|
||||
if not did_change:
|
||||
raise HTTPException(
|
||||
status_code=400, detail="Provide name, parent_id, source_query, tags, and/or max_tokens to update"
|
||||
status_code=400,
|
||||
detail="Provide name, parent_id, source_query, tags, max_tokens, and/or trigger to update",
|
||||
)
|
||||
if updated is None:
|
||||
raise HTTPException(status_code=404, detail=f"Knowledge node '{node_id}' not found")
|
||||
@@ -5861,14 +6022,21 @@ def _register_routes(app: FastAPI):
|
||||
"/v1/default/banks/{bank_id}/directives",
|
||||
response_model=DirectiveListResponse,
|
||||
summary="List directives",
|
||||
description="List hard rules that are injected into prompts.",
|
||||
description="List directive definitions. Unlike reflect, an omitted tag filter returns all directives.",
|
||||
operation_id="list_directives",
|
||||
tags=["Directives"],
|
||||
)
|
||||
async def api_list_directives(
|
||||
bank_id: str,
|
||||
tags_filter: list[str] | None = Query(None, alias="tags", description="Filter by tags"),
|
||||
tags_match: Literal["any", "all", "exact"] = Query("any", description="How to match tags"),
|
||||
tags_filter: list[str] | None = Query(
|
||||
None,
|
||||
alias="tags",
|
||||
description="Filter directives by execution scope. Omit or pass [] to list all directives.",
|
||||
),
|
||||
tags_match: Literal["any", "all", "exact"] = Query(
|
||||
"any",
|
||||
description="How tagged directives match the requested scope. Untagged/global directives are included.",
|
||||
),
|
||||
active_only: bool = Query(True, description="Only return active directives"),
|
||||
limit: int = Query(100, ge=1, le=1000),
|
||||
offset: int = Query(0, ge=0),
|
||||
@@ -5876,7 +6044,7 @@ def _register_routes(app: FastAPI):
|
||||
):
|
||||
"""List directives for a bank."""
|
||||
try:
|
||||
directives = await app.state.memory.list_directives(
|
||||
page = await app.state.memory.list_directives(
|
||||
bank_id=bank_id,
|
||||
tags=tags_filter,
|
||||
tags_match=tags_match,
|
||||
@@ -5885,7 +6053,12 @@ def _register_routes(app: FastAPI):
|
||||
offset=offset,
|
||||
request_context=request_context,
|
||||
)
|
||||
return DirectiveListResponse(items=[DirectiveResponse(**d) for d in directives])
|
||||
return DirectiveListResponse(
|
||||
items=[DirectiveResponse(**d) for d in page.items],
|
||||
total=page.total,
|
||||
limit=limit,
|
||||
offset=offset,
|
||||
)
|
||||
except OperationValidationError as e:
|
||||
raise HTTPException(status_code=e.status_code, detail=e.reason)
|
||||
except (AuthenticationError, HTTPException):
|
||||
@@ -5935,7 +6108,7 @@ def _register_routes(app: FastAPI):
|
||||
"/v1/default/banks/{bank_id}/directives",
|
||||
response_model=DirectiveResponse,
|
||||
summary="Create directive",
|
||||
description="Create a hard rule that will be injected into prompts.",
|
||||
description="Create a global or tag-scoped hard rule for reflect prompts.",
|
||||
operation_id="create_directive",
|
||||
tags=["Directives"],
|
||||
)
|
||||
@@ -6984,12 +7157,13 @@ def _register_routes(app: FastAPI):
|
||||
filtered_overrides = {k: v for k, v in bank_overrides.items() if k in template_config_fields}
|
||||
bank_config = BankTemplateConfig(**filtered_overrides) if filtered_overrides else None
|
||||
|
||||
# Get mental models
|
||||
# Get mental models (limit=None — an export that stopped at the
|
||||
# default page size would silently drop the rest of the bank)
|
||||
mental_models_raw = await app.state.memory.list_mental_models(
|
||||
bank_id=bank_id, request_context=request_context
|
||||
bank_id=bank_id, limit=None, request_context=request_context
|
||||
)
|
||||
template_mental_models: list[BankTemplateMentalModel] = []
|
||||
for mm in mental_models_raw:
|
||||
for mm in mental_models_raw.items:
|
||||
trigger_data = mm.get("trigger", {})
|
||||
trigger = MentalModelTrigger(**trigger_data) if trigger_data else MentalModelTrigger()
|
||||
template_mental_models.append(
|
||||
@@ -7003,12 +7177,12 @@ def _register_routes(app: FastAPI):
|
||||
)
|
||||
)
|
||||
|
||||
# Get directives
|
||||
# Get directives (limit=None for the same reason as the models above)
|
||||
directives_raw = await app.state.memory.list_directives(
|
||||
bank_id=bank_id, active_only=False, request_context=request_context
|
||||
bank_id=bank_id, active_only=False, limit=None, request_context=request_context
|
||||
)
|
||||
template_directives: list[BankTemplateDirective] = []
|
||||
for d in directives_raw:
|
||||
for d in directives_raw.items:
|
||||
template_directives.append(
|
||||
BankTemplateDirective(
|
||||
name=d["name"],
|
||||
@@ -7406,8 +7580,9 @@ def _register_routes(app: FastAPI):
|
||||
"/v1/default/banks/{bank_id}/config",
|
||||
response_model=BankConfigResponse,
|
||||
summary="Update bank configuration",
|
||||
description="Update configuration overrides for a bank. Only hierarchical fields can be overridden (LLM settings, retention parameters, etc.). "
|
||||
"Keys can be provided in Python field format (llm_provider) or environment variable format (HINDSIGHT_API_LLM_PROVIDER).",
|
||||
description="Update configuration overrides for a bank. Only hierarchical behavioral settings can be "
|
||||
"overridden (retention parameters, recall settings, etc.). Keys can be provided in Python field format "
|
||||
"(retain_extraction_mode) or environment variable format (HINDSIGHT_API_RETAIN_EXTRACTION_MODE).",
|
||||
operation_id="update_bank_config",
|
||||
tags=["Banks"],
|
||||
)
|
||||
@@ -7869,6 +8044,7 @@ def _register_routes(app: FastAPI):
|
||||
content_dict["document_id"] = item.document_id
|
||||
if item.entities:
|
||||
content_dict["entities"] = [{"text": e.text, "type": e.type or "CONCEPT"} for e in item.entities]
|
||||
content_dict["resolve_entities"] = item.resolve_entities
|
||||
if item.tags:
|
||||
content_dict["tags"] = item.tags
|
||||
if item.observation_scopes is not None:
|
||||
|
||||
@@ -9,6 +9,7 @@ from fastmcp import FastMCP
|
||||
|
||||
from hindsight_api import MemoryEngine
|
||||
from hindsight_api import __version__ as HINDSIGHT_VERSION
|
||||
from hindsight_api.api.passthrough_headers import collect_passthrough_headers
|
||||
from hindsight_api.config import DEFAULT_MCP_RECALL_DESCRIPTION, DEFAULT_MCP_RETAIN_DESCRIPTION, _get_raw_config
|
||||
from hindsight_api.engine.memory_engine import _current_schema
|
||||
from hindsight_api.extensions import MCPExtension, load_extension
|
||||
@@ -52,6 +53,11 @@ _current_api_key_id: ContextVar[str | None] = ContextVar("current_api_key_id", d
|
||||
# Context variable for MCP pre-authentication flag (set when MCP_AUTH_TOKEN validates)
|
||||
_current_mcp_authenticated: ContextVar[bool] = ContextVar("current_mcp_authenticated", default=False)
|
||||
|
||||
# Context variable for the headers an operator opted into forwarding to extensions
|
||||
# (HINDSIGHT_API_EXTENSION_PASSTHROUGH_HEADERS). Defaults to None rather than {}
|
||||
# so no single dict is shared as a default across requests.
|
||||
_current_extra_headers: ContextVar[dict[str, str] | None] = ContextVar("current_extra_headers", default=None)
|
||||
|
||||
|
||||
def get_current_bank_id() -> str | None:
|
||||
"""Get the current bank_id from context."""
|
||||
@@ -78,6 +84,15 @@ def get_current_mcp_authenticated() -> bool:
|
||||
return _current_mcp_authenticated.get()
|
||||
|
||||
|
||||
def get_current_extra_headers() -> dict[str, str]:
|
||||
"""Get the allowlisted passthrough headers for the current request.
|
||||
|
||||
Returns a copy: every RequestContext built during the request owns its dict,
|
||||
so extension code mutating one cannot alter what the next tool call sees.
|
||||
"""
|
||||
return dict(_current_extra_headers.get() or {})
|
||||
|
||||
|
||||
def _build_mcp_tool_descriptions(extra_instructions: str | None) -> tuple[str | None, str | None]:
|
||||
"""Return custom retain/recall descriptions when server-level MCP instructions are set."""
|
||||
if not isinstance(extra_instructions, str):
|
||||
@@ -159,6 +174,7 @@ def create_mcp_server(memory: MemoryEngine, multi_bank: bool = True) -> FastMCP:
|
||||
tenant_id_resolver=get_current_tenant_id, # Propagate tenant_id for usage metering
|
||||
api_key_id_resolver=get_current_api_key_id, # Propagate api_key_id for usage metering
|
||||
mcp_authenticated_resolver=get_current_mcp_authenticated, # Propagate MCP pre-auth flag
|
||||
extra_headers_resolver=get_current_extra_headers, # Propagate allowlisted headers to extensions
|
||||
include_bank_id_param=multi_bank,
|
||||
tools=base_tools,
|
||||
retain_description=retain_description,
|
||||
@@ -378,6 +394,16 @@ class MCPMiddleware:
|
||||
return header_value.decode()
|
||||
return None
|
||||
|
||||
def _get_extra_headers(self, scope: dict) -> dict[str, str]:
|
||||
"""Collect the headers an operator opted into forwarding to extensions.
|
||||
|
||||
Shares ``collect_passthrough_headers`` with the HTTP transport, so both
|
||||
agree on decoding and on what a duplicated header means. Empty unless
|
||||
HINDSIGHT_API_EXTENSION_PASSTHROUGH_HEADERS names a header the request
|
||||
actually carries.
|
||||
"""
|
||||
return collect_passthrough_headers(scope.get("headers", []), _get_raw_config().extension_passthrough_headers)
|
||||
|
||||
async def __call__(self, scope, receive, send):
|
||||
if scope["type"] != "http":
|
||||
await self.app(scope, receive, send)
|
||||
@@ -412,6 +438,11 @@ class MCPMiddleware:
|
||||
# Support both "Bearer <token>" and direct token
|
||||
auth_token = auth_header[7:].strip() if auth_header.startswith("Bearer ") else auth_header.strip()
|
||||
|
||||
# Resolved before authentication so authenticate_mcp() can read a
|
||||
# passthrough header, not just the bearer token. Named for the request
|
||||
# side: _send_error()'s `extra_headers` below is *response* headers.
|
||||
passthrough_headers = self._get_extra_headers(scope)
|
||||
|
||||
# Authenticate: check legacy MCP_AUTH_TOKEN first, then TenantExtension
|
||||
tenant_context = None
|
||||
auth_tenant_id: str | None = None
|
||||
@@ -431,7 +462,7 @@ class MCPMiddleware:
|
||||
else:
|
||||
# Use TenantExtension.authenticate_mcp() for auth
|
||||
try:
|
||||
auth_context = RequestContext(api_key=auth_token)
|
||||
auth_context = RequestContext(api_key=auth_token, extra_headers=dict(passthrough_headers))
|
||||
tenant_context = await self.tenant_extension.authenticate_mcp(auth_context)
|
||||
# Capture tenant_id and api_key_id set by authenticate() for usage metering
|
||||
auth_tenant_id = auth_context.tenant_id
|
||||
@@ -483,6 +514,8 @@ class MCPMiddleware:
|
||||
api_key_id_token = _current_api_key_id.set(auth_api_key_id) if auth_api_key_id else None
|
||||
# Store MCP pre-authentication flag to skip tenant re-validation
|
||||
mcp_auth_token = _current_mcp_authenticated.set(mcp_pre_authenticated)
|
||||
# Store the allowlisted passthrough headers so per-tool RequestContexts carry them
|
||||
extra_headers_token = _current_extra_headers.set(passthrough_headers)
|
||||
try:
|
||||
new_scope = scope.copy()
|
||||
new_scope["path"] = new_path
|
||||
@@ -528,6 +561,7 @@ class MCPMiddleware:
|
||||
if api_key_id_token is not None:
|
||||
_current_api_key_id.reset(api_key_id_token)
|
||||
_current_mcp_authenticated.reset(mcp_auth_token)
|
||||
_current_extra_headers.reset(extra_headers_token)
|
||||
if schema_token is not None:
|
||||
_current_schema.reset(schema_token)
|
||||
|
||||
|
||||
@@ -0,0 +1,56 @@
|
||||
"""Collection of the request headers an operator forwards to extensions.
|
||||
|
||||
Shared by both transports (the HTTP dependency and the MCP ASGI middleware) so
|
||||
they cannot disagree about which value an extension sees. Both hand over raw
|
||||
ASGI header pairs — Starlette exposes them as ``request.headers.raw``, the MCP
|
||||
middleware reads them straight off the ASGI scope — so one implementation covers
|
||||
decoding, case-folding and duplicate handling for both.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from collections.abc import Iterable, Sequence
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def collect_passthrough_headers(
|
||||
raw_headers: Iterable[tuple[bytes, bytes]],
|
||||
allowlist: Sequence[str],
|
||||
) -> dict[str, str]:
|
||||
"""Pick the allowlisted headers out of a request, keyed by lower-cased name.
|
||||
|
||||
``allowlist`` is ``HindsightConfig.extension_passthrough_headers``, already
|
||||
lower-cased at config load; empty (the default) means nothing is forwarded.
|
||||
|
||||
A header sent more than once is dropped rather than resolved. These headers
|
||||
carry identity for the deployments that enable this, and there is no safe
|
||||
universal rule for picking between copies: a proxy may append its trusted
|
||||
value after a client-supplied one or before it. Dropping turns a duplicate
|
||||
into a loud failure in the extension (which sees no header) instead of a
|
||||
silent choice between a real and a spoofed value.
|
||||
|
||||
Values are decoded as latin-1, matching Starlette and the HTTP/1.1 wire
|
||||
encoding, so a header carrying non-UTF-8 bytes cannot fail the request.
|
||||
"""
|
||||
if not allowlist:
|
||||
return {}
|
||||
|
||||
wanted = set(allowlist)
|
||||
found: dict[str, list[bytes]] = {}
|
||||
for raw_name, raw_value in raw_headers:
|
||||
name = raw_name.decode("latin-1").lower()
|
||||
if name in wanted:
|
||||
found.setdefault(name, []).append(raw_value)
|
||||
|
||||
collected: dict[str, str] = {}
|
||||
for name, values in found.items():
|
||||
if len(values) > 1:
|
||||
logger.warning(
|
||||
"Header '%s' is in HINDSIGHT_API_EXTENSION_PASSTHROUGH_HEADERS but arrived %d times; "
|
||||
"not forwarding it to extensions (no safe way to choose between the copies)",
|
||||
name,
|
||||
len(values),
|
||||
)
|
||||
continue
|
||||
collected[name] = values[0].decode("latin-1")
|
||||
return collected
|
||||
@@ -233,7 +233,7 @@ DEFAULT_LLM_DEFAULT_HEADERS = (
|
||||
)
|
||||
# "auto" is safe as a default because it is an allowlist, not a best-effort probe:
|
||||
# it emits a hint only for hosts documented to accept one (x.ai / grok.com get the
|
||||
# header, native OpenAI / openai.com / Azure OpenAI get the field) and resolves to
|
||||
# header, native OpenAI / openai.com get the field) and resolves to
|
||||
# "none" for every other backend, so vLLM, ollama, groq, openrouter and any custom
|
||||
# OpenAI-compatible endpoint keep receiving byte-identical requests. Measured on a
|
||||
# live xAI backend: 29% of a shared prefix cached without the header vs 99% with it,
|
||||
@@ -377,6 +377,12 @@ ENV_CONSOLIDATION_LLM_EXTRA_BODY = "HINDSIGHT_API_CONSOLIDATION_LLM_EXTRA_BODY"
|
||||
ENV_CONSOLIDATION_LLM_CACHE_AFFINITY = "HINDSIGHT_API_CONSOLIDATION_LLM_CACHE_AFFINITY"
|
||||
|
||||
ENV_EMBEDDINGS_PROVIDER = "HINDSIGHT_API_EMBEDDINGS_PROVIDER"
|
||||
# Provider-agnostic asymmetric prefixes: applied client-side by every provider that
|
||||
# is plain text-in/vector-out (tei, litellm, litellm-sdk, openai-compatible). Providers
|
||||
# with a native asymmetry mechanism (local, zeroentropy) ignore them; onnx has its own
|
||||
# pair below because its defaults are non-empty.
|
||||
ENV_EMBEDDINGS_QUERY_PREFIX = "HINDSIGHT_API_EMBEDDINGS_QUERY_PREFIX"
|
||||
ENV_EMBEDDINGS_PASSAGE_PREFIX = "HINDSIGHT_API_EMBEDDINGS_PASSAGE_PREFIX"
|
||||
ENV_EMBEDDINGS_LOCAL_MODEL = "HINDSIGHT_API_EMBEDDINGS_LOCAL_MODEL"
|
||||
ENV_EMBEDDINGS_LOCAL_FORCE_CPU = "HINDSIGHT_API_EMBEDDINGS_LOCAL_FORCE_CPU"
|
||||
ENV_EMBEDDINGS_LOCAL_ALLOW_MPS = "HINDSIGHT_API_EMBEDDINGS_LOCAL_ALLOW_MPS"
|
||||
@@ -505,6 +511,7 @@ ENV_SEMANTIC_LINK_MIN_SIMILARITY = "HINDSIGHT_API_SEMANTIC_LINK_MIN_SIMILARITY"
|
||||
ENV_RERANKER_FLASHRANK_MODEL = "HINDSIGHT_API_RERANKER_FLASHRANK_MODEL"
|
||||
ENV_RERANKER_FLASHRANK_CACHE_DIR = "HINDSIGHT_API_RERANKER_FLASHRANK_CACHE_DIR"
|
||||
ENV_RERANKER_FLASHRANK_CPU_MEM_ARENA = "HINDSIGHT_API_RERANKER_FLASHRANK_CPU_MEM_ARENA"
|
||||
ENV_RERANKER_FLASHRANK_BATCH_SIZE = "HINDSIGHT_API_RERANKER_FLASHRANK_BATCH_SIZE"
|
||||
|
||||
# ZeroEntropy configuration (reranker only)
|
||||
ENV_RERANKER_ZEROENTROPY_API_KEY = "HINDSIGHT_API_RERANKER_ZEROENTROPY_API_KEY"
|
||||
@@ -557,6 +564,12 @@ ENV_LINK_EXPANSION_PER_ENTITY_LIMIT = "HINDSIGHT_API_LINK_EXPANSION_PER_ENTITY_L
|
||||
ENV_LINK_EXPANSION_TIMEOUT = "HINDSIGHT_API_LINK_EXPANSION_TIMEOUT"
|
||||
ENV_BANK_STATS_CACHE_TTL_SECONDS = "HINDSIGHT_API_BANK_STATS_CACHE_TTL_SECONDS"
|
||||
ENV_BANK_STATS_CACHE_MAX_ENTRIES = "HINDSIGHT_API_BANK_STATS_CACHE_MAX_ENTRIES"
|
||||
# Request headers copied into RequestContext.extra_headers for extensions to read.
|
||||
# Comma-separated, matched case-insensitively. Empty by default: extensions only
|
||||
# ever see headers an operator has explicitly opted in, so a custom
|
||||
# TenantExtension/OperationValidatorExtension can't be handed request data its
|
||||
# author never asked for.
|
||||
ENV_EXTENSION_PASSTHROUGH_HEADERS = "HINDSIGHT_API_EXTENSION_PASSTHROUGH_HEADERS"
|
||||
|
||||
# OpenTelemetry tracing configuration
|
||||
ENV_OTEL_TRACES_ENABLED = "HINDSIGHT_API_OTEL_TRACES_ENABLED"
|
||||
@@ -706,6 +719,7 @@ ENV_DB_COMMAND_TIMEOUT = "HINDSIGHT_API_DB_COMMAND_TIMEOUT"
|
||||
ENV_DB_ACQUIRE_TIMEOUT = "HINDSIGHT_API_DB_ACQUIRE_TIMEOUT"
|
||||
ENV_DB_STATEMENT_TIMEOUT = "HINDSIGHT_API_DB_STATEMENT_TIMEOUT"
|
||||
ENV_DB_MAX_PARALLEL_WORKERS_PER_GATHER = "HINDSIGHT_API_DB_MAX_PARALLEL_WORKERS_PER_GATHER"
|
||||
ENV_DB_SESSION_SETUP_ON_ACQUIRE = "HINDSIGHT_API_DB_SESSION_SETUP_ON_ACQUIRE"
|
||||
ENV_ENTITY_TRGM_SIMILARITY_THRESHOLD = "HINDSIGHT_API_ENTITY_TRGM_SIMILARITY_THRESHOLD"
|
||||
ENV_ENTITY_INTRABATCH_MERGE_SIMILARITY = "HINDSIGHT_API_ENTITY_INTRABATCH_MERGE_SIMILARITY"
|
||||
|
||||
@@ -741,6 +755,7 @@ WORKER_SLOT_TYPE_DEFAULTS: dict[str, int] = {
|
||||
"file_convert_retain": 0,
|
||||
"refresh_mental_model": 0,
|
||||
"graph_maintenance": 0,
|
||||
"vector_index_maintenance": 0,
|
||||
"import_documents": 0,
|
||||
"export_documents": 0,
|
||||
}
|
||||
@@ -792,6 +807,7 @@ ENV_REFLECT_MAX_CONTEXT_TOKENS = "HINDSIGHT_API_REFLECT_MAX_CONTEXT_TOKENS"
|
||||
ENV_REFLECT_WALL_TIMEOUT = "HINDSIGHT_API_REFLECT_WALL_TIMEOUT"
|
||||
ENV_REFLECT_MISSION = "HINDSIGHT_API_REFLECT_MISSION"
|
||||
ENV_REFLECT_SOURCE_FACTS_MAX_TOKENS = "HINDSIGHT_API_REFLECT_SOURCE_FACTS_MAX_TOKENS"
|
||||
ENV_REFLECT_MAX_COMPLETION_TOKENS = "HINDSIGHT_API_REFLECT_MAX_COMPLETION_TOKENS"
|
||||
ENV_RECALL_INCLUDE_CHUNKS = "HINDSIGHT_API_RECALL_INCLUDE_CHUNKS"
|
||||
ENV_RECALL_MAX_TOKENS = "HINDSIGHT_API_RECALL_MAX_TOKENS"
|
||||
ENV_RECALL_CHUNKS_MAX_TOKENS = "HINDSIGHT_API_RECALL_CHUNKS_MAX_TOKENS"
|
||||
@@ -819,6 +835,7 @@ ENV_RECALL_BUDGET_MAX = "HINDSIGHT_API_RECALL_BUDGET_MAX"
|
||||
# Recall candidate gating (per-source cap + BM25 score floor)
|
||||
ENV_BM25_MIN_SCORE = "HINDSIGHT_API_BM25_MIN_SCORE"
|
||||
ENV_BM25_MAX_QUERY_TERMS = "HINDSIGHT_API_BM25_MAX_QUERY_TERMS"
|
||||
ENV_BM25_SELECTIVE_TERMS = "HINDSIGHT_API_BM25_SELECTIVE_TERMS"
|
||||
ENV_RECALL_MAX_CANDIDATES_PER_SOURCE = "HINDSIGHT_API_RECALL_MAX_CANDIDATES_PER_SOURCE"
|
||||
# Per-strategy recall boost. Prioritises specific retrieval arms (semantic,
|
||||
# bm25, graph, temporal) on recall via a human priority level — e.g.
|
||||
@@ -857,6 +874,10 @@ ENV_LLM_TRACE_MAX_CHARS = "HINDSIGHT_API_LLM_TRACE_MAX_CHARS"
|
||||
# Background maintenance settings
|
||||
ENV_CONSOLIDATION_RECONCILE_INTERVAL_SECONDS = "HINDSIGHT_API_CONSOLIDATION_RECONCILE_INTERVAL_SECONDS"
|
||||
ENV_MENTAL_MODEL_REFRESH_TICK_SECONDS = "HINDSIGHT_API_MENTAL_MODEL_REFRESH_TICK_SECONDS"
|
||||
ENV_RETENTION_SWEEP_INTERVAL_SECONDS = "HINDSIGHT_API_RETENTION_SWEEP_INTERVAL_SECONDS"
|
||||
ENV_OPERATION_CLEANUP_INTERVAL_SECONDS = "HINDSIGHT_API_OPERATION_CLEANUP_INTERVAL_SECONDS"
|
||||
ENV_MAINTENANCE_START_JITTER_SECONDS = "HINDSIGHT_API_MAINTENANCE_START_JITTER_SECONDS"
|
||||
ENV_VECTOR_INDEX_MIN_ROWS = "HINDSIGHT_API_VECTOR_INDEX_MIN_ROWS"
|
||||
|
||||
# Disposition settings
|
||||
ENV_DISPOSITION_SKEPTICISM = "HINDSIGHT_API_DISPOSITION_SKEPTICISM"
|
||||
@@ -932,7 +953,6 @@ DEFAULT_LLM_MAX_RETRIES = 3 # Max retry attempts for LLM API calls
|
||||
DEFAULT_LLM_INITIAL_BACKOFF = 1.0 # Initial backoff in seconds for retry exponential backoff
|
||||
DEFAULT_LLM_MAX_BACKOFF = 60.0 # Max backoff cap in seconds for retry exponential backoff
|
||||
DEFAULT_LLM_TIMEOUT = 120.0 # seconds
|
||||
DEFAULT_LLM_REASONING_EFFORT = "low"
|
||||
DEFAULT_LLM_SEND_BANK_AS_USER = False # Opt-in: tag provider calls with user=<bank_id>
|
||||
|
||||
# Vertex AI defaults
|
||||
@@ -955,6 +975,10 @@ DEFAULT_EMBEDDINGS_ONNX_FILE = "onnx/model.onnx"
|
||||
DEFAULT_EMBEDDINGS_ONNX_MAX_TOKENS = 512
|
||||
DEFAULT_EMBEDDINGS_ONNX_POOLING = "mean"
|
||||
DEFAULT_EMBEDDINGS_ONNX_NORMALIZE = True
|
||||
# Empty by default: most hosted embedding models are symmetric, so prefixing is opt-in
|
||||
# for asymmetric models (E5, embeddinggemma, ...) served behind a plain text-in endpoint.
|
||||
DEFAULT_EMBEDDINGS_QUERY_PREFIX = ""
|
||||
DEFAULT_EMBEDDINGS_PASSAGE_PREFIX = ""
|
||||
DEFAULT_EMBEDDINGS_ONNX_QUERY_PREFIX = "query: "
|
||||
DEFAULT_EMBEDDINGS_ONNX_PASSAGE_PREFIX = "passage: "
|
||||
DEFAULT_EMBEDDINGS_OPENAI_MODEL = "text-embedding-3-small"
|
||||
@@ -1004,9 +1028,18 @@ DEFAULT_SEMANTIC_LINK_MIN_SIMILARITY = 0.7
|
||||
# zero-score (non-matching) rows on backends — notably VectorChord — whose
|
||||
# operator ranks every document rather than pre-filtering to term matches.
|
||||
DEFAULT_BM25_MIN_SCORE = 0.0
|
||||
# Native tsvector BM25 can optionally cap the OR tsquery built from normalized
|
||||
# query tokens. 0 preserves the historical uncapped behavior.
|
||||
DEFAULT_BM25_MAX_QUERY_TERMS = 0
|
||||
# Native tsvector BM25 caps the OR tsquery built from normalized query tokens.
|
||||
# Native ranking has no IDF and re-ranks every `@@` match, so an uncapped long
|
||||
# query over common terms scans and ranks a large fraction of the bank and can
|
||||
# time out. When the query has more tokens than this cap, the most selective
|
||||
# terms (lowest tenant-wide document frequency, from pg_stats) are kept and the
|
||||
# rest dropped. 0 restores the historical uncapped behavior.
|
||||
DEFAULT_BM25_MAX_QUERY_TERMS = 16
|
||||
# Whether the cap above selects terms by pg_stats document frequency (keep the
|
||||
# most selective) rather than by position (keep the first N). True is strictly
|
||||
# better for recall at no extra cost when stats exist; set False to opt out of
|
||||
# the catalog read and cap by position instead. Ignored when the cap is 0.
|
||||
DEFAULT_BM25_SELECTIVE_TERMS = True
|
||||
# Per-source candidate cap applied to each retrieval arm (semantic, BM25, graph,
|
||||
# temporal) before RRF, so a single over-expanding backend cannot fill the
|
||||
# reranker's global candidate budget on its own. 0 disables the cap.
|
||||
@@ -1072,6 +1105,11 @@ def _parse_strategy_boosts(raw: str | None) -> dict[str, str]:
|
||||
DEFAULT_RERANKER_FLASHRANK_MODEL = "ms-marco-MiniLM-L-12-v2" # Best balance of speed and quality
|
||||
DEFAULT_RERANKER_FLASHRANK_CACHE_DIR = None # Use default cache directory
|
||||
DEFAULT_RERANKER_FLASHRANK_CPU_MEM_ARENA = False # Disable ONNX CPU memory arena to bound RSS
|
||||
# Passages per FlashRank forward pass. A single pass allocates attention tensors
|
||||
# sized batch * heads * seq^2, so an unbatched rerank of a full candidate pool
|
||||
# costs gigabytes and can OOM the container (issue #3355). Matches the local
|
||||
# reranker's default batch size.
|
||||
DEFAULT_RERANKER_FLASHRANK_BATCH_SIZE = 32
|
||||
|
||||
DEFAULT_EMBEDDINGS_COHERE_MODEL = "embed-english-v3.0"
|
||||
DEFAULT_RERANKER_COHERE_MODEL = "rerank-english-v3.0"
|
||||
@@ -1276,6 +1314,25 @@ DEFAULT_DB_STATEMENT_TIMEOUT = 600 # seconds (Postgres statement_timeout applie
|
||||
# workers buy latency, which background work doesn't need, at the cost of
|
||||
# concurrent CPU footprint, which multi-tenant primaries do care about.
|
||||
DEFAULT_DB_MAX_PARALLEL_WORKERS_PER_GATHER: int | None = None
|
||||
# Whether the per-connection session setup (statement_timeout, hnsw.ef_search,
|
||||
# pg_trgm.similarity_threshold, max_parallel_workers_per_gather, and the vchord
|
||||
# search_path) is re-applied on every pool acquire, not just when a connection is
|
||||
# first opened.
|
||||
#
|
||||
# True (default) is the correct setting for a plain asyncpg pool: releasing a
|
||||
# connection runs RESET ALL, which wipes every SET the init callback applied, so
|
||||
# without the re-apply a reused connection silently runs with server defaults.
|
||||
#
|
||||
# Set False only when those settings are already pinned server-side — ALTER ROLE
|
||||
# / ALTER DATABASE ... SET — because RESET ALL then restores them to exactly the
|
||||
# values we would have re-sent, and the re-apply is a wasted round trip on every
|
||||
# acquire. Behind a transaction-mode pooler that round trip is also its own
|
||||
# server-side transaction, which is what made it visible as commit-rate burn in
|
||||
# #3499. Note that on the vchord text-search backend the set includes
|
||||
# search_path (bm25_catalog, tokenizer_catalog): unlike the tuning GUCs, losing
|
||||
# that one fails recall outright ('type "bm25vector" does not exist') rather
|
||||
# than degrading it, so pin it too before turning this off.
|
||||
DEFAULT_DB_SESSION_SETUP_ON_ACQUIRE = True
|
||||
# pg_trgm similarity threshold applied on every pool connection (SET
|
||||
# pg_trgm.similarity_threshold). Governs how close a name must be for the `%`
|
||||
# operator to treat it as a candidate during entity resolution: lower catches
|
||||
@@ -1323,6 +1380,14 @@ DEFAULT_REFLECT_PROMPT_CACHE_ENABLED = True
|
||||
DEFAULT_REFLECT_MAX_CONTEXT_TOKENS = 100_000 # Max accumulated context tokens before forcing final prompt
|
||||
DEFAULT_REFLECT_WALL_TIMEOUT = 300 # Wall-clock timeout in seconds for the entire reflect operation (5 minutes)
|
||||
DEFAULT_REFLECT_SOURCE_FACTS_MAX_TOKENS = -1 # Token budget for source facts in search_observations (-1 = disabled)
|
||||
# Transport-level output cap (max_completion_tokens) for reflect's final synthesis.
|
||||
# None = uncapped: the model runs to a natural stop and the desired page length is
|
||||
# governed by a prompt directive + the post-hoc rewrite, NOT by truncating the
|
||||
# provider call. This decouples the mental-model/reflect ``max_tokens`` (a page-length
|
||||
# target) from the raw provider budget, which on thinking models is consumed by
|
||||
# reasoning tokens and would otherwise cut pages off mid-word (#3365). Set an integer
|
||||
# only if you want a hard cost ceiling on the synthesis call.
|
||||
DEFAULT_REFLECT_MAX_COMPLETION_TOKENS: int | None = None
|
||||
DEFAULT_RECALL_INCLUDE_CHUNKS = True # Whether internal recall (e.g. mental model refresh) returns raw chunks
|
||||
DEFAULT_RECALL_MAX_TOKENS = 2048 # Token budget for facts returned by internal recall
|
||||
DEFAULT_RECALL_CHUNKS_MAX_TOKENS = 1000 # Token budget for raw chunks returned by internal recall
|
||||
@@ -1396,7 +1461,55 @@ DEFAULT_CONSOLIDATION_RECONCILE_INTERVAL_SECONDS = 300
|
||||
# How often the maintenance loop checks for cron-scheduled mental models that are
|
||||
# due for a refresh. This is the *check* cadence; the actual schedule is the
|
||||
# per-model cron expression in the mental model's trigger. 0 disables the sweep.
|
||||
DEFAULT_MENTAL_MODEL_REFRESH_TICK_SECONDS = 60
|
||||
#
|
||||
# Discovery is one cross-tenant round-trip that probes every schema holding a
|
||||
# mental_models table, so its cost scales with tenant count while the models it
|
||||
# looks for are rare. Five minutes keeps that cost proportionate; the floor it
|
||||
# imposes on cron granularity (a `* * * * *` schedule fires every 5 minutes, not
|
||||
# every minute) is why it stays tunable.
|
||||
DEFAULT_MENTAL_MODEL_REFRESH_TICK_SECONDS = 300
|
||||
|
||||
# How often the audit_log / llm_requests retention sweeps run. Retention windows
|
||||
# are measured in days, so this only sets how promptly expired rows disappear.
|
||||
DEFAULT_RETENTION_SWEEP_INTERVAL_SECONDS = 3600
|
||||
|
||||
# How often terminal async_operations rows past their retention are pruned. One
|
||||
# bounded batch per tenant schema per run, so this also sets the drain rate for a
|
||||
# backlog (batch size: HINDSIGHT_API_OPERATION_CLEANUP_BATCH_SIZE). Like the
|
||||
# retention sweeps it deletes rows whose retention is counted in days, so a slow
|
||||
# cadence costs nothing but avoids a per-tick cross-tenant probe. 0 disables it.
|
||||
DEFAULT_OPERATION_CLEANUP_INTERVAL_SECONDS = 900
|
||||
|
||||
# Upper bound on the random delay applied before a process runs its first
|
||||
# maintenance tick. Every job is due on the first tick, so without this a fleet
|
||||
# started together (deploy, rolling restart) runs every sweep in every process at
|
||||
# the same instant. 0 disables the jitter (deterministic start).
|
||||
DEFAULT_MAINTENANCE_START_JITTER_SECONDS = 60
|
||||
|
||||
# Rows a (bank, fact_type) partition needs before it gets its own partial vector
|
||||
# index. These indexes live on the *shared* memory_units table: PostgreSQL locks
|
||||
# and builds an IndexOptInfo for every index on a relation at plan time, and
|
||||
# opens every one of them for each DML statement, so one bank's index is a cost
|
||||
# paid by every other bank in the deployment. Three per bank exhausts the lock
|
||||
# table at a few thousand banks (issue #3485).
|
||||
#
|
||||
# 0 is the default and means "no minimum": every partition that holds rows gets
|
||||
# an index, which is the behaviour before the threshold existed. Deployments
|
||||
# holding thousands of banks raise it — above the threshold ANN wins, and below
|
||||
# it PostgreSQL answers the same query from the (bank_id, fact_type) B-tree plus
|
||||
# a top-N sort, which is exact rather than approximate *and* faster, because
|
||||
# sorting a few thousand rows by distance costs less than descending an ANN
|
||||
# graph. 10_000 is a reasonable starting point (it is also ScaNN's own build
|
||||
# floor, SCANN_MIN_ROWS_FOR_AUTO_INDEX).
|
||||
DEFAULT_VECTOR_INDEX_MIN_ROWS = 0
|
||||
|
||||
# A partition that falls back below MIN_ROWS * this ratio loses its index. The
|
||||
# gap between the build and drop thresholds is hysteresis: with a single
|
||||
# boundary, consolidation pruning a bank back and forth across it would rebuild
|
||||
# and drop the same ANN index on alternating writes. At the default threshold of
|
||||
# 0 there is no gap and nothing to flap — a partition either holds rows or does
|
||||
# not.
|
||||
VECTOR_INDEX_DROP_RATIO = 0.5
|
||||
|
||||
# Default MCP tool descriptions (can be customized via env vars)
|
||||
DEFAULT_MCP_RETAIN_DESCRIPTION = """Store important information to long-term memory.
|
||||
@@ -1871,6 +1984,7 @@ class RerankerMemberConfig:
|
||||
flashrank_model: str
|
||||
flashrank_cache_dir: str | None
|
||||
flashrank_cpu_mem_arena: bool
|
||||
flashrank_batch_size: int
|
||||
# litellm (proxy)
|
||||
litellm_api_base: str
|
||||
litellm_api_key: str | None
|
||||
@@ -2007,6 +2121,7 @@ def _parse_reranker_members() -> list[RerankerMemberConfig]:
|
||||
flashrank_cpu_mem_arena=_member_bool(
|
||||
base, "FLASHRANK_CPU_MEM_ARENA", DEFAULT_RERANKER_FLASHRANK_CPU_MEM_ARENA
|
||||
),
|
||||
flashrank_batch_size=_member_int(base, "FLASHRANK_BATCH_SIZE", DEFAULT_RERANKER_FLASHRANK_BATCH_SIZE),
|
||||
litellm_api_base=_member_str(base, "LITELLM_API_BASE", DEFAULT_LITELLM_API_BASE),
|
||||
litellm_api_key=_member_opt_str(base, "LITELLM_API_KEY"),
|
||||
litellm_model=_member_str(base, "LITELLM_MODEL", DEFAULT_RERANKER_LITELLM_MODEL),
|
||||
@@ -2105,7 +2220,10 @@ class HindsightConfig:
|
||||
llm_initial_backoff: float
|
||||
llm_max_backoff: float
|
||||
llm_timeout: float
|
||||
llm_reasoning_effort: str
|
||||
# None when unset, and unset means no provider sends a reasoning parameter at all —
|
||||
# each model runs at its own default effort. A configured value is a statement about
|
||||
# the deployment and is sent as given (issue #3449).
|
||||
llm_reasoning_effort: str | None
|
||||
llm_groq_service_tier: str # Groq: "on_demand", "flex", or "auto"
|
||||
llm_openai_service_tier: str | None # OpenAI: None (default) or "flex" (50% cheaper)
|
||||
llm_bedrock_service_tier: str | None # Bedrock: None (default), "flex", "priority", or "reserved"
|
||||
@@ -2496,6 +2614,7 @@ class HindsightConfig:
|
||||
db_acquire_timeout: int
|
||||
db_statement_timeout: int
|
||||
db_max_parallel_workers_per_gather: int | None
|
||||
db_session_setup_on_acquire: bool
|
||||
entity_trgm_similarity_threshold: float
|
||||
entity_intrabatch_merge_similarity: float
|
||||
model_init_timeout: float
|
||||
@@ -2520,6 +2639,7 @@ class HindsightConfig:
|
||||
reflect_max_context_tokens: int
|
||||
reflect_wall_timeout: int
|
||||
reflect_prompt_cache_enabled: bool
|
||||
reflect_max_completion_tokens: int | None
|
||||
|
||||
# OpenTelemetry tracing configuration
|
||||
otel_traces_enabled: bool
|
||||
@@ -2562,6 +2682,9 @@ class HindsightConfig:
|
||||
# How often the maintenance loop checks for cron-scheduled mental models due for
|
||||
# refresh (the per-model schedule lives in the mental model trigger). 0 = disabled.
|
||||
mental_model_refresh_tick_seconds: int
|
||||
# Rows a (bank, fact_type) needs before it gets its own partial vector index.
|
||||
# 0 (default) = no minimum: every partition holding rows is indexed.
|
||||
vector_index_min_rows: int
|
||||
|
||||
# Webhook configuration (static - server-level only, not per-bank)
|
||||
webhook_url: str | None # Global webhook URL (None = disabled)
|
||||
@@ -2573,6 +2696,8 @@ class HindsightConfig:
|
||||
# Keep at the end of the dataclass; Python forbids non-default fields after default fields.
|
||||
embeddings_openai_batch_size: int = DEFAULT_EMBEDDINGS_OPENAI_BATCH_SIZE
|
||||
embeddings_openai_dimensions: int | None = None
|
||||
embeddings_query_prefix: str = DEFAULT_EMBEDDINGS_QUERY_PREFIX
|
||||
embeddings_passage_prefix: str = DEFAULT_EMBEDDINGS_PASSAGE_PREFIX
|
||||
embeddings_zeroentropy_api_key: str | None = None
|
||||
embeddings_zeroentropy_model: str = DEFAULT_EMBEDDINGS_ZEROENTROPY_MODEL
|
||||
embeddings_zeroentropy_base_url: str = DEFAULT_ZEROENTROPY_BASE_URL
|
||||
@@ -2607,6 +2732,7 @@ class HindsightConfig:
|
||||
# embed api_keys/base_urls).
|
||||
reranker_members: list[RerankerMemberConfig] = field(default_factory=list)
|
||||
bm25_max_query_terms: int = DEFAULT_BM25_MAX_QUERY_TERMS
|
||||
bm25_selective_terms: bool = DEFAULT_BM25_SELECTIVE_TERMS
|
||||
|
||||
# Webhook SSRF hardening (static, server-level only — deliberately NOT
|
||||
# per-bank configurable: a tenant must not be able to re-open the private
|
||||
@@ -2614,6 +2740,20 @@ class HindsightConfig:
|
||||
webhook_allowed_hosts: list[str] = field(default_factory=list)
|
||||
webhook_expose_response_body: bool = DEFAULT_WEBHOOK_EXPOSE_RESPONSE_BODY
|
||||
|
||||
# Headers forwarded to extensions via RequestContext.extra_headers (static,
|
||||
# server-level only — deliberately NOT per-bank configurable: a tenant must
|
||||
# not be able to widen the set of request headers its own extension code
|
||||
# sees). Stored lower-cased; empty means no header is ever forwarded.
|
||||
extension_passthrough_headers: list[str] = field(default_factory=list)
|
||||
|
||||
# Background maintenance cadences (static, server-level only). Each sweep's
|
||||
# discovery is one cross-tenant round-trip that probes every schema holding
|
||||
# the relevant table, so its cost scales with tenant count and the cadence is
|
||||
# the lever a large deployment tunes. 0 disables the job.
|
||||
retention_sweep_interval_seconds: int = DEFAULT_RETENTION_SWEEP_INTERVAL_SECONDS
|
||||
operation_cleanup_interval_seconds: int = DEFAULT_OPERATION_CLEANUP_INTERVAL_SECONDS
|
||||
maintenance_start_jitter_seconds: int = DEFAULT_MAINTENANCE_START_JITTER_SECONDS
|
||||
|
||||
# Class-level sets for configuration categorization
|
||||
|
||||
# CREDENTIAL_FIELDS: Never exposed via API, never configurable per-tenant/bank
|
||||
@@ -2777,6 +2917,11 @@ class HindsightConfig:
|
||||
ENV_RERANKER_FLASHRANK_CPU_MEM_ARENA, str(DEFAULT_RERANKER_FLASHRANK_CPU_MEM_ARENA)
|
||||
).lower()
|
||||
in ("true", "1", "yes"),
|
||||
# Tolerate a set-but-empty value the way _member_int does — an unset
|
||||
# `VAR=` in a compose/env file must fall back, not fail config load.
|
||||
flashrank_batch_size=int(
|
||||
os.environ.get(ENV_RERANKER_FLASHRANK_BATCH_SIZE, "").strip() or DEFAULT_RERANKER_FLASHRANK_BATCH_SIZE
|
||||
),
|
||||
litellm_api_base=self.reranker_litellm_api_base,
|
||||
litellm_api_key=self.reranker_litellm_api_key,
|
||||
litellm_model=self.reranker_litellm_model,
|
||||
@@ -3055,7 +3200,7 @@ class HindsightConfig:
|
||||
llm_initial_backoff=float(os.getenv(ENV_LLM_INITIAL_BACKOFF, str(DEFAULT_LLM_INITIAL_BACKOFF))),
|
||||
llm_max_backoff=float(os.getenv(ENV_LLM_MAX_BACKOFF, str(DEFAULT_LLM_MAX_BACKOFF))),
|
||||
llm_timeout=float(os.getenv(ENV_LLM_TIMEOUT, str(DEFAULT_LLM_TIMEOUT))),
|
||||
llm_reasoning_effort=os.getenv(ENV_LLM_REASONING_EFFORT, DEFAULT_LLM_REASONING_EFFORT),
|
||||
llm_reasoning_effort=os.getenv(ENV_LLM_REASONING_EFFORT) or None,
|
||||
llm_groq_service_tier=os.getenv(ENV_LLM_GROQ_SERVICE_TIER, DEFAULT_LLM_GROQ_SERVICE_TIER),
|
||||
llm_openai_service_tier=os.getenv(ENV_LLM_OPENAI_SERVICE_TIER, DEFAULT_LLM_OPENAI_SERVICE_TIER),
|
||||
llm_bedrock_service_tier=os.getenv(ENV_LLM_BEDROCK_SERVICE_TIER) or None,
|
||||
@@ -3279,6 +3424,8 @@ class HindsightConfig:
|
||||
ENV_EMBEDDINGS_OPENAI_DIMENSIONS,
|
||||
os.getenv(ENV_EMBEDDINGS_OPENAI_DIMENSIONS),
|
||||
),
|
||||
embeddings_query_prefix=os.getenv(ENV_EMBEDDINGS_QUERY_PREFIX, DEFAULT_EMBEDDINGS_QUERY_PREFIX),
|
||||
embeddings_passage_prefix=os.getenv(ENV_EMBEDDINGS_PASSAGE_PREFIX, DEFAULT_EMBEDDINGS_PASSAGE_PREFIX),
|
||||
# Cohere embeddings (with backward-compatible fallback to shared API key)
|
||||
embeddings_cohere_api_key=os.getenv(ENV_EMBEDDINGS_COHERE_API_KEY) or os.getenv(ENV_COHERE_API_KEY),
|
||||
embeddings_cohere_model=os.getenv(ENV_EMBEDDINGS_COHERE_MODEL, DEFAULT_EMBEDDINGS_COHERE_MODEL),
|
||||
@@ -3428,6 +3575,7 @@ class HindsightConfig:
|
||||
os.getenv(ENV_BM25_MAX_QUERY_TERMS),
|
||||
DEFAULT_BM25_MAX_QUERY_TERMS,
|
||||
),
|
||||
bm25_selective_terms=_parse_boolean_env(ENV_BM25_SELECTIVE_TERMS, DEFAULT_BM25_SELECTIVE_TERMS),
|
||||
recall_max_candidates_per_source=int(
|
||||
os.getenv(ENV_RECALL_MAX_CANDIDATES_PER_SOURCE, str(DEFAULT_RECALL_MAX_CANDIDATES_PER_SOURCE))
|
||||
),
|
||||
@@ -3735,6 +3883,9 @@ class HindsightConfig:
|
||||
ENV_DB_MAX_PARALLEL_WORKERS_PER_GATHER,
|
||||
os.getenv(ENV_DB_MAX_PARALLEL_WORKERS_PER_GATHER),
|
||||
),
|
||||
db_session_setup_on_acquire=_parse_boolean_env(
|
||||
ENV_DB_SESSION_SETUP_ON_ACQUIRE, DEFAULT_DB_SESSION_SETUP_ON_ACQUIRE
|
||||
),
|
||||
entity_trgm_similarity_threshold=float(
|
||||
os.getenv(ENV_ENTITY_TRGM_SIMILARITY_THRESHOLD, str(DEFAULT_ENTITY_TRGM_SIMILARITY_THRESHOLD))
|
||||
),
|
||||
@@ -3785,6 +3936,11 @@ class HindsightConfig:
|
||||
reflect_source_facts_max_tokens=int(
|
||||
os.getenv(ENV_REFLECT_SOURCE_FACTS_MAX_TOKENS, str(DEFAULT_REFLECT_SOURCE_FACTS_MAX_TOKENS))
|
||||
),
|
||||
reflect_max_completion_tokens=(
|
||||
int(os.getenv(ENV_REFLECT_MAX_COMPLETION_TOKENS))
|
||||
if os.getenv(ENV_REFLECT_MAX_COMPLETION_TOKENS)
|
||||
else DEFAULT_REFLECT_MAX_COMPLETION_TOKENS
|
||||
),
|
||||
enable_temporal_retrieval=os.getenv(
|
||||
ENV_ENABLE_TEMPORAL_RETRIEVAL, str(DEFAULT_ENABLE_TEMPORAL_RETRIEVAL)
|
||||
).lower()
|
||||
@@ -3886,6 +4042,26 @@ class HindsightConfig:
|
||||
str(DEFAULT_MENTAL_MODEL_REFRESH_TICK_SECONDS),
|
||||
)
|
||||
),
|
||||
retention_sweep_interval_seconds=_parse_non_negative_int(
|
||||
ENV_RETENTION_SWEEP_INTERVAL_SECONDS,
|
||||
os.getenv(ENV_RETENTION_SWEEP_INTERVAL_SECONDS),
|
||||
DEFAULT_RETENTION_SWEEP_INTERVAL_SECONDS,
|
||||
),
|
||||
vector_index_min_rows=_parse_non_negative_int(
|
||||
ENV_VECTOR_INDEX_MIN_ROWS,
|
||||
os.getenv(ENV_VECTOR_INDEX_MIN_ROWS),
|
||||
DEFAULT_VECTOR_INDEX_MIN_ROWS,
|
||||
),
|
||||
operation_cleanup_interval_seconds=_parse_non_negative_int(
|
||||
ENV_OPERATION_CLEANUP_INTERVAL_SECONDS,
|
||||
os.getenv(ENV_OPERATION_CLEANUP_INTERVAL_SECONDS),
|
||||
DEFAULT_OPERATION_CLEANUP_INTERVAL_SECONDS,
|
||||
),
|
||||
maintenance_start_jitter_seconds=_parse_non_negative_int(
|
||||
ENV_MAINTENANCE_START_JITTER_SECONDS,
|
||||
os.getenv(ENV_MAINTENANCE_START_JITTER_SECONDS),
|
||||
DEFAULT_MAINTENANCE_START_JITTER_SECONDS,
|
||||
),
|
||||
# Webhook configuration (static, server-level only)
|
||||
webhook_url=os.getenv(ENV_WEBHOOK_URL) or DEFAULT_WEBHOOK_URL,
|
||||
webhook_secret=os.getenv(ENV_WEBHOOK_SECRET) or DEFAULT_WEBHOOK_SECRET,
|
||||
@@ -3904,6 +4080,12 @@ class HindsightConfig:
|
||||
webhook_expose_response_body=_parse_boolean_env(
|
||||
ENV_WEBHOOK_EXPOSE_RESPONSE_BODY, DEFAULT_WEBHOOK_EXPOSE_RESPONSE_BODY
|
||||
),
|
||||
# Lower-cased here so the transports can match incoming header names
|
||||
# case-insensitively (HTTP header names are case-insensitive) without
|
||||
# re-normalising the allowlist on every request.
|
||||
extension_passthrough_headers=[
|
||||
h.lower() for h in _parse_str_list(os.getenv(ENV_EXTENSION_PASSTHROUGH_HEADERS, ""))
|
||||
],
|
||||
)
|
||||
config.validate()
|
||||
return config
|
||||
|
||||
@@ -42,7 +42,13 @@ OPENAI_PROMPT_CACHE_KEY_PARAM = "prompt_cache_key"
|
||||
# Hosts (exact or parent domain) whose backends implement the xAI header.
|
||||
_XAI_DOMAINS = ("x.ai", "grok.com")
|
||||
# Hosts (exact or parent domain) that accept OpenAI's prompt_cache_key field.
|
||||
_OPENAI_DOMAINS = ("openai.com", "openai.azure.com")
|
||||
# Deliberately excludes openai.azure.com: Azure OpenAI itself accepts the field
|
||||
# on GPT deployments, but the same *.openai.azure.com endpoint also fronts
|
||||
# non-OpenAI Foundry models (DeepSeek, Llama, Mistral) that reject it with
|
||||
# `unrecognized_request_argument` (#3518). The host says nothing about which
|
||||
# model family the deployment serves, so `auto` stays off there and an Azure
|
||||
# GPT operator opts in with openai_prompt_cache_key.
|
||||
_OPENAI_DOMAINS = ("openai.com",)
|
||||
|
||||
|
||||
class CacheAffinityMode(StrEnum):
|
||||
@@ -87,7 +93,7 @@ def resolve_cache_affinity(mode: CacheAffinityMode, provider: str, base_url: str
|
||||
|
||||
Non-``auto`` modes are returned unchanged. ``auto`` resolves to
|
||||
``xai_conv_id`` for an x.ai / grok.com host, ``openai_prompt_cache_key`` for
|
||||
native OpenAI (no base URL) or an openai.com / Azure OpenAI host, and
|
||||
native OpenAI (no base URL) or an openai.com host, and
|
||||
``none`` for everything else — an unknown backend gets no unfamiliar field.
|
||||
|
||||
The xAI check is host-only and deliberately provider-independent: the
|
||||
|
||||
@@ -114,19 +114,17 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
|
||||
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
|
||||
month = month_index % 12 + 1
|
||||
day = min(reference_date.day, calendar.monthrange(year, month)[1])
|
||||
return reference_date.replace(year=year, month=month, day=day)
|
||||
def subtract_months(months: int) -> datetime | None:
|
||||
return add_months(reference_date, -months)
|
||||
|
||||
def month_end(year: int, month: int) -> datetime:
|
||||
return datetime(year, month, calendar.monthrange(year, month)[1])
|
||||
|
||||
def add_months(base_date: datetime, months: int) -> datetime:
|
||||
def add_months(base_date: datetime, months: int) -> datetime | None:
|
||||
month_index = base_date.month + months - 1
|
||||
year = base_date.year + month_index // 12
|
||||
if year < datetime.min.year or year > datetime.max.year:
|
||||
return None
|
||||
month = month_index % 12 + 1
|
||||
day = min(base_date.day, calendar.monthrange(year, month)[1])
|
||||
return base_date.replace(year=year, month=month, day=day)
|
||||
@@ -373,7 +371,7 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
|
||||
sat = start + timedelta(days=5)
|
||||
return constraint(sat, sat + timedelta(days=1))
|
||||
|
||||
def relative_month_start(period: str | None) -> datetime:
|
||||
def relative_month_start(period: str | None) -> datetime | None:
|
||||
return add_months(reference_date.replace(day=1), relative_period_offset(period))
|
||||
|
||||
def exact_day_constraint(year: int, month_text: str, day_text: str) -> DateRange | None:
|
||||
@@ -418,6 +416,8 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
|
||||
if day is None:
|
||||
return None
|
||||
start = relative_month_start(period)
|
||||
if start is None:
|
||||
return None
|
||||
if day > calendar.monthrange(start.year, start.month)[1]:
|
||||
return None
|
||||
return datetime(start.year, start.month, day)
|
||||
@@ -472,9 +472,9 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
|
||||
|
||||
def relative_offset_datetime(amount: int, unit: str, direction: int) -> datetime | None:
|
||||
if unit in ("天", "日"):
|
||||
return reference_date + timedelta(days=direction * amount)
|
||||
return add_days(reference_date, direction * amount)
|
||||
if unit in ("周", "星期", "礼拜"):
|
||||
return reference_date + timedelta(weeks=direction * amount)
|
||||
return add_days(reference_date, direction * amount * 7)
|
||||
if unit == "月":
|
||||
return add_months(reference_date, direction * amount)
|
||||
return add_years(reference_date, direction * amount)
|
||||
@@ -616,6 +616,8 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
|
||||
if relative_month_range_match:
|
||||
first = relative_month_start(relative_month_range_match.group(1))
|
||||
second = relative_month_start(relative_month_range_match.group(2))
|
||||
if first is None or second is None:
|
||||
return NO_TEMPORAL_CONSTRAINT
|
||||
start = min(first, second)
|
||||
end = max(first, second)
|
||||
return constraint(start, month_end(end.year, end.month))
|
||||
@@ -839,7 +841,7 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
|
||||
rf"(?<![上下大小])(上上|大上|上|这|本|当|下下|大下|下){_CHINESE_OPTIONAL_PERIOD_MARKER}月{chinese_since_suffix_pattern}"
|
||||
)
|
||||
if month_since_match:
|
||||
return since_constraint(relative_month_start(month_since_match.group(1)))
|
||||
return safe_since_constraint(relative_month_start(month_since_match.group(1)))
|
||||
|
||||
absolute_year_month_since_match = chinese_search(
|
||||
rf"({chinese_year_pattern})\s*年\s*({chinese_month_pattern})\s*月{chinese_since_suffix_pattern}"
|
||||
@@ -1035,7 +1037,7 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
|
||||
|
||||
if chinese_search(r"一年半前"):
|
||||
d = subtract_months(18)
|
||||
return constraint(d, d)
|
||||
return safe_constraint(d, d)
|
||||
|
||||
if chinese_search(r"([一二两三四五六七八九十]+)年半前"):
|
||||
match = chinese_search(r"([一二两三四五六七八九十]+)年半前")
|
||||
@@ -1043,15 +1045,15 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
|
||||
years = parse_chinese_number(match.group(1))
|
||||
if years is not None:
|
||||
d = subtract_months(years * 12 + 6)
|
||||
return constraint(d, d)
|
||||
return safe_constraint(d, d)
|
||||
|
||||
if chinese_search(r"([0-9]+|[一二两三四五六七八九十]+)个?半月前"):
|
||||
match = chinese_search(r"([0-9]+|[一二两三四五六七八九十]+)个?半月前")
|
||||
if match is not None:
|
||||
months = parse_chinese_number(match.group(1))
|
||||
if months is not None:
|
||||
d = subtract_months(months) - timedelta(days=15)
|
||||
return constraint(d, d)
|
||||
d = add_days(subtract_months(months), -15)
|
||||
return safe_constraint(d, d)
|
||||
|
||||
if chinese_search(r"半个?月前"):
|
||||
d = reference_date - timedelta(days=15)
|
||||
@@ -1059,7 +1061,7 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
|
||||
|
||||
if chinese_search(r"半年前"):
|
||||
d = subtract_months(6)
|
||||
return constraint(d, d)
|
||||
return safe_constraint(d, d)
|
||||
|
||||
future_year_half_match = chinese_search(
|
||||
rf"([0-9]+|[{_CHINESE_NUMERAL_CHARS}]+)年半{chinese_relative_future_suffix_pattern}"
|
||||
@@ -1068,7 +1070,7 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
|
||||
years = parse_chinese_number(future_year_half_match.group(1))
|
||||
if years is not None:
|
||||
d = add_months(reference_date, years * 12 + 6)
|
||||
return constraint(d, d)
|
||||
return safe_constraint(d, d)
|
||||
|
||||
future_half_month_match = chinese_search(
|
||||
rf"([0-9]+|[{_CHINESE_NUMERAL_CHARS}]+)个?半月{chinese_relative_future_suffix_pattern}"
|
||||
@@ -1076,8 +1078,8 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
|
||||
if future_half_month_match:
|
||||
months = parse_chinese_number(future_half_month_match.group(1))
|
||||
if months is not None:
|
||||
d = add_months(reference_date, months) + timedelta(days=15)
|
||||
return constraint(d, d)
|
||||
d = add_days(add_months(reference_date, months), 15)
|
||||
return safe_constraint(d, d)
|
||||
|
||||
if chinese_search(rf"半个?月{chinese_relative_future_suffix_pattern}"):
|
||||
d = reference_date + timedelta(days=15)
|
||||
@@ -1085,7 +1087,7 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
|
||||
|
||||
if chinese_search(rf"半年{chinese_relative_future_suffix_pattern}"):
|
||||
d = add_months(reference_date, 6)
|
||||
return constraint(d, d)
|
||||
return safe_constraint(d, d)
|
||||
|
||||
adjacent_fuzzy_future_match = chinese_search(
|
||||
r"(?<![一二三四五六七八九十百千万零\d后])"
|
||||
@@ -1232,14 +1234,14 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
|
||||
unit = rolling_past_half_match.group(2)
|
||||
if unit == "月":
|
||||
return constraint(reference_date - timedelta(days=15), reference_date)
|
||||
return constraint(subtract_months(6), reference_date)
|
||||
return safe_constraint(subtract_months(6), reference_date)
|
||||
|
||||
within_half_match = chinese_search(r"半个?(月|年)(?:以内|之内|内)")
|
||||
if within_half_match:
|
||||
unit = within_half_match.group(1)
|
||||
if unit == "月":
|
||||
return constraint(reference_date - timedelta(days=15), reference_date)
|
||||
return constraint(subtract_months(6), reference_date)
|
||||
return safe_constraint(subtract_months(6), reference_date)
|
||||
|
||||
within_count_match = chinese_search(
|
||||
rf"([0-9]+|[{_CHINESE_NUMERAL_CHARS}]+)(个?)(天|日|周|星期|礼拜|月|年)(?:以内|之内|内)"
|
||||
@@ -1302,7 +1304,7 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
|
||||
unit = rolling_future_half_match.group(2)
|
||||
if unit == "月":
|
||||
return constraint(reference_date, reference_date + timedelta(days=15))
|
||||
return constraint(reference_date, add_months(reference_date, 6))
|
||||
return safe_constraint(reference_date, add_months(reference_date, 6))
|
||||
|
||||
absolute_year_quarter_since_match = chinese_search(
|
||||
rf"({chinese_year_pattern})\s*年\s*(第?[一二三四1-4])季(?:度)?{chinese_since_suffix_pattern}"
|
||||
@@ -1439,6 +1441,8 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
|
||||
)
|
||||
if next_month_phase_since_match:
|
||||
start = add_months(reference_date.replace(day=1), 1)
|
||||
if start is None:
|
||||
return NO_TEMPORAL_CONSTRAINT
|
||||
return since_from_period(month_phase_period(start.year, start.month, next_month_phase_since_match.group(1)))
|
||||
|
||||
second_next_month_phase_since_match = chinese_search(
|
||||
@@ -1447,6 +1451,8 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
|
||||
)
|
||||
if second_next_month_phase_since_match:
|
||||
start = add_months(reference_date.replace(day=1), 2)
|
||||
if start is None:
|
||||
return NO_TEMPORAL_CONSTRAINT
|
||||
return since_from_period(
|
||||
month_phase_period(start.year, start.month, second_next_month_phase_since_match.group(2))
|
||||
)
|
||||
@@ -1465,7 +1471,10 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
|
||||
rf"({chinese_month_phase_pattern}){chinese_since_suffix_pattern}"
|
||||
)
|
||||
if second_previous_month_phase_since_match:
|
||||
start = subtract_months(2).replace(day=1)
|
||||
start = subtract_months(2)
|
||||
if start is None:
|
||||
return NO_TEMPORAL_CONSTRAINT
|
||||
start = start.replace(day=1)
|
||||
return since_from_period(
|
||||
month_phase_period(start.year, start.month, second_previous_month_phase_since_match.group(2))
|
||||
)
|
||||
@@ -1590,6 +1599,8 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
|
||||
)
|
||||
if next_month_phase_match:
|
||||
start = add_months(reference_date.replace(day=1), 1)
|
||||
if start is None:
|
||||
return NO_TEMPORAL_CONSTRAINT
|
||||
return month_phase_period(start.year, start.month, next_month_phase_match.group(1))
|
||||
|
||||
second_next_month_phase_match = chinese_search(
|
||||
@@ -1597,6 +1608,8 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
|
||||
)
|
||||
if second_next_month_phase_match:
|
||||
start = add_months(reference_date.replace(day=1), 2)
|
||||
if start is None:
|
||||
return NO_TEMPORAL_CONSTRAINT
|
||||
return month_phase_period(start.year, start.month, second_next_month_phase_match.group(2))
|
||||
|
||||
previous_month_phase_match = chinese_search(
|
||||
@@ -1611,7 +1624,10 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
|
||||
rf"(?<![上大])(上上|大上){_CHINESE_OPTIONAL_PERIOD_MARKER}月份?\s*({chinese_month_phase_pattern})"
|
||||
)
|
||||
if second_previous_month_phase_match:
|
||||
start = subtract_months(2).replace(day=1)
|
||||
start = subtract_months(2)
|
||||
if start is None:
|
||||
return NO_TEMPORAL_CONSTRAINT
|
||||
start = start.replace(day=1)
|
||||
return month_phase_period(start.year, start.month, second_previous_month_phase_match.group(2))
|
||||
|
||||
bare_specific_month_phase_match = chinese_search(
|
||||
@@ -1730,10 +1746,14 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
|
||||
rf"(?<![下大])(下下|大下){_CHINESE_OPTIONAL_PERIOD_MARKER}月(?!{chinese_month_boundary_suffix_pattern})"
|
||||
):
|
||||
start = add_months(reference_date.replace(day=1), 2)
|
||||
if start is None:
|
||||
return NO_TEMPORAL_CONSTRAINT
|
||||
return constraint(start, month_end(start.year, start.month))
|
||||
|
||||
if chinese_search(rf"(?<![下大])下{_CHINESE_OPTIONAL_PERIOD_MARKER}月(?!{chinese_month_boundary_suffix_pattern})"):
|
||||
start = add_months(reference_date.replace(day=1), 1)
|
||||
if start is None:
|
||||
return NO_TEMPORAL_CONSTRAINT
|
||||
return constraint(start, month_end(start.year, start.month))
|
||||
|
||||
if chinese_search(rf"(下一个年度|下一年度|下年度|下一年|明年)(?!{chinese_boundary_suffix_pattern})"):
|
||||
@@ -1763,7 +1783,10 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
|
||||
if chinese_search(
|
||||
rf"(?<![上大])(上上|大上){_CHINESE_OPTIONAL_PERIOD_MARKER}月(?!{chinese_month_boundary_suffix_pattern})"
|
||||
):
|
||||
start = subtract_months(2).replace(day=1)
|
||||
start = subtract_months(2)
|
||||
if start is None:
|
||||
return NO_TEMPORAL_CONSTRAINT
|
||||
start = start.replace(day=1)
|
||||
return constraint(start, month_end(start.year, start.month))
|
||||
|
||||
if chinese_search(rf"前一个?(周|星期|礼拜)(?!{chinese_boundary_suffix_pattern})"):
|
||||
|
||||
@@ -55,12 +55,51 @@ if TYPE_CHECKING:
|
||||
from asyncpg import Connection
|
||||
|
||||
from ...api.http import RequestContext
|
||||
from ..memories.base import StoredMemory
|
||||
from ..memory_engine import MemoryEngine
|
||||
from ..response_models import MemoryFact, RecallResult
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
async def _gather_or_cancel(coros: list[Any]) -> list[Any]:
|
||||
"""``asyncio.gather`` that leaves no task running behind it.
|
||||
|
||||
Plain ``asyncio.gather`` re-raises the first exception immediately but does
|
||||
NOT cancel its siblings — they keep running detached. In consolidation that
|
||||
is actively harmful: the failure propagates out of ``run_consolidation_job``
|
||||
to the worker, which marks the operation failed and re-queues it with a 5s
|
||||
base backoff, while the orphaned tag groups are still calling the LLM,
|
||||
stamping ``mark_consolidated`` and committing write-groups. The per-scope
|
||||
``scope_locks`` are local to one dispatch, so nothing serialises an orphan
|
||||
against the retry, and the "batches within a group run serially" invariant
|
||||
that keeps two consolidators out of the same observation scope is broken
|
||||
exactly when it matters.
|
||||
|
||||
So: cancel the outstanding tasks and await them before propagating. A
|
||||
cancelled batch's writes stay invisible (its witness row is never
|
||||
committed) and are resolved by the recovery sweep, which is the same state
|
||||
a crash would leave.
|
||||
|
||||
Deliberately not ``asyncio.TaskGroup``: it wraps failures in an
|
||||
``ExceptionGroup``, and the worker's ``_is_non_retryable_task_error`` does
|
||||
``isinstance`` checks on the raised exception — a wrapped
|
||||
``IntegrityConstraintViolationError`` would be misclassified as retryable
|
||||
and retried forever. This helper re-raises the original exception unchanged.
|
||||
"""
|
||||
tasks = [asyncio.ensure_future(c) for c in coros]
|
||||
try:
|
||||
return await asyncio.gather(*tasks)
|
||||
except BaseException:
|
||||
for t in tasks:
|
||||
if not t.done():
|
||||
t.cancel()
|
||||
# Await the cancellations before propagating: returning while they are
|
||||
# still unwinding would reintroduce the very overlap this prevents.
|
||||
await asyncio.gather(*tasks, return_exceptions=True)
|
||||
raise
|
||||
|
||||
|
||||
def _native_search_vector_update(config, param: str) -> str:
|
||||
"""UPDATE-clause fragment that repopulates ``search_vector`` inline, or ''
|
||||
when the backend does not maintain a native tsvector column that way.
|
||||
@@ -183,6 +222,52 @@ def _dedup_active(config: Any) -> bool:
|
||||
return get_config().database_backend != "oracle"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _TemporalBounds:
|
||||
"""The temporal columns an observation inherits from the facts behind it.
|
||||
|
||||
Merging two observations (or an observation and a fresh set of source facts) must widen
|
||||
these, never replace them: ``event_date``/``occurred_start`` keep the earliest known value
|
||||
and ``occurred_end``/``mentioned_at`` the latest, with a missing value on either side
|
||||
ignored. That is exactly the ``_aggregate_source_fields`` rule, and the Python mirror of the
|
||||
``LEAST``/``GREATEST`` the SQL paths apply.
|
||||
|
||||
The SQL spelling differs by reach, deliberately. The dedup folds only ever run on PostgreSQL
|
||||
(``_dedup_active`` disables dedup on Oracle) and use the plain
|
||||
``LEAST(col, COALESCE(x, col))``, which is enough there because PostgreSQL ignores NULL
|
||||
arguments. ``_execute_update_action`` also runs on Oracle, where LEAST/GREATEST return NULL
|
||||
if any argument is NULL, so it wraps the whole expression in one more COALESCE — see the
|
||||
comment there.
|
||||
"""
|
||||
|
||||
event_date: "datetime | None" = None
|
||||
occurred_start: "datetime | None" = None
|
||||
occurred_end: "datetime | None" = None
|
||||
mentioned_at: "datetime | None" = None
|
||||
|
||||
@classmethod
|
||||
def of(cls, row: "StoredMemory | _SourceAggregation") -> "_TemporalBounds":
|
||||
"""The bounds carried by a stored memory or by an aggregation over source facts.
|
||||
|
||||
Deliberately not a recall ``MemoryFact``: that model has no ``event_date`` at all and
|
||||
keeps the rest as ISO strings, so it has to be read field by field where it is used.
|
||||
"""
|
||||
return cls(
|
||||
event_date=row.event_date,
|
||||
occurred_start=row.occurred_start,
|
||||
occurred_end=row.occurred_end,
|
||||
mentioned_at=row.mentioned_at,
|
||||
)
|
||||
|
||||
def merged_with(self, other: "_TemporalBounds") -> "_TemporalBounds":
|
||||
return _TemporalBounds(
|
||||
event_date=_merge_min(self.event_date, other.event_date),
|
||||
occurred_start=_merge_min(self.occurred_start, other.occurred_start),
|
||||
occurred_end=_merge_max(self.occurred_end, other.occurred_end),
|
||||
mentioned_at=_merge_max(self.mentioned_at, other.mentioned_at),
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class _DedupOutcome:
|
||||
"""Result of probing one observation against its in-scope neighbours.
|
||||
@@ -224,7 +309,7 @@ async def _dedup_adjudicate(
|
||||
The embedder and the LLM both run with NO connection held; only the semantic+BM25 probe
|
||||
briefly borrows a short-lived connection.
|
||||
"""
|
||||
from ..search.retrieval import retrieve_semantic_bm25_combined
|
||||
from ..memories import get_memories
|
||||
|
||||
threshold = config.consolidation_dedup_threshold
|
||||
if anchor_emb_str is None:
|
||||
@@ -233,10 +318,19 @@ async def _dedup_adjudicate(
|
||||
return _DedupOutcome(best_id=None, merged_text="", should_merge=False)
|
||||
anchor_emb_str = str(embs[0])
|
||||
tags_match = "all_strict" if tags else "any"
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
grouped = await retrieve_semantic_bm25_combined(
|
||||
conn, anchor_emb_str, anchor_text, bank_id, ["observation"], _DEDUP_TOP_K, tags=tags, tags_match=tags_match
|
||||
)
|
||||
# Dedup only needs the dense/keyword arms over observations — no graph, no temporal window.
|
||||
grouped = await get_memories().recall_unified(
|
||||
conn=pool,
|
||||
bank_id=bank_id,
|
||||
fact_types=["observation"],
|
||||
query_embedding=anchor_emb_str,
|
||||
query_text=anchor_text,
|
||||
limit=_DEDUP_TOP_K,
|
||||
tags=tags,
|
||||
tags_match=tags_match,
|
||||
enable_graph=False,
|
||||
temporal_window=None,
|
||||
)
|
||||
results = grouped["observation"].semantic
|
||||
best_id: str | None = None
|
||||
best_text = ""
|
||||
@@ -275,6 +369,7 @@ async def _dedup_reconcile_create(
|
||||
create_text: str,
|
||||
create_source_ids: list[uuid.UUID],
|
||||
tags: list[str] | None,
|
||||
source_bounds: _TemporalBounds,
|
||||
txn=None,
|
||||
) -> str | None:
|
||||
"""Semantic dedup for a single CREATE (create-time, focused 1-by-1).
|
||||
@@ -283,6 +378,10 @@ async def _dedup_reconcile_create(
|
||||
observation and returns its id (caller skips the CREATE). Returns None when there is
|
||||
no near twin or the LLM keeps them distinct.
|
||||
|
||||
``source_bounds`` are the dates the skipped CREATE would have been stamped with. They are
|
||||
folded into the twin too: this path bypasses the CREATE writer, so without them the twin
|
||||
would cite dated source facts while reporting the dates of its original sources only (#3477).
|
||||
|
||||
The probe/embed/LLM adjudication runs with no connection held; the fold takes a
|
||||
short-lived connection and re-checks source liveness inside the fold transaction.
|
||||
"""
|
||||
@@ -315,6 +414,10 @@ async def _dedup_reconcile_create(
|
||||
SET text = $1,
|
||||
source_memory_ids = (SELECT array_agg(DISTINCT e) FROM unnest(source_memory_ids || $2::uuid[]) e),
|
||||
proof_count = (SELECT count(DISTINCT e) FROM unnest(source_memory_ids || $2::uuid[]) e),
|
||||
event_date = LEAST(event_date, COALESCE($5, event_date)),
|
||||
occurred_start = LEAST(occurred_start, COALESCE($6, occurred_start)),
|
||||
occurred_end = GREATEST(occurred_end, COALESCE($7, occurred_end)),
|
||||
mentioned_at = GREATEST(mentioned_at, COALESCE($8, mentioned_at)),
|
||||
updated_at = now(){search_vector_clause}
|
||||
WHERE id = $3::uuid AND text = $4
|
||||
RETURNING id
|
||||
@@ -323,6 +426,10 @@ async def _dedup_reconcile_create(
|
||||
live_source_ids,
|
||||
uuid.UUID(outcome.best_id),
|
||||
outcome.best_text,
|
||||
source_bounds.event_date,
|
||||
source_bounds.occurred_start,
|
||||
source_bounds.occurred_end,
|
||||
source_bounds.mentioned_at,
|
||||
)
|
||||
if folded is None:
|
||||
# The twin vanished (or was rewritten) during the connection-free LLM window.
|
||||
@@ -335,7 +442,15 @@ async def _dedup_reconcile_create(
|
||||
return None
|
||||
else:
|
||||
await _reconcile_merge_via_store(
|
||||
store, conn, memory_engine, bank_id, outcome.best_id, outcome.merged_text, live_source_ids, txn=txn
|
||||
store,
|
||||
conn,
|
||||
memory_engine,
|
||||
bank_id,
|
||||
outcome.best_id,
|
||||
outcome.merged_text,
|
||||
live_source_ids,
|
||||
source_bounds,
|
||||
txn=txn,
|
||||
)
|
||||
return outcome.best_id
|
||||
|
||||
@@ -379,7 +494,8 @@ async def _dedup_reconcile_update(
|
||||
# Fold the updated observation's live sources into the twin (keeping the twin's embedding, as
|
||||
# in the create path) then delete the now-redundant updated row. The all_strict/any tag match
|
||||
# guarantees twin and updated share scope, so dropping the updated row's tags loses no
|
||||
# visibility. Temporal fields follow the surviving twin (minimal scope; matches create).
|
||||
# visibility. Temporal fields are the UNION of both rows' bounds: the updated row is about to
|
||||
# be deleted, so anything only it knew about would otherwise be lost with it (#3477).
|
||||
# The fold + delete share one short transaction so the twin gains the sources exactly as the
|
||||
# redundant row is removed; the slow adjudication above already ran connection-free.
|
||||
store = get_memories()
|
||||
@@ -422,6 +538,10 @@ async def _dedup_reconcile_update(
|
||||
proof_count = (
|
||||
SELECT count(DISTINCT e) FROM unnest(t.source_memory_ids || $6::uuid[]) e
|
||||
),
|
||||
event_date = LEAST(t.event_date, COALESCE(u.event_date, t.event_date)),
|
||||
occurred_start = LEAST(t.occurred_start, COALESCE(u.occurred_start, t.occurred_start)),
|
||||
occurred_end = GREATEST(t.occurred_end, COALESCE(u.occurred_end, t.occurred_end)),
|
||||
mentioned_at = GREATEST(t.mentioned_at, COALESCE(u.mentioned_at, t.mentioned_at)),
|
||||
updated_at = now(){search_vector_clause}
|
||||
FROM {fq_table("memory_units")} u
|
||||
WHERE t.id = $2::uuid AND u.id = $3::uuid AND t.text = $4 AND u.text = $5
|
||||
@@ -447,7 +567,15 @@ async def _dedup_reconcile_update(
|
||||
if not live_u_sources:
|
||||
return
|
||||
await _reconcile_merge_via_store(
|
||||
store, conn, memory_engine, bank_id, outcome.best_id, outcome.merged_text, live_u_sources, txn=txn
|
||||
store,
|
||||
conn,
|
||||
memory_engine,
|
||||
bank_id,
|
||||
outcome.best_id,
|
||||
outcome.merged_text,
|
||||
live_u_sources,
|
||||
_TemporalBounds.of(updated_obs[0]),
|
||||
txn=txn,
|
||||
)
|
||||
await _execute_delete_action(conn, bank_id, updated_id, txn=txn)
|
||||
logger.info(
|
||||
@@ -954,17 +1082,22 @@ async def _reconcile_merge_via_store(
|
||||
observation_id: str,
|
||||
merged_text: str,
|
||||
add_source_ids: list,
|
||||
add_bounds: _TemporalBounds,
|
||||
txn=None,
|
||||
) -> None:
|
||||
"""Dedup merge for a store that owns its rows: fold the extra source facts and the merged text
|
||||
into the twin observation and re-upsert it, preserving its other fields. Re-embeds the merged
|
||||
text because ``get_memories`` does not return the stored vector (the SQL path reuses it in
|
||||
place instead)."""
|
||||
place instead).
|
||||
|
||||
``add_bounds`` are the folded-in side's dates, widened onto the twin exactly as the SQL
|
||||
path's LEAST/GREATEST does."""
|
||||
current = await store.get_memories(conn=conn, fq_table=fq_table, bank_id=bank_id, unit_ids=[observation_id])
|
||||
cur = current[0] if current else None
|
||||
if cur is None:
|
||||
return
|
||||
merged_sources = list(dict.fromkeys([*(cur.source_memory_ids or []), *(str(s) for s in add_source_ids)]))
|
||||
merged_bounds = _TemporalBounds.of(cur).merged_with(add_bounds)
|
||||
embeddings = await embedding_utils.generate_embeddings_batch(memory_engine.embeddings, [merged_text])
|
||||
await store.upsert_observation(
|
||||
conn=conn,
|
||||
@@ -978,10 +1111,10 @@ async def _reconcile_merge_via_store(
|
||||
tags=list(cur.tags or []),
|
||||
proof_count=len(merged_sources),
|
||||
source_memory_ids=merged_sources,
|
||||
event_date=cur.event_date,
|
||||
occurred_start=cur.occurred_start,
|
||||
occurred_end=cur.occurred_end,
|
||||
mentioned_at=cur.mentioned_at,
|
||||
event_date=merged_bounds.event_date,
|
||||
occurred_start=merged_bounds.occurred_start,
|
||||
occurred_end=merged_bounds.occurred_end,
|
||||
mentioned_at=merged_bounds.mentioned_at,
|
||||
created_at=cur.created_at,
|
||||
),
|
||||
)
|
||||
@@ -1050,12 +1183,65 @@ async def _count_unconsolidated_rows(
|
||||
)
|
||||
|
||||
|
||||
def _as_op_uuid(operation_id: str | uuid.UUID) -> uuid.UUID:
|
||||
return uuid.UUID(operation_id) if isinstance(operation_id, str) else operation_id
|
||||
|
||||
|
||||
async def _persist_pending_refresh_tags(conn, operation_id: str, new_tags: list[str]) -> None:
|
||||
"""Union ``new_tags`` into the consolidation op's durable ``pending_refresh_tags``.
|
||||
|
||||
Called inside each batch's witness transaction, so the tags of an
|
||||
already-consolidated batch are durable the instant that batch is — a mid-round
|
||||
worker crash no longer loses them. On retry the op re-reads ``task_payload`` and the
|
||||
final round still refreshes those models (#3411); without this, a crash after batch 1
|
||||
committed but before the round finished would drop batch 1's tags, because the retry
|
||||
skips its now-consolidated rows and never re-collects them. ``SELECT ... FOR UPDATE``
|
||||
serialises the concurrent batches of one op so their unions don't clobber each other.
|
||||
"""
|
||||
op_uuid = _as_op_uuid(operation_id)
|
||||
row = await conn.fetchrow(
|
||||
f"SELECT task_payload FROM {fq_table('async_operations')} WHERE operation_id = $1 FOR UPDATE",
|
||||
op_uuid,
|
||||
)
|
||||
if row is None:
|
||||
return
|
||||
payload = row["task_payload"]
|
||||
payload = json.loads(payload) if isinstance(payload, str) else (payload or {})
|
||||
existing = set(payload.get("pending_refresh_tags") or [])
|
||||
merged = existing | set(new_tags)
|
||||
if merged == existing:
|
||||
return
|
||||
payload["pending_refresh_tags"] = sorted(merged)
|
||||
await conn.execute(
|
||||
f"UPDATE {fq_table('async_operations')} SET task_payload = $1::jsonb, updated_at = now() "
|
||||
f"WHERE operation_id = $2",
|
||||
json.dumps(payload),
|
||||
op_uuid,
|
||||
)
|
||||
|
||||
|
||||
async def _read_pending_refresh_tags(pool, operation_id: str) -> set[str]:
|
||||
"""Read the op's durably-accumulated ``pending_refresh_tags`` (crash-safe source of
|
||||
truth for the final-round flush)."""
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
row = await conn.fetchrow(
|
||||
f"SELECT task_payload FROM {fq_table('async_operations')} WHERE operation_id = $1",
|
||||
_as_op_uuid(operation_id),
|
||||
)
|
||||
if row is None:
|
||||
return set()
|
||||
payload = row["task_payload"]
|
||||
payload = json.loads(payload) if isinstance(payload, str) else (payload or {})
|
||||
return set(payload.get("pending_refresh_tags") or [])
|
||||
|
||||
|
||||
async def run_consolidation_job(
|
||||
memory_engine: "MemoryEngine",
|
||||
bank_id: str,
|
||||
request_context: "RequestContext",
|
||||
operation_id: str | None = None,
|
||||
observation_scopes: list[list[str]] | None = None,
|
||||
pending_refresh_tags: list[str] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Run consolidation job for a bank.
|
||||
@@ -1070,6 +1256,9 @@ async def run_consolidation_job(
|
||||
observation_scopes: Optional list of tag scopes. When provided, only
|
||||
unconsolidated memories whose tags contain all tags in at least one
|
||||
scope are processed.
|
||||
pending_refresh_tags: Tags of memories consolidated by earlier rounds of this
|
||||
round-limited chain, carried through the re-queue so the final round can
|
||||
refresh every affected mental model exactly once (#3411).
|
||||
|
||||
Returns:
|
||||
Dict with consolidation results
|
||||
@@ -1089,7 +1278,14 @@ async def run_consolidation_job(
|
||||
trace_token = set_trace_context(trace_ctx) if trace_ctx is not None else None
|
||||
try:
|
||||
return await _run_consolidation_job(
|
||||
memory_engine, bank_id, request_context, config, llm_config, operation_id, observation_scopes
|
||||
memory_engine,
|
||||
bank_id,
|
||||
request_context,
|
||||
config,
|
||||
llm_config,
|
||||
operation_id,
|
||||
observation_scopes,
|
||||
pending_refresh_tags,
|
||||
)
|
||||
finally:
|
||||
if trace_token is not None:
|
||||
@@ -1107,6 +1303,7 @@ async def _run_consolidation_job(
|
||||
llm_config: Any,
|
||||
operation_id: str | None = None,
|
||||
observation_scopes: list[list[str]] | None = None,
|
||||
pending_refresh_tags: list[str] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Core consolidation flow. See ``run_consolidation_job`` for the public entrypoint."""
|
||||
perf = ConsolidationPerfLog(bank_id)
|
||||
@@ -1287,21 +1484,60 @@ async def _run_consolidation_job(
|
||||
_txn_provider = get_memories()
|
||||
_batch_txn = await _txn_provider.mint_txn(bank_id=bank_id, mutating=True)
|
||||
|
||||
pending: list[list[dict[str, Any]]] = [llm_batch_local]
|
||||
while pending:
|
||||
sub_batch = pending.pop(0)
|
||||
try:
|
||||
pending: list[list[dict[str, Any]]] = [llm_batch_local]
|
||||
while pending:
|
||||
sub_batch = pending.pop(0)
|
||||
|
||||
# No connection is held across the batch: recall, the main LLM call, the
|
||||
# per-action embeds, and dedup all run connection-free; each helper acquires a
|
||||
# short-lived connection only around its own SQL.
|
||||
obs_tags_list = _resolve_obs_tags_list(sub_batch[0]) if sub_batch else None
|
||||
# No connection is held across the batch: recall, the main LLM call, the
|
||||
# per-action embeds, and dedup all run connection-free; each helper acquires a
|
||||
# short-lived connection only around its own SQL.
|
||||
obs_tags_list = _resolve_obs_tags_list(sub_batch[0]) if sub_batch else None
|
||||
|
||||
sub_deleted: int = 0
|
||||
sub_llm_failed = False
|
||||
if obs_tags_list:
|
||||
sub_results: list[dict[str, Any]] = []
|
||||
for obs_tags in obs_tags_list:
|
||||
pass_results, pass_deleted, pass_failed = await _process_memory_batch(
|
||||
sub_deleted: int = 0
|
||||
sub_llm_failed = False
|
||||
if obs_tags_list:
|
||||
sub_results: list[dict[str, Any]] = []
|
||||
for obs_tags in obs_tags_list:
|
||||
pass_results, pass_deleted, pass_failed = await _process_memory_batch(
|
||||
pool=pool,
|
||||
memory_engine=memory_engine,
|
||||
llm_config=llm_config,
|
||||
bank_id=bank_id,
|
||||
memories=sub_batch,
|
||||
request_context=request_context,
|
||||
perf=batch_perf,
|
||||
config=config,
|
||||
obs_tags_override=obs_tags,
|
||||
txn=_batch_txn,
|
||||
)
|
||||
sub_deleted += pass_deleted
|
||||
sub_llm_failed = sub_llm_failed or pass_failed
|
||||
if not sub_results:
|
||||
sub_results = pass_results
|
||||
else:
|
||||
for i, (existing, new) in enumerate(zip(sub_results, pass_results)):
|
||||
if existing.get("action") == "skipped" and new.get("action") != "skipped":
|
||||
sub_results[i] = new
|
||||
elif existing.get("action") != "skipped" and new.get("action") != "skipped":
|
||||
existing_created = existing.get(
|
||||
"created", 1 if existing.get("action") == "created" else 0
|
||||
)
|
||||
existing_updated = existing.get(
|
||||
"updated", 1 if existing.get("action") == "updated" else 0
|
||||
)
|
||||
new_created = new.get("created", 1 if new.get("action") == "created" else 0)
|
||||
new_updated = new.get("updated", 1 if new.get("action") == "updated" else 0)
|
||||
total = existing_created + existing_updated + new_created + new_updated
|
||||
sub_results[i] = {
|
||||
"action": "multiple",
|
||||
"created": existing_created + new_created,
|
||||
"updated": existing_updated + new_updated,
|
||||
"merged": 0,
|
||||
"total_actions": total,
|
||||
}
|
||||
else:
|
||||
sub_results, sub_deleted, sub_llm_failed = await _process_memory_batch(
|
||||
pool=pool,
|
||||
memory_engine=memory_engine,
|
||||
llm_config=llm_config,
|
||||
@@ -1310,96 +1546,90 @@ async def _run_consolidation_job(
|
||||
request_context=request_context,
|
||||
perf=batch_perf,
|
||||
config=config,
|
||||
obs_tags_override=obs_tags,
|
||||
txn=_batch_txn,
|
||||
)
|
||||
sub_deleted += pass_deleted
|
||||
sub_llm_failed = sub_llm_failed or pass_failed
|
||||
if not sub_results:
|
||||
sub_results = pass_results
|
||||
else:
|
||||
for i, (existing, new) in enumerate(zip(sub_results, pass_results)):
|
||||
if existing.get("action") == "skipped" and new.get("action") != "skipped":
|
||||
sub_results[i] = new
|
||||
elif existing.get("action") != "skipped" and new.get("action") != "skipped":
|
||||
existing_created = existing.get(
|
||||
"created", 1 if existing.get("action") == "created" else 0
|
||||
)
|
||||
existing_updated = existing.get(
|
||||
"updated", 1 if existing.get("action") == "updated" else 0
|
||||
)
|
||||
new_created = new.get("created", 1 if new.get("action") == "created" else 0)
|
||||
new_updated = new.get("updated", 1 if new.get("action") == "updated" else 0)
|
||||
total = existing_created + existing_updated + new_created + new_updated
|
||||
sub_results[i] = {
|
||||
"action": "multiple",
|
||||
"created": existing_created + new_created,
|
||||
"updated": existing_updated + new_updated,
|
||||
"merged": 0,
|
||||
"total_actions": total,
|
||||
}
|
||||
else:
|
||||
sub_results, sub_deleted, sub_llm_failed = await _process_memory_batch(
|
||||
pool=pool,
|
||||
memory_engine=memory_engine,
|
||||
llm_config=llm_config,
|
||||
bank_id=bank_id,
|
||||
memories=sub_batch,
|
||||
request_context=request_context,
|
||||
perf=batch_perf,
|
||||
config=config,
|
||||
txn=_batch_txn,
|
||||
)
|
||||
|
||||
all_deleted += sub_deleted
|
||||
all_deleted += sub_deleted
|
||||
|
||||
if sub_llm_failed and len(sub_batch) > 1:
|
||||
mid = len(sub_batch) // 2
|
||||
if sub_llm_failed and len(sub_batch) > 1:
|
||||
mid = len(sub_batch) // 2
|
||||
logger.warning(
|
||||
f"[CONSOLIDATION] bank={bank_id} LLM failed for sub-batch of {len(sub_batch)},"
|
||||
f" splitting into {mid}/{len(sub_batch) - mid}"
|
||||
)
|
||||
pending[0:0] = [sub_batch[:mid], sub_batch[mid:]]
|
||||
elif sub_llm_failed:
|
||||
failed_ids.append(sub_batch[0]["id"])
|
||||
all_results.append({"action": "failed"})
|
||||
logger.warning(
|
||||
f"[CONSOLIDATION] bank={bank_id} LLM failed for single memory"
|
||||
f" {sub_batch[0]['id']}, marking consolidation_failed_at"
|
||||
)
|
||||
else:
|
||||
succeeded_ids.extend(m["id"] for m in sub_batch)
|
||||
all_results.extend(sub_results)
|
||||
|
||||
# Mark through the store so the flag lands wherever the source facts live — tagged
|
||||
# with this batch's txn, so the marks become visible together with the observations
|
||||
# above. Then record the witness row and commit in this ONE short transaction (no LLM
|
||||
# work inside it): its commit is the batch's fate, and `decide` publishes the group.
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
store = get_memories()
|
||||
now = datetime.now(timezone.utc)
|
||||
if succeeded_ids:
|
||||
await store.mark_consolidated(
|
||||
conn=conn,
|
||||
fq_table=fq_table,
|
||||
bank_id=bank_id,
|
||||
unit_ids=[str(mem_id) for mem_id in succeeded_ids],
|
||||
when=now,
|
||||
failed=False,
|
||||
txn=_batch_txn,
|
||||
)
|
||||
if failed_ids:
|
||||
await store.mark_consolidated(
|
||||
conn=conn,
|
||||
fq_table=fq_table,
|
||||
bank_id=bank_id,
|
||||
unit_ids=[str(mem_id) for mem_id in failed_ids],
|
||||
when=now,
|
||||
failed=True,
|
||||
txn=_batch_txn,
|
||||
)
|
||||
async with conn.transaction():
|
||||
await _txn_provider.write_txn_witness(_batch_txn, conn=conn, fq_table=fq_table)
|
||||
# Persist this batch's mental-model refresh tags atomically with the
|
||||
# witness, so they share the batch's fate: durable iff the batch is
|
||||
# (#3411). Only the succeeded source facts — the ones just marked
|
||||
# consolidated — contribute a tag.
|
||||
if operation_id and succeeded_ids:
|
||||
succeeded_set = {str(mem_id) for mem_id in succeeded_ids}
|
||||
batch_tags = sorted(
|
||||
{
|
||||
t
|
||||
for m in llm_batch_local
|
||||
if str(m["id"]) in succeeded_set
|
||||
for t in (m.get("tags") or [])
|
||||
}
|
||||
)
|
||||
if batch_tags:
|
||||
await _persist_pending_refresh_tags(conn, operation_id, batch_tags)
|
||||
except BaseException:
|
||||
# The witness row was never committed, so this batch's writes are invisible;
|
||||
# discard the write-group rather than leaving it pending for the recovery
|
||||
# sweep. This matters more now that a sibling group's failure cancels this
|
||||
# task mid-batch instead of letting it run to completion. Kept OUTSIDE the
|
||||
# decide(commit=True) below on purpose: once the witness has committed, the
|
||||
# batch's fate is decided and an abort here would discard durable writes.
|
||||
try:
|
||||
await _txn_provider.decide_txn(_batch_txn, commit=False)
|
||||
except Exception:
|
||||
logger.warning(
|
||||
f"[CONSOLIDATION] bank={bank_id} LLM failed for sub-batch of {len(sub_batch)},"
|
||||
f" splitting into {mid}/{len(sub_batch) - mid}"
|
||||
f"[CONSOLIDATION] bank={bank_id} failed to abort write-group for"
|
||||
f" llm_batch #{batch_num_local}; recovery sweep will resolve it",
|
||||
exc_info=True,
|
||||
)
|
||||
pending[0:0] = [sub_batch[:mid], sub_batch[mid:]]
|
||||
elif sub_llm_failed:
|
||||
failed_ids.append(sub_batch[0]["id"])
|
||||
all_results.append({"action": "failed"})
|
||||
logger.warning(
|
||||
f"[CONSOLIDATION] bank={bank_id} LLM failed for single memory"
|
||||
f" {sub_batch[0]['id']}, marking consolidation_failed_at"
|
||||
)
|
||||
else:
|
||||
succeeded_ids.extend(m["id"] for m in sub_batch)
|
||||
all_results.extend(sub_results)
|
||||
|
||||
# Mark through the store so the flag lands wherever the source facts live — tagged
|
||||
# with this batch's txn, so the marks become visible together with the observations
|
||||
# above. Then record the witness row and commit in this ONE short transaction (no LLM
|
||||
# work inside it): its commit is the batch's fate, and `decide` publishes the group.
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
store = get_memories()
|
||||
now = datetime.now(timezone.utc)
|
||||
if succeeded_ids:
|
||||
await store.mark_consolidated(
|
||||
conn=conn,
|
||||
fq_table=fq_table,
|
||||
bank_id=bank_id,
|
||||
unit_ids=[str(mem_id) for mem_id in succeeded_ids],
|
||||
when=now,
|
||||
failed=False,
|
||||
txn=_batch_txn,
|
||||
)
|
||||
if failed_ids:
|
||||
await store.mark_consolidated(
|
||||
conn=conn,
|
||||
fq_table=fq_table,
|
||||
bank_id=bank_id,
|
||||
unit_ids=[str(mem_id) for mem_id in failed_ids],
|
||||
when=now,
|
||||
failed=True,
|
||||
txn=_batch_txn,
|
||||
)
|
||||
async with conn.transaction():
|
||||
await _txn_provider.write_txn_witness(_batch_txn, conn=conn, fq_table=fq_table)
|
||||
raise
|
||||
# Postgres committed the witness: publish the batch's write-group. On a crash before
|
||||
# here the writes stay invisible and the recovery sweep resolves them (spec §5).
|
||||
await _txn_provider.decide_txn(_batch_txn, commit=True)
|
||||
@@ -1561,7 +1791,7 @@ async def _run_consolidation_job(
|
||||
await stack.enter_async_context(scope_locks[s])
|
||||
return await _process_tag_group(group_batches)
|
||||
|
||||
group_results = await asyncio.gather(*(_run_group(g, s) for g, s in zip(numbered_groups, group_scopes)))
|
||||
group_results = await _gather_or_cancel([_run_group(g, s) for g, s in zip(numbered_groups, group_scopes)])
|
||||
batch_results: list[_BatchDeltas] = [d for gd in group_results for d in gd]
|
||||
any_cancelled = any(d.cancelled for d in batch_results)
|
||||
else:
|
||||
@@ -1598,6 +1828,19 @@ async def _run_consolidation_job(
|
||||
# execute_task's retry handler means the op is retried with backoff; on retry the
|
||||
# consolidator skips already-consolidated rows via the consolidated_at filter and
|
||||
# picks up the remainder. Issue #1842.
|
||||
# The affected-tag union for the whole round-limited chain. Refresh fires once, when
|
||||
# the backlog has fully drained (the final round), not once per round — a model's
|
||||
# memories can straddle rounds, and gating on the final round alone (the prior
|
||||
# behaviour) dropped every model consolidated earlier because the final round's tags
|
||||
# no longer named them (#3411). The union is durable: each batch writes its tags into
|
||||
# the op's ``task_payload`` inside the batch's own witness txn (crash-safe), and the
|
||||
# re-queue threads the accumulated set forward to the next round. Prefer that durable
|
||||
# value; fall back to the in-memory union when there is no backing op (a direct
|
||||
# ``run_consolidation_job`` call, e.g. in tests).
|
||||
all_refresh_tags = set(pending_refresh_tags or []) | consolidated_tags
|
||||
if operation_id:
|
||||
all_refresh_tags |= await _read_pending_refresh_tags(pool, operation_id)
|
||||
|
||||
if hit_round_limit:
|
||||
remaining = total_count - stats["memories_processed"]
|
||||
logger.info(
|
||||
@@ -1608,6 +1851,7 @@ async def _run_consolidation_job(
|
||||
bank_id=bank_id,
|
||||
request_context=request_context,
|
||||
observation_scopes=observation_scopes,
|
||||
pending_refresh_tags=sorted(all_refresh_tags) or None,
|
||||
)
|
||||
|
||||
# Build summary
|
||||
@@ -1644,11 +1888,19 @@ async def _run_consolidation_job(
|
||||
if timing_parts:
|
||||
perf.log(f"[4] Timing breakdown: {', '.join(timing_parts)}")
|
||||
|
||||
# Trigger mental model refreshes only on the final round (when all memories are processed).
|
||||
# If we hit the round limit and re-queued, skip MM refresh — the next round will handle it.
|
||||
# Trigger mental-model refreshes once, when the chain has fully drained. On a
|
||||
# round-limited round we skip and carry the affected tags forward (above); the
|
||||
# final round flushes the accumulated union, so a model whose memories were
|
||||
# consolidated in ANY round is refreshed exactly once — deduplicated, not dropped
|
||||
# (#3411). Each model is still refreshed at most once per drain: a strict tagged
|
||||
# model appears once in the trigger's candidate query regardless of how many rounds
|
||||
# its tag spanned.
|
||||
if hit_round_limit:
|
||||
stats["mental_models_refreshed"] = 0
|
||||
logger.info(f"[CONSOLIDATION] bank={bank_id} skipping mental model refresh (round limit hit, re-queued)")
|
||||
logger.info(
|
||||
f"[CONSOLIDATION] bank={bank_id} deferring mental model refresh to the final round "
|
||||
f"(round limit hit; carrying {len(all_refresh_tags)} tags forward)"
|
||||
)
|
||||
else:
|
||||
set_stage("consolidation.refreshing_mental_models")
|
||||
await memory_engine._write_operation_progress(
|
||||
@@ -1662,7 +1914,7 @@ async def _run_consolidation_job(
|
||||
memory_engine=memory_engine,
|
||||
bank_id=bank_id,
|
||||
request_context=request_context,
|
||||
consolidated_tags=list(consolidated_tags) if consolidated_tags else None,
|
||||
consolidated_tags=sorted(all_refresh_tags) or None,
|
||||
perf=perf,
|
||||
)
|
||||
stats["mental_models_refreshed"] = mental_models_refreshed
|
||||
@@ -1726,7 +1978,7 @@ async def _trigger_mental_model_refreshes(
|
||||
if consolidated_tags:
|
||||
candidates = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, name, tags, last_refreshed_at, trigger
|
||||
SELECT id, name, tags, last_refreshed_at, last_memory_seen_at, trigger
|
||||
FROM {fq_table("mental_models")}
|
||||
WHERE bank_id = $1
|
||||
AND (trigger->>'refresh_after_consolidation')::boolean = true
|
||||
@@ -1741,7 +1993,7 @@ async def _trigger_mental_model_refreshes(
|
||||
else:
|
||||
candidates = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, name, tags, last_refreshed_at, trigger
|
||||
SELECT id, name, tags, last_refreshed_at, last_memory_seen_at, trigger
|
||||
FROM {fq_table("mental_models")}
|
||||
WHERE bank_id = $1
|
||||
AND (trigger->>'refresh_after_consolidation')::boolean = true
|
||||
@@ -1772,10 +2024,14 @@ async def _trigger_mental_model_refreshes(
|
||||
for row in rows:
|
||||
mental_model_id = row["id"]
|
||||
try:
|
||||
# skip_if_in_flight: a consolidation chain fires this every round and
|
||||
# overlapping consolidations can run on the same bank, so a model still
|
||||
# pending/processing a refresh must not be enqueued a second time (#3411).
|
||||
await memory_engine.submit_async_refresh_mental_model(
|
||||
bank_id=bank_id,
|
||||
mental_model_id=mental_model_id,
|
||||
request_context=request_context,
|
||||
skip_if_in_flight=True,
|
||||
)
|
||||
refreshed_count += 1
|
||||
logger.info(
|
||||
@@ -1820,8 +2076,6 @@ async def _process_memory_batch(
|
||||
consolidation where a single memory can contribute to observations
|
||||
scoped at different tag levels (e.g., user-level vs session-level).
|
||||
"""
|
||||
import asyncio
|
||||
|
||||
# Map the source memories this batch consumes onto the consolidation trace.
|
||||
record_source_memory_ids([str(m["id"]) for m in memories])
|
||||
|
||||
@@ -1839,7 +2093,11 @@ async def _process_memory_batch(
|
||||
)
|
||||
for m in memories
|
||||
]
|
||||
per_fact_recalls = await asyncio.gather(*recall_tasks)
|
||||
# A failed recall must fail the batch rather than degrade to "no related
|
||||
# observations": proceeding with an empty candidate set would hide an
|
||||
# existing twin from the LLM and turn an UPDATE into a duplicate CREATE.
|
||||
# The batch's memories stay unconsolidated and are picked up on retry.
|
||||
per_fact_recalls = await _gather_or_cancel(recall_tasks)
|
||||
if perf:
|
||||
perf.record_timing("recall", time.time() - t0)
|
||||
|
||||
@@ -1962,9 +2220,7 @@ async def _process_memory_batch(
|
||||
new_text=update.text,
|
||||
observations=union_observations,
|
||||
source_fact_tags=agg.tags,
|
||||
source_occurred_start=agg.occurred_start,
|
||||
source_occurred_end=agg.occurred_end,
|
||||
source_mentioned_at=agg.mentioned_at,
|
||||
source_bounds=_TemporalBounds.of(agg),
|
||||
perf=perf,
|
||||
txn=txn,
|
||||
)
|
||||
@@ -2032,6 +2288,7 @@ async def _process_memory_batch(
|
||||
create.text,
|
||||
create_source_ids,
|
||||
agg.tags,
|
||||
_TemporalBounds.of(agg),
|
||||
txn=txn,
|
||||
)
|
||||
if merged_into is not None:
|
||||
@@ -2166,17 +2423,15 @@ async def _execute_update_action(
|
||||
new_text: str,
|
||||
observations: list["MemoryFact"],
|
||||
source_fact_tags: list[str] | None = None,
|
||||
source_occurred_start: datetime | None = None,
|
||||
source_occurred_end: datetime | None = None,
|
||||
source_mentioned_at: datetime | None = None,
|
||||
source_bounds: _TemporalBounds = _TemporalBounds(),
|
||||
perf: ConsolidationPerfLog | None = None,
|
||||
txn=None,
|
||||
) -> str | None:
|
||||
"""
|
||||
Update an existing observation.
|
||||
|
||||
Extends source_memory_ids with all contributing memories, updates temporal fields
|
||||
(LEAST for occurred_start, GREATEST for occurred_end / mentioned_at), and merges tags.
|
||||
Extends source_memory_ids with all contributing memories, widens the observation's temporal
|
||||
bounds by ``source_bounds`` (see :class:`_TemporalBounds`), and merges tags.
|
||||
|
||||
The embedding is computed off-connection (a slow embedder must never pin a pooled
|
||||
connection); the liveness check + UPDATE + history + observation_sources sync then run
|
||||
@@ -2246,6 +2501,15 @@ async def _execute_update_action(
|
||||
|
||||
t0 = time.time()
|
||||
if store.writes_memory_rows_in_sql_for(bank_id):
|
||||
# Unlike the dedup folds this statement also runs on Oracle, where LEAST/GREATEST
|
||||
# return NULL as soon as ANY argument is NULL (PostgreSQL ignores NULL arguments).
|
||||
# The inner COALESCE covers a NULL *parameter*; the outer one covers a NULL
|
||||
# *column* — an observation with no occurred interval yet, which is precisely the
|
||||
# #3477 case. Without it Oracle would compute LEAST(NULL, <source date>) = NULL and
|
||||
# silently drop the date it was told to inherit. Keep the inner
|
||||
# ``COALESCE($n, col)`` spelled exactly like this: the Oracle driver shim keys its
|
||||
# TIMESTAMP-TZ input-size hint off that pattern (db/oracle.py::_apply_clob_input_sizes),
|
||||
# and a NULL parameter binds as VARCHAR2 (ORA-00932) without it.
|
||||
updated_rows = await conn.execute_rows_affected(
|
||||
f"""
|
||||
UPDATE {fq_table("memory_units")}
|
||||
@@ -2253,11 +2517,12 @@ async def _execute_update_action(
|
||||
embedding = $2::vector,
|
||||
source_memory_ids = $3,
|
||||
proof_count = $4,
|
||||
tags = $9,
|
||||
tags = $10,
|
||||
updated_at = now(),
|
||||
occurred_start = LEAST(occurred_start, COALESCE($6, occurred_start)),
|
||||
occurred_end = GREATEST(occurred_end, COALESCE($7, occurred_end)),
|
||||
mentioned_at = GREATEST(mentioned_at, COALESCE($8, mentioned_at)){search_vector_clause}
|
||||
event_date = COALESCE(LEAST(event_date, COALESCE($6, event_date)), $6),
|
||||
occurred_start = COALESCE(LEAST(occurred_start, COALESCE($7, occurred_start)), $7),
|
||||
occurred_end = COALESCE(GREATEST(occurred_end, COALESCE($8, occurred_end)), $8),
|
||||
mentioned_at = COALESCE(GREATEST(mentioned_at, COALESCE($9, mentioned_at)), $9){search_vector_clause}
|
||||
WHERE id = $5
|
||||
""",
|
||||
new_text,
|
||||
@@ -2265,9 +2530,10 @@ async def _execute_update_action(
|
||||
source_ids,
|
||||
len(source_ids),
|
||||
uuid.UUID(observation_id),
|
||||
source_occurred_start,
|
||||
source_occurred_end,
|
||||
source_mentioned_at,
|
||||
source_bounds.event_date,
|
||||
source_bounds.occurred_start,
|
||||
source_bounds.occurred_end,
|
||||
source_bounds.mentioned_at,
|
||||
merged_tags,
|
||||
)
|
||||
# The source-liveness checks above guard the *source* memories; the
|
||||
@@ -2285,12 +2551,24 @@ async def _execute_update_action(
|
||||
return None
|
||||
else:
|
||||
# Upsert overwrites the whole observation, so start from its current state (fetched
|
||||
# from the store) and apply the same merge the SQL does — LEAST/GREATEST on the times
|
||||
# — while preserving fields the update never touches (event_date, created_at).
|
||||
# from the store) and apply the same merge the SQL does — LEAST/GREATEST on the
|
||||
# times — while preserving fields the update never touches (created_at).
|
||||
current = await store.get_memories(
|
||||
conn=conn, fq_table=fq_table, bank_id=bank_id, unit_ids=[observation_id]
|
||||
)
|
||||
cur = current[0] if current else None
|
||||
# Widen the row the store still holds. If it has vanished, fall back to the
|
||||
# pre-update recall snapshot — ISO strings, and no event_date on that model.
|
||||
current_bounds = (
|
||||
_TemporalBounds.of(cur)
|
||||
if cur
|
||||
else _TemporalBounds(
|
||||
occurred_start=_as_dt(model.occurred_start),
|
||||
occurred_end=_as_dt(model.occurred_end),
|
||||
mentioned_at=_as_dt(model.mentioned_at),
|
||||
)
|
||||
)
|
||||
merged_bounds = current_bounds.merged_with(source_bounds)
|
||||
await store.upsert_observation(
|
||||
conn=conn,
|
||||
bank_id=bank_id,
|
||||
@@ -2303,10 +2581,10 @@ async def _execute_update_action(
|
||||
tags=merged_tags,
|
||||
proof_count=len(source_ids),
|
||||
source_memory_ids=[str(s) for s in source_ids],
|
||||
event_date=cur.event_date if cur else None,
|
||||
occurred_start=_merge_min(model.occurred_start, source_occurred_start),
|
||||
occurred_end=_merge_max(model.occurred_end, source_occurred_end),
|
||||
mentioned_at=_merge_max(model.mentioned_at, source_mentioned_at),
|
||||
event_date=merged_bounds.event_date,
|
||||
occurred_start=merged_bounds.occurred_start,
|
||||
occurred_end=merged_bounds.occurred_end,
|
||||
mentioned_at=merged_bounds.mentioned_at,
|
||||
created_at=cur.created_at if cur else None,
|
||||
),
|
||||
)
|
||||
|
||||
@@ -19,6 +19,7 @@ from ..config import (
|
||||
DEFAULT_LITELLM_API_BASE,
|
||||
DEFAULT_RERANKER_ALIBABA_MODEL,
|
||||
DEFAULT_RERANKER_COHERE_MODEL,
|
||||
DEFAULT_RERANKER_FLASHRANK_BATCH_SIZE,
|
||||
DEFAULT_RERANKER_FLASHRANK_CACHE_DIR,
|
||||
DEFAULT_RERANKER_FLASHRANK_MODEL,
|
||||
DEFAULT_RERANKER_GOOGLE_MODEL,
|
||||
@@ -872,6 +873,7 @@ class FlashRankCrossEncoder(CrossEncoderModel):
|
||||
max_length: int = 512,
|
||||
max_concurrent: int = 4,
|
||||
cpu_mem_arena: bool = False,
|
||||
batch_size: int = DEFAULT_RERANKER_FLASHRANK_BATCH_SIZE,
|
||||
):
|
||||
"""
|
||||
Initialize FlashRank cross-encoder.
|
||||
@@ -885,11 +887,16 @@ class FlashRankCrossEncoder(CrossEncoderModel):
|
||||
When True, ONNX pre-allocates a memory arena that never
|
||||
shrinks, causing RSS to grow monotonically. False trades
|
||||
slightly slower per-call allocation for bounded RSS.
|
||||
batch_size: Passages per forward pass. Default: 32. See
|
||||
``_predict_sync`` for why this must stay bounded.
|
||||
"""
|
||||
self.model_name = model_name or DEFAULT_RERANKER_FLASHRANK_MODEL
|
||||
self.cache_dir = cache_dir or DEFAULT_RERANKER_FLASHRANK_CACHE_DIR
|
||||
self.max_length = max_length
|
||||
self.cpu_mem_arena = cpu_mem_arena
|
||||
# A non-positive size would mean "one pass for everything", which is the
|
||||
# unbounded behaviour this batching exists to prevent.
|
||||
self.batch_size = max(1, batch_size)
|
||||
self._ranker = None
|
||||
self._device_type: str = "cpu" # FlashRank runs on CPU via ONNX Runtime
|
||||
FlashRankCrossEncoder._max_concurrent = max_concurrent
|
||||
@@ -962,7 +969,21 @@ class FlashRankCrossEncoder(CrossEncoderModel):
|
||||
logger.info("Reranker: FlashRank provider initialized (using existing executor)")
|
||||
|
||||
def _predict_sync(self, pairs: list[tuple[str, str]]) -> list[float]:
|
||||
"""Synchronous predict - processes each query group."""
|
||||
"""Synchronous predict — each query group, in bounded batches.
|
||||
|
||||
FlashRank scores every passage of a request in one ONNX forward pass, and
|
||||
that pass allocates attention tensors sized ``batch * heads * seq^2``. At
|
||||
the default reranker candidate cap that is gigabytes per call, which OOM-
|
||||
killed containers on large banks (issue #3355): the burst scales with the
|
||||
candidate pool the retrieval arms produce, not with how much work the
|
||||
caller asked for. FlashRank also pads a request to its longest passage, so
|
||||
one long candidate inflates the sequence length for every other one.
|
||||
|
||||
Splitting into ``batch_size`` chunks bounds the peak the same way the
|
||||
local and TEI providers already do. Scores are identical either way —
|
||||
passages are scored independently, so batching changes only the
|
||||
allocation profile.
|
||||
"""
|
||||
if not pairs:
|
||||
return []
|
||||
|
||||
@@ -979,20 +1000,25 @@ class FlashRankCrossEncoder(CrossEncoderModel):
|
||||
all_scores = [0.0] * len(pairs)
|
||||
|
||||
for query, indexed_texts in query_groups.items():
|
||||
# Build passages list for FlashRank
|
||||
passages = [{"id": i, "text": text} for i, (_, text) in enumerate(indexed_texts)]
|
||||
global_indices = [idx for idx, _ in indexed_texts]
|
||||
|
||||
# Create rerank request
|
||||
request = RerankRequest(query=query, passages=passages)
|
||||
results = self._ranker.rerank(request)
|
||||
for start in range(0, len(indexed_texts), self.batch_size):
|
||||
batch = indexed_texts[start : start + self.batch_size]
|
||||
|
||||
# Map scores back to original positions
|
||||
for result in results:
|
||||
local_idx = result["id"]
|
||||
score = result["score"]
|
||||
global_idx = global_indices[local_idx]
|
||||
all_scores[global_idx] = score
|
||||
# Build passages list for FlashRank. Ids are batch-local, so
|
||||
# `start` shifts them back onto the query group's indices.
|
||||
passages = [{"id": i, "text": text} for i, (_, text) in enumerate(batch)]
|
||||
|
||||
# Create rerank request
|
||||
request = RerankRequest(query=query, passages=passages)
|
||||
results = self._ranker.rerank(request)
|
||||
|
||||
# Map scores back to original positions
|
||||
for result in results:
|
||||
local_idx = result["id"]
|
||||
score = result["score"]
|
||||
global_idx = global_indices[start + local_idx]
|
||||
all_scores[global_idx] = score
|
||||
|
||||
return all_scores
|
||||
finally:
|
||||
@@ -1750,6 +1776,7 @@ def create_cross_encoder(member: RerankerMemberConfig) -> CrossEncoderModel:
|
||||
model_name=member.flashrank_model,
|
||||
cache_dir=member.flashrank_cache_dir,
|
||||
cpu_mem_arena=member.flashrank_cpu_mem_arena,
|
||||
batch_size=member.flashrank_batch_size,
|
||||
)
|
||||
elif provider == "litellm":
|
||||
return LiteLLMCrossEncoder(
|
||||
|
||||
@@ -25,12 +25,65 @@ from .base import DatabaseConnection
|
||||
from .result import ResultRow
|
||||
|
||||
|
||||
def document_serialization_sql(table: str, alias: str) -> str:
|
||||
"""SQL predicate keeping one document to a single in-flight retain.
|
||||
|
||||
A retain that targets exactly one document carries it in
|
||||
``serialization_key``. Appending to a document is a read-modify-write over
|
||||
its whole text, so two concurrent retains for one document can only produce
|
||||
a lost update or a wasted extraction — never more throughput. This
|
||||
predicate makes the queue reflect that: a candidate is claimable only when
|
||||
no peer for the same document is already ``processing``, and only when it
|
||||
is the oldest claimable pending peer for that document.
|
||||
|
||||
Ordering, not just exclusion, is the point. Appends are cumulative, so the
|
||||
order they commit in is the order the document ends up in; claiming them by
|
||||
``(created_at, operation_id)`` makes that the submission order. It also
|
||||
stops a single claim batch from taking several peers at once, which
|
||||
excluding busy documents alone would not prevent.
|
||||
|
||||
Rows with a NULL ``serialization_key`` — multi-document batches, and every
|
||||
non-retain operation — are unaffected, and documents are independent of one
|
||||
another, so this costs no parallelism across a busy bank: only the retains
|
||||
that were racing each other for one document are put in a line.
|
||||
|
||||
A peer wedged in 'processing' holds its document until claim recovery
|
||||
releases it, the same caveat ``graph_maintenance_bank_serialization_sql``
|
||||
carries and the same general gap.
|
||||
|
||||
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}.serialization_key IS NULL OR NOT EXISTS (
|
||||
SELECT 1 FROM {table} doc_peer
|
||||
WHERE doc_peer.bank_id = {alias}.bank_id
|
||||
AND doc_peer.serialization_key = {alias}.serialization_key
|
||||
AND (
|
||||
doc_peer.status = 'processing'
|
||||
OR (doc_peer.status = 'pending'
|
||||
AND doc_peer.task_payload IS NOT NULL
|
||||
AND (doc_peer.next_retry_at IS NULL OR doc_peer.next_retry_at <= NOW())
|
||||
AND (doc_peer.created_at < {alias}.created_at
|
||||
OR (doc_peer.created_at = {alias}.created_at
|
||||
AND doc_peer.operation_id < {alias}.operation_id)))
|
||||
)
|
||||
))
|
||||
"""
|
||||
|
||||
|
||||
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
|
||||
Every graph_maintenance run for a bank is interchangeable — the payload
|
||||
carries only ``bank_id``, and ``run_graph_maintenance_job`` drains that bank's
|
||||
queues — 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.
|
||||
@@ -428,23 +481,11 @@ class DataAccessOps(ABC):
|
||||
|
||||
# -- Bank index management -------------------------------------------
|
||||
|
||||
@abstractmethod
|
||||
async def create_bank_vector_indexes(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
bank_id: str,
|
||||
internal_id: str,
|
||||
index_clause: str,
|
||||
fact_types: dict[str, str],
|
||||
) -> None:
|
||||
"""Create per-bank partial vector indexes.
|
||||
|
||||
PG creates per-(bank, fact_type) partial indexes.
|
||||
Non-PG is a no-op (uses global index).
|
||||
"""
|
||||
...
|
||||
|
||||
# No create counterpart: per-(bank, fact_type) partial vector indexes are
|
||||
# earned by size and built by the maintenance sweep over its own autocommit
|
||||
# connection (see engine/vector_index_health.py), never on a request path.
|
||||
# The drop stays here because bank deletion must remove a large bank's
|
||||
# indexes while it still knows the internal_id they are named after.
|
||||
@abstractmethod
|
||||
async def drop_bank_vector_indexes(
|
||||
self,
|
||||
@@ -593,6 +634,43 @@ class DataAccessOps(ABC):
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def enqueue_entity_maintenance(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
ue_table: str,
|
||||
bank_id: str,
|
||||
unit_ids: list,
|
||||
) -> int:
|
||||
"""Enqueue the entities referenced by ``unit_ids`` as prune candidates.
|
||||
|
||||
Reads the entity ids out of ``unit_entities`` and inserts them into
|
||||
entity_maintenance_queue, deduplicating on the (bank_id, entity_id)
|
||||
primary key. Returns the number of rows the insert added.
|
||||
|
||||
Must run inside the triggering transaction and BEFORE the rows go —
|
||||
once the unit_entities rows are deleted (or cascaded away) there is
|
||||
nothing left to read the entity ids from.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def claim_entity_maintenance_batch(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
bank_id: str,
|
||||
limit: int,
|
||||
) -> list:
|
||||
"""Atomically claim a batch of rows from entity_maintenance_queue and
|
||||
remove them from the table.
|
||||
|
||||
Returns the claimed entity ids. Empty list when the queue for
|
||||
``bank_id`` is drained.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def prune_orphan_entities(
|
||||
self,
|
||||
@@ -600,9 +678,10 @@ class DataAccessOps(ABC):
|
||||
entities_table: str,
|
||||
ue_table: str,
|
||||
bank_id: str,
|
||||
entity_ids: list,
|
||||
) -> int:
|
||||
"""Delete entities in ``bank_id`` that no longer have any unit_entities
|
||||
rows referencing them. Returns the number of rows deleted.
|
||||
"""Delete those of ``entity_ids`` in ``bank_id`` that no longer have any
|
||||
unit_entities rows referencing them. Returns the number of rows deleted.
|
||||
|
||||
FK ON DELETE CASCADE on entity_cooccurrences then removes any
|
||||
cooccurrence row pointing at the pruned entities.
|
||||
@@ -615,11 +694,10 @@ class DataAccessOps(ABC):
|
||||
conn: DatabaseConnection,
|
||||
ec_table: str,
|
||||
ue_table: str,
|
||||
entities_table: str,
|
||||
bank_id: str,
|
||||
entity_ids: list,
|
||||
) -> int:
|
||||
"""Delete entity_cooccurrences rows in ``bank_id`` where the two
|
||||
entities still exist but no current unit references both of them.
|
||||
"""Delete entity_cooccurrences rows incident to ``entity_ids`` where the
|
||||
two entities still exist but no current unit references both of them.
|
||||
|
||||
These are stale-count rows: cooccurrence was real at the time it was
|
||||
recorded, but every memory_unit that witnessed both entities has
|
||||
@@ -665,7 +743,9 @@ class DataAccessOps(ABC):
|
||||
|
||||
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.
|
||||
such row per bank is ever in flight, and :func:`document_serialization_sql`
|
||||
to every query that can return a ``retain`` row, so at most one retain per
|
||||
document is ever in flight.
|
||||
|
||||
Args:
|
||||
consolidation_bank_priority: Per-bank priority for consolidation scheduling.
|
||||
@@ -675,8 +755,48 @@ class DataAccessOps(ABC):
|
||||
When set, consolidation tasks are claimed in priority tiers.
|
||||
None preserves current behavior (pure created_at ordering).
|
||||
|
||||
Returns claimed rows with operation_id, operation_type, task_payload, retry_count.
|
||||
The caller is responsible for building ClaimedTask objects.
|
||||
Returns claimed rows with operation_id, operation_type, task_payload,
|
||||
retry_count, bank_id and serialization_key. The caller is responsible for
|
||||
building ClaimedTask objects.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def fetch_foldable_retain_peers(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
bank_id: str,
|
||||
serialization_key: str,
|
||||
limit: int,
|
||||
) -> list[ResultRow]:
|
||||
"""Lock the pending retains queued behind a just-claimed one, in order.
|
||||
|
||||
Called inside the claim transaction, so the rows come back locked and
|
||||
the caller can fold some of them into the claimed execution and leave
|
||||
the rest pending simply by not marking them (their locks release with
|
||||
the transaction).
|
||||
|
||||
``SKIP LOCKED`` matters here for liveness, not just speed: a peer some
|
||||
other worker is already looking at must never stall this claim.
|
||||
|
||||
Returns rows with operation_id, task_payload and retry_count, ordered by
|
||||
``(created_at, operation_id)`` — the order the fold planner requires.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def mark_operations_processing(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
worker_id: str,
|
||||
operation_ids: list,
|
||||
) -> None:
|
||||
"""Claim the given pending operations for ``worker_id``.
|
||||
|
||||
Used to fold peers into an execution that has already been claimed;
|
||||
runs in the same transaction that locked them.
|
||||
"""
|
||||
...
|
||||
|
||||
|
||||
@@ -15,6 +15,7 @@ from .ops import (
|
||||
LinkExpansionRows,
|
||||
TagListingParts,
|
||||
UpdatedWindow,
|
||||
document_serialization_sql,
|
||||
graph_maintenance_bank_serialization_sql,
|
||||
)
|
||||
from .result import DictResultRow as ResultRow
|
||||
@@ -357,13 +358,80 @@ class OracleOps(DataAccessOps):
|
||||
)
|
||||
return claimed
|
||||
|
||||
async def enqueue_entity_maintenance(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
ue_table: str,
|
||||
bank_id: str,
|
||||
unit_ids: list,
|
||||
) -> int:
|
||||
if not unit_ids:
|
||||
return 0
|
||||
rows = await conn.fetch(
|
||||
f"SELECT DISTINCT entity_id FROM {ue_table} WHERE unit_id = ANY($1::uuid[])",
|
||||
unit_ids,
|
||||
)
|
||||
# Sorted for the same reason as enqueue_graph_maintenance: the MERGE
|
||||
# takes the (bank_id, entity_id) row locks in executemany array order,
|
||||
# and claim_entity_maintenance_batch deletes in that same order, so
|
||||
# overlapping mutation/worker sets cannot cycle.
|
||||
candidates = sorted(str(row["entity_id"]) for row in rows)
|
||||
if not candidates:
|
||||
return 0
|
||||
# MERGE is the Oracle analogue of ON CONFLICT DO UPDATE: WHEN MATCHED
|
||||
# locks the existing queue row (the SET is a no-op preserving
|
||||
# enqueued_at) so a re-enqueue serialises against a concurrent claim
|
||||
# instead of being silently dropped (#3034).
|
||||
await conn.executemany(
|
||||
f"""
|
||||
MERGE INTO {table} q
|
||||
USING (SELECT $1 AS bank_id, $2 AS entity_id FROM dual) s
|
||||
ON (q.bank_id = s.bank_id AND q.entity_id = s.entity_id)
|
||||
WHEN MATCHED THEN UPDATE SET q.enqueued_at = q.enqueued_at
|
||||
WHEN NOT MATCHED THEN INSERT (bank_id, entity_id) VALUES (s.bank_id, s.entity_id)
|
||||
""",
|
||||
[(bank_id, eid) for eid in candidates],
|
||||
)
|
||||
return len(candidates)
|
||||
|
||||
async def claim_entity_maintenance_batch(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
bank_id: str,
|
||||
limit: int,
|
||||
) -> list:
|
||||
# Two-step claim, same as claim_graph_maintenance_batch: Oracle's
|
||||
# DELETE ... RETURNING doesn't accept a multi-row subquery.
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT entity_id FROM {table}
|
||||
WHERE bank_id = $1
|
||||
ORDER BY enqueued_at
|
||||
FETCH FIRST $2 ROWS ONLY
|
||||
""",
|
||||
bank_id,
|
||||
limit,
|
||||
)
|
||||
claimed = sorted(str(row["entity_id"]) for row in rows)
|
||||
if claimed:
|
||||
await conn.executemany(
|
||||
f"DELETE FROM {table} WHERE bank_id = $1 AND entity_id = $2",
|
||||
[(bank_id, eid) for eid in claimed],
|
||||
)
|
||||
return claimed
|
||||
|
||||
async def prune_orphan_entities(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
entities_table: str,
|
||||
ue_table: str,
|
||||
bank_id: str,
|
||||
entity_ids: list,
|
||||
) -> int:
|
||||
if not entity_ids:
|
||||
return 0
|
||||
# The Oracle DatabaseConnection wrapper reshapes ``cursor.rowcount`` into
|
||||
# the same ``"DELETE N"`` status string asyncpg returns, so the same
|
||||
# ``int(deleted.split()[-1])`` parsing works on both dialects.
|
||||
@@ -371,9 +439,11 @@ class OracleOps(DataAccessOps):
|
||||
f"""
|
||||
DELETE FROM {entities_table}
|
||||
WHERE bank_id = $1
|
||||
AND id = ANY($2::uuid[])
|
||||
AND id NOT IN (SELECT DISTINCT entity_id FROM {ue_table})
|
||||
""",
|
||||
bank_id,
|
||||
entity_ids,
|
||||
)
|
||||
return int(deleted.split()[-1]) if isinstance(deleted, str) and deleted.startswith("DELETE") else 0
|
||||
|
||||
@@ -382,26 +452,33 @@ class OracleOps(DataAccessOps):
|
||||
conn: DatabaseConnection,
|
||||
ec_table: str,
|
||||
ue_table: str,
|
||||
entities_table: str,
|
||||
bank_id: str,
|
||||
entity_ids: list,
|
||||
) -> int:
|
||||
if not entity_ids:
|
||||
return 0
|
||||
# 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.
|
||||
#
|
||||
# Scoped to the claimed candidates on either endpoint (#3222). The OR is
|
||||
# safe to write directly here, unlike on Postgres: both endpoint columns
|
||||
# are indexed and Oracle's optimizer expands an OR of two index-driven
|
||||
# predicates into a concatenation, rather than the whole-table scan the
|
||||
# PG planner picks (which is why the PG side spells it as a UNION).
|
||||
deleted = await conn.execute(
|
||||
f"""
|
||||
DELETE FROM {ec_table}
|
||||
WHERE entity_id_1 IN (SELECT id FROM {entities_table} WHERE bank_id = $1)
|
||||
WHERE (entity_id_1 = ANY($1::uuid[]) OR entity_id_2 = ANY($1::uuid[]))
|
||||
AND (entity_id_1, entity_id_2) NOT IN (
|
||||
SELECT u1.entity_id, u2.entity_id
|
||||
FROM {ue_table} u1
|
||||
JOIN {ue_table} u2 ON u1.unit_id = u2.unit_id
|
||||
)
|
||||
""",
|
||||
bank_id,
|
||||
entity_ids,
|
||||
)
|
||||
return int(deleted.split()[-1]) if isinstance(deleted, str) and deleted.startswith("DELETE") else 0
|
||||
|
||||
@@ -619,6 +696,21 @@ class OracleOps(DataAccessOps):
|
||||
# Entity expansion via observation_sources junction table.
|
||||
# Previously used JSON_TABLE to explode source_memory_ids CLOB. The junction
|
||||
# table approach uses standard SQL joins, identical to the PG backend.
|
||||
#
|
||||
# Two PostgreSQL fixes are deliberately NOT mirrored here, because neither
|
||||
# was measured against Oracle and both are tuned to PostgreSQL's planner:
|
||||
# - #3085 made PG score set-wise; the scoring below is still the
|
||||
# correlated per-observation COUNT(*). On Oracle that counts rows of
|
||||
# the indexed observation_sources junction table rather than scanning
|
||||
# an unpruned array, so it is a much weaker version of that problem.
|
||||
# - #3510 replaced PG's `DISTINCT` over a `LATERAL ... LIMIT` with a
|
||||
# row_number() window, because PostgreSQL cannot estimate the row count
|
||||
# of that shape and mis-planned the scoring join into a nested loop.
|
||||
# `connected_sources` below has the same shape, so the same collapse is
|
||||
# structurally possible, but Oracle's cardinality estimation differs and
|
||||
# no Oracle instance was available to measure it.
|
||||
# If observation recall is reported slow on Oracle, start by capturing the
|
||||
# plan for connected_sources and checking its estimated vs actual rows.
|
||||
from ..schema import fq_table
|
||||
|
||||
obs_sources_table = fq_table("observation_sources")
|
||||
@@ -752,23 +844,6 @@ class OracleOps(DataAccessOps):
|
||||
bank_prefix="mu.",
|
||||
)
|
||||
|
||||
async def create_bank_vector_indexes(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
bank_id: str,
|
||||
internal_id: str,
|
||||
index_clause: str,
|
||||
fact_types: dict[str, str],
|
||||
) -> None:
|
||||
# Oracle 23ai supports HNSW vector indexes but does NOT support partial
|
||||
# indexes (WHERE clause on CREATE INDEX for vector indexes). Uses a single
|
||||
# global HNSW index with ORGANIZATION NEIGHBOR PARTITIONS created during
|
||||
# migrations. memory_units is partitioned by LIST (bank_id) AUTOMATIC,
|
||||
# so Oracle creates partitions per bank on INSERT and the optimizer can
|
||||
# prune partitions on bank_id-scoped queries.
|
||||
return
|
||||
|
||||
async def drop_bank_vector_indexes(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
@@ -776,7 +851,12 @@ class OracleOps(DataAccessOps):
|
||||
internal_id: str,
|
||||
fact_types: dict[str, str],
|
||||
) -> None:
|
||||
# Oracle uses a single global vector index (no per-bank indexes to drop).
|
||||
# Oracle uses a single global vector index — it does not support partial
|
||||
# (WHERE-clause) vector indexes, so there are no per-bank ones to drop.
|
||||
# Bank scoping comes from the table itself instead: memory_units is
|
||||
# partitioned LIST (bank_id) AUTOMATIC, so Oracle creates a partition per
|
||||
# bank on INSERT and the optimizer prunes on bank_id. That is why the
|
||||
# size threshold and its sweep are PostgreSQL-only concerns.
|
||||
return
|
||||
|
||||
def get_entity_resolution_strategy(self) -> str:
|
||||
@@ -1144,7 +1224,7 @@ class OracleOps(DataAccessOps):
|
||||
if exclude_ids:
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
SELECT operation_id, operation_type, task_payload, retry_count, serialization_key, bank_id
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
@@ -1163,7 +1243,7 @@ class OracleOps(DataAccessOps):
|
||||
else:
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
SELECT operation_id, operation_type, task_payload, retry_count, serialization_key, bank_id
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
@@ -1181,7 +1261,7 @@ class OracleOps(DataAccessOps):
|
||||
if exclude_ids:
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
SELECT operation_id, operation_type, task_payload, retry_count, serialization_key, bank_id
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
@@ -1198,7 +1278,7 @@ class OracleOps(DataAccessOps):
|
||||
else:
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
SELECT operation_id, operation_type, task_payload, retry_count, serialization_key, bank_id
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
@@ -1239,7 +1319,7 @@ class OracleOps(DataAccessOps):
|
||||
extra = " AND ".join(conditions)
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
SELECT operation_id, operation_type, task_payload, retry_count, serialization_key, bank_id
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
@@ -1286,7 +1366,7 @@ class OracleOps(DataAccessOps):
|
||||
extra_clause = (" AND " + " AND ".join(conditions)) if conditions else ""
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
SELECT operation_id, operation_type, task_payload, retry_count, serialization_key, bank_id
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
@@ -1338,13 +1418,14 @@ class OracleOps(DataAccessOps):
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT o.operation_id, o.operation_type, o.task_payload, o.retry_count
|
||||
SELECT o.operation_id, o.operation_type, o.task_payload, o.retry_count, o.serialization_key, o.bank_id
|
||||
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")}
|
||||
AND {document_serialization_sql(table, "o")}
|
||||
ORDER BY o.created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
@@ -1366,7 +1447,7 @@ class OracleOps(DataAccessOps):
|
||||
if claimed_ids:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT o.operation_id, o.operation_type, o.task_payload, o.retry_count
|
||||
SELECT o.operation_id, o.operation_type, o.task_payload, o.retry_count, o.serialization_key, o.bank_id
|
||||
FROM {table} o
|
||||
WHERE o.status = 'pending'
|
||||
AND o.task_payload IS NOT NULL
|
||||
@@ -1374,6 +1455,7 @@ class OracleOps(DataAccessOps):
|
||||
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")}
|
||||
AND {document_serialization_sql(table, "o")}
|
||||
ORDER BY o.created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
@@ -1384,13 +1466,14 @@ class OracleOps(DataAccessOps):
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT o.operation_id, o.operation_type, o.task_payload, o.retry_count
|
||||
SELECT o.operation_id, o.operation_type, o.task_payload, o.retry_count, o.serialization_key, o.bank_id
|
||||
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")}
|
||||
AND {document_serialization_sql(table, "o")}
|
||||
ORDER BY o.created_at
|
||||
LIMIT $1
|
||||
FOR UPDATE SKIP LOCKED
|
||||
@@ -1431,6 +1514,19 @@ class OracleOps(DataAccessOps):
|
||||
|
||||
# Mark all claimed rows as processing
|
||||
operation_ids = [row["operation_id"] for row in all_rows]
|
||||
await self.mark_operations_processing(conn, table, worker_id, operation_ids)
|
||||
|
||||
return all_rows
|
||||
|
||||
async def mark_operations_processing(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
worker_id: str,
|
||||
operation_ids: list,
|
||||
) -> None:
|
||||
if not operation_ids:
|
||||
return
|
||||
await conn.execute(
|
||||
f"""
|
||||
UPDATE {table}
|
||||
@@ -1441,4 +1537,34 @@ class OracleOps(DataAccessOps):
|
||||
operation_ids,
|
||||
)
|
||||
|
||||
return all_rows
|
||||
async def fetch_foldable_retain_peers(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
bank_id: str,
|
||||
serialization_key: str,
|
||||
limit: int,
|
||||
) -> list[ResultRow]:
|
||||
if limit <= 0:
|
||||
return []
|
||||
# Same ``LIMIT $n ... FOR UPDATE SKIP LOCKED`` shape the claim queries
|
||||
# above use, which the Oracle SQL translation layer rewrites into the
|
||||
# row-limited form Oracle accepts (a bare one raises ORA-02014).
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'retain'
|
||||
AND bank_id = $1
|
||||
AND serialization_key = $2
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
ORDER BY created_at, operation_id
|
||||
LIMIT $3
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
bank_id,
|
||||
serialization_key,
|
||||
limit,
|
||||
)
|
||||
|
||||
@@ -13,6 +13,7 @@ from .ops import (
|
||||
LinkExpansionRows,
|
||||
TagListingParts,
|
||||
UpdatedWindow,
|
||||
document_serialization_sql,
|
||||
graph_maintenance_bank_serialization_sql,
|
||||
)
|
||||
from .result import ResultRow
|
||||
@@ -472,26 +473,126 @@ class PostgreSQLOps(DataAccessOps):
|
||||
)
|
||||
return [str(row["unit_id"]) for row in rows]
|
||||
|
||||
async def enqueue_entity_maintenance(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
ue_table: str,
|
||||
bank_id: str,
|
||||
unit_ids: list,
|
||||
) -> int:
|
||||
# Read the candidates straight out of unit_entities rather than making
|
||||
# callers pass entity ids: every caller runs this immediately before the
|
||||
# rows go, and the join they'd have to write is this one.
|
||||
#
|
||||
# The inner ORDER BY is load-bearing, not cosmetic: it makes the INSERT
|
||||
# take the (bank_id, entity_id) row locks ascending, the same order
|
||||
# claim_entity_maintenance_batch takes them, so a mutation enqueueing an
|
||||
# overlapping candidate set cannot cycle against a worker draining it.
|
||||
# (Same protocol as enqueue_graph_maintenance, which sorts in Python
|
||||
# because its ids arrive as a bind array.)
|
||||
#
|
||||
# DO UPDATE (not DO NOTHING) on a duplicate — #3034. The SET is a
|
||||
# deliberate no-op preserving enqueued_at; its only purpose is to lock
|
||||
# the conflicting row. DO NOTHING does not lock it, so a delete
|
||||
# re-enqueueing an already-queued entity could not block a worker from
|
||||
# claiming that row and evaluating the entity's pre-delete state — it
|
||||
# would find the entity still referenced, keep it, and the re-enqueue
|
||||
# signal would be lost, stranding the orphan until some later delete
|
||||
# happened to name it again.
|
||||
result = await conn.execute(
|
||||
f"""
|
||||
INSERT INTO {table} (bank_id, entity_id)
|
||||
SELECT $1, s.entity_id
|
||||
FROM (
|
||||
SELECT DISTINCT ue.entity_id
|
||||
FROM {ue_table} ue
|
||||
WHERE ue.unit_id = ANY($2::uuid[])
|
||||
ORDER BY 1
|
||||
) s
|
||||
ON CONFLICT (bank_id, entity_id)
|
||||
DO UPDATE SET enqueued_at = {table}.enqueued_at
|
||||
""",
|
||||
bank_id,
|
||||
unit_ids,
|
||||
)
|
||||
return int(result.split()[-1]) if isinstance(result, str) and result.startswith("INSERT") else 0
|
||||
|
||||
async def claim_entity_maintenance_batch(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
bank_id: str,
|
||||
limit: int,
|
||||
) -> list:
|
||||
# Same claim shape as claim_graph_maintenance_batch: pick the oldest
|
||||
# batch by enqueued_at, but acquire the row locks in (bank_id, entity_id)
|
||||
# order — the order enqueue_entity_maintenance takes them — so a
|
||||
# concurrent enqueue can never cycle against this claim. `chosen` is
|
||||
# MATERIALIZED so the enqueued_at pick is fenced from the locking clause,
|
||||
# and `FOR UPDATE OF q ... ORDER BY q.entity_id` puts LockRows above the
|
||||
# Sort.
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
WITH chosen AS MATERIALIZED (
|
||||
SELECT bank_id, entity_id FROM {table}
|
||||
WHERE bank_id = $1
|
||||
ORDER BY enqueued_at
|
||||
LIMIT $2
|
||||
),
|
||||
locked AS (
|
||||
SELECT q.bank_id, q.entity_id
|
||||
FROM {table} q
|
||||
JOIN chosen c ON c.bank_id = q.bank_id AND c.entity_id = q.entity_id
|
||||
ORDER BY q.entity_id
|
||||
FOR UPDATE OF q
|
||||
)
|
||||
DELETE FROM {table} q
|
||||
USING locked l
|
||||
WHERE q.bank_id = l.bank_id AND q.entity_id = l.entity_id
|
||||
RETURNING q.entity_id
|
||||
""",
|
||||
bank_id,
|
||||
limit,
|
||||
)
|
||||
return [row["entity_id"] for row in rows]
|
||||
|
||||
async def prune_orphan_entities(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
entities_table: str,
|
||||
ue_table: str,
|
||||
bank_id: str,
|
||||
entity_ids: list,
|
||||
) -> int:
|
||||
# Scoped by entities.bank_id (indexed). The NOT EXISTS subquery is
|
||||
# backed by idx_ue_entity on unit_entities(entity_id), so this stays
|
||||
# linear in the number of entities in the bank — not in the size of
|
||||
# unit_entities globally.
|
||||
# Scoped to the claimed candidates: primary-key lookups, with the
|
||||
# NOT EXISTS backed by idx_unit_entities_entity_unit. Cost tracks the
|
||||
# batch, not the bank (#3222) — the bank-wide form this replaces probed
|
||||
# once per entity in the bank on every single run.
|
||||
#
|
||||
# Victims are locked in id order before the delete so the locks are
|
||||
# acquired the same way retain's entity upsert takes them
|
||||
# (bulk_upsert_entities locks `ORDER BY id FOR KEY SHARE`), which is what
|
||||
# keeps a prune and a concurrent re-assert from cycling.
|
||||
result = await conn.execute(
|
||||
f"""
|
||||
WITH victims AS (
|
||||
SELECT e.id
|
||||
FROM {entities_table} e
|
||||
WHERE e.bank_id = $1
|
||||
AND e.id = ANY($2::uuid[])
|
||||
AND NOT EXISTS (
|
||||
SELECT 1 FROM {ue_table} ue WHERE ue.entity_id = e.id
|
||||
)
|
||||
ORDER BY e.id
|
||||
FOR UPDATE
|
||||
)
|
||||
DELETE FROM {entities_table} e
|
||||
WHERE e.bank_id = $1
|
||||
AND NOT EXISTS (
|
||||
SELECT 1 FROM {ue_table} ue WHERE ue.entity_id = e.id
|
||||
)
|
||||
USING victims v
|
||||
WHERE e.id = v.id
|
||||
""",
|
||||
bank_id,
|
||||
entity_ids,
|
||||
)
|
||||
# asyncpg returns "DELETE N"
|
||||
return int(result.split()[-1]) if isinstance(result, str) and result.startswith("DELETE") else 0
|
||||
@@ -501,12 +602,18 @@ class PostgreSQLOps(DataAccessOps):
|
||||
conn: DatabaseConnection,
|
||||
ec_table: str,
|
||||
ue_table: str,
|
||||
entities_table: str,
|
||||
bank_id: str,
|
||||
entity_ids: list,
|
||||
) -> int:
|
||||
# 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).
|
||||
# Scoped to cooccurrence rows incident to the claimed candidates. The
|
||||
# two arms are a UNION rather than
|
||||
# `WHERE entity_id_1 = ANY(...) OR entity_id_2 = ANY(...)`: an OR across
|
||||
# two columns of the same table cannot be driven from either index, so
|
||||
# the planner would make entity_cooccurrences the outer relation and
|
||||
# scan it whole — the #3387 shape. As a UNION each arm is an index scan
|
||||
# (the PK for arm 1, idx_entity_cooccurrences_entity2 for arm 2).
|
||||
#
|
||||
# No bank predicate: entities don't span banks and the candidates came
|
||||
# off a bank-scoped queue, so both endpoints are already this bank's.
|
||||
#
|
||||
# Ordered locking (deadlock avoidance, #2529): retain's concurrent
|
||||
# cooccurrence upsert (entity_resolver._flush_pending) locks rows in
|
||||
@@ -522,26 +629,58 @@ class PostgreSQLOps(DataAccessOps):
|
||||
# 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.
|
||||
# Staleness is decided against a SET of currently-live pairs, not with a
|
||||
# per-cooccurrence-row check (#3367). The old form ran a correlated
|
||||
# `NOT EXISTS (… INTERSECT …)` per row, and each evaluation re-scanned a
|
||||
# hub entity's full membership set — cost scaled as (rows judged) x (hub
|
||||
# degree), 88-140s on a real bank with a ~22K-degree hub. #2473 had
|
||||
# swapped an earlier hub-rescanning self-join to that INTERSECT, but only
|
||||
# made each per-row check cheaper; it kept the per-row structure, so the
|
||||
# product blew up again at scale.
|
||||
#
|
||||
# `live` groups unit_entities by unit (self-join on unit_id) to emit every
|
||||
# co-occurring (e1<e2) pair in one materialised pass: its cost is driven
|
||||
# by unit degree (entities per unit — small), never by entity degree, so a
|
||||
# hub contributes only its per-unit membership rather than a rescan per
|
||||
# edge. The victims anti-join then hashes against it. MATERIALIZED keeps
|
||||
# the planner from inlining `live` back into a per-row correlated plan.
|
||||
#
|
||||
# `live` is seeded from the *candidates'* units rather than the whole
|
||||
# bank's (#3222 composed with #3367): graph maintenance is queue-driven,
|
||||
# so this only has to decide about pairs in `incident`, and every such
|
||||
# pair has a candidate as at least one endpoint. Any unit still
|
||||
# witnessing such a pair therefore references a candidate, and so is in
|
||||
# the seeded set — the scoped build cannot miss a live pair. That keeps
|
||||
# the whole statement proportional to the batch instead of re-deriving
|
||||
# every pair in the bank on every run.
|
||||
result = await conn.execute(
|
||||
f"""
|
||||
WITH victims AS (
|
||||
WITH incident AS MATERIALIZED (
|
||||
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
|
||||
)
|
||||
WHERE c.entity_id_1 = ANY($1::uuid[])
|
||||
UNION
|
||||
SELECT c.entity_id_1, c.entity_id_2
|
||||
FROM {ec_table} c
|
||||
WHERE c.entity_id_2 = ANY($1::uuid[])
|
||||
),
|
||||
live AS MATERIALIZED (
|
||||
SELECT u1.entity_id AS e1, u2.entity_id AS e2
|
||||
FROM {ue_table} seed
|
||||
JOIN {ue_table} u1 ON u1.unit_id = seed.unit_id
|
||||
JOIN {ue_table} u2 ON u2.unit_id = u1.unit_id
|
||||
AND u2.entity_id > u1.entity_id
|
||||
WHERE seed.entity_id = ANY($1::uuid[])
|
||||
),
|
||||
victims AS (
|
||||
SELECT c.entity_id_1, c.entity_id_2
|
||||
FROM incident i
|
||||
JOIN {ec_table} c
|
||||
ON c.entity_id_1 = i.entity_id_1 AND c.entity_id_2 = i.entity_id_2
|
||||
WHERE NOT EXISTS (
|
||||
SELECT 1 FROM live l
|
||||
WHERE l.e1 = c.entity_id_1 AND l.e2 = c.entity_id_2
|
||||
)
|
||||
ORDER BY c.entity_id_1, c.entity_id_2
|
||||
FOR UPDATE OF c
|
||||
)
|
||||
@@ -550,7 +689,7 @@ class PostgreSQLOps(DataAccessOps):
|
||||
WHERE c.entity_id_1 = v.entity_id_1
|
||||
AND c.entity_id_2 = v.entity_id_2
|
||||
""",
|
||||
bank_id,
|
||||
entity_ids,
|
||||
)
|
||||
return int(result.split()[-1]) if isinstance(result, str) and result.startswith("DELETE") else 0
|
||||
|
||||
@@ -764,6 +903,38 @@ class PostgreSQLOps(DataAccessOps):
|
||||
# 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.
|
||||
#
|
||||
# The shared-source count is scored set-wise (`scored`), not per candidate row.
|
||||
# It used to be a correlated subquery — COUNT(DISTINCT s) over
|
||||
# unnest(mu.source_memory_ids) filtered by `= ANY(ca.source_ids)` — which
|
||||
# re-scanned the connected-source array linearly for every element of every
|
||||
# candidate's array. Because consolidation appends to source_memory_ids and
|
||||
# never prunes it (issue #1725), that product grows with the bank's age: at
|
||||
# 5k observations averaging 113 sources against ~3k connected sources it was
|
||||
# ~1.7B element comparisons, 2.6s of one saturated backend (issue #3085).
|
||||
# Unnesting once and hash-joining connected_sources makes the work linear in
|
||||
# the number of source ids instead.
|
||||
#
|
||||
# `connected_sources` caps each entity with row_number() rather than the
|
||||
# LATERAL + LIMIT that reads more naturally. Do not "simplify" it back
|
||||
# (issue #3510). The scoring join above is O(C + U) when planned as a hash
|
||||
# join and O(U x C) when planned as a nested loop — 15s and ~15M rejected
|
||||
# rows on a realistically-shaped bank — and PostgreSQL picks between them
|
||||
# from its row estimate for this CTE. Out of a LATERAL + LIMIT subquery the
|
||||
# capped column carries no n_distinct statistic, so DISTINCT over it was
|
||||
# estimated at 2 and the NOT EXISTS took that to 1 against an actual ~3,700;
|
||||
# a 1-row inner side makes the nested loop look free, so it won on cost and
|
||||
# lost by four orders of magnitude at runtime. Ranking with a window keeps
|
||||
# the column traceable to unit_entities.unit_id, so the estimate comes from
|
||||
# real statistics (207-3,449 against 2,242-4,193 actual) and the nested loop
|
||||
# is priced honestly.
|
||||
#
|
||||
# The trade is that this reads every unit_entities row of a matched entity
|
||||
# to rank it, where the LATERAL stopped at per_entity_limit off the index:
|
||||
# O(sum of degree) rather than O(entities x per_entity_limit). Measured at
|
||||
# parity up to ~12k-degree hubs and +50% traversal cost at 38k. If banks
|
||||
# grow hubs far past that, re-measure before assuming this is still the
|
||||
# right shape.
|
||||
|
||||
entity_rows = await conn.fetch(
|
||||
f"""
|
||||
@@ -780,33 +951,52 @@ class PostgreSQLOps(DataAccessOps):
|
||||
),
|
||||
connected_sources AS (
|
||||
SELECT DISTINCT t.unit_id AS source_id
|
||||
FROM source_entities se
|
||||
CROSS JOIN LATERAL (
|
||||
SELECT ue_target.unit_id
|
||||
FROM (
|
||||
SELECT
|
||||
ue_target.unit_id,
|
||||
row_number() OVER (
|
||||
PARTITION BY ue_target.entity_id
|
||||
ORDER BY ue_target.unit_id DESC
|
||||
) AS rn
|
||||
FROM {ue_table} ue_target
|
||||
WHERE ue_target.entity_id = se.entity_id
|
||||
ORDER BY ue_target.unit_id DESC
|
||||
LIMIT {per_entity_limit}
|
||||
JOIN source_entities se ON se.entity_id = ue_target.entity_id
|
||||
) t
|
||||
WHERE NOT EXISTS (
|
||||
SELECT 1 FROM seed_sources ss WHERE ss.source_id = t.unit_id
|
||||
)
|
||||
WHERE t.rn <= {per_entity_limit}
|
||||
AND NOT EXISTS (
|
||||
SELECT 1 FROM seed_sources ss WHERE ss.source_id = t.unit_id
|
||||
)
|
||||
),
|
||||
connected_array AS (
|
||||
SELECT array_agg(source_id) AS source_ids FROM connected_sources
|
||||
),
|
||||
candidates AS (
|
||||
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, mu.chunk_id, mu.tags, mu.proof_count,
|
||||
mu.source_memory_ids
|
||||
FROM {mu_table} mu, connected_array ca
|
||||
WHERE mu.fact_type = 'observation'
|
||||
AND mu.id != ALL($1::uuid[])
|
||||
AND ca.source_ids IS NOT NULL
|
||||
AND mu.source_memory_ids && ca.source_ids
|
||||
{window.clause("mu")}
|
||||
),
|
||||
scored AS (
|
||||
SELECT c.id, COUNT(DISTINCT cs.source_id)::float AS score
|
||||
FROM candidates c
|
||||
CROSS JOIN LATERAL unnest(c.source_memory_ids) AS s(source_id)
|
||||
JOIN connected_sources cs ON cs.source_id = s.source_id
|
||||
GROUP BY c.id
|
||||
)
|
||||
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, mu.chunk_id, mu.tags, mu.proof_count,
|
||||
(SELECT COUNT(DISTINCT s) FROM unnest(mu.source_memory_ids) s WHERE s = ANY(ca.source_ids))::float AS score
|
||||
FROM {mu_table} mu, connected_array ca
|
||||
WHERE mu.fact_type = 'observation'
|
||||
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
|
||||
c.id, c.text, c.context, c.event_date, c.occurred_start,
|
||||
c.occurred_end, c.mentioned_at,
|
||||
c.fact_type, c.document_id, c.chunk_id, c.tags, c.proof_count,
|
||||
sc.score
|
||||
FROM candidates c
|
||||
JOIN scored sc ON sc.id = c.id
|
||||
ORDER BY sc.score DESC
|
||||
LIMIT $2
|
||||
""",
|
||||
seed_ids,
|
||||
@@ -881,26 +1071,6 @@ class PostgreSQLOps(DataAccessOps):
|
||||
bank_prefix="",
|
||||
)
|
||||
|
||||
async def create_bank_vector_indexes(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
bank_id: str,
|
||||
internal_id: str,
|
||||
index_clause: str,
|
||||
fact_types: dict[str, str],
|
||||
) -> None:
|
||||
escaped = bank_id.replace("'", "''")
|
||||
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,
|
||||
conn: DatabaseConnection,
|
||||
@@ -914,8 +1084,14 @@ class PostgreSQLOps(DataAccessOps):
|
||||
# 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`.
|
||||
#
|
||||
# The in-process lock serializes concurrent bank deletes against each
|
||||
# other. It does not cover the maintenance sweep, which reconciles the
|
||||
# same indexes over its own raw connection: an in-process lock could not
|
||||
# help there anyway, since the sweep runs in every process and the real
|
||||
# contention is cross-process. Both paths retry the transient deadlock
|
||||
# (40P01) instead, which is the only lock-free option available — the
|
||||
# project forbids advisory locks (unreliable behind poolers, #2817).
|
||||
async with self._index_ddl_lock(f"{schema}.memory_units"):
|
||||
for ft, suffix in fact_types.items():
|
||||
uid = str(internal_id).replace("-", "")[:16]
|
||||
@@ -1234,7 +1410,7 @@ class PostgreSQLOps(DataAccessOps):
|
||||
if exclude_ids:
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
SELECT operation_id, operation_type, task_payload, retry_count, serialization_key, bank_id
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
@@ -1253,7 +1429,7 @@ class PostgreSQLOps(DataAccessOps):
|
||||
else:
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
SELECT operation_id, operation_type, task_payload, retry_count, serialization_key, bank_id
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
@@ -1271,7 +1447,7 @@ class PostgreSQLOps(DataAccessOps):
|
||||
if exclude_ids:
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
SELECT operation_id, operation_type, task_payload, retry_count, serialization_key, bank_id
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
@@ -1288,7 +1464,7 @@ class PostgreSQLOps(DataAccessOps):
|
||||
else:
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
SELECT operation_id, operation_type, task_payload, retry_count, serialization_key, bank_id
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
@@ -1329,7 +1505,7 @@ class PostgreSQLOps(DataAccessOps):
|
||||
extra = " AND ".join(conditions)
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
SELECT operation_id, operation_type, task_payload, retry_count, serialization_key, bank_id
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
@@ -1376,7 +1552,7 @@ class PostgreSQLOps(DataAccessOps):
|
||||
extra_clause = (" AND " + " AND ".join(conditions)) if conditions else ""
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
SELECT operation_id, operation_type, task_payload, retry_count, serialization_key, bank_id
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
@@ -1427,13 +1603,14 @@ class PostgreSQLOps(DataAccessOps):
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT o.operation_id, o.operation_type, o.task_payload, o.retry_count
|
||||
SELECT o.operation_id, o.operation_type, o.task_payload, o.retry_count, o.serialization_key, o.bank_id
|
||||
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")}
|
||||
AND {document_serialization_sql(table, "o")}
|
||||
ORDER BY o.created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
@@ -1455,7 +1632,7 @@ class PostgreSQLOps(DataAccessOps):
|
||||
if claimed_ids:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT o.operation_id, o.operation_type, o.task_payload, o.retry_count
|
||||
SELECT o.operation_id, o.operation_type, o.task_payload, o.retry_count, o.serialization_key, o.bank_id
|
||||
FROM {table} o
|
||||
WHERE o.status = 'pending'
|
||||
AND o.task_payload IS NOT NULL
|
||||
@@ -1463,6 +1640,7 @@ class PostgreSQLOps(DataAccessOps):
|
||||
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")}
|
||||
AND {document_serialization_sql(table, "o")}
|
||||
ORDER BY o.created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
@@ -1473,13 +1651,14 @@ class PostgreSQLOps(DataAccessOps):
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT o.operation_id, o.operation_type, o.task_payload, o.retry_count
|
||||
SELECT o.operation_id, o.operation_type, o.task_payload, o.retry_count, o.serialization_key, o.bank_id
|
||||
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")}
|
||||
AND {document_serialization_sql(table, "o")}
|
||||
ORDER BY o.created_at
|
||||
LIMIT $1
|
||||
FOR UPDATE SKIP LOCKED
|
||||
@@ -1520,6 +1699,19 @@ class PostgreSQLOps(DataAccessOps):
|
||||
|
||||
# Mark all claimed rows as processing
|
||||
operation_ids = [row["operation_id"] for row in all_rows]
|
||||
await self.mark_operations_processing(conn, table, worker_id, operation_ids)
|
||||
|
||||
return all_rows
|
||||
|
||||
async def mark_operations_processing(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
worker_id: str,
|
||||
operation_ids: list,
|
||||
) -> None:
|
||||
if not operation_ids:
|
||||
return
|
||||
await conn.execute(
|
||||
f"""
|
||||
UPDATE {table}
|
||||
@@ -1530,4 +1722,31 @@ class PostgreSQLOps(DataAccessOps):
|
||||
operation_ids,
|
||||
)
|
||||
|
||||
return all_rows
|
||||
async def fetch_foldable_retain_peers(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
bank_id: str,
|
||||
serialization_key: str,
|
||||
limit: int,
|
||||
) -> list[ResultRow]:
|
||||
if limit <= 0:
|
||||
return []
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'retain'
|
||||
AND bank_id = $1
|
||||
AND serialization_key = $2
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
ORDER BY created_at, operation_id
|
||||
LIMIT $3
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
bank_id,
|
||||
serialization_key,
|
||||
limit,
|
||||
)
|
||||
|
||||
@@ -8,9 +8,10 @@ avoiding Python-level wrapping overhead (~570K __getitem__ calls per
|
||||
"""
|
||||
|
||||
import logging
|
||||
from collections.abc import AsyncIterator
|
||||
from collections.abc import AsyncIterator, Awaitable, Callable
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import Any
|
||||
from urllib.parse import parse_qs, urlparse
|
||||
|
||||
import asyncpg # noqa: F401
|
||||
|
||||
@@ -20,6 +21,42 @@ from .pool_instrumentation import PoolStats, instrument_acquire
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
async def apply_session_settings(conn: asyncpg.Connection, settings: list[tuple[str, str]]) -> None:
|
||||
"""Apply session-scoped GUCs to ``conn`` in a single round trip.
|
||||
|
||||
Unless ``HINDSIGHT_API_DB_SESSION_SETUP_ON_ACQUIRE=false``, the pool passes
|
||||
its init callback as ``setup=`` too, so this runs on *every* acquire, not
|
||||
just on connection creation. Issued as N separate ``SET`` statements that
|
||||
was N round trips — and behind a transaction-mode pooler N server-side
|
||||
transactions — per acquire, which the worker's per-schema acquires
|
||||
multiplied into a sustained commit-rate burn (#3499). One
|
||||
``SELECT set_config(...)`` collapses them into one statement.
|
||||
|
||||
Some of the settings are extension-provided (``hnsw.ef_search``,
|
||||
``pg_trgm.similarity_threshold``) and may not exist on the cluster; a single
|
||||
statement fails as a whole, so on error fall back to applying them one by
|
||||
one, skipping only the ones the server rejects.
|
||||
"""
|
||||
if not settings:
|
||||
return
|
||||
|
||||
args: list[str] = [value for pair in settings for value in pair]
|
||||
projection = ", ".join(f"set_config(${2 * i + 1}, ${2 * i + 2}, false)" for i in range(len(settings)))
|
||||
try:
|
||||
await conn.execute(f"SELECT {projection}", *args)
|
||||
return
|
||||
except asyncpg.exceptions.PostgresError:
|
||||
# Narrow to PostgresError so genuine bugs in the pool/conn layer surface
|
||||
# instead of being silently retried statement-by-statement.
|
||||
logger.debug("Batched session setup failed — applying settings individually")
|
||||
|
||||
for name, value in settings:
|
||||
try:
|
||||
await conn.execute("SELECT set_config($1, $2, false)", name, value)
|
||||
except asyncpg.exceptions.PostgresError:
|
||||
logger.debug("Could not set %s — the server may not know this setting", name)
|
||||
|
||||
|
||||
class PostgresConnection(DatabaseConnection):
|
||||
"""DatabaseConnection wrapper around an asyncpg.Connection."""
|
||||
|
||||
@@ -64,6 +101,54 @@ class PostgresConnection(DatabaseConnection):
|
||||
await self._conn.copy_records_to_table(table_name, records=records, columns=columns, timeout=timeout)
|
||||
|
||||
|
||||
def application_name_from_dsn(dsn: str) -> str | None:
|
||||
"""Extract the ``application_name`` query parameter from a PostgreSQL DSN.
|
||||
|
||||
asyncpg already forwards this to the server in the startup packet (it
|
||||
passes unrecognized DSN query parameters through as ``server_settings``),
|
||||
so a direct connection is labelled correctly in ``pg_stat_activity``.
|
||||
The value is extracted here so it can be re-applied per acquire — see
|
||||
``_application_name_setup``.
|
||||
"""
|
||||
try:
|
||||
values = parse_qs(urlparse(dsn).query).get("application_name")
|
||||
except ValueError:
|
||||
return None
|
||||
if not values:
|
||||
return None
|
||||
# libpq semantics: the last occurrence of a repeated parameter wins.
|
||||
return values[-1] or None
|
||||
|
||||
|
||||
def _application_name_setup(app_name: str, init_callback: Any | None) -> Callable[[Any], Awaitable[None]]:
|
||||
"""Wrap ``init_callback`` so every acquire re-asserts ``application_name``.
|
||||
|
||||
asyncpg runs ``RESET ALL`` when a connection is released back to the pool.
|
||||
Connected straight to PostgreSQL that is harmless: ``RESET ALL`` restores
|
||||
the value from the startup packet, which carried the DSN's name.
|
||||
|
||||
Behind a connection pooler (pgbouncer) it is not. The server connection's
|
||||
startup packet is the *pooler's*, with no application_name; pgbouncer
|
||||
applies the client's value with a ``SET`` when it links client to server.
|
||||
``RESET ALL`` therefore resets it to empty, and pgbouncer — which already
|
||||
believes the value is applied — does not re-issue it. Only the first
|
||||
acquire on each server connection is attributed; every later one reports
|
||||
an empty application_name, which is exactly the sort of gap that shows up
|
||||
in production but never under psql.
|
||||
|
||||
Re-asserting it on every acquire fixes both topologies. ``set_config``
|
||||
rather than ``SET`` because the name is operator-supplied and ``SET`` does
|
||||
not accept bind parameters.
|
||||
"""
|
||||
|
||||
async def _setup(conn: Any) -> None:
|
||||
await conn.execute("SELECT set_config('application_name', $1, false)", app_name)
|
||||
if init_callback is not None:
|
||||
await init_callback(conn)
|
||||
|
||||
return _setup
|
||||
|
||||
|
||||
class PostgreSQLBackend(DatabaseBackend):
|
||||
"""DatabaseBackend implementation wrapping an asyncpg connection pool."""
|
||||
|
||||
@@ -79,6 +164,12 @@ class PostgreSQLBackend(DatabaseBackend):
|
||||
self._pool: asyncpg.Pool | None = None
|
||||
self._acquire_warn_threshold_s: float = 1.0
|
||||
self._acquire_timeout_s: float | None = None
|
||||
self._dsn: str | None = None
|
||||
|
||||
@property
|
||||
def dsn(self) -> str | None:
|
||||
"""The DSN this backend's pool was opened with, if it has been initialized."""
|
||||
return self._dsn
|
||||
|
||||
async def initialize(
|
||||
self,
|
||||
@@ -93,7 +184,15 @@ class PostgreSQLBackend(DatabaseBackend):
|
||||
) -> None:
|
||||
from ...config import get_config
|
||||
|
||||
self._acquire_warn_threshold_s = get_config().db_acquire_warn_threshold_ms / 1000.0
|
||||
config = get_config()
|
||||
# Kept so code that needs its *own* connection — CREATE/DROP INDEX
|
||||
# CONCURRENTLY cannot run on a pooled one inside a transaction — can
|
||||
# reach the database this engine is actually attached to. Re-deriving it
|
||||
# from HINDSIGHT_API_DATABASE_URL is wrong whenever the engine was handed
|
||||
# a DSN directly (embedders, and the test suite, which resolves pg0 in a
|
||||
# fixture and never sets the env var).
|
||||
self._dsn = dsn
|
||||
self._acquire_warn_threshold_s = 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.
|
||||
@@ -101,6 +200,30 @@ class PostgreSQLBackend(DatabaseBackend):
|
||||
# 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
|
||||
# The DSN's application_name survives RESET ALL only on a direct
|
||||
# connection; behind pgbouncer it has to be re-asserted per acquire
|
||||
# (see _application_name_setup).
|
||||
app_name = application_name_from_dsn(dsn)
|
||||
pool_init = _application_name_setup(app_name, init_callback) if app_name else init_callback
|
||||
|
||||
# init runs once per new connection; setup runs on every acquire, after
|
||||
# asyncpg's release-time RESET ALL. Re-running the session GUCs
|
||||
# (hnsw.ef_search, statement_timeout, …) there is what keeps a *reused*
|
||||
# connection from silently falling back to server defaults, so it is the
|
||||
# default. Deployments that pin those GUCs server-side (ALTER ROLE /
|
||||
# ALTER DATABASE ... SET) get them back from RESET ALL anyway, making the
|
||||
# re-apply a wasted round trip on every acquire — and behind a
|
||||
# transaction-mode pooler, a wasted transaction too (#3499); they can
|
||||
# drop it with HINDSIGHT_API_DB_SESSION_SETUP_ON_ACQUIRE=false.
|
||||
# application_name is NOT part of that trade-off: pgbouncer never
|
||||
# re-issues it after RESET ALL, so it keeps its per-acquire hook either
|
||||
# way (#3491).
|
||||
setup_on_acquire = config.db_session_setup_on_acquire
|
||||
if setup_on_acquire:
|
||||
pool_setup = pool_init
|
||||
else:
|
||||
pool_setup = _application_name_setup(app_name, None) if app_name else None
|
||||
|
||||
self._pool = await asyncpg.create_pool(
|
||||
dsn,
|
||||
min_size=min_size,
|
||||
@@ -108,16 +231,13 @@ 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,
|
||||
init=pool_init,
|
||||
setup=pool_setup,
|
||||
)
|
||||
logger.info(
|
||||
f"PostgreSQL pool created (min={min_size}, max={max_size}, "
|
||||
f"cmd_timeout={command_timeout}s, acquire_timeout={acquire_timeout}s)"
|
||||
f"cmd_timeout={command_timeout}s, acquire_timeout={acquire_timeout}s, "
|
||||
f"session_setup_on_acquire={setup_on_acquire})"
|
||||
)
|
||||
|
||||
async def shutdown(self) -> None:
|
||||
|
||||
@@ -125,12 +125,32 @@ class Embeddings(ABC):
|
||||
"""
|
||||
pass
|
||||
|
||||
# Client-side asymmetric prefixes, empty unless a provider populates them from
|
||||
# config. Class-level so providers that never set them are unchanged.
|
||||
query_prefix: str = ""
|
||||
passage_prefix: str = ""
|
||||
|
||||
def encode_query(self, texts: list[str]) -> list[list[float]]:
|
||||
"""Generate embeddings for query text. Providers without asymmetric embeddings use encode()."""
|
||||
return self.encode(texts)
|
||||
"""Generate embeddings for query text, applying the configured query prefix."""
|
||||
return self._encode_prefixed(texts, self.query_prefix)
|
||||
|
||||
def encode_documents(self, texts: list[str]) -> list[list[float]]:
|
||||
"""Generate embeddings for stored document text. Providers without asymmetric embeddings use encode()."""
|
||||
"""Generate embeddings for stored document text, applying the configured passage prefix."""
|
||||
return self._encode_prefixed(texts, self.passage_prefix)
|
||||
|
||||
def _encode_prefixed(self, texts: list[str], prefix: str) -> list[list[float]]:
|
||||
"""Prepend an asymmetric model's instruction before handing text to encode().
|
||||
|
||||
Asymmetric models (E5, embeddinggemma, ...) expect a different instruction in
|
||||
front of a search than in front of stored text. A provider that is plain
|
||||
text-in/vector-out — TEI, LiteLLM, anything behind an OpenAI-compatible
|
||||
/embeddings endpoint — has no other channel to carry that distinction, so the
|
||||
client has to prepend it. Providers with a native mechanism (SentenceTransformers'
|
||||
own prompts, ZeroEntropy's input_type) override encode_query/encode_documents
|
||||
instead and never reach this. Empty prefixes leave the text byte-identical.
|
||||
"""
|
||||
if prefix:
|
||||
return self.encode([f"{prefix}{text}" for text in texts])
|
||||
return self.encode(texts)
|
||||
|
||||
|
||||
@@ -396,17 +416,6 @@ class OnnxEmbeddings(Embeddings):
|
||||
self._dimension = detected
|
||||
logger.info("Embeddings: ONNX provider initialized (dim: %s)", self._dimension)
|
||||
|
||||
def _encode_prefixed(self, texts: list[str], prefix: str) -> list[list[float]]:
|
||||
if prefix:
|
||||
return self.encode([f"{prefix}{text}" for text in texts])
|
||||
return self.encode(texts)
|
||||
|
||||
def encode_query(self, texts: list[str]) -> list[list[float]]:
|
||||
return self._encode_prefixed(texts, self.query_prefix)
|
||||
|
||||
def encode_documents(self, texts: list[str]) -> list[list[float]]:
|
||||
return self._encode_prefixed(texts, self.passage_prefix)
|
||||
|
||||
def encode(self, texts: list[str]) -> list[list[float]]:
|
||||
if self._session is None or self._tokenizer is None:
|
||||
raise RuntimeError("Embeddings not initialized. Call initialize() first.")
|
||||
@@ -469,6 +478,8 @@ class RemoteTEIEmbeddings(Embeddings):
|
||||
batch_size: int = 32,
|
||||
max_retries: int = 3,
|
||||
retry_delay: float = 0.5,
|
||||
query_prefix: str = "",
|
||||
passage_prefix: str = "",
|
||||
):
|
||||
"""
|
||||
Initialize remote TEI embeddings client.
|
||||
@@ -479,12 +490,16 @@ class RemoteTEIEmbeddings(Embeddings):
|
||||
batch_size: Maximum batch size for embedding requests (default: 32)
|
||||
max_retries: Maximum number of retries for failed requests (default: 3)
|
||||
retry_delay: Initial delay between retries in seconds, doubles each retry (default: 0.5)
|
||||
query_prefix: Prefix prepended to recall/search queries (default: none)
|
||||
passage_prefix: Prefix prepended to retained document text (default: none)
|
||||
"""
|
||||
self.base_url = base_url.rstrip("/")
|
||||
self.timeout = timeout
|
||||
self.batch_size = batch_size
|
||||
self.max_retries = max_retries
|
||||
self.retry_delay = retry_delay
|
||||
self.query_prefix = query_prefix
|
||||
self.passage_prefix = passage_prefix
|
||||
self._client: httpx.Client | None = None
|
||||
self._model_id: str | None = None
|
||||
self._dimension: int | None = None
|
||||
@@ -638,6 +653,8 @@ class OpenAIEmbeddings(Embeddings):
|
||||
batch_size: int = 100,
|
||||
dimensions: int | None = None,
|
||||
max_retries: int = 3,
|
||||
query_prefix: str = "",
|
||||
passage_prefix: str = "",
|
||||
):
|
||||
"""
|
||||
Initialize OpenAI embeddings client.
|
||||
@@ -649,6 +666,8 @@ class OpenAIEmbeddings(Embeddings):
|
||||
batch_size: Maximum batch size for embedding requests (default: 100)
|
||||
dimensions: Optional requested output dimensions for OpenAI text-embedding-3 models
|
||||
max_retries: Maximum number of retries for failed requests (default: 3)
|
||||
query_prefix: Prefix prepended to recall/search queries (default: none)
|
||||
passage_prefix: Prefix prepended to retained document text (default: none)
|
||||
"""
|
||||
self.api_key = api_key
|
||||
self.model = model
|
||||
@@ -656,6 +675,8 @@ class OpenAIEmbeddings(Embeddings):
|
||||
self.batch_size = batch_size
|
||||
self.dimensions = dimensions
|
||||
self.max_retries = max_retries
|
||||
self.query_prefix = query_prefix
|
||||
self.passage_prefix = passage_prefix
|
||||
self._client = None
|
||||
self._dimension: int | None = None
|
||||
|
||||
@@ -774,6 +795,8 @@ class CodexOAuthEmbeddings(OpenAIEmbeddings):
|
||||
batch_size: int = 100,
|
||||
dimensions: int | None = None,
|
||||
max_retries: int = 3,
|
||||
query_prefix: str = "",
|
||||
passage_prefix: str = "",
|
||||
):
|
||||
from .providers.codex_auth import CodexAuthManager
|
||||
|
||||
@@ -785,6 +808,8 @@ class CodexOAuthEmbeddings(OpenAIEmbeddings):
|
||||
batch_size=batch_size,
|
||||
dimensions=dimensions,
|
||||
max_retries=max_retries,
|
||||
query_prefix=query_prefix,
|
||||
passage_prefix=passage_prefix,
|
||||
)
|
||||
|
||||
@property
|
||||
@@ -1129,6 +1154,8 @@ class LiteLLMEmbeddings(Embeddings):
|
||||
model: str = DEFAULT_EMBEDDINGS_LITELLM_MODEL,
|
||||
batch_size: int = 100,
|
||||
timeout: float = 60.0,
|
||||
query_prefix: str = "",
|
||||
passage_prefix: str = "",
|
||||
):
|
||||
"""
|
||||
Initialize LiteLLM embeddings client.
|
||||
@@ -1140,12 +1167,16 @@ class LiteLLMEmbeddings(Embeddings):
|
||||
Use provider prefix for non-OpenAI models (e.g., cohere/embed-english-v3.0)
|
||||
batch_size: Maximum batch size for embedding requests (default: 100)
|
||||
timeout: Request timeout in seconds (default: 60.0)
|
||||
query_prefix: Prefix prepended to recall/search queries (default: none)
|
||||
passage_prefix: Prefix prepended to retained document text (default: none)
|
||||
"""
|
||||
self.api_base = api_base.rstrip("/")
|
||||
self.api_key = api_key
|
||||
self.model = model
|
||||
self.batch_size = batch_size
|
||||
self.timeout = timeout
|
||||
self.query_prefix = query_prefix
|
||||
self.passage_prefix = passage_prefix
|
||||
self._client: httpx.Client | None = None
|
||||
self._dimension: int | None = None
|
||||
|
||||
@@ -1245,6 +1276,8 @@ class LiteLLMSDKEmbeddings(Embeddings):
|
||||
batch_size: int = 100,
|
||||
timeout: float = 60.0,
|
||||
encoding_format: str | None = "float",
|
||||
query_prefix: str = "",
|
||||
passage_prefix: str = "",
|
||||
):
|
||||
"""
|
||||
Initialize LiteLLM SDK embeddings client.
|
||||
@@ -1259,6 +1292,8 @@ class LiteLLMSDKEmbeddings(Embeddings):
|
||||
timeout: Request timeout in seconds (default: 60.0)
|
||||
encoding_format: Encoding format for embeddings (default: "float").
|
||||
Set to None or empty string to omit (needed for Voyage AI, Gemini).
|
||||
query_prefix: Prefix prepended to recall/search queries (default: none)
|
||||
passage_prefix: Prefix prepended to retained document text (default: none)
|
||||
"""
|
||||
self.api_key = api_key
|
||||
self.model = model
|
||||
@@ -1267,6 +1302,8 @@ class LiteLLMSDKEmbeddings(Embeddings):
|
||||
self.batch_size = batch_size
|
||||
self.timeout = timeout
|
||||
self.encoding_format = encoding_format or None
|
||||
self.query_prefix = query_prefix
|
||||
self.passage_prefix = passage_prefix
|
||||
self._litellm = None # Will be set during initialization
|
||||
self._dimension: int | None = None
|
||||
|
||||
@@ -1618,11 +1655,24 @@ def create_embeddings_from_env() -> Embeddings:
|
||||
config = get_config()
|
||||
provider = config.embeddings_provider.lower()
|
||||
|
||||
# Asymmetric prefixes are handed only to the providers that are plain
|
||||
# text-in/vector-out. `local` and `zeroentropy` carry the distinction natively
|
||||
# (SentenceTransformers prompts / input_type) and `onnx` has its own pair with
|
||||
# non-empty E5 defaults, so none of them take these.
|
||||
query_prefix = config.embeddings_query_prefix
|
||||
passage_prefix = config.embeddings_passage_prefix
|
||||
if query_prefix or passage_prefix:
|
||||
logger.info(
|
||||
"Embeddings: asymmetric prefixes configured (query=%r, passage=%r)",
|
||||
query_prefix,
|
||||
passage_prefix,
|
||||
)
|
||||
|
||||
if provider == "tei":
|
||||
url = config.embeddings_tei_url
|
||||
if not url:
|
||||
raise ValueError(f"{ENV_EMBEDDINGS_TEI_URL} is required when {ENV_EMBEDDINGS_PROVIDER} is 'tei'")
|
||||
return RemoteTEIEmbeddings(base_url=url)
|
||||
return RemoteTEIEmbeddings(base_url=url, query_prefix=query_prefix, passage_prefix=passage_prefix)
|
||||
elif provider == "local":
|
||||
return LocalSTEmbeddings(
|
||||
model_name=config.embeddings_local_model,
|
||||
@@ -1660,6 +1710,8 @@ def create_embeddings_from_env() -> Embeddings:
|
||||
base_url=base_url,
|
||||
batch_size=config.embeddings_openai_batch_size,
|
||||
dimensions=config.embeddings_openai_dimensions,
|
||||
query_prefix=query_prefix,
|
||||
passage_prefix=passage_prefix,
|
||||
)
|
||||
elif provider == "openai-codex":
|
||||
model = os.environ.get(ENV_EMBEDDINGS_OPENAI_MODEL, DEFAULT_EMBEDDINGS_OPENAI_MODEL)
|
||||
@@ -1667,6 +1719,8 @@ def create_embeddings_from_env() -> Embeddings:
|
||||
model=model,
|
||||
batch_size=config.embeddings_openai_batch_size,
|
||||
dimensions=config.embeddings_openai_dimensions,
|
||||
query_prefix=query_prefix,
|
||||
passage_prefix=passage_prefix,
|
||||
)
|
||||
elif provider == "openrouter":
|
||||
api_key = config.embeddings_openrouter_api_key
|
||||
@@ -1681,6 +1735,8 @@ def create_embeddings_from_env() -> Embeddings:
|
||||
base_url="https://openrouter.ai/api/v1",
|
||||
batch_size=config.embeddings_openai_batch_size,
|
||||
dimensions=config.embeddings_openai_dimensions,
|
||||
query_prefix=query_prefix,
|
||||
passage_prefix=passage_prefix,
|
||||
)
|
||||
elif provider == "requesty":
|
||||
api_key = config.embeddings_requesty_api_key
|
||||
@@ -1695,6 +1751,8 @@ def create_embeddings_from_env() -> Embeddings:
|
||||
base_url="https://router.requesty.ai/v1",
|
||||
batch_size=config.embeddings_openai_batch_size,
|
||||
dimensions=config.embeddings_openai_dimensions,
|
||||
query_prefix=query_prefix,
|
||||
passage_prefix=passage_prefix,
|
||||
)
|
||||
elif provider == "zeroentropy":
|
||||
api_key = config.embeddings_zeroentropy_api_key
|
||||
@@ -1727,6 +1785,8 @@ def create_embeddings_from_env() -> Embeddings:
|
||||
api_base=config.embeddings_litellm_api_base,
|
||||
api_key=config.embeddings_litellm_api_key,
|
||||
model=config.embeddings_litellm_model,
|
||||
query_prefix=query_prefix,
|
||||
passage_prefix=passage_prefix,
|
||||
)
|
||||
elif provider == "litellm-sdk":
|
||||
return LiteLLMSDKEmbeddings(
|
||||
@@ -1735,6 +1795,8 @@ def create_embeddings_from_env() -> Embeddings:
|
||||
api_base=config.embeddings_litellm_sdk_api_base,
|
||||
output_dimensions=config.embeddings_litellm_sdk_output_dimensions,
|
||||
encoding_format=config.embeddings_litellm_sdk_encoding_format,
|
||||
query_prefix=query_prefix,
|
||||
passage_prefix=passage_prefix,
|
||||
)
|
||||
elif provider == "google":
|
||||
vertexai_project_id = config.embeddings_vertexai_project_id
|
||||
|
||||
@@ -45,6 +45,9 @@ class _EntityToCreate:
|
||||
# Also stored on the row as entities.entity_kind so label rows stay out of the
|
||||
# partial trigram index (#3208).
|
||||
is_label: bool = False
|
||||
# False when the caller wrote this name literally: it is created as spelled and never
|
||||
# merged with a same-batch near-duplicate (#3479).
|
||||
resolve: bool = True
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -428,6 +431,18 @@ class EntityResolver:
|
||||
unit_event_date: When this unit was created
|
||||
conn: Optional connection to use (if None, acquires from pool)
|
||||
|
||||
Each mention may carry ``"resolve": False`` to opt out of resolution. The
|
||||
default, True, treats a name as a *guess* at which entity is meant, so
|
||||
similar existing entities are scored on name similarity + co-occurrence +
|
||||
recency and the best above threshold is reused. False takes the name
|
||||
literally: an existing entity is reused only when its canonical name
|
||||
matches case-insensitively, any other name creates its own entity, and it
|
||||
is never merged with a same-batch near-duplicate. Callers who authored the
|
||||
names deliberately want False (#3479) — resolution would otherwise let
|
||||
what the graph already believes outscore, and silently discard, their
|
||||
correction. It is per mention because retain resolves caller-supplied and
|
||||
extracted names in one batch, and only the caller's half is authoritative.
|
||||
|
||||
Returns:
|
||||
Resolved entity identities (id + stored canonical name) in the same
|
||||
order as input.
|
||||
@@ -457,6 +472,25 @@ class EntityResolver:
|
||||
taxonomy_lookup: set[str] | None = None,
|
||||
labels_cfg=None,
|
||||
) -> list[ResolvedEntity]:
|
||||
# `entities_data and` matters: an empty batch must fall through to the normal strategy
|
||||
# dispatch (which the pg_trgm auto-detection hangs off), not take the shortcut vacuously.
|
||||
if entities_data and not any(e.get("resolve", True) for e in entities_data):
|
||||
# Nothing in this batch resolves, so the trigram/UTL_MATCH probe and the
|
||||
# co-occurrence fetch would both be dead work. _resolve_from_candidates routes every
|
||||
# mention straight to its find-or-create path, which matches on LOWER(canonical_name)
|
||||
# equality. A *mixed* batch still probes — the per-mention check below skips the
|
||||
# literal names when scoring, which costs a little wasted lookup but keeps the
|
||||
# common all-resolving case on one code path.
|
||||
return await self._resolve_from_candidates(
|
||||
conn,
|
||||
bank_id,
|
||||
entities_data,
|
||||
unit_event_date,
|
||||
all_candidates={},
|
||||
cooccurrence_map={},
|
||||
taxonomy_lookup=taxonomy_lookup,
|
||||
labels_cfg=labels_cfg,
|
||||
)
|
||||
if self.entity_lookup == "trigram":
|
||||
# Route to backend-specific fuzzy strategy.
|
||||
# Non-PG backends (Oracle) use UTL_MATCH instead of pg_trgm.
|
||||
@@ -872,7 +906,7 @@ class EntityResolver:
|
||||
rep_by_lower: dict[str, str] = {}
|
||||
count_by_lower: dict[str, int] = {}
|
||||
for e in entities_to_create:
|
||||
if e.is_label:
|
||||
if e.is_label or not e.resolve:
|
||||
continue
|
||||
name_lower = e.name.lower()
|
||||
rep_by_lower.setdefault(name_lower, e.name)
|
||||
@@ -904,7 +938,12 @@ class EntityResolver:
|
||||
taxonomy_lookup: set[str] | None = None,
|
||||
labels_cfg=None,
|
||||
) -> list[ResolvedEntity]:
|
||||
"""Shared scoring + upsert logic used by both lookup strategies."""
|
||||
"""Shared scoring + upsert logic used by every lookup strategy.
|
||||
|
||||
A mention carrying ``"resolve": False`` skips the scoring entirely and takes the
|
||||
find-or-create path below, which matches an existing row on ``LOWER(canonical_name)``
|
||||
equality and inserts one otherwise.
|
||||
"""
|
||||
|
||||
# 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
|
||||
@@ -923,6 +962,9 @@ class EntityResolver:
|
||||
# Use per-entity date if available, otherwise fall back to batch-level date
|
||||
entity_event_date = entity_data.get("event_date", unit_event_date)
|
||||
|
||||
# Per mention, not per batch: retain resolves the caller's entities and the
|
||||
# extractor's in one pass, and only the caller's are meant literally (#3479).
|
||||
resolve = entity_data.get("resolve", True)
|
||||
candidates = all_candidates.get(entity_text, [])
|
||||
|
||||
# Backstop truncation for candidate sets that were not capped at the
|
||||
@@ -951,10 +993,14 @@ class EntityResolver:
|
||||
# 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
|
||||
if not resolve or not candidates:
|
||||
# Nothing to score against — or the caller named the entity literally, so
|
||||
# similarity must not get a vote. Either way the find-or-create pass below
|
||||
# reuses an identically-named row and otherwise inserts this exact name.
|
||||
entities_to_create.append(
|
||||
_EntityToCreate(idx=idx, name=entity_text, event_date=entity_event_date, is_label=is_label)
|
||||
_EntityToCreate(
|
||||
idx=idx, name=entity_text, event_date=entity_event_date, is_label=is_label, resolve=resolve
|
||||
)
|
||||
)
|
||||
continue
|
||||
|
||||
@@ -1057,7 +1103,9 @@ class EntityResolver:
|
||||
# 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.
|
||||
# excluded and keep exact grouping, and so are names the caller wrote literally:
|
||||
# "Alice" and "Alice Smith" listed side by side are two entities because they were
|
||||
# written as two (#3479).
|
||||
canonical_by_member = self._intrabatch_canonical_map(entities_to_create)
|
||||
|
||||
@dataclass
|
||||
@@ -1242,8 +1290,38 @@ class EntityResolver:
|
||||
else:
|
||||
return await self._link_units_to_entities_batch_impl(conn, normalized, bank_id)
|
||||
|
||||
async def record_unit_entity_postings(
|
||||
self,
|
||||
unit_entity_pairs: list[tuple[str, str]] | list[tuple[str, str, datetime | None]],
|
||||
bank_id: str | None = None,
|
||||
txn=None,
|
||||
):
|
||||
"""Store-owned variant of :meth:`link_units_to_entities_batch` that touches NO
|
||||
Postgres connection.
|
||||
|
||||
For a memories store that OWNS its memory rows (an external backend), the unit→entity
|
||||
posting is recorded by the store — ``record_unit_entities`` ignores the ``conn`` — and
|
||||
the co-occurrence update only accumulates in memory for the post-transaction flush.
|
||||
Neither needs a database transaction, so the retain orchestrator can run the posting in
|
||||
its connection-free store phase and never hold the data-plane connection across the
|
||||
object-store write. NOT for the Postgres store, whose posting is a real ``unit_entities``
|
||||
INSERT that requires the connection.
|
||||
|
||||
``txn`` is the caller's write-group handle. For a store that keeps the posting on the
|
||||
memory, this re-writes rows the same write-group just created, so it belongs to that
|
||||
group — see :meth:`MemoriesExtension.record_unit_entities`.
|
||||
"""
|
||||
if not unit_entity_pairs:
|
||||
return
|
||||
|
||||
normalized: list[tuple[str, str, datetime | None]] = [
|
||||
(t[0], t[1], t[2] if len(t) >= 3 else None) # type: ignore[misc]
|
||||
for t in unit_entity_pairs
|
||||
]
|
||||
return await self._link_units_to_entities_batch_impl(None, normalized, bank_id, txn=txn)
|
||||
|
||||
async def _link_units_to_entities_batch_impl(
|
||||
self, conn, unit_entity_pairs: list[tuple[str, str, datetime | None]], bank_id: str | None = None
|
||||
self, conn, unit_entity_pairs: list[tuple[str, str, datetime | None]], bank_id: str | None = None, txn=None
|
||||
):
|
||||
# Sorted bulk insert to prevent deadlocks from inconsistent lock ordering
|
||||
# across concurrent transactions on the unit_entities unique index.
|
||||
@@ -1264,6 +1342,7 @@ class EntityResolver:
|
||||
bank_id=bank_id,
|
||||
unit_ids=unit_ids,
|
||||
entity_ids=entity_ids,
|
||||
txn=txn,
|
||||
)
|
||||
|
||||
# Build maps keyed by unit_id:
|
||||
|
||||
@@ -1,29 +1,33 @@
|
||||
"""Async graph maintenance after document/unit deletes.
|
||||
|
||||
Three reconciliation passes run together on every worker invocation:
|
||||
Two queue-driven 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 and insert the missing links.
|
||||
|
||||
2. **Orphan entity prune.** Delete ``entities`` rows in the bank that no
|
||||
longer have any live memory references. FK ON DELETE CASCADE on
|
||||
``entity_cooccurrences`` then removes any cooccurrence row pointing
|
||||
at the pruned entities.
|
||||
2. **Entity prune.** Drain ``entity_maintenance_queue`` (entities a delete
|
||||
may have stranded). Per batch: delete the candidates no ``unit_entities``
|
||||
row references any more — FK ON DELETE CASCADE on ``entity_cooccurrences``
|
||||
takes their cooccurrence rows with them — then delete the cooccurrence rows
|
||||
incident to the survivors that no current memory witnesses, the stale-count
|
||||
case the cascade cannot see.
|
||||
|
||||
3. **Stale cooccurrence prune.** Defensive sweep for cooccurrence rows
|
||||
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.
|
||||
Both passes are *queued work*, not sweeps. Pass 2 used to be two bank-wide
|
||||
statements re-evaluated on every invocation whether or not anything had
|
||||
changed, so its cost tracked the size of the bank instead of the size of the
|
||||
delete; on a multi-million-row bank neither statement could finish inside
|
||||
asyncpg's command timeout and the job failed on every run, forever (#3222).
|
||||
Both queues are now filled inside the deleting transaction, so each run only
|
||||
looks at what that delete actually touched.
|
||||
|
||||
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.
|
||||
out. This module orchestrates them (pass ordering, the time budget, the timing
|
||||
log) 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 both passes are no-ops for it.
|
||||
|
||||
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
|
||||
@@ -32,9 +36,9 @@ 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.
|
||||
concurrent runs would do no extra work anyway — each drains the same two
|
||||
bank-scoped queues — while convoying on each other's row locks and holding a
|
||||
worker slot each.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -60,19 +64,14 @@ if TYPE_CHECKING:
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# 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
|
||||
|
||||
|
||||
@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
|
||||
# Wall-clock budget for one graph_maintenance run. Both passes commit per batch,
|
||||
# so hitting the budget is not a failure: it stops claiming new work, reports
|
||||
# what it did, and the follow-up run resumes from the queue rows still there.
|
||||
# A backlog (a bulk delete, say) then converges over several runs instead of
|
||||
# holding a worker slot for as long as it takes — the failure mode #3222
|
||||
# describes, where the whole run was cancelled and every batch's work was
|
||||
# retried from scratch.
|
||||
_JOB_TIME_BUDGET_SECONDS = 240.0
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -81,15 +80,22 @@ class JobResult:
|
||||
|
||||
relink_units_processed: int = 0
|
||||
relink_links_added: int = 0
|
||||
entities_examined: int = 0
|
||||
orphan_entities_pruned: int = 0
|
||||
stale_cooccurrences_pruned: int = 0
|
||||
# False when the time budget stopped a drain with work still queued. The
|
||||
# caller re-submits so the backlog keeps draining without waiting for the
|
||||
# next delete to trigger a run.
|
||||
queues_drained: bool = True
|
||||
|
||||
def as_dict(self) -> dict[str, int]:
|
||||
def as_dict(self) -> dict[str, int | bool]:
|
||||
return {
|
||||
"relink_units_processed": self.relink_units_processed,
|
||||
"relink_links_added": self.relink_links_added,
|
||||
"entities_examined": self.entities_examined,
|
||||
"orphan_entities_pruned": self.orphan_entities_pruned,
|
||||
"stale_cooccurrences_pruned": self.stale_cooccurrences_pruned,
|
||||
"queues_drained": self.queues_drained,
|
||||
}
|
||||
|
||||
|
||||
@@ -143,19 +149,65 @@ async def enqueue_relink_victims(
|
||||
)
|
||||
|
||||
|
||||
async def enqueue_entity_prune_candidates(
|
||||
conn: DatabaseConnection,
|
||||
bank_id: str,
|
||||
affected_unit_ids: list[str],
|
||||
) -> int:
|
||||
"""Enqueue the entities ``affected_unit_ids`` reference as prune candidates.
|
||||
|
||||
Must run inside the same transaction that removes those units (or replaces
|
||||
their entity postings), *before* the delete or cascade fires: afterwards the
|
||||
``unit_entities`` rows naming the entities are gone, and an entity nothing
|
||||
points at is an orphan nothing will ever look at again.
|
||||
|
||||
Pair this with :func:`enqueue_relink_victims` at every delete site. They
|
||||
capture different things — that one records the *survivors* whose links now
|
||||
dangle, this one the *entities* the doomed units were holding up — and
|
||||
neither substitutes for the other. A site that deletes units without calling
|
||||
this leaks orphan entities and stale cooccurrences until something else
|
||||
happens to enqueue the same entity.
|
||||
|
||||
Over-enqueueing costs nothing: the drain re-checks each candidate and keeps
|
||||
the ones still referenced.
|
||||
|
||||
Delegated to the memories store: a store that never wrote ``unit_entities``
|
||||
has no postings to lose and returns 0.
|
||||
|
||||
Args:
|
||||
conn: Database connection inside the active transaction.
|
||||
bank_id: Bank owning the affected units.
|
||||
affected_unit_ids: Memory_unit IDs whose entity postings are about to
|
||||
be (or are being) removed.
|
||||
|
||||
Returns:
|
||||
Number of candidate entities enqueued (0 for a store with no postings).
|
||||
"""
|
||||
if not affected_unit_ids:
|
||||
return 0
|
||||
|
||||
from .memories import get_memories
|
||||
|
||||
return await get_memories().enqueue_entity_prune_candidates(
|
||||
conn=conn,
|
||||
fq_table=fq_table,
|
||||
bank_id=bank_id,
|
||||
affected_unit_ids=affected_unit_ids,
|
||||
)
|
||||
|
||||
|
||||
async def run_graph_maintenance_job(
|
||||
memory_engine: "MemoryEngine",
|
||||
bank_id: str,
|
||||
request_context: RequestContext,
|
||||
operation_id: str | None = None,
|
||||
) -> dict[str, int]:
|
||||
"""Run all maintenance passes for ``bank_id`` until the relink queue is
|
||||
drained, then sweep entities and cooccurrences once.
|
||||
) -> dict[str, int | bool]:
|
||||
"""Drain both maintenance queues for ``bank_id``, within a time budget.
|
||||
|
||||
Returns:
|
||||
Per-pass counters from :class:`JobResult`.
|
||||
Per-pass counters from :class:`JobResult`. ``queues_drained`` is False
|
||||
when the budget ran out with work still queued — the caller re-submits.
|
||||
"""
|
||||
del request_context # accepted for symmetry with other run_*_job helpers
|
||||
from ..config import get_config
|
||||
from .memories import get_memories
|
||||
|
||||
@@ -165,6 +217,7 @@ async def run_graph_maintenance_job(
|
||||
|
||||
result = JobResult()
|
||||
job_start = time.time()
|
||||
deadline = time.monotonic() + _JOB_TIME_BUDGET_SECONDS
|
||||
|
||||
# --- Pass 1: relink ---
|
||||
# The store owns the whole drain loop: it is a claim → top-up → commit over
|
||||
@@ -173,51 +226,119 @@ async def run_graph_maintenance_job(
|
||||
# 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)
|
||||
relink = await store.relink_pass(
|
||||
backend=backend, fq_table=fq_table, bank_id=bank_id, config=config, deadline=deadline
|
||||
)
|
||||
result.relink_units_processed = relink.units_processed
|
||||
result.relink_links_added = relink.links_added
|
||||
|
||||
# --- 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 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
|
||||
# --- Pass 2: entity prune ---
|
||||
# Same shape as Pass 1 and owned by the store for the same reason: a
|
||||
# claim → prune → commit loop over `entities` / `unit_entities` /
|
||||
# `entity_cooccurrences`, including the ordered locking that keeps its
|
||||
# deletes from cycling against retain's concurrent entity and cooccurrence
|
||||
# upserts. A store that never wrote `unit_entities` returns an empty dict
|
||||
# and this is a no-op. Runs after the relink pass so the remaining budget
|
||||
# is whatever Pass 1 left.
|
||||
prune = await store.entity_prune_pass(backend=backend, fq_table=fq_table, bank_id=bank_id, deadline=deadline)
|
||||
result.entities_examined = prune.entities_examined
|
||||
result.orphan_entities_pruned = prune.orphan_entities_pruned
|
||||
result.stale_cooccurrences_pruned = prune.stale_cooccurrences_pruned
|
||||
result.queues_drained = relink.queue_exhausted and prune.queue_exhausted
|
||||
|
||||
elapsed = time.time() - job_start
|
||||
|
||||
# --- Hand-off: schedule a successor for any work this run leaves behind ---
|
||||
#
|
||||
# Submit-time dedup now treats a *running* graph-maintenance job as covering
|
||||
# the bank (see _submit_async_operation's dedupe_by_bank_includes_processing).
|
||||
# That is what stops one job being queued per triggering operation, but it
|
||||
# means a submit made while this job runs is suppressed. So this job has to
|
||||
# hand off to a successor for any work it leaves behind, or that work strands
|
||||
# until some unrelated future trigger. Both hand-offs below pass
|
||||
# dedupe_excludes_operation_id: the worker only marks the operation completed
|
||||
# after this body returns, so the row is still 'processing' now and the
|
||||
# widened predicate would otherwise dedup the hand-off against its own row and
|
||||
# silently do nothing.
|
||||
from .memory_engine import acquire_with_retry
|
||||
from .task_backend import SyncTaskBackend
|
||||
|
||||
if not result.queues_drained:
|
||||
# Backlog case: the time budget stopped a drain with work still queued, so
|
||||
# more is provably left. Chain a follow-up so the backlog converges
|
||||
# without waiting for the next delete to trigger a run — on a bank that
|
||||
# has gone quiet that may be never. WARNING because a bank that keeps
|
||||
# landing here is producing maintenance faster than one run absorbs it.
|
||||
logger.warning(
|
||||
f"[GRAPH_MAINT] bank={bank_id} hit the {_JOB_TIME_BUDGET_SECONDS:.0f}s budget with work still "
|
||||
f"queued; committed {result.as_dict()} in {elapsed:.2f}s"
|
||||
)
|
||||
# A synchronous task backend (tests, embedded) runs the successor inline,
|
||||
# which would recurse one job per budget window instead of scheduling.
|
||||
# There the caller is already the drain loop and gets the remaining rows
|
||||
# on its next call, so skip the hand-off.
|
||||
if not isinstance(memory_engine._task_backend, SyncTaskBackend):
|
||||
try:
|
||||
await memory_engine.submit_async_graph_maintenance(
|
||||
bank_id=bank_id,
|
||||
request_context=request_context,
|
||||
dedupe_excludes_operation_id=operation_id,
|
||||
)
|
||||
except Exception:
|
||||
# Never fail a completed maintenance run over the hand-off. The
|
||||
# work is still queued and the next trigger picks it up; log
|
||||
# loudly so a persistent failure here is visible, not silent.
|
||||
logger.exception(f"[GRAPH_MAINT] bank={bank_id} follow-up submit failed")
|
||||
else:
|
||||
# Gap case: both queues drained within budget, but new rows can have
|
||||
# landed in the gap between a pass's final claim and this job being marked
|
||||
# completed. Their submits were deduped against this still-'processing'
|
||||
# job, so nothing is scheduled to pick them up. Re-check both queues —
|
||||
# reusing the portable existence check submit uses for its empty-queue
|
||||
# short-circuit (no Postgres-only LIMIT, and covers the relink and
|
||||
# entity-prune queues) — and hand off anything that landed.
|
||||
#
|
||||
# Gated on this run having made progress. A run that consumed nothing and
|
||||
# still sees queued work would hand off to a successor that repeats the
|
||||
# exact outcome — an endless per-bank chain. Requiring progress means the
|
||||
# chain only continues while it is actually draining, so it terminates.
|
||||
# (The backlog branch above is not gated this way: its contract is to
|
||||
# always continue a budgeted backlog so a quiet bank is never stranded.)
|
||||
#
|
||||
# Not guarded against SyncTaskBackend, unlike the backlog branch: this
|
||||
# branch cannot fire on one. A synchronous backend is single-threaded, so
|
||||
# nothing enqueues concurrently and the queues are empty once the passes
|
||||
# (which never enqueue for themselves) return — leaving no gap to close.
|
||||
made_progress = result.relink_units_processed > 0 or result.entities_examined > 0
|
||||
try:
|
||||
backend_check = await memory_engine._get_backend()
|
||||
async with acquire_with_retry(backend_check) as conn:
|
||||
work_remains = bool(
|
||||
await conn.fetchval(
|
||||
f"""
|
||||
SELECT 1 WHERE
|
||||
EXISTS (SELECT 1 FROM {fq_table("graph_maintenance_queue")} WHERE bank_id = $1)
|
||||
OR EXISTS (SELECT 1 FROM {fq_table("entity_maintenance_queue")} WHERE bank_id = $1)
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
)
|
||||
if work_remains and not made_progress:
|
||||
logger.warning(
|
||||
f"[GRAPH_MAINT] bank={bank_id} queue still non-empty after a run that drained "
|
||||
f"nothing; not chaining a successor (it would repeat this outcome)"
|
||||
)
|
||||
elif work_remains:
|
||||
logger.info(f"[GRAPH_MAINT] bank={bank_id} work arrived during the run; submitting a follow-up job")
|
||||
await memory_engine.submit_async_graph_maintenance(
|
||||
bank_id=bank_id,
|
||||
request_context=request_context,
|
||||
dedupe_excludes_operation_id=operation_id,
|
||||
)
|
||||
except Exception:
|
||||
# As above: the queued work survives, so log rather than fail the run.
|
||||
logger.exception(f"[GRAPH_MAINT] bank={bank_id} follow-up submit failed")
|
||||
|
||||
logger.info(
|
||||
f"[GRAPH_MAINT] bank={bank_id} done: {result.as_dict()}, elapsed={elapsed:.2f}s, operation_id={operation_id}"
|
||||
)
|
||||
|
||||
@@ -160,16 +160,22 @@ class MemoryEngineInterface(ABC):
|
||||
async def list_banks(
|
||||
self,
|
||||
*,
|
||||
search_query: str | None = None,
|
||||
limit: int = 100,
|
||||
offset: int = 0,
|
||||
request_context: "RequestContext",
|
||||
) -> list[dict[str, Any]]:
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
List all memory banks.
|
||||
List memory banks, one page at a time.
|
||||
|
||||
Args:
|
||||
search_query: Case-insensitive substring matched against bank ID and name.
|
||||
limit: Maximum number of banks to return (0 returns none).
|
||||
offset: Number of banks to skip.
|
||||
request_context: Request context for authentication.
|
||||
|
||||
Returns:
|
||||
List of bank info dicts.
|
||||
Dict with ``banks`` (the page), ``total``, ``limit`` and ``offset``.
|
||||
"""
|
||||
...
|
||||
|
||||
@@ -322,7 +328,7 @@ class MemoryEngineInterface(ABC):
|
||||
self,
|
||||
bank_id: str,
|
||||
*,
|
||||
fact_type: str | None = None,
|
||||
fact_type: str | list[str] | None = None,
|
||||
search_query: str | None = None,
|
||||
entity_id: str | None = None,
|
||||
created_before: datetime | None = None,
|
||||
@@ -335,7 +341,8 @@ class MemoryEngineInterface(ABC):
|
||||
|
||||
Args:
|
||||
bank_id: The memory bank ID.
|
||||
fact_type: Filter by fact type.
|
||||
fact_type: Filter by fact type. A list matches any of them; an empty
|
||||
list is treated as no filter.
|
||||
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.
|
||||
|
||||
@@ -5,6 +5,7 @@ This module defines the interface that all LLM providers must implement,
|
||||
enabling support for multiple LLM backends (OpenAI, Anthropic, Gemini, Codex, etc.)
|
||||
"""
|
||||
|
||||
import logging
|
||||
from abc import ABC, abstractmethod
|
||||
from contextlib import AbstractAsyncContextManager
|
||||
from dataclasses import dataclass
|
||||
@@ -14,6 +15,8 @@ from typing import Any, Callable, Self
|
||||
|
||||
from .response_models import LLMToolCallResult
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class LLMToolChoiceMode(StrEnum):
|
||||
"""Canonical tool-selection modes shared by every LLM provider."""
|
||||
@@ -68,7 +71,7 @@ class LLMInterface(ABC):
|
||||
api_key: str,
|
||||
base_url: str,
|
||||
model: str,
|
||||
reasoning_effort: str = "low",
|
||||
reasoning_effort: str | None = None,
|
||||
**kwargs: Any,
|
||||
):
|
||||
"""
|
||||
@@ -79,14 +82,37 @@ class LLMInterface(ABC):
|
||||
api_key: API key or authentication token.
|
||||
base_url: Base URL for the API.
|
||||
model: Model name.
|
||||
reasoning_effort: Reasoning effort level for supported providers.
|
||||
reasoning_effort: Reasoning effort level, or None when the operator
|
||||
configured none — in which case no provider sends the parameter and
|
||||
every model runs at its own default effort.
|
||||
**kwargs: Additional provider-specific parameters.
|
||||
"""
|
||||
self.provider = provider.lower()
|
||||
self.api_key = api_key
|
||||
self.base_url = base_url
|
||||
self.model = model
|
||||
self.reasoning_effort = reasoning_effort
|
||||
# None means "the operator said nothing", and nothing is what gets sent: no
|
||||
# provider may invent a level. Hindsight used to resolve unset to "low" here and
|
||||
# ship it to whichever lanes their capability check happened to accept, which
|
||||
# made the setting both invisible (a configured value could be silently dropped —
|
||||
# issue #3449) and presumptuous (an unconfigured one was still transmitted).
|
||||
# An empty string is an unset environment variable, not a level.
|
||||
self.reasoning_effort: str | None = reasoning_effort or None
|
||||
|
||||
def _warn_reasoning_effort_unsupported(self) -> None:
|
||||
"""Report, once at startup, that this provider cannot honour a configured effort.
|
||||
|
||||
Providers with no reasoning knob to turn call this from ``__init__``. Silence is
|
||||
what made issue #3449 expensive: the variable is set, documented and visible in
|
||||
the environment, so every signal the operator has says it is in force. A setting
|
||||
this provider cannot act on has to say so out loud.
|
||||
"""
|
||||
if self.reasoning_effort is None:
|
||||
return
|
||||
logger.warning(
|
||||
f"reasoning_effort={self.reasoning_effort!r} is ignored: the {self.provider} provider "
|
||||
f"has no reasoning-effort control. Remove the setting or switch provider to apply it."
|
||||
)
|
||||
|
||||
@abstractmethod
|
||||
async def verify_connection(self) -> None:
|
||||
|
||||
@@ -297,7 +297,7 @@ def create_llm_provider(
|
||||
api_key: str,
|
||||
base_url: str,
|
||||
model: str,
|
||||
reasoning_effort: str,
|
||||
reasoning_effort: str | None,
|
||||
groq_service_tier: str | None = None,
|
||||
openai_service_tier: str | None = None,
|
||||
bedrock_service_tier: str | None = None,
|
||||
@@ -323,7 +323,9 @@ def create_llm_provider(
|
||||
api_key: API key (may be None for local providers or OAuth providers).
|
||||
base_url: Base URL for the API.
|
||||
model: Model name.
|
||||
reasoning_effort: Reasoning effort level for supported providers.
|
||||
reasoning_effort: Reasoning effort level for supported providers, or None when
|
||||
the operator configured none (providers then fall back to the default level
|
||||
and may skip the parameter entirely).
|
||||
groq_service_tier: Groq service tier (for Groq provider) - "on_demand", "flex", or "auto".
|
||||
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".
|
||||
@@ -513,6 +515,7 @@ def create_llm_provider(
|
||||
base_url=base_url,
|
||||
model=model,
|
||||
reasoning_effort=reasoning_effort,
|
||||
extra_body=extra_body,
|
||||
model_path=config.llamacpp_model_path,
|
||||
gpu_layers=config.llamacpp_gpu_layers,
|
||||
context_size=config.llamacpp_context_size,
|
||||
@@ -636,7 +639,7 @@ class LLMProvider:
|
||||
api_key: str,
|
||||
base_url: str,
|
||||
model: str,
|
||||
reasoning_effort: str = "low",
|
||||
reasoning_effort: str | None = None,
|
||||
groq_service_tier: str | None = None,
|
||||
openai_service_tier: str | None = None,
|
||||
bedrock_service_tier: str | None = None,
|
||||
@@ -665,7 +668,8 @@ class LLMProvider:
|
||||
api_key: API key.
|
||||
base_url: Base URL for the API.
|
||||
model: Model name.
|
||||
reasoning_effort: Reasoning effort level for supported providers.
|
||||
reasoning_effort: Reasoning effort level for supported providers, or None
|
||||
when the operator configured none.
|
||||
groq_service_tier: Groq service tier ("on_demand", "flex", "auto") - from config.
|
||||
openai_service_tier: OpenAI service tier (None or "flex") - from config.
|
||||
bedrock_service_tier: Bedrock service tier (None, "flex", "priority", "reserved") - from config.
|
||||
@@ -1436,7 +1440,6 @@ class LLMProvider:
|
||||
DEFAULT_LLM_OPENAI_SERVICE_TIER,
|
||||
DEFAULT_LLM_PROMPT_CACHE_ENABLED,
|
||||
DEFAULT_LLM_PROVIDER,
|
||||
DEFAULT_LLM_REASONING_EFFORT,
|
||||
DEFAULT_LLM_STRUCTURED_OUTPUT_FORCED_TOOL,
|
||||
DEFAULT_LLM_TIMEOUT,
|
||||
ENV_LLM_API_KEY,
|
||||
@@ -1498,7 +1501,7 @@ class LLMProvider:
|
||||
api_key=api_key,
|
||||
base_url=base_url,
|
||||
model=model,
|
||||
reasoning_effort=os.getenv(ENV_LLM_REASONING_EFFORT, DEFAULT_LLM_REASONING_EFFORT),
|
||||
reasoning_effort=os.getenv(ENV_LLM_REASONING_EFFORT) or None,
|
||||
extra_body=extra_body,
|
||||
default_headers=default_headers,
|
||||
cache_affinity=cache_affinity,
|
||||
|
||||
@@ -3,15 +3,16 @@
|
||||
A single periodic loop that drives all of Hindsight's recurring housekeeping
|
||||
from one place, so we don't spawn a separate ``asyncio`` task per concern:
|
||||
|
||||
- **Retention sweeps** (hourly): delete ``audit_log`` and ``llm_requests`` rows
|
||||
older than their configured retention, across *all* tenant schemas.
|
||||
- **Retention sweeps** (configurable, default hourly): delete ``audit_log`` and
|
||||
``llm_requests`` rows older than their configured retention, across *all*
|
||||
tenant schemas.
|
||||
- **Consolidation reconcile** (configurable, default 5 min): re-schedule
|
||||
consolidation for banks that have eligible-but-unscheduled facts and no
|
||||
in-flight consolidation. This recovers facts that were stranded when a
|
||||
consolidation operation failed terminally and left them with
|
||||
``consolidated_at IS NULL AND consolidation_failed_at IS NULL`` and nothing to
|
||||
re-trigger them.
|
||||
- **Scheduled mental model refresh** (configurable check cadence, default 60s):
|
||||
- **Scheduled mental model refresh** (configurable check cadence, default 5 min):
|
||||
refresh mental models whose ``trigger.refresh_cron`` schedule is due, but only
|
||||
when the model is stale (new memories in its scope since its last refresh), so
|
||||
a scheduled tick never burns an LLM call to regenerate identical content. The
|
||||
@@ -28,15 +29,34 @@ 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``).
|
||||
process. Operation cleanup deletes a bounded batch per schema; the consolidation
|
||||
reconcile and the scheduled mental model refresh both dedupe against in-flight
|
||||
operations inside the inserting transaction (see ``_submit_async_operation``);
|
||||
retention deletes in bounded chunks claimed with SKIP LOCKED, so concurrent
|
||||
sweepers split the work instead of colliding (see ``_purge_table_in_batches``).
|
||||
|
||||
The *work* is therefore safe to run everywhere. The *discovery* in front of it is
|
||||
not free: one round-trip on the wire is still one query per tenant schema inside
|
||||
the routine, every process pays it, and its cost scales with tenant count while
|
||||
the work it finds does not. Two things keep that proportionate, and both are
|
||||
load-bearing rather than incidental:
|
||||
|
||||
- **Cadence is config, not a constant.** Every job's interval is a server-level
|
||||
setting. The jobs that delete rows whose retention is measured in *days*
|
||||
(retention, operation cleanup) have no reason to probe every schema every
|
||||
minute.
|
||||
- **The first tick is jittered** (``maintenance_start_jitter_seconds``). Every job
|
||||
is due on the first tick, so a fleet started together — a deploy, a rolling
|
||||
restart — would fire every cross-tenant probe in every process at the same
|
||||
instant. SKIP LOCKED keeps that correct but not cheap: the probes are reads, and
|
||||
N of them land at once. Steady state self-staggers; startup does not.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import random
|
||||
import time
|
||||
from collections.abc import Coroutine
|
||||
from datetime import datetime, timedelta, timezone
|
||||
@@ -54,12 +74,6 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
# Short tick so jobs with different cadences share one loop without per-job tasks.
|
||||
_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
|
||||
@@ -69,6 +83,32 @@ _TXN_RECOVERY_INTERVAL_SECONDS = 300
|
||||
# unwitnessed one — the writer may still be mid-flight (PendingTxn carries no timestamp).
|
||||
_TXN_RECOVERY_GRACE_SECONDS = 300
|
||||
|
||||
# ── retention sweep pacing ────────────────────────────────────────────────────
|
||||
# Retention used to issue one unbounded `DELETE FROM <table> WHERE started_at <
|
||||
# cutoff` per schema. On a table with a real backlog that is a single statement
|
||||
# holding row locks for minutes while it reads the whole expired range — and the
|
||||
# maintenance loop runs in every API/worker process with no leader election, so
|
||||
# every pod issued it at the same hourly boundary. Observed as two concurrent
|
||||
# 330s+ deletes pinned on IO.DataFileRead, blocking each other on row locks,
|
||||
# saturating RDS I/O and tripling recall latency.
|
||||
#
|
||||
# The fix is to design the collision out rather than elect one sweeper: each chunk
|
||||
# claims its rows with FOR UPDATE SKIP LOCKED, so concurrent sweepers take
|
||||
# *disjoint* chunks instead of waiting on each other, and the total work stays the
|
||||
# number of expired rows however many pods join in. Chunks are short, index-driven
|
||||
# transactions with a pause between them, so no statement holds locks for long and
|
||||
# the deletes never monopolise disk I/O.
|
||||
_RETENTION_BATCH_SIZE = 2000
|
||||
# Ceiling on chunks per table per schema per run: a backstop against looping
|
||||
# forever on a table that is filling faster than it drains. A full run therefore
|
||||
# removes at most 2M rows per schema, and the next sweep (whose cadence is
|
||||
# HINDSIGHT_API_RETENTION_SWEEP_INTERVAL_SECONDS, default hourly) continues.
|
||||
_RETENTION_MAX_BATCHES = 1000
|
||||
# Breather between chunks. Paces one process at ~8k rows/s worst case; several
|
||||
# pods sweeping at once multiply that, which is still orders of magnitude gentler
|
||||
# than the unbounded delete this replaces.
|
||||
_RETENTION_BATCH_PAUSE_SECONDS = 0.25
|
||||
|
||||
|
||||
class MaintenanceLoop:
|
||||
"""Owns the single periodic maintenance task for a :class:`MemoryEngine`."""
|
||||
@@ -122,10 +162,11 @@ class MaintenanceLoop:
|
||||
# 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
|
||||
sweep_on = cfg.retention_sweep_interval_seconds > 0
|
||||
audit_on = sweep_on and cfg.audit_log_retention_days > 0
|
||||
llm_on = sweep_on and cfg.llm_trace_enabled and cfg.llm_trace_retention_days > 0
|
||||
mm_refresh_on = cfg.mental_model_refresh_tick_seconds > 0
|
||||
op_cleanup_on = cfg.operation_retention_days > 0
|
||||
op_cleanup_on = cfg.operation_cleanup_interval_seconds > 0 and cfg.operation_retention_days > 0
|
||||
return (
|
||||
reconcile_on
|
||||
or audit_on
|
||||
@@ -154,6 +195,8 @@ class MaintenanceLoop:
|
||||
# ── loop ───────────────────────────────────────────────────────────────
|
||||
|
||||
async def _run(self) -> None:
|
||||
if not await self._wait_start_jitter():
|
||||
return
|
||||
while not self._stop.is_set():
|
||||
try:
|
||||
await self._tick()
|
||||
@@ -164,6 +207,26 @@ class MaintenanceLoop:
|
||||
except asyncio.TimeoutError:
|
||||
pass
|
||||
|
||||
async def _wait_start_jitter(self) -> bool:
|
||||
"""Delay the first tick by a random offset. Returns False if stopped while waiting.
|
||||
|
||||
Every job is due the first time ``_is_due`` sees it, so N processes started
|
||||
together would run all of them at the same instant — the one moment where
|
||||
redundant cross-tenant discovery and overlapping DELETEs actually collide.
|
||||
Spreading the *first* tick is enough: from then on each process keeps its
|
||||
own phase.
|
||||
"""
|
||||
jitter = get_config().maintenance_start_jitter_seconds
|
||||
if jitter <= 0:
|
||||
return True
|
||||
delay = random.uniform(0, jitter)
|
||||
logger.debug(f"Maintenance loop: delaying first tick by {delay:.1f}s")
|
||||
try:
|
||||
await asyncio.wait_for(self._stop.wait(), timeout=delay)
|
||||
except asyncio.TimeoutError:
|
||||
return True
|
||||
return False
|
||||
|
||||
def _is_due(self, job: str, interval_seconds: int) -> bool:
|
||||
"""True if ``job`` has never run or its interval has elapsed; marks it run now."""
|
||||
now = time.monotonic()
|
||||
@@ -175,7 +238,8 @@ class MaintenanceLoop:
|
||||
|
||||
async def _tick(self) -> None:
|
||||
cfg = get_config()
|
||||
if self._is_due("retention", _RETENTION_INTERVAL_SECONDS):
|
||||
retention_interval = cfg.retention_sweep_interval_seconds
|
||||
if retention_interval > 0 and self._is_due("retention", retention_interval):
|
||||
await self._run_timed("retention", self._run_retention(cfg))
|
||||
interval = cfg.consolidation_reconcile_interval_seconds
|
||||
if interval > 0 and self._is_due("reconcile", interval):
|
||||
@@ -183,7 +247,12 @@ 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):
|
||||
cleanup_interval = cfg.operation_cleanup_interval_seconds
|
||||
if (
|
||||
cleanup_interval > 0
|
||||
and cfg.operation_retention_days > 0
|
||||
and self._is_due("operation_cleanup", cleanup_interval)
|
||||
):
|
||||
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())
|
||||
@@ -213,26 +282,86 @@ class MaintenanceLoop:
|
||||
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)
|
||||
|
||||
async def _purge_expired(self, table: str, ts_col: str, days: int) -> None:
|
||||
"""Delete rows older than ``days`` from ``table`` across every tenant schema."""
|
||||
async def _purge_expired(self, table: str, ts_col: str, days: int) -> int:
|
||||
"""Delete rows older than ``days`` from ``table`` across every tenant schema.
|
||||
|
||||
Only for the retention tables (``audit_log``, ``llm_requests``): chunking
|
||||
deletes by primary key assumes the ``id`` column both of them carry.
|
||||
Returns the number of rows deleted by *this* process.
|
||||
"""
|
||||
backend = self._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_rows')}($1, $2, $3)", table, ts_col, days
|
||||
)
|
||||
for row in rows:
|
||||
schema = row[0]
|
||||
# schema names come from pg_class; quote defensively all the same.
|
||||
qschema = '"' + schema.replace('"', '""') + '"'
|
||||
result = await conn.execute(
|
||||
f"DELETE FROM {qschema}.{table} WHERE {ts_col} < NOW() - make_interval(days => $1)",
|
||||
days,
|
||||
)
|
||||
if result and result != "DELETE 0":
|
||||
logger.info(f"Retention sweep {schema}.{table}: {result}")
|
||||
except Exception as e:
|
||||
logger.warning(f"Retention sweep failed for {table}: {e}")
|
||||
logger.warning(f"Retention sweep discovery failed for {table}: {e}")
|
||||
return 0
|
||||
|
||||
# One cutoff for the whole sweep: a per-chunk NOW() would let the window
|
||||
# creep forward mid-run, which makes "deleted < batch means done" wrong.
|
||||
cutoff = datetime.now(timezone.utc) - timedelta(days=days)
|
||||
total = 0
|
||||
for row in rows:
|
||||
schema = row[0]
|
||||
# schema names come from pg_class; quote defensively all the same.
|
||||
qschema = '"' + schema.replace('"', '""') + '"'
|
||||
try:
|
||||
deleted = await self._purge_table_in_batches(f"{qschema}.{table}", ts_col, cutoff)
|
||||
except Exception as e:
|
||||
logger.warning(f"Retention sweep failed for {schema}.{table}: {e}")
|
||||
continue
|
||||
if deleted:
|
||||
total += deleted
|
||||
logger.info(f"Retention sweep {schema}.{table}: DELETE {deleted}")
|
||||
return total
|
||||
|
||||
async def _purge_table_in_batches(self, table: str, ts_col: str, cutoff: datetime) -> int:
|
||||
"""Delete expired rows from one qualified table in bounded chunks.
|
||||
|
||||
Rows are claimed oldest-first off the ``(started_at)`` index with FOR UPDATE
|
||||
SKIP LOCKED, which is what makes a leaderless fleet safe: a chunk never
|
||||
waits on another sweeper (or on a writer still finishing its own row), and
|
||||
two processes sweeping the same table take disjoint chunks rather than
|
||||
redoing each other's work. Each chunk commits on its own, so no transaction
|
||||
holds locks longer than one batch.
|
||||
"""
|
||||
backend = self._engine._backend
|
||||
deleted = 0
|
||||
for batch in range(_RETENTION_MAX_BATCHES):
|
||||
if self._stop.is_set():
|
||||
break
|
||||
if batch:
|
||||
await asyncio.sleep(_RETENTION_BATCH_PAUSE_SECONDS)
|
||||
async with acquire_with_retry(backend, max_retries=1) as conn, conn.transaction():
|
||||
removed = await conn.fetchval(
|
||||
f"""
|
||||
WITH expired AS (
|
||||
SELECT id FROM {table}
|
||||
WHERE {ts_col} < $1
|
||||
ORDER BY {ts_col}
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
), removed AS (
|
||||
DELETE FROM {table} t USING expired e WHERE t.id = e.id RETURNING 1
|
||||
)
|
||||
SELECT count(*) FROM removed
|
||||
""",
|
||||
cutoff,
|
||||
_RETENTION_BATCH_SIZE,
|
||||
)
|
||||
deleted += removed
|
||||
# A short chunk means the expired range is drained — or that another
|
||||
# sweeper holds the rest, which is equally a reason to stop.
|
||||
if removed < _RETENTION_BATCH_SIZE:
|
||||
break
|
||||
else:
|
||||
logger.warning(
|
||||
f"Retention sweep hit its per-run batch ceiling on {table} after {deleted} row(s); "
|
||||
"the remainder is left for the next run"
|
||||
)
|
||||
return deleted
|
||||
|
||||
# ── terminal operation cleanup ─────────────────────────────────────────
|
||||
|
||||
@@ -289,7 +418,12 @@ class MaintenanceLoop:
|
||||
async with acquire_with_retry(backend, max_retries=1) as conn:
|
||||
# Delete export archives owned by rows about to be pruned first,
|
||||
# so the file-storage blobs don't outlive their operation row.
|
||||
await engine.purge_expired_export_archives(conn, table, cutoff)
|
||||
# Same batch bound as the prune below: the two walk the same
|
||||
# ordered window so a backlog doesn't re-purge already-deleted
|
||||
# archives on every cycle.
|
||||
await engine.purge_expired_export_archives(
|
||||
conn, table, cutoff, batch_size=cfg.operation_cleanup_batch_size
|
||||
)
|
||||
async with conn.transaction():
|
||||
deleted = await backend.ops.prune_terminal_operations(
|
||||
conn, table, cutoff, batch_size=cfg.operation_cleanup_batch_size
|
||||
@@ -484,7 +618,8 @@ class MaintenanceLoop:
|
||||
# the row under the bank's schema context.
|
||||
async with acquire_with_retry(engine._backend, max_retries=1) as conn:
|
||||
mm_row = await conn.fetchrow(
|
||||
f"SELECT id, tags, trigger, last_refreshed_at FROM {fq_table('mental_models')} "
|
||||
f"SELECT id, tags, trigger, last_refreshed_at, last_memory_seen_at "
|
||||
f"FROM {fq_table('mental_models')} "
|
||||
"WHERE bank_id = $1 AND id = $2",
|
||||
bank_id,
|
||||
mm_id,
|
||||
|
||||
@@ -18,6 +18,7 @@ from .base import (
|
||||
FactRecord,
|
||||
MemoriesExtension,
|
||||
MemoryPatch,
|
||||
RecallArms,
|
||||
ScanPage,
|
||||
StoredMemory,
|
||||
build_fact_records,
|
||||
@@ -75,6 +76,7 @@ __all__ = [
|
||||
"FactRecord",
|
||||
"MemoriesExtension",
|
||||
"MemoryPatch",
|
||||
"RecallArms",
|
||||
"ScanPage",
|
||||
"StoredMemory",
|
||||
"build_fact_records",
|
||||
|
||||
@@ -47,7 +47,7 @@ from typing import TYPE_CHECKING, Any
|
||||
from ...extensions.base import Extension
|
||||
|
||||
if TYPE_CHECKING: # pragma: no cover - typing only
|
||||
from ..search.retrieval import GraphRetriever, SemanticBm25Result
|
||||
from ..search.retrieval import GraphRetriever
|
||||
|
||||
|
||||
class MemoryTxn:
|
||||
@@ -63,6 +63,21 @@ class MemoryTxn:
|
||||
deliberately empty: only the store that minted a handle interprets it."""
|
||||
|
||||
|
||||
class StoreWriteUnavailable(RuntimeError):
|
||||
"""The store cannot accept writes for this bank *right now*, but will shortly.
|
||||
|
||||
Distinct from a failure: nothing is wrong, the bank is briefly closed to writes — a store
|
||||
migrating a bank between backends holds it for a few seconds while it takes the final delta
|
||||
and flips. The caller should retry rather than surface an error, which is why the API maps
|
||||
this to 503 with a `Retry-After` rather than a 5xx that reads as a bug.
|
||||
|
||||
Raised from :meth:`MemoriesExtension.assert_writable` and from bank-scoped write methods.
|
||||
"""
|
||||
|
||||
#: Seconds a caller should wait before retrying. A cutover freeze is drain + a reconcile.
|
||||
retry_after: int = 30
|
||||
|
||||
|
||||
# Keys used in an implementation's opaque metadata bag for the `memory_units`
|
||||
# columns it has no first-class model of. These round-trip verbatim: they are
|
||||
# stored without interpretation and returned on every hit, which is what lets
|
||||
@@ -77,6 +92,30 @@ META_METADATA_JSON = "metadata_json"
|
||||
META_OBSERVATION_SCOPES = "observation_scopes"
|
||||
META_TEXT_SIGNALS = "text_signals"
|
||||
META_CREATED_AT = "created_at"
|
||||
#: When the memory last changed, and the contract every write path owes it (#3490):
|
||||
#: a write that changes what the memory *is* — text, context, dates, fact_type, tags,
|
||||
#: metadata, embedding, an observation's sources — stamps ``updated_at``, so a consumer
|
||||
#: chasing ``WHERE updated_at > watermark`` sees the change. Those consumers are
|
||||
#: incremental export, cache invalidation, the mental-model staleness check
|
||||
#: (:meth:`any_memory_updated_since`) and its delta refresh — and recall's own
|
||||
#: ``created_after`` / ``created_before`` window, which despite the name filters on this
|
||||
#: column, so what stamps it also decides what a date-bounded recall returns.
|
||||
#:
|
||||
#: The consolidation *scheduler* is the one deliberate exception: when a pass records
|
||||
#: that it folded a fact (or requeues one whose observation went away) it writes only
|
||||
#: ``consolidated_at`` / ``consolidation_failed_at``, which are scheduler state rather
|
||||
#: than the memory. Stamping there would make every pass look like an edit to every fact
|
||||
#: it folded — re-flagging mental models stale and re-feeding unchanged facts to a delta
|
||||
#: refresh. :meth:`MemoriesExtension.mark_consolidated` and the requeue sites that clear
|
||||
#: the markers inline therefore leave the column alone.
|
||||
#:
|
||||
#: The exemption is that *situation*, not the two columns: a write that clears the markers
|
||||
#: as part of a real change to the memory still stamps — :meth:`restore_memory` brings an
|
||||
#: archived memory back and resets it for re-consolidation in one statement, and that is an
|
||||
#: edit. A store that owns memories itself is expected to keep the same contract.
|
||||
#:
|
||||
#: No timestamp can report a hard delete; a consumer that must catch those needs a
|
||||
#: content fingerprint, not a watermark.
|
||||
META_UPDATED_AT = "updated_at"
|
||||
# Observation bookkeeping. `source_memory_ids` is a JSON list: an implementation
|
||||
# with no edge relation carries an observation's sources denormalised.
|
||||
@@ -359,6 +398,50 @@ def build_fact_records(
|
||||
return records
|
||||
|
||||
|
||||
@dataclass
|
||||
class RelinkPassResult:
|
||||
"""What one relink drain got through.
|
||||
|
||||
``queue_exhausted`` is False when the pass stopped on its deadline (or the
|
||||
runaway-iteration cap) with rows still queued — not a failure, since every
|
||||
batch commits before the next is claimed, but the caller needs to know the
|
||||
queue is not empty so it can arrange for the rest to be picked up.
|
||||
"""
|
||||
|
||||
units_processed: int = 0
|
||||
links_added: int = 0
|
||||
queue_exhausted: bool = True
|
||||
|
||||
|
||||
@dataclass
|
||||
class EntityPrunePassResult:
|
||||
"""What one entity-prune drain got through.
|
||||
|
||||
``entities_examined`` counts candidates claimed, not rows deleted: most
|
||||
candidates turn out to be alive and are kept, which is the pass working as
|
||||
intended rather than wasted effort.
|
||||
"""
|
||||
|
||||
entities_examined: int = 0
|
||||
orphan_entities_pruned: int = 0
|
||||
stale_cooccurrences_pruned: int = 0
|
||||
queue_exhausted: bool = True
|
||||
|
||||
|
||||
@dataclass
|
||||
class RecallArms:
|
||||
"""One fact_type's per-arm candidate lists from :meth:`MemoriesExtension.recall_unified`.
|
||||
|
||||
Each list holds ``RetrievalResult`` items, unfused — RRF/rerank happen downstream.
|
||||
``temporal`` is empty unless a window was given; ``graph`` is empty when that arm is off.
|
||||
"""
|
||||
|
||||
semantic: list = field(default_factory=list)
|
||||
bm25: list = field(default_factory=list)
|
||||
graph: list = field(default_factory=list)
|
||||
temporal: list = field(default_factory=list)
|
||||
|
||||
|
||||
class MemoriesExtension(Extension, ABC):
|
||||
"""Storage + retrieval for memory units and their links, behind one interface.
|
||||
|
||||
@@ -408,6 +491,21 @@ class MemoriesExtension(Extension, ABC):
|
||||
that keeps some banks in a separate backend overrides it. See :meth:`writes_memory_rows_in_sql_for`."""
|
||||
return self.owns_document_store
|
||||
|
||||
async def assert_writable(self, bank_id: str) -> None:
|
||||
"""Refuse the operation if the store cannot take writes for this bank right now.
|
||||
|
||||
Called at the entry to a *multi-store* operation — retain, which writes documents, chunks
|
||||
and entities through paths that are not this interface at all. Every write that does go
|
||||
through a store method is already covered by the method itself; this exists for the ones
|
||||
that are not, so a store can close a bank completely rather than only partly.
|
||||
|
||||
The default is a no-op, so no existing store needs a change. A store that migrates banks
|
||||
between backends raises :class:`StoreWriteUnavailable` while a bank is mid-cutover: the
|
||||
window is seconds, and a retain that started before it and writes after it would land in
|
||||
the store that is about to stop being authoritative.
|
||||
"""
|
||||
return None
|
||||
|
||||
# ------------------------------------------------------------------ lifecycle
|
||||
|
||||
async def initialize(self) -> None:
|
||||
@@ -620,10 +718,10 @@ class MemoriesExtension(Extension, ABC):
|
||||
"""Apply partial updates. Only the fields set on each patch change."""
|
||||
raise NotImplementedError
|
||||
|
||||
# ------------------------------------------------------------------ recall arms
|
||||
# ------------------------------------------------------------------ recall
|
||||
|
||||
@abstractmethod
|
||||
async def search(
|
||||
async def recall_unified(
|
||||
self,
|
||||
*,
|
||||
conn,
|
||||
@@ -632,6 +730,8 @@ class MemoriesExtension(Extension, ABC):
|
||||
query_embedding: str,
|
||||
query_text: str,
|
||||
limit: int,
|
||||
temporal_window: "tuple[datetime, datetime] | None" = None,
|
||||
temporal_semantic_threshold: float = 0.1,
|
||||
tags: list[str] | None = None,
|
||||
tags_match: str = "any",
|
||||
tag_groups: list | None = None,
|
||||
@@ -639,40 +739,24 @@ class MemoriesExtension(Extension, ABC):
|
||||
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]":
|
||||
"""Run the semantic + BM25 arms.
|
||||
enable_graph: bool = True,
|
||||
) -> "dict[str, RecallArms]":
|
||||
"""Run ALL retrieval arms for every fact_type — the whole recall interface, in one call.
|
||||
|
||||
Returns ``{fact_type: SemanticBm25Result(semantic, bm25, graph_seeds)}`` of
|
||||
``RetrievalResult`` — the contract ``retrieve_semantic_bm25_combined`` has.
|
||||
``graph_seed_min_similarity`` restricts which semantic hits seed the graph
|
||||
arm (Postgres populates ``graph_seeds``; a store with its own graph arm
|
||||
leaves it ``None``).
|
||||
"""
|
||||
Returns ``{fact_type: RecallArms(semantic, bm25, graph, temporal)}`` of
|
||||
``RetrievalResult``: the four per-arm candidate lists, unfused (RRF/rerank happen
|
||||
downstream, unchanged). ``temporal`` is empty unless ``temporal_window`` is given;
|
||||
``graph`` is empty when ``enable_graph`` is False.
|
||||
|
||||
@abstractmethod
|
||||
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]:
|
||||
"""Run the temporal arm over ``[start_date, end_date]``.
|
||||
This is the ONE method recall goes through — how a store answers the arms is entirely its
|
||||
own business. Postgres runs the split per-arm SQL orchestration behind this (a dense+BM25
|
||||
UNION query, a graph retriever per type, a temporal query); a store that owns its index
|
||||
answers every arm from a single query with no per-arm round-trips. Either way the caller
|
||||
sees only this method and its per-arm result.
|
||||
|
||||
Returns ``{fact_type: [RetrievalResult]}``: entry points whose effective
|
||||
time — ``COALESCE(occurred_start, mentioned_at, occurred_end)`` — falls in
|
||||
the window, spread one hop and scored by proximity to it.
|
||||
``conn`` is the store's connection handle for the call. Postgres treats it as the pool it
|
||||
acquires its own connections from and runs the graph arm on; a store that reaches its index
|
||||
another way (e.g. over the network) ignores it.
|
||||
"""
|
||||
|
||||
def graph_retriever(self) -> "GraphRetriever | None":
|
||||
@@ -822,6 +906,9 @@ class MemoriesExtension(Extension, ABC):
|
||||
|
||||
``failed`` stamps the failure marker instead, so a memory the LLM could
|
||||
not consolidate is not retried forever.
|
||||
|
||||
This is scheduler state, not an edit: it must leave the memory's
|
||||
``updated_at`` alone (see :data:`META_UPDATED_AT`).
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
@@ -849,6 +936,19 @@ class MemoriesExtension(Extension, ABC):
|
||||
inherits its source memories' entities, so a hit reads the same either way.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
async def resolve_entity_names(self, *, conn, fq_table, bank_id: str, entity_ids: list[str]) -> dict[str, str]:
|
||||
"""``{entity_id: canonical_name}`` for the given ids, from the ``entities`` registry.
|
||||
|
||||
The label half of :meth:`entity_map_for_units`, split out so a backend that
|
||||
already carries a unit's entity ids on the recalled result can turn those ids
|
||||
into names without re-fetching the memories — recall then builds the entity map
|
||||
from the result's ids plus this one lookup. Bank-scoped, and ids with no registry
|
||||
row are simply absent from the result. The concrete SQL is the store's, next to
|
||||
:meth:`entity_map_for_units`, because the query dialect belongs to the backend,
|
||||
not this interface.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
async def any_memory_updated_since(
|
||||
self,
|
||||
@@ -889,7 +989,7 @@ class MemoriesExtension(Extension, ABC):
|
||||
``last_memory_write_at`` is the newest write time (``updated_at``) across
|
||||
the bank's memories, or None for an empty bank. It is the bank-wide
|
||||
counterpart of :meth:`any_memory_updated_since`: a mental model whose
|
||||
``last_refreshed_at`` is at or after it cannot be stale, whatever its
|
||||
``last_memory_seen_at`` is at or after it cannot be stale, whatever its
|
||||
scope — which is how the stats and knowledge-tree surfaces answer "is
|
||||
this up to date" for many models without a scoped scan each.
|
||||
"""
|
||||
@@ -943,7 +1043,7 @@ class MemoriesExtension(Extension, ABC):
|
||||
ops,
|
||||
fq_table,
|
||||
bank_id: str,
|
||||
fact_type: str | None = None,
|
||||
fact_type: str | list[str] | None = None,
|
||||
search_query: str | None = None,
|
||||
consolidation_state: str | None = None,
|
||||
state: str | None = None,
|
||||
@@ -1005,6 +1105,9 @@ class MemoriesExtension(Extension, ABC):
|
||||
|
||||
Returns the restored memory (so the caller can recompute its embedding —
|
||||
the archive need not keep one), or ``None`` if it was not archived.
|
||||
|
||||
Bringing a memory back is an edit, so this stamps ``updated_at`` even though
|
||||
it also resets the consolidation markers (see :data:`META_UPDATED_AT`).
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
@@ -1015,6 +1118,10 @@ class MemoriesExtension(Extension, ABC):
|
||||
the store whose write is the row itself — reverting or editing a memory has
|
||||
to put a freshly computed vector back on it, so this is a real write for
|
||||
both. ``embedding`` is a float list or the pgvector literal.
|
||||
|
||||
The vector is part of the memory, so this stamps ``updated_at`` itself rather
|
||||
than leaning on the edit statement its in-tree callers happen to pair it with
|
||||
(see :data:`META_UPDATED_AT`).
|
||||
"""
|
||||
|
||||
async def clear_unit_entities(self, *, conn, fq_table, bank_id: str, unit_id: str) -> None:
|
||||
@@ -1132,7 +1239,15 @@ class MemoriesExtension(Extension, ABC):
|
||||
# join table to sweep, so those passes are no-ops for it.
|
||||
|
||||
async def record_unit_entities(
|
||||
self, *, conn, ops, fq_table, bank_id: str | None = None, unit_ids: list[Any], entity_ids: list[Any]
|
||||
self,
|
||||
*,
|
||||
conn,
|
||||
ops,
|
||||
fq_table,
|
||||
bank_id: str | None = None,
|
||||
unit_ids: list[Any],
|
||||
entity_ids: list[Any],
|
||||
txn: "MemoryTxn | None" = None,
|
||||
) -> None:
|
||||
"""Record the unit→entity postings for a batch of memories.
|
||||
|
||||
@@ -1143,6 +1258,14 @@ class MemoriesExtension(Extension, ABC):
|
||||
the memory (rather than in a global join table) needs to know which
|
||||
namespace the units live in — the Postgres join is keyed by global unit id
|
||||
and ignores it.
|
||||
|
||||
``txn`` is the caller's write-group handle. For a store that keeps the
|
||||
posting ON the memory this call is a re-write of rows the same write-group
|
||||
already created, so it belongs to that group: passing the handle keeps the
|
||||
two writes atomic together and — for a store that records what its groups
|
||||
wrote — keeps this write inside the group's accounting. Ignored by the
|
||||
Postgres store, whose posting is an ordinary row in the caller's own
|
||||
transaction.
|
||||
"""
|
||||
|
||||
async def enqueue_relink_victims(
|
||||
@@ -1157,17 +1280,28 @@ class MemoriesExtension(Extension, ABC):
|
||||
"""
|
||||
return 0
|
||||
|
||||
async def relink_pass(self, *, backend, fq_table, bank_id: str, config) -> dict:
|
||||
"""Top up links for queued victims. ``{}`` when there is nothing to relink."""
|
||||
return {}
|
||||
async def relink_pass(
|
||||
self, *, backend, fq_table, bank_id: str, config, deadline: float | None = None
|
||||
) -> "RelinkPassResult":
|
||||
"""Top up links for queued victims. All-zero when there is nothing to relink."""
|
||||
return RelinkPassResult()
|
||||
|
||||
async def prune_orphan_entities(self, *, conn, fq_table, bank_id: str) -> int:
|
||||
"""Delete `entities` rows no live memory references. Returns the count."""
|
||||
async def enqueue_entity_prune_candidates(self, *, conn, fq_table, bank_id: str, affected_unit_ids: list) -> int:
|
||||
"""Queue the entities ``affected_unit_ids`` reference as prune candidates.
|
||||
|
||||
Zero for a store that never wrote `unit_entities`: it has no entity
|
||||
postings to lose, so nothing can become an orphan.
|
||||
"""
|
||||
return 0
|
||||
|
||||
async def prune_stale_cooccurrences(self, *, conn, fq_table, bank_id: str) -> int:
|
||||
"""Delete co-occurrence rows whose witnessing memories are all gone."""
|
||||
return 0
|
||||
async def entity_prune_pass(
|
||||
self, *, backend, fq_table, bank_id: str, deadline: float | None = None
|
||||
) -> "EntityPrunePassResult":
|
||||
"""Prune queued candidate entities and the co-occurrences they stranded.
|
||||
|
||||
All-zero when the store keeps no entity postings and so queues nothing.
|
||||
"""
|
||||
return EntityPrunePassResult()
|
||||
|
||||
|
||||
__all__ = [
|
||||
@@ -1187,10 +1321,12 @@ __all__ = [
|
||||
"META_UPDATED_AT",
|
||||
"CausalEdgeRecord",
|
||||
"DeletePredicate",
|
||||
"EntityPrunePassResult",
|
||||
"FactRecord",
|
||||
"MemoriesExtension",
|
||||
"MemoryPatch",
|
||||
"MemoryTxn",
|
||||
"RelinkPassResult",
|
||||
"ScanPage",
|
||||
"StoredMemory",
|
||||
"build_fact_records",
|
||||
|
||||
@@ -17,6 +17,13 @@ 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.
|
||||
|
||||
``pending`` and ``failed`` are disjoint: pending carries the consolidator's
|
||||
own candidate predicate (``consolidated_at IS NULL AND consolidation_failed_at
|
||||
IS NULL``, see ``reads.find_unconsolidated``), so it reads as "work the
|
||||
consolidator will still do" and drains to zero. A fact the LLM could not
|
||||
handle is counted once, under ``failed``, and only leaves that bucket via the
|
||||
consolidation-recovery endpoint.
|
||||
|
||||
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)
|
||||
@@ -29,7 +36,11 @@ async def consolidation_freshness(*, conn, fq_table: Callable[[str], str], bank_
|
||||
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 consolidated_at IS NULL
|
||||
AND consolidation_failed_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
|
||||
|
||||
@@ -84,7 +84,7 @@ async def list_memory_units(
|
||||
ops,
|
||||
fq_table,
|
||||
bank_id: str,
|
||||
fact_type: str | None = None,
|
||||
fact_type: str | list[str] | None = None,
|
||||
search_query: str | None = None,
|
||||
consolidation_state: str | None = None,
|
||||
state: str | None = None,
|
||||
@@ -104,7 +104,8 @@ async def list_memory_units(
|
||||
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)
|
||||
fact_type: Filter by fact type (world, experience). A list matches any of
|
||||
them; an empty list is treated as no filter.
|
||||
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
|
||||
@@ -154,8 +155,14 @@ async def list_memory_units(
|
||||
|
||||
if fact_type:
|
||||
param_count += 1
|
||||
query_conditions.append(f"fact_type = ${param_count}")
|
||||
query_params.append(fact_type)
|
||||
if isinstance(fact_type, str):
|
||||
query_conditions.append(f"fact_type = ${param_count}")
|
||||
query_params.append(fact_type)
|
||||
else:
|
||||
# A list is "any of these" — one array parameter rather than an IN list
|
||||
# whose placeholder count varies with the caller's argument.
|
||||
query_conditions.append(f"fact_type = ANY(${param_count}::text[])")
|
||||
query_params.append(list(fact_type))
|
||||
|
||||
if document_id:
|
||||
param_count += 1
|
||||
@@ -240,7 +247,8 @@ async def list_memory_units(
|
||||
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}
|
||||
tags, metadata, consolidated_at, consolidation_failed_at, edited_at,
|
||||
updated_at, source_memory_ids, {curation_cols}
|
||||
FROM {source_table}
|
||||
{where_clause}
|
||||
ORDER BY mentioned_at DESC NULLS LAST, created_at DESC
|
||||
@@ -304,6 +312,12 @@ async def list_memory_units(
|
||||
"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,
|
||||
# Both come off the row already selected above, so neither adds a
|
||||
# query: updated_at is the write watermark curation and freshness
|
||||
# checks compare against, and source_memory_ids is an observation's
|
||||
# lineage (empty for a source fact).
|
||||
"updated_at": row["updated_at"].isoformat() if row["updated_at"] else None,
|
||||
"source_memory_ids": [str(sid) for sid in row["source_memory_ids"] or []],
|
||||
}
|
||||
)
|
||||
|
||||
@@ -463,7 +477,7 @@ async def list_entities(
|
||||
# Get paginated entities
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, canonical_name, mention_count, first_seen, last_seen, metadata
|
||||
SELECT id, canonical_name, entity_kind, mention_count, first_seen, last_seen, metadata
|
||||
FROM {fq_table("entities")}
|
||||
WHERE {where_clause}
|
||||
ORDER BY mention_count DESC, last_seen DESC, id ASC
|
||||
@@ -490,6 +504,9 @@ async def list_entities(
|
||||
{
|
||||
"id": str(row["id"]),
|
||||
"canonical_name": row["canonical_name"],
|
||||
# How the entity was classified (label vs free-form, etc.); same row,
|
||||
# so listing it costs nothing extra.
|
||||
"entity_kind": row["entity_kind"],
|
||||
"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,
|
||||
|
||||
@@ -11,11 +11,11 @@ Two groups of callers:
|
||||
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.
|
||||
* **The graph-maintenance job.** :func:`enqueue_relink_victims` and
|
||||
:func:`enqueue_entity_prune_candidates` run inside the delete transaction;
|
||||
:func:`relink_pass` and :func:`entity_prune_pass` are the two drain loops the
|
||||
job drives. The job keeps the orchestration (pass ordering, the time budget,
|
||||
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.
|
||||
@@ -28,8 +28,10 @@ answers them with zeroes rather than with SQL.
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import time
|
||||
import uuid as uuid_module
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
from ....config import get_config
|
||||
@@ -40,6 +42,7 @@ from ...retain.link_utils import (
|
||||
_normalize_datetime,
|
||||
compute_semantic_links_ann,
|
||||
)
|
||||
from ..base import EntityPrunePassResult, RelinkPassResult
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -59,6 +62,37 @@ _DRAIN_BATCH_SIZE = 50
|
||||
# iteration that's 500k targets, far beyond any realistic single-bank backlog.
|
||||
_RELINK_ITERATION_CAP = 10000
|
||||
|
||||
# Candidate entities claimed per entity-prune iteration.
|
||||
#
|
||||
# The binding cost is the cooccurrence prune: a candidate drags in every
|
||||
# cooccurrence pair it appears in, and the statement builds a set of currently
|
||||
# live pairs (#3367) seeded from those candidates' units to judge them against.
|
||||
# Measured on a deliberately dense fixture (100k entities, 1.5M unit_entities,
|
||||
# 2.86M cooccurrences, endpoints holding 150-400 postings each):
|
||||
#
|
||||
# batch 50 → 16-65ms <- here
|
||||
# batch 500 → 265s <- the planner flips to a per-row plan and the
|
||||
# statement blows the 60s command timeout
|
||||
#
|
||||
# So the batch size, not the bank size, is what has to stay bounded — and it has
|
||||
# to stay small enough that the planner keeps choosing the hash anti-join.
|
||||
# Raising it trades away three orders of magnitude of margin; don't, without
|
||||
# re-measuring against a bank with hub entities.
|
||||
#
|
||||
# Also stays under Oracle's 1000-element IN-list limit, since ops_oracle expands
|
||||
# ``= ANY(...)`` into an explicit list.
|
||||
_ENTITY_PRUNE_BATCH_SIZE = 50
|
||||
|
||||
# Deadlock retries per entity-prune batch. The batch is idempotent, so a retry
|
||||
# only re-deletes what is still dead; a handful of attempts clears the
|
||||
# contention window a concurrent retain opens.
|
||||
_PRUNE_BATCH_MAX_RETRIES = 3
|
||||
|
||||
# Unit ids per candidate-lookup round-trip when enqueueing. Bounded by Oracle's
|
||||
# 1000-element IN-list limit (ops_oracle expands ``= ANY(...)`` into a literal
|
||||
# list), which a bulk delete would otherwise blow straight through.
|
||||
_ENQUEUE_LOOKUP_CHUNK = 500
|
||||
|
||||
# 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
|
||||
@@ -445,6 +479,46 @@ async def entity_map_for_units(
|
||||
return by_unit
|
||||
|
||||
|
||||
async def resolve_entity_names(
|
||||
*,
|
||||
conn: DatabaseConnection,
|
||||
fq_table: Callable[[str], str],
|
||||
bank_id: str,
|
||||
entity_ids: list[str],
|
||||
) -> dict[str, str]:
|
||||
"""``{entity_id: canonical_name}`` for the given ids, scoped to ``bank_id``.
|
||||
|
||||
The label half of :func:`entity_map_for_units`, for a backend that already
|
||||
carries a unit's entity ids on the recalled result: recall builds the
|
||||
unit->entity map from those ids and needs only the names, so this resolves the
|
||||
``entities`` registry once without re-fetching any memory. Bank-scoped like the
|
||||
sibling registry reads (the ``entities`` table has a ``bank_id`` column).
|
||||
|
||||
Ids that don't parse as UUIDs are dropped rather than raised — a malformed id
|
||||
on a store's result payload must not turn into a DB error mid-recall — and ids
|
||||
with no registry row (or in another bank) are simply absent from the result.
|
||||
"""
|
||||
if not entity_ids:
|
||||
return {}
|
||||
# Bind ``uuid.UUID`` objects for the ``uuid[]`` param (repo convention, see
|
||||
# ``_as_uuids``), but coerce defensively: skip anything unparseable instead of
|
||||
# letting the whole resolve raise.
|
||||
uuids: list = []
|
||||
for raw in {str(e) for e in entity_ids}:
|
||||
try:
|
||||
uuids.append(uuid_module.UUID(raw))
|
||||
except (ValueError, AttributeError, TypeError):
|
||||
continue
|
||||
if not uuids:
|
||||
return {}
|
||||
rows = await conn.fetch(
|
||||
f"SELECT id, canonical_name FROM {fq_table('entities')} WHERE id = ANY($1::uuid[]) AND bank_id = $2",
|
||||
uuids,
|
||||
bank_id,
|
||||
)
|
||||
return {str(row["id"]): row["canonical_name"] for row in rows}
|
||||
|
||||
|
||||
# --------------------------------------------------------------- maintenance
|
||||
|
||||
|
||||
@@ -528,7 +602,8 @@ async def relink_pass(
|
||||
fq_table: Callable[[str], str],
|
||||
bank_id: str,
|
||||
config: Any,
|
||||
) -> dict:
|
||||
deadline: float | None = None,
|
||||
) -> RelinkPassResult:
|
||||
"""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
|
||||
@@ -549,8 +624,13 @@ async def relink_pass(
|
||||
means) and never reads it; it is accepted so a store that *does* tune its
|
||||
relinking gets it.
|
||||
|
||||
``deadline`` is a ``time.monotonic()`` value past which no new batch is
|
||||
claimed. Each batch commits before the next is claimed, so stopping early
|
||||
keeps the work already done and leaves the rest queued for the next run.
|
||||
|
||||
Returns:
|
||||
``{"relink_units_processed": int, "relink_links_added": int}``.
|
||||
A :class:`RelinkPassResult`. ``queue_exhausted`` is False when the
|
||||
deadline (or the iteration cap) stopped the drain with rows still queued.
|
||||
"""
|
||||
del config # accepted for symmetry with stores that tune their own relinking
|
||||
ops = backend.ops
|
||||
@@ -558,7 +638,12 @@ async def relink_pass(
|
||||
units_processed = 0
|
||||
links_added = 0
|
||||
iterations = 0
|
||||
drained = True
|
||||
while True:
|
||||
if deadline is not None and time.monotonic() >= deadline:
|
||||
drained = False
|
||||
break
|
||||
|
||||
from ...memory_engine import acquire_with_retry
|
||||
|
||||
async with acquire_with_retry(backend) as conn:
|
||||
@@ -584,9 +669,14 @@ async def relink_pass(
|
||||
f"[GRAPH_MAINT] bank={bank_id} hit iteration cap ({iterations}); aborting relink "
|
||||
f"(units_processed={units_processed}, links_added={links_added})"
|
||||
)
|
||||
drained = False
|
||||
break
|
||||
|
||||
return {"relink_units_processed": units_processed, "relink_links_added": links_added}
|
||||
return RelinkPassResult(
|
||||
units_processed=units_processed,
|
||||
links_added=links_added,
|
||||
queue_exhausted=drained,
|
||||
)
|
||||
|
||||
|
||||
async def _relink_batch(
|
||||
@@ -716,69 +806,192 @@ async def _relink_batch(
|
||||
return len(new_links)
|
||||
|
||||
|
||||
async def prune_orphan_entities(
|
||||
async def enqueue_entity_prune_candidates(
|
||||
*,
|
||||
conn: DatabaseConnection,
|
||||
fq_table: Callable[[str], str],
|
||||
bank_id: str,
|
||||
affected_unit_ids: list,
|
||||
) -> int:
|
||||
"""Delete ``entities`` rows in the bank with no remaining ``unit_entities``
|
||||
references. Returns the number pruned.
|
||||
"""Enqueue the entities ``affected_unit_ids`` reference as prune candidates.
|
||||
|
||||
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.
|
||||
Must run inside the same transaction that removes those units (or their
|
||||
``unit_entities`` rows), *before* the delete or cascade fires — afterwards
|
||||
there is no posting left to read the entity ids from, and the entity is
|
||||
stranded as an orphan nothing will ever look at again.
|
||||
|
||||
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.
|
||||
Enqueueing an entity that turns out to still be referenced is free: the
|
||||
drain re-checks and keeps it. Over-enqueueing is always the safe direction.
|
||||
|
||||
Returns:
|
||||
Number of candidate entities enqueued.
|
||||
"""
|
||||
if not affected_unit_ids:
|
||||
return 0
|
||||
|
||||
ops = _ops_for(conn)
|
||||
return await ops.prune_orphan_entities(
|
||||
conn,
|
||||
fq_table("entities"),
|
||||
fq_table("unit_entities"),
|
||||
bank_id,
|
||||
)
|
||||
queue_table = fq_table("entity_maintenance_queue")
|
||||
ue_table = fq_table("unit_entities")
|
||||
unit_uuids = _as_uuids(list(affected_unit_ids))
|
||||
|
||||
# Chunked because a bulk delete can hand in thousands of unit ids and the
|
||||
# lookup binds them with `= ANY(...)`, which ops_oracle expands into a
|
||||
# literal IN list — Oracle caps those at 1000 elements.
|
||||
enqueued = 0
|
||||
for start in range(0, len(unit_uuids), _ENQUEUE_LOOKUP_CHUNK):
|
||||
enqueued += await ops.enqueue_entity_maintenance(
|
||||
conn,
|
||||
queue_table,
|
||||
ue_table,
|
||||
bank_id,
|
||||
unit_uuids[start : start + _ENQUEUE_LOOKUP_CHUNK],
|
||||
)
|
||||
return enqueued
|
||||
|
||||
|
||||
async def prune_stale_cooccurrences(
|
||||
@dataclass
|
||||
class _PruneBatch:
|
||||
"""One entity-prune iteration's counters (avoids a bare tuple return)."""
|
||||
|
||||
claimed: int
|
||||
orphan_entities_pruned: int
|
||||
stale_cooccurrences_pruned: int
|
||||
|
||||
|
||||
async def entity_prune_pass(
|
||||
*,
|
||||
conn: DatabaseConnection,
|
||||
backend: Any,
|
||||
fq_table: Callable[[str], str],
|
||||
bank_id: str,
|
||||
) -> int:
|
||||
"""Delete cooccurrence rows no current memory witnesses. Returns the count.
|
||||
deadline: float | None = None,
|
||||
) -> EntityPrunePassResult:
|
||||
"""Drain ``entity_maintenance_queue`` for ``bank_id``, pruning what died.
|
||||
|
||||
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.
|
||||
Per-iteration loop: claim → prune → commit, mirroring :func:`relink_pass`.
|
||||
Each iteration does two deletes over the claimed batch:
|
||||
|
||||
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.
|
||||
1. **Orphan entities** — candidates with no remaining ``unit_entities`` row.
|
||||
FK ON DELETE CASCADE on ``entity_cooccurrences`` takes their cooccurrence
|
||||
rows with them, which is why this runs first.
|
||||
2. **Stale cooccurrences** — pairs incident to a surviving candidate where
|
||||
both entities still exist but no current unit witnesses them together.
|
||||
The cooccurrence was real when recorded; every unit that saw it has since
|
||||
been deleted. The FK cascade above cannot see this case.
|
||||
|
||||
Both deletes are scoped to the claimed batch. They used to be bank-wide
|
||||
statements re-run on every invocation — the orphan prune probing once per
|
||||
entity in the bank, the cooccurrence prune evaluating an INTERSECT per
|
||||
cooccurrence row in the bank — so their cost tracked the size of the bank
|
||||
rather than the size of the delete, and past a few million rows they could
|
||||
no longer finish inside asyncpg's command timeout. The job then failed on
|
||||
every run, forever, on exactly the banks that most needed it (#3222).
|
||||
|
||||
Committing per batch is what makes the pass resumable: work already done
|
||||
stays done when ``deadline`` cuts the drain short or the task dies, and the
|
||||
next run picks up the remaining queue rows.
|
||||
|
||||
Args:
|
||||
backend: Database backend — the loop spans a transaction per batch, so
|
||||
it acquires its own connections.
|
||||
fq_table: Schema-qualifier for table names.
|
||||
bank_id: Bank to drain.
|
||||
deadline: ``time.monotonic()`` value past which no new batch is claimed.
|
||||
``None`` drains to empty.
|
||||
|
||||
Returns:
|
||||
An :class:`EntityPrunePassResult`. ``queue_exhausted`` is False when the
|
||||
deadline stopped the drain with rows still queued.
|
||||
"""
|
||||
ops = _ops_for(conn)
|
||||
return await ops.prune_stale_cooccurrences(
|
||||
conn,
|
||||
fq_table("entity_cooccurrences"),
|
||||
fq_table("unit_entities"),
|
||||
fq_table("entities"),
|
||||
bank_id,
|
||||
from ...db_utils import retry_with_backoff
|
||||
from ...memory_engine import acquire_with_retry
|
||||
|
||||
examined = 0
|
||||
orphans_pruned = 0
|
||||
stale_pruned = 0
|
||||
drained = True
|
||||
|
||||
while True:
|
||||
if deadline is not None and time.monotonic() >= deadline:
|
||||
drained = False
|
||||
break
|
||||
|
||||
async def _run_batch() -> _PruneBatch:
|
||||
async with acquire_with_retry(backend) as conn:
|
||||
async with conn.transaction():
|
||||
ops = backend.ops
|
||||
entity_ids = await ops.claim_entity_maintenance_batch(
|
||||
conn,
|
||||
fq_table("entity_maintenance_queue"),
|
||||
bank_id,
|
||||
_ENTITY_PRUNE_BATCH_SIZE,
|
||||
)
|
||||
if not entity_ids:
|
||||
return _PruneBatch(claimed=0, orphan_entities_pruned=0, stale_cooccurrences_pruned=0)
|
||||
|
||||
orphaned = await ops.prune_orphan_entities(
|
||||
conn,
|
||||
fq_table("entities"),
|
||||
fq_table("unit_entities"),
|
||||
bank_id,
|
||||
entity_ids,
|
||||
)
|
||||
# The orphan prune above cascades cooccurrences via FK. This
|
||||
# second delete catches the *stale-count* case: both entities
|
||||
# still exist but no current unit witnesses them together.
|
||||
stale = await ops.prune_stale_cooccurrences(
|
||||
conn,
|
||||
fq_table("entity_cooccurrences"),
|
||||
fq_table("unit_entities"),
|
||||
entity_ids,
|
||||
)
|
||||
return _PruneBatch(
|
||||
claimed=len(entity_ids),
|
||||
orphan_entities_pruned=orphaned,
|
||||
stale_cooccurrences_pruned=stale,
|
||||
)
|
||||
|
||||
# Retry the batch on deadlock. Both deletes take their row locks in the
|
||||
# same order the concurrent retain writers do (entity id for the entity
|
||||
# upsert, (entity_id_1, entity_id_2) for the cooccurrence upsert), so a
|
||||
# cycle should not form on Postgres at all; this stays as the backstop
|
||||
# for the paths that ordering can't cover — the FK cascade out of the
|
||||
# orphan prune, and Oracle, whose DELETE can't carry the ordered-lock
|
||||
# CTE. Both deletes are idempotent, so re-running the batch is safe.
|
||||
#
|
||||
# Deliberately narrower than the budget the bank-wide sweep used (8).
|
||||
# `retry_with_backoff` treats a TimeoutError as transient, which was
|
||||
# ruinous while the statement was O(bank): a sweep that could never
|
||||
# finish inside the command timeout was re-run nine times, burning ten
|
||||
# minutes of a worker slot per task attempt (#3222). A bounded batch
|
||||
# that times out is not slow work, it is a sick database — retry a few
|
||||
# times and let the failure surface.
|
||||
batch = await retry_with_backoff(_run_batch, max_retries=_PRUNE_BATCH_MAX_RETRIES)
|
||||
if batch.claimed == 0:
|
||||
break
|
||||
|
||||
examined += batch.claimed
|
||||
orphans_pruned += batch.orphan_entities_pruned
|
||||
stale_pruned += batch.stale_cooccurrences_pruned
|
||||
|
||||
return EntityPrunePassResult(
|
||||
entities_examined=examined,
|
||||
orphan_entities_pruned=orphans_pruned,
|
||||
stale_cooccurrences_pruned=stale_pruned,
|
||||
queue_exhausted=drained,
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"MAX_SEMANTIC_LINKS_PER_UNIT",
|
||||
"enqueue_entity_prune_candidates",
|
||||
"enqueue_relink_victims",
|
||||
"entities_for_units",
|
||||
"entity_map_for_units",
|
||||
"entity_memory_counts",
|
||||
"entity_prune_pass",
|
||||
"graph_direct_links",
|
||||
"graph_entity_rows",
|
||||
"graph_units",
|
||||
"prune_orphan_entities",
|
||||
"prune_stale_cooccurrences",
|
||||
"relink_pass",
|
||||
"resolve_entity_names",
|
||||
]
|
||||
|
||||
@@ -461,8 +461,9 @@ async def mark_consolidated(
|
||||
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
|
||||
``updated_at`` is deliberately left alone — the one exception to the contract
|
||||
documented on ``META_UPDATED_AT`` (``memories.base``) that every other write
|
||||
path owes the column. 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.
|
||||
"""
|
||||
|
||||
@@ -293,6 +293,8 @@ async def delete_stale_observations(
|
||||
)
|
||||
|
||||
if remaining_source_ids:
|
||||
# Requeue: consolidation bookkeeping, so `updated_at` is deliberately not
|
||||
# stamped (see META_UPDATED_AT in ..base) — nothing about these facts changed.
|
||||
await conn.execute(
|
||||
f"""
|
||||
UPDATE {fq_table("memory_units")}
|
||||
@@ -420,7 +422,8 @@ async def invalidate_memory(*, conn, fq_table, bank_id: str, unit_id: str, reaso
|
||||
|
||||
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",
|
||||
f"UPDATE {fq_table('invalidated_memory_units')} SET invalidation_reason = $3, updated_at = now() "
|
||||
f"WHERE id = $1 AND bank_id = $2",
|
||||
str(unit_id),
|
||||
bank_id,
|
||||
reason,
|
||||
@@ -498,7 +501,8 @@ async def restore_memory(*, conn, fq_table, bank_id: str, unit_id: str) -> Store
|
||||
|
||||
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",
|
||||
f"UPDATE {fq_table('memory_units')} SET embedding = $3::vector, updated_at = now() "
|
||||
f"WHERE id = $1 AND bank_id = $2",
|
||||
str(unit_id),
|
||||
bank_id,
|
||||
embedding,
|
||||
|
||||
@@ -23,7 +23,16 @@ from __future__ import annotations
|
||||
from datetime import datetime
|
||||
from typing import Any
|
||||
|
||||
from .base import DeletePredicate, MemoriesExtension, MemoryPatch, ScanPage, StoredMemory
|
||||
from .base import (
|
||||
DeletePredicate,
|
||||
EntityPrunePassResult,
|
||||
MemoriesExtension,
|
||||
MemoryPatch,
|
||||
RecallArms,
|
||||
RelinkPassResult,
|
||||
ScanPage,
|
||||
StoredMemory,
|
||||
)
|
||||
from .pg import counts, curation, graph, reads, writes
|
||||
|
||||
|
||||
@@ -71,7 +80,135 @@ class PostgresMemories(MemoriesExtension):
|
||||
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
|
||||
# ------------------------------------------------------------------ recall
|
||||
|
||||
async def recall_unified(
|
||||
self,
|
||||
*,
|
||||
conn,
|
||||
bank_id: str,
|
||||
fact_types: list[str],
|
||||
query_embedding: str,
|
||||
query_text: str,
|
||||
limit: int,
|
||||
temporal_window: "tuple[datetime, datetime] | None" = None,
|
||||
temporal_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,
|
||||
min_semantic: float | None = None,
|
||||
min_keyword: float | None = None,
|
||||
enable_graph: bool = True,
|
||||
) -> "dict[str, RecallArms]":
|
||||
"""Run every recall arm for Postgres by orchestrating the split per-arm SQL internally.
|
||||
|
||||
The per-arm split is Postgres's own business, kept off the interface: this reproduces the
|
||||
exact orchestration recall used before it was unified — one dense+BM25 UNION query and the
|
||||
temporal query share a single connection, then the graph retriever runs per fact_type on the
|
||||
pool in parallel, seeded by the same dense over-fetch. Result is byte-identical to running
|
||||
the arms separately; fusion/rerank still happen downstream.
|
||||
"""
|
||||
import asyncio
|
||||
|
||||
from ..db_utils import acquire_with_retry
|
||||
from ..search.retrieval import get_default_graph_retriever
|
||||
|
||||
# `conn` is the connection pool: this store owns the per-arm orchestration and acquires its
|
||||
# own connections from it (and runs the graph arm on it).
|
||||
pool = conn
|
||||
|
||||
# graph_seed_min_similarity restricts which dense hits seed the graph arm; only the graph
|
||||
# arm consumes the seeds, so it is resolved only when that arm runs. It does not affect the
|
||||
# semantic/bm25 lists, so the dense+BM25 result is identical whether or not it is passed.
|
||||
graph_seed_min_similarity = None
|
||||
retriever = None
|
||||
if enable_graph:
|
||||
from ...config import get_config
|
||||
|
||||
graph_seed_min_similarity = get_config().graph_seed_min_similarity
|
||||
# Resolving the retriever can lazily construct one, so only do it when the arm is on.
|
||||
retriever = get_default_graph_retriever()
|
||||
|
||||
# Semantic + BM25 (+ temporal) share ONE connection, exactly as before: the dense/keyword
|
||||
# UNION runs first, then the temporal query on the same connection, which is then released
|
||||
# before the graph arm opens its own connections.
|
||||
async with acquire_with_retry(pool) as db_conn:
|
||||
semantic_bm25 = await self.search(
|
||||
conn=db_conn,
|
||||
bank_id=bank_id,
|
||||
fact_types=fact_types,
|
||||
query_embedding=query_embedding,
|
||||
query_text=query_text,
|
||||
limit=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,
|
||||
)
|
||||
|
||||
temporal_by_ft: dict[str, list] = {}
|
||||
if temporal_window is not None:
|
||||
start_date, end_date = temporal_window
|
||||
temporal_by_ft = await self.temporal_search(
|
||||
conn=db_conn,
|
||||
bank_id=bank_id,
|
||||
fact_types=fact_types,
|
||||
query_embedding=query_embedding,
|
||||
start_date=start_date,
|
||||
end_date=end_date,
|
||||
limit=limit,
|
||||
semantic_threshold=temporal_semantic_threshold,
|
||||
tags=tags,
|
||||
tags_match=tags_match,
|
||||
tag_groups=tag_groups,
|
||||
created_after=created_after,
|
||||
created_before=created_before,
|
||||
)
|
||||
|
||||
# Graph per fact_type in parallel, on the pool, after the dense connection is released —
|
||||
# seeded by the dense over-fetch (preselected_semantic_seeds), matching the prior path.
|
||||
graph_by_ft: dict[str, list] = {ft: [] for ft in fact_types}
|
||||
if enable_graph:
|
||||
assert retriever is not None # only resolved when the arm is on
|
||||
|
||||
async def _run_graph(ft: str) -> list:
|
||||
results, _timing = await retriever.retrieve(
|
||||
pool=pool,
|
||||
query_embedding_str=query_embedding,
|
||||
bank_id=bank_id,
|
||||
fact_type=ft,
|
||||
budget=limit,
|
||||
query_text=query_text,
|
||||
tags=tags,
|
||||
tags_match=tags_match,
|
||||
tag_groups=tag_groups,
|
||||
created_after=created_after,
|
||||
created_before=created_before,
|
||||
preselected_semantic_seeds=semantic_bm25[ft].graph_seeds,
|
||||
)
|
||||
return results
|
||||
|
||||
# gather preserves input order, so zip back onto fact_types positionally.
|
||||
graph_lists = await asyncio.gather(*[_run_graph(ft) for ft in fact_types])
|
||||
graph_by_ft = dict(zip(fact_types, graph_lists))
|
||||
|
||||
return {
|
||||
ft: RecallArms(
|
||||
semantic=semantic_bm25[ft].semantic,
|
||||
bm25=semantic_bm25[ft].bm25,
|
||||
graph=graph_by_ft.get(ft, []),
|
||||
temporal=temporal_by_ft.get(ft, []),
|
||||
)
|
||||
for ft in fact_types
|
||||
}
|
||||
|
||||
# ---- per-arm SQL helpers, private to Postgres (called only by recall_unified) ----
|
||||
|
||||
async def search(
|
||||
self,
|
||||
@@ -322,7 +459,7 @@ class PostgresMemories(MemoriesExtension):
|
||||
ops,
|
||||
fq_table,
|
||||
bank_id: str,
|
||||
fact_type: str | None = None,
|
||||
fact_type: str | list[str] | None = None,
|
||||
search_query: str | None = None,
|
||||
consolidation_state: str | None = None,
|
||||
state: str | None = None,
|
||||
@@ -477,12 +614,25 @@ class PostgresMemories(MemoriesExtension):
|
||||
) -> 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)
|
||||
|
||||
async def resolve_entity_names(self, *, conn, fq_table, bank_id: str, entity_ids: list[str]) -> dict[str, str]:
|
||||
return await graph.resolve_entity_names(conn=conn, fq_table=fq_table, bank_id=bank_id, entity_ids=entity_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]
|
||||
self,
|
||||
*,
|
||||
conn,
|
||||
ops,
|
||||
fq_table,
|
||||
bank_id: str | None = None,
|
||||
unit_ids: list[Any],
|
||||
entity_ids: list[Any],
|
||||
txn=None,
|
||||
) -> None:
|
||||
# The join is keyed by global unit id, so bank_id is not needed here.
|
||||
# The join is keyed by global unit id, so bank_id is not needed here. `txn` is inert: this
|
||||
# posting is an ordinary INSERT in the caller's own transaction, which is already the unit
|
||||
# of atomicity — there is no second store to coordinate with.
|
||||
await ops.bulk_insert_unit_entities(conn, fq_table("unit_entities"), unit_ids, entity_ids)
|
||||
|
||||
async def enqueue_relink_victims(
|
||||
@@ -496,14 +646,25 @@ class PostgresMemories(MemoriesExtension):
|
||||
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 relink_pass(
|
||||
self, *, backend, fq_table, bank_id: str, config, deadline: float | None = None
|
||||
) -> RelinkPassResult:
|
||||
return await graph.relink_pass(
|
||||
backend=backend, fq_table=fq_table, bank_id=bank_id, config=config, deadline=deadline
|
||||
)
|
||||
|
||||
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 enqueue_entity_prune_candidates(self, *, conn, fq_table, bank_id: str, affected_unit_ids: list) -> int:
|
||||
return await graph.enqueue_entity_prune_candidates(
|
||||
conn=conn,
|
||||
fq_table=fq_table,
|
||||
bank_id=bank_id,
|
||||
affected_unit_ids=affected_unit_ids,
|
||||
)
|
||||
|
||||
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)
|
||||
async def entity_prune_pass(
|
||||
self, *, backend, fq_table, bank_id: str, deadline: float | None = None
|
||||
) -> EntityPrunePassResult:
|
||||
return await graph.entity_prune_pass(backend=backend, fq_table=fq_table, bank_id=bank_id, deadline=deadline)
|
||||
|
||||
|
||||
__all__ = ["PostgresMemories"]
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -68,20 +68,22 @@ class MentalModelRefreshWindow(BaseModel):
|
||||
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."
|
||||
"Lower bound on when a memory last changed. Set only in delta mode, where it is the "
|
||||
"model's last_memory_seen_at — so a delta refresh only sees memories written or edited "
|
||||
"since the newest one the previous refresh saw."
|
||||
),
|
||||
)
|
||||
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."
|
||||
"Database-time snapshot bounding the refresh. Memories written or edited 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 last_memory_seen_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."
|
||||
),
|
||||
)
|
||||
|
||||
@@ -0,0 +1,35 @@
|
||||
"""Canonical handling of user-supplied memory metadata (issue #3209).
|
||||
|
||||
Metadata is accepted as arbitrary JSON at ingest (file retain, the MCP tools and
|
||||
direct engine calls all take ``dict[str, Any]``) but is stored in a JSONB column
|
||||
and read back through models that declare ``dict[str, str]``. A JSON ``null``
|
||||
value therefore sailed through the write path and then failed validation on
|
||||
every read that returned the affected rows.
|
||||
|
||||
Both ends normalize through the helpers here so the rule lives in one place:
|
||||
|
||||
* ``drop_null_values`` — the write contract. Null-valued keys are dropped;
|
||||
every other value is stored as given (the read contract stringifies).
|
||||
* ``as_string_metadata`` — the read contract. Null-valued keys are dropped and
|
||||
the rest are coerced to strings, so rows written before this normalization
|
||||
existed stay readable without a data migration.
|
||||
"""
|
||||
|
||||
from collections.abc import Mapping
|
||||
from typing import Any
|
||||
|
||||
|
||||
def drop_null_values(metadata: Mapping[str, Any] | None) -> dict[str, Any]:
|
||||
"""Return ``metadata`` without its null-valued keys (empty dict for no metadata)."""
|
||||
if not metadata:
|
||||
return {}
|
||||
return {k: v for k, v in metadata.items() if v is not None}
|
||||
|
||||
|
||||
def as_string_metadata(metadata: Mapping[str, Any] | None) -> dict[str, str]:
|
||||
"""Coerce a stored metadata bag to the ``dict[str, str]`` read contract.
|
||||
|
||||
Drops null-valued keys and stringifies the rest (JSONB round-trips integers
|
||||
as integers, e.g. ``{"original_id": 348}``).
|
||||
"""
|
||||
return {str(k): str(v) for k, v in drop_null_values(metadata).items()}
|
||||
@@ -88,7 +88,7 @@ class AnthropicLLM(LLMInterface):
|
||||
api_key: str,
|
||||
base_url: str,
|
||||
model: str,
|
||||
reasoning_effort: str = "low",
|
||||
reasoning_effort: str | None = None,
|
||||
timeout: float = 300.0,
|
||||
default_headers: dict[str, str] | None = None,
|
||||
extra_body: dict[str, Any] | None = None,
|
||||
@@ -115,6 +115,7 @@ class AnthropicLLM(LLMInterface):
|
||||
**kwargs: Additional provider-specific parameters.
|
||||
"""
|
||||
super().__init__(provider, api_key, base_url, model, reasoning_effort, **kwargs)
|
||||
self._warn_reasoning_effort_unsupported()
|
||||
|
||||
if not self.api_key:
|
||||
raise ValueError("API key is required for Anthropic provider")
|
||||
|
||||
@@ -79,11 +79,12 @@ class ClaudeCodeLLM(LLMInterface):
|
||||
api_key: str, # Will be ignored, uses CLI auth
|
||||
base_url: str,
|
||||
model: str,
|
||||
reasoning_effort: str = "low",
|
||||
reasoning_effort: str | None = None,
|
||||
**kwargs: Any,
|
||||
):
|
||||
"""Initialize Claude Code LLM provider."""
|
||||
super().__init__(provider, api_key, base_url, model, reasoning_effort, **kwargs)
|
||||
self._warn_reasoning_effort_unsupported()
|
||||
|
||||
# Verify Claude Agent SDK is available
|
||||
try:
|
||||
|
||||
@@ -126,7 +126,7 @@ class CodexLLM(LLMInterface):
|
||||
api_key: str, # Will be ignored, reads from the Codex auth.json (CODEX_HOME or ~/.codex)
|
||||
base_url: str,
|
||||
model: str,
|
||||
reasoning_effort: str = "low",
|
||||
reasoning_effort: str | None = None,
|
||||
extra_body: dict[str, Any] | None = None,
|
||||
**kwargs: Any,
|
||||
):
|
||||
@@ -329,12 +329,26 @@ class CodexLLM(LLMInterface):
|
||||
except CodexRefreshExpiredError:
|
||||
raise
|
||||
|
||||
def _map_reasoning_effort(self, effort: str) -> str:
|
||||
def _reasoning_payload(self, summary: str) -> dict[str, str]:
|
||||
"""Build the ``reasoning`` request object.
|
||||
|
||||
``effort`` is present only when the operator configured one: an unset
|
||||
HINDSIGHT_API_*_REASONING_EFFORT means the model runs at its own default
|
||||
effort, and Hindsight does not pick one on the operator's behalf.
|
||||
"""
|
||||
payload = {"summary": summary}
|
||||
if self.reasoning_effort is not None:
|
||||
payload["effort"] = self.reasoning_effort
|
||||
return payload
|
||||
|
||||
def _map_reasoning_effort(self, effort: str | None) -> str:
|
||||
"""
|
||||
Map standard reasoning effort to Codex reasoning summary format.
|
||||
|
||||
Args:
|
||||
effort: Standard effort level ("low", "medium", "high", "xhigh").
|
||||
effort: Standard effort level ("low", "medium", "high", "xhigh"), or None
|
||||
when unconfigured — the summary then stays "auto", the same neutral
|
||||
presentation an unrecognised level gets.
|
||||
|
||||
Returns:
|
||||
Codex reasoning summary: "concise", "detailed", or "auto".
|
||||
@@ -345,7 +359,7 @@ class CodexLLM(LLMInterface):
|
||||
"high": "detailed",
|
||||
"xhigh": "detailed",
|
||||
}
|
||||
return mapping.get(effort.lower(), "auto")
|
||||
return mapping.get(effort.lower(), "auto") if effort else "auto"
|
||||
|
||||
async def verify_connection(self) -> None:
|
||||
"""Verify Codex connection by making a simple test call."""
|
||||
@@ -453,7 +467,7 @@ class CodexLLM(LLMInterface):
|
||||
"tools": [],
|
||||
"tool_choice": "auto",
|
||||
"parallel_tool_calls": True,
|
||||
"reasoning": {"effort": self.reasoning_effort, "summary": reasoning_summary},
|
||||
"reasoning": self._reasoning_payload(reasoning_summary),
|
||||
"store": False, # Codex uses stateless mode
|
||||
"stream": True, # SSE streaming
|
||||
"include": ["reasoning.encrypted_content"],
|
||||
@@ -842,7 +856,7 @@ class CodexLLM(LLMInterface):
|
||||
else tool_choice.mode.value
|
||||
),
|
||||
"parallel_tool_calls": True,
|
||||
"reasoning": {"effort": self.reasoning_effort, "summary": reasoning_summary},
|
||||
"reasoning": self._reasoning_payload(reasoning_summary),
|
||||
"store": False,
|
||||
"stream": True,
|
||||
"include": ["reasoning.encrypted_content"],
|
||||
|
||||
@@ -63,7 +63,7 @@ class FireworksLLM(OpenAICompatibleLLM):
|
||||
api_key: str,
|
||||
base_url: str = "",
|
||||
model: str,
|
||||
reasoning_effort: str = "low",
|
||||
reasoning_effort: str | None = None,
|
||||
account_id: str | None = None,
|
||||
batch_base_url: str | None = None,
|
||||
max_wait_seconds: int | None = None,
|
||||
|
||||
@@ -174,11 +174,12 @@ class GeminiLLM(LLMInterface):
|
||||
api_key: str,
|
||||
base_url: str,
|
||||
model: str,
|
||||
reasoning_effort: str = "low",
|
||||
reasoning_effort: str | None = None,
|
||||
**kwargs: Any,
|
||||
):
|
||||
"""Initialize Gemini/VertexAI LLM provider."""
|
||||
super().__init__(provider, api_key, base_url, model, reasoning_effort, **kwargs)
|
||||
self._warn_reasoning_effort_unsupported()
|
||||
|
||||
self._client = None
|
||||
self._is_vertexai = self.provider == "vertexai"
|
||||
@@ -510,6 +511,21 @@ class GeminiLLM(LLMInterface):
|
||||
if hasattr(response, "candidates") and response.candidates:
|
||||
if hasattr(response.candidates[0], "finish_reason"):
|
||||
finish_reason = str(response.candidates[0].finish_reason)
|
||||
|
||||
# Surface silent truncation. A non-empty response that stopped on
|
||||
# MAX_TOKENS was cut off (often mid-word) yet still returns as a
|
||||
# success — on thinking models the reasoning tokens can consume the
|
||||
# whole max_output_tokens budget, leaving the visible answer
|
||||
# truncated (#3365). Make it visible in the logs rather than let a
|
||||
# half-written page look healthy.
|
||||
if finish_reason and "MAX_TOKENS" in finish_reason and content:
|
||||
logger.warning(
|
||||
"Gemini response truncated at max_output_tokens "
|
||||
f"(scope={scope}, model={self.model}, max_output_tokens={max_completion_tokens}, "
|
||||
f"output_tokens={output_tokens}, thoughts_tokens={thoughts_tokens}). The visible "
|
||||
"output was cut off; raise the cap or leave it unset for reasoning models."
|
||||
)
|
||||
|
||||
span_recorder = get_span_recorder()
|
||||
from hindsight_api.tracing import _serialize_for_span
|
||||
|
||||
|
||||
@@ -97,7 +97,7 @@ class LiteLLMLLM(LLMInterface):
|
||||
api_key: str,
|
||||
base_url: str,
|
||||
model: str,
|
||||
reasoning_effort: str = "low",
|
||||
reasoning_effort: str | None = None,
|
||||
timeout: float | None = None,
|
||||
extra_body: dict[str, Any] | None = None,
|
||||
bedrock_service_tier: str | None = None,
|
||||
@@ -185,6 +185,14 @@ class LiteLLMLLM(LLMInterface):
|
||||
kwargs["max_completion_tokens"] = self._cap_max_completion_tokens(max_completion_tokens)
|
||||
if temperature is not None:
|
||||
kwargs["temperature"] = temperature
|
||||
# LiteLLM translates reasoning_effort per target provider (Anthropic thinking
|
||||
# budgets, Gemini thinking config, OpenAI's flat param), so forwarding the
|
||||
# operator's setting is all that is needed to honour it here — dropping it was
|
||||
# a silent no-op on every model behind this lane (issue #3449). Only sent when
|
||||
# configured; ``litellm.drop_params = True`` discards it for models that have
|
||||
# no reasoning knob rather than raising.
|
||||
if self.reasoning_effort is not None:
|
||||
kwargs["reasoning_effort"] = self.reasoning_effort
|
||||
|
||||
# User-configured extras fill in only where the caller didn't set a value,
|
||||
# so explicit per-call params (model, messages, temperature, …) always win.
|
||||
|
||||
@@ -66,7 +66,7 @@ class LiteLLMRouterLLM(LiteLLMLLM):
|
||||
base_url: str,
|
||||
model: str,
|
||||
config: dict[str, Any],
|
||||
reasoning_effort: str = "low",
|
||||
reasoning_effort: str | None = None,
|
||||
timeout: float | None = None,
|
||||
**kwargs: Any,
|
||||
):
|
||||
@@ -146,6 +146,13 @@ class LiteLLMRouterLLM(LiteLLMLLM):
|
||||
kwargs["max_completion_tokens"] = self._cap_max_completion_tokens(max_completion_tokens)
|
||||
if temperature is not None:
|
||||
kwargs["temperature"] = temperature
|
||||
# Like api_key/base_url, per-deployment reasoning could live in the Router config,
|
||||
# but the operator-level setting is cross-cutting and LiteLLM translates it per
|
||||
# target provider — so this override forwards it exactly as the base provider does.
|
||||
# Omitting it made HINDSIGHT_API_*_REASONING_EFFORT a no-op on the router lane
|
||||
# alone (issue #3449); only sent when configured.
|
||||
if self.reasoning_effort is not None:
|
||||
kwargs["reasoning_effort"] = self.reasoning_effort
|
||||
|
||||
# Forward operator-configured default headers as ``extra_headers`` so they
|
||||
# reach the provider behind the Router (proxies / request-tracing middleware).
|
||||
|
||||
@@ -273,7 +273,8 @@ class LlamaCppLLM(LLMInterface):
|
||||
api_key: str,
|
||||
base_url: str,
|
||||
model: str,
|
||||
reasoning_effort: str = "low",
|
||||
reasoning_effort: str | None = None,
|
||||
extra_body: dict[str, Any] | None = None,
|
||||
model_path: str | None = None,
|
||||
gpu_layers: int = -1,
|
||||
context_size: int = 8192,
|
||||
@@ -289,6 +290,7 @@ class LlamaCppLLM(LLMInterface):
|
||||
model=model or DEFAULT_LLAMACPP_MODEL_ALIAS,
|
||||
reasoning_effort=reasoning_effort,
|
||||
)
|
||||
self._extra_body = extra_body
|
||||
self._model_path_str = model_path
|
||||
self._gpu_layers = gpu_layers
|
||||
self._context_size = context_size
|
||||
@@ -336,7 +338,10 @@ class LlamaCppLLM(LLMInterface):
|
||||
api_key="llamacpp",
|
||||
base_url=self._server.base_url,
|
||||
model=self.model,
|
||||
# None (unconfigured) must stay None so the delegate omits the parameter
|
||||
# rather than inventing a level for the local model.
|
||||
reasoning_effort=self.reasoning_effort,
|
||||
extra_body=self._extra_body,
|
||||
)
|
||||
|
||||
self._initialized = True
|
||||
|
||||
@@ -48,7 +48,7 @@ class MockLLM(LLMInterface):
|
||||
api_key: str,
|
||||
base_url: str,
|
||||
model: str,
|
||||
reasoning_effort: str = "low",
|
||||
reasoning_effort: str | None = None,
|
||||
**kwargs: Any,
|
||||
):
|
||||
"""
|
||||
|
||||
@@ -48,7 +48,7 @@ class NousLLM(OpenAICompatibleLLM):
|
||||
api_key: str, # Ignored — the token is read from ~/.hermes/auth.json
|
||||
base_url: str,
|
||||
model: str,
|
||||
reasoning_effort: str = "low",
|
||||
reasoning_effort: str | None = None,
|
||||
**kwargs: Any,
|
||||
):
|
||||
try:
|
||||
|
||||
@@ -190,8 +190,8 @@ def _strip_code_fences(content: str) -> str:
|
||||
# Reasoning/thinking tags emitted by extended-thinking models. Some providers
|
||||
# (e.g. MiniMax-M3) leak the chain-of-thought wrapped in these tags into the
|
||||
# response body instead of a separate reasoning_content field. Each entry is
|
||||
# (open_tag, close_tag); the open tag also matches when the close tag is missing
|
||||
# (truncated output) so a dangling block is removed to end-of-string.
|
||||
# (open_tag, close_tag); a line-start open tag also matches when the close tag is
|
||||
# missing (truncated output) so a dangling block is removed to end-of-string.
|
||||
_REASONING_TAG_PAIRS: tuple[tuple[str, str], ...] = (
|
||||
("<think>", "</think>"),
|
||||
("<thinking>", "</thinking>"),
|
||||
@@ -214,7 +214,11 @@ def _strip_reasoning_tags(text: str) -> str:
|
||||
Handles two cases:
|
||||
1. Closed blocks: ``<think>...</think>`` removed wherever they appear.
|
||||
2. Unclosed blocks: a dangling ``<think>`` with no closing tag (model output
|
||||
truncated mid-thought) is removed from the open tag to end-of-string.
|
||||
truncated mid-thought) is removed to end-of-string, but only when it starts
|
||||
its own line (line-start, possibly indented). Inline occurrences (e.g. a
|
||||
JSON value quoting ``<think>`` verbatim) are real content and must be kept
|
||||
-- an unanchored greedy ``.*`` to end-of-string deleted every inline tag
|
||||
plus all following content, surfacing as ``Unterminated string`` in retain.
|
||||
|
||||
Returns the input unchanged (modulo surrounding whitespace) when no tags are
|
||||
present.
|
||||
@@ -224,9 +228,14 @@ def _strip_reasoning_tags(text: str) -> str:
|
||||
for open_tag, close_tag in _REASONING_TAG_PAIRS:
|
||||
open_re = re.escape(open_tag)
|
||||
close_re = re.escape(close_tag)
|
||||
# Closed blocks first, then any remaining unclosed (truncated) block.
|
||||
# Closed blocks first.
|
||||
text = re.sub(rf"{open_re}.*?{close_re}", "", text, flags=re.DOTALL)
|
||||
text = re.sub(rf"{open_re}.*", "", text, flags=re.DOTALL)
|
||||
# Unclosed (truncated) blocks: strip only when the open tag starts its own
|
||||
# line, from there to end-of-string. The line-start anchor preserves inline
|
||||
# literals (e.g. a JSON value quoting ``<think>``) so valid JSON is not
|
||||
# corrupted; DOTALL-to-end still removes a multi-line truncated block whole,
|
||||
# so leaked reasoning does not survive past its first line.
|
||||
text = re.sub(rf"(^|\n)[ \t]*{open_re}.*", "", text, flags=re.DOTALL)
|
||||
return text.strip()
|
||||
|
||||
|
||||
@@ -532,7 +541,7 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
api_key: str,
|
||||
base_url: str,
|
||||
model: str,
|
||||
reasoning_effort: str = "low",
|
||||
reasoning_effort: str | None = None,
|
||||
timeout: float | None = None,
|
||||
groq_service_tier: str | None = None,
|
||||
extra_body: dict[str, Any] | None = None,
|
||||
@@ -684,8 +693,18 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
self._client = AsyncOpenAI(**client_kwargs)
|
||||
logger.info(
|
||||
f"OpenAI-compatible client initialized: provider={self.provider}, model={self.model}, "
|
||||
f"base_url={self.base_url or 'default'}"
|
||||
f"base_url={self.base_url or 'default'}, "
|
||||
f"reasoning_effort={self.reasoning_effort if self._sends_reasoning_effort() else 'not sent'}"
|
||||
)
|
||||
if self.reasoning_effort is not None and not self._sends_reasoning_effort():
|
||||
# Never drop a configured value silently: the variable is set, documented and
|
||||
# visible in the environment, so every signal the operator has says it is in
|
||||
# force. Saying so once at startup is what turns this into a seconds-long
|
||||
# diagnosis instead of a source-reading exercise (issue #3449).
|
||||
logger.warning(
|
||||
f"reasoning_effort={self.reasoning_effort!r} is not sent to the model: "
|
||||
f"{self.model!r} is a known non-reasoning model that rejects the parameter"
|
||||
)
|
||||
logger.debug(
|
||||
f"Cache affinity resolved: provider={self.provider}, base_url={self.base_url or 'default'}, "
|
||||
f"mode={self._cache_affinity.value}"
|
||||
@@ -727,8 +746,45 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Connection verification failed for {self.provider}/{self.model}: {e}") from e
|
||||
|
||||
def _sends_reasoning_effort(self) -> bool:
|
||||
"""Whether ``reasoning_effort`` is attached to requests.
|
||||
|
||||
The operator decides, not a model name. ``provider=openai`` with a custom base_url
|
||||
can serve any model under any name — vLLM, Ollama, llama.cpp, TGI — so the name
|
||||
carries no capability signal, and gating on it made every
|
||||
``HINDSIGHT_API_*_REASONING_EFFORT`` variable a silent no-op on exactly those
|
||||
deployments (issue #3449). Unset means unset: no level is invented for a model
|
||||
just because its name is recognisable.
|
||||
"""
|
||||
return self.reasoning_effort is not None and not self._rejects_reasoning_effort()
|
||||
|
||||
def _rejects_reasoning_effort(self) -> bool:
|
||||
"""Whether the model is a known product that rejects ``reasoning_effort`` outright.
|
||||
|
||||
The one place a name still overrides an explicit setting, and it matches only
|
||||
OpenAI's own non-reasoning products — names invented by OpenAI, so a self-hosted
|
||||
model is not going to collide with one by accident. Sending the parameter to
|
||||
gpt-4o is an immediate HTTP 400, so honouring the setting there would trade a
|
||||
silently ignored value for a hard failure. The drop is logged at startup.
|
||||
"""
|
||||
model_lower = self.model.lower()
|
||||
return any(x in model_lower for x in ["gpt-4o", "gpt-4.1", "gpt-4-", "gpt-3.5"])
|
||||
|
||||
def _supports_reasoning_model(self) -> bool:
|
||||
"""Check if the current model is a reasoning model (o1, o3, GPT-5, DeepSeek)."""
|
||||
"""Check if the current model is a reasoning model (o1, o3, GPT-5, DeepSeek).
|
||||
|
||||
**Deprecated as a capability check — this list is frozen. Do not add models to
|
||||
it.** Guessing capability from a name never worked outside OpenAI's own products:
|
||||
``provider=openai`` with a custom base_url serves anything under any name, so the
|
||||
list could only ever grow stale while silently discarding what operators asked
|
||||
for (issue #3449). Reasoning effort is now purely the operator's call, via
|
||||
``HINDSIGHT_API_LLM_REASONING_EFFORT`` and its per-operation variants — a new
|
||||
model needs configuration, not a new substring here.
|
||||
|
||||
All that is left is the request *shape* a recognised OpenAI reasoning model
|
||||
requires regardless of effort: the max-completion-tokens floor, the parameter
|
||||
name, temperature suppression.
|
||||
"""
|
||||
model_lower = self.model.lower()
|
||||
if "deepseek" in model_lower:
|
||||
# DeepSeek v4-flash is the non-thinking route. Treating every
|
||||
@@ -867,8 +923,8 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
temperature = max(0.01, min(temperature, 1.0))
|
||||
call_params["temperature"] = temperature
|
||||
|
||||
# Set reasoning_effort for reasoning models
|
||||
if is_reasoning_model:
|
||||
# Set reasoning_effort when configured, or for models recognised as reasoning models
|
||||
if self._sends_reasoning_effort():
|
||||
call_params["reasoning_effort"] = self.reasoning_effort
|
||||
|
||||
# Provider-specific parameters
|
||||
@@ -1304,7 +1360,7 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
# 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():
|
||||
if self._sends_reasoning_effort():
|
||||
call_params["reasoning_effort"] = self.reasoning_effort
|
||||
|
||||
# Provider-specific parameters
|
||||
|
||||
@@ -176,7 +176,7 @@ class OpenAIResponsesLLM(LLMInterface):
|
||||
api_key: str,
|
||||
base_url: str,
|
||||
model: str,
|
||||
reasoning_effort: str = "low",
|
||||
reasoning_effort: str | None = None,
|
||||
timeout: float | None = None,
|
||||
extra_body: dict[str, Any] | None = None,
|
||||
default_headers: dict[str, str] | None = None,
|
||||
@@ -455,7 +455,9 @@ class OpenAIResponsesLLM(LLMInterface):
|
||||
params["max_output_tokens"] = max_completion_tokens
|
||||
if temperature is not None and not is_reasoning_model:
|
||||
params["temperature"] = temperature
|
||||
if is_reasoning_model:
|
||||
# Only when the operator configured a level: unset means the model runs at the
|
||||
# Responses API's own default effort rather than one Hindsight picked.
|
||||
if is_reasoning_model and self.reasoning_effort is not None:
|
||||
params["reasoning"] = {"effort": self.reasoning_effort}
|
||||
if self.openai_service_tier:
|
||||
params["service_tier"] = self.openai_service_tier
|
||||
@@ -571,7 +573,9 @@ class OpenAIResponsesLLM(LLMInterface):
|
||||
params["max_output_tokens"] = max_completion_tokens
|
||||
if temperature is not None and not is_reasoning_model:
|
||||
params["temperature"] = temperature
|
||||
if is_reasoning_model:
|
||||
# Only when the operator configured a level: unset means the model runs at the
|
||||
# Responses API's own default effort rather than one Hindsight picked.
|
||||
if is_reasoning_model and self.reasoning_effort is not None:
|
||||
params["reasoning"] = {"effort": self.reasoning_effort}
|
||||
if self.openai_service_tier:
|
||||
params["service_tier"] = self.openai_service_tier
|
||||
|
||||
@@ -362,7 +362,7 @@ class XaiOAuthLLM(LLMInterface):
|
||||
api_key: str, # Ignored: the credential is the OAuth grant in the token store.
|
||||
base_url: str,
|
||||
model: str,
|
||||
reasoning_effort: str = "low",
|
||||
reasoning_effort: str | None = None,
|
||||
timeout: float | None = None,
|
||||
auth_manager: XaiOAuthManager | None = None,
|
||||
**kwargs: Any,
|
||||
|
||||
@@ -10,6 +10,7 @@ import re
|
||||
from abc import ABC, abstractmethod
|
||||
from datetime import datetime, timedelta
|
||||
|
||||
from dateparser.conf import Settings, apply_settings
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from hindsight_api.engine.temporal_periods import (
|
||||
@@ -91,6 +92,39 @@ _PERIOD_WORDS = {
|
||||
}
|
||||
|
||||
|
||||
# Every token _date_match_score can award points for, in one alternation.
|
||||
# Derived from the same four sets the scorer uses so the two cannot drift apart
|
||||
# (``test_prefilter_matches_scorer`` fails if a word is added to one and not the
|
||||
# other).
|
||||
_SCOREABLE_WORDS = _MONTH_WORDS | _RELATIVE_WORDS | _WEEKDAY_WORDS | _PERIOD_WORDS
|
||||
_SCOREABLE_RE = re.compile("[0-9]|" + "|".join(sorted(_SCOREABLE_WORDS)))
|
||||
_NON_ALNUM_RE = re.compile(r"[^a-z0-9]+")
|
||||
|
||||
|
||||
def _query_can_score(query: str) -> bool:
|
||||
"""Whether any span of ``query`` could score above zero.
|
||||
|
||||
``search_dates`` returns substrings of the *original* text (``translate_search``
|
||||
keeps parallel original/translated token streams and reports the original),
|
||||
and ``_date_match_score`` awards points only for an ASCII digit or one of the
|
||||
four English word sets. So if the query contains none of those anywhere, every
|
||||
match it could possibly return scores zero and ``analyze`` returns None —
|
||||
which means the entire dateparser search can be skipped without changing the
|
||||
answer.
|
||||
|
||||
This is deliberately an over-approximation: substring matching (rather than
|
||||
tokenised matching) means "maybe" counts as containing "may" and we take the
|
||||
slow path unnecessarily. That costs time, never correctness. The compacted
|
||||
second check covers the case where dateparser joins adjacent tokens and drops
|
||||
the separator between them, which could surface a word that is not contiguous
|
||||
in the raw text.
|
||||
"""
|
||||
low = query.lower()
|
||||
if _SCOREABLE_RE.search(low):
|
||||
return True
|
||||
return bool(_SCOREABLE_RE.search(_NON_ALNUM_RE.sub("", low)))
|
||||
|
||||
|
||||
def _date_match_score(text: str) -> int:
|
||||
"""Score how strong a temporal signal a matched span carries.
|
||||
|
||||
@@ -198,30 +232,76 @@ class DateparserQueryAnalyzer(QueryAnalyzer):
|
||||
(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}
|
||||
self._loaded = False
|
||||
self._locales = None
|
||||
self._exact_search = None
|
||||
|
||||
def load(self) -> None:
|
||||
"""Load dateparser and warm up internal data structures.
|
||||
|
||||
Triggers the real initialization cost (regex tables, timezone data) at
|
||||
load time so the first actual recall doesn't pay the cold-start penalty.
|
||||
Triggers the real initialization cost (locale dictionaries, timezone
|
||||
tables, the cached character tables used by detection) at load time so
|
||||
the first actual recall doesn't pay the cold-start penalty.
|
||||
"""
|
||||
if self._search_dates is None:
|
||||
from dateparser.search import search_dates
|
||||
if self._loaded:
|
||||
return
|
||||
|
||||
self._search_dates = search_dates
|
||||
# Warm up: fire a dummy call to trigger lazy-loaded internal tables.
|
||||
self._search_dates("today", **self._search_kwargs())
|
||||
from dateparser.conf import settings as dateparser_settings
|
||||
from dateparser.search import _search_with_detection
|
||||
from dateparser.search.search import _ExactLanguageSearch
|
||||
|
||||
available = _search_with_detection.available_language_map
|
||||
if self._languages is None:
|
||||
self._locales = list(available.values())
|
||||
else:
|
||||
unknown = set(self._languages) - set(available)
|
||||
if unknown:
|
||||
raise ValueError("Unknown language(s): %s" % ", ".join(map(repr, sorted(unknown))))
|
||||
self._locales = [available[code] for code in self._languages]
|
||||
|
||||
# Our own instance rather than dateparser's module-level singleton:
|
||||
# _ExactLanguageSearch caches the "current" locale on itself, so sharing
|
||||
# it across callers is a data race the moment this runs off the event
|
||||
# loop thread.
|
||||
self._exact_search = _ExactLanguageSearch(_search_with_detection.loader)
|
||||
self._loaded = True
|
||||
|
||||
# Warm the lazily-built locale dictionaries and the character tables.
|
||||
self._find_dates("today", settings=dateparser_settings)
|
||||
|
||||
@apply_settings
|
||||
def _find_dates(self, query: str, settings: "Settings | None" = None) -> list[tuple[str, datetime]] | None:
|
||||
"""``dateparser.search.search_dates`` without its redundant work.
|
||||
|
||||
Same three steps as upstream — preprocess, detect the language, parse the
|
||||
detected language's date expressions — but detection goes through
|
||||
:mod:`hindsight_api.engine.temporal_language_detection`, which is the same
|
||||
algorithm with the per-locale recomputation hoisted and memoised. See that
|
||||
module for why each step is equivalence-preserving, and
|
||||
``tests/test_temporal_extraction.py`` for the differential proof.
|
||||
"""
|
||||
from dateparser.conf import check_settings
|
||||
from dateparser.conf import settings as dateparser_defaults
|
||||
from dateparser.search import _search_with_detection
|
||||
|
||||
from .temporal_language_detection import best_language
|
||||
|
||||
# @apply_settings always injects a Settings (converting a dict if the
|
||||
# caller passed one); the None default only exists to satisfy its
|
||||
# keyword-argument contract. Fall back rather than assert so a direct
|
||||
# call without the decorator still behaves like dateparser's own entry
|
||||
# points.
|
||||
settings = settings or dateparser_defaults
|
||||
check_settings(settings)
|
||||
text = _search_with_detection.preprocess_text(query, self._languages)
|
||||
# Settings populates its attributes dynamically, so this is a getattr
|
||||
# rather than a plain access: the name is not statically visible.
|
||||
default_languages = getattr(settings, "DEFAULT_LANGUAGES", None)
|
||||
language = best_language(text, self._locales) or (default_languages[0] if default_languages else None)
|
||||
if not language:
|
||||
return None
|
||||
return self._exact_search.search_parse(language, text, settings=settings) or None
|
||||
|
||||
def analyze(self, query: str, reference_date: datetime | None = None) -> QueryAnalysis:
|
||||
"""
|
||||
@@ -249,6 +329,14 @@ class DateparserQueryAnalyzer(QueryAnalyzer):
|
||||
start_date, end_date = period_result
|
||||
return QueryAnalysis(temporal_constraint=TemporalConstraint(start_date=start_date, end_date=end_date))
|
||||
|
||||
# Cheap sound rejection before the expensive search. dateparser's
|
||||
# search_dates spends ~98% of its time detecting which of 205 locales the
|
||||
# text is in, and it is *slowest* when there is no date to find (every
|
||||
# locale runs to completion before concluding nothing matched). When no
|
||||
# span could score above zero, that entire cost buys a guaranteed None.
|
||||
if not _query_can_score(query):
|
||||
return QueryAnalysis(temporal_constraint=None)
|
||||
|
||||
# Lazy load dateparser (only imports on first call, then cached)
|
||||
self.load()
|
||||
|
||||
@@ -266,7 +354,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, **self._search_kwargs())
|
||||
results = self._find_dates(query, settings=settings)
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
"dateparser raised %s on query (treating as no temporal constraint): %s",
|
||||
|
||||
@@ -17,10 +17,15 @@ from ...config import get_config
|
||||
from ..llm_interface import LLM_TOOL_CHOICE_AUTO, LLMToolChoice
|
||||
from .models import DirectiveInfo, LLMCall, ReflectAgentResult, StructuredOutputResult, TokenUsageSummary, ToolCall
|
||||
from .prompts import (
|
||||
_SPLIT_SYNTHESIS_WARN_CHUNKS,
|
||||
CLAIMS_SYSTEM_PROMPT,
|
||||
_extract_directive_rules,
|
||||
build_chunk_claims_prompt,
|
||||
build_final_prompt,
|
||||
build_final_system_prompt,
|
||||
build_reduce_prompt,
|
||||
build_system_prompt_for_tools,
|
||||
split_context_history,
|
||||
)
|
||||
from .tokenization import count_cl100k_tokens
|
||||
from .tools_schema import get_reflect_tools
|
||||
@@ -462,7 +467,12 @@ async def _run_reflect_agent_inner(
|
||||
expand_fn: Tool callback for expand (memory_ids, depth) -> result
|
||||
context: Optional additional context
|
||||
max_iterations: Maximum number of iterations before forcing response
|
||||
max_tokens: Maximum tokens for the final response
|
||||
max_tokens: Desired *visible* length of the final answer. Communicated to
|
||||
the model as a soft directive and enforced by the post-hoc rewrite --
|
||||
NOT passed as the provider's ``max_completion_tokens``, which on
|
||||
thinking models is consumed by reasoning tokens and would truncate the
|
||||
answer mid-word (#3365). The transport-level cost cap is a separate,
|
||||
uncapped-by-default config (``reflect_max_completion_tokens``).
|
||||
response_schema: Optional JSON Schema for structured output in final response
|
||||
directives: Optional list of directive mental models to inject as hard rules
|
||||
|
||||
@@ -471,6 +481,13 @@ async def _run_reflect_agent_inner(
|
||||
"""
|
||||
start_time = time.time()
|
||||
|
||||
# Transport-level output cap for the synthesis calls. Decoupled from
|
||||
# ``max_tokens`` (a page-length target enforced via prompt + rewrite): None by
|
||||
# default so reasoning models run to a natural stop instead of truncating the
|
||||
# visible page mid-word (#3365). An operator can set a hard cost ceiling via
|
||||
# HINDSIGHT_API_REFLECT_MAX_COMPLETION_TOKENS.
|
||||
synthesis_max_completion_tokens = get_config().reflect_max_completion_tokens
|
||||
|
||||
# Build directives_applied for the trace
|
||||
directives_applied = _build_directives_applied(directives)
|
||||
|
||||
@@ -638,6 +655,106 @@ async def _run_reflect_agent_inner(
|
||||
f"total={elapsed_ms}ms"
|
||||
)
|
||||
|
||||
async def _tracked_llm_call(prompt: str, trace_scope: str, system_prompt: str, completion_cap: int | None) -> str:
|
||||
"""One tool-less LLM call with usage/trace accounting folded in."""
|
||||
nonlocal total_input_tokens, total_output_tokens, total_cached_tokens, total_thoughts_tokens
|
||||
llm_start = time.time()
|
||||
response, usage = await llm_config.call(
|
||||
messages=[
|
||||
{"role": "system", "content": system_prompt},
|
||||
{"role": "user", "content": prompt},
|
||||
],
|
||||
scope="reflect",
|
||||
max_completion_tokens=completion_cap,
|
||||
return_usage=True,
|
||||
)
|
||||
llm_duration = int((time.time() - llm_start) * 1000)
|
||||
total_input_tokens += usage.input_tokens
|
||||
total_output_tokens += usage.output_tokens
|
||||
total_cached_tokens += getattr(usage, "cached_tokens", 0) or 0
|
||||
total_thoughts_tokens += getattr(usage, "thoughts_tokens", 0) or 0
|
||||
llm_trace.append(
|
||||
{
|
||||
"scope": trace_scope,
|
||||
"duration_ms": llm_duration,
|
||||
"input_tokens": usage.input_tokens,
|
||||
"output_tokens": usage.output_tokens,
|
||||
}
|
||||
)
|
||||
return response.strip()
|
||||
|
||||
async def _forced_final_synthesis(iterations_completed: int) -> ReflectAgentResult:
|
||||
"""Answer without tools from the accumulated tool results.
|
||||
|
||||
When the accumulated results fit the prompt budget this is one LLM call,
|
||||
exactly as before. When they exceed it, they are SPLIT — not truncated:
|
||||
each budget-sized chunk is compressed in parallel into dated, cited
|
||||
claims, and one reduce call synthesizes the answer from every chunk's
|
||||
claims. The old behavior dropped any over-budget block whole (plus all
|
||||
older ones), which produced confident "no information" answers carrying
|
||||
hundreds of citations the synthesis model never saw (#3122).
|
||||
"""
|
||||
nonlocal total_input_tokens, total_output_tokens, total_cached_tokens, total_thoughts_tokens
|
||||
final_system = build_final_system_prompt(bank_profile.get("mission"), llm_output_language, directives)
|
||||
chunks = split_context_history(context_history, max_context_tokens)
|
||||
# Every call below uses the transport-level cap, never the caller's
|
||||
# max_tokens: that is a visible-length target carried as a prompt
|
||||
# directive (#3365), and capping the transport with it would truncate
|
||||
# thinking models mid-word — or, on the map calls, starve the evidence
|
||||
# extraction.
|
||||
if len(chunks) <= 1:
|
||||
prompt = build_final_prompt(
|
||||
query,
|
||||
context_history,
|
||||
bank_profile,
|
||||
context,
|
||||
max_context_tokens=max_context_tokens,
|
||||
max_tokens=max_tokens,
|
||||
)
|
||||
answer = await _tracked_llm_call(prompt, "final", final_system, synthesis_max_completion_tokens)
|
||||
else:
|
||||
log = logger.warning if len(chunks) > _SPLIT_SYNTHESIS_WARN_CHUNKS else logger.info
|
||||
log(
|
||||
f"[REFLECT {reflect_id}] Retrieved data exceeds the context budget; "
|
||||
f"split synthesis over {len(chunks)} chunks."
|
||||
)
|
||||
# Map: each chunk in parallel.
|
||||
claim_sections = await asyncio.gather(
|
||||
*(
|
||||
_tracked_llm_call(
|
||||
build_chunk_claims_prompt(query, chunk),
|
||||
f"final_map_{i}",
|
||||
CLAIMS_SYSTEM_PROMPT,
|
||||
synthesis_max_completion_tokens,
|
||||
)
|
||||
for i, chunk in enumerate(chunks, 1)
|
||||
)
|
||||
)
|
||||
# Reduce: one synthesis call over every chunk's claims.
|
||||
prompt = build_reduce_prompt(query, list(claim_sections), bank_profile, context, max_tokens=max_tokens)
|
||||
answer = await _tracked_llm_call(prompt, "final", final_system, synthesis_max_completion_tokens)
|
||||
|
||||
structured_output = None
|
||||
if response_schema and answer:
|
||||
struct = await _generate_structured_output(answer, response_schema, llm_config, reflect_id, max_tokens)
|
||||
structured_output = struct.structured_output
|
||||
total_input_tokens += struct.input_tokens
|
||||
total_output_tokens += struct.output_tokens
|
||||
total_cached_tokens += struct.cached_tokens
|
||||
total_thoughts_tokens += struct.thoughts_tokens
|
||||
|
||||
_log_completion(answer, iterations_completed, forced=True)
|
||||
return ReflectAgentResult(
|
||||
text=answer,
|
||||
structured_output=structured_output,
|
||||
iterations=iterations_completed,
|
||||
tools_called=total_tools_called,
|
||||
tool_trace=tool_trace,
|
||||
llm_trace=_get_llm_trace(),
|
||||
usage=_get_usage(),
|
||||
directives_applied=directives_applied,
|
||||
)
|
||||
|
||||
consecutive_errors = 0
|
||||
# When a forced ``search_mental_models`` returns fresh, usable models on a
|
||||
# low/mid-budget call, we stop forcing the lower retrieval layers from this
|
||||
@@ -656,60 +773,7 @@ async def _run_reflect_agent_inner(
|
||||
|
||||
if is_last:
|
||||
# Force text response on last iteration - no tools
|
||||
prompt = build_final_prompt(
|
||||
query, context_history, bank_profile, context, max_context_tokens=max_context_tokens
|
||||
)
|
||||
llm_start = time.time()
|
||||
response, usage = await llm_config.call(
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": build_final_system_prompt(
|
||||
bank_profile.get("mission"), llm_output_language, directives
|
||||
),
|
||||
},
|
||||
{"role": "user", "content": prompt},
|
||||
],
|
||||
scope="reflect",
|
||||
max_completion_tokens=max_tokens,
|
||||
return_usage=True,
|
||||
)
|
||||
llm_duration = int((time.time() - llm_start) * 1000)
|
||||
total_input_tokens += usage.input_tokens
|
||||
total_output_tokens += usage.output_tokens
|
||||
total_cached_tokens += getattr(usage, "cached_tokens", 0) or 0
|
||||
total_thoughts_tokens += getattr(usage, "thoughts_tokens", 0) or 0
|
||||
llm_trace.append(
|
||||
{
|
||||
"scope": "final",
|
||||
"duration_ms": llm_duration,
|
||||
"input_tokens": usage.input_tokens,
|
||||
"output_tokens": usage.output_tokens,
|
||||
}
|
||||
)
|
||||
answer = response.strip()
|
||||
|
||||
# Generate structured output if schema provided
|
||||
structured_output = None
|
||||
if response_schema and answer:
|
||||
struct = await _generate_structured_output(answer, response_schema, llm_config, reflect_id, max_tokens)
|
||||
structured_output = struct.structured_output
|
||||
total_input_tokens += struct.input_tokens
|
||||
total_output_tokens += struct.output_tokens
|
||||
total_cached_tokens += struct.cached_tokens
|
||||
total_thoughts_tokens += struct.thoughts_tokens
|
||||
|
||||
_log_completion(answer, iteration + 1, forced=True)
|
||||
return ReflectAgentResult(
|
||||
text=answer,
|
||||
structured_output=structured_output,
|
||||
iterations=iteration + 1,
|
||||
tools_called=total_tools_called,
|
||||
tool_trace=tool_trace,
|
||||
llm_trace=_get_llm_trace(),
|
||||
usage=_get_usage(),
|
||||
directives_applied=directives_applied,
|
||||
)
|
||||
return await _forced_final_synthesis(iteration + 1)
|
||||
|
||||
# Proactive context-window guard: if accumulated messages would exceed the
|
||||
# configured token budget, bail out early and synthesize from what we have.
|
||||
@@ -721,59 +785,7 @@ async def _run_reflect_agent_inner(
|
||||
f"[REFLECT {reflect_id}] Context budget exceeded on iteration {iteration + 1}: "
|
||||
f"~{estimated_tokens} tokens >= {max_context_tokens} limit. Forcing final synthesis."
|
||||
)
|
||||
prompt = build_final_prompt(
|
||||
query, context_history, bank_profile, context, max_context_tokens=max_context_tokens
|
||||
)
|
||||
llm_start = time.time()
|
||||
response, usage = await llm_config.call(
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": build_final_system_prompt(
|
||||
bank_profile.get("mission"), llm_output_language, directives
|
||||
),
|
||||
},
|
||||
{"role": "user", "content": prompt},
|
||||
],
|
||||
scope="reflect",
|
||||
max_completion_tokens=max_tokens,
|
||||
return_usage=True,
|
||||
)
|
||||
llm_duration = int((time.time() - llm_start) * 1000)
|
||||
total_input_tokens += usage.input_tokens
|
||||
total_output_tokens += usage.output_tokens
|
||||
total_cached_tokens += getattr(usage, "cached_tokens", 0) or 0
|
||||
total_thoughts_tokens += getattr(usage, "thoughts_tokens", 0) or 0
|
||||
llm_trace.append(
|
||||
{
|
||||
"scope": "final",
|
||||
"duration_ms": llm_duration,
|
||||
"input_tokens": usage.input_tokens,
|
||||
"output_tokens": usage.output_tokens,
|
||||
}
|
||||
)
|
||||
answer = response.strip()
|
||||
|
||||
structured_output = None
|
||||
if response_schema and answer:
|
||||
struct = await _generate_structured_output(answer, response_schema, llm_config, reflect_id, max_tokens)
|
||||
structured_output = struct.structured_output
|
||||
total_input_tokens += struct.input_tokens
|
||||
total_output_tokens += struct.output_tokens
|
||||
total_cached_tokens += struct.cached_tokens
|
||||
total_thoughts_tokens += struct.thoughts_tokens
|
||||
|
||||
_log_completion(answer, iteration + 1, forced=True)
|
||||
return ReflectAgentResult(
|
||||
text=answer,
|
||||
structured_output=structured_output,
|
||||
iterations=iteration + 1,
|
||||
tools_called=total_tools_called,
|
||||
tool_trace=tool_trace,
|
||||
llm_trace=_get_llm_trace(),
|
||||
usage=_get_usage(),
|
||||
directives_applied=directives_applied,
|
||||
)
|
||||
return await _forced_final_synthesis(iteration + 1)
|
||||
|
||||
# Call LLM with tools
|
||||
llm_start = time.time()
|
||||
@@ -864,60 +876,7 @@ async def _run_reflect_agent_inner(
|
||||
# For other errors: retry if no evidence yet (but cap consecutive errors to avoid long hangs)
|
||||
elif not has_gathered_evidence and iteration < max_iterations - 1 and consecutive_errors < 2:
|
||||
continue
|
||||
prompt = build_final_prompt(
|
||||
query, context_history, bank_profile, context, max_context_tokens=max_context_tokens
|
||||
)
|
||||
llm_start = time.time()
|
||||
response, usage = await llm_config.call(
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": build_final_system_prompt(
|
||||
bank_profile.get("mission"), llm_output_language, directives
|
||||
),
|
||||
},
|
||||
{"role": "user", "content": prompt},
|
||||
],
|
||||
scope="reflect",
|
||||
max_completion_tokens=max_tokens,
|
||||
return_usage=True,
|
||||
)
|
||||
llm_duration = int((time.time() - llm_start) * 1000)
|
||||
total_input_tokens += usage.input_tokens
|
||||
total_output_tokens += usage.output_tokens
|
||||
total_cached_tokens += getattr(usage, "cached_tokens", 0) or 0
|
||||
total_thoughts_tokens += getattr(usage, "thoughts_tokens", 0) or 0
|
||||
llm_trace.append(
|
||||
{
|
||||
"scope": "final",
|
||||
"duration_ms": llm_duration,
|
||||
"input_tokens": usage.input_tokens,
|
||||
"output_tokens": usage.output_tokens,
|
||||
}
|
||||
)
|
||||
answer = response.strip()
|
||||
|
||||
# Generate structured output if schema provided
|
||||
structured_output = None
|
||||
if response_schema and answer:
|
||||
struct = await _generate_structured_output(answer, response_schema, llm_config, reflect_id, max_tokens)
|
||||
structured_output = struct.structured_output
|
||||
total_input_tokens += struct.input_tokens
|
||||
total_output_tokens += struct.output_tokens
|
||||
total_cached_tokens += struct.cached_tokens
|
||||
total_thoughts_tokens += struct.thoughts_tokens
|
||||
|
||||
_log_completion(answer, iteration + 1, forced=True)
|
||||
return ReflectAgentResult(
|
||||
text=answer,
|
||||
structured_output=structured_output,
|
||||
iterations=iteration + 1,
|
||||
tools_called=total_tools_called,
|
||||
tool_trace=tool_trace,
|
||||
llm_trace=_get_llm_trace(),
|
||||
usage=_get_usage(),
|
||||
directives_applied=directives_applied,
|
||||
)
|
||||
return await _forced_final_synthesis(iteration + 1)
|
||||
|
||||
# No tool calls this turn.
|
||||
if not result.tool_calls:
|
||||
@@ -942,60 +901,7 @@ async def _run_reflect_agent_inner(
|
||||
)
|
||||
# Model tool-called earlier and is now stopping: fall through to a clean
|
||||
# forced final synthesis (tools disabled, prose expected).
|
||||
prompt = build_final_prompt(
|
||||
query, context_history, bank_profile, context, max_context_tokens=max_context_tokens
|
||||
)
|
||||
llm_start = time.time()
|
||||
response, usage = await llm_config.call(
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": build_final_system_prompt(
|
||||
bank_profile.get("mission"), llm_output_language, directives
|
||||
),
|
||||
},
|
||||
{"role": "user", "content": prompt},
|
||||
],
|
||||
scope="reflect",
|
||||
max_completion_tokens=max_tokens,
|
||||
return_usage=True,
|
||||
)
|
||||
llm_duration = int((time.time() - llm_start) * 1000)
|
||||
total_input_tokens += usage.input_tokens
|
||||
total_output_tokens += usage.output_tokens
|
||||
total_cached_tokens += getattr(usage, "cached_tokens", 0) or 0
|
||||
total_thoughts_tokens += getattr(usage, "thoughts_tokens", 0) or 0
|
||||
llm_trace.append(
|
||||
{
|
||||
"scope": "final",
|
||||
"duration_ms": llm_duration,
|
||||
"input_tokens": usage.input_tokens,
|
||||
"output_tokens": usage.output_tokens,
|
||||
}
|
||||
)
|
||||
answer = response.strip()
|
||||
|
||||
# Generate structured output if schema provided
|
||||
structured_output = None
|
||||
if response_schema and answer:
|
||||
struct = await _generate_structured_output(answer, response_schema, llm_config, reflect_id, max_tokens)
|
||||
structured_output = struct.structured_output
|
||||
total_input_tokens += struct.input_tokens
|
||||
total_output_tokens += struct.output_tokens
|
||||
total_cached_tokens += struct.cached_tokens
|
||||
total_thoughts_tokens += struct.thoughts_tokens
|
||||
|
||||
_log_completion(answer, iteration + 1, forced=True)
|
||||
return ReflectAgentResult(
|
||||
text=answer,
|
||||
structured_output=structured_output,
|
||||
iterations=iteration + 1,
|
||||
tools_called=total_tools_called,
|
||||
tool_trace=tool_trace,
|
||||
llm_trace=_get_llm_trace(),
|
||||
usage=_get_usage(),
|
||||
directives_applied=directives_applied,
|
||||
)
|
||||
return await _forced_final_synthesis(iteration + 1)
|
||||
|
||||
# The model produced at least one tool call reflect could parse: it can
|
||||
# drive the loop, so a later text-only turn is a legitimate stop, not a
|
||||
@@ -1287,6 +1193,10 @@ async def _process_done_tool(
|
||||
final_usage = usage
|
||||
if llm_config and max_tokens is not None and count_cl100k_tokens(answer) > max_tokens:
|
||||
rewrite_start = time.time()
|
||||
# The token budget is enforced via the prompt, not a hard provider cap:
|
||||
# on thinking models a hard cap is eaten by reasoning tokens and would
|
||||
# truncate the rewrite mid-word (#3365). Cost is bounded by the separate
|
||||
# reflect_max_completion_tokens config (uncapped by default).
|
||||
rewritten, rewrite_usage = await llm_config.call(
|
||||
messages=[
|
||||
{
|
||||
@@ -1303,7 +1213,7 @@ async def _process_done_tool(
|
||||
},
|
||||
],
|
||||
scope="reflect",
|
||||
max_completion_tokens=max_tokens,
|
||||
max_completion_tokens=get_config().reflect_max_completion_tokens,
|
||||
return_usage=True,
|
||||
)
|
||||
answer = rewritten.strip()
|
||||
|
||||
@@ -443,25 +443,153 @@ def build_system_prompt_for_tools(
|
||||
return "\n".join(parts)
|
||||
|
||||
|
||||
def build_final_prompt(
|
||||
query: str,
|
||||
context_history: list[dict],
|
||||
bank_profile: dict,
|
||||
additional_context: str | None = None,
|
||||
max_context_tokens: int = 100_000,
|
||||
) -> str:
|
||||
"""Build the final prompt when forcing a text response (no tools)."""
|
||||
parts = []
|
||||
#: Result-list keys a tool output can carry; an over-budget block is split on
|
||||
#: these entry boundaries so no retrieved evidence is dropped.
|
||||
_SPLITTABLE_RESULT_KEYS = ("observations", "memories", "results")
|
||||
|
||||
# Bank identity
|
||||
#: Above this many synthesis chunks the retrieval volume is pathological
|
||||
#: (each chunk is ~0.8 * max_context_tokens); we still process everything,
|
||||
#: but loudly, so the real cause (an unbounded tool result) gets looked at.
|
||||
_SPLIT_SYNTHESIS_WARN_CHUNKS = 4
|
||||
|
||||
#: Floor for the per-chunk budget during splitting. A tiny configured
|
||||
#: ``max_context_tokens`` (tests use 1) would otherwise shred the history into
|
||||
#: one chunk per result entry — an LLM call per fact. A ~1k-token prompt is
|
||||
#: safe for any real model, so the floor caps fan-out without dropping data.
|
||||
_MIN_SPLIT_CHUNK_TOKENS = 1024
|
||||
|
||||
_FINAL_INSTRUCTIONS = (
|
||||
"Provide a thoughtful answer by synthesizing and reasoning from the retrieved data above. "
|
||||
"You can make reasonable inferences from the memories, but don't completely fabricate information. "
|
||||
"If the exact answer isn't stated, use what IS stated to give the best possible answer. "
|
||||
"Only say 'I don't have information' if the retrieved data is truly unrelated to the question.\n\n"
|
||||
"IMPORTANT: Output ONLY the final answer. Do NOT include meta-commentary like "
|
||||
'"I\'ll search..." or "Let me analyze...". Do NOT explain your reasoning process. '
|
||||
"Just provide the direct synthesized answer."
|
||||
)
|
||||
|
||||
|
||||
def _render_history_block(entry: dict) -> str:
|
||||
"""Render one context-history entry as a fenced JSON block."""
|
||||
tool = entry["tool"]
|
||||
output = entry["output"]
|
||||
try:
|
||||
output_str = json.dumps(output, indent=2, default=str, ensure_ascii=False)
|
||||
except (TypeError, ValueError):
|
||||
output_str = str(output)
|
||||
return f"\n### From {tool}:\n```json\n{output_str}\n```"
|
||||
|
||||
|
||||
def _cut_entry_to_budget(entry: dict, token_budget: int) -> dict:
|
||||
"""Token-bound one indivisible over-budget entry by cutting its serialized text.
|
||||
|
||||
Only reachable when a single result entry (or a list-less output like a
|
||||
document expand) alone exceeds the whole per-chunk budget — the one case
|
||||
where "split, don't drop" cannot be honored without exceeding the model's
|
||||
window. The cut text is wrapped back into an output dict so the entry
|
||||
renders like any other block.
|
||||
"""
|
||||
output = entry["output"]
|
||||
try:
|
||||
output_str = json.dumps(output, indent=2, default=str, ensure_ascii=False)
|
||||
except (TypeError, ValueError):
|
||||
output_str = str(output)
|
||||
tokens = count_cl100k_tokens(output_str)
|
||||
while output_str and tokens > token_budget:
|
||||
# Proportional shrink with a safety margin; the loop guards against the
|
||||
# estimate landing high, and always makes progress.
|
||||
keep = min(len(output_str) - 1, max(1, int(len(output_str) * token_budget / tokens * 0.95)))
|
||||
output_str = output_str[:keep]
|
||||
tokens = count_cl100k_tokens(output_str)
|
||||
return {**entry, "output": {"truncated": True, "content": output_str}}
|
||||
|
||||
|
||||
def split_context_history(context_history: list[dict], max_context_tokens: int) -> list[list[dict]]:
|
||||
"""Partition tool-result history into chunks that each fit the prompt budget.
|
||||
|
||||
Greedy chronological packing: blocks keep their order, and a chunk closes
|
||||
when the next block would push its rendered size past the budget. A single
|
||||
block bigger than the whole budget is split on result-entry boundaries
|
||||
(``observations``/``memories``/``results``) into synthetic partial blocks,
|
||||
so evidence is split across chunks rather than dropped — the failure mode
|
||||
of the old ``break`` was answering from nothing while citing everything
|
||||
(#3122). Only an *indivisible* over-budget entry gets token-cut.
|
||||
|
||||
Returns at least one chunk when history is non-empty; every original
|
||||
result entry appears in exactly one chunk.
|
||||
"""
|
||||
budget = max(_MIN_SPLIT_CHUNK_TOKENS, int(max_context_tokens * _FINAL_PROMPT_CONTEXT_FRACTION))
|
||||
chunks: list[list[dict]] = []
|
||||
current: list[dict] = []
|
||||
current_tokens = 0
|
||||
|
||||
def _close_current() -> None:
|
||||
nonlocal current, current_tokens
|
||||
if current:
|
||||
chunks.append(current)
|
||||
current = []
|
||||
current_tokens = 0
|
||||
|
||||
def _append_block(entry: dict, tokens: int) -> None:
|
||||
nonlocal current_tokens
|
||||
if current and current_tokens + tokens > budget:
|
||||
_close_current()
|
||||
current.append(entry)
|
||||
current_tokens += tokens
|
||||
|
||||
for entry in context_history:
|
||||
tokens = count_cl100k_tokens(_render_history_block(entry))
|
||||
if tokens <= budget:
|
||||
_append_block(entry, tokens)
|
||||
continue
|
||||
|
||||
# Over-budget block: split it on result-entry boundaries.
|
||||
output = entry["output"]
|
||||
split_key = next(
|
||||
(
|
||||
k
|
||||
for k in _SPLITTABLE_RESULT_KEYS
|
||||
if isinstance(output, dict) and isinstance(output.get(k), list) and output.get(k)
|
||||
),
|
||||
None,
|
||||
)
|
||||
if split_key is None:
|
||||
cut = _cut_entry_to_budget(entry, budget)
|
||||
_append_block(cut, count_cl100k_tokens(_render_history_block(cut)))
|
||||
continue
|
||||
|
||||
items = output[split_key]
|
||||
piece: list = []
|
||||
for item in items:
|
||||
candidate = {**entry, "output": {**output, split_key: piece + [item]}}
|
||||
if piece and count_cl100k_tokens(_render_history_block(candidate)) > budget:
|
||||
partial = {**entry, "output": {**output, split_key: piece}}
|
||||
_append_block(partial, count_cl100k_tokens(_render_history_block(partial)))
|
||||
piece = []
|
||||
candidate = {**entry, "output": {**output, split_key: [item]}}
|
||||
single_tokens = count_cl100k_tokens(_render_history_block(candidate))
|
||||
if not piece and single_tokens > budget:
|
||||
cut = _cut_entry_to_budget({**entry, "output": {**output, split_key: [item]}}, budget)
|
||||
_append_block(cut, count_cl100k_tokens(_render_history_block(cut)))
|
||||
else:
|
||||
piece.append(item)
|
||||
if piece:
|
||||
partial = {**entry, "output": {**output, split_key: piece}}
|
||||
_append_block(partial, count_cl100k_tokens(_render_history_block(partial)))
|
||||
|
||||
_close_current()
|
||||
return chunks
|
||||
|
||||
|
||||
def _bank_identity_section(bank_profile: dict, additional_context: str | None) -> list[str]:
|
||||
"""The shared bank-identity/disposition/context head of a synthesis prompt."""
|
||||
name = bank_profile.get("name", "Assistant")
|
||||
mission = bank_profile.get("mission", "")
|
||||
|
||||
parts.append(f"## Memory Bank Context\nName: {name}")
|
||||
parts = [f"## Memory Bank Context\nName: {name}"]
|
||||
if mission:
|
||||
parts.append(f"Mission: {mission}")
|
||||
|
||||
# Disposition traits if present
|
||||
disposition = bank_profile.get("disposition", {})
|
||||
if disposition:
|
||||
traits = []
|
||||
@@ -474,9 +602,51 @@ def build_final_prompt(
|
||||
if traits:
|
||||
parts.append(f"Disposition: {', '.join(traits)}")
|
||||
|
||||
# Additional context from caller
|
||||
if additional_context:
|
||||
parts.append(f"\n## Additional Context\n{additional_context}")
|
||||
return parts
|
||||
|
||||
|
||||
def _length_directive(max_tokens: int | None) -> str | None:
|
||||
"""Soft visible-length directive for a synthesis prompt, or None.
|
||||
|
||||
``max_tokens`` is the desired *visible* length of the answer (e.g. a mental
|
||||
model page's ``max_tokens``). It is communicated as a prompt directive
|
||||
rather than enforced by truncating the provider call: on thinking models the
|
||||
provider budget is consumed by reasoning tokens, so a hard cap cuts the page
|
||||
off mid-word (#3365). The hard length guarantee is the post-hoc rewrite in
|
||||
the agent; this directive just steers the model toward the target so the
|
||||
rewrite rarely has to fire.
|
||||
"""
|
||||
if max_tokens is None:
|
||||
return None
|
||||
return (
|
||||
"\n## Length\n"
|
||||
f"Aim for a complete, self-contained answer of approximately {max_tokens} tokens. "
|
||||
"Finishing cleanly matters more than length: end on a complete sentence and NEVER stop "
|
||||
"mid-word, mid-list, or mid-code-fence. If you near the budget, wrap up gracefully rather "
|
||||
"than cutting off."
|
||||
)
|
||||
|
||||
|
||||
def build_final_prompt(
|
||||
query: str,
|
||||
context_history: list[dict],
|
||||
bank_profile: dict,
|
||||
additional_context: str | None = None,
|
||||
max_context_tokens: int = 100_000,
|
||||
max_tokens: int | None = None,
|
||||
) -> str:
|
||||
"""Build the final prompt when forcing a text response (no tools).
|
||||
|
||||
``max_tokens`` is the soft visible-length target (see ``_length_directive``).
|
||||
|
||||
Callers overflow-proof this via ``split_context_history``: when the whole
|
||||
history fits one chunk this renders it directly, and the per-block budget
|
||||
walk below never trims. (The walk is kept as a defensive bound for direct
|
||||
callers that skip splitting.)
|
||||
"""
|
||||
parts = _bank_identity_section(bank_profile, additional_context)
|
||||
|
||||
# Tool call history — include as many entries as fit within the token budget,
|
||||
# preferring the most recent calls (they tend to be the most targeted).
|
||||
@@ -487,13 +657,7 @@ def build_final_prompt(
|
||||
rendered: list[str] = []
|
||||
truncated = False
|
||||
for entry in reversed(context_history):
|
||||
tool = entry["tool"]
|
||||
output = entry["output"]
|
||||
try:
|
||||
output_str = json.dumps(output, indent=2, default=str, ensure_ascii=False)
|
||||
except (TypeError, ValueError):
|
||||
output_str = str(output)
|
||||
block = f"\n### From {tool}:\n```json\n{output_str}\n```"
|
||||
block = _render_history_block(entry)
|
||||
block_tokens = count_cl100k_tokens(block)
|
||||
if block_tokens > token_budget:
|
||||
truncated = True
|
||||
@@ -511,16 +675,92 @@ def build_final_prompt(
|
||||
parts.append(f"\n## Question\n{query}")
|
||||
|
||||
# Final instructions
|
||||
parts.append("\n## Instructions\n" + _FINAL_INSTRUCTIONS)
|
||||
|
||||
length_directive = _length_directive(max_tokens)
|
||||
if length_directive is not None:
|
||||
parts.append(length_directive)
|
||||
|
||||
return "\n".join(parts)
|
||||
|
||||
|
||||
#: System prompt for the intermediate (map) calls of split synthesis. They do
|
||||
#: NOT answer the question — they compress one chunk of retrieved data into
|
||||
#: dated, cited claims that the reduce call can reason over. Dates and ids are
|
||||
#: mandatory because conflicting facts can land in different chunks: only the
|
||||
#: reduce call sees every chunk's claims, and it needs each claim's
|
||||
#: ``mentioned_at`` to apply the latest-statement-wins supersession rule.
|
||||
CLAIMS_SYSTEM_PROMPT = (
|
||||
"You extract evidence from retrieved memory data. You MUST ONLY use information "
|
||||
"from the provided data. NEVER make up names, people, events, or entities.\n\n"
|
||||
"Output a markdown bulleted list of factual claims relevant to the question. For EVERY claim:\n"
|
||||
"- state the fact in one sentence, in the same language as the question;\n"
|
||||
"- append its provenance in parentheses, exactly: "
|
||||
"(mentioned_at: <ISO date or unknown>; occurred: <ISO date/range or unknown>; memory_ids: <comma-separated ids>)\n\n"
|
||||
"Rules:\n"
|
||||
"- Be exhaustive over RELEVANT evidence; skip clearly irrelevant entries.\n"
|
||||
"- Do NOT synthesize, conclude, resolve conflicts, or answer the question — report conflicting "
|
||||
"claims as separate bullets with their dates; a later pass reconciles them.\n"
|
||||
"- Copy memory ids exactly as they appear in the data.\n"
|
||||
"- If nothing in the data is relevant, output exactly: (no relevant evidence)"
|
||||
)
|
||||
|
||||
|
||||
def build_chunk_claims_prompt(query: str, chunk: list[dict]) -> str:
|
||||
"""Build the user prompt for one intermediate (map) call of split synthesis."""
|
||||
parts = ["## Retrieved Data (extract relevant claims from this data)"]
|
||||
for entry in chunk:
|
||||
parts.append(_render_history_block(entry))
|
||||
parts.append(f"\n## Question\n{query}")
|
||||
parts.append(
|
||||
"\n## Instructions\n"
|
||||
"Provide a thoughtful answer by synthesizing and reasoning from the retrieved data above. "
|
||||
"You can make reasonable inferences from the memories, but don't completely fabricate information. "
|
||||
"If the exact answer isn't stated, use what IS stated to give the best possible answer. "
|
||||
"Only say 'I don't have information' if the retrieved data is truly unrelated to the question.\n\n"
|
||||
"IMPORTANT: Output ONLY the final answer. Do NOT include meta-commentary like "
|
||||
'"I\'ll search..." or "Let me analyze...". Do NOT explain your reasoning process. '
|
||||
"Just provide the direct synthesized answer."
|
||||
"List every claim in the retrieved data relevant to the question, one bullet per claim, "
|
||||
"each with its (mentioned_at: ...; occurred: ...; memory_ids: ...) provenance. "
|
||||
"Do not answer the question."
|
||||
)
|
||||
return "\n".join(parts)
|
||||
|
||||
|
||||
def build_reduce_prompt(
|
||||
query: str,
|
||||
claim_sections: list[str],
|
||||
bank_profile: dict,
|
||||
additional_context: str | None = None,
|
||||
max_tokens: int | None = None,
|
||||
) -> str:
|
||||
"""Build the final prompt that synthesizes the answer from per-chunk claims.
|
||||
|
||||
The retrieved data exceeded the context budget, so it was split into chunks
|
||||
and each chunk was compressed to dated, cited claims by a parallel LLM call.
|
||||
This prompt hands ALL the claim sets to one model. Conflicting facts may sit
|
||||
in different sections — that is why the claims carry ``mentioned_at``: the
|
||||
supersession rule (latest statement wins) must be applied across sections,
|
||||
not within one.
|
||||
"""
|
||||
parts = _bank_identity_section(bank_profile, additional_context)
|
||||
|
||||
parts.append(
|
||||
"\n## Retrieved Evidence (synthesize and reason from these claims)\n"
|
||||
"The retrieved data was processed in parallel passes; each section below holds one pass's "
|
||||
"extracted claims with provenance dates and memory ids. Treat the sections as ONE evidence "
|
||||
"pool: related and conflicting claims may appear in different sections."
|
||||
)
|
||||
for i, section in enumerate(claim_sections, 1):
|
||||
parts.append(f"\n### Evidence pass {i}:\n{section}")
|
||||
|
||||
parts.append(f"\n## Question\n{query}")
|
||||
|
||||
parts.append(
|
||||
"\n## Instructions\n"
|
||||
"When claims about the same fact conflict, the claim with the LATEST mentioned_at date is "
|
||||
"authoritative — later statements supersede earlier ones, regardless of which section they "
|
||||
"appear in. If equally-recent claims disagree and nothing resolves them, say so explicitly "
|
||||
"rather than picking one.\n\n" + _FINAL_INSTRUCTIONS
|
||||
)
|
||||
|
||||
length_directive = _length_directive(max_tokens)
|
||||
if length_directive is not None:
|
||||
parts.append(length_directive)
|
||||
|
||||
return "\n".join(parts)
|
||||
|
||||
|
||||
@@ -111,7 +111,7 @@ async def tool_search_mental_models(
|
||||
query_embedding: Pre-computed embedding for semantic search
|
||||
max_results: Maximum number of mental models to return
|
||||
tags: Optional tags to filter mental models
|
||||
tags_match: How to match tags - "any" (OR), "all" (AND)
|
||||
tags_match: How to match tags - "any", "all", "any_strict", "all_strict", or "exact"
|
||||
exclude_ids: Optional list of mental model IDs to exclude (e.g., when refreshing a mental model)
|
||||
last_memory_write_at: The bank's newest memory write, resolved once per reflect. Skips the
|
||||
per-model staleness query for any model refreshed at or after it.
|
||||
@@ -150,7 +150,7 @@ async def tool_search_mental_models(
|
||||
f"""
|
||||
SELECT
|
||||
id, name, content,
|
||||
tags, created_at, last_refreshed_at, trigger,
|
||||
tags, created_at, last_refreshed_at, last_memory_seen_at, trigger,
|
||||
1 - (embedding <=> $2::vector) as relevance
|
||||
FROM {fq_table("mental_models")}
|
||||
WHERE bank_id = $1 AND embedding IS NOT NULL {filters}
|
||||
@@ -167,6 +167,12 @@ async def tool_search_mental_models(
|
||||
if last_refreshed_at and last_refreshed_at.tzinfo is None:
|
||||
last_refreshed_at = last_refreshed_at.replace(tzinfo=timezone.utc)
|
||||
|
||||
# How far through the bank's memories this model is written — the cheap
|
||||
# bank-wide check below compares against that, not against when it last ran.
|
||||
last_memory_seen_at = row["last_memory_seen_at"] or last_refreshed_at
|
||||
if last_memory_seen_at and last_memory_seen_at.tzinfo is None:
|
||||
last_memory_seen_at = last_memory_seen_at.replace(tzinfo=timezone.utc)
|
||||
|
||||
# Per-MM staleness: new in-scope memories since last refresh (includes pending).
|
||||
# The scoped query has no index to use and scans the bank's memories in full, so
|
||||
# skip it for a model the bank-wide watermark already proves current: nothing was
|
||||
@@ -174,7 +180,7 @@ async def tool_search_mental_models(
|
||||
# model still gets the exact answer — the agent trusts a model without a verifying
|
||||
# recall() only on `is_stale is False`, so guessing conservatively here would buy
|
||||
# LLM turns to save a query. No watermark (absent, or an empty bank) → ask.
|
||||
if last_memory_write_at is not None and not _may_need_refresh(last_refreshed_at, last_memory_write_at):
|
||||
if last_memory_write_at is not None and not _may_need_refresh(last_memory_seen_at, last_memory_write_at):
|
||||
is_stale = False
|
||||
else:
|
||||
is_stale = await memory_engine.compute_mental_model_is_stale(conn, bank_id, row)
|
||||
@@ -228,7 +234,7 @@ async def tool_search_observations(
|
||||
request_context: Request context for authentication
|
||||
max_tokens: Maximum tokens for results (default 5000)
|
||||
tags: Optional tags to filter observations
|
||||
tags_match: How to match tags - "any" (OR), "all" (AND)
|
||||
tags_match: How to match tags - "any", "all", "any_strict", "all_strict", or "exact"
|
||||
last_consolidated_at: When consolidation last ran (for staleness check)
|
||||
pending_consolidation: Number of memories waiting to be consolidated
|
||||
source_facts_max_tokens: Token budget for source facts (-1 = disabled, 0+ = enabled with limit)
|
||||
@@ -318,7 +324,7 @@ async def tool_recall(
|
||||
request_context: Request context for authentication
|
||||
max_tokens: Maximum tokens for results (default 2048)
|
||||
tags: Filter by tags (includes untagged memories)
|
||||
tags_match: How to match tags - "any" (OR), "all" (AND), or "exact"
|
||||
tags_match: How to match tags - "any", "all", "any_strict", "all_strict", or "exact"
|
||||
connection_budget: Max DB connections for this recall (default 1 for internal ops)
|
||||
max_chunk_tokens: Maximum tokens for raw source chunk text (default 1000)
|
||||
fact_types: Optional filter for fact types to retrieve. Defaults to ["experience", "world"].
|
||||
|
||||
@@ -10,6 +10,8 @@ from typing import Any
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, field_validator
|
||||
|
||||
from .metadata_utils import as_string_metadata
|
||||
|
||||
VALID_RECALL_FACT_TYPES = frozenset(["world", "experience", "observation"])
|
||||
|
||||
|
||||
@@ -260,6 +262,9 @@ class MemoryFact(BaseModel):
|
||||
Also coerces non-string dict values (e.g., integer IDs stored in JSONB)
|
||||
to strings, preventing ValidationError when consolidation encounters
|
||||
metadata like {"original_id": 348} instead of {"original_id": "348"}.
|
||||
Null-valued keys are dropped rather than stringified to "None" (issue
|
||||
#3209), so rows written before retain normalized its input stay
|
||||
readable without a data migration.
|
||||
"""
|
||||
if v is None:
|
||||
return None
|
||||
@@ -268,7 +273,7 @@ class MemoryFact(BaseModel):
|
||||
|
||||
v = json.loads(v)
|
||||
if isinstance(v, dict):
|
||||
return {str(k): str(val) for k, val in v.items()}
|
||||
return as_string_metadata(v)
|
||||
return v
|
||||
|
||||
chunk_id: str | None = Field(
|
||||
@@ -353,6 +358,14 @@ class RecallResult(BaseModel):
|
||||
source_facts: dict[str, MemoryFact] | None = Field(
|
||||
None, description="Source facts for observation-type results, keyed by fact ID"
|
||||
)
|
||||
source_facts_truncated: bool | None = Field(
|
||||
None,
|
||||
description=(
|
||||
"Whether the source_facts map was cut short by the token budget. When true, some IDs in "
|
||||
"results[].source_fact_ids have no entry in source_facts — the budget ran out, the "
|
||||
"references are not dangling. Only set when source facts were requested."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class ReflectResult(BaseModel):
|
||||
|
||||
@@ -45,36 +45,6 @@ def _vector_index_clause() -> str | None:
|
||||
return index_using_clause(ext)
|
||||
|
||||
|
||||
async def create_bank_vector_indexes(conn, bank_id: str, internal_id: str, ops=None) -> None:
|
||||
"""Create per-(bank, fact_type) partial vector indexes for a newly created bank.
|
||||
|
||||
Respects the HINDSIGHT_API_VECTOR_EXTENSION config to use the appropriate
|
||||
index type (HNSW for pgvector, DiskANN for pgvectorscale, vchordrq for vchord).
|
||||
|
||||
AlloyDB ScaNN uses global vector indexes with filtered vector search; it
|
||||
cannot safely create per-bank indexes at bank-creation time because new
|
||||
banks have no embedding rows.
|
||||
bank_id is escaped for SQL literal safety (apostrophes doubled).
|
||||
|
||||
On Oracle 23ai, this is a no-op — Oracle uses a single global vector index
|
||||
created during migrations. Partial indexes (WHERE clause) are not supported
|
||||
for Oracle vector indexes.
|
||||
"""
|
||||
index_clause = _vector_index_clause()
|
||||
if index_clause is None:
|
||||
logger.debug("Skipping per-bank vector indexes for configured backend")
|
||||
return
|
||||
|
||||
await ops.create_bank_vector_indexes(
|
||||
conn,
|
||||
fq_table("memory_units"),
|
||||
bank_id,
|
||||
internal_id,
|
||||
index_clause,
|
||||
_BANK_INDEX_FACT_TYPES,
|
||||
)
|
||||
|
||||
|
||||
async def drop_bank_vector_indexes(conn, internal_id: str, ops=None) -> None:
|
||||
"""Drop per-(bank, fact_type) partial vector indexes for a bank being deleted.
|
||||
|
||||
@@ -190,12 +160,12 @@ async def get_or_create_bank_profile(pool, bank_id: str) -> BankProfileResult:
|
||||
``get_or_create_bank_profile_on_conn`` instead.
|
||||
"""
|
||||
|
||||
# A fresh bank builds its per-(bank, fact_type) partial vector indexes with
|
||||
# a plain CREATE INDEX (it must — this runs inside the bank-create tx, and
|
||||
# CONCURRENTLY cannot). That CREATE takes a ShareLock on the shared
|
||||
# memory_units table, which can deadlock with concurrent writers. The build
|
||||
# is idempotent (INSERT ... ON CONFLICT + CREATE INDEX IF NOT EXISTS), so a
|
||||
# transient deadlock (40P01 / ORA-00060) is safe to retry as a whole tx.
|
||||
# Retried as a whole transaction. This used to guard the per-bank CREATE
|
||||
# INDEX that ran inline here and took a ShareLock on the shared memory_units
|
||||
# table; that DDL is gone (#3485), but the lazy create can still lose a
|
||||
# deadlock (40P01 / ORA-00060) to a concurrent writer touching the same
|
||||
# bank row, and the body is idempotent (INSERT ... ON CONFLICT DO NOTHING),
|
||||
# so retrying stays correct and cheap.
|
||||
async def _create() -> BankProfileResult:
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
async with conn.transaction():
|
||||
@@ -242,10 +212,16 @@ async def get_or_create_bank_profile_on_conn(conn, bank_id: str, *, ops) -> Bank
|
||||
created=False,
|
||||
)
|
||||
|
||||
# Bank doesn't exist, create with defaults.
|
||||
# Generate internal_id here so we control the value and can use it
|
||||
# immediately for vector index creation without a RETURNING round-trip.
|
||||
internal_id = uuid.uuid4()
|
||||
# Bank doesn't exist, create with defaults. internal_id is minted here rather
|
||||
# than defaulted server-side so its value is known without a RETURNING
|
||||
# round-trip; the vector-index sweep derives index names from it.
|
||||
#
|
||||
# No vector-index DDL here. A fresh bank holds no rows, so it cannot meet
|
||||
# the size threshold that earns a per-(bank, fact_type) partial index; the
|
||||
# maintenance sweep builds one if and when the bank grows into it. Keeping
|
||||
# DDL out of this path also takes CREATE INDEX's ShareLock on the shared
|
||||
# memory_units table off the retain hot path, where it deadlocked against
|
||||
# concurrent writers. See issue #3485.
|
||||
inserted = await conn.fetchval(
|
||||
f"""
|
||||
INSERT INTO {fq_table("banks")} (bank_id, name, disposition, mission, internal_id)
|
||||
@@ -257,14 +233,10 @@ async def get_or_create_bank_profile_on_conn(conn, bank_id: str, *, ops) -> Bank
|
||||
bank_id, # Default name is the bank_id
|
||||
json.dumps(DEFAULT_DISPOSITION),
|
||||
"",
|
||||
internal_id,
|
||||
uuid.uuid4(),
|
||||
)
|
||||
|
||||
created = inserted is not None
|
||||
if created:
|
||||
# Fresh insert — create per-bank vector indexes (instant on empty bank)
|
||||
await create_bank_vector_indexes(conn, bank_id, str(internal_id), ops=ops)
|
||||
|
||||
return BankProfileResult(
|
||||
profile=BankProfile(name=bank_id, disposition=DispositionTraits(**DEFAULT_DISPOSITION), mission=""),
|
||||
created=created,
|
||||
@@ -424,9 +396,9 @@ def _as_utc(ts: datetime | None) -> datetime | None:
|
||||
return ts if ts.tzinfo is not None else ts.replace(tzinfo=UTC)
|
||||
|
||||
|
||||
async def list_banks(pool) -> list:
|
||||
async def list_banks(pool, *, search_query: str | None = None) -> list:
|
||||
"""
|
||||
List all banks in the system with summary stats.
|
||||
List banks with summary stats, optionally narrowed by a search string.
|
||||
|
||||
``last_document_at`` is document *ingestion* time (when a document first
|
||||
landed), while ``last_write_at`` is the last time anything was written to
|
||||
@@ -434,8 +406,14 @@ async def list_banks(pool) -> list:
|
||||
to a long-lived document does not move ``last_document_at``, which is why
|
||||
the two differ and why UIs showing "last write" must use ``last_write_at``.
|
||||
|
||||
``fact_count`` comes from the ``memory_units`` join, which is empty for a bank
|
||||
whose memories live outside SQL. Those banks need :func:`apply_store_fact_counts`
|
||||
to get a real count; callers run it on the page they actually return so the live
|
||||
per-bank count query doesn't fire for every bank in the system.
|
||||
|
||||
Args:
|
||||
pool: Database connection pool
|
||||
search_query: Case-insensitive substring matched against bank ID and name
|
||||
|
||||
Returns:
|
||||
List of dicts with bank info and stats (fact_count, last_document_at, last_write_at),
|
||||
@@ -445,6 +423,15 @@ async def list_banks(pool) -> list:
|
||||
docs_table = fq_table("documents")
|
||||
mu_table = fq_table("memory_units")
|
||||
|
||||
# Spelled out as UPPER(...) LIKE UPPER(...) rather than ILIKE: the Oracle
|
||||
# rewriter only recognizes ILIKE on an unqualified column, and these are
|
||||
# alias-qualified.
|
||||
where_clause = ""
|
||||
params: list[str] = []
|
||||
if search_query:
|
||||
where_clause = "WHERE (UPPER(b.bank_id) LIKE UPPER($1) OR UPPER(COALESCE(b.name, '')) LIKE UPPER($2))"
|
||||
params = [f"%{search_query}%", f"%{search_query}%"]
|
||||
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
@@ -470,19 +457,16 @@ async def list_banks(pool) -> list:
|
||||
FROM {mu_table}
|
||||
GROUP BY bank_id
|
||||
) m ON m.bank_id = b.bank_id
|
||||
{where_clause}
|
||||
ORDER BY b.bank_id
|
||||
"""
|
||||
""",
|
||||
*params,
|
||||
)
|
||||
|
||||
result = []
|
||||
# Banks are ordered by last write in Python rather than SQL: GREATEST() has
|
||||
# different NULL semantics on PostgreSQL vs Oracle, and the bank list is small.
|
||||
sort_keys: dict[str, datetime] = {}
|
||||
# A store that keeps memories outside SQL leaves the memory_units join empty, so its
|
||||
# per-bank fact_count comes from the store instead (one live count per bank).
|
||||
from ..memories import get_memories
|
||||
|
||||
_store = get_memories()
|
||||
|
||||
for row in rows:
|
||||
disposition_data = row["disposition"]
|
||||
@@ -498,12 +482,6 @@ async def list_banks(pool) -> list:
|
||||
write_times = [t for t in (_as_utc(row["last_document_write_at"]), _as_utc(row["last_fact_at"])) if t]
|
||||
last_write = max(write_times) if write_times else None
|
||||
|
||||
fact_count = row["fact_count"]
|
||||
if not _store.writes_memory_rows_in_sql_for(row["bank_id"]):
|
||||
fact_count = sum(
|
||||
(await _store.count_memories(conn=conn, fq_table=fq_table, bank_id=row["bank_id"])).values()
|
||||
)
|
||||
|
||||
sort_keys[row["bank_id"]] = last_write or created_at or _UNIX_EPOCH
|
||||
result.append(
|
||||
{
|
||||
@@ -513,7 +491,7 @@ async def list_banks(pool) -> list:
|
||||
"mission": row["mission"] or "",
|
||||
"created_at": created_at.isoformat() if created_at else None,
|
||||
"updated_at": updated_at.isoformat() if updated_at else None,
|
||||
"fact_count": fact_count,
|
||||
"fact_count": row["fact_count"],
|
||||
"last_document_at": last_doc.isoformat() if last_doc else None,
|
||||
"last_write_at": last_write.isoformat() if last_write else None,
|
||||
}
|
||||
@@ -521,3 +499,23 @@ async def list_banks(pool) -> list:
|
||||
|
||||
result.sort(key=lambda bank: sort_keys[bank["bank_id"]], reverse=True)
|
||||
return result
|
||||
|
||||
|
||||
async def apply_store_fact_counts(pool, banks: list[dict]) -> None:
|
||||
"""Replace ``fact_count`` in-place for banks that keep their memories outside SQL.
|
||||
|
||||
Those banks leave the ``memory_units`` join empty, so the count has to come
|
||||
from the store — one live count per bank, which is why this runs on a single
|
||||
page of :func:`list_banks` rather than on every bank in the system.
|
||||
"""
|
||||
from ..memories import get_memories
|
||||
|
||||
store = get_memories()
|
||||
external = [bank for bank in banks if not store.writes_memory_rows_in_sql_for(bank["bank_id"])]
|
||||
if not external:
|
||||
return
|
||||
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
for bank in external:
|
||||
counts = await store.count_memories(conn=conn, fq_table=fq_table, bank_id=bank["bank_id"])
|
||||
bank["fact_count"] = sum(counts.values())
|
||||
|
||||
@@ -136,9 +136,26 @@ async def delete_chunks_by_ids(conn, chunk_ids: list[str], bank_id: str | None =
|
||||
if bank_id:
|
||||
outgoing_unit_ids = await memory_ids_for_chunks(conn, bank_id, chunk_ids)
|
||||
if outgoing_unit_ids:
|
||||
from ..graph_maintenance import enqueue_entity_prune_candidates
|
||||
from .fact_storage import delete_stale_observations_for_memories
|
||||
|
||||
invalidated = await delete_stale_observations_for_memories(conn, bank_id, outgoing_unit_ids, ops=ops)
|
||||
# Queue the entities these facts reference BEFORE the cascade takes
|
||||
# their unit_entities rows: afterwards an entity whose last posting
|
||||
# was here is unreachable garbage. Delta retain deletes facts only
|
||||
# through this cascade, so this is the one place that can catch them
|
||||
# (the full-replace path enqueues in ``handle_document_tracking``).
|
||||
await enqueue_entity_prune_candidates(conn, bank_id, outgoing_unit_ids)
|
||||
|
||||
# Capture surviving units whose temporal/semantic links point at
|
||||
# the outgoing facts before the link/chunk cascade below removes
|
||||
# the evidence needed to find them. Full document replacement does
|
||||
# the same in ``handle_document_tracking``; without it, a delta
|
||||
# edit leaves survivors permanently below their configured link
|
||||
# caps even though retain submits graph maintenance afterwards.
|
||||
from ..graph_maintenance import enqueue_relink_victims
|
||||
|
||||
await enqueue_relink_victims(conn, bank_id, outgoing_unit_ids)
|
||||
|
||||
# The chunks->memory_units FK cascade below does not reach a store that keeps memories
|
||||
# outside SQL (its memory_units is empty), so drop the memories carrying each deleted
|
||||
@@ -155,6 +172,19 @@ async def delete_chunks_by_ids(conn, chunk_ids: list[str], bank_id: str | None =
|
||||
# memory_links in opposite orders and deadlock. Delete links explicitly in a
|
||||
# total order before deleting chunks so every writer takes row locks the same
|
||||
# way; the FK cascade still handles anything inserted later in this txn.
|
||||
#
|
||||
# ``matched_links`` collects the endpoints as a UNION of two single-column joins
|
||||
# rather than the one ``tu.id = ml.from_unit_id OR tu.id = ml.to_unit_id`` predicate
|
||||
# it replaces. An OR spanning two columns of ``ml`` is not indexable: the planner
|
||||
# cannot drive it from either endpoint index, so it made memory_links the outer
|
||||
# relation of a nested-loop semi join and sequentially scanned the whole table once
|
||||
# per delete — O(rows_in_memory_links x target_units). Past a few million links that
|
||||
# exceeded the asyncpg command timeout and delta retain failed with a bare
|
||||
# TimeoutError (issue #3387). Split in two, each half is an index scan on
|
||||
# idx_memory_links_from_type_weight / idx_memory_links_to_type_weight.
|
||||
# The UNION yields the identical row set; the deterministic ORDER BY and
|
||||
# FOR UPDATE that #2570 added stay in ``ordered_links``, which locks the rows in
|
||||
# that order after the endpoints have been found.
|
||||
await conn.execute(
|
||||
f"""
|
||||
WITH target_units AS MATERIALIZED (
|
||||
@@ -162,14 +192,19 @@ async def delete_chunks_by_ids(conn, chunk_ids: list[str], bank_id: str | None =
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE chunk_id = ANY($1::text[])
|
||||
),
|
||||
matched_links AS MATERIALIZED (
|
||||
SELECT ml.ctid AS link_ctid
|
||||
FROM {fq_table("memory_links")} ml
|
||||
JOIN target_units tu ON tu.id = ml.from_unit_id
|
||||
UNION
|
||||
SELECT ml.ctid AS link_ctid
|
||||
FROM {fq_table("memory_links")} ml
|
||||
JOIN target_units tu ON tu.id = ml.to_unit_id
|
||||
),
|
||||
ordered_links AS MATERIALIZED (
|
||||
SELECT ml.ctid
|
||||
FROM {fq_table("memory_links")} ml
|
||||
WHERE EXISTS (
|
||||
SELECT 1
|
||||
FROM target_units tu
|
||||
WHERE tu.id = ml.from_unit_id OR tu.id = ml.to_unit_id
|
||||
)
|
||||
JOIN matched_links ON ml.ctid = matched_links.link_ctid
|
||||
ORDER BY
|
||||
LEAST(ml.from_unit_id, ml.to_unit_id),
|
||||
GREATEST(ml.from_unit_id, ml.to_unit_id),
|
||||
|
||||
@@ -7,18 +7,23 @@ Handles entity extraction and resolution for stored facts.
|
||||
import logging
|
||||
|
||||
from . import link_utils
|
||||
from .types import EntityResolutionResult, ProcessedFact
|
||||
from .types import EntityResolutionResult, ProcessedFact, UserEntities
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _prepare_facts_for_entity_processing(
|
||||
facts: list[ProcessedFact],
|
||||
user_entities_per_content: dict[int, list[dict]] | None = None,
|
||||
user_entities_per_content: dict[int, UserEntities] | None = None,
|
||||
) -> tuple[list[str], list, list[list[dict]]]:
|
||||
"""
|
||||
Extract fact texts, dates, and merged entity lists from ProcessedFact objects.
|
||||
|
||||
Extracted names always carry ``resolve=True`` — they are the extractor's guess at a name, so
|
||||
matching them onto the bank's existing entities is the point. Caller-supplied names carry the
|
||||
content item's ``resolve_entities`` flag, so a caller can have their own names taken literally
|
||||
without turning off resolution for the extractor's (#3479).
|
||||
|
||||
Returns:
|
||||
Tuple of (fact_texts, fact_dates, entities_per_fact)
|
||||
"""
|
||||
@@ -29,20 +34,29 @@ def _prepare_facts_for_entity_processing(
|
||||
|
||||
entities_per_fact = []
|
||||
for fact in facts:
|
||||
llm_entities = [{"text": entity.name, "type": "CONCEPT"} for entity in (fact.entities or [])]
|
||||
llm_entities = [{"text": entity.name, "type": "CONCEPT", "resolve": True} for entity in (fact.entities or [])]
|
||||
|
||||
user_entities = user_entities_per_content.get(fact.content_index, [])
|
||||
supplied = user_entities_per_content.get(fact.content_index)
|
||||
user_entities = supplied.entities if supplied else []
|
||||
user_resolve = supplied.resolve if supplied else True
|
||||
|
||||
seen_texts = {e["text"].lower() for e in llm_entities}
|
||||
by_text = {e["text"].lower(): e for e in llm_entities}
|
||||
for user_entity in user_entities:
|
||||
if user_entity["text"].lower() not in seen_texts:
|
||||
llm_entities.append(
|
||||
{
|
||||
"text": user_entity["text"],
|
||||
"type": user_entity.get("type", "CONCEPT"),
|
||||
}
|
||||
)
|
||||
seen_texts.add(user_entity["text"].lower())
|
||||
text_lower = user_entity["text"].lower()
|
||||
existing = by_text.get(text_lower)
|
||||
if existing is None:
|
||||
entity = {
|
||||
"text": user_entity["text"],
|
||||
"type": user_entity.get("type", "CONCEPT"),
|
||||
"resolve": user_resolve,
|
||||
}
|
||||
llm_entities.append(entity)
|
||||
by_text[text_lower] = entity
|
||||
else:
|
||||
# The extractor produced this name too. The caller still authored it, so their
|
||||
# intent wins: a literal name must not become resolvable just because extraction
|
||||
# happened to agree on the spelling.
|
||||
existing["resolve"] = existing["resolve"] and user_resolve
|
||||
|
||||
entities_per_fact.append(llm_entities)
|
||||
|
||||
@@ -56,7 +70,7 @@ async def resolve_entities(
|
||||
unit_ids: list[str],
|
||||
facts: list[ProcessedFact],
|
||||
log_buffer: list[str] = None,
|
||||
user_entities_per_content: dict[int, list[dict]] = None,
|
||||
user_entities_per_content: dict[int, UserEntities] | None = None,
|
||||
entity_labels: list | None = None,
|
||||
) -> EntityResolutionResult:
|
||||
"""
|
||||
@@ -72,7 +86,8 @@ async def resolve_entities(
|
||||
unit_ids: Placeholder unit IDs (used only for grouping)
|
||||
facts: List of ProcessedFact objects
|
||||
log_buffer: Optional buffer for detailed logging
|
||||
user_entities_per_content: Dict mapping content_index to user-provided entities
|
||||
user_entities_per_content: Dict mapping content_index to the caller-supplied
|
||||
entities for that content item and whether to resolve them
|
||||
entity_labels: Optional entity label taxonomy
|
||||
|
||||
Returns:
|
||||
|
||||
@@ -12,7 +12,8 @@ from typing import Any
|
||||
|
||||
from ...config import _get_raw_config
|
||||
from ..memory_engine import fq_table
|
||||
from .bank_utils import DEFAULT_DISPOSITION, create_bank_vector_indexes
|
||||
from ..metadata_utils import drop_null_values
|
||||
from .bank_utils import DEFAULT_DISPOSITION
|
||||
from .fact_extraction import _sanitize_text
|
||||
from .types import ProcessedFact
|
||||
|
||||
@@ -131,25 +132,26 @@ async def ensure_bank_exists(conn, bank_id: str, ops=None) -> None:
|
||||
conn: Database connection
|
||||
bank_id: Bank identifier
|
||||
"""
|
||||
# Generate internal_id here so we control the value and can use it
|
||||
# immediately for HNSW index creation without a RETURNING round-trip.
|
||||
internal_id = uuid.uuid4()
|
||||
inserted = await conn.fetchval(
|
||||
# internal_id is generated here rather than defaulted server-side so the
|
||||
# value is known without a RETURNING round-trip; the vector-index sweep
|
||||
# derives index names from it.
|
||||
#
|
||||
# No vector-index DDL on this path. A fresh bank holds no rows, so it cannot
|
||||
# meet the size threshold that earns a per-(bank, fact_type) partial index;
|
||||
# the maintenance sweep builds one if the bank later grows into it. See
|
||||
# issue #3485.
|
||||
await conn.execute(
|
||||
f"""
|
||||
INSERT INTO {fq_table("banks")} (bank_id, name, disposition, mission, internal_id)
|
||||
VALUES ($1, $2, $3::jsonb, $4, $5)
|
||||
ON CONFLICT (bank_id) DO NOTHING
|
||||
RETURNING bank_id
|
||||
""",
|
||||
bank_id,
|
||||
bank_id, # Default name is the bank_id (matches get_or_create_bank_profile)
|
||||
json.dumps(DEFAULT_DISPOSITION),
|
||||
"",
|
||||
internal_id,
|
||||
uuid.uuid4(),
|
||||
)
|
||||
if inserted:
|
||||
# Fresh insert — create per-bank vector indexes
|
||||
await create_bank_vector_indexes(conn, bank_id, str(internal_id), ops=ops)
|
||||
|
||||
|
||||
async def delete_stale_observations_for_memories(
|
||||
@@ -284,9 +286,16 @@ async def handle_document_tracking(
|
||||
# those links. ``ops`` may be None for older callers that haven't
|
||||
# been wired up — skip enqueue in that case rather than crash.
|
||||
if ops is not None:
|
||||
from ..graph_maintenance import enqueue_relink_victims
|
||||
from ..graph_maintenance import enqueue_entity_prune_candidates, enqueue_relink_victims
|
||||
|
||||
await enqueue_relink_victims(conn, bank_id, [str(uid) for uid in existing_unit_ids])
|
||||
doomed_ids = [str(uid) for uid in existing_unit_ids]
|
||||
await enqueue_relink_victims(conn, bank_id, doomed_ids)
|
||||
# Same timing, different target: the entities these units are
|
||||
# about to stop referencing may have no other posting. The
|
||||
# re-ingest re-resolves entities from scratch, so the ones the
|
||||
# new facts don't name again are orphans the moment this
|
||||
# cascade lands.
|
||||
await enqueue_entity_prune_candidates(conn, bank_id, doomed_ids)
|
||||
|
||||
# Explicitly delete memory_units by document_id BEFORE deleting the
|
||||
# document row. The CASCADE from documents→chunks→memory_units only
|
||||
@@ -416,6 +425,12 @@ async def update_memory_units_metadata_and_tags(
|
||||
current document tags and metadata so its optimized result matches a full
|
||||
replace.
|
||||
|
||||
``metadata`` arrives as the raw retain_params bag (the document row keeps
|
||||
the caller's input verbatim), so null-valued keys are dropped here — the
|
||||
same normalization ``RetainContent`` applies to freshly extracted facts
|
||||
(issue #3209). Without it a re-retain would leave surviving units carrying
|
||||
nulls while the units around them do not.
|
||||
|
||||
Returns:
|
||||
Number of memory units updated.
|
||||
"""
|
||||
@@ -442,7 +457,7 @@ async def update_memory_units_metadata_and_tags(
|
||||
bank_id,
|
||||
document_id,
|
||||
tags or [],
|
||||
json.dumps(metadata or {}),
|
||||
json.dumps(drop_null_values(metadata)),
|
||||
)
|
||||
# result is a status string like "UPDATE 5"
|
||||
try:
|
||||
|
||||
@@ -0,0 +1,228 @@
|
||||
"""Coalescing several queued retains for one document into a single execution.
|
||||
|
||||
Appends to one document are serialized by the claim predicate
|
||||
(``document_serialization_sql``), which is what makes them correct — but on its
|
||||
own it also makes them slow: a client that buffers 50 turns offline and flushes
|
||||
them would run 50 sequential retains, each re-reading and reprocessing the
|
||||
document. Folding collapses that burst into one execution over the concatenated
|
||||
turns.
|
||||
|
||||
The folding happens at **claim** time, over rows the claim transaction already
|
||||
holds ``FOR UPDATE``, and it never rewrites an operation's ``task_payload``.
|
||||
That is the whole reason this is tractable:
|
||||
|
||||
* Rows are immutable after insert, so there is no submit-vs-claim race to
|
||||
reason about — the alternative (merging new content into a pending row at
|
||||
submit time) races the worker reading that row and reintroduces exactly the
|
||||
lost-update bug this is meant to fix.
|
||||
* Operation identity survives. Each submission keeps its own ``operation_id``,
|
||||
status, and post-retain hook, so nothing downstream of the queue — polling,
|
||||
webhooks, cancellation, metering — has to learn about folding.
|
||||
|
||||
Because of that, folding is a pure optimization: delete it and the system is
|
||||
still correct, only slower. A bug here can cost latency or LLM spend; it cannot
|
||||
lose a turn, because ``ConcurrentAppendConflict`` and the claim predicate carry
|
||||
correctness independently.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Peers folded into one execution on a first attempt. Bounded so a document
|
||||
# that has accumulated a long backlog still commits in reasonable steps rather
|
||||
# than one enormous transaction; the remainder is claimed on the next cycle.
|
||||
DEFAULT_MAX_FOLD_PEERS = 16
|
||||
|
||||
|
||||
@dataclass
|
||||
class FoldMember:
|
||||
"""One submitted operation participating in a folded execution."""
|
||||
|
||||
operation_id: str
|
||||
contents: list[dict[str, Any]]
|
||||
tenant_id: str | None = None
|
||||
api_key_id: str | None = None
|
||||
document_tags: list[str] | None = None
|
||||
strategy: str | None = None
|
||||
has_file_metadata: bool = False
|
||||
"""Whether the submission carries ``_file_metadata`` (a converted upload)."""
|
||||
|
||||
@property
|
||||
def is_append_only(self) -> bool:
|
||||
"""Whether every item this operation submitted appends to its document.
|
||||
|
||||
Only appends are foldable, because only appends are cumulative — see
|
||||
:func:`plan_retain_fold`.
|
||||
"""
|
||||
return bool(self.contents) and all(item.get("update_mode") == "append" for item in self.contents)
|
||||
|
||||
|
||||
def _can_join(primary: FoldMember, peer: FoldMember) -> str | None:
|
||||
"""Why ``peer`` cannot join ``primary``'s execution, or None if it can.
|
||||
|
||||
Everything an execution applies once, from the primary's payload, has to
|
||||
match — otherwise folding would silently apply the primary's value to the
|
||||
peer's content.
|
||||
"""
|
||||
if not peer.is_append_only:
|
||||
return "not an append"
|
||||
if peer.has_file_metadata:
|
||||
return "carries file metadata"
|
||||
if peer.tenant_id != primary.tenant_id or peer.api_key_id != primary.api_key_id:
|
||||
# Usage is attributed to the members of a fold, so mixing credentials
|
||||
# inside one execution would bill one caller for another's extraction.
|
||||
return "different tenant/api key"
|
||||
if sorted(peer.document_tags or []) != sorted(primary.document_tags or []):
|
||||
# The execution applies the primary's document_tags to the whole
|
||||
# document; folding a differently-tagged peer would drop its tags.
|
||||
return "different document_tags"
|
||||
if peer.strategy != primary.strategy:
|
||||
return "different strategy"
|
||||
return None
|
||||
|
||||
|
||||
@dataclass
|
||||
class FoldMemberRef:
|
||||
"""A fold member as it survives the trip through the task payload.
|
||||
|
||||
The fold is decided at claim time and consumed by the engine, with the task
|
||||
payload (JSON) in between — so this is the boundary type: the payload
|
||||
carries plain dicts and they are parsed back into this the moment the engine
|
||||
reads them, rather than being passed around as loose dicts.
|
||||
|
||||
Only what the engine needs downstream: which operation, and how many content
|
||||
items it contributed, which is what slices the execution's results back
|
||||
apart for the per-operation post-retain hooks.
|
||||
"""
|
||||
|
||||
operation_id: str
|
||||
items_count: int
|
||||
|
||||
def to_payload(self) -> dict[str, Any]:
|
||||
return {"operation_id": self.operation_id, "items_count": self.items_count}
|
||||
|
||||
@classmethod
|
||||
def from_payload(cls, raw: dict[str, Any]) -> "FoldMemberRef":
|
||||
return cls(operation_id=str(raw["operation_id"]), items_count=int(raw["items_count"]))
|
||||
|
||||
@classmethod
|
||||
def list_from_payload(cls, raw: list[dict[str, Any]] | None) -> list["FoldMemberRef"] | None:
|
||||
return None if raw is None else [cls.from_payload(item) for item in raw]
|
||||
|
||||
|
||||
@dataclass
|
||||
class FoldPlan:
|
||||
"""Which queued peers join this execution, and which stay pending."""
|
||||
|
||||
members: list[FoldMember] = field(default_factory=list)
|
||||
"""The primary operation first, then the peers folded into it, in submission order."""
|
||||
|
||||
deferred: list[str] = field(default_factory=list)
|
||||
"""Operation ids left pending for a later claim, with the reason logged."""
|
||||
|
||||
@property
|
||||
def peer_ids(self) -> list[str]:
|
||||
"""Ids of the folded peers — everything but the primary."""
|
||||
return [m.operation_id for m in self.members[1:]]
|
||||
|
||||
|
||||
def max_fold_peers_for_retry(retry_count: int, base: int = DEFAULT_MAX_FOLD_PEERS) -> int:
|
||||
"""Fold width for an operation on its ``retry_count``-th attempt.
|
||||
|
||||
A folded execution is all-or-nothing: one poisonous turn (content that makes
|
||||
extraction fail every time) would otherwise take every turn folded with it
|
||||
down on every retry, and those turns would be re-folded with it again on the
|
||||
next attempt — a burst of good content held hostage by one bad item.
|
||||
|
||||
Halving the width per retry makes the fold converge on the poison: by the
|
||||
time the operation has failed a few times it runs alone, fails alone, and
|
||||
is dead-lettered alone while its neighbours proceed. Expressed as a rule
|
||||
rather than a special case, so there is no "is this the bad one" heuristic
|
||||
to get wrong.
|
||||
"""
|
||||
if retry_count <= 0:
|
||||
return base
|
||||
return max(0, base >> retry_count)
|
||||
|
||||
|
||||
def plan_retain_fold(
|
||||
primary: FoldMember,
|
||||
peers: list[FoldMember],
|
||||
*,
|
||||
max_peers: int,
|
||||
token_budget: int,
|
||||
count_tokens,
|
||||
) -> FoldPlan:
|
||||
"""Choose the contiguous run of ``peers`` that folds into ``primary``.
|
||||
|
||||
``peers`` must already be ordered by ``(created_at, operation_id)`` — the
|
||||
order the claim query imposes and the order the document accumulates in.
|
||||
|
||||
Folding stops at the **first** peer that cannot join, and takes nothing
|
||||
after it. Appends are cumulative, so skipping one and taking the next would
|
||||
commit turns out of order; a contiguous prefix is the only shape that keeps
|
||||
the document's text in submission order. Peers left behind stay pending and
|
||||
are claimed on a later cycle, still in order.
|
||||
|
||||
**Only appends fold.** Appends are cumulative: running two of them as one
|
||||
execution over the concatenated turns produces the same document as running
|
||||
them in sequence. Replace is not — it means "this body supersedes what is
|
||||
stored", so folding two replaces would store ``body1 + body2`` where the
|
||||
correct answer is ``body2``, and folding an append behind a replace would
|
||||
turn a document-wiping submission into a concatenation. An operation that
|
||||
is not append-only therefore runs alone, as primary and as peer.
|
||||
|
||||
A peer cannot join when:
|
||||
|
||||
* it would push the execution past ``token_budget``. Beyond that the engine
|
||||
splits the work into sub-batches, and several sub-batches carrying
|
||||
different bodies for one document defeat the streaming ownership check
|
||||
(#3282) — so the fold stays inside a single orchestrator pass by
|
||||
construction. The primary alone is always allowed through, however large.
|
||||
* ``_can_join`` rejects it: not an append, carrying file metadata, or
|
||||
differing in anything the execution applies once from the primary's
|
||||
payload (tenant, API key, document_tags, strategy).
|
||||
* ``max_peers`` is already reached.
|
||||
"""
|
||||
plan = FoldPlan(members=[primary])
|
||||
if max_peers <= 0 or not primary.is_append_only or primary.has_file_metadata:
|
||||
# A non-append primary is not a fold base at all: whatever queued behind
|
||||
# it must wait for it to finish and then be re-evaluated against the
|
||||
# document it leaves behind.
|
||||
plan.deferred = [p.operation_id for p in peers]
|
||||
return plan
|
||||
|
||||
used = sum(count_tokens(item.get("content", "")) for item in primary.contents)
|
||||
|
||||
for index, peer in enumerate(peers):
|
||||
if len(plan.members) > max_peers:
|
||||
plan.deferred = [p.operation_id for p in peers[index:]]
|
||||
break
|
||||
rejection = _can_join(primary, peer)
|
||||
if rejection is not None:
|
||||
logger.debug("Not folding %s: %s", peer.operation_id, rejection)
|
||||
plan.deferred = [p.operation_id for p in peers[index:]]
|
||||
break
|
||||
peer_tokens = sum(count_tokens(item.get("content", "")) for item in peer.contents)
|
||||
if used + peer_tokens > token_budget:
|
||||
plan.deferred = [p.operation_id for p in peers[index:]]
|
||||
break
|
||||
used += peer_tokens
|
||||
plan.members.append(peer)
|
||||
|
||||
return plan
|
||||
|
||||
|
||||
def merge_fold_contents(members: list[FoldMember]) -> list[dict[str, Any]]:
|
||||
"""Flatten a fold's submissions into the content list for one execution.
|
||||
|
||||
Order is the members' order, which is submission order — the document ends
|
||||
up carrying the turns exactly as the callers sent them.
|
||||
"""
|
||||
merged: list[dict[str, Any]] = []
|
||||
for member in members:
|
||||
merged.extend(member.contents)
|
||||
return merged
|
||||
@@ -43,6 +43,16 @@ def _normalize_entity_name(name: str) -> str:
|
||||
return _WHITESPACE_RUN_RE.sub(" ", name).strip()
|
||||
|
||||
|
||||
def _entity_resolve_flag(ent) -> bool:
|
||||
"""Whether this candidate name should be resolved against existing entities.
|
||||
|
||||
Defaults to True (extraction's behaviour). Only dict candidates can opt out, which is how
|
||||
retain marks the entities its *caller* supplied: those are authoritative names, not guesses
|
||||
at which entity is meant (#3479).
|
||||
"""
|
||||
return bool(ent.get("resolve", True)) if isinstance(ent, dict) else True
|
||||
|
||||
|
||||
# Maximum number of temporal links to keep per unit (from_unit_id).
|
||||
# Retrieval only reads top 10-20 per unit via LATERAL join, so keeping
|
||||
# more is wasted storage and write amplification.
|
||||
@@ -79,6 +89,26 @@ def _cap_links_per_unit(links: list[tuple], max_per_unit: int = MAX_TEMPORAL_LIN
|
||||
return result
|
||||
|
||||
|
||||
def _lock_order_key(lnk: tuple) -> tuple[str, str, str, str]:
|
||||
"""Canonical lock-order key for a link row, shared by every writer.
|
||||
|
||||
Mirrors the total order that ``chunk_storage.delete_chunks_by_ids`` uses when
|
||||
it locks ``memory_links`` before a cascade delete:
|
||||
|
||||
(LEAST(from, to), GREATEST(from, to), link_type, COALESCE(entity_id, nil))
|
||||
|
||||
Direction is normalised so ``(A, B)`` and ``(B, A)`` sort adjacent, and the
|
||||
key covers the full unique index — including ``link_type`` and ``entity_id``
|
||||
— so two edges sharing a ``(from, to)`` pair can't be locked in opposite
|
||||
orders by concurrent inserts. UUID string ordering matches the DB's ``uuid``
|
||||
ordering because the ids are canonical lowercase-hex form.
|
||||
"""
|
||||
a, b = str(lnk[0]), str(lnk[1])
|
||||
low, high = (a, b) if a <= b else (b, a)
|
||||
entity = str(lnk[4]) if lnk[4] is not None else _NIL_ENTITY_UUID
|
||||
return (low, high, str(lnk[2]), entity)
|
||||
|
||||
|
||||
async def _bulk_insert_links(
|
||||
conn,
|
||||
links: list[tuple],
|
||||
@@ -89,8 +119,9 @@ async def _bulk_insert_links(
|
||||
) -> None:
|
||||
"""Bulk-insert links using sorted INSERT FROM unnest().
|
||||
|
||||
Sorting by (from_unit_id, to_unit_id) ensures all concurrent transactions
|
||||
acquire index locks in the same order, eliminating circular-wait deadlocks.
|
||||
Sorting on the full, direction-normalised unique key ensures all concurrent
|
||||
writers — inserts and deletes alike — acquire index locks in the same order,
|
||||
eliminating circular-wait deadlocks. See :func:`_lock_order_key`.
|
||||
|
||||
Args:
|
||||
conn: Database connection (must be inside a transaction).
|
||||
@@ -106,9 +137,9 @@ async def _bulk_insert_links(
|
||||
if not links:
|
||||
return
|
||||
|
||||
# Sort by (from_unit_id, to_unit_id) to guarantee consistent lock ordering
|
||||
# across concurrent transactions — prevents deadlocks.
|
||||
sorted_links = sorted(links, key=lambda lnk: (str(lnk[0]), str(lnk[1])))
|
||||
# Sort on the canonical lock-order key so every concurrent writer takes the
|
||||
# index locks in the same order — prevents circular-wait deadlocks.
|
||||
sorted_links = sorted(links, key=_lock_order_key)
|
||||
|
||||
exists_clause = ""
|
||||
if not skip_exists_check:
|
||||
@@ -192,7 +223,7 @@ def _prepare_entities_for_resolution(
|
||||
# own entity list) identical, and the upstream dedup in
|
||||
# entity_processing runs on the raw text. Without this, the same entity
|
||||
# would be resolved twice for one fact and its mention_count bumped twice.
|
||||
seen_in_fact: set[str] = set()
|
||||
seen_in_fact: dict[str, dict] = {}
|
||||
for ent in entity_list:
|
||||
if hasattr(ent, "text"):
|
||||
raw_text, entity_type = ent.text, "CONCEPT"
|
||||
@@ -209,11 +240,19 @@ def _prepare_entities_for_resolution(
|
||||
dropped_empty += 1
|
||||
continue
|
||||
|
||||
if normalized_text.lower() in seen_in_fact:
|
||||
resolve = _entity_resolve_flag(ent)
|
||||
kept = seen_in_fact.get(normalized_text.lower())
|
||||
if kept is not None:
|
||||
# Same name after normalization. Keep the first spelling but carry the stricter
|
||||
# flag: entity_processing dedups on the RAW text, so a caller's literal
|
||||
# "Acme Corp" and the extractor's "Acme\nCorp" both reach here, and dropping the
|
||||
# caller's outright would let the name be resolved away after all (#3479).
|
||||
kept["resolve"] = kept["resolve"] and resolve
|
||||
continue
|
||||
seen_in_fact.add(normalized_text.lower())
|
||||
|
||||
formatted_entities.append({"text": normalized_text, "type": entity_type})
|
||||
entity = {"text": normalized_text, "type": entity_type, "resolve": resolve}
|
||||
seen_in_fact[normalized_text.lower()] = entity
|
||||
formatted_entities.append(entity)
|
||||
all_entities.append(formatted_entities)
|
||||
|
||||
if dropped_empty:
|
||||
@@ -242,6 +281,7 @@ def _prepare_entities_for_resolution(
|
||||
{
|
||||
"text": entity["text"],
|
||||
"type": entity["type"],
|
||||
"resolve": entity["resolve"],
|
||||
"nearby_entities": entities,
|
||||
}
|
||||
)
|
||||
|
||||
@@ -275,16 +275,23 @@ from . import (
|
||||
from .types import (
|
||||
CausalRelation,
|
||||
ChunkMetadata,
|
||||
ConcurrentAppendConflict,
|
||||
ExtractedFact,
|
||||
Phase1Result,
|
||||
ProcessedFact,
|
||||
ResolvedEntity,
|
||||
RetainContent,
|
||||
RetainContentDict,
|
||||
UserEntities,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Sentinel append base: the append read found no document row at all. Distinct
|
||||
# from None (not an append) and from any real content_hash, so the write gate
|
||||
# can tell "nobody had written this document yet" apart from "we didn't look".
|
||||
_APPEND_BASE_ABSENT = "__absent__"
|
||||
|
||||
RetainOutboxCallback = Callable[[asyncpg.Connection], Awaitable[None]]
|
||||
RetainOutboxCallbackFactory = Callable[[list[RetainContentDict]], RetainOutboxCallback | None]
|
||||
|
||||
@@ -397,7 +404,11 @@ async def _pre_resolve_phase1(
|
||||
set_stage("retain.phase1.resolve")
|
||||
from .link_utils import compute_semantic_links_ann
|
||||
|
||||
user_entities_per_content = {idx: content.entities for idx, content in enumerate(contents) if content.entities}
|
||||
user_entities_per_content = {
|
||||
idx: UserEntities(entities=content.entities, resolve=content.resolve_entities)
|
||||
for idx, content in enumerate(contents)
|
||||
if content.entities
|
||||
}
|
||||
|
||||
# Use placeholder unit_ids for grouping during resolution. The actual
|
||||
# unit_ids are created later by insert_facts_batch inside the transaction,
|
||||
@@ -579,6 +590,367 @@ async def _insert_facts_and_links(
|
||||
return result_unit_ids
|
||||
|
||||
|
||||
@dataclass
|
||||
class _ExtStreamingWriteResult:
|
||||
"""Outcome of :func:`_streaming_batch_write_ext`."""
|
||||
|
||||
aborted: bool
|
||||
batch_result_ids: list[list[str]]
|
||||
|
||||
|
||||
@dataclass
|
||||
class _ExtDeltaWriteResult:
|
||||
"""Outcome of :func:`_delta_batch_write_ext`."""
|
||||
|
||||
fell_back: bool
|
||||
result_unit_ids: list[list[str]]
|
||||
|
||||
|
||||
async def _streaming_batch_write_ext(
|
||||
*,
|
||||
provider,
|
||||
ext_txn,
|
||||
pool,
|
||||
bank_id: str,
|
||||
fq_table,
|
||||
entity_resolver,
|
||||
phase1,
|
||||
batch_contents: list,
|
||||
batch_extracted: list,
|
||||
batch_processed: list,
|
||||
batch_chunk_meta: list,
|
||||
effective_doc_id: str,
|
||||
config,
|
||||
log_buffer: list[str],
|
||||
is_recovery: bool,
|
||||
is_first_batch: bool,
|
||||
is_last: bool,
|
||||
doc_tracking_done: list[bool],
|
||||
pipeline_aborted: list[bool],
|
||||
append_base_hash,
|
||||
new_content_hash,
|
||||
combined_content: str,
|
||||
retain_params,
|
||||
merged_tags,
|
||||
outbox_callback,
|
||||
assert_append_base_unchanged,
|
||||
p2_start: float,
|
||||
) -> _ExtStreamingWriteResult:
|
||||
"""Streaming batch write for a store that OWNS its memory rows in a SEPARATE system.
|
||||
|
||||
Unlike the Postgres path (one long transaction that also carries the memory write), this
|
||||
NEVER holds the data-plane connection across the object-store write. It runs in two phases:
|
||||
|
||||
1. STORE PHASE — no connection. Mint ids and stage the memory records (facts + causal edges,
|
||||
then a re-write carrying the resolved entity ids) to the object store, each tagged with
|
||||
``ext_txn`` so they stay INVISIBLE until :meth:`decide_txn`. Co-occurrence only accumulates
|
||||
in memory (flushed post-batch). No Postgres transaction is open.
|
||||
2. CONNECTION PHASE — a SHORT transaction: the document/chunk metadata rows, the entity
|
||||
registry reassert, the transactional-outbox row, and finally the commit witness. On commit
|
||||
the witness is the group's proof; ``decide_txn(commit=True)`` (a connection-free object-store
|
||||
marker) then publishes it. A crash before the witness commits leaves the staged writes for
|
||||
the recovery sweep to abort; a crash after leaves them for the sweep to commit.
|
||||
|
||||
The PG link writers are intentionally skipped: temporal/semantic links would touch zero rows
|
||||
(no ``memory_units`` for this org), and causal edges already travel on the memory record —
|
||||
writing them to PG ``memory_links`` would violate its deferrable FK to ``memory_units``.
|
||||
|
||||
``aborted`` in the result is True when a later batch lost the document to a concurrent
|
||||
takeover (the staged write is discarded); the call may also raise
|
||||
:class:`ConcurrentAppendConflict` for a lost append race, exactly like the Postgres path —
|
||||
the staged writes are discarded on that path too.
|
||||
"""
|
||||
# ---- STORE PHASE (no connection held) ----
|
||||
# Chunk ids are a deterministic function of identity (mirrors chunk_storage.store_chunks_batch),
|
||||
# so facts can be tagged with document_id + chunk_id before the metadata rows are written.
|
||||
chunk_id_by_index = {}
|
||||
if batch_chunk_meta:
|
||||
chunk_id_by_index = {
|
||||
cm.chunk_index: f"{bank_id}_{effective_doc_id}_{cm.chunk_index}" for cm in batch_chunk_meta
|
||||
}
|
||||
for fact, processed_fact in zip(batch_extracted, batch_processed):
|
||||
processed_fact.document_id = effective_doc_id
|
||||
if batch_chunk_meta and fact.chunk_index is not None:
|
||||
cid = chunk_id_by_index.get(fact.chunk_index)
|
||||
if cid:
|
||||
processed_fact.chunk_id = cid
|
||||
|
||||
# Stage the memory records to the store (conn unused by a store-owned backend), tagged with
|
||||
# ext_txn. This is the slow object-store write we are keeping OUT of the connection window.
|
||||
unit_ids = await fact_storage.insert_facts_batch(None, bank_id, batch_processed, ops=pool.ops, txn=ext_txn)
|
||||
batch_result_ids = _map_results_to_contents(batch_contents, batch_processed, unit_ids if unit_ids else [])
|
||||
|
||||
if unit_ids:
|
||||
# Remap Phase-1 placeholder ids onto the real unit ids, then re-write each memory with its
|
||||
# entity ids attached — also connection-free for a store-owned backend.
|
||||
resolved_entity_ids = [entity.entity_id for entity in phase1.entities.resolved_entities]
|
||||
remapped_entity_to_unit, _remapped_unit_to_entity_ids, _remapped_semantic = _remap_phase1_results(
|
||||
resolved_entity_ids, phase1.entities.entity_to_unit, phase1.entities.unit_to_entity_ids, [], unit_ids
|
||||
)
|
||||
unit_entity_pairs = [
|
||||
(unit_id, resolved_entity_ids[idx], fact_date)
|
||||
for idx, (unit_id, _local_idx, fact_date) in enumerate(remapped_entity_to_unit)
|
||||
]
|
||||
await entity_resolver.record_unit_entity_postings(unit_entity_pairs, bank_id=bank_id, txn=ext_txn)
|
||||
|
||||
# ---- CONNECTION PHASE (short transaction: local metadata + commit witness) ----
|
||||
try:
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
async with conn.transaction():
|
||||
# Ownership gate: lock the document row (serializes concurrent same-document
|
||||
# writers) and read its pre-existing hash for the takeover check.
|
||||
existing_hash = await pool.ops.lock_document_for_write(
|
||||
conn, fq_table("documents"), effective_doc_id, bank_id
|
||||
)
|
||||
|
||||
if not doc_tracking_done[0]:
|
||||
# Append compare-and-swap under the row lock (same as the Postgres path).
|
||||
assert_append_base_unchanged(existing_hash)
|
||||
if is_recovery:
|
||||
await fact_storage.upsert_document_metadata(
|
||||
conn,
|
||||
bank_id,
|
||||
effective_doc_id,
|
||||
combined_content,
|
||||
retain_params,
|
||||
merged_tags,
|
||||
store_document_text=getattr(config, "store_document_text", True),
|
||||
)
|
||||
log_buffer.append(
|
||||
f"[streaming] Document {effective_doc_id} updated (recovery, preserving existing chunks)"
|
||||
)
|
||||
else:
|
||||
await fact_storage.handle_document_tracking(
|
||||
conn,
|
||||
bank_id,
|
||||
effective_doc_id,
|
||||
combined_content,
|
||||
is_first_batch,
|
||||
retain_params,
|
||||
merged_tags,
|
||||
ops=pool.ops,
|
||||
store_document_text=getattr(config, "store_document_text", True),
|
||||
txn=ext_txn,
|
||||
)
|
||||
log_buffer.append(f"[streaming] Document {effective_doc_id} tracked (full content)")
|
||||
doc_tracking_done[0] = True
|
||||
else:
|
||||
# Later batches: verify we still own the document.
|
||||
if existing_hash is not None and existing_hash != new_content_hash:
|
||||
log_buffer.append(
|
||||
f"[streaming] Document {effective_doc_id} taken over by "
|
||||
f"concurrent request (hash mismatch) — aborting remaining batches"
|
||||
)
|
||||
logger.info("\n" + "\n".join(log_buffer) + "\n")
|
||||
if append_base_hash is not None:
|
||||
# The BaseException handler below discards the staged writes.
|
||||
raise ConcurrentAppendConflict(
|
||||
f"Document {effective_doc_id} was taken over by a concurrent "
|
||||
f"retain while this append was storing its batches"
|
||||
)
|
||||
# Discard the staged store writes rather than leave them for the sweep.
|
||||
await provider.decide_txn(ext_txn, commit=False)
|
||||
pipeline_aborted[0] = True
|
||||
return _ExtStreamingWriteResult(aborted=True, batch_result_ids=batch_result_ids)
|
||||
|
||||
# Chunk metadata rows (bulky bodies were already stored before this call).
|
||||
if batch_chunk_meta:
|
||||
await chunk_storage.store_chunks_batch(
|
||||
conn,
|
||||
bank_id,
|
||||
effective_doc_id,
|
||||
batch_chunk_meta,
|
||||
ops=pool.ops,
|
||||
store_document_text=getattr(config, "store_document_text", True),
|
||||
)
|
||||
|
||||
# Entity registry reassert (Postgres `entities`): re-create the resolved parents
|
||||
# this txn so a concurrent prune can't leave the postings dangling (#2662).
|
||||
if unit_ids:
|
||||
await entity_resolver.reassert_entities_batch(bank_id, phase1.entities.resolved_entities, conn=conn)
|
||||
|
||||
# Transactional-outbox row — must ride this Postgres transaction.
|
||||
if is_last and outbox_callback is not None:
|
||||
await outbox_callback(conn)
|
||||
|
||||
# The commit witness: its presence at commit is what the recovery sweep consults.
|
||||
await provider.write_txn_witness(ext_txn, conn=conn, fq_table=fq_table)
|
||||
|
||||
# Postgres committed the witness: publish the write-group (object-store marker, no conn).
|
||||
await provider.decide_txn(ext_txn, commit=True)
|
||||
logger.info(f"[streaming] Phase 2 (ext write txn): {time.time() - p2_start:.3f}s")
|
||||
except BaseException:
|
||||
# The witness never committed (this also covers a lost-append ConcurrentAppendConflict)
|
||||
# → make sure the staged store writes don't linger; the recovery sweep is the backstop
|
||||
# if this best-effort abort also fails.
|
||||
try:
|
||||
await provider.decide_txn(ext_txn, commit=False)
|
||||
except Exception:
|
||||
logger.warning(f"[streaming] best-effort abort of ext txn for {effective_doc_id} failed", exc_info=True)
|
||||
raise
|
||||
|
||||
return _ExtStreamingWriteResult(aborted=False, batch_result_ids=batch_result_ids)
|
||||
|
||||
|
||||
async def _delta_batch_write_ext(
|
||||
*,
|
||||
provider,
|
||||
ext_txn,
|
||||
pool,
|
||||
bank_id: str,
|
||||
fq_table,
|
||||
entity_resolver,
|
||||
phase1,
|
||||
effective_doc_id: str,
|
||||
config,
|
||||
log_buffer: list[str],
|
||||
processed_facts: list,
|
||||
extracted_facts: list,
|
||||
delta_contents: list,
|
||||
contents_dicts: list,
|
||||
document_tags,
|
||||
document_body_override,
|
||||
doc_hash_at_load,
|
||||
new_chunk_metadata: list,
|
||||
delta_chunk_map: dict,
|
||||
new_chunks_with_contents: dict,
|
||||
existing_by_index: dict,
|
||||
changed_indices: list,
|
||||
removed_indices: list,
|
||||
outbox_callback,
|
||||
) -> _ExtDeltaWriteResult:
|
||||
"""Delta re-retain write for a store that OWNS its memory rows in a SEPARATE system.
|
||||
|
||||
Same connection-management contract as :func:`_streaming_batch_write_ext`: the slow object-store
|
||||
writes (the new facts, then their entity re-write) plus the document-body upload are staged with
|
||||
NO connection held; the connection is taken only for the SHORT transaction that records the
|
||||
document/chunk metadata, the chunk tombstones, and the commit witness. ``fell_back`` True in
|
||||
the result means the document moved underneath us and the caller must redo the work on the
|
||||
streaming path.
|
||||
"""
|
||||
# ---- STORE PHASE (no connection held) ----
|
||||
if document_body_override is not None:
|
||||
combined_content = document_body_override
|
||||
else:
|
||||
combined_content = "\n".join([c.get("content", "") for c in contents_dicts])
|
||||
retain_params, merged_tags = _build_retain_params(contents_dicts, document_tags)
|
||||
|
||||
# Re-upload the document bodies (dedup by hash — only what changed moves). A store write, so it
|
||||
# belongs in the connection-free phase.
|
||||
await _store_document_bodies(
|
||||
bank_id=bank_id,
|
||||
document_id=effective_doc_id,
|
||||
combined_content=combined_content,
|
||||
chunk_texts=[new_chunks_with_contents[i] for i in sorted(new_chunks_with_contents)],
|
||||
merged_tags=merged_tags,
|
||||
config=config,
|
||||
)
|
||||
|
||||
# Deterministic chunk ids for the new/changed chunks (mirrors chunk_storage.store_chunks_batch
|
||||
# after the delta remap), so facts can be tagged before the metadata rows are written.
|
||||
remapped_new_indices = {delta_chunk_map.get(cm.chunk_index, cm.chunk_index) for cm in new_chunk_metadata}
|
||||
for ef, pf in zip(extracted_facts, processed_facts):
|
||||
pf.document_id = effective_doc_id
|
||||
if ef.chunk_index is not None:
|
||||
original_idx = delta_chunk_map.get(ef.chunk_index, ef.chunk_index)
|
||||
if original_idx in remapped_new_indices:
|
||||
pf.chunk_id = f"{bank_id}_{effective_doc_id}_{original_idx}"
|
||||
|
||||
# Stage the memory writes to the store (conn unused), tagged with ext_txn.
|
||||
unit_ids = await fact_storage.insert_facts_batch(None, bank_id, processed_facts, ops=pool.ops, txn=ext_txn)
|
||||
result_unit_ids = _map_results_to_contents(delta_contents, processed_facts, unit_ids if unit_ids else [])
|
||||
|
||||
if unit_ids:
|
||||
resolved_entity_ids = [entity.entity_id for entity in phase1.entities.resolved_entities]
|
||||
remapped_entity_to_unit, _r_u2e, _r_sem = _remap_phase1_results(
|
||||
resolved_entity_ids, phase1.entities.entity_to_unit, phase1.entities.unit_to_entity_ids, [], unit_ids
|
||||
)
|
||||
unit_entity_pairs = [
|
||||
(unit_id, resolved_entity_ids[idx], fact_date)
|
||||
for idx, (unit_id, _local_idx, fact_date) in enumerate(remapped_entity_to_unit)
|
||||
]
|
||||
await entity_resolver.record_unit_entity_postings(unit_entity_pairs, bank_id=bank_id, txn=ext_txn)
|
||||
|
||||
# ---- CONNECTION PHASE (short transaction: local metadata + tombstones + witness) ----
|
||||
try:
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
async with conn.transaction():
|
||||
# Ownership recheck: the delta diff was computed against a snapshot taken outside
|
||||
# this txn; if the document was replaced since, the diff is stale — fall back.
|
||||
current_hash = await conn.fetchval(
|
||||
f"SELECT content_hash FROM {fq_table('documents')} WHERE id = $1 AND bank_id = $2 FOR UPDATE",
|
||||
effective_doc_id,
|
||||
bank_id,
|
||||
)
|
||||
if current_hash is not None and doc_hash_at_load is not None and current_hash != doc_hash_at_load:
|
||||
log_buffer.append(
|
||||
f"[delta] Document {effective_doc_id} was modified by concurrent request "
|
||||
f"since chunks were loaded — aborting delta, falling back to full retain"
|
||||
)
|
||||
logger.info("\n" + "\n".join(log_buffer) + "\n")
|
||||
await provider.decide_txn(ext_txn, commit=False)
|
||||
return _ExtDeltaWriteResult(fell_back=True, result_unit_ids=result_unit_ids)
|
||||
|
||||
await fact_storage.upsert_document_metadata(
|
||||
conn, bank_id, effective_doc_id, combined_content, retain_params, merged_tags
|
||||
)
|
||||
|
||||
# Tombstone the changed/removed chunks' memories (store delete tagged ext_txn +
|
||||
# Postgres observation invalidation) — same write-group as the new facts above.
|
||||
chunks_to_delete = [
|
||||
existing_by_index[idx].chunk_id
|
||||
for idx in changed_indices + removed_indices
|
||||
if idx in existing_by_index
|
||||
]
|
||||
await chunk_storage.delete_chunks_by_ids(conn, chunks_to_delete, bank_id, txn=ext_txn, ops=pool.ops)
|
||||
|
||||
# Sync tags/metadata onto unchanged survivors (zero rows for a store-owned backend).
|
||||
await fact_storage.update_memory_units_metadata_and_tags(
|
||||
conn, bank_id, effective_doc_id, merged_tags, retain_params.get("metadata", {})
|
||||
)
|
||||
|
||||
# New/changed chunk metadata rows.
|
||||
if new_chunk_metadata:
|
||||
remapped_chunks = [
|
||||
ChunkMetadata(
|
||||
chunk_text=cm.chunk_text,
|
||||
fact_count=cm.fact_count,
|
||||
content_index=cm.content_index,
|
||||
chunk_index=delta_chunk_map.get(cm.chunk_index, cm.chunk_index),
|
||||
)
|
||||
for cm in new_chunk_metadata
|
||||
]
|
||||
await chunk_storage.store_chunks_batch(
|
||||
conn,
|
||||
bank_id,
|
||||
effective_doc_id,
|
||||
remapped_chunks,
|
||||
ops=pool.ops,
|
||||
store_document_text=getattr(config, "store_document_text", True),
|
||||
)
|
||||
|
||||
# Entity registry reassert (Postgres `entities`) — see the streaming path (#2662).
|
||||
if unit_ids:
|
||||
await entity_resolver.reassert_entities_batch(bank_id, phase1.entities.resolved_entities, conn=conn)
|
||||
|
||||
# Transactional-outbox row — must ride this Postgres transaction.
|
||||
if outbox_callback is not None:
|
||||
await outbox_callback(conn)
|
||||
|
||||
# The commit witness.
|
||||
await provider.write_txn_witness(ext_txn, conn=conn, fq_table=fq_table)
|
||||
|
||||
await provider.decide_txn(ext_txn, commit=True)
|
||||
except BaseException:
|
||||
try:
|
||||
await provider.decide_txn(ext_txn, commit=False)
|
||||
except Exception:
|
||||
logger.warning(f"[delta] best-effort abort of ext txn for {effective_doc_id} failed", exc_info=True)
|
||||
raise
|
||||
|
||||
return _ExtDeltaWriteResult(fell_back=False, result_unit_ids=result_unit_ids)
|
||||
|
||||
|
||||
async def _extract_and_embed(
|
||||
contents: list[RetainContent],
|
||||
llm_config,
|
||||
@@ -754,6 +1126,15 @@ async def retain_batch(
|
||||
didn't dedup (caller should treat as "bill full submitted content").
|
||||
See ``RetainResult.processed_content_tokens`` for details.
|
||||
"""
|
||||
# Before anything is written. A retain is the one operation that writes to BOTH stores —
|
||||
# documents, chunks and entities go to SQL through paths that never touch the memories
|
||||
# interface — so a store that needs a bank closed to writes (a backend cutover) cannot enforce
|
||||
# it from its own methods alone. Checked here, at the single entry every retain passes through,
|
||||
# rather than at each of the writes it fans out into.
|
||||
from ..memories import get_memories
|
||||
|
||||
await get_memories().assert_writable(bank_id)
|
||||
|
||||
start_time = time.time()
|
||||
total_chars = sum(len(item.get("content", "")) for item in contents_dicts)
|
||||
|
||||
@@ -973,9 +1354,23 @@ async def retain_batch(
|
||||
update_mode = item_mode
|
||||
break
|
||||
|
||||
if update_mode == "append" and effective_doc_id and is_first_batch:
|
||||
# The document version this append was built on. Captured with the text it
|
||||
# reads so the write path can prove nothing else appended in between — see
|
||||
# ``ConcurrentAppendConflict`` and the gate in ``_streaming_retain_batch``.
|
||||
# ``_APPEND_BASE_ABSENT`` distinguishes "read a document that wasn't there"
|
||||
# from "not an append", which None alone cannot express.
|
||||
append_base_hash: str | None = None
|
||||
is_append = update_mode == "append" and bool(effective_doc_id) and is_first_batch
|
||||
|
||||
if is_append:
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
existing_text = await fact_storage.get_document_content(conn, bank_id, effective_doc_id)
|
||||
base_row = await conn.fetchrow(
|
||||
f"SELECT original_text, content_hash FROM {fq_table('documents')} WHERE id = $1 AND bank_id = $2",
|
||||
effective_doc_id,
|
||||
bank_id,
|
||||
)
|
||||
existing_text = base_row["original_text"] if base_row else None
|
||||
append_base_hash = base_row["content_hash"] if base_row else _APPEND_BASE_ABSENT
|
||||
if existing_text:
|
||||
# Prepend existing text as a new content item at the beginning
|
||||
existing_content: RetainContentDict = {"content": existing_text}
|
||||
@@ -1041,6 +1436,20 @@ async def retain_batch(
|
||||
if doc_row and doc_row["updated_at"]:
|
||||
doc_updated = doc_row["updated_at"].timestamp()
|
||||
if doc_updated > start_time:
|
||||
# Under replace semantics dropping this request is right: a newer
|
||||
# submission of the same document already superseded it. Under
|
||||
# append semantics it is data loss — our content is a turn the
|
||||
# winner never saw — so raise and let the caller retry on top of
|
||||
# the newer document instead.
|
||||
if is_append:
|
||||
log_buffer.append(
|
||||
f"[append] Document {effective_doc_id} advanced before extraction — "
|
||||
f"retrying this append on the newer document"
|
||||
)
|
||||
logger.info("\n" + "\n".join(log_buffer) + "\n")
|
||||
raise ConcurrentAppendConflict(
|
||||
f"Document {effective_doc_id} was updated by a concurrent retain after this append read its content"
|
||||
)
|
||||
log_buffer.append(
|
||||
f"[stale] Skipping retain: document {effective_doc_id} was updated at "
|
||||
f"{doc_row['updated_at'].isoformat()} (after this request started at "
|
||||
@@ -1074,6 +1483,7 @@ async def retain_batch(
|
||||
outbox_callback,
|
||||
db_semaphore,
|
||||
document_body_override=document_body_override,
|
||||
append_base_hash=append_base_hash,
|
||||
)
|
||||
if delta_result is not None:
|
||||
return delta_result
|
||||
@@ -1138,6 +1548,7 @@ async def retain_batch(
|
||||
document_body_override=document_body_override,
|
||||
chunk_index_offset=chunk_index_offset,
|
||||
progress_callback=progress_callback,
|
||||
append_base_hash=append_base_hash,
|
||||
)
|
||||
|
||||
|
||||
@@ -1331,6 +1742,7 @@ async def _streaming_retain_batch(
|
||||
document_body_override: str | None = None,
|
||||
chunk_index_offset: int = 0,
|
||||
progress_callback: "Callable[..., Awaitable[None]] | None" = None,
|
||||
append_base_hash: str | None = None,
|
||||
) -> tuple[list[list[str]], TokenUsage]:
|
||||
"""
|
||||
Process a large document in streaming mini-batches to bound memory usage.
|
||||
@@ -1463,6 +1875,34 @@ async def _streaming_retain_batch(
|
||||
# stops processing further batches.
|
||||
pipeline_aborted: list[bool] = [False]
|
||||
|
||||
def _assert_append_base_unchanged(existing_hash: str | None) -> None:
|
||||
"""Fail the append if the document moved since it read its base text.
|
||||
|
||||
Called under the document row lock, on the write that establishes
|
||||
ownership. ``append_base_hash`` is the ``content_hash`` the append read
|
||||
alongside the text it concatenated onto; the row can only still carry
|
||||
that hash if no one else committed in between. A freshly created row
|
||||
reads back ``'__pending__'``, which is the expected value exactly when
|
||||
the append found no document at all.
|
||||
|
||||
No-op for replace-mode retains (``append_base_hash is None``), whose
|
||||
last-writer-wins semantics make a moved document the correct outcome
|
||||
rather than a conflict.
|
||||
"""
|
||||
if append_base_hash is None or existing_hash is None:
|
||||
return
|
||||
expected = "__pending__" if append_base_hash == _APPEND_BASE_ABSENT else append_base_hash
|
||||
if existing_hash == expected:
|
||||
return
|
||||
log_buffer.append(
|
||||
f"[append] Document {effective_doc_id} moved between the append read and this "
|
||||
f"write (expected {expected[:12]}, found {existing_hash[:12]}) — retrying on the newer document"
|
||||
)
|
||||
logger.info("\n" + "\n".join(log_buffer) + "\n")
|
||||
raise ConcurrentAppendConflict(
|
||||
f"Document {effective_doc_id} was updated by a concurrent retain while this append was extracting"
|
||||
)
|
||||
|
||||
# ---- LLM Producer ----
|
||||
# Fires all chunk extractions as concurrent tasks (bounded by the LLM
|
||||
# semaphore inside fact_extraction to 32 concurrent). As each completes
|
||||
@@ -1476,6 +1916,7 @@ async def _streaming_retain_batch(
|
||||
event_date=source.event_date,
|
||||
metadata=source.metadata,
|
||||
entities=source.entities,
|
||||
resolve_entities=source.resolve_entities,
|
||||
tags=source.tags,
|
||||
observation_scopes=source.observation_scopes,
|
||||
)
|
||||
@@ -1678,19 +2119,17 @@ async def _streaming_retain_batch(
|
||||
_edge_txn = None
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
async with conn.transaction():
|
||||
await conn.execute(
|
||||
f"INSERT INTO {fq_table('documents')} (id, bank_id, original_text, content_hash) "
|
||||
f"VALUES ($1, $2, '', '__pending__') "
|
||||
f"ON CONFLICT (id, bank_id) DO NOTHING",
|
||||
effective_doc_id,
|
||||
bank_id,
|
||||
)
|
||||
await conn.fetchval(
|
||||
f"SELECT content_hash FROM {fq_table('documents')} "
|
||||
f"WHERE id = $1 AND bank_id = $2 FOR UPDATE",
|
||||
# Same create-and-lock the fact-bearing path uses. Routed
|
||||
# through the ops layer so this branch takes the row lock
|
||||
# on Oracle too, and so the append gate below sees the
|
||||
# pre-existing hash rather than discarding it.
|
||||
existing_hash = await pool.ops.lock_document_for_write(
|
||||
conn,
|
||||
fq_table("documents"),
|
||||
effective_doc_id,
|
||||
bank_id,
|
||||
)
|
||||
_assert_append_base_unchanged(existing_hash)
|
||||
if is_recovery:
|
||||
await fact_storage.upsert_document_metadata(
|
||||
conn,
|
||||
@@ -1719,6 +2158,10 @@ async def _streaming_retain_batch(
|
||||
store_document_text=getattr(config, "store_document_text", True),
|
||||
txn=_edge_txn,
|
||||
)
|
||||
# Re-record the witness now that the group's writes have happened, so
|
||||
# the row carries what they actually wrote. `begin_txn` recorded it
|
||||
# before any write existed; the upsert widens rather than replaces.
|
||||
await _edge_provider.write_txn_witness(_edge_txn, conn=conn, fq_table=fq_table)
|
||||
doc_tracking_done[0] = True
|
||||
# Memory: combined_content has been persisted; release
|
||||
# it now so the rest of the consumer loop doesn't pin
|
||||
@@ -1783,6 +2226,65 @@ async def _streaming_retain_batch(
|
||||
|
||||
p2_start = time.time()
|
||||
batch_result_ids = None
|
||||
|
||||
# A store that owns its memory rows in a SEPARATE system returns a write-group handle
|
||||
# from mint_txn (Postgres returns None). For that store we must NOT hold the data-plane
|
||||
# connection across the object-store write, so we run a distinct connection-management
|
||||
# path. Postgres falls through to the single-transaction path below, unchanged.
|
||||
from ..memories import get_memories
|
||||
|
||||
_ext_provider = get_memories()
|
||||
_ext_txn = await _ext_provider.mint_txn(bank_id=bank_id, mutating=True)
|
||||
if _ext_txn is not None:
|
||||
ext_result = await _streaming_batch_write_ext(
|
||||
provider=_ext_provider,
|
||||
ext_txn=_ext_txn,
|
||||
pool=pool,
|
||||
bank_id=bank_id,
|
||||
fq_table=fq_table,
|
||||
entity_resolver=entity_resolver,
|
||||
phase1=phase1,
|
||||
batch_contents=batch_contents,
|
||||
batch_extracted=batch_extracted,
|
||||
batch_processed=batch_processed,
|
||||
batch_chunk_meta=batch_chunk_meta,
|
||||
effective_doc_id=effective_doc_id,
|
||||
config=config,
|
||||
log_buffer=log_buffer,
|
||||
is_recovery=is_recovery,
|
||||
is_first_batch=is_first_batch,
|
||||
is_last=is_last,
|
||||
doc_tracking_done=doc_tracking_done,
|
||||
pipeline_aborted=pipeline_aborted,
|
||||
append_base_hash=append_base_hash,
|
||||
new_content_hash=new_content_hash,
|
||||
combined_content=combined_content,
|
||||
retain_params=retain_params,
|
||||
merged_tags=merged_tags,
|
||||
outbox_callback=outbox_callback,
|
||||
assert_append_base_unchanged=_assert_append_base_unchanged,
|
||||
p2_start=p2_start,
|
||||
)
|
||||
# Doc-tracking consumed combined_content on the first batch; release it (mirrors
|
||||
# the Postgres path's first-batch reset).
|
||||
combined_content = ""
|
||||
if not ext_result.aborted:
|
||||
# The short txn above committed the transactional-outbox row; record it so
|
||||
# the post-loop fallback doesn't queue a duplicate delivery.
|
||||
if is_last and outbox_callback is not None:
|
||||
outbox_fired[0] = True
|
||||
# Deferred-stats flush + unit collection — mirrors the shared tail the Postgres
|
||||
# path reaches after its connection block exits.
|
||||
try:
|
||||
await entity_resolver.flush_pending_stats()
|
||||
except Exception:
|
||||
logger.warning(
|
||||
f"Entity stats flush (consumer batch {consumer_batch_idx + 1}) failed", exc_info=True
|
||||
)
|
||||
for content_ids in ext_result.batch_result_ids:
|
||||
all_unit_ids.extend(content_ids)
|
||||
return
|
||||
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
async with conn.transaction():
|
||||
# --- Document ownership gate ---
|
||||
@@ -1801,6 +2303,12 @@ async def _streaming_retain_batch(
|
||||
bank_id,
|
||||
)
|
||||
|
||||
# Append compare-and-swap, under the row lock and before any
|
||||
# write-group opens: an append that lost its read-modify-write
|
||||
# race must abort here rather than commit over the winner.
|
||||
if not doc_tracking_done[0]:
|
||||
_assert_append_base_unchanged(existing_hash)
|
||||
|
||||
# Open the cross-store write-group txn INSIDE this batch's transaction,
|
||||
# before the first-batch replace deletes any outgoing memories: the delete
|
||||
# and this batch's writes must ride the same txn so they commit together.
|
||||
@@ -1856,12 +2364,21 @@ async def _streaming_retain_batch(
|
||||
f"concurrent request (hash mismatch) — aborting remaining batches"
|
||||
)
|
||||
logger.info("\n" + "\n".join(log_buffer) + "\n")
|
||||
# Signal the consumer to stop processing further batches
|
||||
pipeline_aborted[0] = True
|
||||
# Abort the write-group we just opened rather than leaving it for the
|
||||
# recovery sweep — we wrote nothing this batch and are bailing. No-op for
|
||||
# the Postgres store (begin_txn returned None).
|
||||
await _provider.decide_txn(_group_txn, commit=False)
|
||||
# Discarding the rest is only acceptable under replace
|
||||
# semantics, where the winner's content supersedes ours.
|
||||
# An append's remaining batches carry content nobody
|
||||
# else has, so raise and redo the whole append instead.
|
||||
if append_base_hash is not None:
|
||||
raise ConcurrentAppendConflict(
|
||||
f"Document {effective_doc_id} was taken over by a concurrent "
|
||||
f"retain while this append was storing its batches"
|
||||
)
|
||||
# Signal the consumer to stop processing further batches
|
||||
pipeline_aborted[0] = True
|
||||
return
|
||||
|
||||
# Store chunks with correct global indices
|
||||
@@ -1909,6 +2426,11 @@ async def _streaming_retain_batch(
|
||||
txn=_group_txn,
|
||||
)
|
||||
|
||||
# Last thing inside the transaction: re-record the witness now that this
|
||||
# batch's writes have happened, so the row carries what they actually wrote.
|
||||
# `begin_txn` above recorded it before any write existed; the upsert widens.
|
||||
await _provider.write_txn_witness(_group_txn, conn=conn, fq_table=fq_table)
|
||||
|
||||
# Postgres committed this batch: publish its write-group. If it had aborted,
|
||||
# this is skipped and the recovery sweep resolves the undecided txn (spec §5).
|
||||
await _provider.decide_txn(_group_txn, commit=True)
|
||||
@@ -2087,6 +2609,10 @@ async def _streaming_retain_batch(
|
||||
store_document_text=getattr(config, "store_document_text", True),
|
||||
txn=_edge_txn,
|
||||
)
|
||||
# Re-record the witness now that the group's writes have happened, so the
|
||||
# row carries what they actually wrote. `begin_txn` recorded it before any
|
||||
# write existed; the upsert widens rather than replaces.
|
||||
await _edge_provider.write_txn_witness(_edge_txn, conn=conn, fq_table=fq_table)
|
||||
doc_tracking_done[0] = True
|
||||
# Memory: combined_content has been persisted and won't be
|
||||
# read again — release the per-document text now.
|
||||
@@ -2263,6 +2789,7 @@ async def _try_delta_retain(
|
||||
db_semaphore: "asyncio.Semaphore | None" = None,
|
||||
*,
|
||||
document_body_override: str | None = None,
|
||||
append_base_hash: str | None = None,
|
||||
) -> tuple[list[list[str]], TokenUsage, int | None] | None:
|
||||
"""
|
||||
Attempt delta retain for a document upsert. Returns result tuple if delta
|
||||
@@ -2307,6 +2834,19 @@ async def _try_delta_retain(
|
||||
# the extraction freshness recheck below) forces a streaming fallback.
|
||||
existing_chunks = await chunk_storage.load_existing_chunks(conn, bank_id, effective_doc_id)
|
||||
|
||||
# For an append, the document this delta plans against must still be the one
|
||||
# whose text the append concatenated onto. Every write below is gated on
|
||||
# ``doc_hash_at_load``, so a document that already moved would let the delta
|
||||
# commit content assembled from a stale base — losing the turn that moved it.
|
||||
# Cheapest possible place to notice: before any chunking or extraction.
|
||||
if append_base_hash is not None:
|
||||
expected_base = "__pending__" if append_base_hash == _APPEND_BASE_ABSENT else append_base_hash
|
||||
if doc_hash_at_load is not None and doc_hash_at_load != expected_base:
|
||||
raise ConcurrentAppendConflict(
|
||||
f"Document {effective_doc_id} was updated by a concurrent retain "
|
||||
f"between this append's read and its delta plan"
|
||||
)
|
||||
|
||||
if not existing_chunks:
|
||||
return None
|
||||
|
||||
@@ -2488,7 +3028,14 @@ async def _try_delta_retain(
|
||||
result_unit_ids: list[list[str]] = []
|
||||
log_buffer_pre_db = len(log_buffer)
|
||||
|
||||
async def _run_delta_db_work() -> None:
|
||||
async def _run_delta_db_work() -> bool:
|
||||
"""Write this delta. Returns False when the document moved underneath it.
|
||||
|
||||
The caller must translate False into "fall back to the streaming path" —
|
||||
this used to be declared ``-> None`` with a bare ``return None`` on the
|
||||
abort branch, so the guard logged that it was falling back while the
|
||||
delta actually committed on top of the concurrent writer.
|
||||
"""
|
||||
nonlocal result_unit_ids
|
||||
del log_buffer[log_buffer_pre_db:]
|
||||
for pf in processed_facts:
|
||||
@@ -2501,6 +3048,51 @@ async def _try_delta_retain(
|
||||
pool, entity_resolver, bank_id, delta_contents, processed_facts, config, log_buffer
|
||||
)
|
||||
|
||||
# A store that owns its rows in a separate system uses a distinct connection-management
|
||||
# path (mint_txn returns a handle; Postgres returns None and takes the path below,
|
||||
# unchanged) so the data-plane connection is not held across the object-store write.
|
||||
from ..memories import get_memories
|
||||
|
||||
_ext_provider = get_memories()
|
||||
_ext_txn = await _ext_provider.mint_txn(bank_id=bank_id, mutating=True)
|
||||
if _ext_txn is not None:
|
||||
delta_result = await _delta_batch_write_ext(
|
||||
provider=_ext_provider,
|
||||
ext_txn=_ext_txn,
|
||||
pool=pool,
|
||||
bank_id=bank_id,
|
||||
fq_table=fq_table,
|
||||
entity_resolver=entity_resolver,
|
||||
phase1=phase1,
|
||||
effective_doc_id=effective_doc_id,
|
||||
config=config,
|
||||
log_buffer=log_buffer,
|
||||
processed_facts=processed_facts,
|
||||
extracted_facts=extracted_facts,
|
||||
delta_contents=delta_contents,
|
||||
contents_dicts=contents_dicts,
|
||||
document_tags=document_tags,
|
||||
document_body_override=document_body_override,
|
||||
doc_hash_at_load=doc_hash_at_load,
|
||||
new_chunk_metadata=new_chunk_metadata,
|
||||
delta_chunk_map=delta_chunk_map,
|
||||
new_chunks_with_contents=new_chunks_with_contents,
|
||||
existing_by_index=existing_by_index,
|
||||
changed_indices=changed_indices,
|
||||
removed_indices=removed_indices,
|
||||
outbox_callback=outbox_callback,
|
||||
)
|
||||
if delta_result.fell_back:
|
||||
return False
|
||||
result_unit_ids = delta_result.result_unit_ids
|
||||
log_buffer.append(f"DELTA RETAIN COMPLETE (ext store): {len(processed_facts)} new units")
|
||||
logger.info("\n" + "\n".join(log_buffer) + "\n")
|
||||
try:
|
||||
await entity_resolver.flush_pending_stats()
|
||||
except Exception:
|
||||
logger.warning("Entity stats flush failed — retrieval unaffected", exc_info=True)
|
||||
return True
|
||||
|
||||
# PHASE 2 — Core Write Transaction (atomic)
|
||||
# Lock the document row and verify ownership. Delta loaded existing
|
||||
# chunks OUTSIDE this TXN, so a concurrent retain may have cascade-deleted
|
||||
@@ -2521,8 +3113,9 @@ async def _try_delta_retain(
|
||||
f"since chunks were loaded — aborting delta, falling back to full retain"
|
||||
)
|
||||
logger.info("\n" + "\n".join(log_buffer) + "\n")
|
||||
# Return None to fall back to streaming (which has full FOR UPDATE protection)
|
||||
return None
|
||||
# Fall back to streaming, which re-locks the document and (for
|
||||
# an append) verifies the base this content was built on.
|
||||
return False
|
||||
|
||||
# Update document metadata (no delete)
|
||||
step_start = time.time()
|
||||
@@ -2656,6 +3249,12 @@ async def _try_delta_retain(
|
||||
txn=_group_txn,
|
||||
)
|
||||
|
||||
# Last thing inside the transaction: re-record the witness now that the group's
|
||||
# writes have happened, so the row carries what they actually wrote. `begin_txn`
|
||||
# above recorded it before any write existed; the upsert widens rather than
|
||||
# replaces.
|
||||
await _provider.write_txn_witness(_group_txn, conn=conn, fq_table=fq_table)
|
||||
|
||||
# Postgres has committed: publish the write-group so its writes become visible.
|
||||
# If the transaction had aborted instead, this line is skipped and the recovery
|
||||
# sweep resolves the undecided txn against the (absent) witness row (spec §5).
|
||||
@@ -2679,11 +3278,18 @@ async def _try_delta_retain(
|
||||
except Exception:
|
||||
logger.warning("Entity stats flush failed — retrieval unaffected", exc_info=True)
|
||||
|
||||
return True
|
||||
|
||||
if db_semaphore is not None:
|
||||
async with db_semaphore:
|
||||
await _run_delta_db_work()
|
||||
delta_committed = await _run_delta_db_work()
|
||||
else:
|
||||
await _run_delta_db_work()
|
||||
delta_committed = await _run_delta_db_work()
|
||||
if not delta_committed:
|
||||
# The document moved while this delta was extracting. Nothing was
|
||||
# written; the streaming path redoes the work under its own lock, and
|
||||
# for an append its base check turns the loss into a retry.
|
||||
return None
|
||||
await _record_retain_document_outcome(pool, bank_id, effective_doc_id, sum(len(ids) for ids in result_unit_ids))
|
||||
# Count content + context tokens that actually went through extraction.
|
||||
# ``delta_contents`` holds the per-chunk RetainContent items for the
|
||||
@@ -2783,6 +3389,7 @@ def _build_contents(contents_dicts: list[RetainContentDict], document_tags: list
|
||||
event_date=event_date_value,
|
||||
metadata=item.get("metadata", {}),
|
||||
entities=item.get("entities", []),
|
||||
resolve_entities=item.get("resolve_entities", True),
|
||||
tags=merged_tags,
|
||||
observation_scopes=item.get("observation_scopes"),
|
||||
)
|
||||
@@ -2844,6 +3451,7 @@ def _build_delta_contents(
|
||||
event_date=template_content.event_date,
|
||||
metadata=template_content.metadata,
|
||||
entities=template_content.entities,
|
||||
resolve_entities=template_content.resolve_entities,
|
||||
tags=template_content.tags,
|
||||
observation_scopes=template_content.observation_scopes,
|
||||
)
|
||||
|
||||
@@ -11,6 +11,8 @@ from datetime import datetime
|
||||
from typing import Literal, TypedDict
|
||||
from uuid import UUID
|
||||
|
||||
from ..metadata_utils import drop_null_values
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@@ -24,6 +26,8 @@ class RetainContentDict(TypedDict, total=False):
|
||||
metadata: Custom key-value metadata (optional)
|
||||
document_id: Document ID for this content item (optional)
|
||||
entities: User-provided entities to merge with extracted entities (optional)
|
||||
resolve_entities: Whether the supplied `entities` are resolved against the bank's
|
||||
existing entities (optional, default True). False takes them literally.
|
||||
tags: Visibility scope tags for this content item (optional)
|
||||
observation_scopes: How to scope observations for consolidation (optional).
|
||||
"per_tag" runs one pass per individual tag; "combined" (default) runs a
|
||||
@@ -41,6 +45,7 @@ class RetainContentDict(TypedDict, total=False):
|
||||
metadata: dict[str, str]
|
||||
document_id: str
|
||||
entities: list[dict[str, str]] # [{"text": "...", "type": "..."}]
|
||||
resolve_entities: bool
|
||||
tags: list[str] # Visibility scope tags
|
||||
observation_scopes: (
|
||||
Literal["per_tag", "combined", "all_combinations", "shared"] | list[list[str]]
|
||||
@@ -48,6 +53,18 @@ class RetainContentDict(TypedDict, total=False):
|
||||
update_mode: Literal["replace", "append"]
|
||||
|
||||
|
||||
@dataclass
|
||||
class UserEntities:
|
||||
"""The entities a caller supplied for one retain content item, and how to match them.
|
||||
|
||||
Kept together so the resolution choice travels with the names it applies to: retain merges
|
||||
these with the extractor's own entities into one batch, and only these are authoritative.
|
||||
"""
|
||||
|
||||
entities: list[dict[str, str]]
|
||||
resolve: bool = True
|
||||
|
||||
|
||||
@dataclass
|
||||
class RetainContent:
|
||||
"""
|
||||
@@ -61,11 +78,23 @@ class RetainContent:
|
||||
event_date: datetime | None = None
|
||||
metadata: dict[str, str] = field(default_factory=dict)
|
||||
entities: list[dict[str, str]] = field(default_factory=list) # User-provided entities
|
||||
# Whether the supplied `entities` are matched against the bank's existing entities. False
|
||||
# takes them literally; extracted entities are always resolved either way (#3479).
|
||||
resolve_entities: bool = True
|
||||
tags: list[str] = field(default_factory=list) # Visibility scope tags
|
||||
observation_scopes: Literal["per_tag", "combined", "all_combinations", "shared"] | list[list[str]] | None = (
|
||||
None # Observation scopes
|
||||
)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
# Drop null-valued metadata keys (issue #3209): the retain API accepts
|
||||
# arbitrary JSON metadata, and a null value stored verbatim poisons the
|
||||
# read path, which validates MemoryFact.metadata as dict[str, str].
|
||||
# Non-string values are preserved; the read path coerces them. An
|
||||
# explicit ``"metadata": null`` in the request normalizes to {} so the
|
||||
# field always matches its declared type.
|
||||
self.metadata = drop_null_values(self.metadata)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ChunkMetadata:
|
||||
@@ -348,3 +377,19 @@ class RetainBatch:
|
||||
|
||||
# Results (populated after storage)
|
||||
unit_ids_by_content: list[list[str]] = field(default_factory=list)
|
||||
|
||||
|
||||
class ConcurrentAppendConflict(Exception):
|
||||
"""An append-mode retain lost its read-modify-write race for a document.
|
||||
|
||||
``update_mode="append"`` reads ``documents.original_text``, concatenates the
|
||||
new content onto it, and reprocesses the result. That read and the write
|
||||
that follows it are separated by LLM extraction, so a second append
|
||||
committing in between would make this request overwrite a turn it never
|
||||
saw. Every write path that can observe the document having moved raises
|
||||
this instead of dropping the submission, so a lost race costs a retry
|
||||
rather than the caller's content.
|
||||
|
||||
Retryable by construction: the retry re-reads the (now newer) stored text
|
||||
and re-appends the same submission on top of it.
|
||||
"""
|
||||
|
||||
@@ -0,0 +1,111 @@
|
||||
"""IDF-aware BM25 query-term selection for the native tsvector backend.
|
||||
|
||||
A long recall query is tokenized and OR-joined into one ``tsquery``
|
||||
(``tok1 | tok2 | ...``). On the native backend the ``@@`` gate then matches a
|
||||
large fraction of the bank and ``ts_rank_cd`` — which has no IDF and is not
|
||||
index-backed — is computed for *every* matched row, so ``ORDER BY ... LIMIT``
|
||||
cannot prune before ranking. A query that matches thousands of memories times
|
||||
out (the production +60s BM25 hang this module addresses).
|
||||
|
||||
Blindly truncating to the first N tokens is the wrong cut: it keeps whichever
|
||||
terms happen to come first, which are usually the common, low-signal words that
|
||||
drive the fan-out, and drops the discriminative ones. Instead we keep the N
|
||||
tokens with the *lowest* corpus document frequency — the most selective,
|
||||
highest-signal terms, which is exactly what BM25's IDF weighting favours.
|
||||
|
||||
The document frequencies come from ``pg_stats.most_common_elems`` for
|
||||
``memory_units.search_vector``, a statistic PostgreSQL's ``ANALYZE`` maintains
|
||||
for free (autovacuum-refreshed). No new table, no index change, no reindex.
|
||||
|
||||
Caveats, by design:
|
||||
- The stats are per *table*, i.e. tenant-global across all banks in the schema,
|
||||
not per bank. A term rare tenant-wide but hot in a single bank is not caught
|
||||
here; that residual case is what a statement-timeout backstop is for.
|
||||
- Only the most-common lexemes are tracked, so a query term absent from the
|
||||
stats is treated as rare (df 0) and kept — precisely what we want.
|
||||
- On any failure (stats absent on a freshly loaded table, permissions, an
|
||||
unexpected catalog shape) we fall back to keeping the first N tokens, so
|
||||
selection can never block recall.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Map each query token to the maximum document frequency among its lexemes (as
|
||||
# produced by the same text-search config that built ``search_vector``), reading
|
||||
# the per-lexeme frequencies ANALYZE stored in ``pg_stats.most_common_elems``.
|
||||
#
|
||||
# ``most_common_elem_freqs`` carries the per-element frequencies followed by
|
||||
# trailing summary values (min/max/null frequency), so it is sliced to the
|
||||
# length of ``most_common_elems``. The ``::text::text[]`` round-trip is the
|
||||
# standard way to unnest the view's ``anyarray`` element column.
|
||||
#
|
||||
# Tokens whose text yields no lexeme (stopwords) come back with
|
||||
# ``has_lexeme = false`` and are dropped by the caller; tokens whose lexemes are
|
||||
# absent from the stats get df 0 (rare → kept first).
|
||||
_TOKEN_DF_SQL = """
|
||||
WITH toks AS (
|
||||
SELECT ord, tok
|
||||
FROM unnest($1::text[]) WITH ORDINALITY AS u(tok, ord)
|
||||
),
|
||||
lex AS (
|
||||
SELECT t.ord, l.lexeme
|
||||
FROM toks t
|
||||
LEFT JOIN LATERAL unnest(tsvector_to_array(to_tsvector($4::regconfig, t.tok))) AS l(lexeme) ON true
|
||||
),
|
||||
stats AS (
|
||||
SELECT unnest(most_common_elems::text::text[]) AS lexeme,
|
||||
unnest((most_common_elem_freqs)[1:array_length(most_common_elems::text::text[], 1)]) AS freq
|
||||
FROM pg_stats
|
||||
WHERE schemaname = $2 AND tablename = $3 AND attname = 'search_vector'
|
||||
)
|
||||
SELECT lex.ord AS ord,
|
||||
bool_or(lex.lexeme IS NOT NULL) AS has_lexeme,
|
||||
COALESCE(MAX(s.freq), 0)::float8 AS df
|
||||
FROM lex
|
||||
LEFT JOIN stats s ON s.lexeme = lex.lexeme
|
||||
GROUP BY lex.ord
|
||||
ORDER BY lex.ord
|
||||
"""
|
||||
|
||||
|
||||
async def select_selective_bm25_tokens(
|
||||
conn,
|
||||
tokens: list[str],
|
||||
*,
|
||||
schema: str,
|
||||
table: str,
|
||||
language: str,
|
||||
max_terms: int,
|
||||
) -> list[str]:
|
||||
"""Return at most ``max_terms`` tokens, preferring the lowest-df (most selective).
|
||||
|
||||
The returned tokens keep their original relative order (the OR ``tsquery`` is
|
||||
order-insensitive; preserving order keeps traces and logs readable). Falls
|
||||
back to the first ``max_terms`` tokens whenever document-frequency stats
|
||||
cannot be read, so recall is never blocked by a stats problem.
|
||||
"""
|
||||
if max_terms <= 0 or len(tokens) <= max_terms:
|
||||
return tokens
|
||||
|
||||
try:
|
||||
rows = await conn.fetch(_TOKEN_DF_SQL, tokens, schema, table, language)
|
||||
except Exception:
|
||||
logger.debug("BM25 term-df lookup failed; falling back to first-N cap", exc_info=True)
|
||||
return tokens[:max_terms]
|
||||
|
||||
if not rows:
|
||||
return tokens[:max_terms]
|
||||
|
||||
# ``ord`` is 1-based into ``tokens``. Keep only real search terms (drop
|
||||
# stopwords), then take the ``max_terms`` lowest-df, ties broken by position.
|
||||
scored = [(row["df"], row["ord"]) for row in rows if row["has_lexeme"]]
|
||||
if not scored:
|
||||
return tokens[:max_terms]
|
||||
|
||||
scored.sort(key=lambda df_ord: df_ord)
|
||||
kept_ords = sorted(ord_ for _, ord_ in scored[:max_terms])
|
||||
return [tokens[ord_ - 1] for ord_ in kept_ords]
|
||||
@@ -64,23 +64,25 @@ async def _find_semantic_seeds(
|
||||
tag_groups_param_start = 6 + (1 if tags else 0)
|
||||
groups_clause, groups_params, _ = build_tag_groups_where_clause(tag_groups, tag_groups_param_start)
|
||||
|
||||
# created_after/created_before filter `updated_at`, matching the other recall arms
|
||||
# (see retrieval.py) so a window narrows every arm the same way.
|
||||
_next_idx = tag_groups_param_start + len(groups_params)
|
||||
created_range_clause = ""
|
||||
created_range_params: list[Any] = []
|
||||
updated_range_clause = ""
|
||||
updated_range_params: list[Any] = []
|
||||
if created_after is not None:
|
||||
created_range_params.append(created_after)
|
||||
created_range_clause += f" AND updated_at > ${_next_idx}"
|
||||
updated_range_params.append(created_after)
|
||||
updated_range_clause += f" AND updated_at > ${_next_idx}"
|
||||
_next_idx += 1
|
||||
if created_before is not None:
|
||||
created_range_params.append(created_before)
|
||||
created_range_clause += f" AND updated_at < ${_next_idx}"
|
||||
updated_range_params.append(created_before)
|
||||
updated_range_clause += f" AND updated_at < ${_next_idx}"
|
||||
_next_idx += 1
|
||||
|
||||
params = [query_embedding_str, bank_id, fact_type, threshold, limit]
|
||||
if tags:
|
||||
params.append(tags)
|
||||
params.extend(groups_params)
|
||||
params.extend(created_range_params)
|
||||
params.extend(updated_range_params)
|
||||
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
@@ -94,7 +96,7 @@ async def _find_semantic_seeds(
|
||||
AND (1 - (embedding <=> $1::vector)) >= $4
|
||||
{tags_clause}
|
||||
{groups_clause}
|
||||
{created_range_clause}
|
||||
{updated_range_clause}
|
||||
ORDER BY embedding <=> $1::vector
|
||||
LIMIT $5
|
||||
""",
|
||||
|
||||
@@ -8,18 +8,17 @@ Implements:
|
||||
4. Temporal retrieval (time-aware search with spreading)
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import re
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import UTC, datetime
|
||||
from typing import TYPE_CHECKING, Any, Optional
|
||||
|
||||
from ...config import DEFAULT_BM25_MAX_QUERY_TERMS, DEFAULT_TEMPORAL_SEMANTIC_MIN_SIMILARITY, get_config
|
||||
from ...config import DEFAULT_BM25_MAX_QUERY_TERMS, get_config
|
||||
from ..db.ops import UpdatedWindow
|
||||
from ..db_utils import acquire_with_retry
|
||||
from ..memory_engine import fq_table
|
||||
from ..memory_engine import fq_table, get_current_schema
|
||||
from ..sql import create_sql_dialect
|
||||
from .bm25_term_selection import select_selective_bm25_tokens
|
||||
from .graph_retrieval import GraphRetriever
|
||||
from .link_expansion_retrieval import GRAPH_SEED_LIMIT, LinkExpansionRetriever
|
||||
from .tags import TagGroup, TagsMatch, build_tag_groups_where_clause, build_tags_where_clause_simple
|
||||
@@ -120,47 +119,6 @@ def set_default_graph_retriever(retriever: GraphRetriever | None) -> None:
|
||||
_default_graph_retriever = retriever
|
||||
|
||||
|
||||
async def retrieve_semantic_bm25_combined(
|
||||
conn,
|
||||
query_emb_str: str,
|
||||
query_text: str,
|
||||
bank_id: str,
|
||||
fact_types: list[str],
|
||||
limit: int,
|
||||
tags: list[str] | None = None,
|
||||
tags_match: TagsMatch = "any",
|
||||
tag_groups: list[TagGroup] | 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]:
|
||||
"""Combined semantic + BM25 retrieval, run by the configured memories store.
|
||||
|
||||
With the default Postgres store this calls straight through to
|
||||
:func:`retrieve_semantic_bm25_combined_sql` below — same query, same results.
|
||||
"""
|
||||
from ..memories import get_memories
|
||||
|
||||
return await get_memories().search(
|
||||
conn=conn,
|
||||
bank_id=bank_id,
|
||||
fact_types=fact_types,
|
||||
query_embedding=query_emb_str,
|
||||
query_text=query_text,
|
||||
limit=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 retrieve_semantic_bm25_combined_sql(
|
||||
conn,
|
||||
query_emb_str: str,
|
||||
@@ -256,19 +214,22 @@ async def retrieve_semantic_bm25_combined_sql(
|
||||
tag_groups_param_start = tags_param_idx + (1 if tags else 0)
|
||||
groups_clause, groups_params, _ = build_tag_groups_where_clause(tag_groups, tag_groups_param_start)
|
||||
|
||||
# --- created_at time range filter (appended after tags/groups) ---
|
||||
# --- created_after/created_before time range filter (appended after tags/groups) ---
|
||||
# The bounds are named for creation but filter `updated_at` — "memories that changed
|
||||
# in this window", so an edited fact re-enters it. That is what the mental-model delta
|
||||
# refresh needs from its watermark; see META_UPDATED_AT in engine/memories/base.py.
|
||||
# Param indices are computed relative to the final params list built below,
|
||||
# so we pre-compute the next available index after all preceding params.
|
||||
_next_idx = tag_groups_param_start + len(groups_params)
|
||||
created_range_clause = ""
|
||||
created_range_params: list[Any] = []
|
||||
updated_range_clause = ""
|
||||
updated_range_params: list[Any] = []
|
||||
if created_after is not None:
|
||||
created_range_params.append(created_after)
|
||||
created_range_clause += f" AND updated_at > ${_next_idx}"
|
||||
updated_range_params.append(created_after)
|
||||
updated_range_clause += f" AND updated_at > ${_next_idx}"
|
||||
_next_idx += 1
|
||||
if created_before is not None:
|
||||
created_range_params.append(created_before)
|
||||
created_range_clause += f" AND updated_at < ${_next_idx}"
|
||||
updated_range_params.append(created_before)
|
||||
updated_range_clause += f" AND updated_at < ${_next_idx}"
|
||||
_next_idx += 1
|
||||
|
||||
# --- Semantic UNION ALL arms (one per fact_type) ---
|
||||
@@ -285,7 +246,7 @@ async def retrieve_semantic_bm25_combined_sql(
|
||||
min_similarity=sem_min,
|
||||
tags_clause=tags_clause,
|
||||
groups_clause=groups_clause,
|
||||
extra_where=created_range_clause,
|
||||
extra_where=updated_range_clause,
|
||||
)
|
||||
for ft in fact_types
|
||||
]
|
||||
@@ -293,11 +254,36 @@ async def retrieve_semantic_bm25_combined_sql(
|
||||
# --- BM25 UNION ALL arms (one per fact_type, only when tokens present) ---
|
||||
if _include_bm25:
|
||||
text_ext = config.text_search_extension
|
||||
max_query_terms = getattr(config, "bm25_max_query_terms", DEFAULT_BM25_MAX_QUERY_TERMS)
|
||||
bm25_tokens = tokens
|
||||
# Native tsvector has no IDF and ranks every `@@` match, so a long OR
|
||||
# query over common terms scans and ranks a large fraction of the bank
|
||||
# (the +60s prod timeout). Keep only the most selective terms — lowest
|
||||
# tenant-wide document frequency, read for free from pg_stats — which
|
||||
# bounds both the match set and the per-row rank cost while preserving
|
||||
# the high-signal terms a blunt first-N cap would discard. PG-native
|
||||
# only; best-effort (falls back to first-N when stats are unavailable).
|
||||
# Opt out via bm25_selective_terms to cap by position instead.
|
||||
if (
|
||||
text_ext == "native"
|
||||
and max_query_terms > 0
|
||||
and len(tokens) > max_query_terms
|
||||
and getattr(config, "bm25_selective_terms", True)
|
||||
and getattr(conn, "backend_type", "postgresql") == "postgresql"
|
||||
):
|
||||
bm25_tokens = await select_selective_bm25_tokens(
|
||||
conn,
|
||||
tokens,
|
||||
schema=get_current_schema(),
|
||||
table="memory_units",
|
||||
language=config.text_search_extension_native_language,
|
||||
max_terms=max_query_terms,
|
||||
)
|
||||
bm25_text_param: str = dialect.prepare_bm25_text(
|
||||
tokens,
|
||||
bm25_tokens,
|
||||
query_text,
|
||||
text_search_extension=text_ext,
|
||||
max_query_terms=getattr(config, "bm25_max_query_terms", DEFAULT_BM25_MAX_QUERY_TERMS),
|
||||
max_query_terms=max_query_terms,
|
||||
)
|
||||
for i, ft in enumerate(fact_types):
|
||||
arms.append(
|
||||
@@ -314,7 +300,7 @@ async def retrieve_semantic_bm25_combined_sql(
|
||||
text_search_extension=text_ext,
|
||||
bm25_language=config.text_search_extension_native_language,
|
||||
bm25_min_score=bm25_min,
|
||||
extra_where=created_range_clause,
|
||||
extra_where=updated_range_clause,
|
||||
)
|
||||
)
|
||||
|
||||
@@ -327,7 +313,7 @@ async def retrieve_semantic_bm25_combined_sql(
|
||||
if tags:
|
||||
params.append(tags)
|
||||
params.extend(groups_params)
|
||||
params.extend(created_range_params)
|
||||
params.extend(updated_range_params)
|
||||
|
||||
try:
|
||||
rows = await conn.fetch(query, *params)
|
||||
@@ -346,12 +332,12 @@ async def retrieve_semantic_bm25_combined_sql(
|
||||
fb_groups_start = fb_tags_idx + (1 if tags else 0)
|
||||
fb_groups_clause, _, _ = build_tag_groups_where_clause(tag_groups, fb_groups_start)
|
||||
fb_next_idx = fb_groups_start + len(groups_params)
|
||||
fb_created_clause = ""
|
||||
fb_updated_clause = ""
|
||||
if created_after is not None:
|
||||
fb_created_clause += f" AND updated_at > ${fb_next_idx}"
|
||||
fb_updated_clause += f" AND updated_at > ${fb_next_idx}"
|
||||
fb_next_idx += 1
|
||||
if created_before is not None:
|
||||
fb_created_clause += f" AND updated_at < ${fb_next_idx}"
|
||||
fb_updated_clause += f" AND updated_at < ${fb_next_idx}"
|
||||
fb_next_idx += 1
|
||||
fb_arms = [
|
||||
dialect.build_semantic_arm(
|
||||
@@ -364,7 +350,7 @@ async def retrieve_semantic_bm25_combined_sql(
|
||||
min_similarity=sem_min,
|
||||
tags_clause=fb_tags_clause,
|
||||
groups_clause=fb_groups_clause,
|
||||
extra_where=fb_created_clause,
|
||||
extra_where=fb_updated_clause,
|
||||
)
|
||||
for ft in fact_types
|
||||
]
|
||||
@@ -373,7 +359,7 @@ async def retrieve_semantic_bm25_combined_sql(
|
||||
if tags:
|
||||
fb_params.append(tags)
|
||||
fb_params.extend(groups_params)
|
||||
fb_params.extend(created_range_params)
|
||||
fb_params.extend(updated_range_params)
|
||||
rows = await conn.fetch(fb_query, *fb_params)
|
||||
else:
|
||||
raise
|
||||
@@ -474,45 +460,6 @@ def _select_with_temporal_coverage(
|
||||
return selected
|
||||
|
||||
|
||||
async def retrieve_temporal_combined(
|
||||
conn,
|
||||
query_emb_str: str,
|
||||
bank_id: str,
|
||||
fact_types: list[str],
|
||||
start_date: datetime,
|
||||
end_date: datetime,
|
||||
budget: int,
|
||||
semantic_threshold: float = DEFAULT_TEMPORAL_SEMANTIC_MIN_SIMILARITY,
|
||||
tags: list[str] | None = None,
|
||||
tags_match: TagsMatch = "any",
|
||||
tag_groups: list[TagGroup] | None = None,
|
||||
created_after: datetime | None = None,
|
||||
created_before: datetime | None = None,
|
||||
) -> dict[str, list[RetrievalResult]]:
|
||||
"""Temporal retrieval, run by the configured memories store.
|
||||
|
||||
The timestamps live with the memories, so whoever holds them runs the arm.
|
||||
With the default Postgres store this is :func:`retrieve_temporal_combined_sql`.
|
||||
"""
|
||||
from ..memories import get_memories
|
||||
|
||||
return await get_memories().temporal_search(
|
||||
conn=conn,
|
||||
bank_id=bank_id,
|
||||
fact_types=fact_types,
|
||||
query_embedding=query_emb_str,
|
||||
start_date=start_date,
|
||||
end_date=end_date,
|
||||
limit=budget,
|
||||
semantic_threshold=semantic_threshold,
|
||||
tags=tags,
|
||||
tags_match=tags_match,
|
||||
tag_groups=tag_groups,
|
||||
created_after=created_after,
|
||||
created_before=created_before,
|
||||
)
|
||||
|
||||
|
||||
async def retrieve_temporal_combined_sql(
|
||||
conn,
|
||||
query_emb_str: str,
|
||||
@@ -559,29 +506,30 @@ async def retrieve_temporal_combined_sql(
|
||||
# Entry-point query: fixed params are $1-$5 (emb, bank, start, end, threshold), tags at $6.
|
||||
# fact_type is inlined as a literal per UNION ALL arm (not a bind) — this avoids `unnest`,
|
||||
# which has no Oracle equivalent (the `<=>` operator and LIMIT are translated to Oracle by
|
||||
# the backend on execute, but `unnest` is not). Mirrors retrieve_semantic_bm25_combined.
|
||||
# the backend on execute, but `unnest` is not). Mirrors retrieve_semantic_bm25_combined_sql.
|
||||
tags_clause = build_tags_where_clause_simple(tags, 6, match=tags_match)
|
||||
tag_groups_param_start = 6 + (1 if tags else 0)
|
||||
groups_clause, groups_params, _ = build_tag_groups_where_clause(tag_groups, tag_groups_param_start)
|
||||
|
||||
# created_at time range filter (after tags/groups)
|
||||
# created_after/created_before time range filter (after tags/groups) — filters
|
||||
# `updated_at`, as above.
|
||||
_next_idx = tag_groups_param_start + len(groups_params)
|
||||
created_range_clause = ""
|
||||
created_range_params: list[Any] = []
|
||||
updated_range_clause = ""
|
||||
updated_range_params: list[Any] = []
|
||||
if created_after is not None:
|
||||
created_range_params.append(created_after)
|
||||
created_range_clause += f" AND updated_at > ${_next_idx}"
|
||||
updated_range_params.append(created_after)
|
||||
updated_range_clause += f" AND updated_at > ${_next_idx}"
|
||||
_next_idx += 1
|
||||
if created_before is not None:
|
||||
created_range_params.append(created_before)
|
||||
created_range_clause += f" AND updated_at < ${_next_idx}"
|
||||
updated_range_params.append(created_before)
|
||||
updated_range_clause += f" AND updated_at < ${_next_idx}"
|
||||
_next_idx += 1
|
||||
|
||||
params: list = [query_emb_str, bank_id, start_date, end_date, semantic_threshold]
|
||||
if tags:
|
||||
params.append(tags)
|
||||
params.extend(groups_params)
|
||||
params.extend(created_range_params)
|
||||
params.extend(updated_range_params)
|
||||
|
||||
# Entry-point selection: similarity-gated, window-filtered, then narrowed for coverage.
|
||||
#
|
||||
@@ -608,7 +556,7 @@ async def retrieve_temporal_combined_sql(
|
||||
# One similarity-ranked, window-filtered arm per fact_type, UNION ALL'd — each arm has its
|
||||
# own ORDER BY ... LIMIT so the per-(bank, fact_type) vector index can serve it. fact_type
|
||||
# is inlined as a literal (controlled internal enum, never user input), matching
|
||||
# retrieve_semantic_bm25_combined; this keeps the query free of `unnest`/LATERAL, which the
|
||||
# retrieve_semantic_bm25_combined_sql; this keeps the query free of `unnest`/LATERAL, which the
|
||||
# Oracle backend cannot translate.
|
||||
pool_cols = (
|
||||
"id, text, context, event_date, occurred_start, occurred_end, mentioned_at, "
|
||||
@@ -635,7 +583,7 @@ async def retrieve_temporal_combined_sql(
|
||||
AND (1 - (embedding <=> $1::vector)) >= $5
|
||||
{tags_clause}
|
||||
{groups_clause}
|
||||
{created_range_clause}
|
||||
{updated_range_clause}
|
||||
ORDER BY embedding <=> $1::vector
|
||||
LIMIT {_TEMPORAL_POOL_SIZE}
|
||||
)"""
|
||||
@@ -851,7 +799,6 @@ async def retrieve_all_fact_types_parallel(
|
||||
thinking_budget: int,
|
||||
question_date: datetime | None = None,
|
||||
query_analyzer: Optional["QueryAnalyzer"] = None,
|
||||
graph_retriever: GraphRetriever | None = None,
|
||||
tags: list[str] | None = None,
|
||||
tags_match: TagsMatch = "any",
|
||||
tag_groups: list[TagGroup] | None = None,
|
||||
@@ -863,15 +810,16 @@ async def retrieve_all_fact_types_parallel(
|
||||
enable_graph_retrieval: bool = True,
|
||||
) -> MultiFactTypeRetrievalResult:
|
||||
"""
|
||||
Optimized retrieval for multiple fact types using batched queries.
|
||||
Retrieve every recall arm for all fact types, through the memories store.
|
||||
|
||||
This reduces database round-trips by:
|
||||
1. Combining semantic + BM25 into one CTE query for ALL fact types (1 query instead of 2N)
|
||||
2. Running graph retrieval per fact type in parallel (N parallel tasks)
|
||||
3. Running temporal retrieval per fact type in parallel (N parallel tasks)
|
||||
Extracts the temporal constraint (CPU-only), then hands the whole recall off to the
|
||||
store's single ``recall_unified`` method — the one recall interface. How the arms are
|
||||
run (a per-arm SQL orchestration for Postgres, a single index query for a store that
|
||||
owns its index) is the store's business; this only assembles the per-arm result it
|
||||
returns into :class:`MultiFactTypeRetrievalResult`. Fusion/rerank happen downstream.
|
||||
|
||||
Args:
|
||||
pool: Database connection pool
|
||||
pool: Database connection pool, handed to the store as its connection handle.
|
||||
query_text: Query text
|
||||
query_embedding_str: Query embedding as string
|
||||
bank_id: Bank ID
|
||||
@@ -879,7 +827,6 @@ async def retrieve_all_fact_types_parallel(
|
||||
thinking_budget: Budget for graph traversal and retrieval limits
|
||||
question_date: Optional date when question was asked (for temporal filtering)
|
||||
query_analyzer: Query analyzer to use (defaults to TransformerQueryAnalyzer)
|
||||
graph_retriever: Graph retrieval strategy (defaults to configured retriever)
|
||||
enable_temporal_retrieval: Run the temporal arm. False also skips the date-aware
|
||||
query analysis that feeds it (no constraint means nothing to filter on).
|
||||
enable_graph_retrieval: Run the entity/link graph arm. False skips those queries
|
||||
@@ -890,155 +837,72 @@ async def retrieve_all_fact_types_parallel(
|
||||
"""
|
||||
import time
|
||||
|
||||
# Resolving the retriever can lazily construct one, so skip it when the arm is off.
|
||||
retriever = (graph_retriever or get_default_graph_retriever()) if enable_graph_retrieval else None
|
||||
config = get_config()
|
||||
start_time = time.time()
|
||||
timings: dict[str, float] = {}
|
||||
|
||||
# Step 1: Extract temporal constraint first (CPU work, no DB)
|
||||
# Do this before DB queries so we know if we need temporal retrieval
|
||||
# Do this before the store call so we know whether the temporal arm is needed at all.
|
||||
temporal_extraction_start = time.time()
|
||||
temporal_constraint = None
|
||||
if enable_temporal_retrieval:
|
||||
from .temporal_extraction import extract_temporal_constraint
|
||||
from .temporal_extraction import extract_temporal_constraint_async
|
||||
|
||||
temporal_constraint = extract_temporal_constraint(
|
||||
# Off the event loop: this is pure CPU and would otherwise stall every
|
||||
# other in-flight request in the process, not just this recall.
|
||||
temporal_constraint = await extract_temporal_constraint_async(
|
||||
query_text, reference_date=question_date, analyzer=query_analyzer
|
||||
)
|
||||
temporal_extraction_time = time.time() - temporal_extraction_start
|
||||
timings["temporal_extraction"] = temporal_extraction_time
|
||||
|
||||
# Step 2: Run semantic + BM25 + temporal combined in ONE connection!
|
||||
# This reduces connection usage from 2 to 1 for these operations
|
||||
semantic_bm25_start = time.time()
|
||||
temporal_results_by_ft: dict[str, list[RetrievalResult]] = {}
|
||||
temporal_time = 0.0
|
||||
# Step 2: Run every arm for every fact type through the store's single recall method.
|
||||
from ..memories import RecallArms, get_memories
|
||||
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
conn_wait = time.time() - semantic_bm25_start
|
||||
unified = await get_memories().recall_unified(
|
||||
conn=pool,
|
||||
bank_id=bank_id,
|
||||
fact_types=fact_types,
|
||||
query_embedding=query_embedding_str,
|
||||
query_text=query_text,
|
||||
limit=thinking_budget,
|
||||
temporal_window=temporal_constraint,
|
||||
temporal_semantic_threshold=config.temporal_semantic_min_similarity,
|
||||
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,
|
||||
enable_graph=enable_graph_retrieval,
|
||||
)
|
||||
|
||||
# Semantic + BM25 combined
|
||||
semantic_bm25_results = await retrieve_semantic_bm25_combined(
|
||||
conn,
|
||||
query_embedding_str,
|
||||
query_text,
|
||||
bank_id,
|
||||
fact_types,
|
||||
thinking_budget,
|
||||
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=config.graph_seed_min_similarity,
|
||||
)
|
||||
semantic_bm25_time = time.time() - semantic_bm25_start
|
||||
|
||||
# Temporal combined (if constraint detected) - same connection!
|
||||
if temporal_constraint:
|
||||
tc_start, tc_end = temporal_constraint
|
||||
temporal_start = time.time()
|
||||
temporal_results_by_ft = await retrieve_temporal_combined(
|
||||
conn,
|
||||
query_embedding_str,
|
||||
bank_id,
|
||||
fact_types,
|
||||
tc_start,
|
||||
tc_end,
|
||||
budget=thinking_budget,
|
||||
semantic_threshold=config.temporal_semantic_min_similarity,
|
||||
tags=tags,
|
||||
tags_match=tags_match,
|
||||
tag_groups=tag_groups,
|
||||
created_after=created_after,
|
||||
created_before=created_before,
|
||||
)
|
||||
temporal_time = time.time() - temporal_start
|
||||
|
||||
timings["semantic_bm25_combined"] = semantic_bm25_time
|
||||
timings["temporal_combined"] = temporal_time
|
||||
|
||||
# Step 3: Run graph retrieval for each fact type in parallel
|
||||
async def run_graph_for_fact_type(
|
||||
ft: str,
|
||||
) -> tuple[str, list[RetrievalResult], float, GraphRetrievalTimings | None]:
|
||||
graph_start = time.time()
|
||||
assert retriever is not None # only scheduled when enable_graph_retrieval is True
|
||||
results, graph_timing = await retriever.retrieve(
|
||||
pool=pool,
|
||||
query_embedding_str=query_embedding_str,
|
||||
bank_id=bank_id,
|
||||
fact_type=ft,
|
||||
budget=thinking_budget,
|
||||
query_text=query_text,
|
||||
tags=tags,
|
||||
tags_match=tags_match,
|
||||
tag_groups=tag_groups,
|
||||
created_after=created_after,
|
||||
created_before=created_before,
|
||||
preselected_semantic_seeds=semantic_bm25_results[ft].graph_seeds,
|
||||
)
|
||||
return ft, results, time.time() - graph_start, graph_timing
|
||||
|
||||
# Run graph for all fact types in parallel (skipped entirely when the arm is off)
|
||||
graph_results_list: list[tuple[str, list[RetrievalResult], float, GraphRetrievalTimings | None]] = []
|
||||
if enable_graph_retrieval:
|
||||
graph_tasks = [run_graph_for_fact_type(ft) for ft in fact_types]
|
||||
graph_results_list = await asyncio.gather(*graph_tasks)
|
||||
|
||||
# Organize results by fact type
|
||||
results_by_fact_type: dict[str, ParallelRetrievalResult] = {}
|
||||
max_conn_wait = conn_wait # Single connection for semantic+bm25+temporal
|
||||
all_graph_timings: list[GraphRetrievalTimings] = []
|
||||
|
||||
for ft in fact_types:
|
||||
# Get semantic + bm25 results for this fact type
|
||||
semantic_results = semantic_bm25_results[ft].semantic
|
||||
bm25_results = semantic_bm25_results[ft].bm25
|
||||
|
||||
# Find graph results for this fact type
|
||||
graph_results = []
|
||||
graph_time = 0.0
|
||||
graph_timing = None
|
||||
for gr in graph_results_list:
|
||||
if gr[0] == ft:
|
||||
graph_results = gr[1]
|
||||
graph_time = gr[2]
|
||||
graph_timing = gr[3]
|
||||
if graph_timing:
|
||||
all_graph_timings.append(graph_timing)
|
||||
break
|
||||
|
||||
# Get temporal results for this fact type from combined result
|
||||
temporal_results = temporal_results_by_ft.get(ft) if temporal_constraint else None
|
||||
if temporal_results is not None and len(temporal_results) == 0:
|
||||
temporal_results = None
|
||||
|
||||
arms = unified.get(ft) or RecallArms()
|
||||
# An empty temporal list collapses to None — the "no temporal arm" signal downstream.
|
||||
temporal_arm = arms.temporal or None
|
||||
results_by_fact_type[ft] = ParallelRetrievalResult(
|
||||
semantic=semantic_results,
|
||||
bm25=bm25_results,
|
||||
graph=graph_results,
|
||||
temporal=temporal_results,
|
||||
semantic=arms.semantic,
|
||||
bm25=arms.bm25,
|
||||
graph=arms.graph,
|
||||
temporal=temporal_arm,
|
||||
timings={
|
||||
"semantic": semantic_bm25_time / 2, # Approximate split
|
||||
"bm25": semantic_bm25_time / 2,
|
||||
"graph": graph_time,
|
||||
"temporal": temporal_time, # Same for all fact types (single query)
|
||||
"semantic": 0.0,
|
||||
"bm25": 0.0,
|
||||
"graph": 0.0,
|
||||
"temporal": 0.0,
|
||||
"temporal_extraction": temporal_extraction_time,
|
||||
},
|
||||
temporal_constraint=temporal_constraint,
|
||||
graph_timings=[graph_timing] if graph_timing else [],
|
||||
max_conn_wait=max_conn_wait,
|
||||
graph_timings=[],
|
||||
max_conn_wait=0.0,
|
||||
)
|
||||
|
||||
total_time = time.time() - start_time
|
||||
timings["total"] = total_time
|
||||
|
||||
timings["total"] = time.time() - start_time
|
||||
return MultiFactTypeRetrievalResult(
|
||||
results_by_fact_type=results_by_fact_type,
|
||||
timings=timings,
|
||||
max_conn_wait=max_conn_wait,
|
||||
max_conn_wait=0.0,
|
||||
)
|
||||
|
||||
@@ -4,13 +4,58 @@ Temporal extraction for time-aware search queries.
|
||||
Handles natural language temporal expressions using transformer-based query analysis.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import atexit
|
||||
import logging
|
||||
import threading
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from datetime import datetime
|
||||
|
||||
from hindsight_api.engine.query_analyzer import DateparserQueryAnalyzer, QueryAnalyzer
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Temporal extraction is pure CPU, and recall calls it from an async request
|
||||
# path. Running it inline blocks the event loop for the whole duration, which
|
||||
# stalls every other in-flight request in the process — not just the recall
|
||||
# doing the work. Measured with 16 concurrent extractions over document-sized
|
||||
# text: the loop got a single scheduler tick in 1.3 seconds.
|
||||
#
|
||||
# It is offloaded to a thread instead. The pool is deliberately **one worker**:
|
||||
# the work is pure-Python and holds the GIL, so widening it does not add
|
||||
# parallelism, it just makes threads fight over the GIL. Measured, same 16
|
||||
# extractions:
|
||||
#
|
||||
# inline total= 1318ms loop stall max=1318ms
|
||||
# max_workers=1 total= 1438ms loop stall max= 2.8ms
|
||||
# max_workers=2 total= 2091ms loop stall max= 4.5ms
|
||||
# max_workers=4 total= 4751ms loop stall max= 6.3ms
|
||||
# unbounded total=16688ms loop stall max= 33.8ms
|
||||
#
|
||||
# One worker keeps throughput (+9%) while the loop stays responsive (470x), and
|
||||
# preserves the serialisation the inline version already had. Anything wider
|
||||
# trades throughput away for nothing.
|
||||
_executor: ThreadPoolExecutor | None = None
|
||||
_executor_lock = threading.Lock()
|
||||
|
||||
|
||||
def _get_executor() -> ThreadPoolExecutor:
|
||||
global _executor
|
||||
if _executor is None:
|
||||
with _executor_lock:
|
||||
if _executor is None:
|
||||
_executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="temporal-extract")
|
||||
atexit.register(_shutdown_executor)
|
||||
return _executor
|
||||
|
||||
|
||||
def _shutdown_executor() -> None:
|
||||
global _executor
|
||||
executor, _executor = _executor, None
|
||||
if executor is not None:
|
||||
executor.shutdown(wait=False)
|
||||
|
||||
|
||||
# Global default analyzer instance
|
||||
# Can be overridden by passing a custom analyzer to extract_temporal_constraint
|
||||
_default_analyzer: QueryAnalyzer | None = None
|
||||
@@ -52,10 +97,48 @@ def extract_temporal_constraint(
|
||||
if analyzer is None:
|
||||
analyzer = get_default_analyzer()
|
||||
|
||||
analysis = analyzer.analyze(query, reference_date)
|
||||
# Recall must never fail because temporal analysis choked on the query text.
|
||||
# Consolidation recalls with stored fact text as the query, so a single
|
||||
# pathological phrase (e.g. "十万年前" → year -97974) would otherwise fail
|
||||
# every recall touching that bank, deterministically (issue #3217). Degrade
|
||||
# to "no temporal signal" here — the one entry point the recall path uses —
|
||||
# while analyze() itself stays strict so parser bugs still surface in tests
|
||||
# and to direct callers.
|
||||
try:
|
||||
analysis = analyzer.analyze(query, reference_date)
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
"Temporal query analysis raised %s (treating as no temporal constraint): %s",
|
||||
type(e).__name__,
|
||||
e,
|
||||
)
|
||||
return None
|
||||
|
||||
if analysis.temporal_constraint:
|
||||
result = (analysis.temporal_constraint.start_date, analysis.temporal_constraint.end_date)
|
||||
return result
|
||||
|
||||
return None
|
||||
|
||||
|
||||
async def extract_temporal_constraint_async(
|
||||
query: str,
|
||||
reference_date: datetime | None = None,
|
||||
analyzer: QueryAnalyzer | None = None,
|
||||
) -> tuple[datetime, datetime] | None:
|
||||
"""Async form of :func:`extract_temporal_constraint`, off the event loop.
|
||||
|
||||
Same result as the sync function — this only changes *where* the CPU work
|
||||
runs. Use this from request paths; the sync form remains for callers that
|
||||
are not already async.
|
||||
|
||||
Safe to run off-thread as of the detector rewrite: the analyzer owns its
|
||||
``_ExactLanguageSearch`` rather than sharing dateparser's module-level
|
||||
singleton (which caches state on itself per call), and the character-table
|
||||
cache is lock-guarded.
|
||||
"""
|
||||
loop = asyncio.get_running_loop()
|
||||
return await loop.run_in_executor(
|
||||
_get_executor(),
|
||||
lambda: extract_temporal_constraint(query, reference_date=reference_date, analyzer=analyzer),
|
||||
)
|
||||
|
||||
@@ -50,6 +50,22 @@ class RetrievalResult:
|
||||
metadata: dict[str, str] | None = None # User-provided metadata
|
||||
proof_count: int | None = None # Number of supporting memories (observations only)
|
||||
|
||||
# Entity postings the backend already resolved for this unit, if any.
|
||||
# ``None`` means "this backend does not carry entity ids on the result" (the default
|
||||
# store, which resolves them later via ``entity_map_for_units``); a list — possibly
|
||||
# empty — means the backend already resolved the unit->entity posting inline, so
|
||||
# recall can build the entity map directly instead of re-fetching the memories.
|
||||
#
|
||||
# CONTRACT: a backend that populates this for an OBSERVATION MUST include the
|
||||
# entities it inherits from its source memories, not only any it carries directly.
|
||||
# Recall builds the entity map straight from this list and does NOT resolve
|
||||
# observation-from-source inheritance itself (the default store, which leaves this
|
||||
# ``None``, resolves that inheritance inside ``entity_map_for_units`` instead). A
|
||||
# backend that owns its index and resolves the inheritance at write time — so the
|
||||
# stored record's entity ids are already the complete set — satisfies this; one that
|
||||
# only stores direct postings must leave this ``None`` for observations.
|
||||
entity_ids: list[str] | None = None
|
||||
|
||||
# Retrieval-specific scores (only one will be set depending on retrieval method)
|
||||
similarity: float | None = None # Semantic retrieval
|
||||
bm25_score: float | None = None # BM25 retrieval
|
||||
|
||||
@@ -0,0 +1,94 @@
|
||||
"""Token-budget selection for the recall ``source_facts`` map.
|
||||
|
||||
Recall returns each observation's ``source_fact_ids`` in full and resolves them
|
||||
through a ``source_facts`` map that is filled up to a token budget. Which facts
|
||||
make it into that map is what this module decides.
|
||||
|
||||
Two properties matter to callers resolving provenance (issue #3221):
|
||||
|
||||
* the budget is spent in **observation-rank order**, so truncation hits the tail
|
||||
of the result list and the top-ranked results keep their provenance;
|
||||
* one oversized fact skips itself only — it does not evict every shorter fact
|
||||
behind it — and any skip is reported so the caller can tell a budget-truncated
|
||||
map from a dangling reference.
|
||||
"""
|
||||
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SourceFactSelection:
|
||||
"""The source facts that fit the budget, plus whether anything was dropped."""
|
||||
|
||||
ids: list[str]
|
||||
"""Selected source fact IDs, in the order they were considered."""
|
||||
|
||||
truncated: bool
|
||||
"""True when at least one resolvable source fact was skipped for budget."""
|
||||
|
||||
|
||||
def select_source_facts_within_budget(
|
||||
*,
|
||||
source_ids_ordered: list[str],
|
||||
source_fact_ids_by_obs: dict[str, list[str]],
|
||||
text_by_id: dict[str, str],
|
||||
max_total_tokens: int,
|
||||
max_tokens_per_observation: int,
|
||||
count_tokens: Callable[[str], int],
|
||||
) -> SourceFactSelection:
|
||||
"""Pick the source facts that fit the requested token budget.
|
||||
|
||||
``source_ids_ordered`` and ``source_fact_ids_by_obs`` must both be in
|
||||
observation-rank order — the budget is spent front to back, so their order
|
||||
decides which results keep their provenance when the budget runs out.
|
||||
|
||||
A non-negative ``max_tokens_per_observation`` gives every observation its own
|
||||
budget (the global one is then unused); otherwise a single ``max_total_tokens``
|
||||
budget is shared, with a negative value meaning unlimited. IDs absent from
|
||||
``text_by_id`` are skipped without counting as truncation: those rows did not
|
||||
resolve at all, which is a different condition from running out of budget.
|
||||
"""
|
||||
selected: list[str] = []
|
||||
seen: set[str] = set()
|
||||
# Skipped for budget, not necessarily lost: under a per-observation cap the same
|
||||
# source can be dropped by one observation's budget and still fit another's. Only
|
||||
# a skip that leaves the fact out of the final map counts as truncation, so the
|
||||
# flag means exactly what it says — some advertised ID has no entry.
|
||||
skipped: set[str] = set()
|
||||
|
||||
def _take(sid: str) -> None:
|
||||
if sid not in seen:
|
||||
seen.add(sid)
|
||||
selected.append(sid)
|
||||
|
||||
if max_tokens_per_observation >= 0:
|
||||
# Per-observation capping: each observation independently selects source
|
||||
# facts up to its own budget.
|
||||
for sids in source_fact_ids_by_obs.values():
|
||||
obs_tokens = 0
|
||||
for sid in sids:
|
||||
text = text_by_id.get(sid)
|
||||
if text is None:
|
||||
continue
|
||||
fact_tokens = count_tokens(text)
|
||||
if obs_tokens + fact_tokens > max_tokens_per_observation:
|
||||
skipped.add(sid)
|
||||
continue
|
||||
obs_tokens += fact_tokens
|
||||
_take(sid)
|
||||
return SourceFactSelection(ids=selected, truncated=bool(skipped - seen))
|
||||
|
||||
total_tokens = 0
|
||||
for sid in source_ids_ordered:
|
||||
text = text_by_id.get(sid)
|
||||
if text is None:
|
||||
continue
|
||||
fact_tokens = count_tokens(text)
|
||||
if max_total_tokens >= 0 and total_tokens + fact_tokens > max_total_tokens:
|
||||
skipped.add(sid)
|
||||
continue
|
||||
total_tokens += fact_tokens
|
||||
_take(sid)
|
||||
|
||||
return SourceFactSelection(ids=selected, truncated=bool(skipped - seen))
|
||||
@@ -0,0 +1,155 @@
|
||||
"""Behaviour-identical, allocation-free replacement for dateparser's search-time
|
||||
language detection.
|
||||
|
||||
``dateparser.search.search_dates`` spends ~98% of its time in
|
||||
``FullTextLanguageDetector._best_language`` deciding which of 205 locales the
|
||||
text is in. Three costs in that function are pure waste, and none of them are
|
||||
fixed upstream (``search/text_detection.py`` on dateparser master is byte-identical
|
||||
to the 1.2.2 copy we ship):
|
||||
|
||||
1. ``get_unique_characters()`` does an O(locales²) set-difference sweep — 42,025
|
||||
set operations — to work out which characters are unique to each locale. The
|
||||
result depends only on the locale set, which never changes, but a fresh
|
||||
``FullTextLanguageDetector`` is constructed per call so it is recomputed every
|
||||
time. **Cached here.**
|
||||
|
||||
2. Each locale that scores zero is retried with ``strip_timezone=True``, and that
|
||||
retry calls ``pop_tz_offset_from_string`` on *the same string* — up to 199
|
||||
times per query, each one a linear scan of 773 compiled timezone regexes
|
||||
(~154k regex searches per query). The result cannot differ between locales.
|
||||
**Hoisted to one call.**
|
||||
|
||||
3. When stripping the timezone does not change the string, the retry re-runs a
|
||||
computation whose result is already known to be ``[0, 0]``. **Skipped.**
|
||||
|
||||
Every transformation here is equivalence-preserving by construction, and
|
||||
``tests/test_temporal_extraction.py`` asserts that directly: it runs this
|
||||
implementation and dateparser's side by side over the corpus plus thousands of
|
||||
randomly generated strings and requires identical answers.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import threading
|
||||
from dataclasses import dataclass
|
||||
|
||||
from dateparser.conf import Settings
|
||||
from dateparser.conf import settings as default_settings
|
||||
from dateparser.languages.locale import Locale
|
||||
from dateparser.timezone_parser import pop_tz_offset_from_string
|
||||
from dateparser.utils import normalize_unicode
|
||||
|
||||
# Mirrors FullTextLanguageDetector.character_check. A string made only of these
|
||||
# characters carries no language signal, so detection short-circuits to the
|
||||
# first candidate locale.
|
||||
_SYMBOL_SET = frozenset("0123456789 /-)(.:\\,'")
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class LocaleCharTables:
|
||||
"""Character evidence per locale, in locale order.
|
||||
|
||||
``language_chars[i]`` is every character locale *i* uses in its date
|
||||
vocabulary; ``unique_chars[i]`` is the subset no other candidate locale uses,
|
||||
which lets a single character decide the language outright.
|
||||
"""
|
||||
|
||||
language_chars: list[set[str]]
|
||||
unique_chars: list[set[str]]
|
||||
|
||||
|
||||
_char_table_lock = threading.Lock()
|
||||
# id(locale tuple) is not stable, so key on the locale shortnames.
|
||||
_char_table_cache: dict[tuple[str, ...], LocaleCharTables] = {}
|
||||
|
||||
|
||||
def _char_tables(languages: list[Locale], settings: Settings) -> LocaleCharTables:
|
||||
"""Per-locale character sets, and the characters unique to each locale.
|
||||
|
||||
This is ``FullTextLanguageDetector.get_unique_characters`` with its result
|
||||
memoised. It is a pure function of the locale set: ``get_wordchars_for_detection``
|
||||
caches on the (singleton) ``Locale``, and the O(n²) difference sweep over
|
||||
those sets is therefore deterministic.
|
||||
"""
|
||||
key = tuple(lang.shortname for lang in languages)
|
||||
cached = _char_table_cache.get(key)
|
||||
if cached is not None:
|
||||
return cached
|
||||
|
||||
detection_settings = settings.replace(NORMALIZE=False)
|
||||
language_chars = [lang.get_wordchars_for_detection(settings=detection_settings) for lang in languages]
|
||||
|
||||
unique_chars = []
|
||||
for char_set in language_chars:
|
||||
remaining = char_set
|
||||
for other in language_chars:
|
||||
if other != char_set:
|
||||
remaining = remaining - other
|
||||
unique_chars.append(remaining)
|
||||
|
||||
tables = LocaleCharTables(language_chars=language_chars, unique_chars=unique_chars)
|
||||
with _char_table_lock:
|
||||
_char_table_cache[key] = tables
|
||||
return tables
|
||||
|
||||
|
||||
def _character_check(text: str, languages: list[Locale], settings: Settings) -> list[Locale]:
|
||||
"""Narrow the candidate locales by character evidence.
|
||||
|
||||
Faithful to ``FullTextLanguageDetector.character_check``, including its
|
||||
quirks: it inspects the *original-cased* string, and a single unique
|
||||
character decides the language outright.
|
||||
"""
|
||||
lowered = text.lower()
|
||||
text_chars = set(lowered)
|
||||
if text_chars & _SYMBOL_SET == text_chars:
|
||||
return languages[:1]
|
||||
|
||||
tables = _char_tables(languages, settings)
|
||||
|
||||
for i, chars in enumerate(tables.unique_chars):
|
||||
for char in chars:
|
||||
if char.lower() in lowered:
|
||||
return [languages[i]]
|
||||
|
||||
return [lang for i, lang in enumerate(languages) if text_chars & tables.language_chars[i]]
|
||||
|
||||
|
||||
def best_language(text: str, languages: list[Locale], settings: Settings | None = None) -> str | None:
|
||||
"""Return the detected locale shortname, or None.
|
||||
|
||||
Equivalent to ``FullTextLanguageDetector(languages)._best_language(text)``.
|
||||
|
||||
Note the argument order dateparser uses and that we must preserve: the
|
||||
character check runs on the raw string, and only afterwards is the text
|
||||
lowercased and unicode-normalised for the applicability counts.
|
||||
"""
|
||||
settings = settings or default_settings
|
||||
|
||||
candidates = _character_check(text, languages, settings)
|
||||
text = normalize_unicode(text.lower())
|
||||
if len(candidates) == 1:
|
||||
return candidates[0].shortname
|
||||
|
||||
# Hoisted out of the per-locale loop: the timezone strip depends only on the
|
||||
# text, so all 199 retries would compute the same thing.
|
||||
stripped, _ = pop_tz_offset_from_string(text, as_offset=False)
|
||||
tz_changed = stripped != text
|
||||
|
||||
applicable: list[tuple[str, list[int]]] = []
|
||||
for language in candidates:
|
||||
counts = language.count_applicability(text, strip_timezone=False, settings=settings)
|
||||
if counts[0] > 0 or counts[1] > 0:
|
||||
applicable.append((language.shortname, counts))
|
||||
continue
|
||||
if not tz_changed:
|
||||
# The retry strips a timezone that is not there, so it would re-run
|
||||
# the identical computation and get the identical [0, 0].
|
||||
continue
|
||||
counts = language.count_applicability(stripped, strip_timezone=False, settings=settings)
|
||||
if counts[0] > 0 or counts[1] > 0:
|
||||
applicable.append((language.shortname, counts))
|
||||
|
||||
if not applicable:
|
||||
return None
|
||||
return max(applicable, key=lambda pair: (pair[1][0], pair[1][1]))[0]
|
||||
@@ -50,7 +50,9 @@ def _month_end(year: int, month: int) -> datetime:
|
||||
return datetime(year, month, calendar.monthrange(year, month)[1])
|
||||
|
||||
|
||||
def _extract_non_chinese_period(query: str, reference_date: datetime) -> DateRange | None:
|
||||
def _extract_non_chinese_period(
|
||||
query: str, reference_date: datetime
|
||||
) -> DateRange | NoTemporalConstraintSentinel | None:
|
||||
if re.search(r"\b(yesterday|ayer|ieri|hier|gestern|вчера)\b", query, re.IGNORECASE):
|
||||
d = reference_date - timedelta(days=1)
|
||||
return _constraint(d, d)
|
||||
@@ -162,6 +164,11 @@ def _extract_non_chinese_period(query: str, reference_date: datetime) -> DateRan
|
||||
match = re.search(rf"\b({pattern})\s+(\d{{4}})\b", query, re.IGNORECASE)
|
||||
if match:
|
||||
year = int(match.group(2))
|
||||
if year < datetime.min.year:
|
||||
# "june 0000" — an explicit but unrepresentable year. Treat it
|
||||
# as no constraint rather than crashing recall or letting the
|
||||
# dateparser fallback invent a different date (issue #3217).
|
||||
return NO_TEMPORAL_CONSTRAINT
|
||||
start = datetime(year, month_num, 1)
|
||||
return _constraint(start, _month_end(year, month_num))
|
||||
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user