Compare commits
1
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 Ben
|
82c6db52ad |
@@ -77,18 +77,6 @@ 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.
|
||||
@@ -164,28 +152,6 @@ 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:
|
||||
@@ -211,26 +177,6 @@ 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:
|
||||
@@ -247,48 +193,6 @@ 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`:
|
||||
@@ -350,12 +254,9 @@ 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
|
||||
|
||||
+6
-69
@@ -7,10 +7,7 @@ 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: 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.
|
||||
# Reasoning effort for providers/models that support it. Examples: low, medium, high, xhigh.
|
||||
# HINDSIGHT_API_LLM_REASONING_EFFORT=low
|
||||
|
||||
# Sampling temperature for internal LLM calls. Set a number in [0.0, 2.0], or `none`
|
||||
@@ -57,13 +54,6 @@ 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.
|
||||
@@ -169,7 +159,6 @@ 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.
|
||||
@@ -177,34 +166,12 @@ 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
|
||||
|
||||
# Let a vector index scan resume until the query's LIMIT is satisfied, instead of
|
||||
# stopping when its first candidate list drains (pgvector: hnsw.ef_search, 200) — with
|
||||
# it off, a larger recall budget cannot retrieve more rows. Needs pgvector 0.8.0+;
|
||||
# older servers reject it and it is dropped automatically. Set false and restart as a
|
||||
# quick revert to the previous retrieval depth, with no code change.
|
||||
# HINDSIGHT_API_ANN_ITERATIVE_SCAN=true
|
||||
# Ceiling on tuples one resumed scan may visit. Bounds the CPU and memory a selective
|
||||
# query can spend resuming (filters are applied after the scan, so it resumes often).
|
||||
# Lower it to trade depth back for latency. pgvector's own default is 20000.
|
||||
# HINDSIGHT_API_ANN_MAX_SCAN_TUPLES=4000
|
||||
# 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
|
||||
@@ -217,18 +184,11 @@ HINDSIGHT_API_LOG_LEVEL=info
|
||||
# chinese_lindera/lindera(chinese), japanese_lindera/lindera(japanese),
|
||||
# korean_lindera/lindera(korean), ngram(min,max), edge_ngram(min,max)
|
||||
# HINDSIGHT_API_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER=
|
||||
# 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
|
||||
# Optional cap on the number of terms in the native PostgreSQL BM25 tsquery.
|
||||
# Long queries OR-join every normalized token, which can match too much of a
|
||||
# large bank. 0 (default) keeps the historical uncapped behavior; a positive
|
||||
# value bounds only the native backend (other BM25 backends get the raw query).
|
||||
# HINDSIGHT_API_BM25_MAX_QUERY_TERMS=0
|
||||
|
||||
# File Parser (Optional - uses markitdown by default)
|
||||
# HINDSIGHT_API_FILE_PARSER=markitdown
|
||||
@@ -274,12 +234,6 @@ 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:
|
||||
@@ -336,10 +290,6 @@ 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
|
||||
@@ -387,19 +337,6 @@ 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: 41 KiB After Width: | Height: | Size: 40 KiB |
@@ -1016,38 +1016,6 @@
|
||||
{
|
||||
"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
|
||||
},
|
||||
{
|
||||
"date": "2026-08-19",
|
||||
"stars": 20214
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
@@ -327,50 +327,17 @@ 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/*
|
||||
path: artifacts/${{ matrix.asset_name }}
|
||||
retention-days: 1
|
||||
|
||||
release-docker-images:
|
||||
@@ -634,11 +601,6 @@ 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,7 +32,6 @@ 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 }}
|
||||
@@ -144,8 +143,6 @@ 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:
|
||||
@@ -888,37 +885,6 @@ 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: >-
|
||||
@@ -1455,30 +1421,6 @@ 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
|
||||
@@ -4644,8 +4586,6 @@ jobs:
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 45
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
@@ -4670,8 +4610,6 @@ 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
|
||||
|
||||
@@ -20,9 +20,6 @@ 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
|
||||
|
||||
@@ -2,14 +2,13 @@
|
||||
|
||||
|
||||
|
||||
[Documentation](https://hindsight.vectorize.io) • [Integrations](https://hindsight.vectorize.io/integrations) • [Cookbook](https://hindsight.vectorize.io/cookbook) • [Benchmarks](https://benchmarks.hindsight.vectorize.io/) • [Paper](https://arxiv.org/abs/2512.12818) • [Hindsight Cloud](https://ui.hindsight.vectorize.io/signup)
|
||||
[Documentation](https://hindsight.vectorize.io) • [Paper](https://arxiv.org/abs/2512.12818) • [Cookbook](https://hindsight.vectorize.io/cookbook) • [Hindsight Cloud](https://ui.hindsight.vectorize.io/signup)
|
||||
|
||||
[](https://github.com/vectorize-io/hindsight/actions/workflows/release.yml)
|
||||
[](https://pypi.org/project/hindsight-api/)
|
||||
[](https://pypi.org/project/hindsight-client/)
|
||||
[](https://www.npmjs.com/package/@vectorize-io/hindsight-client)
|
||||
[](https://github.com/vectorize-io/hindsight/actions/workflows/release.yml)
|
||||
[](https://join.slack.com/t/hindsight-space/shared_invite/zt-3nhbm4w29-LeSJ5Ixi6j8PdiYOCPlOgg)
|
||||
[](https://opensource.org/licenses/MIT)
|
||||
|
||||
|
||||
<br/>
|
||||
|
||||
<a href="https://trendshift.io/repositories/15603" target="_blank"><img src="https://trendshift.io/api/badge/repositories/15603" alt="vectorize-io%2Fhindsight | Trendshift" style="width: 250px; height: 55px;" width="250" height="55"/></a>
|
||||
@@ -21,33 +20,28 @@
|
||||
|
||||
Hindsight™ is an agent memory system built to create smarter agents that learn over time. Most agent memory systems focus on recalling conversation history. Hindsight is focused on making agents that learn, not just remember.
|
||||
|
||||
|
||||
<video src="https://github.com/user-attachments/assets/923b798d-3581-4897-bb62-9cfa5a931682" controls></video>
|
||||
|
||||
It eliminates the shortcomings of alternative techniques such as RAG and knowledge graph and delivers state-of-the-art performance on long term memory tasks.
|
||||
|
||||
**Contents**
|
||||
|
||||
- [Memory Performance & Accuracy](#memory-performance--accuracy)
|
||||
- [Quick Start](#quick-start) — [server](#1-start-a-server) · [clients](#2-connect-a-client) · [platforms](#supported-platforms) · [embedded](#python-embedded-no-server-required)
|
||||
- [Adding Hindsight to Your Agent](#adding-hindsight-to-your-agent) — [LLM Wrapper](#llm-wrapper-2-lines-of-code) · [integrations](#integrations) · [coding agents](#coding-agents) · [MCP](#mcp-server)
|
||||
- [Core Concepts](#core-concepts) — [memory types](#memory-types) · [retain / recall / reflect](#the-three-operations) · [observations](#observations) · [mental models & knowledge pages](#mental-models--knowledge-pages) · [banks](#memory-banks)
|
||||
- [Use Cases](#use-cases)
|
||||
- [Running in Production](#running-in-production)
|
||||
- [Resources](#resources)
|
||||
|
||||
---
|
||||
|
||||
## Memory Performance & Accuracy
|
||||
|
||||
Hindsight is the most accurate agent memory system ever tested according to benchmark performance. It has achieved state-of-the-art performance on the LongMemEval benchmark, widely used to assess memory system performance across a variety of conversational AI scenarios. The current reported performance of Hindsight and other agent memory solutions as of January 2026 is shown here:
|
||||
|
||||
|
||||
|
||||
> Live, continuously updated results — including per-model accuracy, latency and cost — are published at [benchmarks.hindsight.vectorize.io](https://benchmarks.hindsight.vectorize.io/).
|
||||
|
||||
The benchmark performance data for Hindsight has been independently reproduced by research collaborators at the Virginia Tech [Sanghani Center for Artificial Intelligence and Data Analytics](https://sanghani.cs.vt.edu/) and The Washington Post. Other scores are self-reported by software vendors.
|
||||
|
||||
Hindsight is being used in production at Fortune 500 enterprises and by a growing number of AI startups.
|
||||
Hindsight is being used in production at Fortune 500 enterprises and by a growing number of AI startups.
|
||||
|
||||
## Adding Hindsight to Your AI Agents
|
||||
|
||||
The easiest way to use Hindsight with an existing agent is with the LLM Wrapper. You can add memory to your agent with 2 lines of code. That will swap your current LLM client out with the Hindsight wrapper. After that, memories will be stored and retrieved automatically as you make LLM calls.
|
||||
|
||||
If you need more control over how and when your agent stores and recalls memories, there's also a simple API you can integrate with using the SDKs or directly via HTTP.
|
||||
|
||||
|
||||
|
||||
---
|
||||
|
||||
@@ -59,11 +53,10 @@ Hindsight is being used in production at Fortune 500 enterprises and by a growin
|
||||
|
||||
---
|
||||
|
||||
|
||||
## Quick Start
|
||||
|
||||
### 1. Start a server
|
||||
|
||||
#### Docker (recommended)
|
||||
### Docker (recommended)
|
||||
|
||||
```bash
|
||||
export OPENAI_API_KEY=sk-xxx
|
||||
@@ -77,52 +70,31 @@ docker run -it --pull always --name hindsight --restart unless-stopped -p 8888:8
|
||||
>API: http://localhost:8888
|
||||
>UI: http://localhost:9999
|
||||
|
||||
Hindsight works with **25+ LLM providers** via `HINDSIGHT_API_LLM_PROVIDER` — hosted (`openai`, `anthropic`, `gemini`, `groq`, `bedrock`, `vertexai`, `minimax`, `deepseek`, `atlas`, …), fully local (`ollama`, `lmstudio`, `llamacpp`), any OpenAI-compatible endpoint, and gateways (`litellm`, `litellmrouter`) that reach the rest. Existing subscriptions work too: `openai-codex` (ChatGPT Plus/Pro) and `claude-code` (Claude Pro/Max) need no API key. See [supported models](https://hindsight.vectorize.io/developer/models).
|
||||
You can modify the LLM provider by setting `HINDSIGHT_API_LLM_PROVIDER`. Valid options are `openai`, `anthropic`, `gemini`, `groq`, `ollama`, `lmstudio`, `minimax`, and `atlas` ([Atlas Cloud](https://www.atlascloud.ai/?utm_source=github&utm_medium=link&utm_campaign=hindsight)). The documentation provides more details on [supported models](https://hindsight.vectorize.io/developer/models).
|
||||
|
||||
#### Docker (external PostgreSQL)
|
||||
|
||||
|
||||
### Docker (external PostgreSQL)
|
||||
|
||||
```bash
|
||||
export OPENAI_API_KEY=sk-xxx
|
||||
export HINDSIGHT_DB_PASSWORD=choose-a-password
|
||||
cd docker/docker-compose
|
||||
docker compose up
|
||||
docker compose up
|
||||
```
|
||||
|
||||
> Oracle AI Database is also supported for enterprise deployments with full feature parity. See the [storage documentation](https://hindsight.vectorize.io/developer/storage) for details.
|
||||
|
||||
#### Bare metal (pip)
|
||||
|
||||
>API: http://localhost:8888
|
||||
>UI: http://localhost:9999
|
||||
|
||||
### Client
|
||||
|
||||
```bash
|
||||
pip install hindsight-api
|
||||
export HINDSIGHT_API_LLM_API_KEY=sk-xxx
|
||||
|
||||
hindsight-api
|
||||
```
|
||||
|
||||
#### Kubernetes (Helm)
|
||||
|
||||
```bash
|
||||
helm install hindsight oci://ghcr.io/vectorize-io/charts/hindsight \
|
||||
--set api.llm.provider=openai \
|
||||
--set api.llm.apiKey=sk-xxx \
|
||||
--set postgresql.enabled=true
|
||||
```
|
||||
|
||||
#### Managed (no server)
|
||||
|
||||
[Hindsight Cloud](https://vectorize.io/pricing) is the hosted option: managed infrastructure that scales automatically, plus a dashboard, backups, team collaboration and a 99.9% uptime SLA. Billing is usage-based with free credits to start — no fixed monthly or per-seat fee. Point any client at `https://api.hindsight.vectorize.io` with your API key and skip the deployment entirely.
|
||||
|
||||
[Compare self-hosted, Cloud and Enterprise →](https://vectorize.io/pricing) · [Sign up →](https://ui.hindsight.vectorize.io/signup)
|
||||
|
||||
All options, including Windows and air-gapped setups, are covered in the [installation guide](https://hindsight.vectorize.io/developer/installation).
|
||||
|
||||
### 2. Connect a client
|
||||
|
||||
```bash
|
||||
pip install hindsight-client -U # Python
|
||||
npm install @vectorize-io/hindsight-client # Node.js / TypeScript
|
||||
go get github.com/vectorize-io/hindsight/hindsight-clients/go # Go
|
||||
curl -fsSL https://hindsight.vectorize.io/get-cli | bash # CLI
|
||||
pip install hindsight-client -U
|
||||
# or
|
||||
npm install @vectorize-io/hindsight-client
|
||||
```
|
||||
|
||||
#### Python
|
||||
@@ -144,6 +116,10 @@ client.reflect(bank_id="my-bank", query="Tell me about Alice")
|
||||
|
||||
#### Node.js / TypeScript
|
||||
|
||||
```bash
|
||||
npm install @vectorize-io/hindsight-client
|
||||
```
|
||||
|
||||
```javascript
|
||||
const { HindsightClient } = require('@vectorize-io/hindsight-client');
|
||||
|
||||
@@ -159,18 +135,6 @@ const main = async () => {
|
||||
main();
|
||||
```
|
||||
|
||||
Full reference: [Python](https://hindsight.vectorize.io/sdks/python) · [Node.js](https://hindsight.vectorize.io/sdks/nodejs) · [Go](https://hindsight.vectorize.io/sdks/go) · [CLI](https://hindsight.vectorize.io/sdks/cli) · [REST API](https://hindsight.vectorize.io/api-reference)
|
||||
|
||||
### Supported Platforms
|
||||
|
||||
| Platform | Docker | Bare Metal (pip) | Embedded DB (pg0) |
|
||||
|----------|--------|------------------|--------------------|
|
||||
| **Linux** (x86_64, ARM64) | ✅ | ✅ | ✅ |
|
||||
| **macOS** (Apple Silicon / arm64) | ✅ | ✅ | ✅ |
|
||||
| **macOS** (Intel / x86_64) | ✅ | ⚠️ | ✅ |
|
||||
| **Windows** (x86_64) | ✅ | ✅ | ✅ |
|
||||
|
||||
⚠️ Intel Macs: use `hindsight-all-slim` — see the [installation guide](https://hindsight.vectorize.io/developer/installation#supported-platforms) for details.
|
||||
|
||||
### Python Embedded (no server required)
|
||||
|
||||
@@ -186,7 +150,7 @@ from hindsight import HindsightServer, HindsightClient
|
||||
|
||||
with HindsightServer(
|
||||
llm_provider="openai",
|
||||
llm_model="gpt-5-mini",
|
||||
llm_model="gpt-5-mini",
|
||||
llm_api_key=os.environ["OPENAI_API_KEY"]
|
||||
) as server:
|
||||
client = HindsightClient(base_url=server.url)
|
||||
@@ -194,182 +158,12 @@ with HindsightServer(
|
||||
results = client.recall(bank_id="my-bank", query="Where does Alice work?")
|
||||
```
|
||||
|
||||
A [Node.js equivalent](https://hindsight.vectorize.io/sdks/hindsight-all-npm) and a [daemon CLI](https://hindsight.vectorize.io/sdks/embed) are also available.
|
||||
|
||||
---
|
||||
|
||||
## Adding Hindsight to Your Agent
|
||||
|
||||
### LLM Wrapper (2 lines of code)
|
||||
|
||||
The easiest way to add memory to an existing agent is the LLM Wrapper. Swap your LLM client for a wrapped one — memories are then stored and retrieved automatically on every call, with no other changes to your code.
|
||||
|
||||
```bash
|
||||
pip install hindsight-litellm
|
||||
```
|
||||
|
||||
```python
|
||||
from openai import OpenAI
|
||||
from hindsight_litellm import wrap_openai
|
||||
|
||||
# Wrap your existing LLM client and you're done.
|
||||
# Defaults to Hindsight Cloud; pass hindsight_api_url for a self-hosted server.
|
||||
client = wrap_openai(
|
||||
OpenAI(),
|
||||
bank_id="user-123",
|
||||
hindsight_api_url="http://localhost:8888",
|
||||
)
|
||||
|
||||
# Hindsight recalls relevant memories before the call
|
||||
# and retains the conversation after it.
|
||||
response = client.chat.completions.create(
|
||||
model="gpt-5-mini",
|
||||
messages=[{"role": "user", "content": "What do you know about me?"}],
|
||||
)
|
||||
```
|
||||
|
||||
`wrap_anthropic()` does the same for the Anthropic SDK, and every setting — bank, recall budget, fact types, reflect instead of recall — can be overridden per call with `hindsight_*` kwargs. LiteLLM sits underneath, so the same integration covers **100+ models**. See the [LiteLLM integration](https://hindsight.vectorize.io/sdks/integrations/litellm).
|
||||
|
||||
If you need explicit control over *when* memories are stored and recalled, use the [SDKs or REST API](#2-connect-a-client) directly instead.
|
||||
|
||||
### Integrations
|
||||
|
||||
**60+ integrations** — most need no code changes.
|
||||
|
||||
| | |
|
||||
|---|---|
|
||||
| **Coding agents** | [Claude Code](https://hindsight.vectorize.io/sdks/integrations/claude-code) · [Codex](https://hindsight.vectorize.io/sdks/integrations/codex) · [Cursor](https://hindsight.vectorize.io/sdks/integrations/cursor) · [GitHub Copilot](https://hindsight.vectorize.io/sdks/integrations/github-copilot) · [opencode](https://hindsight.vectorize.io/sdks/integrations/opencode) · [Cline](https://hindsight.vectorize.io/sdks/integrations/cline) · [Aider](https://hindsight.vectorize.io/sdks/integrations/aider) · [Zed](https://hindsight.vectorize.io/sdks/integrations/zed) · [Continue](https://hindsight.vectorize.io/sdks/integrations/continue) · [Roo Code](https://hindsight.vectorize.io/sdks/integrations/roo-code) · [OpenHands](https://hindsight.vectorize.io/sdks/integrations/openhands) |
|
||||
| **Agent frameworks** | [LangGraph / LangChain](https://hindsight.vectorize.io/sdks/integrations/langgraph) · [LlamaIndex](https://hindsight.vectorize.io/sdks/integrations/llamaindex) · [CrewAI](https://hindsight.vectorize.io/sdks/integrations/crewai) · [Pydantic AI](https://hindsight.vectorize.io/sdks/integrations/pydantic-ai) · [OpenAI Agents SDK](https://hindsight.vectorize.io/sdks/integrations/openai-agents) · [Google ADK](https://hindsight.vectorize.io/sdks/integrations/google-adk) · [Agno](https://hindsight.vectorize.io/sdks/integrations/agno) · [Strands](https://hindsight.vectorize.io/sdks/integrations/strands) · [AutoGen](https://hindsight.vectorize.io/sdks/integrations/autogen) · [Microsoft Agent Framework](https://hindsight.vectorize.io/sdks/integrations/agent-framework) · [Vercel AI SDK](https://hindsight.vectorize.io/sdks/integrations/ai-sdk) · [Haystack](https://hindsight.vectorize.io/sdks/integrations/haystack) |
|
||||
| **No-code / low-code** | [n8n](https://hindsight.vectorize.io/sdks/integrations/n8n) · [Zapier](https://hindsight.vectorize.io/sdks/integrations/zapier) · [Dify](https://hindsight.vectorize.io/sdks/integrations/dify) · [Flowise](https://hindsight.vectorize.io/sdks/integrations/flowise) |
|
||||
| **Apps & tools** | [ChatGPT](https://hindsight.vectorize.io/sdks/integrations/chatgpt) · [Perplexity](https://hindsight.vectorize.io/sdks/integrations/perplexity) · [Obsidian](https://hindsight.vectorize.io/sdks/integrations/obsidian) · [Pipecat](https://hindsight.vectorize.io/sdks/integrations/pipecat) · [Vapi](https://hindsight.vectorize.io/sdks/integrations/vapi) |
|
||||
|
||||
👉 [**Browse all integrations**](https://hindsight.vectorize.io/integrations)
|
||||
|
||||
### Coding Agents
|
||||
|
||||
One package gives CLI coding agents long-term project memory: a per-repo bank built automatically from git history and past sessions, injected into the agent as it starts working, plus curated knowledge pages covering architecture, conventions and in-flight work.
|
||||
|
||||
```bash
|
||||
npx @vectorize-io/hindsight-coding-agents install all # every detected agent, wired natively
|
||||
npx @vectorize-io/hindsight-coding-agents install claude-code # or just one
|
||||
```
|
||||
|
||||
Supports Claude Code, Codex CLI, Cursor CLI, GitHub Copilot CLI, opencode, Kilo CLI, Cline CLI, Antigravity CLI, Devin CLI, Prime Agent, Grok Build and DeepSeek Harness. Ingestion is automatic — there is no setup command. See the [coding agents integration](https://hindsight.vectorize.io/sdks/integrations/coding-agents).
|
||||
|
||||
### MCP Server
|
||||
|
||||
Every server ships a built-in [Model Context Protocol](https://modelcontextprotocol.io/) endpoint, one per bank, enabled by default:
|
||||
|
||||
```
|
||||
http://localhost:8888/mcp/{bank_id}/
|
||||
```
|
||||
|
||||
Point any MCP client at it to expose retain, recall and reflect as tools. See the [MCP server docs](https://hindsight.vectorize.io/developer/mcp-server).
|
||||
|
||||
---
|
||||
|
||||
## Core Concepts
|
||||
|
||||
|
||||
|
||||
### Memory Types
|
||||
|
||||
Most agent memory implementations rely on basic vector search or sometimes use a knowledge graph. Hindsight uses biomimetic data structures to organize agent memories in a way that is more like how human memory works:
|
||||
|
||||
- **World facts:** facts about the world ("The stove gets hot")
|
||||
- **Experiences:** the agent's own experiences ("I touched the stove and it really hurt")
|
||||
- **Observations:** consolidated, evidence-backed beliefs formed from many memories
|
||||
- **Mental models:** learned understanding of the agent's world, synthesized from observations and facts
|
||||
|
||||
Memories live in **banks**. When memories are added, they are pushed into either the world facts or the experiences pathway, then represented as a combination of entities, relationships, and time series with sparse/dense vector representations to aid in later recall.
|
||||
|
||||
### The Three Operations
|
||||
|
||||
#### Retain
|
||||
|
||||
The `retain` operation is used to push new memories into Hindsight. It tells Hindsight to _retain_ the information you pass in as an input.
|
||||
|
||||
```python
|
||||
client.retain(
|
||||
bank_id="my-bank",
|
||||
content="Alice got promoted to senior engineer",
|
||||
context="career update",
|
||||
timestamp="2025-06-15T10:00:00Z",
|
||||
)
|
||||
```
|
||||
|
||||
Behind the scenes, retain uses an LLM to extract key facts, temporal data, entities, and relationships. It passes these through a normalization process to transform extracted data into canonical entities, time series, and search indexes along with metadata. These representations create the pathways for accurate memory retrieval in the recall and reflect operations.
|
||||
|
||||
|
||||
|
||||
[Retain docs →](https://hindsight.vectorize.io/developer/retain)
|
||||
|
||||
#### Recall
|
||||
|
||||
The recall operation is used to retrieve memories. These memories can come from any of the memory types (world, experiences, etc.)
|
||||
|
||||
```python
|
||||
client.recall(bank_id="my-bank", query="What does Alice do?")
|
||||
client.recall(bank_id="my-bank", query="What happened in June?") # temporal
|
||||
```
|
||||
|
||||
Recall performs 4 retrieval strategies in parallel:
|
||||
- Semantic: Vector similarity
|
||||
- Keyword: BM25 exact matching
|
||||
- Graph: Entity/temporal/causal links
|
||||
- Temporal: Time range filtering
|
||||
|
||||
|
||||
|
||||
The individual results are merged, ordered by relevance using reciprocal rank fusion and a cross-encoder reranking model, then trimmed as needed to fit within the token limit.
|
||||
|
||||
[Recall docs →](https://hindsight.vectorize.io/developer/retrieval)
|
||||
|
||||
#### Reflect
|
||||
|
||||
The reflect operation performs a more thorough analysis of existing memories. This allows the agent to form new connections between memories and build a more thorough understanding of its world — or to answer a question that needs deep thinking rather than lookup.
|
||||
|
||||
```python
|
||||
client.reflect(bank_id="my-bank", query="What should I know about Alice?")
|
||||
```
|
||||
|
||||
For example, reflect supports use cases such as:
|
||||
|
||||
- An **AI Project Manager** reflecting on what risks need to be mitigated on a project.
|
||||
- A **Sales Agent** reflecting on why certain outreach messages have gotten responses while others haven't.
|
||||
- A **Support Agent** reflecting on opportunities where customers have questions not answered by current product documentation.
|
||||
|
||||
|
||||
|
||||
[Reflect docs →](https://hindsight.vectorize.io/developer/reflect)
|
||||
|
||||
### Observations
|
||||
|
||||
Retained facts don't stay a flat pile. In the background, Hindsight consolidates related facts into **observations** — deduplicated beliefs the bank has built up over time. Each observation keeps its supporting evidence with exact quotes and a proof count, and is *refined* rather than overwritten when new evidence arrives, so new information strengthens, weakens or extends an existing belief instead of silently replacing it.
|
||||
|
||||
[Observations docs →](https://hindsight.vectorize.io/developer/observations)
|
||||
|
||||
### Mental Models & Knowledge Pages
|
||||
|
||||
A **mental model** is a standing answer to a question about a bank ("What are this user's preferences?"). You define the question once; Hindsight writes the answer, stores it, and rewrites it in the background as the bank learns more. Reading one is a database read — no retrieval, no LLM call — so an agent can boot with a page of settled knowledge instead of rediscovering it every session.
|
||||
|
||||
**Knowledge pages** are mental models with the mechanics hidden: living documents a bank writes about itself, organized in folders like a wiki, searchable, and projectable onto disk as ordinary markdown. Supply a name and a question; every other decision is a default you can override.
|
||||
|
||||
[Mental models →](https://hindsight.vectorize.io/developer/mental-models) · [Knowledge pages →](https://hindsight.vectorize.io/developer/knowledge-pages)
|
||||
|
||||
### Memory Banks
|
||||
|
||||
A **bank** is an isolated memory store — one "brain" for one user, agent, or project. Isolation is strict: no cross-bank leakage. Banks carry background context and **disposition traits** (skepticism, literalism, empathy) that shape how reflect reasons over their memories, and can be created from declarative [bank templates](https://hindsight.vectorize.io/developer/api/bank-templates).
|
||||
|
||||
Two more things worth knowing:
|
||||
|
||||
- **Multilingual by default.** Input language is detected and preserved end to end — facts stay in their original language and entities keep their native script (张伟 stays 张伟, not "Zhang Wei"). [Docs →](https://hindsight.vectorize.io/developer/multilingual)
|
||||
- **Memory Defense.** An opt-in, per-bank policy that scans every retain for secrets and PII against 45 patterns and either redacts the match (`[REDACTED:github_token]`) or blocks the item before it reaches storage. [Docs →](https://hindsight.vectorize.io/developer/memory-defense)
|
||||
|
||||
---
|
||||
|
||||
## Use Cases
|
||||
|
||||
|
||||
Hindsight is built to support conversational AI agents as well as agents that are intended to perform tasks autonomously. The ideal use case for Hindsight are agents that require a blend of these features such as AI employees that need to handle open-ended tasks, change behavior based on user feedback, and learn to perform complex tasks to automate work at a level that approximates a human work. Hindsight can be used with simple AI workflows like those built with n8n and other similar tools, but may be overkill for such applications.
|
||||
|
||||
### Per-User Memories and Chat History
|
||||
@@ -382,46 +176,141 @@ The requirements for this use case usually look something like this:
|
||||
|
||||
<video src="https://github.com/user-attachments/assets/4805e8e1-e7d1-47c6-a4f8-2344a5ec8906" controls></video>
|
||||
|
||||
Satisfying these requirements in Hindsight is straightforward. When new user inputs and tool calls are ingested into Hindsight using the retain operation, custom metadata can be used to enrich the new memories. Metadata provides a convenient way to isolate memories that need to be restricted to a given user. Once these are fed into the retain operation, any raw memories and mental models that get created can be filtered when retrieving relevant memories.
|
||||
Satisfying these requirements in Hindsight is straightforward. When new user inputs and tool calls are ingested into Hindsight using the retain operation, custom metadata can be used to enrich the new memories. Metadata provides a convenient way to isolate memories that need to be restricted to a given user. Once these are fed into the retain operation, any raw memories and mental models that get created can be filtered when retrieving relevant memories.
|
||||
|
||||
|
||||
|
||||
More patterns in the [Cookbook](https://hindsight.vectorize.io/cookbook) and [Best Practices](https://hindsight.vectorize.io/best-practices).
|
||||
|
||||
---
|
||||
|
||||
## Running in Production
|
||||
## Architecture & Operations
|
||||
|
||||
| | |
|
||||
|---|---|
|
||||
| **Storage** | PostgreSQL + pgvector, or Oracle AI Database 23ai with full feature parity — [storage](https://hindsight.vectorize.io/developer/storage) |
|
||||
| **Configuration** | Hierarchical: global env vars → per-tenant → per-bank — [configuration](https://hindsight.vectorize.io/developer/configuration) |
|
||||
| **Monitoring** | Prometheus metrics and dashboards for LLM calls, tokens and latency — [monitoring](https://hindsight.vectorize.io/developer/monitoring) |
|
||||
| **Operations** | Admin CLI for migrations, bank repair and stuck operations — [admin CLI](https://hindsight.vectorize.io/developer/admin-cli) |
|
||||
| **Events** | Webhooks for retain, consolidation and refresh lifecycle events — [webhooks](https://hindsight.vectorize.io/developer/api/webhooks) |
|
||||
| **Extensibility** | Tenant, auth and storage extension points — [extensions](https://hindsight.vectorize.io/developer/extensions) |
|
||||
| **Managed** | Skip all of it with [Hindsight Cloud](https://vectorize.io/pricing) — managed, usage-based, 99.9% uptime SLA |
|
||||
|
||||
|
||||
Most agent memory implementations rely on basic vector search or sometimes use a knowledge graph. Hindsight uses biomimetic data structures to organize agent memories in a way that is more like how human memory works:
|
||||
|
||||
- **World:** Facts about the world ("The stove gets hot")
|
||||
- **Experiences:** Agent's own experiences ("I touched the stove and it really hurt")
|
||||
- **Mental Models:** Learned understanding of the agent's world formed by reflecting on raw memories and experiences.
|
||||
|
||||
Memories in Hindsight are stored in banks (i.e. memory banks). When memories are added to Hindsight, they are pushed into either the world facts or experiences memory pathway. They are then represented as a combination of entities, relationships, and time series with sparse/dense vector representations to aid in later recall.
|
||||
|
||||
Hindsight provides three simple methods to interact with the system:
|
||||
|
||||
- **Retain:** Provide information to Hindsight that you want it to remember
|
||||
- **Recall:** Retrieve memories from Hindsight
|
||||
- **Reflect:** Reflect on memories and experiences to generate new observations and insights from existing memories.
|
||||
|
||||
### Retain
|
||||
|
||||
The `retain` operation is used to push new memories into Hindsight. It tells Hindsight to _retain_ the information you pass in as an input.
|
||||
|
||||
```python
|
||||
from hindsight_client import Hindsight
|
||||
|
||||
client = Hindsight(base_url="http://localhost:8888")
|
||||
|
||||
# Simple
|
||||
client.retain(
|
||||
bank_id="my-bank",
|
||||
content="Alice works at Google as a software engineer"
|
||||
)
|
||||
|
||||
# With context and timestamp
|
||||
client.retain(
|
||||
bank_id="my-bank",
|
||||
content="Alice got promoted to senior engineer",
|
||||
context="career update",
|
||||
timestamp="2025-06-15T10:00:00Z"
|
||||
)
|
||||
```
|
||||
|
||||
Behind the scenes, the retain operation uses an LLM to extract key facts, temporal data, entities, and relationships. It passes these through a normalization process to transform extracted data into canonical entities, time series, and search indexes along with metadata. These representations create the pathways for accurate memory retrieval in the recall and reflect operations.
|
||||
|
||||
|
||||
|
||||
### Recall
|
||||
|
||||
The recall operation is used to retrieve memories. These memories can come from any of the memory types (world, experiences, etc.)
|
||||
|
||||
```python
|
||||
from hindsight_client import Hindsight
|
||||
|
||||
client = Hindsight(base_url="http://localhost:8888")
|
||||
|
||||
# Simple
|
||||
client.recall(bank_id="my-bank", query="What does Alice do?")
|
||||
|
||||
# Temporal
|
||||
client.recall(bank_id="my-bank", query="What happened in June?")
|
||||
```
|
||||
|
||||
Recall performs 4 retrieval strategies in parallel:
|
||||
- Semantic: Vector similarity
|
||||
- Keyword: BM25 exact matching
|
||||
- Graph: Entity/temporal/causal links
|
||||
- Temporal: Time range filtering
|
||||
|
||||
|
||||
|
||||
The individual results from the retrievals are merged, then ordered by relevance using reciprocal rank fusion and a cross-encoder reranking model.
|
||||
|
||||
The final output is trimmed as needed to fit within the token limit.
|
||||
|
||||
### Reflect
|
||||
|
||||
The reflect operation is used to perform a more thorough analysis of existing memories. This allows the agent to form new connections between memories and build a more thorough understanding of its world.
|
||||
|
||||
For example, the `reflect` operation can be used to support use cases such as:
|
||||
|
||||
- An **AI Project Manager** reflecting on what risks need to be mitigated on a project.
|
||||
- A **Sales Agent** reflecting on why certain outreach messages have gotten responses while others haven't.
|
||||
- A **Support Agent** reflecting on opportunities where customers have questions not answered by current product documentation.
|
||||
|
||||
The `reflect` operation can also be used to handle on-demand question answering or analysis which require more deep thinking.
|
||||
|
||||
```python
|
||||
from hindsight_client import Hindsight
|
||||
|
||||
client = Hindsight(base_url="http://localhost:8888")
|
||||
|
||||
client.reflect(bank_id="my-bank", query="What should I know about Alice?")
|
||||
```
|
||||
|
||||
|
||||
|
||||
---
|
||||
|
||||
## Resources
|
||||
|
||||
**Documentation:**
|
||||
- [Docs](https://hindsight.vectorize.io) · [FAQ](https://hindsight.vectorize.io/faq) · [Best Practices](https://hindsight.vectorize.io/best-practices) · [Cookbook](https://hindsight.vectorize.io/cookbook) · [Blog](https://hindsight.vectorize.io/blog)
|
||||
- [Paper](https://arxiv.org/abs/2512.12818) · [Benchmarks](https://benchmarks.hindsight.vectorize.io/) · [RAG vs Memory](https://hindsight.vectorize.io/developer/rag-vs-hindsight)
|
||||
**Documentation:**
|
||||
- [https://hindsight.vectorize.io](https://hindsight.vectorize.io)
|
||||
|
||||
**Clients:**
|
||||
- [Python](https://hindsight.vectorize.io/sdks/python) · [Node.js](https://hindsight.vectorize.io/sdks/nodejs) · [Go](https://hindsight.vectorize.io/sdks/go) · [CLI](https://hindsight.vectorize.io/sdks/cli) · [REST API](https://hindsight.vectorize.io/api-reference)
|
||||
- [Python](http://hindsight.vectorize.io/sdks/python)
|
||||
- [Node.js](http://hindsight.vectorize.io/sdks/nodejs)
|
||||
- [REST API](https://hindsight.vectorize.io/api-reference)
|
||||
- [CLI](https://hindsight.vectorize.io/sdks/cli)
|
||||
|
||||
**Community:**
|
||||
- [Slack](https://join.slack.com/t/hindsight-space/shared_invite/zt-3nhbm4w29-LeSJ5Ixi6j8PdiYOCPlOgg)
|
||||
- [GitHub Issues](https://github.com/vectorize-io/hindsight/issues)
|
||||
|
||||
---
|
||||
|
||||
## Star History
|
||||
|
||||
[](https://github.com/vectorize-io/hindsight/stargazers)
|
||||
---
|
||||
|
||||
## Supported Platforms
|
||||
|
||||
| Platform | Docker | Bare Metal (pip) | Embedded DB (pg0) |
|
||||
|----------|--------|------------------|--------------------|
|
||||
| **Linux** (x86_64, ARM64) | ✅ | ✅ | ✅ |
|
||||
| **macOS** (Apple Silicon / arm64) | ✅ | ✅ | ✅ |
|
||||
| **macOS** (Intel / x86_64) | ✅ | ⚠️ | ✅ |
|
||||
| **Windows** (x86_64) | ✅ | ✅ | ✅ |
|
||||
|
||||
⚠️ Intel Macs: use `hindsight-all-slim` — see the [installation guide](https://hindsight.vectorize.io/developer/installation#supported-platforms) for details.
|
||||
|
||||
---
|
||||
|
||||
|
||||
@@ -65,7 +65,7 @@ services:
|
||||
environment:
|
||||
# LLM Configuration
|
||||
HINDSIGHT_API_LLM_PROVIDER: ${HINDSIGHT_API_LLM_PROVIDER:-openai}
|
||||
HINDSIGHT_API_LLM_API_KEY: ${HINDSIGHT_API_LLM_API_KEY:-}
|
||||
HINDSIGHT_API_LLM_API_KEY: ${OPENAI_API_KEY:-your-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 HINDSIGHT_API_LLM_API_KEY=sk-xxx
|
||||
export OPENAI_API_KEY=sk-xxx
|
||||
|
||||
docker compose -f docker/docker-compose/custom-models/docker-compose.yaml up --build
|
||||
```
|
||||
|
||||
@@ -3,12 +3,11 @@ name: hindsight-custom-models
|
||||
# in at build time, so pod startup does not depend on HuggingFace at runtime.
|
||||
#
|
||||
# Quick start:
|
||||
# export HINDSIGHT_API_LLM_API_KEY=sk-xxx
|
||||
# export OPENAI_API_KEY=sk-xxx
|
||||
# docker compose -f docker/docker-compose/custom-models/docker-compose.yaml up --build
|
||||
#
|
||||
# Required environment variables:
|
||||
# - HINDSIGHT_API_LLM_API_KEY (pair it with HINDSIGHT_API_LLM_PROVIDER to use
|
||||
# a provider other than the default openai)
|
||||
# - OPENAI_API_KEY (or configure another LLM provider via HINDSIGHT_API_LLM_*)
|
||||
|
||||
services:
|
||||
hindsight:
|
||||
@@ -26,7 +25,7 @@ services:
|
||||
- "9999:9999"
|
||||
environment:
|
||||
HINDSIGHT_API_LLM_PROVIDER: ${HINDSIGHT_API_LLM_PROVIDER:-openai}
|
||||
HINDSIGHT_API_LLM_API_KEY: ${HINDSIGHT_API_LLM_API_KEY:-}
|
||||
HINDSIGHT_API_LLM_API_KEY: ${OPENAI_API_KEY:-your-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=${HINDSIGHT_API_LLM_API_KEY:?Please set the HINDSIGHT_API_LLM_API_KEY env variable}
|
||||
- HINDSIGHT_API_LLM_API_KEY=${OPENAI_API_KEY?Please set the OPENAI_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: ${HINDSIGHT_API_LLM_API_KEY:-}
|
||||
HINDSIGHT_API_LLM_API_KEY: ${OPENAI_API_KEY:-your-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: ${HINDSIGHT_API_LLM_API_KEY:-}
|
||||
HINDSIGHT_API_LLM_API_KEY: ${OPENAI_API_KEY:-your-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: ${HINDSIGHT_API_LLM_API_KEY:-}
|
||||
HINDSIGHT_API_LLM_API_KEY: ${OPENAI_API_KEY:-your-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=${HINDSIGHT_API_LLM_API_KEY:?Please set the HINDSIGHT_API_LLM_API_KEY env variable}
|
||||
- HINDSIGHT_API_LLM_API_KEY=${OPENAI_API_KEY?Please set the OPENAI_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
|
||||
|
||||
@@ -1,107 +0,0 @@
|
||||
# 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.
|
||||
@@ -1,113 +0,0 @@
|
||||
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,12 +8,17 @@ HINDSIGHT_VERSION=latest
|
||||
|
||||
# LLM Configuration
|
||||
HINDSIGHT_API_LLM_PROVIDER=openai
|
||||
HINDSIGHT_API_LLM_API_KEY=your-openai-api-key-here
|
||||
OPENAI_API_KEY=your-openai-api-key-here
|
||||
|
||||
# Alternative LLM providers (uncomment and set the key above accordingly):
|
||||
# Alternative LLM providers (uncomment and configure as needed):
|
||||
# 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
|
||||
- An OpenAI API key (or a key for another LLM provider)
|
||||
- OpenAI API key (or another LLM provider)
|
||||
|
||||
## Quick Start
|
||||
|
||||
```bash
|
||||
# Set environment variables
|
||||
export HINDSIGHT_DB_PASSWORD="your-secure-password"
|
||||
export HINDSIGHT_API_LLM_API_KEY="your-openai-api-key"
|
||||
export OPENAI_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` |
|
||||
| `HINDSIGHT_API_LLM_API_KEY` | API key for the LLM provider | (required) |
|
||||
| `OPENAI_API_KEY` | OpenAI API key | (required) |
|
||||
| `HINDSIGHT_API_LLM_PROVIDER` | LLM provider | `openai` |
|
||||
|
||||
### Why Timescale Extensions?
|
||||
|
||||
@@ -8,8 +8,7 @@ name: hindsight
|
||||
#
|
||||
# Required environment variables:
|
||||
# - HINDSIGHT_DB_PASSWORD: Password for the PostgreSQL user
|
||||
# - HINDSIGHT_API_LLM_API_KEY (pair it with HINDSIGHT_API_LLM_PROVIDER to use
|
||||
# a provider other than the default openai)
|
||||
# - OPENAI_API_KEY (or configure another LLM provider)
|
||||
#
|
||||
# Optional environment variables with defaults:
|
||||
# - HINDSIGHT_VERSION: Hindsight application version (default: latest)
|
||||
@@ -81,7 +80,7 @@ services:
|
||||
environment:
|
||||
# LLM Configuration
|
||||
HINDSIGHT_API_LLM_PROVIDER: ${HINDSIGHT_API_LLM_PROVIDER:-openai}
|
||||
HINDSIGHT_API_LLM_API_KEY: ${HINDSIGHT_API_LLM_API_KEY:-}
|
||||
HINDSIGHT_API_LLM_API_KEY: ${OPENAI_API_KEY:-your-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: ${HINDSIGHT_API_LLM_API_KEY:-}
|
||||
HINDSIGHT_API_LLM_API_KEY: ${OPENAI_API_KEY:-your-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}
|
||||
|
||||
@@ -0,0 +1,341 @@
|
||||
# v2 Knowledge Pages — Implementation Plan
|
||||
|
||||
> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking.
|
||||
|
||||
**Goal:** Make v2 knowledge pages a reliable, cleanly-tiered "wiki" surface: passive `entity_labels` tier-tagging + tag-scoped seeded pages, a `hindsight_*` MCP surface with one active `capture_initiative` verb that creates per-initiative pages linked by tag, and SessionStart/UserPromptSubmit page-roster injection.
|
||||
|
||||
**Architecture:** Shared TS core (`hindsight-integrations/hindsight-coding-agents`). Extraction stays blind to "pages"; classification is intrinsic (`knowledge:<tier>` tags), pages are tag-scoped saved views. Per-initiative navigation via a `relatedPageId:<id>` tag the synthesizer renders into `[[page:<id>]]`, with the Initiatives folder/roster as the guaranteed fallback.
|
||||
|
||||
**Tech Stack:** TypeScript, vitest, tsup bundling. Hindsight REST API (`/knowledge-base/*`, `/mental-models`, `/memories`, bank `/config`).
|
||||
|
||||
**Spec:** `docs/superpowers/specs/2026-07-25-v2-knowledge-pages-design.md`
|
||||
|
||||
**Working dir for all commands:** `hindsight-integrations/hindsight-coding-agents`
|
||||
**Test command:** `npx vitest run <file>` (fast suite; excludes `*.live.test.ts`). Full check: `npx vitest run && npx tsc --noEmit`.
|
||||
|
||||
**Conventions to follow (existing patterns):**
|
||||
- `HindsightClient` HTTP via `this.req("METHOD", this.bankUrl(path), body?)`; JSON via `await r.json()`.
|
||||
- MCP tools are SDK-free `ToolSpec { name, description, inputSchema (ZodRawShape), handler }`; wrap handler bodies in `guarded(...)`; `ok(value)` / `err(e)` result helpers.
|
||||
- Fail-open everywhere in hooks; pure logic separated from stdin/stdout plumbing.
|
||||
- Do NOT add a Claude co-author trailer to any commit.
|
||||
|
||||
---
|
||||
|
||||
## Task 1: Config field `pageRefreshEveryTurns`
|
||||
|
||||
**Files:**
|
||||
- Modify: `src/core/config.ts`
|
||||
- Test: `src/core/config.test.ts`
|
||||
|
||||
- [ ] **Step 1: Write failing test** — assert the default resolves to 10 and an override wins.
|
||||
|
||||
```ts
|
||||
it("pageRefreshEveryTurns defaults to 10 and is overridable", () => {
|
||||
expect(loadConfig({ harness: "claude-code", projectDir: process.cwd() }).pageRefreshEveryTurns).toBe(10);
|
||||
});
|
||||
```
|
||||
(Add an override case mirroring the existing override tests in this file.)
|
||||
|
||||
- [ ] **Step 2: Run** `npx vitest run src/core/config.test.ts` → FAIL (property missing).
|
||||
- [ ] **Step 3: Implement** — add `pageRefreshEveryTurns: number` to the `Config` type and default `10` in the same place `recallMaxTokens`/`reflectTimeoutMs` are defined/merged. Follow the exact merge/layering pattern already used for numeric fields.
|
||||
- [ ] **Step 4: Run** the test → PASS.
|
||||
- [ ] **Step 5: Commit** `git add src/core/config.ts src/core/config.test.ts && git commit -m "feat(core): add pageRefreshEveryTurns config (default 10)"`
|
||||
|
||||
---
|
||||
|
||||
## Task 2: `knowledge-injection.ts` — roster/preamble formatting (pure, new)
|
||||
|
||||
**Files:**
|
||||
- Create: `src/core/knowledge-injection.ts`
|
||||
- Test: `src/core/knowledge-injection.test.ts`
|
||||
|
||||
Pure, SDK-free, no network. Parses the `listPages()` payload and formats the two injections.
|
||||
|
||||
- [ ] **Step 1: Write failing tests**
|
||||
|
||||
```ts
|
||||
import { describe, expect, it } from "vitest";
|
||||
import { parsePageList, buildKnowledgePreamble, buildRosterRefresh } from "./knowledge-injection";
|
||||
|
||||
describe("parsePageList", () => {
|
||||
it("extracts {id,title} from the mental-model list shape, tolerating junk", () => {
|
||||
const raw = { items: [{ id: "p1", name: "Component map" }, { id: "p2", name: "Core concepts" }, { nope: 1 }] };
|
||||
expect(parsePageList(raw)).toEqual([{ id: "p1", title: "Component map" }, { id: "p2", title: "Core concepts" }]);
|
||||
});
|
||||
it("returns [] for null/garbage", () => {
|
||||
expect(parsePageList(null)).toEqual([]);
|
||||
expect(parsePageList(42 as unknown)).toEqual([]);
|
||||
});
|
||||
});
|
||||
|
||||
describe("buildKnowledgePreamble", () => {
|
||||
it("includes guidance, a roster of pages, and a refresh note", () => {
|
||||
const out = buildKnowledgePreamble([{ id: "p1", title: "Component map" }]);
|
||||
expect(out).toContain("<hindsight_knowledge>");
|
||||
expect(out).toContain("Component map");
|
||||
expect(out).toContain("p1");
|
||||
expect(out).toMatch(/hindsight_read_knowledge_page/);
|
||||
});
|
||||
it("has an empty-state line when there are no pages", () => {
|
||||
const out = buildKnowledgePreamble([]);
|
||||
expect(out).toMatch(/no knowledge pages yet|still learning/i);
|
||||
});
|
||||
});
|
||||
|
||||
describe("buildRosterRefresh", () => {
|
||||
it("is a compact 'current pages' block listing ids+titles", () => {
|
||||
const out = buildRosterRefresh([{ id: "p1", title: "Component map" }]);
|
||||
expect(out).toContain("Component map");
|
||||
expect(out).toContain("p1");
|
||||
});
|
||||
it("returns undefined when there are no pages (nothing to refresh)", () => {
|
||||
expect(buildRosterRefresh([])).toBeUndefined();
|
||||
});
|
||||
});
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Run** `npx vitest run src/core/knowledge-injection.test.ts` → FAIL.
|
||||
- [ ] **Step 3: Implement**
|
||||
|
||||
```ts
|
||||
export interface PageRef { id: string; title: string; }
|
||||
|
||||
/** Defensive parse of HindsightClient.listPages() (GET /mental-models?detail=metadata → {items:[{id,name}]}). */
|
||||
export function parsePageList(raw: unknown): PageRef[] {
|
||||
const items = (raw as { items?: unknown })?.items;
|
||||
if (!Array.isArray(items)) return [];
|
||||
const out: PageRef[] = [];
|
||||
for (const it of items) {
|
||||
const id = (it as { id?: unknown })?.id;
|
||||
const name = (it as { name?: unknown })?.name;
|
||||
if (typeof id === "string" && typeof name === "string") out.push({ id, title: name });
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
function roster(pages: PageRef[]): string {
|
||||
return pages.map((p) => `- ${p.title} (${p.id})`).join("\n");
|
||||
}
|
||||
|
||||
/** SessionStart: teach when/why to use pages + list what exists. Empty-state aware. */
|
||||
export function buildKnowledgePreamble(pages: PageRef[]): string {
|
||||
const body = pages.length
|
||||
? `Knowledge pages available in this repository:\n${roster(pages)}`
|
||||
: "No knowledge pages yet — Hindsight is still learning this repo; they'll appear as it processes.";
|
||||
return (
|
||||
"<hindsight_knowledge>\n" +
|
||||
"This repository has a Hindsight knowledge base: curated, continuously-updated pages summarizing its " +
|
||||
"durable engineering knowledge (architecture, components, conventions, key decisions, and in-flight initiatives).\n" +
|
||||
"Before substantial work, consult the relevant pages instead of re-deriving understanding from the code: read " +
|
||||
"Conventions before writing new code, the Component map before changing a subsystem, and an initiative's page " +
|
||||
"before continuing that feature.\n" +
|
||||
`${body}\n` +
|
||||
"Read one with hindsight_read_knowledge_page(page_id). Follow any [[page:<id>]] links you see. The list is " +
|
||||
"re-injected for you periodically as it changes.\n" +
|
||||
"</hindsight_knowledge>"
|
||||
);
|
||||
}
|
||||
|
||||
/** Periodic UserPromptSubmit refresh — compact, or undefined when there's nothing to show. */
|
||||
export function buildRosterRefresh(pages: PageRef[]): string | undefined {
|
||||
if (!pages.length) return undefined;
|
||||
return (
|
||||
"<hindsight_knowledge_refresh>\n" +
|
||||
`Current Hindsight knowledge pages (may have changed):\n${roster(pages)}\n` +
|
||||
"Read any with hindsight_read_knowledge_page(page_id).\n" +
|
||||
"</hindsight_knowledge_refresh>"
|
||||
);
|
||||
}
|
||||
```
|
||||
|
||||
- [ ] **Step 4: Run** the test → PASS.
|
||||
- [ ] **Step 5: Commit** `git add src/core/knowledge-injection.ts src/core/knowledge-injection.test.ts && git commit -m "feat(core): knowledge-injection roster/preamble formatting"`
|
||||
|
||||
---
|
||||
|
||||
## Task 3: `entity_labels` tier vocabulary + configureBank wiring
|
||||
|
||||
**Files:**
|
||||
- Modify: `src/core/missions.ts` (add `KNOWLEDGE_LABELS`)
|
||||
- Modify: `src/core/hindsight.ts` (`configureBank` PATCH sets `entity_labels`)
|
||||
- Test: `src/core/hindsight.*.test.ts` (add/extend a config test with a mock client)
|
||||
|
||||
- [ ] **Step 1: Write failing test** — assert `configureBank` PATCHes `/config` with `entity_labels` containing the `knowledge` group and its five values, and `entities_allow_free_form: true`. Use the existing fetch/req mock pattern from `hindsight.*.test.ts`; capture the PATCH body to `/config` and assert on it.
|
||||
|
||||
- [ ] **Step 2: Run** → FAIL.
|
||||
- [ ] **Step 3: Implement**
|
||||
- In `missions.ts`, export `KNOWLEDGE_LABELS` — the exact object from the spec §4 (`key:"knowledge"`, `type:"multi-values"`, `optional:true`, `tag:true`, the verbose group `description`, and the five value `{value,description}` entries: feature-work, decision, convention, component, concept).
|
||||
- In `hindsight.ts::configureBank`, extend the existing `PATCH .../config` `updates` object to include `entity_labels: [KNOWLEDGE_LABELS]` and `entities_allow_free_form: true`. Import `KNOWLEDGE_LABELS`.
|
||||
- Update the `[bank] configured …` log to mention `entity_labels`.
|
||||
- [ ] **Step 4: Run** → PASS.
|
||||
- [ ] **Step 5: Commit** `git add src/core/missions.ts src/core/hindsight.ts src/core/hindsight.*.test.ts && git commit -m "feat(core): passive knowledge entity_labels tier vocabulary + configureBank wiring"`
|
||||
|
||||
---
|
||||
|
||||
## Task 4: Tag-scoped seeded pages + Initiatives folder + link source_query
|
||||
|
||||
**Files:**
|
||||
- Modify: `src/core/missions.ts` (`PAGES` gain `tags`; Initiatives `source_query` link instruction)
|
||||
- Modify: `src/core/hindsight.ts` (`ensureFolder`, `createPages` sets page `tags` + parents Initiatives under the folder)
|
||||
- Test: `src/core/hindsight.pages.test.ts`
|
||||
|
||||
- [ ] **Step 1: Write failing tests** (mock client `req`):
|
||||
- Each seeded page POST to `/knowledge-base/pages` includes `tags: ["knowledge:<tier>"]` mapped per the spec §5 table.
|
||||
- The Initiatives page is created with `parent_id` equal to the id returned by an Initiatives folder POST to `/knowledge-base/folders`.
|
||||
- `ensureFolder("Initiatives")` returns an existing root folder's id when the tree already contains it (GET `/knowledge-base/tree`) and does NOT POST a duplicate.
|
||||
|
||||
- [ ] **Step 2: Run** → FAIL.
|
||||
- [ ] **Step 3: Implement**
|
||||
- `missions.ts`: add `tags: string[]` to each `PAGES` entry (feature-work/decision/convention/component/concept mapping). Append to the Initiatives `source_query`: *"When a source memory carries a tag of the form `relatedPageId:<id>`, include a Markdown link `[[page:<id>]]` to that page in the summary, so each initiative links to its detailed page."*
|
||||
- `hindsight.ts`: add
|
||||
```ts
|
||||
/** Find a root folder by name (case-insensitive) or create it; returns its id. */
|
||||
async ensureFolder(name: string): Promise<string | undefined> {
|
||||
try {
|
||||
const tree = (await (await this.req("GET", this.bankUrl("/knowledge-base/tree"))).json()) as
|
||||
{ roots?: { id?: string; kind?: string; name?: string }[] };
|
||||
const hit = (tree.roots || []).find((n) => n.kind === "folder" && (n.name || "").toLowerCase() === name.toLowerCase());
|
||||
if (hit?.id) return hit.id;
|
||||
} catch { /* fall through to create */ }
|
||||
try {
|
||||
const r = await this.req("POST", this.bankUrl("/knowledge-base/folders"), { name });
|
||||
return ((await r.json()) as { id?: string }).id;
|
||||
} catch { return undefined; }
|
||||
}
|
||||
```
|
||||
- In `createPages()`: before the loop, `const initiativesFolderId = await this.ensureFolder("Initiatives");`. For each page, build body `{ name, source_query, tags: p.tags, parent_id: <initiativesFolderId if this is the Initiatives page else undefined>, trigger: { fact_types:[...], refresh_after_consolidation:true } }`. (Page-level `tags` drives synthesis scoping via `RefreshTagFiltering`; `tags_match` defaults to `all_strict` when tags present.)
|
||||
- [ ] **Step 4: Run** → PASS.
|
||||
- [ ] **Step 5: Commit** `git add src/core/missions.ts src/core/hindsight.ts src/core/hindsight.pages.test.ts && git commit -m "feat(core): tag-scope seeded pages, Initiatives folder, relatedPageId link source_query"`
|
||||
|
||||
---
|
||||
|
||||
## Task 5: Client helpers — per-initiative page + marker retain
|
||||
|
||||
**Files:**
|
||||
- Modify: `src/core/hindsight.ts` (`captureInitiative`)
|
||||
- Test: `src/core/hindsight.pages.test.ts`
|
||||
|
||||
- [ ] **Step 1: Write failing tests** (mock `req`):
|
||||
- `captureInitiative({title:"Retry backoff for the uploader", summary:"…"})` → derives slug `retry-backoff-for-the-uploader`, POSTs a page id `initiative-<slug>` to `/knowledge-base/pages` with `parent_id` = the Initiatives folder and `tags: ["knowledge:feature-work"]`, AND POSTs a marker to `/memories` (via `retain`) tagged `["knowledge:feature-work","relatedPageId:initiative-<slug>"]`, strategy `session` or `document` (pick `document`), `async:true`. Returns `{ page_id: "initiative-<slug>" }`.
|
||||
- Slug is deterministic and identical between the page id and the `relatedPageId:` tag value.
|
||||
- Enhancement path: `captureInitiative({title, summary, relatesToPageId:"initiative-x"})` POSTs NO new page; marker tagged `relatedPageId:initiative-x`; returns `{ page_id: "initiative-x" }`.
|
||||
|
||||
- [ ] **Step 2: Run** → FAIL.
|
||||
- [ ] **Step 3: Implement**
|
||||
|
||||
```ts
|
||||
private slugify(s: string): string {
|
||||
return s.toLowerCase().replace(/[^a-z0-9]+/g, "-").replace(/^-+|-+$/g, "").slice(0, 60) || "initiative";
|
||||
}
|
||||
|
||||
/** Active-path capture: register a major feature as a per-initiative page + a tagged marker memory. */
|
||||
async captureInitiative(args: { title: string; summary: string; relatesToPageId?: string }): Promise<{ page_id: string }> {
|
||||
const pageId = args.relatesToPageId ?? `initiative-${this.slugify(args.title)}`;
|
||||
if (!args.relatesToPageId) {
|
||||
const folderId = await this.ensureFolder("Initiatives");
|
||||
await this.req("POST", this.bankUrl("/knowledge-base/pages"), {
|
||||
name: args.title,
|
||||
source_query: `Summarize the "${args.title}" initiative: what is being built or changed and why, and its current state — drawn from the project's memory.`,
|
||||
parent_id: folderId,
|
||||
tags: ["knowledge:feature-work", `relatedPageId:${pageId}`],
|
||||
trigger: { fact_types: ["world", "experience", "observation"], refresh_after_consolidation: true },
|
||||
});
|
||||
}
|
||||
const verb = args.relatesToPageId ? "Enhancement to an existing initiative" : "New initiative";
|
||||
const content = `${verb}: ${args.title}. ${args.summary}`;
|
||||
await this.retain(content, "initiative marker", pageId /* not a stable doc id requirement; see note */,
|
||||
["knowledge:feature-work", `relatedPageId:${pageId}`], "document", { async: true });
|
||||
return { page_id: pageId };
|
||||
}
|
||||
```
|
||||
- NOTE: use a UNIQUE document id per marker (e.g. `initiative-marker-<slug>-<n>`), NOT `pageId`, so repeated enhancement captures accrue instead of replacing. Since `Date.now()` is fine here (runtime, not a workflow script), suffix with a timestamp: `initiative-marker-${this.slugify(args.title)}-${Date.now()}`. Keep the `relatedPageId` tag equal to `pageId`.
|
||||
- Confirm `retain(content, context, documentId, tags, strategy, opts)` signature matches current `HindsightClient.retain`.
|
||||
|
||||
- [ ] **Step 4: Run** → PASS.
|
||||
- [ ] **Step 5: Commit** `git add src/core/hindsight.ts src/core/hindsight.pages.test.ts && git commit -m "feat(core): captureInitiative — per-initiative page + relatedPageId marker"`
|
||||
|
||||
---
|
||||
|
||||
## Task 6: MCP surface — `hindsight_*` grounding + `capture_initiative`; drop page CRUD
|
||||
|
||||
**Files:**
|
||||
- Modify: `src/core/knowledge-tools.ts`
|
||||
- Modify: `src/mcp-server.ts` (only if it references removed tool names)
|
||||
- Test: `src/core/knowledge-tools.test.ts`, `src/mcp-server.test.ts` (tool-count assertions)
|
||||
|
||||
- [ ] **Step 1: Write failing tests**
|
||||
- `buildKnowledgeTools(client, bankId)` returns exactly these tool names: `hindsight_get_current_bank`, `hindsight_list_knowledge_pages`, `hindsight_read_knowledge_page`, `hindsight_search_memory`, `hindsight_capture_initiative`, `hindsight_ingest_document`. (Assert the set; update any count assertion.)
|
||||
- `hindsight_capture_initiative` handler calls `client.captureInitiative` with `{title, summary, relatesToPageId?}` and returns the page id (mock client).
|
||||
- No `create_page` / `update_page` / `delete_page` tools are present.
|
||||
- Each tool still fails closed via `guarded` (a thrown client error → `isError:true`, no throw).
|
||||
|
||||
- [ ] **Step 2: Run** → FAIL.
|
||||
- [ ] **Step 3: Implement**
|
||||
- Rebuild the `buildKnowledgeTools` list: rename read/recall/ingest/bank tools to the `hindsight_*` names; drop `create_page`/`update_page`/`delete_page`; add `hindsight_capture_initiative` with `inputSchema { title: z.string(), summary: z.string(), relates_to_page_id: z.string().optional() }` calling `client.captureInitiative({ title, summary, relatesToPageId: relates_to_page_id })`.
|
||||
- Use the **verbatim agent-facing `description` strings** from the spec §6 / the brainstorm (grounding tools + the explicit WHEN/WHEN-NOT `capture_initiative` description).
|
||||
- Update `mcp-server.ts` only if it enumerates tool names; otherwise it consumes `buildKnowledgeTools` generically and needs no change.
|
||||
- [ ] **Step 4: Run** `npx vitest run src/core/knowledge-tools.test.ts src/mcp-server.test.ts` → PASS.
|
||||
- [ ] **Step 5: Commit** `git add src/core/knowledge-tools.ts src/mcp-server.ts src/core/knowledge-tools.test.ts src/mcp-server.test.ts && git commit -m "feat(mcp): hindsight_* grounding tools + capture_initiative; remove raw page CRUD from agent"`
|
||||
|
||||
---
|
||||
|
||||
## Task 7: SessionStart — preamble + roster
|
||||
|
||||
**Files:**
|
||||
- Modify: `src/core/session-start.ts`
|
||||
- Test: `src/core/session-start.test.ts`
|
||||
|
||||
- [ ] **Step 1: Write failing tests**
|
||||
- `buildSessionStartContext` now fetches pages via the client and injects `buildKnowledgePreamble(...)` instead of the static `KNOWLEDGE_MISSION`. Extend the `SeedContextClient` interface with `listPages(): Promise<unknown>`; the mock returns `{items:[{id:"p1",name:"Component map"}]}` and the output contains "Component map".
|
||||
- listPages failure is fail-open: the preamble still renders (empty-state) and the seed logic is unaffected.
|
||||
|
||||
- [ ] **Step 2: Run** → FAIL.
|
||||
- [ ] **Step 3: Implement**
|
||||
- Add `listPages` to `SeedContextClient`.
|
||||
- Replace the `parts.push(KNOWLEDGE_MISSION)` line with: fetch `const pages = parsePageList(await client.listPages().catch(() => null));` then `parts.push(buildKnowledgePreamble(pages));`. Import from `./knowledge-injection`.
|
||||
- Remove the now-unused `KNOWLEDGE_MISSION` export if nothing else references it (grep first; keep if referenced).
|
||||
- [ ] **Step 4: Run** → PASS.
|
||||
- [ ] **Step 5: Commit** `git add src/core/session-start.ts src/core/session-start.test.ts && git commit -m "feat(core): SessionStart injects page roster + guidance preamble"`
|
||||
|
||||
---
|
||||
|
||||
## Task 8: UserPromptSubmit — hook-counted periodic roster refresh
|
||||
|
||||
**Files:**
|
||||
- Modify: `src/core/hook.ts`
|
||||
- Test: `src/core/hook.test.ts`
|
||||
|
||||
- [ ] **Step 1: Write failing tests**
|
||||
- The session cache round-trips `{answer, turns}`; each `buildHookOutput` call increments `turns`.
|
||||
- Add `listPages` to the `HookClient` interface. On a turn where `turns % cfg.pageRefreshEveryTurns === 0`, the output includes `buildRosterRefresh(...)` content (assert "Component map" appears); on other turns it does not.
|
||||
- Refresh is fail-open (a `listPages` rejection doesn't break recall/injection).
|
||||
- First-turn behavior (reflect) unchanged.
|
||||
|
||||
- [ ] **Step 2: Run** → FAIL.
|
||||
- [ ] **Step 3: Implement**
|
||||
- Extend the cache read/write to `{ answer?: string; turns?: number }`. Compute `const turns = (cached.turns ?? 0) + 1;` and persist it alongside `answer`.
|
||||
- Add `listPages(): Promise<unknown>` to `HookClient`.
|
||||
- After computing `memBlock`, if `cfg.pageRefreshEveryTurns > 0 && turns % cfg.pageRefreshEveryTurns === 0`, `try { const refresh = buildRosterRefresh(parsePageList(await client.listPages())); if (refresh) blocks.push(refresh); } catch { /* fail-open */ }`. Kick the `listPages` call off concurrently with recall to avoid added latency.
|
||||
- Import from `./knowledge-injection`.
|
||||
- [ ] **Step 4: Run** `npx vitest run src/core/hook.test.ts` → PASS.
|
||||
- [ ] **Step 5: Commit** `git add src/core/hook.ts src/core/hook.test.ts && git commit -m "feat(core): UserPromptSubmit hook-counted periodic page-roster refresh"`
|
||||
|
||||
---
|
||||
|
||||
## Task 9: Full check + LLM behavior (live) verification
|
||||
|
||||
**Files:**
|
||||
- Modify: `src/system.live.test.ts` (add coverage; runs only under `HINDSIGHT_LIVE_E2E=1`)
|
||||
|
||||
- [ ] **Step 1: Full fast suite + types** — `npx vitest run && npx tsc --noEmit` → all green.
|
||||
- [ ] **Step 2: Add a live assertion** (guarded by the existing live env flag) that after seeding a small repo + one `captureInitiative`, the Initiatives page content contains a `[[page:initiative-…]]` link (verifies the `relatedPageId` → link rendering end-to-end). Keep it in the live suite; do not run in the fast job.
|
||||
- [ ] **Step 3: Manual/live run** (optional, operator): `HINDSIGHT_API_URL=http://localhost:8888 npm run test:live`.
|
||||
- [ ] **Step 4: Commit** `git add src/system.live.test.ts && git commit -m "test(live): initiative page renders relatedPageId link end-to-end"`
|
||||
|
||||
---
|
||||
|
||||
## Final review
|
||||
|
||||
- [ ] Dispatch a final code-reviewer over the whole change set against the spec (`docs/superpowers/specs/2026-07-25-v2-knowledge-pages-design.md`).
|
||||
- [ ] Rebuild + dev-install the `claude-code-v2` bundle so the running plugin picks up the new hooks/MCP (`bash scripts/dev-install.sh`); do not push/PR without explicit consent.
|
||||
- [ ] Note deferred follow-ups: session drill-down tag, `capture_decision`, `gotcha` tier, older-bank reseed requirement.
|
||||
@@ -0,0 +1,135 @@
|
||||
# v2 Knowledge Pages — Design Spec
|
||||
|
||||
**Status:** approved in brainstorm (2026-07-25), pending implementation plan
|
||||
**Scope:** `hindsight-integrations/hindsight-coding-agents` (shared TS core) + `claude-code-v2` wrapper
|
||||
**Motivation:** make knowledge pages a real, trustworthy "wiki" surface for the vectorize-crm demo (the `knowledge-pages-as-trust-surface` principle) — the agent reliably knows what pages exist, pages are cleanly tiered instead of blended, and major initiatives become first-class, linkable pages.
|
||||
|
||||
---
|
||||
|
||||
## 1. Problem
|
||||
|
||||
Three gaps in the current v2 branch:
|
||||
|
||||
1. **Page discovery is a blind fetch.** SessionStart injects a static `KNOWLEDGE_MISSION` telling the agent to call `agent_knowledge_list_pages`, but hands it **no roster** — the agent never learns a page exists unless it independently decides to call the tool. Per-turn recall injects facts, not pages.
|
||||
2. **Pages are blended.** Neither the seeded `PAGES` nor the agent's `create_page` tool scope synthesis by tag, so every page synthesizes from the whole bank filtered only by `fact_type`. Git-log, session, and survey memories all bleed into every page.
|
||||
3. **No first-class initiative tracking / linking.** Hindsight has no native page-to-page links. A "major feature" leaves no durable, navigable page a future session can pick up.
|
||||
|
||||
## 2. Principles applied
|
||||
|
||||
- Automatic/visible value; zero out-of-band CLI; memory beats code search; knowledge pages as a trust surface; minimal post-setup burden.
|
||||
- Modular units, small files, follow existing patterns (per-hook specs, fail-open, unit-testable pure cores).
|
||||
- The **memory extractor never knows what a "page" is.** Classification is by the fact's *intrinsic* nature; pages are application-side saved views. No abstraction leak into extraction.
|
||||
|
||||
## 3. Architecture overview
|
||||
|
||||
Two complementary curation paths + a discovery layer:
|
||||
|
||||
- **Passive (automatic):** `entity_labels` schema-forces the extractor to tag qualifying facts `knowledge:<tier>`. Seeded **tier pages** each filter on one tier tag. No agent effort.
|
||||
- **Active (high-signal):** one intent-named MCP verb, `hindsight_capture_initiative`, lets the agent register a major feature as a **per-initiative page** with a tag-based link back from the aggregate Initiatives page.
|
||||
- **Discovery:** SessionStart injects guidance + the page roster; the UserPromptSubmit hook re-injects a fresh roster on a fixed cadence (hook-counted, not model-counted).
|
||||
|
||||
## 4. `entity_labels` — passive tier tagging
|
||||
|
||||
One hierarchical bank config group, set by `configureBank` at seed time:
|
||||
|
||||
```jsonc
|
||||
{
|
||||
"key": "knowledge",
|
||||
"type": "multi-values", // 0, 1, or several — empty is normal
|
||||
"optional": true,
|
||||
"tag": true, // emits knowledge:<value> onto the fact's tags
|
||||
"description": "Routing labels for this project's Hindsight KNOWLEDGE PAGES — curated, human-readable summaries of the repo's DURABLE engineering knowledge (architecture, key decisions, conventions, ongoing initiatives), each page rebuilt automatically from the facts labeled for it. Mark a fact only when it is durable, reusable knowledge a developer would still want surfaced in future sessions. IMPORTANT: leave this EMPTY for routine, transient, or operational facts — a passing test, a one-off command, a status update, a debugging dead-end. MOST facts should get no label here. Assign more than one value only when the fact genuinely fits several.",
|
||||
"values": [
|
||||
{ "value": "feature-work", "description": "A new feature, initiative, or enhancement being planned or built — the capability being added and the intent behind it. Not routine bug-fixes or chores." },
|
||||
{ "value": "decision", "description": "A technical decision that will constrain future work, with its rationale — why this approach was chosen over alternatives, or a rule deliberately adopted." },
|
||||
{ "value": "convention", "description": "An established way this project does things — naming, structure, testing, error handling, or another recurring pattern a contributor is expected to follow." },
|
||||
{ "value": "component", "description": "What a specific module, file, service, or subsystem is responsible for, or how components depend on and connect to one another." },
|
||||
{ "value": "concept", "description": "A domain concept, key abstraction, or piece of project vocabulary a new contributor must understand to work effectively." }
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
Notes:
|
||||
- `tag: true` → `_inject_label_tags` copies each `knowledge:<value>` onto the fact's `tags` (no extra query infra).
|
||||
- Selectivity (multi-values + "mostly empty" instruction) prevents force-fitting routine facts into a tier.
|
||||
|
||||
## 5. Seeded tier pages (tag-scoped)
|
||||
|
||||
Created via `/knowledge-base/pages` (supports `tags`, `trigger`, `parent_id`) — **not** `/mental-models`. Each `PAGES` entry gains a `trigger.tags` pin:
|
||||
|
||||
| Page | `trigger.tags` |
|
||||
| --- | --- |
|
||||
| Initiatives and enhancements | `["knowledge:feature-work"]` |
|
||||
| Key decisions and rationale | `["knowledge:decision"]` |
|
||||
| Conventions and patterns | `["knowledge:convention"]` |
|
||||
| Component map | `["knowledge:component"]` |
|
||||
| Core concepts | `["knowledge:concept"]` |
|
||||
|
||||
`tags_match` strict enough to exclude untagged facts (`all_strict`/`any_strict`). Tag matching is exact set-ops (no wildcards) — this is *why* the vocabulary is fixed, not per-feature.
|
||||
|
||||
## 6. MCP surface
|
||||
|
||||
Raw page CRUD (`create_page`/`update_page`/`delete_page`) is **removed** from the agent. The agent sees grounding tools + one capture verb. Naming convention: `hindsight_*`.
|
||||
|
||||
**Grounding**
|
||||
- `hindsight_list_knowledge_pages` `{}` — roster: id, title, one-line coverage. (agent-facing description as drafted in brainstorm)
|
||||
- `hindsight_read_knowledge_page` `{ page_id }` — full page content; follow `[[page:<id>]]` links by re-calling.
|
||||
- `hindsight_search_memory` `{ query, max_tokens? }` — raw fact recall for specifics pages don't cover.
|
||||
- `hindsight_get_current_bank` `{}` — minor introspection (kept).
|
||||
|
||||
**Capture**
|
||||
- `hindsight_capture_initiative` `{ title, summary, relates_to_page_id? }` — the one active verb. Explicit WHEN / WHEN-NOT description (as drafted). Returns the initiative page id.
|
||||
- `hindsight_ingest_document` `{ title, content }` — existing `agent_knowledge_ingest`, reframed.
|
||||
|
||||
(Full agent-facing descriptions are captured verbatim in the brainstorm thread and will be reproduced in the implementation plan.)
|
||||
|
||||
## 7. `hindsight_capture_initiative` mechanism
|
||||
|
||||
- Derive one slug `S` from `title`. Page id = `initiative-<S>`. **The slug in the tag and the page id are the same token, derived once** (cannot drift).
|
||||
- **New initiative** (`relates_to_page_id` omitted):
|
||||
1. Create page `initiative-<S>` (title from `title`, `source_query` about that initiative) under an **"Initiatives" folder** (tag-scoped).
|
||||
2. Retain a marker memory (text = title + summary) tagged `["knowledge:feature-work", "relatedPageId:initiative-<S>"]`. **No session tag** (decided — the MCP server has no Claude session id; faking one wouldn't link to the Stop write-back's `conversation:<sessionId>` doc anyway).
|
||||
- **Enhancement** (`relates_to_page_id` given): marker only, `relatedPageId = relates_to_page_id`; no new page. Re-invoking for the same initiative accrues markers → the page re-synthesizes with progress.
|
||||
|
||||
### Link survival (why `relatedPageId` as a tag, not in prose)
|
||||
|
||||
A tag is set directly via the retain `tags` param — it **bypasses LLM extraction entirely**, so it's guaranteed present verbatim (no REF-ID-style preservation needed at extraction). Verified: the reflect/synthesis path SELECTs `tags` and serializes facts via `_prune_nulls(model_dump())`, which keeps non-empty tags → **the synthesis LLM sees the tag.** The **Initiatives page `source_query`** instructs: *"when a memory carries a `relatedPageId:<id>` tag, emit a `[[page:<id>]]` link to it."* The link id is generated from the tag value at synthesis time, so it always matches the created page id.
|
||||
|
||||
- Only **Stage 2 (synthesis)** is probabilistic now (bounded token budget may omit some entries when there are many).
|
||||
- **Guaranteed fallback:** the per-initiative page always exists (created via API, independent of any LLM stage) and appears in the **Initiatives folder / injected roster**, so navigation works even if a synthesized inline link drops.
|
||||
|
||||
## 8. Page-access injection
|
||||
|
||||
- **SessionStart** (`session-start.ts`): replace static `KNOWLEDGE_MISSION` with a preamble = (a) guidance on *when/why* to consult pages, (b) the roster fetched via `client.listPages()` (`- <title> (<id>)`, empty-state aware), (c) a note that the list refreshes periodically. Cold repo → empty roster line; roster comes alive mid-session as seeding/survey complete.
|
||||
- **UserPromptSubmit** (`hook.ts`): extend the per-session cache (`{answer}` → `{answer, turns}`); the **hook** counts user turns and, roughly every `pageRefreshEveryTurns` (default 10, approximate), calls `listPages()` and injects a compact roster refresh. Runs concurrently with recall; **fail-open** (a refresh error never blocks the turn).
|
||||
- **Shared formatting** (new `core/knowledge-injection.ts`, SDK-free/unit-testable): `parsePageList(raw) -> {id,title}[]`, `buildKnowledgePreamble(pages)`, `buildRosterRefresh(pages)`.
|
||||
- **Config:** `pageRefreshEveryTurns` (default 10).
|
||||
|
||||
## 9. Non-goals / deferred
|
||||
|
||||
- Session drill-down tag on captured markers (dropped — see §7).
|
||||
- `hindsight_capture_decision` and other capture verbs (passive path covers those tiers; revisit if the aggregate pages aren't sharp enough).
|
||||
- A `gotcha`/`pitfall` tier (five tiers for now).
|
||||
- Native page-to-page links / backlinks (Hindsight has none; we approximate via folder tree + `relatedPageId`-driven `[[page:<id>]]`).
|
||||
|
||||
## 10. Risks / migration
|
||||
|
||||
- **Older banks** need re-seeding to pick up the new `entity_labels`, the `session` retain strategy, and the tag-scoped page triggers (`configureBank` sets them). User is starting fresh with v2 banks, so acceptable; live retain fails open otherwise.
|
||||
- **Stage-2 synthesis omission** for large initiative counts — mitigated by the folder/roster fallback.
|
||||
- **Instruction adherence** for the `source_query` link-rendering and the label selectivity — both are LLM-following behaviors; cover with an `hs_llm_core` judge test, and the deterministic mechanics (tag injection, roster formatting, slug/id equality, hook turn-counting) with fast unit tests.
|
||||
|
||||
## 11. Testing
|
||||
|
||||
- **Deterministic unit tests:** `knowledge-injection` formatting + empty-state; hook turn-counter + cadence; `capture_initiative` slug→id→tag equality and request shape (mock client); tag-scoped page request bodies; entity_labels config emitted by `configureBank`.
|
||||
- **LLM judge test (`hs_llm_core`):** label selectivity (routine facts get no `knowledge:*`), and `relatedPageId` → `[[page:<id>]]` rendering in a synthesized Initiatives page.
|
||||
|
||||
## 12. File map (anticipated)
|
||||
|
||||
- `src/core/knowledge-injection.ts` (new) — roster/preamble formatting.
|
||||
- `src/core/session-start.ts` — preamble + roster.
|
||||
- `src/core/hook.ts` — cache `{answer,turns}` + periodic roster refresh.
|
||||
- `src/core/config.ts` — `pageRefreshEveryTurns`.
|
||||
- `src/core/missions.ts` — `entity_labels` group; tag-scoped `PAGES`; Initiatives `source_query` link instruction.
|
||||
- `src/core/hindsight.ts` — `configureBank` sets `entity_labels`; `createPages` pins `trigger.tags` + Initiatives folder; new `createInitiativePage`/marker retain helpers.
|
||||
- `src/core/knowledge-tools.ts` — new `hindsight_*` grounding + `capture_initiative` tools; remove raw page CRUD from agent surface.
|
||||
- Tests alongside each.
|
||||
@@ -0,0 +1,139 @@
|
||||
# Reflect + Pages Runtime — Design Spec
|
||||
|
||||
**Status:** decided (2026-07-27), reconciles the earlier reflect-based runtime with the recall-based v2 into one opinionated path
|
||||
**Scope:** `hindsight-integrations/hindsight-coding-agents` (shared TS core) + `claude-code-v2` wrapper
|
||||
**Motivation:** the 33-task coding benchmark showed the v2 recall-per-prompt runtime *underperforms no memory* (35.0 mean corrections vs 32.0 baseline), while the earlier reflect-injection runtime beats baseline by 22% (25.0). This spec restores reflect as the only deep-memory path and replaces raw per-turn recall with lightweight injection from knowledge pages — "fast like recall, organized like reflect" — keeping v2's page/curation machinery where it earned its place and deleting it where it didn't.
|
||||
|
||||
---
|
||||
|
||||
## 1. Problem
|
||||
|
||||
Two prior iterations, each half right:
|
||||
|
||||
1. **Reflect runtime (v1):** one agentic REFLECT over the bank at session start, cached and re-injected every turn. Benchmark-proven (25.0 mean corrections) — but nothing surfaced mid-session; a task that drifted away from the first message got stale context.
|
||||
2. **Recall runtime (v2):** per-prompt recall injection for turn-by-turn visibility, plus knowledge pages as a trust surface. But raw recall injects unsynthesized fact fragments — noise that *hurt*: 35.0 mean corrections, worse than running with no memory at all.
|
||||
|
||||
| Runtime | Mean corrections (33-task benchmark) | vs no-memory (32.0) |
|
||||
| --- | --- | --- |
|
||||
| Reflect-injection (v1) | **25.0** | **−22%** |
|
||||
| Recall-per-prompt (v2) | 35.0 | +9% (regression) |
|
||||
| No memory | 32.0 | baseline |
|
||||
|
||||
The reconciliation: keep reflect's synthesis quality as the deep path, keep v2's per-turn visibility principle, but source the per-turn material from the already-synthesized knowledge pages instead of raw recall.
|
||||
|
||||
## 2. Decisions
|
||||
|
||||
Explicit, decided — not options:
|
||||
|
||||
1. **Reflect restored** as the only deep-memory path (session-start, agentic synthesis, cached + re-injected every turn).
|
||||
2. **Recall removed from the runtime** entirely. No per-prompt `recall` call.
|
||||
3. **No `memoryMode` flag.** One opinionated path; config is for environment, naming, and harness wiring only — never behavior selection.
|
||||
4. **Sections, not pages, are the per-turn injection unit** — locally matched, budget-trimmed, provenance-labeled.
|
||||
5. **JSON turn transcripts** replace the markdown tool-call transcript in the Stop-hook write-back, with compact action entries.
|
||||
6. **No tags / no `entity_labels`.** The server re-synthesizes pages after consolidation; "living pages" needs no client-side tagging machinery.
|
||||
|
||||
## 3. Runtime path — session start
|
||||
|
||||
Three steps, in order, all inside existing hooks (no out-of-band CLI):
|
||||
|
||||
### 3a. Cold-repo bootstrap (kept from v2)
|
||||
|
||||
On a bank with no prior memories: automatic shallow gitlog seed + codebase survey, exactly as v2 does it. The user never runs a setup command; the first session self-seeds. (Deep ingestion of that history is §7 — the seed here stays instant.)
|
||||
|
||||
### 3b. REFLECT once, on the first task message
|
||||
|
||||
The benchmark-proven core:
|
||||
|
||||
- On the first user prompt of the session, run one **REFLECT** — agentic synthesis over the whole bank, prompted to return the *root-cause decision with exact values* (concrete file paths, config values, version numbers — not summaries of summaries).
|
||||
- Cache the result per session; **re-inject it every turn**. It is the session's durable deep context.
|
||||
- One LLM-backed call per session, on the message that actually states the task — not on session-open, where there is nothing to reflect about.
|
||||
|
||||
### 3c. Page index build
|
||||
|
||||
Fetch all knowledge pages once (existing `listPages` + page reads), split each page at headings into **sections**, and build a **local section index** in the hook process. This index is what every subsequent turn matches against (§4) — no further server calls on the hot path.
|
||||
|
||||
## 4. Runtime path — every turn
|
||||
|
||||
Per-turn visibility, satisfied at ~zero latency and ~zero cost. Injection sources from **knowledge pages, not raw recall** — the material is already synthesized and organized; the turn hook only *selects* from it.
|
||||
|
||||
Mechanism (local, deterministic — no server call, no LLM call):
|
||||
|
||||
| Aspect | Design |
|
||||
| --- | --- |
|
||||
| Unit | Page **sections** (pages split at headings at index-build time) |
|
||||
| Matching | Lexical: prompt scored against each section by weighted term overlap; **heading hits weighted higher** than body hits |
|
||||
| Selection | Top 2–3 sections |
|
||||
| Budget | Trimmed to a **~700-token total** |
|
||||
| Provenance | Each snippet labeled `From <page> › <section>` + a tool pointer to read the full page |
|
||||
| Floor | A minimum-score threshold below which **nothing is injected** — silence over noise |
|
||||
| Refresh | Section index rebuilt on the existing 10-turn roster cadence (`pageRefreshEveryTurns`) |
|
||||
|
||||
The score floor is load-bearing: the benchmark showed that injecting weak matches is worse than injecting nothing (v2's regression). An empty injection is a correct outcome, not a failure mode.
|
||||
|
||||
## 5. Write-back
|
||||
|
||||
The Stop-hook session retain is **kept** — same trigger, same fail-open behavior. What changes is the transcript format handed to extraction:
|
||||
|
||||
- **JSON turns**, not markdown: an array of `{ "role": "user" | "assistant", "text": ... }` entries for the conversational content.
|
||||
- Tool calls collapse to **compact one-line action entries**: `{ "role": "action", "text": "Edit boltons/strutils.py" }` — tool name + primary target only, **no arguments, no outputs**.
|
||||
|
||||
Rationale: extraction keeps the concrete artifacts (which files were touched, what actions occurred) without the transcript noise of full tool payloads — the markdown tool-call dumps were volume without signal.
|
||||
|
||||
## 6. Knowledge pages
|
||||
|
||||
Simplified from the v2 spec:
|
||||
|
||||
- **Dropped: tags and `entity_labels`** (v2 spec §4–5). The server already re-synthesizes pages after consolidation, so pages stay "living" with no client-side routing machinery. The extractor-never-knows-about-pages principle now holds trivially — there is nothing to route.
|
||||
- **Creation paths:**
|
||||
1. **Seeded taxonomy** at bank creation (the fixed page set, as today, minus tag triggers).
|
||||
2. **Agent-driven `capture_initiative`** at plan approval — the one active capture verb survives from v2.
|
||||
3. **Organic splitting** of pages that outgrow their scope is a **server/curator concern**, not a client feature.
|
||||
|
||||
## 7. Ingestion — progressive background deepening
|
||||
|
||||
*Status: design accepted, implementation phased separately.*
|
||||
|
||||
Replaces the manual backfill CLI as the user-facing path (the CLI was out-of-band burden; nobody runs it). The principle: converge to full-depth history through normal usage, with zero user action.
|
||||
|
||||
1. **Instant shallow seed** — the gitlog seed from §3a; the session is useful immediately.
|
||||
2. **Background deepening** — a background worker deep-ingests **per-commit-with-diffs, incrementally**, never blocking a turn.
|
||||
3. **Working-set prioritization** — commits are ingested in order of relevance to what the agent is actually doing: files the agent reads/edits get their commit histories ingested **first**. Depth arrives where it pays off.
|
||||
4. **Checkpointing** — progress persists across sessions; each session resumes deepening where the last left off, converging to full depth over normal usage.
|
||||
|
||||
The **backfill CLI survives as an internal tool** (benchmark setup, CI bank preparation) — it is no longer a documented user path.
|
||||
|
||||
## 8. Gap analysis — v2 principles under this design
|
||||
|
||||
| v2 principle | How this design satisfies it |
|
||||
| --- | --- |
|
||||
| See-it-working (automatic, visible value) | Reflect answer visible from turn 1; page-section snippets appear with explicit `From <page> › <section>` provenance, so the user sees memory working — and the score floor keeps it from visibly misfiring. |
|
||||
| No out-of-band CLI | Cold-repo auto-seed kept (§3a); backfill CLI demoted to internal-only, replaced by background deepening (§7). Nothing requires a terminal command. |
|
||||
| Reuse-over-reinvent | Reflect, `listPages`, Stop-hook retain, `capture_initiative`, and the 10-turn refresh cadence are all existing machinery recombined; the only new code is the local section index and matcher — deliberately dumb (lexical, no LLM). |
|
||||
| Preserve-intent | Reflect is prompted for root-cause decisions with exact values; JSON transcripts keep concrete action artifacts; per-commit-with-diffs deepening captures *why* the code changed, not just that it did. |
|
||||
| Near-zero-burden | No config flags to choose, no CLI to run, no tags to maintain; one LLM call per session start, everything else local. |
|
||||
|
||||
## 9. Verification gates
|
||||
|
||||
Ship gates, in order:
|
||||
|
||||
1. **Reflect-restored benchmark:** the restored runtime must recover **~25 mean corrections at n=2 on identical banks** to the original reflect run. This proves the restoration is faithful before anything is layered on.
|
||||
2. **Reflect+pages benchmark:** with per-turn section injection enabled, the score **must not regress** vs reflect-alone. Section injection earns its place by not hurting; any regression points at the floor/budget tuning.
|
||||
3. **Live system suite:** existing hook/integration suite updated for the new path — reflect caching + per-turn re-injection, section index build/refresh, score-floor silence, JSON transcript shape, action-entry compaction. Deterministic pieces (matcher scoring, budget trim, provenance formatting, transcript serialization) as fast unit tests.
|
||||
|
||||
## 10. Non-goals / deferred
|
||||
|
||||
- Any per-turn LLM or server call for injection (explicitly excluded — the local matcher is the whole point).
|
||||
- Semantic/embedding-based section matching (revisit only if lexical matching demonstrably misses; start dumb).
|
||||
- Client-side page splitting or curation (server/curator concern, §6).
|
||||
- Progressive-deepening implementation details (worker scheduling, checkpoint format) — phased separately per §7.
|
||||
|
||||
## 11. File map (anticipated)
|
||||
|
||||
- `src/core/reflect.ts` (restored) — session reflect call + per-session cache.
|
||||
- `src/core/section-index.ts` (new) — page → sections split, lexical scorer, budget trim, provenance formatting; pure/unit-testable.
|
||||
- `src/core/hook.ts` — drop recall; inject cached reflect + matched sections; index refresh on roster cadence.
|
||||
- `src/core/session-start.ts` — cold-repo seed (unchanged) + reflect trigger wiring + initial index build.
|
||||
- `src/core/transcript.ts` (new or reworked) — JSON turn serialization + action-entry compaction for the Stop hook.
|
||||
- `src/core/missions.ts` / `src/core/hindsight.ts` — remove `entity_labels` and tag-scoped page triggers; keep seeded taxonomy + `capture_initiative`.
|
||||
- `src/core/config.ts` — remove any behavior flags; keep env/naming/harness + `pageRefreshEveryTurns`.
|
||||
- Tests alongside each.
|
||||
@@ -2,8 +2,8 @@ apiVersion: v2
|
||||
name: hindsight
|
||||
description: Hindsight helm chart
|
||||
type: application
|
||||
version: 0.9.1
|
||||
appVersion: "0.9.1"
|
||||
version: 0.9.0
|
||||
appVersion: "0.9.0"
|
||||
keywords:
|
||||
- ai
|
||||
- memory
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@vectorize-io/hindsight-all",
|
||||
"version": "0.9.1",
|
||||
"version": "0.9.0",
|
||||
"description": "Node.js programmatic lifecycle manager for Hindsight — embeds a local hindsight daemon in a Node application. Pair with @vectorize-io/hindsight-client for memory operations.",
|
||||
"main": "dist/index.js",
|
||||
"types": "dist/index.d.ts",
|
||||
|
||||
@@ -4,15 +4,15 @@ build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "hindsight-all-slim"
|
||||
version = "0.9.1"
|
||||
version = "0.9.0"
|
||||
description = "Hindsight: Agent Memory That Works Like Human Memory - Slim All-in-One Bundle"
|
||||
license = "MIT"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.11"
|
||||
dependencies = [
|
||||
"hindsight-api-slim==0.9.1",
|
||||
"hindsight-api-slim==0.9.0",
|
||||
"hindsight-client>=0.0.7",
|
||||
"hindsight-embed==0.9.1",
|
||||
"hindsight-embed==0.9.0",
|
||||
]
|
||||
|
||||
[tool.uv.sources]
|
||||
|
||||
@@ -64,26 +64,15 @@ class HindsightEmbedded:
|
||||
- create_directive(), list_directives(), etc.
|
||||
- And all async variants (aretain, arecall, areflect, etc.)
|
||||
|
||||
Only the settings you pass explicitly are forwarded to the daemon. Anything
|
||||
left at its default is resolved by the daemon instead, in this order: the
|
||||
profile's .env file, then the parent process environment, then the daemon's
|
||||
own default. That is what lets a client constructed without credentials run
|
||||
against a profile (or a shell) that already has them configured, rather than
|
||||
overwriting them with placeholders (#3253).
|
||||
|
||||
Args:
|
||||
profile: Profile name for data isolation (default: "default")
|
||||
llm_provider: LLM provider ("groq", "openai", "ollama", "gemini", "anthropic",
|
||||
"lmstudio"). Omit to inherit; the server default is "openai".
|
||||
llm_api_key: API key for the LLM provider. Omit to inherit; pass "" to
|
||||
explicitly run without a key (local services that need no auth).
|
||||
llm_model: Model name to use. Omit to inherit; the server picks a default
|
||||
for the resolved provider.
|
||||
llm_provider: LLM provider ("groq", "openai", "ollama", "gemini", "anthropic", "lmstudio")
|
||||
llm_api_key: API key for the LLM provider
|
||||
llm_model: Model name to use
|
||||
llm_base_url: Optional custom base URL for LLM API
|
||||
database_url: Optional database URL override (default: profile-specific pg0)
|
||||
idle_timeout: Seconds before daemon auto-exits when idle. Omit to inherit
|
||||
(daemon default: 0, disabled).
|
||||
log_level: Daemon log level. Omit to inherit (daemon default: "info").
|
||||
idle_timeout: Seconds before daemon auto-exits when idle (default: 0, disabled)
|
||||
log_level: Daemon log level (default: "info")
|
||||
ui: Whether to start the control plane web UI alongside the daemon (default: False)
|
||||
ui_port: Port for the UI. Defaults to daemon_port + 10000.
|
||||
ui_hostname: Hostname to bind the UI to. Defaults to "0.0.0.0".
|
||||
@@ -92,13 +81,13 @@ class HindsightEmbedded:
|
||||
def __init__(
|
||||
self,
|
||||
profile: str = "default",
|
||||
llm_provider: Optional[str] = None,
|
||||
llm_api_key: Optional[str] = None,
|
||||
llm_model: Optional[str] = None,
|
||||
llm_provider: str = "groq",
|
||||
llm_api_key: str = "",
|
||||
llm_model: str = "openai/gpt-oss-120b",
|
||||
llm_base_url: Optional[str] = None,
|
||||
database_url: Optional[str] = None,
|
||||
idle_timeout: Optional[int] = None,
|
||||
log_level: Optional[str] = None,
|
||||
idle_timeout: int = 0,
|
||||
log_level: str = "info",
|
||||
ui: bool = False,
|
||||
ui_port: Optional[int] = None,
|
||||
ui_hostname: str = "0.0.0.0",
|
||||
@@ -106,50 +95,29 @@ class HindsightEmbedded:
|
||||
"""
|
||||
Initialize the embedded client (daemon starts on first use).
|
||||
|
||||
Every LLM/daemon setting left as None is omitted from the daemon config so
|
||||
the daemon resolves it from the profile .env, then the parent environment,
|
||||
then its own default.
|
||||
|
||||
Args:
|
||||
profile: Profile name for data isolation
|
||||
llm_provider: LLM provider. Omit to inherit.
|
||||
llm_api_key: API key for the LLM provider. Omit to inherit; pass "" to
|
||||
explicitly run without a key.
|
||||
llm_model: Model name to use. Omit to inherit.
|
||||
llm_provider: LLM provider
|
||||
llm_api_key: API key for the LLM provider
|
||||
llm_model: Model name to use
|
||||
llm_base_url: Optional custom base URL for LLM API
|
||||
database_url: Optional database URL override
|
||||
idle_timeout: Seconds before daemon auto-exits when idle (0 = disabled).
|
||||
Omit to inherit.
|
||||
log_level: Daemon log level. Omit to inherit.
|
||||
idle_timeout: Seconds before daemon auto-exits when idle (0 = disabled)
|
||||
log_level: Daemon log level
|
||||
ui: Whether to start the control plane web UI alongside the daemon
|
||||
ui_port: Port for the UI (defaults to daemon_port + 10000)
|
||||
ui_hostname: Hostname to bind the UI to (defaults to "0.0.0.0")
|
||||
"""
|
||||
self.profile = profile
|
||||
|
||||
# Build the config dict for the daemon (matches CLI format), omitting
|
||||
# every setting the caller did not specify. An omitted key is inherited
|
||||
# by the daemon from the profile .env / parent environment; sending a
|
||||
# placeholder instead would overwrite it, and _register_profile would
|
||||
# then persist that placeholder into the profile's .env file (#3253).
|
||||
# An explicit "" is still an override — that is how a local LLM service
|
||||
# with no authentication clears an inherited API key.
|
||||
self.config: dict[str, str] = {}
|
||||
|
||||
if llm_provider is not None:
|
||||
self.config["HINDSIGHT_API_LLM_PROVIDER"] = llm_provider
|
||||
|
||||
if llm_api_key is not None:
|
||||
self.config["HINDSIGHT_API_LLM_API_KEY"] = llm_api_key
|
||||
|
||||
if llm_model is not None:
|
||||
self.config["HINDSIGHT_API_LLM_MODEL"] = llm_model
|
||||
|
||||
if log_level is not None:
|
||||
self.config["HINDSIGHT_API_LOG_LEVEL"] = log_level
|
||||
|
||||
if idle_timeout is not None:
|
||||
self.config["HINDSIGHT_EMBED_DAEMON_IDLE_TIMEOUT"] = str(idle_timeout)
|
||||
# Build config dict for daemon (matches CLI format)
|
||||
self.config = {
|
||||
"HINDSIGHT_API_LLM_PROVIDER": llm_provider,
|
||||
"HINDSIGHT_API_LLM_API_KEY": llm_api_key,
|
||||
"HINDSIGHT_API_LLM_MODEL": llm_model,
|
||||
"HINDSIGHT_API_LOG_LEVEL": log_level,
|
||||
"HINDSIGHT_EMBED_DAEMON_IDLE_TIMEOUT": str(idle_timeout),
|
||||
}
|
||||
|
||||
if llm_base_url:
|
||||
self.config["HINDSIGHT_API_LLM_BASE_URL"] = llm_base_url
|
||||
|
||||
@@ -4,15 +4,15 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "hindsight-all"
|
||||
version = "0.9.1"
|
||||
version = "0.9.0"
|
||||
description = "Hindsight: Agent Memory That Works Like Human Memory - All-in-One Bundle"
|
||||
license = "MIT"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.11"
|
||||
dependencies = [
|
||||
"hindsight-api-slim[all]==0.9.1",
|
||||
"hindsight-api-slim[all]==0.9.0",
|
||||
"hindsight-client>=0.0.7",
|
||||
"hindsight-embed==0.9.1",
|
||||
"hindsight-embed==0.9.0",
|
||||
]
|
||||
|
||||
[tool.uv.sources]
|
||||
@@ -22,7 +22,7 @@ hindsight-embed = { workspace = true }
|
||||
|
||||
[project.optional-dependencies]
|
||||
local-llm = [
|
||||
"hindsight-api-slim[local-llm]==0.9.1",
|
||||
"hindsight-api-slim[local-llm]==0.9.0",
|
||||
]
|
||||
test = [
|
||||
"pytest>=7.0.0",
|
||||
|
||||
@@ -1,181 +0,0 @@
|
||||
"""Configuration forwarding rules for HindsightEmbedded.
|
||||
|
||||
Regression coverage for #3253: a setting the caller does not pass must be left
|
||||
out of the daemon config, so the daemon can resolve it from the profile's .env
|
||||
file or the parent environment instead of receiving a client-side placeholder
|
||||
that overwrites it — and that the daemon then persists back into the profile.
|
||||
"""
|
||||
|
||||
import json
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight import HindsightEmbedded
|
||||
from hindsight_embed.daemon_embed_manager import DaemonEmbedManager
|
||||
|
||||
LLM_PROVIDER = "HINDSIGHT_API_LLM_PROVIDER"
|
||||
LLM_API_KEY = "HINDSIGHT_API_LLM_API_KEY"
|
||||
LLM_MODEL = "HINDSIGHT_API_LLM_MODEL"
|
||||
LOG_LEVEL = "HINDSIGHT_API_LOG_LEVEL"
|
||||
IDLE_TIMEOUT = "HINDSIGHT_EMBED_DAEMON_IDLE_TIMEOUT"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def temp_home(tmp_path, monkeypatch):
|
||||
"""Isolate HOME so profile .env files never touch the real user profile.
|
||||
|
||||
USERPROFILE is set as well because Path.home() consults it on Windows.
|
||||
"""
|
||||
home = tmp_path / "home"
|
||||
home.mkdir()
|
||||
monkeypatch.setenv("HOME", str(home))
|
||||
monkeypatch.setenv("USERPROFILE", str(home))
|
||||
return home
|
||||
|
||||
|
||||
def _write_profile(home, name, port, env_contents=None):
|
||||
"""Create a registered profile, optionally with a pre-populated .env file."""
|
||||
profile_dir = home / ".hindsight" / "profiles"
|
||||
profile_dir.mkdir(parents=True, exist_ok=True)
|
||||
(profile_dir / "metadata.json").write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"version": 1,
|
||||
"profiles": {
|
||||
name: {
|
||||
"port": port,
|
||||
"created_at": "2024-01-01T00:00:00+00:00",
|
||||
"last_used": "2024-01-01T00:00:00+00:00",
|
||||
}
|
||||
},
|
||||
}
|
||||
)
|
||||
)
|
||||
env_path = profile_dir / f"{name}.env"
|
||||
if env_contents is not None:
|
||||
env_path.write_text(env_contents)
|
||||
return env_path
|
||||
|
||||
|
||||
def _daemon_env(client):
|
||||
"""Run the real daemon start path with Popen stubbed, returning the child env.
|
||||
|
||||
Asserting on client.config alone would not catch a regression in how the
|
||||
embed manager merges that config with the profile and the parent
|
||||
environment, which is where the reported bug actually surfaced.
|
||||
"""
|
||||
manager = DaemonEmbedManager()
|
||||
captured: dict[str, dict[str, str]] = {}
|
||||
spawned = [False]
|
||||
|
||||
def fake_popen(cmd, env, **kwargs):
|
||||
captured["env"] = env
|
||||
spawned[0] = True
|
||||
process = MagicMock()
|
||||
process.pid = 12345
|
||||
return process
|
||||
|
||||
with (
|
||||
patch("hindsight_embed.daemon_embed_manager.subprocess.Popen", side_effect=fake_popen),
|
||||
patch("hindsight_embed.daemon_embed_manager.time.sleep"),
|
||||
patch.object(manager, "_clear_port", return_value=True),
|
||||
patch.object(manager, "_find_api_command", return_value=["hindsight-api"]),
|
||||
patch.object(manager, "is_running", side_effect=lambda profile="": spawned[0]),
|
||||
patch("hindsight_embed.daemon_embed_manager.platform.system", return_value="Linux"),
|
||||
):
|
||||
assert manager.ensure_running(client.config, client.profile)
|
||||
|
||||
return captured["env"]
|
||||
|
||||
|
||||
def test_nothing_is_forwarded_when_nothing_is_specified(temp_home):
|
||||
assert HindsightEmbedded(profile="test").config == {}
|
||||
|
||||
|
||||
def test_explicitly_passed_settings_are_forwarded(temp_home):
|
||||
client = HindsightEmbedded(
|
||||
profile="test",
|
||||
llm_provider="openai",
|
||||
llm_api_key="sk-real",
|
||||
llm_model="gpt-4o-mini",
|
||||
log_level="debug",
|
||||
idle_timeout=300,
|
||||
)
|
||||
|
||||
assert client.config == {
|
||||
LLM_PROVIDER: "openai",
|
||||
LLM_API_KEY: "sk-real",
|
||||
LLM_MODEL: "gpt-4o-mini",
|
||||
LOG_LEVEL: "debug",
|
||||
IDLE_TIMEOUT: "300",
|
||||
}
|
||||
|
||||
|
||||
def test_empty_api_key_is_forwarded_as_an_override(temp_home):
|
||||
"""An empty string is an explicit choice, not an omission.
|
||||
|
||||
Local LLM services that need no authentication rely on it to clear a key
|
||||
inherited from the environment.
|
||||
"""
|
||||
assert HindsightEmbedded(profile="test", llm_api_key="").config[LLM_API_KEY] == ""
|
||||
|
||||
|
||||
def test_idle_timeout_zero_is_forwarded(temp_home):
|
||||
"""0 is falsy but meaningful ("never auto-exit"), so it must survive."""
|
||||
assert HindsightEmbedded(profile="test", idle_timeout=0).config[IDLE_TIMEOUT] == "0"
|
||||
|
||||
|
||||
def test_omitted_key_inherits_the_parent_environment(temp_home, monkeypatch):
|
||||
monkeypatch.setenv(LLM_API_KEY, "sk-parent")
|
||||
_write_profile(temp_home, "inherit-env", 9871)
|
||||
|
||||
env = _daemon_env(HindsightEmbedded(profile="inherit-env", llm_provider="openai"))
|
||||
|
||||
assert env[LLM_API_KEY] == "sk-parent"
|
||||
|
||||
|
||||
def test_omitted_settings_inherit_the_profile_env(temp_home, monkeypatch):
|
||||
for var in (LLM_PROVIDER, LLM_API_KEY, LLM_MODEL):
|
||||
monkeypatch.delenv(var, raising=False)
|
||||
env_path = _write_profile(
|
||||
temp_home,
|
||||
"prod",
|
||||
9872,
|
||||
"HINDSIGHT_API_LLM_PROVIDER=anthropic\n"
|
||||
"HINDSIGHT_API_LLM_MODEL=claude-sonnet-4-20250514\n"
|
||||
"HINDSIGHT_API_LLM_API_KEY=sk-ant-prod\n",
|
||||
)
|
||||
|
||||
env = _daemon_env(HindsightEmbedded(profile="prod"))
|
||||
|
||||
assert env[LLM_PROVIDER] == "anthropic"
|
||||
assert env[LLM_MODEL] == "claude-sonnet-4-20250514"
|
||||
assert env[LLM_API_KEY] == "sk-ant-prod"
|
||||
|
||||
# A successful start rewrites the profile's .env; it must not come back with
|
||||
# client-side placeholders in place of the configured values.
|
||||
persisted = env_path.read_text()
|
||||
assert "HINDSIGHT_API_LLM_PROVIDER=anthropic" in persisted
|
||||
assert "HINDSIGHT_API_LLM_MODEL=claude-sonnet-4-20250514" in persisted
|
||||
assert "HINDSIGHT_API_LLM_API_KEY=sk-ant-prod" in persisted
|
||||
|
||||
|
||||
def test_explicit_empty_key_overrides_the_parent_environment(temp_home, monkeypatch):
|
||||
monkeypatch.setenv(LLM_API_KEY, "sk-parent")
|
||||
_write_profile(temp_home, "no-auth", 9873)
|
||||
|
||||
env = _daemon_env(
|
||||
HindsightEmbedded(profile="no-auth", llm_provider="lmstudio", llm_api_key="")
|
||||
)
|
||||
|
||||
assert env[LLM_API_KEY] == ""
|
||||
|
||||
|
||||
def test_explicit_settings_still_win_over_the_profile(temp_home, monkeypatch):
|
||||
monkeypatch.delenv(LLM_PROVIDER, raising=False)
|
||||
_write_profile(temp_home, "override", 9874, "HINDSIGHT_API_LLM_PROVIDER=anthropic\n")
|
||||
|
||||
env = _daemon_env(HindsightEmbedded(profile="override", llm_provider="openai"))
|
||||
|
||||
assert env[LLM_PROVIDER] == "openai"
|
||||
@@ -53,4 +53,4 @@ __all__ = [
|
||||
"RemoteTEICrossEncoder",
|
||||
"LLMConfig",
|
||||
]
|
||||
__version__ = "0.9.1"
|
||||
__version__ = "0.9.0"
|
||||
|
||||
@@ -47,64 +47,13 @@ _INDEX_TYPE_KEYWORDS = {
|
||||
"scann": "scann",
|
||||
}
|
||||
|
||||
# Ceiling on how many tuples one resumed ANN scan may visit (hnsw.max_scan_tuples).
|
||||
# Only iterative scans consult it, and it is approximate — the initial round is not
|
||||
# counted. pgvector defaults to 20000; this is deliberately lower.
|
||||
#
|
||||
# The filters that thin a semantic arm (the similarity floor, tags, date ranges) are
|
||||
# applied *after* the index scan, so a selective query resumes repeatedly to fill its
|
||||
# LIMIT. Unbounded, that turns the cheapest queries today into the most expensive:
|
||||
# ~20x the standing batch is enough to fill even a large recall budget on an
|
||||
# unfiltered query, and caps the pathological filtered case at a scan that returns
|
||||
# short — which is exactly what those queries did before iterative scans were on.
|
||||
# The GUCs that make a scan resumable — dropped wholesale when the operator turns the
|
||||
# behaviour off, so a connection is left exactly as it was before it existed (and a
|
||||
# pgvector too old to define them is never sent them either).
|
||||
_ITERATIVE_SCAN_GUCS = frozenset({"hnsw.iterative_scan", "hnsw.max_scan_tuples"})
|
||||
|
||||
|
||||
def iterative_scan_enabled() -> bool:
|
||||
"""Whether ANN scans may resume to satisfy a query's LIMIT.
|
||||
|
||||
Turning it off restores the previous depth exactly: a scan stops when its first
|
||||
candidate list drains, so no recall retrieves more rows than that list holds,
|
||||
whatever its budget.
|
||||
|
||||
Resolved through the config object rather than read from the environment, so a
|
||||
value set any other way — a CLI override applied with dataclasses.replace, a
|
||||
programmatically built config — is honoured, and the parsing and validation live
|
||||
in one place. Imported inside the function because config imports this module.
|
||||
"""
|
||||
from .config import get_config
|
||||
|
||||
return get_config().ann_iterative_scan
|
||||
|
||||
|
||||
def ann_max_scan_tuples() -> int:
|
||||
"""Ceiling on tuples one resumed scan may visit (hnsw.max_scan_tuples).
|
||||
|
||||
This is the knob that governs the cost of the behaviour. It bounds the CPU a
|
||||
selective query can spend resuming, and with it the scan's memory — pgvector
|
||||
otherwise caps that at ``work_mem * hnsw.scan_mem_multiplier``, but at this
|
||||
default the memory ceiling is never approached: squeezing work_mem to 256kB
|
||||
changes neither the rows returned nor the latency.
|
||||
|
||||
Approximate, and the initial scan is not counted, so even 1 leaves intact the
|
||||
depth a query had before scans could resume.
|
||||
"""
|
||||
from .config import get_config
|
||||
|
||||
return get_config().ann_max_scan_tuples
|
||||
|
||||
|
||||
# Per-backend ANN search-time tuning GUCs. Each entry is a tuple of
|
||||
# (guc_name, value) pairs the caller can apply with SET or SET LOCAL.
|
||||
#
|
||||
# - pgvector exposes hnsw.ef_search. The 60 / 200 pair is unchanged from the
|
||||
# pre-dispatcher code (internal benchmarks tuned around our embedding count
|
||||
# and recall floor; see the link_utils / pool init call sites for the
|
||||
# latency-vs-recall framing). With iterative scans on (below) the ef value is a
|
||||
# batch size rather than a ceiling, so a query's own LIMIT decides its depth.
|
||||
# latency-vs-recall framing).
|
||||
# - vchord exposes vchordrq.probes, but its shape must match the index's
|
||||
# build.internal.lists hierarchy. VectorChord 1.1 added per-index fallback
|
||||
# parameters for this reason: a session GUC overrides every vchordrq index,
|
||||
@@ -114,30 +63,11 @@ def ann_max_scan_tuples() -> int:
|
||||
# indexes should attach probes to the index storage parameters instead.
|
||||
# - pgvectorscale / pg_diskann / scann do not expose an equivalent per-statement
|
||||
# knob in the engine today, so the dispatcher returns no statements for them.
|
||||
#
|
||||
# hnsw.iterative_scan is what makes ef_search a *batch* size rather than a ceiling.
|
||||
# With it off (pgvector's default, and what Hindsight ran until now) the ground-layer
|
||||
# search runs once and the scan ends when its list drains, so a query could never get
|
||||
# more rows than ef_search however large its LIMIT — the recall budget moved the SQL
|
||||
# and nothing else. With it on, the scan resumes in ef_search-sized rounds until the
|
||||
# LIMIT is met, so each query gets the depth it asks for with no per-query setting.
|
||||
# strict_order, not relaxed_order: the arms are trimmed in Python on the assumption
|
||||
# that rows arrive ordered by distance.
|
||||
#
|
||||
# Retain-side link probing wants the opposite — it is tuned for latency, not depth,
|
||||
# and resuming past its small candidate list would defeat that — so the low-latency
|
||||
# profile pins it off. Both profiles set it explicitly rather than relying on the
|
||||
# server default, so neither depends on what the other last left on the connection.
|
||||
_ANN_TUNING_LOW_LATENCY: dict[str, tuple[tuple[str, str], ...]] = {
|
||||
"pgvector": (("hnsw.ef_search", "60"), ("hnsw.iterative_scan", "off")),
|
||||
"pgvector": (("hnsw.ef_search", "60"),),
|
||||
}
|
||||
_ANN_TUNING_HIGH_RECALL: dict[str, tuple[tuple[str, str], ...]] = {
|
||||
"pgvector": (
|
||||
("hnsw.ef_search", "200"),
|
||||
("hnsw.iterative_scan", "strict_order"),
|
||||
# Value filled in per call by ann_search_tuning_settings().
|
||||
("hnsw.max_scan_tuples", ""),
|
||||
),
|
||||
"pgvector": (("hnsw.ef_search", "200"),),
|
||||
}
|
||||
|
||||
_EXTENSION_INSTALL_SQL = {
|
||||
@@ -237,12 +167,7 @@ def ann_search_tuning_settings(ext: str, *, kind: str) -> tuple[tuple[str, str],
|
||||
table = _ANN_TUNING_HIGH_RECALL
|
||||
else:
|
||||
raise ValueError(f"Unknown ANN tuning kind: {kind!r}")
|
||||
settings = table.get(_normalize_resolved(ext), ())
|
||||
if not iterative_scan_enabled():
|
||||
return tuple(pair for pair in settings if pair[0] not in _ITERATIVE_SCAN_GUCS)
|
||||
return tuple(
|
||||
(name, str(ann_max_scan_tuples()) if name == "hnsw.max_scan_tuples" else value) for name, value in settings
|
||||
)
|
||||
return table.get(_normalize_resolved(ext), ())
|
||||
|
||||
|
||||
def uses_per_bank_vector_indexes(ext: str) -> bool:
|
||||
@@ -250,83 +175,6 @@ 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,12 +23,7 @@ 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 (
|
||||
BankIndexResult,
|
||||
drop_orphaned_bank_indexes,
|
||||
list_bank_ids,
|
||||
reconcile_bank_vector_indexes,
|
||||
)
|
||||
from ..engine.vector_index_health import SchemaVectorIndexResult, repair_vector_indexes
|
||||
from ..extensions import TenantExtension, load_extension
|
||||
from ..pg0 import parse_pg0_url, resolve_database_url
|
||||
|
||||
@@ -73,7 +68,6 @@ BACKUP_TABLES = [
|
||||
"audit_log",
|
||||
"llm_requests",
|
||||
"graph_maintenance_queue",
|
||||
"entity_maintenance_queue",
|
||||
]
|
||||
|
||||
MANIFEST_VERSION = "2"
|
||||
@@ -623,14 +617,11 @@ async def _run_repair_bank(
|
||||
schema: str | None,
|
||||
bank_id: str | None,
|
||||
dry_run: bool,
|
||||
) -> list[BankIndexResult]:
|
||||
) -> list[SchemaVectorIndexResult]:
|
||||
"""Reconcile per-(bank, fact_type) vector index coverage over a raw connection.
|
||||
|
||||
A single autocommit connection is used because ``CREATE INDEX CONCURRENTLY``
|
||||
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.
|
||||
(used by ``repair_vector_indexes``) cannot run inside a transaction block.
|
||||
"""
|
||||
schemas = [schema] if schema else await _resolve_schemas(base_schema)
|
||||
index_clause = _vector_index_clause()
|
||||
@@ -639,38 +630,13 @@ async def _run_repair_bank(
|
||||
assert index_clause is not None
|
||||
|
||||
conn = await _admin_connect(db_url)
|
||||
results: list[BankIndexResult] = []
|
||||
try:
|
||||
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)"
|
||||
)
|
||||
results = await repair_vector_indexes(conn, schemas, index_clause, dry_run=dry_run, bank_id=bank_id)
|
||||
for result in results:
|
||||
typer.echo(
|
||||
f" schema '{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"
|
||||
f" schema '{result.schema}': {result.banks_scanned} bank(s) scanned, "
|
||||
f"{result.already_present} present, {result.created} created, "
|
||||
f"{result.skipped} to-create (dry-run), {result.failed} failed"
|
||||
)
|
||||
return results
|
||||
finally:
|
||||
@@ -702,23 +668,17 @@ def repair_bank(
|
||||
help="Report what would be repaired without creating or dropping any index.",
|
||||
),
|
||||
):
|
||||
"""Reconcile per-(bank, fact_type) vector index coverage against the size threshold.
|
||||
"""Verify and repair a bank's per-(bank, fact_type) vector index coverage.
|
||||
|
||||
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.
|
||||
Per-bank partial vector indexes are created when a bank is first created
|
||||
(instant on an empty bank). Banks that arrive populated — via logical
|
||||
restore, a cross-version upgrade, or a vector-extension switch — never hit
|
||||
that path, so their recall silently falls back to a global index +
|
||||
post-filter (slower, under-returning). This command detects missing OR
|
||||
invalid coverage (an INVALID leftover or an index whose access method
|
||||
drifted counts as missing) and rebuilds it with CREATE INDEX CONCURRENTLY,
|
||||
so it never blocks the live fleet. Idempotent and safe to re-run — the
|
||||
escape hatch after a restore, upgrade, or backend switch.
|
||||
"""
|
||||
if bool(bank_id) == all_banks:
|
||||
typer.echo("Error: pass exactly one of --bank <id> or --all.", err=True)
|
||||
@@ -752,18 +712,15 @@ def repair_bank(
|
||||
)
|
||||
)
|
||||
|
||||
total_banks = len(results)
|
||||
total_banks = sum(r.banks_scanned for r in results)
|
||||
total_present = sum(r.already_present for r in results)
|
||||
total_created = sum(r.created for r in results)
|
||||
total_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, {total_dropped} dropped, "
|
||||
f"{total_skipped} to-create (dry-run), {total_would_drop} to-drop (dry-run), "
|
||||
f"{total_failed} failed"
|
||||
f"{total_present} already present, {total_created} created, "
|
||||
f"{total_skipped} to-create (dry-run), {total_failed} failed"
|
||||
)
|
||||
if total_failed:
|
||||
failed_names = [name for r in results for name in r.failed_indexes]
|
||||
|
||||
-118
@@ -1,118 +0,0 @@
|
||||
"""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
@@ -1,300 +0,0 @@
|
||||
"""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
@@ -1,95 +0,0 @@
|
||||
"""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
@@ -1,99 +0,0 @@
|
||||
"""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
@@ -1,72 +0,0 @@
|
||||
"""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,7 +21,6 @@ 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,
|
||||
@@ -599,14 +598,6 @@ 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):
|
||||
@@ -654,17 +645,6 @@ 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.",
|
||||
@@ -987,21 +967,17 @@ class ReflectRequest(BaseModel):
|
||||
)
|
||||
tags: list[str] | None = Field(
|
||||
default=None,
|
||||
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.",
|
||||
description="Filter memories by tags during reflection. If not specified, all memories are considered.",
|
||||
)
|
||||
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), or "
|
||||
"'exact' (set equality). Untagged directives remain global in every mode.",
|
||||
"'any_strict' (OR, excludes untagged), 'all_strict' (AND, excludes untagged).",
|
||||
)
|
||||
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: ...}. "
|
||||
"Mutually exclusive with tags.",
|
||||
"Each group is a leaf {tags, match} or compound {and: [...]}, {or: [...]}, {not: ...}.",
|
||||
)
|
||||
apply_all_directives: bool = Field(
|
||||
default=False,
|
||||
@@ -1293,7 +1269,7 @@ class BankListItem(BaseModel):
|
||||
|
||||
|
||||
class BankListResponse(BaseModel):
|
||||
"""Response model for listing banks, one page at a time."""
|
||||
"""Response model for listing all banks."""
|
||||
|
||||
model_config = ConfigDict(
|
||||
json_schema_extra={
|
||||
@@ -1310,18 +1286,12 @@ 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):
|
||||
@@ -1468,7 +1438,8 @@ class BankConfigUpdate(BaseModel):
|
||||
json_schema_extra={
|
||||
"example": {
|
||||
"updates": {
|
||||
"retain_extraction_mode": "custom",
|
||||
"llm_model": "claude-sonnet-4-5",
|
||||
"retain_extraction_mode": "verbose",
|
||||
"retain_custom_instructions": "Extract technical details carefully",
|
||||
}
|
||||
}
|
||||
@@ -1476,8 +1447,8 @@ class BankConfigUpdate(BaseModel):
|
||||
)
|
||||
|
||||
updates: dict[str, Any] = Field(
|
||||
description="Configuration overrides. Keys can be in Python field format (retain_extraction_mode) "
|
||||
"or environment variable format (HINDSIGHT_API_RETAIN_EXTRACTION_MODE). "
|
||||
description="Configuration overrides. Keys can be in Python field format (llm_provider) "
|
||||
"or environment variable format (HINDSIGHT_API_LLM_PROVIDER). "
|
||||
"Only hierarchical fields can be overridden per-bank."
|
||||
)
|
||||
|
||||
@@ -1490,10 +1461,12 @@ 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",
|
||||
},
|
||||
}
|
||||
@@ -1810,19 +1783,8 @@ class UpdateMemoryRequest(BaseModel):
|
||||
)
|
||||
entities: list[str] | None = Field(
|
||||
default=None,
|
||||
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.",
|
||||
description="Replace the fact's entities. Names are resolved/find-or-created "
|
||||
"the same way retain does; '[]' detaches all entities. Omit to leave unchanged.",
|
||||
)
|
||||
state: str | None = Field(
|
||||
default=None,
|
||||
@@ -2018,19 +1980,12 @@ 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_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 "
|
||||
"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 "
|
||||
"mental-model read can confirm."
|
||||
),
|
||||
)
|
||||
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."
|
||||
),
|
||||
)
|
||||
pending_consolidation: int = Field(default=0, description="Number of memories not yet processed into observations")
|
||||
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.",
|
||||
@@ -2138,9 +2093,6 @@ 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):
|
||||
@@ -2150,10 +2102,7 @@ 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="Directive execution scope. Empty means global; non-empty requires a matching reflect scope.",
|
||||
)
|
||||
tags: list[str] = FieldWithDefault(list, description="Tags for filtering")
|
||||
|
||||
|
||||
class UpdateDirectiveRequest(BaseModel):
|
||||
@@ -2324,26 +2273,7 @@ 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 = 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."
|
||||
),
|
||||
)
|
||||
last_refreshed_at: str | None = None
|
||||
created_at: str | None = None
|
||||
reflect_response: dict | None = Field(
|
||||
default=None,
|
||||
@@ -2353,9 +2283,9 @@ class MentalModelResponse(BaseModel):
|
||||
default=None,
|
||||
description=(
|
||||
"True when memories matching this mental model's tag/fact_type scope have been written "
|
||||
"since last_memory_seen_at. Exact, and costly to compute, so it is populated only by the "
|
||||
"since last_refreshed_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_memory_seen_at` against the bank's `last_memory_write_at` from GET /stats: "
|
||||
"each `last_refreshed_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."
|
||||
),
|
||||
)
|
||||
@@ -2365,9 +2295,6 @@ 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")
|
||||
|
||||
|
||||
# =========================================================================
|
||||
@@ -2399,15 +2326,6 @@ 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)
|
||||
|
||||
|
||||
@@ -2446,16 +2364,6 @@ 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):
|
||||
@@ -2524,7 +2432,6 @@ 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,
|
||||
)
|
||||
|
||||
|
||||
@@ -2994,21 +2901,17 @@ 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, limit=None, request_context=request_context)
|
||||
existing_by_id = {m["id"]: m for m in existing.items}
|
||||
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_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, limit=None, request_context=request_context
|
||||
bank_id=bank_id, active_only=False, request_context=request_context
|
||||
)
|
||||
existing_by_name = {d["name"]: d for d in existing_directives.items}
|
||||
existing_by_name = {d["name"]: d for d in existing_directives}
|
||||
|
||||
bank_writes: list[BankTemplateImportWrite] = []
|
||||
if config_updates:
|
||||
@@ -3052,10 +2955,9 @@ 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.items}
|
||||
provisioned_by_id = {item["id"]: item for item in provisioned}
|
||||
existing_by_id.update(
|
||||
{
|
||||
item_id: provisioned_by_id[item_id]
|
||||
@@ -3067,10 +2969,9 @@ 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.items}
|
||||
provisioned_by_name = {item["name"]: item for item in provisioned}
|
||||
existing_by_name.update(
|
||||
{name: provisioned_by_name[name] for name in projected_directive_names & provisioned_by_name.keys()}
|
||||
)
|
||||
@@ -3098,18 +2999,17 @@ 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, limit=None, request_context=request_context)
|
||||
existing_by_id = {model["id"]: model for model in existing.items}
|
||||
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_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.items}
|
||||
existing_by_name = {directive["name"]: directive for directive in existing_directives}
|
||||
|
||||
await _apply_bank_template_resources(
|
||||
memory,
|
||||
@@ -3267,15 +3167,6 @@ 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,
|
||||
@@ -3337,7 +3228,6 @@ 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,
|
||||
@@ -3426,7 +3316,7 @@ class OperationStatusResponse(BaseModel):
|
||||
"example": {
|
||||
"operation_id": "550e8400-e29b-41d4-a716-446655440000",
|
||||
"status": "completed",
|
||||
"operation_type": "refresh_mental_model",
|
||||
"operation_type": "refresh_mental_models",
|
||||
"created_at": "2024-01-15T10:30:00Z",
|
||||
"updated_at": "2024-01-15T10:31:30Z",
|
||||
"completed_at": "2024-01-15T10:31:30Z",
|
||||
@@ -3512,19 +3402,15 @@ class VersionResponse(BaseModel):
|
||||
model_config = ConfigDict(
|
||||
json_schema_extra={
|
||||
"example": {
|
||||
"api_version": "0.9.0",
|
||||
"api_version": "0.4.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,
|
||||
},
|
||||
}
|
||||
}
|
||||
@@ -4068,20 +3954,15 @@ 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(request: Request, authorization: str | None = Header(default=None)) -> RequestContext:
|
||||
def get_request_context(authorization: str | None = Header(default=None)) -> RequestContext:
|
||||
"""
|
||||
Extract request context from the Authorization header.
|
||||
Extract request context from 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:
|
||||
@@ -4089,8 +3970,7 @@ def _register_routes(app: FastAPI):
|
||||
api_key = authorization[7:].strip()
|
||||
else:
|
||||
api_key = authorization.strip()
|
||||
extra_headers = collect_passthrough_headers(request.headers.raw, get_config().extension_passthrough_headers)
|
||||
return RequestContext(api_key=api_key, extra_headers=extra_headers)
|
||||
return RequestContext(api_key=api_key)
|
||||
|
||||
def precheck_for(operation: PrecheckOperation):
|
||||
"""
|
||||
@@ -4162,19 +4042,6 @@ 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()
|
||||
@@ -4552,7 +4419,6 @@ 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,
|
||||
@@ -4771,7 +4637,6 @@ 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
|
||||
@@ -4966,26 +4831,16 @@ def _register_routes(app: FastAPI):
|
||||
@app.get(
|
||||
"/v1/default/banks",
|
||||
response_model=BankListResponse,
|
||||
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."
|
||||
),
|
||||
summary="List all memory banks",
|
||||
description="Get a list of all agents with their profiles",
|
||||
operation_id="list_banks",
|
||||
tags=["Banks"],
|
||||
)
|
||||
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."""
|
||||
async def api_list_banks(request_context: RequestContext = Depends(get_request_context)):
|
||||
"""Get list of all banks with their profiles."""
|
||||
try:
|
||||
data = await app.state.memory.list_banks(
|
||||
search_query=q, limit=limit, offset=offset, request_context=request_context
|
||||
)
|
||||
return BankListResponse(**data)
|
||||
banks = await app.state.memory.list_banks(request_context=request_context)
|
||||
return BankListResponse(banks=banks)
|
||||
except (AuthenticationError, HTTPException):
|
||||
raise
|
||||
except Exception as e:
|
||||
@@ -5289,7 +5144,7 @@ def _register_routes(app: FastAPI):
|
||||
):
|
||||
"""List mental models for a bank."""
|
||||
try:
|
||||
page = await app.state.memory.list_mental_models(
|
||||
mental_models = await app.state.memory.list_mental_models(
|
||||
bank_id=bank_id,
|
||||
tags=tags_filter,
|
||||
tags_match=tags_match,
|
||||
@@ -5298,12 +5153,7 @@ def _register_routes(app: FastAPI):
|
||||
offset=offset,
|
||||
request_context=request_context,
|
||||
)
|
||||
return MentalModelListResponse(
|
||||
items=[MentalModelResponse(**m) for m in page.items],
|
||||
total=page.total,
|
||||
limit=limit,
|
||||
offset=offset,
|
||||
)
|
||||
return MentalModelListResponse(items=[MentalModelResponse(**m) for m in mental_models])
|
||||
except OperationValidationError as e:
|
||||
raise HTTPException(status_code=e.status_code, detail=e.reason)
|
||||
except (AuthenticationError, HTTPException):
|
||||
@@ -5757,11 +5607,7 @@ def _register_routes(app: FastAPI):
|
||||
parent_id=body.parent_id,
|
||||
tags=body.tags if body.tags else None,
|
||||
max_tokens=body.max_tokens,
|
||||
# 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,
|
||||
trigger=body.trigger.model_dump() if body.trigger else None,
|
||||
request_context=request_context,
|
||||
)
|
||||
if node is None:
|
||||
@@ -5909,10 +5755,8 @@ 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`, `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.",
|
||||
"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.",
|
||||
operation_id="update_knowledge_node",
|
||||
tags=["Knowledge Base"],
|
||||
)
|
||||
@@ -5940,7 +5784,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", "trigger"} & body.model_fields_set
|
||||
page_fields = {"source_query", "tags", "max_tokens"} & body.model_fields_set
|
||||
if page_fields:
|
||||
did_change = True
|
||||
updated = await app.state.memory.update_knowledge_page(
|
||||
@@ -5949,10 +5793,6 @@ 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.
|
||||
@@ -5964,8 +5804,7 @@ def _register_routes(app: FastAPI):
|
||||
)
|
||||
if not did_change:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="Provide name, parent_id, source_query, tags, max_tokens, and/or trigger to update",
|
||||
status_code=400, detail="Provide name, parent_id, source_query, tags, and/or max_tokens to update"
|
||||
)
|
||||
if updated is None:
|
||||
raise HTTPException(status_code=404, detail=f"Knowledge node '{node_id}' not found")
|
||||
@@ -6022,21 +5861,14 @@ def _register_routes(app: FastAPI):
|
||||
"/v1/default/banks/{bank_id}/directives",
|
||||
response_model=DirectiveListResponse,
|
||||
summary="List directives",
|
||||
description="List directive definitions. Unlike reflect, an omitted tag filter returns all directives.",
|
||||
description="List hard rules that are injected into prompts.",
|
||||
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 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.",
|
||||
),
|
||||
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"),
|
||||
active_only: bool = Query(True, description="Only return active directives"),
|
||||
limit: int = Query(100, ge=1, le=1000),
|
||||
offset: int = Query(0, ge=0),
|
||||
@@ -6044,7 +5876,7 @@ def _register_routes(app: FastAPI):
|
||||
):
|
||||
"""List directives for a bank."""
|
||||
try:
|
||||
page = await app.state.memory.list_directives(
|
||||
directives = await app.state.memory.list_directives(
|
||||
bank_id=bank_id,
|
||||
tags=tags_filter,
|
||||
tags_match=tags_match,
|
||||
@@ -6053,12 +5885,7 @@ def _register_routes(app: FastAPI):
|
||||
offset=offset,
|
||||
request_context=request_context,
|
||||
)
|
||||
return DirectiveListResponse(
|
||||
items=[DirectiveResponse(**d) for d in page.items],
|
||||
total=page.total,
|
||||
limit=limit,
|
||||
offset=offset,
|
||||
)
|
||||
return DirectiveListResponse(items=[DirectiveResponse(**d) for d in directives])
|
||||
except OperationValidationError as e:
|
||||
raise HTTPException(status_code=e.status_code, detail=e.reason)
|
||||
except (AuthenticationError, HTTPException):
|
||||
@@ -6108,7 +5935,7 @@ def _register_routes(app: FastAPI):
|
||||
"/v1/default/banks/{bank_id}/directives",
|
||||
response_model=DirectiveResponse,
|
||||
summary="Create directive",
|
||||
description="Create a global or tag-scoped hard rule for reflect prompts.",
|
||||
description="Create a hard rule that will be injected into prompts.",
|
||||
operation_id="create_directive",
|
||||
tags=["Directives"],
|
||||
)
|
||||
@@ -7157,13 +6984,12 @@ 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 (limit=None — an export that stopped at the
|
||||
# default page size would silently drop the rest of the bank)
|
||||
# Get mental models
|
||||
mental_models_raw = await app.state.memory.list_mental_models(
|
||||
bank_id=bank_id, limit=None, request_context=request_context
|
||||
bank_id=bank_id, request_context=request_context
|
||||
)
|
||||
template_mental_models: list[BankTemplateMentalModel] = []
|
||||
for mm in mental_models_raw.items:
|
||||
for mm in mental_models_raw:
|
||||
trigger_data = mm.get("trigger", {})
|
||||
trigger = MentalModelTrigger(**trigger_data) if trigger_data else MentalModelTrigger()
|
||||
template_mental_models.append(
|
||||
@@ -7177,12 +7003,12 @@ def _register_routes(app: FastAPI):
|
||||
)
|
||||
)
|
||||
|
||||
# Get directives (limit=None for the same reason as the models above)
|
||||
# Get directives
|
||||
directives_raw = await app.state.memory.list_directives(
|
||||
bank_id=bank_id, active_only=False, limit=None, request_context=request_context
|
||||
bank_id=bank_id, active_only=False, request_context=request_context
|
||||
)
|
||||
template_directives: list[BankTemplateDirective] = []
|
||||
for d in directives_raw.items:
|
||||
for d in directives_raw:
|
||||
template_directives.append(
|
||||
BankTemplateDirective(
|
||||
name=d["name"],
|
||||
@@ -7580,9 +7406,8 @@ 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 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).",
|
||||
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).",
|
||||
operation_id="update_bank_config",
|
||||
tags=["Banks"],
|
||||
)
|
||||
@@ -8044,7 +7869,6 @@ 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,7 +9,6 @@ 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
|
||||
@@ -53,11 +52,6 @@ _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."""
|
||||
@@ -84,15 +78,6 @@ 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):
|
||||
@@ -154,13 +139,6 @@ def create_mcp_server(memory: MemoryEngine, multi_bank: bool = True) -> FastMCP:
|
||||
"update_bank",
|
||||
"delete_bank",
|
||||
"clear_memories",
|
||||
"get_knowledge_base_tree",
|
||||
"search_knowledge_base",
|
||||
"get_knowledge_page",
|
||||
"create_knowledge_folder",
|
||||
"create_knowledge_page",
|
||||
"update_knowledge_node",
|
||||
"delete_knowledge_node",
|
||||
}
|
||||
)
|
||||
base_tools: frozenset[str] | None = None if multi_bank else _SINGLE_BANK_TOOLS
|
||||
@@ -181,7 +159,6 @@ 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,
|
||||
@@ -401,16 +378,6 @@ 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)
|
||||
@@ -445,11 +412,6 @@ 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
|
||||
@@ -469,7 +431,7 @@ class MCPMiddleware:
|
||||
else:
|
||||
# Use TenantExtension.authenticate_mcp() for auth
|
||||
try:
|
||||
auth_context = RequestContext(api_key=auth_token, extra_headers=dict(passthrough_headers))
|
||||
auth_context = RequestContext(api_key=auth_token)
|
||||
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
|
||||
@@ -521,8 +483,6 @@ 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
|
||||
@@ -568,7 +528,6 @@ 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)
|
||||
|
||||
|
||||
@@ -1,56 +0,0 @@
|
||||
"""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 get the field) and resolves to
|
||||
# header, native OpenAI / openai.com / Azure OpenAI 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,12 +377,6 @@ 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"
|
||||
@@ -511,7 +505,6 @@ 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"
|
||||
@@ -533,8 +526,6 @@ ENV_RERANKER_GOOGLE_PROJECT_ID = "HINDSIGHT_API_RERANKER_GOOGLE_PROJECT_ID"
|
||||
ENV_RERANKER_GOOGLE_SERVICE_ACCOUNT_KEY = "HINDSIGHT_API_RERANKER_GOOGLE_SERVICE_ACCOUNT_KEY"
|
||||
|
||||
ENV_VECTOR_EXTENSION = "HINDSIGHT_API_VECTOR_EXTENSION"
|
||||
ENV_ANN_ITERATIVE_SCAN = "HINDSIGHT_API_ANN_ITERATIVE_SCAN"
|
||||
ENV_ANN_MAX_SCAN_TUPLES = "HINDSIGHT_API_ANN_MAX_SCAN_TUPLES"
|
||||
ENV_TEXT_SEARCH_EXTENSION = "HINDSIGHT_API_TEXT_SEARCH_EXTENSION"
|
||||
ENV_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE = "HINDSIGHT_API_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE"
|
||||
ENV_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER = "HINDSIGHT_API_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER"
|
||||
@@ -566,12 +557,6 @@ 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"
|
||||
@@ -721,7 +706,6 @@ 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"
|
||||
|
||||
@@ -757,7 +741,6 @@ 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,
|
||||
}
|
||||
@@ -809,7 +792,6 @@ 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"
|
||||
@@ -837,7 +819,6 @@ 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.
|
||||
@@ -876,10 +857,6 @@ 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"
|
||||
@@ -955,6 +932,7 @@ 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
|
||||
@@ -977,10 +955,6 @@ 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"
|
||||
@@ -1030,18 +1004,9 @@ 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 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
|
||||
# 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
|
||||
# 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.
|
||||
@@ -1107,11 +1072,6 @@ 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"
|
||||
@@ -1146,22 +1106,6 @@ DEFAULT_RERANKER_GOOGLE_MODEL = "semantic-ranker-default-004"
|
||||
|
||||
# Vector extension (pgvector, vchord, pgvectorscale, or AlloyDB ScaNN)
|
||||
DEFAULT_VECTOR_EXTENSION = "pgvector" # Options: "pgvector", "vchord", "pgvectorscale", "scann"
|
||||
# Let an ANN scan resume until the query's LIMIT is met, instead of stopping when its
|
||||
# first candidate list drains. Off, a recall can never retrieve more rows than the
|
||||
# candidate list holds (pgvector: hnsw.ef_search, 200), so a larger recall budget
|
||||
# widens the SQL and changes nothing. On is the intended behaviour; this exists as an
|
||||
# operational kill switch, because turning it off restores exactly the previous
|
||||
# retrieval depth without a deploy.
|
||||
DEFAULT_ANN_ITERATIVE_SCAN = True
|
||||
# Ceiling on how many tuples one resumed scan may visit. Bounds both the CPU a
|
||||
# selective query can spend resuming (the filters that thin a result are applied after
|
||||
# the index scan, so a selective one resumes repeatedly) and the scan's memory, which
|
||||
# pgvector otherwise caps at work_mem * hnsw.scan_mem_multiplier. Measured at this
|
||||
# value the memory ceiling is never approached — squeezing work_mem to 256kB changes
|
||||
# neither rows nor latency — so this is the knob that governs the cost, not work_mem.
|
||||
# Lower it to trade retrieval depth back for latency; the initial scan is not counted,
|
||||
# so even 1 leaves the pre-existing behaviour intact. pgvector's own default is 20000.
|
||||
DEFAULT_ANN_MAX_SCAN_TUPLES = 4000
|
||||
|
||||
# Text search extension (native PostgreSQL, vchord BM25, Timescale pg_textsearch,
|
||||
# pgroonga, or ParadeDB pg_search)
|
||||
@@ -1332,25 +1276,6 @@ 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
|
||||
@@ -1398,14 +1323,6 @@ 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
|
||||
@@ -1479,55 +1396,7 @@ 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.
|
||||
#
|
||||
# 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_MENTAL_MODEL_REFRESH_TICK_SECONDS = 60
|
||||
|
||||
# Default MCP tool descriptions (can be customized via env vars)
|
||||
DEFAULT_MCP_RETAIN_DESCRIPTION = """Store important information to long-term memory.
|
||||
@@ -2002,7 +1871,6 @@ 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
|
||||
@@ -2139,7 +2007,6 @@ 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),
|
||||
@@ -2208,8 +2075,6 @@ class HindsightConfig:
|
||||
migration_database_url: str | None
|
||||
database_schema: str
|
||||
vector_extension: str # "pgvector", "vchord", "pgvectorscale", or "scann"
|
||||
ann_iterative_scan: bool
|
||||
ann_max_scan_tuples: int
|
||||
text_search_extension: str # "native", "vchord", "pg_textsearch", "pgroonga", or "pg_search"
|
||||
# PostgreSQL text search dictionary for the "native" backend (ignored by
|
||||
# other backends). Only the "native" backend reads this field; pgroonga
|
||||
@@ -2240,10 +2105,7 @@ class HindsightConfig:
|
||||
llm_initial_backoff: float
|
||||
llm_max_backoff: float
|
||||
llm_timeout: float
|
||||
# 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_reasoning_effort: str
|
||||
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"
|
||||
@@ -2634,7 +2496,6 @@ 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
|
||||
@@ -2659,7 +2520,6 @@ 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
|
||||
@@ -2702,9 +2562,6 @@ 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)
|
||||
@@ -2716,8 +2573,6 @@ 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
|
||||
@@ -2752,7 +2607,6 @@ 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
|
||||
@@ -2760,20 +2614,6 @@ 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
|
||||
@@ -2937,11 +2777,6 @@ 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,
|
||||
@@ -3027,12 +2862,6 @@ class HindsightConfig:
|
||||
# Validate vector_extension
|
||||
validate_extension(self.vector_extension)
|
||||
|
||||
if self.ann_iterative_scan and self.ann_max_scan_tuples < 1:
|
||||
raise ValueError(
|
||||
f"Invalid ann_max_scan_tuples: {self.ann_max_scan_tuples}. Must be >= 1 when "
|
||||
f"iterative ANN scans are enabled (set {ENV_ANN_ITERATIVE_SCAN}=false to disable them)"
|
||||
)
|
||||
|
||||
# pg_trgm requires the similarity threshold in (0, 1]. Fail fast here
|
||||
# rather than let an out-of-range value raise on every pool connection's
|
||||
# setup (which would leave the API unable to serve any request).
|
||||
@@ -3201,10 +3030,6 @@ class HindsightConfig:
|
||||
migration_database_url=os.getenv(ENV_MIGRATION_DATABASE_URL) or None,
|
||||
database_schema=os.getenv(ENV_DATABASE_SCHEMA, DEFAULT_DATABASE_SCHEMA),
|
||||
vector_extension=os.getenv(ENV_VECTOR_EXTENSION, DEFAULT_VECTOR_EXTENSION).lower(),
|
||||
ann_iterative_scan=_parse_boolean_env(ENV_ANN_ITERATIVE_SCAN, DEFAULT_ANN_ITERATIVE_SCAN),
|
||||
ann_max_scan_tuples=_parse_non_negative_int(
|
||||
ENV_ANN_MAX_SCAN_TUPLES, os.getenv(ENV_ANN_MAX_SCAN_TUPLES), DEFAULT_ANN_MAX_SCAN_TUPLES
|
||||
),
|
||||
text_search_extension=os.getenv(ENV_TEXT_SEARCH_EXTENSION, DEFAULT_TEXT_SEARCH_EXTENSION).lower(),
|
||||
text_search_extension_native_language=os.getenv(
|
||||
ENV_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE,
|
||||
@@ -3230,7 +3055,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) or None,
|
||||
llm_reasoning_effort=os.getenv(ENV_LLM_REASONING_EFFORT, DEFAULT_LLM_REASONING_EFFORT),
|
||||
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,
|
||||
@@ -3454,8 +3279,6 @@ 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),
|
||||
@@ -3605,7 +3428,6 @@ 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))
|
||||
),
|
||||
@@ -3913,9 +3735,6 @@ 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))
|
||||
),
|
||||
@@ -3966,11 +3785,6 @@ 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()
|
||||
@@ -4072,26 +3886,6 @@ 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,
|
||||
@@ -4110,12 +3904,6 @@ 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,13 +42,7 @@ 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.
|
||||
# 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",)
|
||||
_OPENAI_DOMAINS = ("openai.com", "openai.azure.com")
|
||||
|
||||
|
||||
class CacheAffinityMode(StrEnum):
|
||||
@@ -93,7 +87,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 host, and
|
||||
native OpenAI (no base URL) or an openai.com / Azure OpenAI 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,17 +114,19 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
|
||||
return NO_TEMPORAL_CONSTRAINT
|
||||
return constraint(start, end)
|
||||
|
||||
def subtract_months(months: int) -> datetime | None:
|
||||
return add_months(reference_date, -months)
|
||||
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 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 | None:
|
||||
def add_months(base_date: datetime, months: int) -> datetime:
|
||||
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)
|
||||
@@ -371,7 +373,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 | None:
|
||||
def relative_month_start(period: str | None) -> datetime:
|
||||
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:
|
||||
@@ -416,8 +418,6 @@ 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 add_days(reference_date, direction * amount)
|
||||
return reference_date + timedelta(days=direction * amount)
|
||||
if unit in ("周", "星期", "礼拜"):
|
||||
return add_days(reference_date, direction * amount * 7)
|
||||
return reference_date + timedelta(weeks=direction * amount)
|
||||
if unit == "月":
|
||||
return add_months(reference_date, direction * amount)
|
||||
return add_years(reference_date, direction * amount)
|
||||
@@ -616,8 +616,6 @@ 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))
|
||||
@@ -841,7 +839,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 safe_since_constraint(relative_month_start(month_since_match.group(1)))
|
||||
return 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}"
|
||||
@@ -1037,7 +1035,7 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
|
||||
|
||||
if chinese_search(r"一年半前"):
|
||||
d = subtract_months(18)
|
||||
return safe_constraint(d, d)
|
||||
return constraint(d, d)
|
||||
|
||||
if chinese_search(r"([一二两三四五六七八九十]+)年半前"):
|
||||
match = chinese_search(r"([一二两三四五六七八九十]+)年半前")
|
||||
@@ -1045,15 +1043,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 safe_constraint(d, d)
|
||||
return 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 = add_days(subtract_months(months), -15)
|
||||
return safe_constraint(d, d)
|
||||
d = subtract_months(months) - timedelta(days=15)
|
||||
return constraint(d, d)
|
||||
|
||||
if chinese_search(r"半个?月前"):
|
||||
d = reference_date - timedelta(days=15)
|
||||
@@ -1061,7 +1059,7 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
|
||||
|
||||
if chinese_search(r"半年前"):
|
||||
d = subtract_months(6)
|
||||
return safe_constraint(d, d)
|
||||
return constraint(d, d)
|
||||
|
||||
future_year_half_match = chinese_search(
|
||||
rf"([0-9]+|[{_CHINESE_NUMERAL_CHARS}]+)年半{chinese_relative_future_suffix_pattern}"
|
||||
@@ -1070,7 +1068,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 safe_constraint(d, d)
|
||||
return constraint(d, d)
|
||||
|
||||
future_half_month_match = chinese_search(
|
||||
rf"([0-9]+|[{_CHINESE_NUMERAL_CHARS}]+)个?半月{chinese_relative_future_suffix_pattern}"
|
||||
@@ -1078,8 +1076,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_days(add_months(reference_date, months), 15)
|
||||
return safe_constraint(d, d)
|
||||
d = add_months(reference_date, months) + timedelta(days=15)
|
||||
return constraint(d, d)
|
||||
|
||||
if chinese_search(rf"半个?月{chinese_relative_future_suffix_pattern}"):
|
||||
d = reference_date + timedelta(days=15)
|
||||
@@ -1087,7 +1085,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 safe_constraint(d, d)
|
||||
return constraint(d, d)
|
||||
|
||||
adjacent_fuzzy_future_match = chinese_search(
|
||||
r"(?<![一二三四五六七八九十百千万零\d后])"
|
||||
@@ -1234,14 +1232,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 safe_constraint(subtract_months(6), reference_date)
|
||||
return 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 safe_constraint(subtract_months(6), reference_date)
|
||||
return constraint(subtract_months(6), reference_date)
|
||||
|
||||
within_count_match = chinese_search(
|
||||
rf"([0-9]+|[{_CHINESE_NUMERAL_CHARS}]+)(个?)(天|日|周|星期|礼拜|月|年)(?:以内|之内|内)"
|
||||
@@ -1304,7 +1302,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 safe_constraint(reference_date, add_months(reference_date, 6))
|
||||
return 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}"
|
||||
@@ -1441,8 +1439,6 @@ 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(
|
||||
@@ -1451,8 +1447,6 @@ 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))
|
||||
)
|
||||
@@ -1471,10 +1465,7 @@ 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)
|
||||
if start is None:
|
||||
return NO_TEMPORAL_CONSTRAINT
|
||||
start = start.replace(day=1)
|
||||
start = subtract_months(2).replace(day=1)
|
||||
return since_from_period(
|
||||
month_phase_period(start.year, start.month, second_previous_month_phase_since_match.group(2))
|
||||
)
|
||||
@@ -1599,8 +1590,6 @@ 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(
|
||||
@@ -1608,8 +1597,6 @@ 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(
|
||||
@@ -1624,10 +1611,7 @@ 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)
|
||||
if start is None:
|
||||
return NO_TEMPORAL_CONSTRAINT
|
||||
start = start.replace(day=1)
|
||||
start = subtract_months(2).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(
|
||||
@@ -1746,14 +1730,10 @@ 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})"):
|
||||
@@ -1783,10 +1763,7 @@ 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)
|
||||
if start is None:
|
||||
return NO_TEMPORAL_CONSTRAINT
|
||||
start = start.replace(day=1)
|
||||
start = subtract_months(2).replace(day=1)
|
||||
return constraint(start, month_end(start.year, start.month))
|
||||
|
||||
if chinese_search(rf"前一个?(周|星期|礼拜)(?!{chinese_boundary_suffix_pattern})"):
|
||||
|
||||
@@ -55,51 +55,12 @@ 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.
|
||||
@@ -222,52 +183,6 @@ 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.
|
||||
@@ -309,7 +224,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 ..memories import get_memories
|
||||
from ..search.retrieval import retrieve_semantic_bm25_combined
|
||||
|
||||
threshold = config.consolidation_dedup_threshold
|
||||
if anchor_emb_str is None:
|
||||
@@ -318,19 +233,10 @@ 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"
|
||||
# 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,
|
||||
)
|
||||
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
|
||||
)
|
||||
results = grouped["observation"].semantic
|
||||
best_id: str | None = None
|
||||
best_text = ""
|
||||
@@ -369,7 +275,6 @@ 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).
|
||||
@@ -378,10 +283,6 @@ 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.
|
||||
"""
|
||||
@@ -414,10 +315,6 @@ 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
|
||||
@@ -426,10 +323,6 @@ 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.
|
||||
@@ -442,15 +335,7 @@ 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,
|
||||
source_bounds,
|
||||
txn=txn,
|
||||
store, conn, memory_engine, bank_id, outcome.best_id, outcome.merged_text, live_source_ids, txn=txn
|
||||
)
|
||||
return outcome.best_id
|
||||
|
||||
@@ -494,8 +379,7 @@ 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 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).
|
||||
# visibility. Temporal fields follow the surviving twin (minimal scope; matches create).
|
||||
# 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()
|
||||
@@ -538,10 +422,6 @@ 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
|
||||
@@ -567,15 +447,7 @@ 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,
|
||||
_TemporalBounds.of(updated_obs[0]),
|
||||
txn=txn,
|
||||
store, conn, memory_engine, bank_id, outcome.best_id, outcome.merged_text, live_u_sources, txn=txn
|
||||
)
|
||||
await _execute_delete_action(conn, bank_id, updated_id, txn=txn)
|
||||
logger.info(
|
||||
@@ -1082,22 +954,17 @@ 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).
|
||||
|
||||
``add_bounds`` are the folded-in side's dates, widened onto the twin exactly as the SQL
|
||||
path's LEAST/GREATEST does."""
|
||||
place instead)."""
|
||||
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,
|
||||
@@ -1111,10 +978,10 @@ async def _reconcile_merge_via_store(
|
||||
tags=list(cur.tags or []),
|
||||
proof_count=len(merged_sources),
|
||||
source_memory_ids=merged_sources,
|
||||
event_date=merged_bounds.event_date,
|
||||
occurred_start=merged_bounds.occurred_start,
|
||||
occurred_end=merged_bounds.occurred_end,
|
||||
mentioned_at=merged_bounds.mentioned_at,
|
||||
event_date=cur.event_date,
|
||||
occurred_start=cur.occurred_start,
|
||||
occurred_end=cur.occurred_end,
|
||||
mentioned_at=cur.mentioned_at,
|
||||
created_at=cur.created_at,
|
||||
),
|
||||
)
|
||||
@@ -1183,65 +1050,12 @@ 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.
|
||||
@@ -1256,9 +1070,6 @@ 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
|
||||
@@ -1278,14 +1089,7 @@ 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,
|
||||
pending_refresh_tags,
|
||||
memory_engine, bank_id, request_context, config, llm_config, operation_id, observation_scopes
|
||||
)
|
||||
finally:
|
||||
if trace_token is not None:
|
||||
@@ -1303,7 +1107,6 @@ 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)
|
||||
@@ -1484,60 +1287,21 @@ async def _run_consolidation_job(
|
||||
_txn_provider = get_memories()
|
||||
_batch_txn = await _txn_provider.mint_txn(bank_id=bank_id, mutating=True)
|
||||
|
||||
try:
|
||||
pending: list[list[dict[str, Any]]] = [llm_batch_local]
|
||||
while pending:
|
||||
sub_batch = pending.pop(0)
|
||||
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(
|
||||
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(
|
||||
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,
|
||||
@@ -1546,90 +1310,96 @@ async def _run_consolidation_job(
|
||||
request_context=request_context,
|
||||
perf=batch_perf,
|
||||
config=config,
|
||||
obs_tags_override=obs_tags,
|
||||
txn=_batch_txn,
|
||||
)
|
||||
|
||||
all_deleted += sub_deleted
|
||||
|
||||
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} failed to abort write-group for"
|
||||
f" llm_batch #{batch_num_local}; recovery sweep will resolve it",
|
||||
exc_info=True,
|
||||
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,
|
||||
)
|
||||
raise
|
||||
|
||||
all_deleted += sub_deleted
|
||||
|
||||
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)
|
||||
# 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)
|
||||
@@ -1791,7 +1561,7 @@ async def _run_consolidation_job(
|
||||
await stack.enter_async_context(scope_locks[s])
|
||||
return await _process_tag_group(group_batches)
|
||||
|
||||
group_results = await _gather_or_cancel([_run_group(g, s) for g, s in zip(numbered_groups, group_scopes)])
|
||||
group_results = await asyncio.gather(*(_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:
|
||||
@@ -1828,19 +1598,6 @@ 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(
|
||||
@@ -1851,7 +1608,6 @@ 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
|
||||
@@ -1888,19 +1644,11 @@ async def _run_consolidation_job(
|
||||
if timing_parts:
|
||||
perf.log(f"[4] Timing breakdown: {', '.join(timing_parts)}")
|
||||
|
||||
# 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.
|
||||
# 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.
|
||||
if hit_round_limit:
|
||||
stats["mental_models_refreshed"] = 0
|
||||
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)"
|
||||
)
|
||||
logger.info(f"[CONSOLIDATION] bank={bank_id} skipping mental model refresh (round limit hit, re-queued)")
|
||||
else:
|
||||
set_stage("consolidation.refreshing_mental_models")
|
||||
await memory_engine._write_operation_progress(
|
||||
@@ -1914,7 +1662,7 @@ async def _run_consolidation_job(
|
||||
memory_engine=memory_engine,
|
||||
bank_id=bank_id,
|
||||
request_context=request_context,
|
||||
consolidated_tags=sorted(all_refresh_tags) or None,
|
||||
consolidated_tags=list(consolidated_tags) if consolidated_tags else None,
|
||||
perf=perf,
|
||||
)
|
||||
stats["mental_models_refreshed"] = mental_models_refreshed
|
||||
@@ -1978,7 +1726,7 @@ async def _trigger_mental_model_refreshes(
|
||||
if consolidated_tags:
|
||||
candidates = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, name, tags, last_refreshed_at, last_memory_seen_at, trigger
|
||||
SELECT id, name, tags, last_refreshed_at, trigger
|
||||
FROM {fq_table("mental_models")}
|
||||
WHERE bank_id = $1
|
||||
AND (trigger->>'refresh_after_consolidation')::boolean = true
|
||||
@@ -1993,7 +1741,7 @@ async def _trigger_mental_model_refreshes(
|
||||
else:
|
||||
candidates = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, name, tags, last_refreshed_at, last_memory_seen_at, trigger
|
||||
SELECT id, name, tags, last_refreshed_at, trigger
|
||||
FROM {fq_table("mental_models")}
|
||||
WHERE bank_id = $1
|
||||
AND (trigger->>'refresh_after_consolidation')::boolean = true
|
||||
@@ -2024,14 +1772,10 @@ 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(
|
||||
@@ -2076,6 +1820,8 @@ 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])
|
||||
|
||||
@@ -2093,11 +1839,7 @@ async def _process_memory_batch(
|
||||
)
|
||||
for m in memories
|
||||
]
|
||||
# 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)
|
||||
per_fact_recalls = await asyncio.gather(*recall_tasks)
|
||||
if perf:
|
||||
perf.record_timing("recall", time.time() - t0)
|
||||
|
||||
@@ -2220,7 +1962,9 @@ async def _process_memory_batch(
|
||||
new_text=update.text,
|
||||
observations=union_observations,
|
||||
source_fact_tags=agg.tags,
|
||||
source_bounds=_TemporalBounds.of(agg),
|
||||
source_occurred_start=agg.occurred_start,
|
||||
source_occurred_end=agg.occurred_end,
|
||||
source_mentioned_at=agg.mentioned_at,
|
||||
perf=perf,
|
||||
txn=txn,
|
||||
)
|
||||
@@ -2288,7 +2032,6 @@ async def _process_memory_batch(
|
||||
create.text,
|
||||
create_source_ids,
|
||||
agg.tags,
|
||||
_TemporalBounds.of(agg),
|
||||
txn=txn,
|
||||
)
|
||||
if merged_into is not None:
|
||||
@@ -2423,15 +2166,17 @@ async def _execute_update_action(
|
||||
new_text: str,
|
||||
observations: list["MemoryFact"],
|
||||
source_fact_tags: list[str] | None = None,
|
||||
source_bounds: _TemporalBounds = _TemporalBounds(),
|
||||
source_occurred_start: datetime | None = None,
|
||||
source_occurred_end: datetime | None = None,
|
||||
source_mentioned_at: datetime | None = None,
|
||||
perf: ConsolidationPerfLog | None = None,
|
||||
txn=None,
|
||||
) -> str | None:
|
||||
"""
|
||||
Update an existing observation.
|
||||
|
||||
Extends source_memory_ids with all contributing memories, widens the observation's temporal
|
||||
bounds by ``source_bounds`` (see :class:`_TemporalBounds`), and merges tags.
|
||||
Extends source_memory_ids with all contributing memories, updates temporal fields
|
||||
(LEAST for occurred_start, GREATEST for occurred_end / mentioned_at), 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
|
||||
@@ -2501,15 +2246,6 @@ 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")}
|
||||
@@ -2517,12 +2253,11 @@ async def _execute_update_action(
|
||||
embedding = $2::vector,
|
||||
source_memory_ids = $3,
|
||||
proof_count = $4,
|
||||
tags = $10,
|
||||
tags = $9,
|
||||
updated_at = now(),
|
||||
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}
|
||||
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}
|
||||
WHERE id = $5
|
||||
""",
|
||||
new_text,
|
||||
@@ -2530,10 +2265,9 @@ async def _execute_update_action(
|
||||
source_ids,
|
||||
len(source_ids),
|
||||
uuid.UUID(observation_id),
|
||||
source_bounds.event_date,
|
||||
source_bounds.occurred_start,
|
||||
source_bounds.occurred_end,
|
||||
source_bounds.mentioned_at,
|
||||
source_occurred_start,
|
||||
source_occurred_end,
|
||||
source_mentioned_at,
|
||||
merged_tags,
|
||||
)
|
||||
# The source-liveness checks above guard the *source* memories; the
|
||||
@@ -2551,24 +2285,12 @@ 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 (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 (event_date, 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,
|
||||
@@ -2581,10 +2303,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=merged_bounds.event_date,
|
||||
occurred_start=merged_bounds.occurred_start,
|
||||
occurred_end=merged_bounds.occurred_end,
|
||||
mentioned_at=merged_bounds.mentioned_at,
|
||||
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),
|
||||
created_at=cur.created_at if cur else None,
|
||||
),
|
||||
)
|
||||
|
||||
@@ -19,7 +19,6 @@ 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,
|
||||
@@ -873,7 +872,6 @@ 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.
|
||||
@@ -887,16 +885,11 @@ 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
|
||||
@@ -969,21 +962,7 @@ 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 — 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.
|
||||
"""
|
||||
"""Synchronous predict - processes each query group."""
|
||||
if not pairs:
|
||||
return []
|
||||
|
||||
@@ -1000,25 +979,20 @@ 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]
|
||||
|
||||
for start in range(0, len(indexed_texts), self.batch_size):
|
||||
batch = indexed_texts[start : start + self.batch_size]
|
||||
# Create rerank request
|
||||
request = RerankRequest(query=query, passages=passages)
|
||||
results = self._ranker.rerank(request)
|
||||
|
||||
# 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
|
||||
# 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
|
||||
|
||||
return all_scores
|
||||
finally:
|
||||
@@ -1776,7 +1750,6 @@ 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,65 +25,12 @@ 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 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
|
||||
Every graph_maintenance run is the same bank-wide sweep — the payload carries
|
||||
only ``bank_id``, and ``run_graph_maintenance_job`` drains the whole queue —
|
||||
so a second concurrent run for one bank adds no work. It is worse than
|
||||
useless: ``claim_graph_maintenance_batch`` locks queue rows ``FOR UPDATE``
|
||||
*without* ``SKIP LOCKED`` (it is written assuming a single runner per bank),
|
||||
so the runs convoy on each other's row locks while each holds a worker slot.
|
||||
@@ -481,11 +428,23 @@ class DataAccessOps(ABC):
|
||||
|
||||
# -- Bank index management -------------------------------------------
|
||||
|
||||
# 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 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).
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def drop_bank_vector_indexes(
|
||||
self,
|
||||
@@ -634,43 +593,6 @@ 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,
|
||||
@@ -678,10 +600,9 @@ class DataAccessOps(ABC):
|
||||
entities_table: str,
|
||||
ue_table: str,
|
||||
bank_id: str,
|
||||
entity_ids: list,
|
||||
) -> int:
|
||||
"""Delete those of ``entity_ids`` in ``bank_id`` that no longer have any
|
||||
unit_entities rows referencing them. Returns the number of rows deleted.
|
||||
"""Delete entities 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.
|
||||
@@ -694,10 +615,11 @@ class DataAccessOps(ABC):
|
||||
conn: DatabaseConnection,
|
||||
ec_table: str,
|
||||
ue_table: str,
|
||||
entity_ids: list,
|
||||
entities_table: str,
|
||||
bank_id: str,
|
||||
) -> int:
|
||||
"""Delete entity_cooccurrences rows incident to ``entity_ids`` where the
|
||||
two entities still exist but no current unit references both of them.
|
||||
"""Delete entity_cooccurrences rows in ``bank_id`` 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
|
||||
@@ -743,9 +665,7 @@ 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, 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.
|
||||
such row per bank is ever in flight.
|
||||
|
||||
Args:
|
||||
consolidation_bank_priority: Per-bank priority for consolidation scheduling.
|
||||
@@ -755,48 +675,8 @@ 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, 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.
|
||||
Returns claimed rows with operation_id, operation_type, task_payload, retry_count.
|
||||
The caller is responsible for building ClaimedTask objects.
|
||||
"""
|
||||
...
|
||||
|
||||
|
||||
@@ -15,7 +15,6 @@ from .ops import (
|
||||
LinkExpansionRows,
|
||||
TagListingParts,
|
||||
UpdatedWindow,
|
||||
document_serialization_sql,
|
||||
graph_maintenance_bank_serialization_sql,
|
||||
)
|
||||
from .result import DictResultRow as ResultRow
|
||||
@@ -358,80 +357,13 @@ 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.
|
||||
@@ -439,11 +371,9 @@ 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
|
||||
|
||||
@@ -452,33 +382,26 @@ class OracleOps(DataAccessOps):
|
||||
conn: DatabaseConnection,
|
||||
ec_table: str,
|
||||
ue_table: str,
|
||||
entity_ids: list,
|
||||
entities_table: str,
|
||||
bank_id: str,
|
||||
) -> 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 = ANY($1::uuid[]) OR entity_id_2 = ANY($1::uuid[]))
|
||||
WHERE entity_id_1 IN (SELECT id FROM {entities_table} WHERE bank_id = $1)
|
||||
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
|
||||
)
|
||||
""",
|
||||
entity_ids,
|
||||
bank_id,
|
||||
)
|
||||
return int(deleted.split()[-1]) if isinstance(deleted, str) and deleted.startswith("DELETE") else 0
|
||||
|
||||
@@ -696,21 +619,6 @@ 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")
|
||||
@@ -844,6 +752,23 @@ 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,
|
||||
@@ -851,12 +776,7 @@ class OracleOps(DataAccessOps):
|
||||
internal_id: str,
|
||||
fact_types: dict[str, str],
|
||||
) -> None:
|
||||
# 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.
|
||||
# Oracle uses a single global vector index (no per-bank indexes to drop).
|
||||
return
|
||||
|
||||
def get_entity_resolution_strategy(self) -> str:
|
||||
@@ -1224,7 +1144,7 @@ class OracleOps(DataAccessOps):
|
||||
if exclude_ids:
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count, serialization_key, bank_id
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
@@ -1243,7 +1163,7 @@ class OracleOps(DataAccessOps):
|
||||
else:
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count, serialization_key, bank_id
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
@@ -1261,7 +1181,7 @@ class OracleOps(DataAccessOps):
|
||||
if exclude_ids:
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count, serialization_key, bank_id
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
@@ -1278,7 +1198,7 @@ class OracleOps(DataAccessOps):
|
||||
else:
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count, serialization_key, bank_id
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
@@ -1319,7 +1239,7 @@ class OracleOps(DataAccessOps):
|
||||
extra = " AND ".join(conditions)
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count, serialization_key, bank_id
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
@@ -1366,7 +1286,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, serialization_key, bank_id
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
@@ -1418,14 +1338,13 @@ class OracleOps(DataAccessOps):
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT o.operation_id, o.operation_type, o.task_payload, o.retry_count, o.serialization_key, o.bank_id
|
||||
SELECT o.operation_id, o.operation_type, o.task_payload, o.retry_count
|
||||
FROM {table} o
|
||||
WHERE o.status = 'pending'
|
||||
AND o.task_payload IS NOT NULL
|
||||
AND o.operation_type = $1
|
||||
AND (o.next_retry_at IS NULL OR o.next_retry_at <= NOW())
|
||||
AND {graph_maintenance_bank_serialization_sql(table, "o")}
|
||||
AND {document_serialization_sql(table, "o")}
|
||||
ORDER BY o.created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
@@ -1447,7 +1366,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, o.serialization_key, o.bank_id
|
||||
SELECT o.operation_id, o.operation_type, o.task_payload, o.retry_count
|
||||
FROM {table} o
|
||||
WHERE o.status = 'pending'
|
||||
AND o.task_payload IS NOT NULL
|
||||
@@ -1455,7 +1374,6 @@ 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
|
||||
@@ -1466,14 +1384,13 @@ class OracleOps(DataAccessOps):
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT o.operation_id, o.operation_type, o.task_payload, o.retry_count, o.serialization_key, o.bank_id
|
||||
SELECT o.operation_id, o.operation_type, o.task_payload, o.retry_count
|
||||
FROM {table} o
|
||||
WHERE o.status = 'pending'
|
||||
AND o.task_payload IS NOT NULL
|
||||
AND o.operation_type != 'consolidation'
|
||||
AND (o.next_retry_at IS NULL OR o.next_retry_at <= NOW())
|
||||
AND {graph_maintenance_bank_serialization_sql(table, "o")}
|
||||
AND {document_serialization_sql(table, "o")}
|
||||
ORDER BY o.created_at
|
||||
LIMIT $1
|
||||
FOR UPDATE SKIP LOCKED
|
||||
@@ -1514,19 +1431,6 @@ 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}
|
||||
@@ -1537,34 +1441,4 @@ class OracleOps(DataAccessOps):
|
||||
operation_ids,
|
||||
)
|
||||
|
||||
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,
|
||||
)
|
||||
return all_rows
|
||||
|
||||
@@ -13,7 +13,6 @@ from .ops import (
|
||||
LinkExpansionRows,
|
||||
TagListingParts,
|
||||
UpdatedWindow,
|
||||
document_serialization_sql,
|
||||
graph_maintenance_bank_serialization_sql,
|
||||
)
|
||||
from .result import ResultRow
|
||||
@@ -473,126 +472,26 @@ 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 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.
|
||||
# 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.
|
||||
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
|
||||
USING victims v
|
||||
WHERE e.id = v.id
|
||||
WHERE e.bank_id = $1
|
||||
AND NOT EXISTS (
|
||||
SELECT 1 FROM {ue_table} ue WHERE ue.entity_id = e.id
|
||||
)
|
||||
""",
|
||||
bank_id,
|
||||
entity_ids,
|
||||
)
|
||||
# asyncpg returns "DELETE N"
|
||||
return int(result.split()[-1]) if isinstance(result, str) and result.startswith("DELETE") else 0
|
||||
@@ -602,18 +501,12 @@ class PostgreSQLOps(DataAccessOps):
|
||||
conn: DatabaseConnection,
|
||||
ec_table: str,
|
||||
ue_table: str,
|
||||
entity_ids: list,
|
||||
entities_table: str,
|
||||
bank_id: str,
|
||||
) -> int:
|
||||
# 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.
|
||||
# Scope by joining through entities.bank_id (entity_cooccurrences itself
|
||||
# has no bank_id column — entities don't span banks, so scoping via
|
||||
# entity_id_1 is sufficient).
|
||||
#
|
||||
# Ordered locking (deadlock avoidance, #2529): retain's concurrent
|
||||
# cooccurrence upsert (entity_resolver._flush_pending) locks rows in
|
||||
@@ -629,58 +522,26 @@ 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).
|
||||
#
|
||||
# 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.
|
||||
# The staleness predicate is an INTERSECT of the two entities' unit sets
|
||||
# rather than the equivalent `unit_entities u1 JOIN u2 ON u1.unit_id =
|
||||
# u2.unit_id` self-join (#2473): both INTERSECT branches resolve as Index
|
||||
# Only Scans on idx_unit_entities_entity_unit (entity_id, unit_id), so the
|
||||
# per-pair cost is bounded by the two entities' degrees. The self-join let
|
||||
# the planner pick an anti-join that rescanned a high-degree hub entity's
|
||||
# membership set for every pair — 28-30min on a bank with a ~100K-membership
|
||||
# hub, even when zero rows were stale. Don't "simplify" it back.
|
||||
result = await conn.execute(
|
||||
f"""
|
||||
WITH incident AS MATERIALIZED (
|
||||
WITH victims AS (
|
||||
SELECT c.entity_id_1, c.entity_id_2
|
||||
FROM {ec_table} c
|
||||
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
|
||||
)
|
||||
JOIN {entities_table} e ON e.id = c.entity_id_1
|
||||
WHERE e.bank_id = $1
|
||||
AND NOT EXISTS (
|
||||
SELECT unit_id FROM {ue_table} WHERE entity_id = c.entity_id_1
|
||||
INTERSECT
|
||||
SELECT unit_id FROM {ue_table} WHERE entity_id = c.entity_id_2
|
||||
)
|
||||
ORDER BY c.entity_id_1, c.entity_id_2
|
||||
FOR UPDATE OF c
|
||||
)
|
||||
@@ -689,7 +550,7 @@ class PostgreSQLOps(DataAccessOps):
|
||||
WHERE c.entity_id_1 = v.entity_id_1
|
||||
AND c.entity_id_2 = v.entity_id_2
|
||||
""",
|
||||
entity_ids,
|
||||
bank_id,
|
||||
)
|
||||
return int(result.split()[-1]) if isinstance(result, str) and result.startswith("DELETE") else 0
|
||||
|
||||
@@ -903,38 +764,6 @@ 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"""
|
||||
@@ -951,52 +780,33 @@ class PostgreSQLOps(DataAccessOps):
|
||||
),
|
||||
connected_sources AS (
|
||||
SELECT DISTINCT t.unit_id AS source_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 source_entities se
|
||||
CROSS JOIN LATERAL (
|
||||
SELECT ue_target.unit_id
|
||||
FROM {ue_table} ue_target
|
||||
JOIN source_entities se ON se.entity_id = ue_target.entity_id
|
||||
WHERE ue_target.entity_id = se.entity_id
|
||||
ORDER BY ue_target.unit_id DESC
|
||||
LIMIT {per_entity_limit}
|
||||
) t
|
||||
WHERE t.rn <= {per_entity_limit}
|
||||
AND NOT EXISTS (
|
||||
SELECT 1 FROM seed_sources ss WHERE ss.source_id = t.unit_id
|
||||
)
|
||||
WHERE 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
|
||||
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
|
||||
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
|
||||
LIMIT $2
|
||||
""",
|
||||
seed_ids,
|
||||
@@ -1071,6 +881,26 @@ 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,
|
||||
@@ -1084,14 +914,8 @@ 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 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).
|
||||
# The lock key must match create_bank_vector_indexes', whose `table`
|
||||
# is the fq name this reconstructs from `schema`.
|
||||
async with self._index_ddl_lock(f"{schema}.memory_units"):
|
||||
for ft, suffix in fact_types.items():
|
||||
uid = str(internal_id).replace("-", "")[:16]
|
||||
@@ -1410,7 +1234,7 @@ class PostgreSQLOps(DataAccessOps):
|
||||
if exclude_ids:
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count, serialization_key, bank_id
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
@@ -1429,7 +1253,7 @@ class PostgreSQLOps(DataAccessOps):
|
||||
else:
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count, serialization_key, bank_id
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
@@ -1447,7 +1271,7 @@ class PostgreSQLOps(DataAccessOps):
|
||||
if exclude_ids:
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count, serialization_key, bank_id
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
@@ -1464,7 +1288,7 @@ class PostgreSQLOps(DataAccessOps):
|
||||
else:
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count, serialization_key, bank_id
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
@@ -1505,7 +1329,7 @@ class PostgreSQLOps(DataAccessOps):
|
||||
extra = " AND ".join(conditions)
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count, serialization_key, bank_id
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
@@ -1552,7 +1376,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, serialization_key, bank_id
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
@@ -1603,14 +1427,13 @@ class PostgreSQLOps(DataAccessOps):
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT o.operation_id, o.operation_type, o.task_payload, o.retry_count, o.serialization_key, o.bank_id
|
||||
SELECT o.operation_id, o.operation_type, o.task_payload, o.retry_count
|
||||
FROM {table} o
|
||||
WHERE o.status = 'pending'
|
||||
AND o.task_payload IS NOT NULL
|
||||
AND o.operation_type = $1
|
||||
AND (o.next_retry_at IS NULL OR o.next_retry_at <= NOW())
|
||||
AND {graph_maintenance_bank_serialization_sql(table, "o")}
|
||||
AND {document_serialization_sql(table, "o")}
|
||||
ORDER BY o.created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
@@ -1632,7 +1455,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, o.serialization_key, o.bank_id
|
||||
SELECT o.operation_id, o.operation_type, o.task_payload, o.retry_count
|
||||
FROM {table} o
|
||||
WHERE o.status = 'pending'
|
||||
AND o.task_payload IS NOT NULL
|
||||
@@ -1640,7 +1463,6 @@ 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
|
||||
@@ -1651,14 +1473,13 @@ class PostgreSQLOps(DataAccessOps):
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT o.operation_id, o.operation_type, o.task_payload, o.retry_count, o.serialization_key, o.bank_id
|
||||
SELECT o.operation_id, o.operation_type, o.task_payload, o.retry_count
|
||||
FROM {table} o
|
||||
WHERE o.status = 'pending'
|
||||
AND o.task_payload IS NOT NULL
|
||||
AND o.operation_type != 'consolidation'
|
||||
AND (o.next_retry_at IS NULL OR o.next_retry_at <= NOW())
|
||||
AND {graph_maintenance_bank_serialization_sql(table, "o")}
|
||||
AND {document_serialization_sql(table, "o")}
|
||||
ORDER BY o.created_at
|
||||
LIMIT $1
|
||||
FOR UPDATE SKIP LOCKED
|
||||
@@ -1699,19 +1520,6 @@ 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}
|
||||
@@ -1722,31 +1530,4 @@ class PostgreSQLOps(DataAccessOps):
|
||||
operation_ids,
|
||||
)
|
||||
|
||||
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,
|
||||
)
|
||||
return all_rows
|
||||
|
||||
@@ -8,10 +8,9 @@ avoiding Python-level wrapping overhead (~570K __getitem__ calls per
|
||||
"""
|
||||
|
||||
import logging
|
||||
from collections.abc import AsyncIterator, Awaitable, Callable
|
||||
from collections.abc import AsyncIterator
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import Any
|
||||
from urllib.parse import parse_qs, urlparse
|
||||
|
||||
import asyncpg # noqa: F401
|
||||
|
||||
@@ -20,72 +19,6 @@ from .pool_instrumentation import PoolStats, instrument_acquire
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# GUC names this server rejected as unknown. Process-wide and never cleared: the
|
||||
# server's extension set does not change under a running process, and re-probing
|
||||
# would reintroduce the per-acquire cost this exists to avoid.
|
||||
_unsupported_settings: set[str] = set()
|
||||
|
||||
|
||||
def setting_rejected_by_server(name: str) -> bool:
|
||||
"""Whether this server has already rejected ``name`` as an unknown GUC.
|
||||
|
||||
For callers that apply a setting outside this helper — notably retain's link
|
||||
probing, which uses SET LOCAL inside its own transaction so the value cannot leak
|
||||
onto a pooled backend. Such a caller cannot simply let the statement fail: an error
|
||||
inside a transaction poisons it, so an unknown GUC would abort its work rather than
|
||||
merely fail to apply. The pool's setup runs on acquire and names the same GUCs, so
|
||||
by the time one of those callers runs, an unknown one is already recorded here.
|
||||
"""
|
||||
return name in _unsupported_settings
|
||||
|
||||
|
||||
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.
|
||||
"""
|
||||
settings = [pair for pair in settings if pair[0] not in _unsupported_settings]
|
||||
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.UndefinedObjectError:
|
||||
# The server does not define this GUC — an extension we tune for is absent
|
||||
# or predates it (hnsw.iterative_scan needs pgvector 0.8+, and pgvector
|
||||
# reserves the "hnsw." prefix, so an older one rejects it rather than
|
||||
# accepting a placeholder). Remember it: otherwise every acquire from here
|
||||
# on re-pays a failed batch plus one statement per setting, which behind a
|
||||
# transaction-mode pooler is a server-side transaction each — the burn
|
||||
# #3499 removed. Narrow to UndefinedObjectError so a transient failure
|
||||
# does not disable a setting the server does support.
|
||||
logger.info("Server does not know %s — not sending it again on this process", name)
|
||||
_unsupported_settings.add(name)
|
||||
except asyncpg.exceptions.PostgresError:
|
||||
logger.debug("Could not set %s — retrying it on the next acquire", name)
|
||||
|
||||
|
||||
class PostgresConnection(DatabaseConnection):
|
||||
"""DatabaseConnection wrapper around an asyncpg.Connection."""
|
||||
@@ -131,54 +64,6 @@ 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."""
|
||||
|
||||
@@ -194,12 +79,6 @@ 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,
|
||||
@@ -214,15 +93,7 @@ class PostgreSQLBackend(DatabaseBackend):
|
||||
) -> None:
|
||||
from ...config import get_config
|
||||
|
||||
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
|
||||
self._acquire_warn_threshold_s = get_config().db_acquire_warn_threshold_ms / 1000.0
|
||||
# Kept for acquire() below: asyncpg's ``timeout`` create_pool kwarg is a
|
||||
# *connect* kwarg (how long establishing a new connection may take), and
|
||||
# ``Pool.acquire()`` defaults to waiting for a free connection forever.
|
||||
@@ -230,30 +101,6 @@ 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,
|
||||
@@ -261,13 +108,16 @@ class PostgreSQLBackend(DatabaseBackend):
|
||||
command_timeout=command_timeout,
|
||||
statement_cache_size=statement_cache_size,
|
||||
timeout=acquire_timeout,
|
||||
init=pool_init,
|
||||
setup=pool_setup,
|
||||
# init runs once per new connection; setup runs on every acquire,
|
||||
# after asyncpg's release-time RESET ALL. Passing init_callback as
|
||||
# both keeps the per-connection session GUCs (hnsw.ef_search, etc.)
|
||||
# applied after a connection is reused, not just on first creation.
|
||||
init=init_callback,
|
||||
setup=init_callback,
|
||||
)
|
||||
logger.info(
|
||||
f"PostgreSQL pool created (min={min_size}, max={max_size}, "
|
||||
f"cmd_timeout={command_timeout}s, acquire_timeout={acquire_timeout}s, "
|
||||
f"session_setup_on_acquire={setup_on_acquire})"
|
||||
f"cmd_timeout={command_timeout}s, acquire_timeout={acquire_timeout}s)"
|
||||
)
|
||||
|
||||
async def shutdown(self) -> None:
|
||||
|
||||
@@ -125,32 +125,12 @@ 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, applying the configured query prefix."""
|
||||
return self._encode_prefixed(texts, self.query_prefix)
|
||||
"""Generate embeddings for query text. Providers without asymmetric embeddings use encode()."""
|
||||
return self.encode(texts)
|
||||
|
||||
def encode_documents(self, texts: list[str]) -> list[list[float]]:
|
||||
"""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])
|
||||
"""Generate embeddings for stored document text. Providers without asymmetric embeddings use encode()."""
|
||||
return self.encode(texts)
|
||||
|
||||
|
||||
@@ -416,6 +396,17 @@ 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.")
|
||||
@@ -478,8 +469,6 @@ 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.
|
||||
@@ -490,16 +479,12 @@ 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
|
||||
@@ -653,8 +638,6 @@ 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.
|
||||
@@ -666,8 +649,6 @@ 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
|
||||
@@ -675,8 +656,6 @@ 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
|
||||
|
||||
@@ -795,8 +774,6 @@ 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
|
||||
|
||||
@@ -808,8 +785,6 @@ class CodexOAuthEmbeddings(OpenAIEmbeddings):
|
||||
batch_size=batch_size,
|
||||
dimensions=dimensions,
|
||||
max_retries=max_retries,
|
||||
query_prefix=query_prefix,
|
||||
passage_prefix=passage_prefix,
|
||||
)
|
||||
|
||||
@property
|
||||
@@ -1154,8 +1129,6 @@ 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.
|
||||
@@ -1167,16 +1140,12 @@ 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
|
||||
|
||||
@@ -1276,8 +1245,6 @@ 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.
|
||||
@@ -1292,8 +1259,6 @@ 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
|
||||
@@ -1302,8 +1267,6 @@ 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
|
||||
|
||||
@@ -1655,24 +1618,11 @@ 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, query_prefix=query_prefix, passage_prefix=passage_prefix)
|
||||
return RemoteTEIEmbeddings(base_url=url)
|
||||
elif provider == "local":
|
||||
return LocalSTEmbeddings(
|
||||
model_name=config.embeddings_local_model,
|
||||
@@ -1710,8 +1660,6 @@ 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)
|
||||
@@ -1719,8 +1667,6 @@ 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
|
||||
@@ -1735,8 +1681,6 @@ 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
|
||||
@@ -1751,8 +1695,6 @@ 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
|
||||
@@ -1785,8 +1727,6 @@ 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(
|
||||
@@ -1795,8 +1735,6 @@ 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,9 +45,6 @@ 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
|
||||
@@ -431,18 +428,6 @@ 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.
|
||||
@@ -472,25 +457,6 @@ 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.
|
||||
@@ -906,7 +872,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 or not e.resolve:
|
||||
if e.is_label:
|
||||
continue
|
||||
name_lower = e.name.lower()
|
||||
rep_by_lower.setdefault(name_lower, e.name)
|
||||
@@ -938,12 +904,7 @@ class EntityResolver:
|
||||
taxonomy_lookup: set[str] | None = None,
|
||||
labels_cfg=None,
|
||||
) -> list[ResolvedEntity]:
|
||||
"""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.
|
||||
"""
|
||||
"""Shared scoring + upsert logic used by both lookup strategies."""
|
||||
|
||||
# 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
|
||||
@@ -962,9 +923,6 @@ 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
|
||||
@@ -993,14 +951,10 @@ 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 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.
|
||||
if not candidates:
|
||||
# Will create new entity
|
||||
entities_to_create.append(
|
||||
_EntityToCreate(
|
||||
idx=idx, name=entity_text, event_date=entity_event_date, is_label=is_label, resolve=resolve
|
||||
)
|
||||
_EntityToCreate(idx=idx, name=entity_text, event_date=entity_event_date, is_label=is_label)
|
||||
)
|
||||
continue
|
||||
|
||||
@@ -1103,9 +1057,7 @@ 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, 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).
|
||||
# excluded and keep exact grouping.
|
||||
canonical_by_member = self._intrabatch_canonical_map(entities_to_create)
|
||||
|
||||
@dataclass
|
||||
@@ -1290,38 +1242,8 @@ 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, txn=None
|
||||
self, conn, unit_entity_pairs: list[tuple[str, str, datetime | None]], bank_id: str | None = None
|
||||
):
|
||||
# Sorted bulk insert to prevent deadlocks from inconsistent lock ordering
|
||||
# across concurrent transactions on the unit_entities unique index.
|
||||
@@ -1342,7 +1264,6 @@ class EntityResolver:
|
||||
bank_id=bank_id,
|
||||
unit_ids=unit_ids,
|
||||
entity_ids=entity_ids,
|
||||
txn=txn,
|
||||
)
|
||||
|
||||
# Build maps keyed by unit_id:
|
||||
|
||||
@@ -1,33 +1,29 @@
|
||||
"""Async graph maintenance after document/unit deletes.
|
||||
|
||||
Two queue-driven passes run together on every worker invocation:
|
||||
Three reconciliation passes run together on every worker invocation:
|
||||
|
||||
1. **Relink top-up.** Drain ``graph_maintenance_queue`` (units whose
|
||||
outgoing temporal/semantic links lost a neighbour to a delete). For
|
||||
each, count current outgoing links per type; if below cap, run the
|
||||
same probes retain uses and insert the missing links.
|
||||
|
||||
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.
|
||||
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.
|
||||
|
||||
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.
|
||||
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.
|
||||
|
||||
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 (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.
|
||||
out. This module orchestrates them (drain the queue, wrap the sweep in a
|
||||
deadlock-retry) and asks the store to do the part that touches storage. A store
|
||||
whose links travel inside its memories has no `memory_links` to dangle and no
|
||||
join table to sweep, so its relink and cooccurrence passes are no-ops and the
|
||||
job simply prunes the orphan `entities` rows, which stay in Postgres regardless.
|
||||
|
||||
The worker dedupes on bank: a second job for the same bank is dropped
|
||||
while one is pending. Once processing starts, a new job becomes the
|
||||
@@ -36,9 +32,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 drains the same two
|
||||
bank-scoped queues — while convoying on each other's row locks and holding a
|
||||
worker slot each.
|
||||
concurrent runs would do no extra work anyway — each is this same bank-wide
|
||||
sweep — while convoying on each other's row locks and holding a worker slot
|
||||
each.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -64,14 +60,19 @@ if TYPE_CHECKING:
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# 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
|
||||
# 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
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -80,22 +81,15 @@ 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 | bool]:
|
||||
def as_dict(self) -> dict[str, int]:
|
||||
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,
|
||||
}
|
||||
|
||||
|
||||
@@ -149,65 +143,19 @@ 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 | bool]:
|
||||
"""Drain both maintenance queues for ``bank_id``, within a time budget.
|
||||
) -> dict[str, int]:
|
||||
"""Run all maintenance passes for ``bank_id`` until the relink queue is
|
||||
drained, then sweep entities and cooccurrences once.
|
||||
|
||||
Returns:
|
||||
Per-pass counters from :class:`JobResult`. ``queues_drained`` is False
|
||||
when the budget ran out with work still queued — the caller re-submits.
|
||||
Per-pass counters from :class:`JobResult`.
|
||||
"""
|
||||
del request_context # accepted for symmetry with other run_*_job helpers
|
||||
from ..config import get_config
|
||||
from .memories import get_memories
|
||||
|
||||
@@ -217,7 +165,6 @@ 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
|
||||
@@ -226,119 +173,51 @@ 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, deadline=deadline
|
||||
)
|
||||
result.relink_units_processed = relink.units_processed
|
||||
result.relink_links_added = relink.links_added
|
||||
relink = await store.relink_pass(backend=backend, fq_table=fq_table, bank_id=bank_id, config=config)
|
||||
result.relink_units_processed = relink.get("relink_units_processed", 0)
|
||||
result.relink_links_added = relink.get("relink_links_added", 0)
|
||||
|
||||
# --- Pass 2: 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
|
||||
# --- 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
|
||||
|
||||
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,22 +160,16 @@ class MemoryEngineInterface(ABC):
|
||||
async def list_banks(
|
||||
self,
|
||||
*,
|
||||
search_query: str | None = None,
|
||||
limit: int = 100,
|
||||
offset: int = 0,
|
||||
request_context: "RequestContext",
|
||||
) -> dict[str, Any]:
|
||||
) -> list[dict[str, Any]]:
|
||||
"""
|
||||
List memory banks, one page at a time.
|
||||
List all memory banks.
|
||||
|
||||
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:
|
||||
Dict with ``banks`` (the page), ``total``, ``limit`` and ``offset``.
|
||||
List of bank info dicts.
|
||||
"""
|
||||
...
|
||||
|
||||
@@ -328,7 +322,7 @@ class MemoryEngineInterface(ABC):
|
||||
self,
|
||||
bank_id: str,
|
||||
*,
|
||||
fact_type: str | list[str] | None = None,
|
||||
fact_type: str | None = None,
|
||||
search_query: str | None = None,
|
||||
entity_id: str | None = None,
|
||||
created_before: datetime | None = None,
|
||||
@@ -341,8 +335,7 @@ class MemoryEngineInterface(ABC):
|
||||
|
||||
Args:
|
||||
bank_id: The memory bank ID.
|
||||
fact_type: Filter by fact type. A list matches any of them; an empty
|
||||
list is treated as no filter.
|
||||
fact_type: Filter by fact type.
|
||||
search_query: Full-text search query.
|
||||
entity_id: Filter to memory units linked to this entity ID.
|
||||
created_before: Keep units with ``created_at`` before this instant.
|
||||
|
||||
@@ -5,7 +5,6 @@ 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
|
||||
@@ -15,8 +14,6 @@ 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."""
|
||||
@@ -71,7 +68,7 @@ class LLMInterface(ABC):
|
||||
api_key: str,
|
||||
base_url: str,
|
||||
model: str,
|
||||
reasoning_effort: str | None = None,
|
||||
reasoning_effort: str = "low",
|
||||
**kwargs: Any,
|
||||
):
|
||||
"""
|
||||
@@ -82,37 +79,14 @@ 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, or None when the operator
|
||||
configured none — in which case no provider sends the parameter and
|
||||
every model runs at its own default effort.
|
||||
reasoning_effort: Reasoning effort level for supported providers.
|
||||
**kwargs: Additional provider-specific parameters.
|
||||
"""
|
||||
self.provider = provider.lower()
|
||||
self.api_key = api_key
|
||||
self.base_url = base_url
|
||||
self.model = model
|
||||
# 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."
|
||||
)
|
||||
self.reasoning_effort = reasoning_effort
|
||||
|
||||
@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 | None,
|
||||
reasoning_effort: str,
|
||||
groq_service_tier: str | None = None,
|
||||
openai_service_tier: str | None = None,
|
||||
bedrock_service_tier: str | None = None,
|
||||
@@ -323,9 +323,7 @@ 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, or None when
|
||||
the operator configured none (providers then fall back to the default level
|
||||
and may skip the parameter entirely).
|
||||
reasoning_effort: Reasoning effort level for supported providers.
|
||||
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".
|
||||
@@ -515,7 +513,6 @@ 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,
|
||||
@@ -639,7 +636,7 @@ class LLMProvider:
|
||||
api_key: str,
|
||||
base_url: str,
|
||||
model: str,
|
||||
reasoning_effort: str | None = None,
|
||||
reasoning_effort: str = "low",
|
||||
groq_service_tier: str | None = None,
|
||||
openai_service_tier: str | None = None,
|
||||
bedrock_service_tier: str | None = None,
|
||||
@@ -668,8 +665,7 @@ class LLMProvider:
|
||||
api_key: API key.
|
||||
base_url: Base URL for the API.
|
||||
model: Model name.
|
||||
reasoning_effort: Reasoning effort level for supported providers, or None
|
||||
when the operator configured none.
|
||||
reasoning_effort: Reasoning effort level for supported providers.
|
||||
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.
|
||||
@@ -1440,6 +1436,7 @@ 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,
|
||||
@@ -1501,7 +1498,7 @@ class LLMProvider:
|
||||
api_key=api_key,
|
||||
base_url=base_url,
|
||||
model=model,
|
||||
reasoning_effort=os.getenv(ENV_LLM_REASONING_EFFORT) or None,
|
||||
reasoning_effort=os.getenv(ENV_LLM_REASONING_EFFORT, DEFAULT_LLM_REASONING_EFFORT),
|
||||
extra_body=extra_body,
|
||||
default_headers=default_headers,
|
||||
cache_affinity=cache_affinity,
|
||||
|
||||
@@ -3,16 +3,15 @@
|
||||
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** (configurable, default hourly): delete ``audit_log`` and
|
||||
``llm_requests`` rows older than their configured retention, across *all*
|
||||
tenant schemas.
|
||||
- **Retention sweeps** (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 5 min):
|
||||
- **Scheduled mental model refresh** (configurable check cadence, default 60s):
|
||||
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
|
||||
@@ -29,34 +28,15 @@ 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. 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.
|
||||
process. Retention and operation cleanup are deletes; the consolidation reconcile
|
||||
and the scheduled mental model refresh both dedupe against in-flight operations
|
||||
inside the inserting transaction (see ``_submit_async_operation``).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import random
|
||||
import time
|
||||
from collections.abc import Coroutine
|
||||
from datetime import datetime, timedelta, timezone
|
||||
@@ -74,6 +54,12 @@ 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
|
||||
@@ -83,32 +69,6 @@ _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`."""
|
||||
@@ -162,11 +122,10 @@ 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.
|
||||
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
|
||||
audit_on = cfg.audit_log_retention_days > 0
|
||||
llm_on = cfg.llm_trace_enabled and cfg.llm_trace_retention_days > 0
|
||||
mm_refresh_on = cfg.mental_model_refresh_tick_seconds > 0
|
||||
op_cleanup_on = cfg.operation_cleanup_interval_seconds > 0 and cfg.operation_retention_days > 0
|
||||
op_cleanup_on = cfg.operation_retention_days > 0
|
||||
return (
|
||||
reconcile_on
|
||||
or audit_on
|
||||
@@ -195,8 +154,6 @@ 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()
|
||||
@@ -207,26 +164,6 @@ 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()
|
||||
@@ -238,8 +175,7 @@ class MaintenanceLoop:
|
||||
|
||||
async def _tick(self) -> None:
|
||||
cfg = get_config()
|
||||
retention_interval = cfg.retention_sweep_interval_seconds
|
||||
if retention_interval > 0 and self._is_due("retention", retention_interval):
|
||||
if self._is_due("retention", _RETENTION_INTERVAL_SECONDS):
|
||||
await self._run_timed("retention", self._run_retention(cfg))
|
||||
interval = cfg.consolidation_reconcile_interval_seconds
|
||||
if interval > 0 and self._is_due("reconcile", interval):
|
||||
@@ -247,12 +183,7 @@ 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())
|
||||
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)
|
||||
):
|
||||
if cfg.operation_retention_days > 0 and self._is_due("operation_cleanup", _OPERATION_CLEANUP_INTERVAL_SECONDS):
|
||||
await self._run_timed("operation cleanup", self._run_operation_cleanup(cfg))
|
||||
if self._cross_store_recovery_enabled() and self._is_due("txn_recovery", _TXN_RECOVERY_INTERVAL_SECONDS):
|
||||
await self._run_timed("cross-store txn recovery", self._run_txn_recovery())
|
||||
@@ -282,86 +213,26 @@ 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) -> 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.
|
||||
"""
|
||||
async def _purge_expired(self, table: str, ts_col: str, days: int) -> None:
|
||||
"""Delete rows older than ``days`` from ``table`` across every tenant schema."""
|
||||
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
|
||||
)
|
||||
except Exception as 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
|
||||
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,
|
||||
)
|
||||
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
|
||||
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}")
|
||||
|
||||
# ── terminal operation cleanup ─────────────────────────────────────────
|
||||
|
||||
@@ -418,12 +289,7 @@ 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.
|
||||
# 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
|
||||
)
|
||||
await engine.purge_expired_export_archives(conn, table, cutoff)
|
||||
async with conn.transaction():
|
||||
deleted = await backend.ops.prune_terminal_operations(
|
||||
conn, table, cutoff, batch_size=cfg.operation_cleanup_batch_size
|
||||
@@ -618,8 +484,7 @@ 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, last_memory_seen_at "
|
||||
f"FROM {fq_table('mental_models')} "
|
||||
f"SELECT id, tags, trigger, last_refreshed_at FROM {fq_table('mental_models')} "
|
||||
"WHERE bank_id = $1 AND id = $2",
|
||||
bank_id,
|
||||
mm_id,
|
||||
|
||||
@@ -18,7 +18,6 @@ from .base import (
|
||||
FactRecord,
|
||||
MemoriesExtension,
|
||||
MemoryPatch,
|
||||
RecallArms,
|
||||
ScanPage,
|
||||
StoredMemory,
|
||||
build_fact_records,
|
||||
@@ -76,7 +75,6 @@ __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
|
||||
from ..search.retrieval import GraphRetriever, SemanticBm25Result
|
||||
|
||||
|
||||
class MemoryTxn:
|
||||
@@ -63,21 +63,6 @@ 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
|
||||
@@ -92,30 +77,6 @@ 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.
|
||||
@@ -398,50 +359,6 @@ 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.
|
||||
|
||||
@@ -491,21 +408,6 @@ 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:
|
||||
@@ -718,10 +620,10 @@ class MemoriesExtension(Extension, ABC):
|
||||
"""Apply partial updates. Only the fields set on each patch change."""
|
||||
raise NotImplementedError
|
||||
|
||||
# ------------------------------------------------------------------ recall
|
||||
# ------------------------------------------------------------------ recall arms
|
||||
|
||||
@abstractmethod
|
||||
async def recall_unified(
|
||||
async def search(
|
||||
self,
|
||||
*,
|
||||
conn,
|
||||
@@ -730,8 +632,6 @@ 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,
|
||||
@@ -739,24 +639,40 @@ class MemoriesExtension(Extension, ABC):
|
||||
created_before: datetime | None = None,
|
||||
min_semantic: float | None = None,
|
||||
min_keyword: float | None = None,
|
||||
enable_graph: bool = True,
|
||||
) -> "dict[str, RecallArms]":
|
||||
"""Run ALL retrieval arms for every fact_type — the whole recall interface, in one call.
|
||||
graph_seed_min_similarity: float | None = None,
|
||||
) -> "dict[str, SemanticBm25Result]":
|
||||
"""Run the semantic + BM25 arms.
|
||||
|
||||
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.
|
||||
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``).
|
||||
"""
|
||||
|
||||
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.
|
||||
@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]``.
|
||||
|
||||
``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.
|
||||
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.
|
||||
"""
|
||||
|
||||
def graph_retriever(self) -> "GraphRetriever | None":
|
||||
@@ -906,9 +822,6 @@ 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
|
||||
@@ -936,19 +849,6 @@ 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,
|
||||
@@ -989,7 +889,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_memory_seen_at`` is at or after it cannot be stale, whatever its
|
||||
``last_refreshed_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.
|
||||
"""
|
||||
@@ -1043,7 +943,7 @@ class MemoriesExtension(Extension, ABC):
|
||||
ops,
|
||||
fq_table,
|
||||
bank_id: str,
|
||||
fact_type: str | list[str] | None = None,
|
||||
fact_type: str | None = None,
|
||||
search_query: str | None = None,
|
||||
consolidation_state: str | None = None,
|
||||
state: str | None = None,
|
||||
@@ -1105,9 +1005,6 @@ 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
|
||||
@@ -1118,10 +1015,6 @@ 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:
|
||||
@@ -1239,15 +1132,7 @@ 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],
|
||||
txn: "MemoryTxn | None" = None,
|
||||
self, *, conn, ops, fq_table, bank_id: str | None = None, unit_ids: list[Any], entity_ids: list[Any]
|
||||
) -> None:
|
||||
"""Record the unit→entity postings for a batch of memories.
|
||||
|
||||
@@ -1258,14 +1143,6 @@ 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(
|
||||
@@ -1280,28 +1157,17 @@ class MemoriesExtension(Extension, ABC):
|
||||
"""
|
||||
return 0
|
||||
|
||||
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 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 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.
|
||||
"""
|
||||
async def prune_orphan_entities(self, *, conn, fq_table, bank_id: str) -> int:
|
||||
"""Delete `entities` rows no live memory references. Returns the count."""
|
||||
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()
|
||||
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
|
||||
|
||||
|
||||
__all__ = [
|
||||
@@ -1321,12 +1187,10 @@ __all__ = [
|
||||
"META_UPDATED_AT",
|
||||
"CausalEdgeRecord",
|
||||
"DeletePredicate",
|
||||
"EntityPrunePassResult",
|
||||
"FactRecord",
|
||||
"MemoriesExtension",
|
||||
"MemoryPatch",
|
||||
"MemoryTxn",
|
||||
"RelinkPassResult",
|
||||
"ScanPage",
|
||||
"StoredMemory",
|
||||
"build_fact_records",
|
||||
|
||||
@@ -17,13 +17,6 @@ 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)
|
||||
@@ -36,11 +29,7 @@ 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 consolidation_failed_at IS NULL
|
||||
AND fact_type IN ('experience', 'world')
|
||||
) AS pending,
|
||||
COUNT(*) FILTER (WHERE consolidated_at IS NULL AND fact_type IN ('experience', 'world')) AS pending,
|
||||
COUNT(*) FILTER (WHERE consolidation_failed_at IS NOT NULL AND fact_type IN ('experience', 'world')) AS failed
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $1
|
||||
|
||||
@@ -84,7 +84,7 @@ async def list_memory_units(
|
||||
ops,
|
||||
fq_table,
|
||||
bank_id: str,
|
||||
fact_type: str | list[str] | None = None,
|
||||
fact_type: str | None = None,
|
||||
search_query: str | None = None,
|
||||
consolidation_state: str | None = None,
|
||||
state: str | None = None,
|
||||
@@ -104,8 +104,7 @@ 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). A list matches any of
|
||||
them; an empty list is treated as no filter.
|
||||
fact_type: Filter by fact type (world, experience)
|
||||
search_query: Full-text search query (searches text and context fields)
|
||||
document_id: Optional filter to a single source document.
|
||||
tags: Optional list of tag names to filter by. When omitted, no tag
|
||||
@@ -155,14 +154,8 @@ async def list_memory_units(
|
||||
|
||||
if fact_type:
|
||||
param_count += 1
|
||||
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))
|
||||
query_conditions.append(f"fact_type = ${param_count}")
|
||||
query_params.append(fact_type)
|
||||
|
||||
if document_id:
|
||||
param_count += 1
|
||||
@@ -247,8 +240,7 @@ 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,
|
||||
updated_at, source_memory_ids, {curation_cols}
|
||||
tags, metadata, consolidated_at, consolidation_failed_at, edited_at, {curation_cols}
|
||||
FROM {source_table}
|
||||
{where_clause}
|
||||
ORDER BY mentioned_at DESC NULLS LAST, created_at DESC
|
||||
@@ -312,12 +304,6 @@ 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 []],
|
||||
}
|
||||
)
|
||||
|
||||
@@ -477,7 +463,7 @@ async def list_entities(
|
||||
# Get paginated entities
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, canonical_name, entity_kind, mention_count, first_seen, last_seen, metadata
|
||||
SELECT id, canonical_name, mention_count, first_seen, last_seen, metadata
|
||||
FROM {fq_table("entities")}
|
||||
WHERE {where_clause}
|
||||
ORDER BY mention_count DESC, last_seen DESC, id ASC
|
||||
@@ -504,9 +490,6 @@ 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` 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.
|
||||
* **The graph-maintenance job.** :func:`enqueue_relink_victims` runs inside the
|
||||
delete transaction; :func:`relink_pass`, :func:`prune_orphan_entities` and
|
||||
:func:`prune_stale_cooccurrences` are the three reconciliation passes the job
|
||||
drives. The job keeps the orchestration (pass ordering, the deadlock retry
|
||||
around the sweeps, the timing log); each function here does the pass's work.
|
||||
|
||||
:func:`entity_memory_counts` and :func:`entities_for_units` are the two entity
|
||||
postings reads that are not part of the graph view but read the same join table.
|
||||
@@ -28,10 +28,8 @@ 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
|
||||
@@ -42,7 +40,6 @@ from ...retain.link_utils import (
|
||||
_normalize_datetime,
|
||||
compute_semantic_links_ann,
|
||||
)
|
||||
from ..base import EntityPrunePassResult, RelinkPassResult
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -62,37 +59,6 @@ _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
|
||||
@@ -479,46 +445,6 @@ 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
|
||||
|
||||
|
||||
@@ -602,8 +528,7 @@ async def relink_pass(
|
||||
fq_table: Callable[[str], str],
|
||||
bank_id: str,
|
||||
config: Any,
|
||||
deadline: float | None = None,
|
||||
) -> RelinkPassResult:
|
||||
) -> dict:
|
||||
"""Drain ``graph_maintenance_queue`` for ``bank_id``, topping up lost links.
|
||||
|
||||
Per-iteration loop: claim → top up → commit. We rely on at most one job per
|
||||
@@ -624,13 +549,8 @@ 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:
|
||||
A :class:`RelinkPassResult`. ``queue_exhausted`` is False when the
|
||||
deadline (or the iteration cap) stopped the drain with rows still queued.
|
||||
``{"relink_units_processed": int, "relink_links_added": int}``.
|
||||
"""
|
||||
del config # accepted for symmetry with stores that tune their own relinking
|
||||
ops = backend.ops
|
||||
@@ -638,12 +558,7 @@ 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:
|
||||
@@ -669,14 +584,9 @@ 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 RelinkPassResult(
|
||||
units_processed=units_processed,
|
||||
links_added=links_added,
|
||||
queue_exhausted=drained,
|
||||
)
|
||||
return {"relink_units_processed": units_processed, "relink_links_added": links_added}
|
||||
|
||||
|
||||
async def _relink_batch(
|
||||
@@ -806,192 +716,69 @@ async def _relink_batch(
|
||||
return len(new_links)
|
||||
|
||||
|
||||
async def enqueue_entity_prune_candidates(
|
||||
async def prune_orphan_entities(
|
||||
*,
|
||||
conn: DatabaseConnection,
|
||||
fq_table: Callable[[str], str],
|
||||
bank_id: str,
|
||||
affected_unit_ids: list,
|
||||
) -> int:
|
||||
"""Enqueue the entities ``affected_unit_ids`` reference as prune candidates.
|
||||
"""Delete ``entities`` rows in the bank with no remaining ``unit_entities``
|
||||
references. Returns the number pruned.
|
||||
|
||||
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.
|
||||
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.
|
||||
|
||||
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.
|
||||
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.
|
||||
"""
|
||||
if not affected_unit_ids:
|
||||
return 0
|
||||
|
||||
ops = _ops_for(conn)
|
||||
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
|
||||
return await ops.prune_orphan_entities(
|
||||
conn,
|
||||
fq_table("entities"),
|
||||
fq_table("unit_entities"),
|
||||
bank_id,
|
||||
)
|
||||
|
||||
|
||||
@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(
|
||||
async def prune_stale_cooccurrences(
|
||||
*,
|
||||
backend: Any,
|
||||
conn: DatabaseConnection,
|
||||
fq_table: Callable[[str], str],
|
||||
bank_id: str,
|
||||
deadline: float | None = None,
|
||||
) -> EntityPrunePassResult:
|
||||
"""Drain ``entity_maintenance_queue`` for ``bank_id``, pruning what died.
|
||||
) -> int:
|
||||
"""Delete cooccurrence rows no current memory witnesses. Returns the count.
|
||||
|
||||
Per-iteration loop: claim → prune → commit, mirroring :func:`relink_pass`.
|
||||
Each iteration does two deletes over the claimed batch:
|
||||
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.
|
||||
|
||||
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.
|
||||
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.
|
||||
"""
|
||||
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,
|
||||
ops = _ops_for(conn)
|
||||
return await ops.prune_stale_cooccurrences(
|
||||
conn,
|
||||
fq_table("entity_cooccurrences"),
|
||||
fq_table("unit_entities"),
|
||||
fq_table("entities"),
|
||||
bank_id,
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"MAX_SEMANTIC_LINKS_PER_UNIT",
|
||||
"enqueue_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,9 +461,8 @@ 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 — 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
|
||||
``updated_at`` is deliberately left alone, matching the consolidator's own
|
||||
statements: consolidation bookkeeping is not an edit to the memory, and
|
||||
bumping it would make every consolidation pass look like a write to the
|
||||
staleness check below.
|
||||
"""
|
||||
|
||||
@@ -293,8 +293,6 @@ 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")}
|
||||
@@ -422,8 +420,7 @@ 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, updated_at = now() "
|
||||
f"WHERE id = $1 AND bank_id = $2",
|
||||
f"UPDATE {fq_table('invalidated_memory_units')} SET invalidation_reason = $3 WHERE id = $1 AND bank_id = $2",
|
||||
str(unit_id),
|
||||
bank_id,
|
||||
reason,
|
||||
@@ -501,8 +498,7 @@ 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, updated_at = now() "
|
||||
f"WHERE id = $1 AND bank_id = $2",
|
||||
f"UPDATE {fq_table('memory_units')} SET embedding = $3::vector WHERE id = $1 AND bank_id = $2",
|
||||
str(unit_id),
|
||||
bank_id,
|
||||
embedding,
|
||||
|
||||
@@ -23,16 +23,7 @@ from __future__ import annotations
|
||||
from datetime import datetime
|
||||
from typing import Any
|
||||
|
||||
from .base import (
|
||||
DeletePredicate,
|
||||
EntityPrunePassResult,
|
||||
MemoriesExtension,
|
||||
MemoryPatch,
|
||||
RecallArms,
|
||||
RelinkPassResult,
|
||||
ScanPage,
|
||||
StoredMemory,
|
||||
)
|
||||
from .base import DeletePredicate, MemoriesExtension, MemoryPatch, ScanPage, StoredMemory
|
||||
from .pg import counts, curation, graph, reads, writes
|
||||
|
||||
|
||||
@@ -80,135 +71,7 @@ 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
|
||||
|
||||
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 results. 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 results (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) ----
|
||||
# ------------------------------------------------------------------ recall arms
|
||||
|
||||
async def search(
|
||||
self,
|
||||
@@ -228,14 +91,6 @@ class PostgresMemories(MemoriesExtension):
|
||||
min_keyword: float | None = None,
|
||||
graph_seed_min_similarity: float | None = None,
|
||||
) -> "dict[str, SemanticBm25Result]":
|
||||
"""The dense + keyword arms, as one UNION query.
|
||||
|
||||
How deep the ANN scan goes is not decided here: the connection carries
|
||||
``hnsw.iterative_scan``, which lets the scan resume until this query's own LIMIT
|
||||
is met (see ``_ANN_TUNING_HIGH_RECALL``). Before that was enabled the scan
|
||||
stopped at ``hnsw.ef_search`` rows — a fixed 200 — so a larger recall budget
|
||||
widened the SQL and changed nothing.
|
||||
"""
|
||||
# Imported here: retrieval imports this package, so a module-level import
|
||||
# would close the cycle.
|
||||
from ..search.retrieval import retrieve_semantic_bm25_combined_sql
|
||||
@@ -467,7 +322,7 @@ class PostgresMemories(MemoriesExtension):
|
||||
ops,
|
||||
fq_table,
|
||||
bank_id: str,
|
||||
fact_type: str | list[str] | None = None,
|
||||
fact_type: str | None = None,
|
||||
search_query: str | None = None,
|
||||
consolidation_state: str | None = None,
|
||||
state: str | None = None,
|
||||
@@ -622,25 +477,12 @@ 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],
|
||||
txn=None,
|
||||
self, *, conn, ops, fq_table, bank_id: str | None = None, unit_ids: list[Any], entity_ids: list[Any]
|
||||
) -> None:
|
||||
# The join is keyed by global unit id, so bank_id is not needed here. `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.
|
||||
# The join is keyed by global unit id, so bank_id is not needed here.
|
||||
await ops.bulk_insert_unit_entities(conn, fq_table("unit_entities"), unit_ids, entity_ids)
|
||||
|
||||
async def enqueue_relink_victims(
|
||||
@@ -654,25 +496,14 @@ class PostgresMemories(MemoriesExtension):
|
||||
include_affected_units=include_affected_units,
|
||||
)
|
||||
|
||||
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 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 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_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 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)
|
||||
async def prune_stale_cooccurrences(self, *, conn, fq_table, bank_id: str) -> int:
|
||||
return await graph.prune_stale_cooccurrences(conn=conn, fq_table=fq_table, bank_id=bank_id)
|
||||
|
||||
|
||||
__all__ = ["PostgresMemories"]
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -68,22 +68,20 @@ class MentalModelRefreshWindow(BaseModel):
|
||||
created_after: datetime | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"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."
|
||||
"Lower bound on memory creation time. Set only in delta mode, where it is the model's "
|
||||
"last_refreshed_at — so a delta refresh only sees memories newer than the last one."
|
||||
),
|
||||
)
|
||||
created_before: datetime = Field(
|
||||
description=(
|
||||
"Database-time snapshot bounding the refresh. Memories written or edited 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 committed after this are not read, "
|
||||
"so they stay newer than the persisted watermark and are caught by the next refresh."
|
||||
)
|
||||
)
|
||||
watermark: datetime | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"The last_memory_seen_at a real refresh would persist: the newest in-scope memory visible at "
|
||||
"The last_refreshed_at a real refresh would persist: the newest in-scope memory visible at "
|
||||
"the snapshot, not now(). Null means no in-scope memory was visible."
|
||||
),
|
||||
)
|
||||
|
||||
@@ -1,35 +0,0 @@
|
||||
"""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 | None = None,
|
||||
reasoning_effort: str = "low",
|
||||
timeout: float = 300.0,
|
||||
default_headers: dict[str, str] | None = None,
|
||||
extra_body: dict[str, Any] | None = None,
|
||||
@@ -115,7 +115,6 @@ 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,12 +79,11 @@ class ClaudeCodeLLM(LLMInterface):
|
||||
api_key: str, # Will be ignored, uses CLI auth
|
||||
base_url: str,
|
||||
model: str,
|
||||
reasoning_effort: str | None = None,
|
||||
reasoning_effort: str = "low",
|
||||
**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 | None = None,
|
||||
reasoning_effort: str = "low",
|
||||
extra_body: dict[str, Any] | None = None,
|
||||
**kwargs: Any,
|
||||
):
|
||||
@@ -329,26 +329,12 @@ class CodexLLM(LLMInterface):
|
||||
except CodexRefreshExpiredError:
|
||||
raise
|
||||
|
||||
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:
|
||||
def _map_reasoning_effort(self, effort: str) -> str:
|
||||
"""
|
||||
Map standard reasoning effort to Codex reasoning summary format.
|
||||
|
||||
Args:
|
||||
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.
|
||||
effort: Standard effort level ("low", "medium", "high", "xhigh").
|
||||
|
||||
Returns:
|
||||
Codex reasoning summary: "concise", "detailed", or "auto".
|
||||
@@ -359,7 +345,7 @@ class CodexLLM(LLMInterface):
|
||||
"high": "detailed",
|
||||
"xhigh": "detailed",
|
||||
}
|
||||
return mapping.get(effort.lower(), "auto") if effort else "auto"
|
||||
return mapping.get(effort.lower(), "auto")
|
||||
|
||||
async def verify_connection(self) -> None:
|
||||
"""Verify Codex connection by making a simple test call."""
|
||||
@@ -467,7 +453,7 @@ class CodexLLM(LLMInterface):
|
||||
"tools": [],
|
||||
"tool_choice": "auto",
|
||||
"parallel_tool_calls": True,
|
||||
"reasoning": self._reasoning_payload(reasoning_summary),
|
||||
"reasoning": {"effort": self.reasoning_effort, "summary": reasoning_summary},
|
||||
"store": False, # Codex uses stateless mode
|
||||
"stream": True, # SSE streaming
|
||||
"include": ["reasoning.encrypted_content"],
|
||||
@@ -856,7 +842,7 @@ class CodexLLM(LLMInterface):
|
||||
else tool_choice.mode.value
|
||||
),
|
||||
"parallel_tool_calls": True,
|
||||
"reasoning": self._reasoning_payload(reasoning_summary),
|
||||
"reasoning": {"effort": self.reasoning_effort, "summary": 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 | None = None,
|
||||
reasoning_effort: str = "low",
|
||||
account_id: str | None = None,
|
||||
batch_base_url: str | None = None,
|
||||
max_wait_seconds: int | None = None,
|
||||
|
||||
@@ -174,12 +174,11 @@ class GeminiLLM(LLMInterface):
|
||||
api_key: str,
|
||||
base_url: str,
|
||||
model: str,
|
||||
reasoning_effort: str | None = None,
|
||||
reasoning_effort: str = "low",
|
||||
**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"
|
||||
@@ -511,21 +510,6 @@ 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 | None = None,
|
||||
reasoning_effort: str = "low",
|
||||
timeout: float | None = None,
|
||||
extra_body: dict[str, Any] | None = None,
|
||||
bedrock_service_tier: str | None = None,
|
||||
@@ -185,14 +185,6 @@ 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 | None = None,
|
||||
reasoning_effort: str = "low",
|
||||
timeout: float | None = None,
|
||||
**kwargs: Any,
|
||||
):
|
||||
@@ -146,13 +146,6 @@ 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,8 +273,7 @@ class LlamaCppLLM(LLMInterface):
|
||||
api_key: str,
|
||||
base_url: str,
|
||||
model: str,
|
||||
reasoning_effort: str | None = None,
|
||||
extra_body: dict[str, Any] | None = None,
|
||||
reasoning_effort: str = "low",
|
||||
model_path: str | None = None,
|
||||
gpu_layers: int = -1,
|
||||
context_size: int = 8192,
|
||||
@@ -290,7 +289,6 @@ 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
|
||||
@@ -338,10 +336,7 @@ 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 | None = None,
|
||||
reasoning_effort: str = "low",
|
||||
**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 | None = None,
|
||||
reasoning_effort: str = "low",
|
||||
**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); 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.
|
||||
# (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.
|
||||
_REASONING_TAG_PAIRS: tuple[tuple[str, str], ...] = (
|
||||
("<think>", "</think>"),
|
||||
("<thinking>", "</thinking>"),
|
||||
@@ -214,11 +214,7 @@ 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 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.
|
||||
truncated mid-thought) is removed from the open tag to end-of-string.
|
||||
|
||||
Returns the input unchanged (modulo surrounding whitespace) when no tags are
|
||||
present.
|
||||
@@ -228,14 +224,9 @@ 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.
|
||||
# Closed blocks first, then any remaining unclosed (truncated) block.
|
||||
text = re.sub(rf"{open_re}.*?{close_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)
|
||||
text = re.sub(rf"{open_re}.*", "", text, flags=re.DOTALL)
|
||||
return text.strip()
|
||||
|
||||
|
||||
@@ -541,7 +532,7 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
api_key: str,
|
||||
base_url: str,
|
||||
model: str,
|
||||
reasoning_effort: str | None = None,
|
||||
reasoning_effort: str = "low",
|
||||
timeout: float | None = None,
|
||||
groq_service_tier: str | None = None,
|
||||
extra_body: dict[str, Any] | None = None,
|
||||
@@ -693,18 +684,8 @@ 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"reasoning_effort={self.reasoning_effort if self._sends_reasoning_effort() else 'not sent'}"
|
||||
f"base_url={self.base_url or 'default'}"
|
||||
)
|
||||
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}"
|
||||
@@ -746,45 +727,8 @@ 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).
|
||||
|
||||
**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.
|
||||
"""
|
||||
"""Check if the current model is a reasoning model (o1, o3, GPT-5, DeepSeek)."""
|
||||
model_lower = self.model.lower()
|
||||
if "deepseek" in model_lower:
|
||||
# DeepSeek v4-flash is the non-thinking route. Treating every
|
||||
@@ -923,8 +867,8 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
temperature = max(0.01, min(temperature, 1.0))
|
||||
call_params["temperature"] = temperature
|
||||
|
||||
# Set reasoning_effort when configured, or for models recognised as reasoning models
|
||||
if self._sends_reasoning_effort():
|
||||
# Set reasoning_effort for reasoning models
|
||||
if is_reasoning_model:
|
||||
call_params["reasoning_effort"] = self.reasoning_effort
|
||||
|
||||
# Provider-specific parameters
|
||||
@@ -1360,7 +1304,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._sends_reasoning_effort():
|
||||
if self._supports_reasoning_model():
|
||||
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 | None = None,
|
||||
reasoning_effort: str = "low",
|
||||
timeout: float | None = None,
|
||||
extra_body: dict[str, Any] | None = None,
|
||||
default_headers: dict[str, str] | None = None,
|
||||
@@ -455,9 +455,7 @@ class OpenAIResponsesLLM(LLMInterface):
|
||||
params["max_output_tokens"] = max_completion_tokens
|
||||
if temperature is not None and not is_reasoning_model:
|
||||
params["temperature"] = temperature
|
||||
# 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:
|
||||
if is_reasoning_model:
|
||||
params["reasoning"] = {"effort": self.reasoning_effort}
|
||||
if self.openai_service_tier:
|
||||
params["service_tier"] = self.openai_service_tier
|
||||
@@ -573,9 +571,7 @@ class OpenAIResponsesLLM(LLMInterface):
|
||||
params["max_output_tokens"] = max_completion_tokens
|
||||
if temperature is not None and not is_reasoning_model:
|
||||
params["temperature"] = temperature
|
||||
# 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:
|
||||
if is_reasoning_model:
|
||||
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 | None = None,
|
||||
reasoning_effort: str = "low",
|
||||
timeout: float | None = None,
|
||||
auth_manager: XaiOAuthManager | None = None,
|
||||
**kwargs: Any,
|
||||
|
||||
@@ -10,7 +10,6 @@ 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 (
|
||||
@@ -92,39 +91,6 @@ _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.
|
||||
|
||||
@@ -232,76 +198,30 @@ 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
|
||||
self._loaded = False
|
||||
self._locales = None
|
||||
self._exact_search = None
|
||||
|
||||
def _search_kwargs(self) -> dict:
|
||||
"""Extra kwargs for search_dates, shared by load() and analyze().
|
||||
|
||||
Both call sites must use the same locale set: warming up under
|
||||
auto-detection while running restricted (or vice versa) leaves part of
|
||||
the lazy-load cost on the first real query.
|
||||
"""
|
||||
return {} if self._languages is None else {"languages": self._languages}
|
||||
|
||||
def load(self) -> None:
|
||||
"""Load dateparser and warm up internal data structures.
|
||||
|
||||
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.
|
||||
Triggers the real initialization cost (regex tables, timezone data) at
|
||||
load time so the first actual recall doesn't pay the cold-start penalty.
|
||||
"""
|
||||
if self._loaded:
|
||||
return
|
||||
if self._search_dates is None:
|
||||
from dateparser.search import search_dates
|
||||
|
||||
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
|
||||
self._search_dates = search_dates
|
||||
# Warm up: fire a dummy call to trigger lazy-loaded internal tables.
|
||||
self._search_dates("today", **self._search_kwargs())
|
||||
|
||||
def analyze(self, query: str, reference_date: datetime | None = None) -> QueryAnalysis:
|
||||
"""
|
||||
@@ -329,14 +249,6 @@ 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()
|
||||
|
||||
@@ -354,7 +266,7 @@ class DateparserQueryAnalyzer(QueryAnalyzer):
|
||||
# treat any failure as "no temporal constraint found" so the caller
|
||||
# can fall back to non-temporal retrieval.
|
||||
try:
|
||||
results = self._find_dates(query, settings=settings)
|
||||
results = self._search_dates(query, settings=settings, **self._search_kwargs())
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
"dateparser raised %s on query (treating as no temporal constraint): %s",
|
||||
|
||||
@@ -17,15 +17,10 @@ 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
|
||||
@@ -467,12 +462,7 @@ 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: 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``).
|
||||
max_tokens: Maximum tokens for the final response
|
||||
response_schema: Optional JSON Schema for structured output in final response
|
||||
directives: Optional list of directive mental models to inject as hard rules
|
||||
|
||||
@@ -481,13 +471,6 @@ 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)
|
||||
|
||||
@@ -655,106 +638,6 @@ 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
|
||||
@@ -773,7 +656,60 @@ async def _run_reflect_agent_inner(
|
||||
|
||||
if is_last:
|
||||
# Force text response on last iteration - no tools
|
||||
return await _forced_final_synthesis(iteration + 1)
|
||||
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,
|
||||
)
|
||||
|
||||
# Proactive context-window guard: if accumulated messages would exceed the
|
||||
# configured token budget, bail out early and synthesize from what we have.
|
||||
@@ -785,7 +721,59 @@ 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."
|
||||
)
|
||||
return await _forced_final_synthesis(iteration + 1)
|
||||
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,
|
||||
)
|
||||
|
||||
# Call LLM with tools
|
||||
llm_start = time.time()
|
||||
@@ -876,7 +864,60 @@ 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
|
||||
return await _forced_final_synthesis(iteration + 1)
|
||||
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,
|
||||
)
|
||||
|
||||
# No tool calls this turn.
|
||||
if not result.tool_calls:
|
||||
@@ -901,7 +942,60 @@ 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).
|
||||
return await _forced_final_synthesis(iteration + 1)
|
||||
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,
|
||||
)
|
||||
|
||||
# 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
|
||||
@@ -1193,10 +1287,6 @@ 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=[
|
||||
{
|
||||
@@ -1213,7 +1303,7 @@ async def _process_done_tool(
|
||||
},
|
||||
],
|
||||
scope="reflect",
|
||||
max_completion_tokens=get_config().reflect_max_completion_tokens,
|
||||
max_completion_tokens=max_tokens,
|
||||
return_usage=True,
|
||||
)
|
||||
answer = rewritten.strip()
|
||||
|
||||
@@ -443,153 +443,25 @@ def build_system_prompt_for_tools(
|
||||
return "\n".join(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")
|
||||
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 = []
|
||||
|
||||
#: 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."""
|
||||
# Bank identity
|
||||
name = bank_profile.get("name", "Assistant")
|
||||
mission = bank_profile.get("mission", "")
|
||||
|
||||
parts = [f"## Memory Bank Context\nName: {name}"]
|
||||
parts.append(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 = []
|
||||
@@ -602,51 +474,9 @@ def _bank_identity_section(bank_profile: dict, additional_context: str | None) -
|
||||
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).
|
||||
@@ -657,7 +487,13 @@ def build_final_prompt(
|
||||
rendered: list[str] = []
|
||||
truncated = False
|
||||
for entry in reversed(context_history):
|
||||
block = _render_history_block(entry)
|
||||
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_tokens = count_cl100k_tokens(block)
|
||||
if block_tokens > token_budget:
|
||||
truncated = True
|
||||
@@ -675,92 +511,16 @@ 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"
|
||||
"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."
|
||||
"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."
|
||||
)
|
||||
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", "all", "any_strict", "all_strict", or "exact"
|
||||
tags_match: How to match tags - "any" (OR), "all" (AND)
|
||||
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, last_memory_seen_at, trigger,
|
||||
tags, created_at, last_refreshed_at, trigger,
|
||||
1 - (embedding <=> $2::vector) as relevance
|
||||
FROM {fq_table("mental_models")}
|
||||
WHERE bank_id = $1 AND embedding IS NOT NULL {filters}
|
||||
@@ -167,12 +167,6 @@ 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
|
||||
@@ -180,7 +174,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_memory_seen_at, last_memory_write_at):
|
||||
if last_memory_write_at is not None and not _may_need_refresh(last_refreshed_at, last_memory_write_at):
|
||||
is_stale = False
|
||||
else:
|
||||
is_stale = await memory_engine.compute_mental_model_is_stale(conn, bank_id, row)
|
||||
@@ -234,7 +228,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", "all", "any_strict", "all_strict", or "exact"
|
||||
tags_match: How to match tags - "any" (OR), "all" (AND)
|
||||
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)
|
||||
@@ -324,7 +318,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", "all", "any_strict", "all_strict", or "exact"
|
||||
tags_match: How to match tags - "any" (OR), "all" (AND), 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,8 +10,6 @@ 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"])
|
||||
|
||||
|
||||
@@ -262,9 +260,6 @@ 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
|
||||
@@ -273,7 +268,7 @@ class MemoryFact(BaseModel):
|
||||
|
||||
v = json.loads(v)
|
||||
if isinstance(v, dict):
|
||||
return as_string_metadata(v)
|
||||
return {str(k): str(val) for k, val in v.items()}
|
||||
return v
|
||||
|
||||
chunk_id: str | None = Field(
|
||||
@@ -358,14 +353,6 @@ 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,6 +45,36 @@ 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.
|
||||
|
||||
@@ -160,12 +190,12 @@ async def get_or_create_bank_profile(pool, bank_id: str) -> BankProfileResult:
|
||||
``get_or_create_bank_profile_on_conn`` instead.
|
||||
"""
|
||||
|
||||
# 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.
|
||||
# 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.
|
||||
async def _create() -> BankProfileResult:
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
async with conn.transaction():
|
||||
@@ -212,16 +242,10 @@ async def get_or_create_bank_profile_on_conn(conn, bank_id: str, *, ops) -> Bank
|
||||
created=False,
|
||||
)
|
||||
|
||||
# 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.
|
||||
# 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()
|
||||
inserted = await conn.fetchval(
|
||||
f"""
|
||||
INSERT INTO {fq_table("banks")} (bank_id, name, disposition, mission, internal_id)
|
||||
@@ -233,10 +257,14 @@ 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),
|
||||
"",
|
||||
uuid.uuid4(),
|
||||
internal_id,
|
||||
)
|
||||
|
||||
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,
|
||||
@@ -396,9 +424,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, *, search_query: str | None = None) -> list:
|
||||
async def list_banks(pool) -> list:
|
||||
"""
|
||||
List banks with summary stats, optionally narrowed by a search string.
|
||||
List all banks in the system with summary stats.
|
||||
|
||||
``last_document_at`` is document *ingestion* time (when a document first
|
||||
landed), while ``last_write_at`` is the last time anything was written to
|
||||
@@ -406,14 +434,8 @@ async def list_banks(pool, *, search_query: str | None = None) -> 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),
|
||||
@@ -423,15 +445,6 @@ async def list_banks(pool, *, search_query: str | None = None) -> 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"""
|
||||
@@ -457,16 +470,19 @@ async def list_banks(pool, *, search_query: str | None = None) -> 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"]
|
||||
@@ -482,6 +498,12 @@ async def list_banks(pool, *, search_query: str | None = None) -> 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(
|
||||
{
|
||||
@@ -491,7 +513,7 @@ async def list_banks(pool, *, search_query: str | None = None) -> 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": row["fact_count"],
|
||||
"fact_count": fact_count,
|
||||
"last_document_at": last_doc.isoformat() if last_doc else None,
|
||||
"last_write_at": last_write.isoformat() if last_write else None,
|
||||
}
|
||||
@@ -499,23 +521,3 @@ async def list_banks(pool, *, search_query: str | None = None) -> 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,26 +136,9 @@ 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
|
||||
@@ -172,19 +155,6 @@ 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 (
|
||||
@@ -192,19 +162,14 @@ 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
|
||||
JOIN matched_links ON ml.ctid = matched_links.link_ctid
|
||||
WHERE EXISTS (
|
||||
SELECT 1
|
||||
FROM target_units tu
|
||||
WHERE tu.id = ml.from_unit_id OR tu.id = ml.to_unit_id
|
||||
)
|
||||
ORDER BY
|
||||
LEAST(ml.from_unit_id, ml.to_unit_id),
|
||||
GREATEST(ml.from_unit_id, ml.to_unit_id),
|
||||
|
||||
@@ -7,23 +7,18 @@ Handles entity extraction and resolution for stored facts.
|
||||
import logging
|
||||
|
||||
from . import link_utils
|
||||
from .types import EntityResolutionResult, ProcessedFact, UserEntities
|
||||
from .types import EntityResolutionResult, ProcessedFact
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _prepare_facts_for_entity_processing(
|
||||
facts: list[ProcessedFact],
|
||||
user_entities_per_content: dict[int, UserEntities] | None = None,
|
||||
user_entities_per_content: dict[int, list[dict]] | 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)
|
||||
"""
|
||||
@@ -34,29 +29,20 @@ def _prepare_facts_for_entity_processing(
|
||||
|
||||
entities_per_fact = []
|
||||
for fact in facts:
|
||||
llm_entities = [{"text": entity.name, "type": "CONCEPT", "resolve": True} for entity in (fact.entities or [])]
|
||||
llm_entities = [{"text": entity.name, "type": "CONCEPT"} for entity in (fact.entities or [])]
|
||||
|
||||
supplied = user_entities_per_content.get(fact.content_index)
|
||||
user_entities = supplied.entities if supplied else []
|
||||
user_resolve = supplied.resolve if supplied else True
|
||||
user_entities = user_entities_per_content.get(fact.content_index, [])
|
||||
|
||||
by_text = {e["text"].lower(): e for e in llm_entities}
|
||||
seen_texts = {e["text"].lower() for e in llm_entities}
|
||||
for user_entity in user_entities:
|
||||
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
|
||||
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())
|
||||
|
||||
entities_per_fact.append(llm_entities)
|
||||
|
||||
@@ -70,7 +56,7 @@ async def resolve_entities(
|
||||
unit_ids: list[str],
|
||||
facts: list[ProcessedFact],
|
||||
log_buffer: list[str] = None,
|
||||
user_entities_per_content: dict[int, UserEntities] | None = None,
|
||||
user_entities_per_content: dict[int, list[dict]] = None,
|
||||
entity_labels: list | None = None,
|
||||
) -> EntityResolutionResult:
|
||||
"""
|
||||
@@ -86,8 +72,7 @@ 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 the caller-supplied
|
||||
entities for that content item and whether to resolve them
|
||||
user_entities_per_content: Dict mapping content_index to user-provided entities
|
||||
entity_labels: Optional entity label taxonomy
|
||||
|
||||
Returns:
|
||||
|
||||
@@ -12,8 +12,7 @@ from typing import Any
|
||||
|
||||
from ...config import _get_raw_config
|
||||
from ..memory_engine import fq_table
|
||||
from ..metadata_utils import drop_null_values
|
||||
from .bank_utils import DEFAULT_DISPOSITION
|
||||
from .bank_utils import DEFAULT_DISPOSITION, create_bank_vector_indexes
|
||||
from .fact_extraction import _sanitize_text
|
||||
from .types import ProcessedFact
|
||||
|
||||
@@ -132,26 +131,25 @@ async def ensure_bank_exists(conn, bank_id: str, ops=None) -> None:
|
||||
conn: Database connection
|
||||
bank_id: Bank identifier
|
||||
"""
|
||||
# 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(
|
||||
# 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(
|
||||
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),
|
||||
"",
|
||||
uuid.uuid4(),
|
||||
internal_id,
|
||||
)
|
||||
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(
|
||||
@@ -286,16 +284,9 @@ 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_entity_prune_candidates, enqueue_relink_victims
|
||||
from ..graph_maintenance import enqueue_relink_victims
|
||||
|
||||
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)
|
||||
await enqueue_relink_victims(conn, bank_id, [str(uid) for uid in existing_unit_ids])
|
||||
|
||||
# Explicitly delete memory_units by document_id BEFORE deleting the
|
||||
# document row. The CASCADE from documents→chunks→memory_units only
|
||||
@@ -425,12 +416,6 @@ 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.
|
||||
"""
|
||||
@@ -457,7 +442,7 @@ async def update_memory_units_metadata_and_tags(
|
||||
bank_id,
|
||||
document_id,
|
||||
tags or [],
|
||||
json.dumps(drop_null_values(metadata)),
|
||||
json.dumps(metadata or {}),
|
||||
)
|
||||
# result is a status string like "UPDATE 5"
|
||||
try:
|
||||
|
||||
@@ -1,228 +0,0 @@
|
||||
"""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
|
||||
@@ -17,7 +17,6 @@ from ..causal_links import (
|
||||
)
|
||||
from ..db.base import DatabaseConnection
|
||||
from ..db.ops import DataAccessOps
|
||||
from ..db.postgresql import setting_rejected_by_server
|
||||
from ..memory_engine import fq_table
|
||||
from .types import CausalRelation, EntityResolutionResult
|
||||
|
||||
@@ -44,16 +43,6 @@ 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.
|
||||
@@ -90,26 +79,6 @@ 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],
|
||||
@@ -120,9 +89,8 @@ async def _bulk_insert_links(
|
||||
) -> None:
|
||||
"""Bulk-insert links using sorted INSERT FROM unnest().
|
||||
|
||||
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`.
|
||||
Sorting by (from_unit_id, to_unit_id) ensures all concurrent transactions
|
||||
acquire index locks in the same order, eliminating circular-wait deadlocks.
|
||||
|
||||
Args:
|
||||
conn: Database connection (must be inside a transaction).
|
||||
@@ -138,9 +106,9 @@ async def _bulk_insert_links(
|
||||
if not links:
|
||||
return
|
||||
|
||||
# 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)
|
||||
# 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])))
|
||||
|
||||
exists_clause = ""
|
||||
if not skip_exists_check:
|
||||
@@ -224,7 +192,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: dict[str, dict] = {}
|
||||
seen_in_fact: set[str] = set()
|
||||
for ent in entity_list:
|
||||
if hasattr(ent, "text"):
|
||||
raw_text, entity_type = ent.text, "CONCEPT"
|
||||
@@ -241,19 +209,11 @@ def _prepare_entities_for_resolution(
|
||||
dropped_empty += 1
|
||||
continue
|
||||
|
||||
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
|
||||
if normalized_text.lower() in seen_in_fact:
|
||||
continue
|
||||
seen_in_fact.add(normalized_text.lower())
|
||||
|
||||
entity = {"text": normalized_text, "type": entity_type, "resolve": resolve}
|
||||
seen_in_fact[normalized_text.lower()] = entity
|
||||
formatted_entities.append(entity)
|
||||
formatted_entities.append({"text": normalized_text, "type": entity_type})
|
||||
all_entities.append(formatted_entities)
|
||||
|
||||
if dropped_empty:
|
||||
@@ -282,7 +242,6 @@ def _prepare_entities_for_resolution(
|
||||
{
|
||||
"text": entity["text"],
|
||||
"type": entity["type"],
|
||||
"resolve": entity["resolve"],
|
||||
"nearby_entities": entities,
|
||||
}
|
||||
)
|
||||
@@ -594,14 +553,7 @@ async def compute_semantic_links_ann(
|
||||
# are safe to apply at session/transaction scope for the configured
|
||||
# backend. VectorChord probe values are index-shaped, so vchordrq uses
|
||||
# index storage fallback parameters instead of a blanket SET LOCAL.
|
||||
#
|
||||
# A GUC the server has already rejected is skipped rather than attempted:
|
||||
# hnsw.iterative_scan needs pgvector 0.8+, and pgvector reserves the "hnsw."
|
||||
# prefix, so an older server errors on it — which inside this transaction would
|
||||
# abort the whole link computation rather than merely fail to apply.
|
||||
for guc, value in ann_search_tuning_settings(configured_vector_extension(), kind="low_latency"):
|
||||
if setting_rejected_by_server(guc):
|
||||
continue
|
||||
await conn.execute(f"SET LOCAL {guc} = {value}")
|
||||
|
||||
t_setup = time_mod.time()
|
||||
|
||||
@@ -275,23 +275,16 @@ 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]
|
||||
|
||||
@@ -404,11 +397,7 @@ async def _pre_resolve_phase1(
|
||||
set_stage("retain.phase1.resolve")
|
||||
from .link_utils import compute_semantic_links_ann
|
||||
|
||||
user_entities_per_content = {
|
||||
idx: UserEntities(entities=content.entities, resolve=content.resolve_entities)
|
||||
for idx, content in enumerate(contents)
|
||||
if content.entities
|
||||
}
|
||||
user_entities_per_content = {idx: content.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,
|
||||
@@ -590,367 +579,6 @@ 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,
|
||||
@@ -1126,15 +754,6 @@ 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)
|
||||
|
||||
@@ -1354,23 +973,9 @@ async def retain_batch(
|
||||
update_mode = item_mode
|
||||
break
|
||||
|
||||
# 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:
|
||||
if update_mode == "append" and effective_doc_id and is_first_batch:
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
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
|
||||
existing_text = await fact_storage.get_document_content(conn, bank_id, effective_doc_id)
|
||||
if existing_text:
|
||||
# Prepend existing text as a new content item at the beginning
|
||||
existing_content: RetainContentDict = {"content": existing_text}
|
||||
@@ -1436,20 +1041,6 @@ 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 "
|
||||
@@ -1483,7 +1074,6 @@ 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
|
||||
@@ -1548,7 +1138,6 @@ 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,
|
||||
)
|
||||
|
||||
|
||||
@@ -1742,7 +1331,6 @@ 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.
|
||||
@@ -1875,34 +1463,6 @@ 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
|
||||
@@ -1916,7 +1476,6 @@ 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,
|
||||
)
|
||||
@@ -2119,17 +1678,19 @@ async def _streaming_retain_batch(
|
||||
_edge_txn = None
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
async with conn.transaction():
|
||||
# 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"),
|
||||
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",
|
||||
effective_doc_id,
|
||||
bank_id,
|
||||
)
|
||||
_assert_append_base_unchanged(existing_hash)
|
||||
if is_recovery:
|
||||
await fact_storage.upsert_document_metadata(
|
||||
conn,
|
||||
@@ -2158,10 +1719,6 @@ 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
|
||||
@@ -2226,65 +1783,6 @@ 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 ---
|
||||
@@ -2303,12 +1801,6 @@ 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.
|
||||
@@ -2364,21 +1856,12 @@ 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
|
||||
@@ -2426,11 +1909,6 @@ 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)
|
||||
@@ -2609,10 +2087,6 @@ 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.
|
||||
@@ -2789,7 +2263,6 @@ 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
|
||||
@@ -2834,19 +2307,6 @@ 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
|
||||
|
||||
@@ -3028,14 +2488,7 @@ async def _try_delta_retain(
|
||||
result_unit_ids: list[list[str]] = []
|
||||
log_buffer_pre_db = len(log_buffer)
|
||||
|
||||
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.
|
||||
"""
|
||||
async def _run_delta_db_work() -> None:
|
||||
nonlocal result_unit_ids
|
||||
del log_buffer[log_buffer_pre_db:]
|
||||
for pf in processed_facts:
|
||||
@@ -3048,51 +2501,6 @@ 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
|
||||
@@ -3113,9 +2521,8 @@ 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")
|
||||
# Fall back to streaming, which re-locks the document and (for
|
||||
# an append) verifies the base this content was built on.
|
||||
return False
|
||||
# Return None to fall back to streaming (which has full FOR UPDATE protection)
|
||||
return None
|
||||
|
||||
# Update document metadata (no delete)
|
||||
step_start = time.time()
|
||||
@@ -3249,12 +2656,6 @@ 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).
|
||||
@@ -3278,18 +2679,11 @@ 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:
|
||||
delta_committed = await _run_delta_db_work()
|
||||
await _run_delta_db_work()
|
||||
else:
|
||||
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 _run_delta_db_work()
|
||||
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
|
||||
@@ -3389,7 +2783,6 @@ 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"),
|
||||
)
|
||||
@@ -3451,7 +2844,6 @@ 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,8 +11,6 @@ from datetime import datetime
|
||||
from typing import Literal, TypedDict
|
||||
from uuid import UUID
|
||||
|
||||
from ..metadata_utils import drop_null_values
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@@ -26,8 +24,6 @@ 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
|
||||
@@ -45,7 +41,6 @@ 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]]
|
||||
@@ -53,18 +48,6 @@ 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:
|
||||
"""
|
||||
@@ -78,23 +61,11 @@ 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:
|
||||
@@ -377,19 +348,3 @@ 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.
|
||||
"""
|
||||
|
||||
@@ -1,111 +0,0 @@
|
||||
"""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,25 +64,23 @@ 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)
|
||||
updated_range_clause = ""
|
||||
updated_range_params: list[Any] = []
|
||||
created_range_clause = ""
|
||||
created_range_params: list[Any] = []
|
||||
if created_after is not None:
|
||||
updated_range_params.append(created_after)
|
||||
updated_range_clause += f" AND updated_at > ${_next_idx}"
|
||||
created_range_params.append(created_after)
|
||||
created_range_clause += f" AND updated_at > ${_next_idx}"
|
||||
_next_idx += 1
|
||||
if created_before is not None:
|
||||
updated_range_params.append(created_before)
|
||||
updated_range_clause += f" AND updated_at < ${_next_idx}"
|
||||
created_range_params.append(created_before)
|
||||
created_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(updated_range_params)
|
||||
params.extend(created_range_params)
|
||||
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
@@ -96,7 +94,7 @@ async def _find_semantic_seeds(
|
||||
AND (1 - (embedding <=> $1::vector)) >= $4
|
||||
{tags_clause}
|
||||
{groups_clause}
|
||||
{updated_range_clause}
|
||||
{created_range_clause}
|
||||
ORDER BY embedding <=> $1::vector
|
||||
LIMIT $5
|
||||
""",
|
||||
|
||||
@@ -8,17 +8,18 @@ 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, get_config
|
||||
from ...config import DEFAULT_BM25_MAX_QUERY_TERMS, DEFAULT_TEMPORAL_SEMANTIC_MIN_SIMILARITY, get_config
|
||||
from ..db.ops import UpdatedWindow
|
||||
from ..memory_engine import fq_table, get_current_schema
|
||||
from ..db_utils import acquire_with_retry
|
||||
from ..memory_engine import fq_table
|
||||
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
|
||||
@@ -119,6 +120,47 @@ 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,
|
||||
@@ -147,14 +189,10 @@ async def retrieve_semantic_bm25_combined_sql(
|
||||
idx_mu_emb_observation, idx_mu_emb_experience), created automatically by
|
||||
Alembic migration a3b4c5d6e7f8_add_partial_hnsw_indexes.py.
|
||||
|
||||
Each semantic arm asks for exactly ``limit`` rows. It used to ask for ``limit * 5``
|
||||
and trim back to ``limit`` in Python "to compensate for HNSW approximation", but that
|
||||
could never work: the rows arrive already ordered by distance within their arm, so
|
||||
keeping the first ``limit`` of ``limit * 5`` returns precisely what ``LIMIT limit``
|
||||
would have — the extra rows were fetched, decoded and dropped, unread. What actually
|
||||
governs ANN quality is the size of the candidate list the scan explores, which is a
|
||||
connection setting, not a row count; the caller sizes it for this query (see
|
||||
``PostgresMemories.search``) rather than over-fetching rows here.
|
||||
HNSW is approximate — semantic arms over-fetch by 5x (min 100) and trim to
|
||||
limit in Python to compensate. ef_search=200 is set globally on pool
|
||||
connections at init time (see memory_engine.py) to improve recall on sparse
|
||||
graphs.
|
||||
|
||||
fact_type values are inlined as literals (safe: they come from a controlled
|
||||
internal enum, never from user input).
|
||||
@@ -184,17 +222,8 @@ async def retrieve_semantic_bm25_combined_sql(
|
||||
sem_min = min_semantic if min_semantic is not None else config.semantic_min_similarity
|
||||
bm25_min = min_keyword if min_keyword is not None else config.bm25_min_score
|
||||
|
||||
# How many semantic rows each arm must return. Two consumers read them: the semantic
|
||||
# list itself (``limit``), and — when the dense rows also clear the graph arm's
|
||||
# threshold — its entry points (``GRAPH_SEED_LIMIT``), derived from the same ordered
|
||||
# rows instead of a duplicate ANN query per fact type. A budget below GRAPH_SEED_LIMIT
|
||||
# would otherwise starve the graph arm of seeds.
|
||||
graph_seed_threshold = (
|
||||
graph_seed_min_similarity
|
||||
if graph_seed_min_similarity is not None and sem_min <= graph_seed_min_similarity
|
||||
else None
|
||||
)
|
||||
semantic_fetch = max(limit, GRAPH_SEED_LIMIT if graph_seed_threshold is not None else 0)
|
||||
# Over-fetch for HNSW approximation; semantic results trimmed to limit in Python.
|
||||
hnsw_fetch = max(limit * 5, 100)
|
||||
|
||||
cols = (
|
||||
"id, text, context, event_date, occurred_start, occurred_end, mentioned_at, "
|
||||
@@ -212,7 +241,7 @@ async def retrieve_semantic_bm25_combined_sql(
|
||||
# $1 = query_emb_str (semantic arms)
|
||||
# $2 = bank_id
|
||||
# When tokens present:
|
||||
# $3 = limit (BM25 LIMIT; semantic inlines the same limit as a literal)
|
||||
# $3 = limit (BM25 LIMIT; semantic uses inlined hnsw_fetch literal)
|
||||
# $4 = bm25_text
|
||||
# $5 = tags (if present)
|
||||
# $6+ = tag_groups params (one per leaf)
|
||||
@@ -227,22 +256,19 @@ 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_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.
|
||||
# --- created_at time range filter (appended after tags/groups) ---
|
||||
# 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)
|
||||
updated_range_clause = ""
|
||||
updated_range_params: list[Any] = []
|
||||
created_range_clause = ""
|
||||
created_range_params: list[Any] = []
|
||||
if created_after is not None:
|
||||
updated_range_params.append(created_after)
|
||||
updated_range_clause += f" AND updated_at > ${_next_idx}"
|
||||
created_range_params.append(created_after)
|
||||
created_range_clause += f" AND updated_at > ${_next_idx}"
|
||||
_next_idx += 1
|
||||
if created_before is not None:
|
||||
updated_range_params.append(created_before)
|
||||
updated_range_clause += f" AND updated_at < ${_next_idx}"
|
||||
created_range_params.append(created_before)
|
||||
created_range_clause += f" AND updated_at < ${_next_idx}"
|
||||
_next_idx += 1
|
||||
|
||||
# --- Semantic UNION ALL arms (one per fact_type) ---
|
||||
@@ -255,11 +281,11 @@ async def retrieve_semantic_bm25_combined_sql(
|
||||
fact_type=ft,
|
||||
embedding_param="$1",
|
||||
bank_id_param="$2",
|
||||
fetch_limit=semantic_fetch,
|
||||
fetch_limit=hnsw_fetch,
|
||||
min_similarity=sem_min,
|
||||
tags_clause=tags_clause,
|
||||
groups_clause=groups_clause,
|
||||
extra_where=updated_range_clause,
|
||||
extra_where=created_range_clause,
|
||||
)
|
||||
for ft in fact_types
|
||||
]
|
||||
@@ -267,36 +293,11 @@ 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(
|
||||
bm25_tokens,
|
||||
tokens,
|
||||
query_text,
|
||||
text_search_extension=text_ext,
|
||||
max_query_terms=max_query_terms,
|
||||
max_query_terms=getattr(config, "bm25_max_query_terms", DEFAULT_BM25_MAX_QUERY_TERMS),
|
||||
)
|
||||
for i, ft in enumerate(fact_types):
|
||||
arms.append(
|
||||
@@ -313,7 +314,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=updated_range_clause,
|
||||
extra_where=created_range_clause,
|
||||
)
|
||||
)
|
||||
|
||||
@@ -326,7 +327,7 @@ async def retrieve_semantic_bm25_combined_sql(
|
||||
if tags:
|
||||
params.append(tags)
|
||||
params.extend(groups_params)
|
||||
params.extend(updated_range_params)
|
||||
params.extend(created_range_params)
|
||||
|
||||
try:
|
||||
rows = await conn.fetch(query, *params)
|
||||
@@ -345,12 +346,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_updated_clause = ""
|
||||
fb_created_clause = ""
|
||||
if created_after is not None:
|
||||
fb_updated_clause += f" AND updated_at > ${fb_next_idx}"
|
||||
fb_created_clause += f" AND updated_at > ${fb_next_idx}"
|
||||
fb_next_idx += 1
|
||||
if created_before is not None:
|
||||
fb_updated_clause += f" AND updated_at < ${fb_next_idx}"
|
||||
fb_created_clause += f" AND updated_at < ${fb_next_idx}"
|
||||
fb_next_idx += 1
|
||||
fb_arms = [
|
||||
dialect.build_semantic_arm(
|
||||
@@ -359,11 +360,11 @@ async def retrieve_semantic_bm25_combined_sql(
|
||||
fact_type=ft,
|
||||
embedding_param="$1",
|
||||
bank_id_param="$2",
|
||||
fetch_limit=semantic_fetch,
|
||||
fetch_limit=hnsw_fetch,
|
||||
min_similarity=sem_min,
|
||||
tags_clause=fb_tags_clause,
|
||||
groups_clause=fb_groups_clause,
|
||||
extra_where=fb_updated_clause,
|
||||
extra_where=fb_created_clause,
|
||||
)
|
||||
for ft in fact_types
|
||||
]
|
||||
@@ -372,12 +373,22 @@ async def retrieve_semantic_bm25_combined_sql(
|
||||
if tags:
|
||||
fb_params.append(tags)
|
||||
fb_params.extend(groups_params)
|
||||
fb_params.extend(updated_range_params)
|
||||
fb_params.extend(created_range_params)
|
||||
rows = await conn.fetch(fb_query, *fb_params)
|
||||
else:
|
||||
raise
|
||||
|
||||
# Group results, converting only the prefix either consumer can observe.
|
||||
# Group results. The semantic SQL deliberately over-fetches for HNSW recall;
|
||||
# when that pool also covers the graph threshold, derive graph entry points
|
||||
# from the same ordered rows instead of issuing one duplicate ANN query per
|
||||
# fact type. Convert only the prefix either consumer can observe, not the
|
||||
# entire HNSW over-fetch pool.
|
||||
graph_seed_threshold = (
|
||||
graph_seed_min_similarity
|
||||
if graph_seed_min_similarity is not None and sem_min <= graph_seed_min_similarity
|
||||
else None
|
||||
)
|
||||
semantic_candidate_limit = max(limit, GRAPH_SEED_LIMIT if graph_seed_threshold is not None else 0)
|
||||
semantic_candidates: dict[str, list[RetrievalResult]] = {ft: [] for ft in fact_types}
|
||||
for r in rows:
|
||||
row = dict(r)
|
||||
@@ -386,7 +397,7 @@ async def retrieve_semantic_bm25_combined_sql(
|
||||
if ft not in result_dict:
|
||||
continue
|
||||
if source == "semantic":
|
||||
if len(semantic_candidates[ft]) < semantic_fetch:
|
||||
if len(semantic_candidates[ft]) < semantic_candidate_limit:
|
||||
semantic_candidates[ft].append(RetrievalResult.from_db_row(row))
|
||||
else:
|
||||
result_dict[ft].bm25.append(RetrievalResult.from_db_row(row))
|
||||
@@ -463,6 +474,45 @@ 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,
|
||||
@@ -509,30 +559,29 @@ 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_sql.
|
||||
# the backend on execute, but `unnest` is not). Mirrors retrieve_semantic_bm25_combined.
|
||||
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_after/created_before time range filter (after tags/groups) — filters
|
||||
# `updated_at`, as above.
|
||||
# created_at time range filter (after tags/groups)
|
||||
_next_idx = tag_groups_param_start + len(groups_params)
|
||||
updated_range_clause = ""
|
||||
updated_range_params: list[Any] = []
|
||||
created_range_clause = ""
|
||||
created_range_params: list[Any] = []
|
||||
if created_after is not None:
|
||||
updated_range_params.append(created_after)
|
||||
updated_range_clause += f" AND updated_at > ${_next_idx}"
|
||||
created_range_params.append(created_after)
|
||||
created_range_clause += f" AND updated_at > ${_next_idx}"
|
||||
_next_idx += 1
|
||||
if created_before is not None:
|
||||
updated_range_params.append(created_before)
|
||||
updated_range_clause += f" AND updated_at < ${_next_idx}"
|
||||
created_range_params.append(created_before)
|
||||
created_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(updated_range_params)
|
||||
params.extend(created_range_params)
|
||||
|
||||
# Entry-point selection: similarity-gated, window-filtered, then narrowed for coverage.
|
||||
#
|
||||
@@ -559,7 +608,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_sql; this keeps the query free of `unnest`/LATERAL, which the
|
||||
# retrieve_semantic_bm25_combined; 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, "
|
||||
@@ -586,7 +635,7 @@ async def retrieve_temporal_combined_sql(
|
||||
AND (1 - (embedding <=> $1::vector)) >= $5
|
||||
{tags_clause}
|
||||
{groups_clause}
|
||||
{updated_range_clause}
|
||||
{created_range_clause}
|
||||
ORDER BY embedding <=> $1::vector
|
||||
LIMIT {_TEMPORAL_POOL_SIZE}
|
||||
)"""
|
||||
@@ -802,6 +851,7 @@ 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,
|
||||
@@ -813,16 +863,15 @@ async def retrieve_all_fact_types_parallel(
|
||||
enable_graph_retrieval: bool = True,
|
||||
) -> MultiFactTypeRetrievalResult:
|
||||
"""
|
||||
Retrieve every recall arm for all fact types, through the memories store.
|
||||
Optimized retrieval for multiple fact types using batched queries.
|
||||
|
||||
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.
|
||||
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)
|
||||
|
||||
Args:
|
||||
pool: Database connection pool, handed to the store as its connection handle.
|
||||
pool: Database connection pool
|
||||
query_text: Query text
|
||||
query_embedding_str: Query embedding as string
|
||||
bank_id: Bank ID
|
||||
@@ -830,6 +879,7 @@ 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
|
||||
@@ -840,72 +890,155 @@ 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 the store call so we know whether the temporal arm is needed at all.
|
||||
# Do this before DB queries so we know if we need temporal retrieval
|
||||
temporal_extraction_start = time.time()
|
||||
temporal_constraint = None
|
||||
if enable_temporal_retrieval:
|
||||
from .temporal_extraction import extract_temporal_constraint_async
|
||||
from .temporal_extraction import 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(
|
||||
temporal_constraint = extract_temporal_constraint(
|
||||
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 every arm for every fact type through the store's single recall method.
|
||||
from ..memories import RecallArms, get_memories
|
||||
# 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
|
||||
|
||||
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,
|
||||
)
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
conn_wait = time.time() - semantic_bm25_start
|
||||
|
||||
# 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:
|
||||
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
|
||||
# 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
|
||||
|
||||
results_by_fact_type[ft] = ParallelRetrievalResult(
|
||||
semantic=arms.semantic,
|
||||
bm25=arms.bm25,
|
||||
graph=arms.graph,
|
||||
temporal=temporal_arm,
|
||||
semantic=semantic_results,
|
||||
bm25=bm25_results,
|
||||
graph=graph_results,
|
||||
temporal=temporal_results,
|
||||
timings={
|
||||
"semantic": 0.0,
|
||||
"bm25": 0.0,
|
||||
"graph": 0.0,
|
||||
"temporal": 0.0,
|
||||
"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)
|
||||
"temporal_extraction": temporal_extraction_time,
|
||||
},
|
||||
temporal_constraint=temporal_constraint,
|
||||
graph_timings=[],
|
||||
max_conn_wait=0.0,
|
||||
graph_timings=[graph_timing] if graph_timing else [],
|
||||
max_conn_wait=max_conn_wait,
|
||||
)
|
||||
|
||||
timings["total"] = time.time() - start_time
|
||||
total_time = time.time() - start_time
|
||||
timings["total"] = total_time
|
||||
|
||||
return MultiFactTypeRetrievalResult(
|
||||
results_by_fact_type=results_by_fact_type,
|
||||
timings=timings,
|
||||
max_conn_wait=0.0,
|
||||
max_conn_wait=max_conn_wait,
|
||||
)
|
||||
|
||||
@@ -4,58 +4,13 @@ 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
|
||||
@@ -97,48 +52,10 @@ def extract_temporal_constraint(
|
||||
if analyzer is None:
|
||||
analyzer = get_default_analyzer()
|
||||
|
||||
# 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
|
||||
analysis = analyzer.analyze(query, reference_date)
|
||||
|
||||
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,22 +50,6 @@ 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
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user